Skip to content

Commit

Permalink
Merge remote-tracking branch 'materialsproject/main' into uuid_patch
Browse files Browse the repository at this point in the history
  • Loading branch information
jmmshn committed Jan 22, 2024
2 parents 923468b + 9bd99a8 commit 82ea1b5
Show file tree
Hide file tree
Showing 2 changed files with 20 additions and 7 deletions.
24 changes: 17 additions & 7 deletions src/jobflow/managers/local.py
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,7 @@ def run_locally(
root_dir: str | Path | None = None,
ensure_success: bool = False,
allow_external_references: bool = False,
raise_immediately: bool = False,
) -> dict[str, dict[int, jobflow.Response]]:
"""
Run a :obj:`Job` or :obj:`Flow` locally.
Expand All @@ -46,6 +47,10 @@ def run_locally(
allow_external_references : bool
If False all the references to other outputs should be from other Jobs
of the Flow.
raise_immediately : bool
If True, raise an exception immediately if a job fails. If False, continue
running the flow and only raise an exception at the end if the flow did not
finish running successfully.
Returns
-------
Expand Down Expand Up @@ -102,14 +107,19 @@ def _run_job(job: jobflow.Job, parents):
errored.add(job.uuid)
return None, False

try:
if raise_immediately:
response = job.run(store=store)
except Exception:
import traceback

logger.info(f"{job.name} failed with exception:\n{traceback.format_exc()}")
errored.add(job.uuid)
return None, False
else:
try:
response = job.run(store=store)
except Exception:
import traceback

logger.info(
f"{job.name} failed with exception:\n{traceback.format_exc()}"
)
errored.add(job.uuid)
return None, False

responses[job.uuid][job.index] = response

Expand Down
3 changes: 3 additions & 0 deletions tests/managers/test_local.py
Original file line number Diff line number Diff line change
Expand Up @@ -331,6 +331,9 @@ def test_error_flow(memory_jobstore, clean_dir, error_flow, capsys):
with pytest.raises(RuntimeError):
run_locally(flow, store=memory_jobstore, ensure_success=True)

with pytest.raises(ValueError, match="errored"):
run_locally(flow, store=memory_jobstore, raise_immediately=True)


def test_ensure_success_with_replace(memory_jobstore, error_replace_flow, capsys):
from jobflow import run_locally
Expand Down

0 comments on commit 82ea1b5

Please sign in to comment.