diff --git a/src/hipercow/root.py b/src/hipercow/root.py index 4857a94..6cfb8c3 100644 --- a/src/hipercow/root.py +++ b/src/hipercow/root.py @@ -28,6 +28,18 @@ def __init__(self, path: str | Path) -> None: raise Exception(msg) self.path = path + def path_task(self, task_id: str) -> Path: + return self.path / "tasks" / task_id[:2] / task_id[2:] + + def path_task_times(self, task_id: str) -> Path: + return self.path_task(task_id) / "times" + + def path_task_data(self, task_id: str) -> Path: + return self.path_task(task_id) / "data" + + def path_task_result(self, task_id: str) -> Path: + return self.path_task(task_id) / "result" + def open_root(path: None | str | Path = None) -> Root: root = find_file_descend("hipercow", path or Path.cwd()) diff --git a/src/hipercow/task.py b/src/hipercow/task.py index e6891cc..c281de0 100644 --- a/src/hipercow/task.py +++ b/src/hipercow/task.py @@ -44,13 +44,13 @@ class TaskTimes: finished: float def write(self, root: Root, task_id: str): - with open(root.path / "tasks" / task_id / "times", "wb") as f: + with root.path_task_times(task_id).open("wb") as f: pickle.dump(self, f) def task_status(root: Root, task_id: str) -> TaskStatus: # check_task_id(task_id) - path = root.path / "tasks" / task_id + path = root.path_task(task_id) if not path.exists(): return TaskStatus.MISSING for v, p in STATUS_FILE_MAP.items(): @@ -60,7 +60,7 @@ def task_status(root: Root, task_id: str) -> TaskStatus: def set_task_status(root: Root, task_id: str, status: TaskStatus): - file_create(root.path / "tasks" / task_id / STATUS_FILE_MAP[status]) + file_create(root.path_task(task_id) / STATUS_FILE_MAP[status]) @dataclass @@ -80,12 +80,13 @@ def read(root: Root, task_id: str): def task_data_write(root: Root, data: TaskData) -> None: - path_task_dir = root.path / "tasks" / data.task_id + task_id = data.task_id + path_task_dir = root.path_task(task_id) path_task_dir.mkdir(parents=True, exist_ok=True) - with open(path_task_dir / "data", "wb") as f: + with root.path_task_data(task_id).open("wb") as f: pickle.dump(data, f) def task_data_read(root: Root, task_id: str) -> TaskData: - with open(root.path / "tasks" / task_id / "data", "rb") as f: + with root.path_task_data(task_id).open("rb") as f: return pickle.load(f) diff --git a/src/hipercow/task_eval.py b/src/hipercow/task_eval.py index e256175..047bb73 100644 --- a/src/hipercow/task_eval.py +++ b/src/hipercow/task_eval.py @@ -32,7 +32,7 @@ def task_eval_data(root: Root, data: TaskData) -> None: msg = f"Can't run '{task_id}', which has status '{status}'" raise Exception(msg) - t_created = (root.path / "tasks" / task_id / "data").stat().st_ctime + t_created = root.path_task_data(task_id).stat().st_ctime t_start = time.time() set_task_status(root, task_id, TaskStatus.RUNNING) @@ -43,7 +43,7 @@ def task_eval_data(root: Root, data: TaskData) -> None: t_end = time.time() status = TaskStatus.SUCCESS if res.success else TaskStatus.FAILURE - with open(root.path / "tasks" / task_id / "result", "wb") as f: + with root.path_task_result(task_id).open("wb") as f: pickle.dump(res.data, f) times = TaskTimes(t_created, t_start, t_end) diff --git a/tests/test_create.py b/tests/test_create.py index d6b0e2b..779b021 100644 --- a/tests/test_create.py +++ b/tests/test_create.py @@ -11,7 +11,7 @@ def test_create_simple_task(tmp_path): with transient_working_directory(tmp_path): tid = tc.task_create_shell(["echo", "hello world"]) assert re.match("^[0-9a-f]{32}$", tid) - assert (tmp_path / "tasks" / tid / "data").exists() + assert (tmp_path / "tasks" / tid[:2] / tid[2:] / "data").exists() d = TaskData.read(root.open_root(tmp_path), tid) assert isinstance(d, TaskData) assert d.task_id == tid