ckpt handler

mdagent/utils/ckpt_handler.py

@@ -0,0 +1,37 @@
+import os
+
+
+def find_repo_root(start_path="current"):
+    """
+    Finds the folder containing setup.py
+    """
+    if start_path == "current":
+        start_path = os.getcwd()
+    current_path = start_path
+    while current_path != os.path.dirname(current_path):
+        if "setup.py" in os.listdir(current_path):
+            return current_path
+        current_path = os.path.dirname(current_path)
+    raise FileNotFoundError(
+        "Could not find a folder containing setup.py in the directory tree."
+    )
+
+
+def get_ckpt_folder_path(ckpt_dir="ckpt", start_path="current"):
+    """Returns the path to the ckpt folder in the repository."""
+    repo_root = find_repo_root(start_path)
+    return os.path.join(repo_root, ckpt_dir)
+
+
+def make_ckpt_path():
+    root_repo = find_repo_root()
+    for i in range(10):
+        ckpt_path = os.path.join(root_repo, "ckpt", f"ckpt_{i}")
+        if not os.path.exists(ckpt_path):
+            # make that folder
+            os.makedirs(ckpt_path)
+            return ckpt_path
+
+
+if __name__ == "__main__":
+    print(make_ckpt_path())