Merge pull request #486 from materialsproject/run-locally-root-dir

`run_locally()` add `root_dir: str | Path | None` keyword

src/jobflow/core/reference.py

@@ -177,7 +177,11 @@ def resolve(
 
         for attr_type, attr in self.attributes:
             # i means index else use attribute access
-            data = data[attr] if attr_type == "i" else getattr(data, attr)
+            data = (
+                data[attr]
+                if attr_type == "i" or isinstance(data, dict)
+                else getattr(data, attr)
+            )
 
         return data
 

src/jobflow/managers/local.py

@@ -6,8 +6,9 @@
 import typing
 
 if typing.TYPE_CHECKING:
-    import jobflow
+    from pathlib import Path
 
+    import jobflow
 
 logger = logging.getLogger(__name__)
 
@@ -17,6 +18,7 @@ def run_locally(
     log: bool = True,
     store: jobflow.JobStore = None,
     create_folders: bool = False,
+    root_dir: str | Path | None = None,
     ensure_success: bool = False,
     allow_external_references: bool = False,
 ) -> dict[str, dict[int, jobflow.Response]]:
@@ -25,25 +27,29 @@ def run_locally(
 
     Parameters
     ----------
-    flow
+    flow : Flow | Job | list[Job]
         A job or flow.
-    log
+    log : bool
         Whether to print log messages.
-    store
+    store : JobStore
         A job store. If a job store is not specified then
         :obj:`JobflowSettings.JOB_STORE` will be used. By default this is a maggma
         ``MemoryStore`` but can be customised by setting the jobflow configuration file.
-    create_folders
+    create_folders : bool
         Whether to run each job in a new folder.
-    ensure_success
+    root_dir : str | Path | None
+        The root directory to run the jobs in or where to create new subfolders if
+            ``create_folders`` is True. If None then the current working
+            directory will be used.
+    ensure_success : bool
         Raise an error if the flow was not executed successfully.
-    allow_external_references
+    allow_external_references : bool
         If False all the references to other outputs should be from other Jobs
         of the Flow.
 
     Returns
     -------
-    Dict[str, Dict[int, Response]]
+    dict[str, dict[int, Response]]
         The responses of the jobs, as a dict of ``{uuid: {index: response}}``.
     """
     from collections import defaultdict
@@ -60,6 +66,9 @@ def run_locally(
     if store is None:
         store = SETTINGS.JOB_STORE
 
+    root_dir = Path.cwd() if root_dir is None else Path(root_dir).resolve()
+    root_dir.mkdir(exist_ok=True)
+
     store.connect()
 
     if log:
@@ -72,8 +81,6 @@ def run_locally(
     responses: dict[str, dict[int, jobflow.Response]] = defaultdict(dict)
     stop_jobflow = False
 
-    root_dir = Path.cwd()
-
     def _run_job(job: jobflow.Job, parents):
         nonlocal stop_jobflow
 

tests/managers/test_local.py

@@ -62,6 +62,15 @@ def test_simple_flow(memory_jobstore, clean_dir, simple_flow, capsys):
     folders = list(Path(".").glob("job_*/"))
     assert len(folders) == 1
 
+    # run with folders and root_dir
+    assert Path(root_dir := "test").exists() is False
+    responses = run_locally(
+        flow, store=memory_jobstore, create_folders=True, root_dir=root_dir
+    )
+    assert responses[uuid][1].output == "12345_end"
+    folders = list(Path(root_dir).glob("job_*/"))
+    assert len(folders) == 1
+
 
 def test_connected_flow(memory_jobstore, clean_dir, connected_flow):
     from jobflow import run_locally