From 82319af422b8495cb338754e0e84126db1ee4af9 Mon Sep 17 00:00:00 2001 From: Antonio Bellotta Date: Thu, 1 Feb 2024 15:49:41 +0100 Subject: [PATCH] Added tests for pickle import --- neurodamus/utils/memory.py | 15 +++++++++++++++ tests/integration-e2e/test_dry_run_worflow.py | 10 ++++++++++ 2 files changed, 25 insertions(+) diff --git a/neurodamus/utils/memory.py b/neurodamus/utils/memory.py index c0ece6643..c850abf96 100644 --- a/neurodamus/utils/memory.py +++ b/neurodamus/utils/memory.py @@ -299,6 +299,21 @@ def export_allocation_stats(rank_allocation, filename): logging.warning("Unable to export allocation stats: {}".format(e)) +@run_only_rank0 +def import_allocation_stats(filename): + """ + Import allocation dictionary from serialized pickle file. + """ + import pickle + try: + with open(filename, 'rb') as f: + rank_allocation = pickle.load(f) + return rank_allocation + except Exception as e: + logging.warning("Unable to import allocation stats: {}".format(e)) + return None + + class SynapseMemoryUsage: ''' A small class that works as a lookup table for the memory used by each type of synapse. diff --git a/tests/integration-e2e/test_dry_run_worflow.py b/tests/integration-e2e/test_dry_run_worflow.py index cdd217b71..420b1db68 100644 --- a/tests/integration-e2e/test_dry_run_worflow.py +++ b/tests/integration-e2e/test_dry_run_worflow.py @@ -23,3 +23,13 @@ def test_dry_run_workflow(USECASE3): } assert nd._dry_run_stats.metype_counts == expected_items assert nd._dry_run_stats.suggest_nodes(0.3) > 0 + + # Test that the dry run mode works with a pickle file + import pickle + try: + with open((str(USECASE3 / "allocation.bin")), 'rb') as f: + rank_allocation = pickle.load(f) + except Exception as e: + print("Unable to import allocation stats: {}".format(e)) + expected_items = {'NodeA': {0: [3]}, 'NodeB': {1: [2]}} + assert rank_allocation == expected_items