diff --git a/executorlib/cache/backend.py b/executorlib/cache/backend.py index fb78fae7..0e0bf779 100644 --- a/executorlib/cache/backend.py +++ b/executorlib/cache/backend.py @@ -46,7 +46,7 @@ def backend_write_file(file_name: str, output: Any, runtime: float) -> None: os.rename(file_name, file_name_out + ".h5ready") dump( file_name=file_name_out + ".h5ready", - data_dict={"output": output, "time": runtime}, + data_dict={"output": output, "runtime": runtime}, ) os.rename(file_name_out + ".h5ready", file_name_out + ".h5out") diff --git a/executorlib/standalone/hdf.py b/executorlib/standalone/hdf.py index 06048f6b..c0dd0609 100644 --- a/executorlib/standalone/hdf.py +++ b/executorlib/standalone/hdf.py @@ -18,6 +18,7 @@ def dump(file_name: str, data_dict: dict) -> None: "args": "input_args", "kwargs": "input_kwargs", "output": "output", + "runtime": "runtime", "queue_id": "queue_id", } with h5py.File(file_name, "a") as fname: @@ -73,6 +74,23 @@ def get_output(file_name: str) -> Tuple[bool, object]: return False, None +def get_runtime(file_name: str) -> float: + """ + Get run time from HDF5 file + + Args: + file_name (str): file name of the HDF5 file as absolute path + + Returns: + float: run time from the execution of the python function + """ + with h5py.File(file_name, "r") as hdf: + if "runtime" in hdf: + return cloudpickle.loads(np.void(hdf["/runtime"])) + else: + return 0.0 + + def get_queue_id(file_name: str) -> Optional[int]: with h5py.File(file_name, "r") as hdf: if "queue_id" in hdf: diff --git a/tests/test_cache_shared.py b/tests/test_cache_shared.py index 544eb73f..5fef76af 100644 --- a/tests/test_cache_shared.py +++ b/tests/test_cache_shared.py @@ -7,7 +7,7 @@ try: from executorlib.cache.backend import backend_execute_task_in_file from executorlib.cache.shared import _check_task_output, FutureItem - from executorlib.standalone.hdf import dump + from executorlib.standalone.hdf import dump, get_runtime from executorlib.standalone.serialize import serialize_funct_h5 skip_h5io_test = False @@ -40,6 +40,7 @@ def test_execute_function_mixed(self): ) self.assertTrue(future_obj.done()) self.assertEqual(future_obj.result(), 3) + self.assertTrue(get_runtime(file_name=os.path.join(cache_directory, task_key + ".h5out")) > 0.0) future_file_obj = FutureItem( file_name=os.path.join(cache_directory, task_key + ".h5out") ) @@ -63,6 +64,7 @@ def test_execute_function_args(self): ) self.assertTrue(future_obj.done()) self.assertEqual(future_obj.result(), 3) + self.assertTrue(get_runtime(file_name=os.path.join(cache_directory, task_key + ".h5out")) > 0.0) future_file_obj = FutureItem( file_name=os.path.join(cache_directory, task_key + ".h5out") ) @@ -86,6 +88,7 @@ def test_execute_function_kwargs(self): ) self.assertTrue(future_obj.done()) self.assertEqual(future_obj.result(), 3) + self.assertTrue(get_runtime(file_name=os.path.join(cache_directory, task_key + ".h5out")) > 0.0) future_file_obj = FutureItem( file_name=os.path.join(cache_directory, task_key + ".h5out") )