Skip to content

Fabi dev #3

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
FabiPi3 wants to merge 20 commits into
base: main
Choose a base branch
from
Open

Fabi dev #3

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
20 commits
Select commit Hold shift + click to select a range
d9ea096
fix pass manager config for dynamic jobs
Dec 6, 2023
13b834f
fix tests with removing decorator
Dec 7, 2023
1933d34
add doc
Dec 8, 2023
1619bb0
apply black
Dec 8, 2023
393e71d
Merge branch 'main' into fix_manager_config_dynamic_jobs
Dec 11, 2023
6b224c1
Merge branch 'main' into fix_manager_config_dynamic_jobs
Dec 12, 2023
50321d7
apply black and reformat import
Dec 12, 2023
9bea0c8
Merge branch 'materialsproject:main' into fix_manager_config_dynamic_…
FabiPi3 Jan 16, 2024
53ff019
Merge branch 'main' into fix_manager_config_dynamic_jobs
Feb 28, 2024
d413017
add type hint
Feb 28, 2024
72081f9
Merge branch 'main' into fix_manager_config_dynamic_jobs
Mar 13, 2024
e6e92d0
remove enum value keywords
Mar 13, 2024
c536056
Merge branch 'fix_enum_serialization' into fabi_dev
Mar 13, 2024
978e66f
fix ruff
Mar 13, 2024
7e1a73d
Merge branch 'fix_enum_serialization' into fabi_dev
Mar 13, 2024
2c19c67
fix typing
Apr 2, 2024
4ac1d9d
fix ruff
Apr 3, 2024
ee58872
Merge branch 'add_type_hints' into fabi_dev
Apr 3, 2024
65fd809
fix @job annotation
Apr 5, 2024
78f8780
Merge branch 'fix_job_annotation' into fabi_dev
Apr 5, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
26 changes: 14 additions & 12 deletions src/jobflow/core/job.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -78,7 +78,9 @@ def job(method: Callable = None, **job_kwargs) -> Callable[..., Callable[..., Jo
pass


def job(method: Callable = None, **job_kwargs):
def job(
method: Callable = None, **job_kwargs
) -> Callable[..., Job] | Callable[..., Callable[..., Job]]:
"""
Wrap a function to produce a :obj:`Job`.

Expand Down Expand Up @@ -638,9 +640,7 @@ def run(self, store: jobflow.JobStore, job_dir: Path = None) -> Response:
pass_manager_config(response.replace, passed_config)

try:
output = jsanitize(
response.output, strict=True, enum_values=True, allow_bson=True
)
output = jsanitize(response.output, strict=True, allow_bson=True)
except AttributeError as err:
raise RuntimeError(
"Job output contained an object that is not MSONable and therefore "
Expand Down Expand Up @@ -1121,7 +1121,7 @@ def as_dict(self) -> dict:

# fireworks can't serialize functions and classes, so explicitly serialize to
# the job recursively using monty to avoid issues
return jsanitize(d, strict=True, enum_values=True, allow_bson=True)
return jsanitize(d, strict=True, allow_bson=True)

def __setattr__(self, key, value):
"""Handle setting attributes. Implements a special case for job name."""
Expand Down Expand Up @@ -1288,7 +1288,6 @@ def apply_schema(output: Any, schema: type[BaseModel] | None):
return schema(**output)


@job(config=JobConfig(resolve_references=False, on_missing_references=OnMissing.NONE))
def store_inputs(inputs: Any) -> Any:
"""
Job to store inputs.
Expand Down Expand Up @@ -1338,7 +1337,12 @@ def prepare_replace(
# add a job with same UUID as the current job to store the outputs of the
# flow; this job will inherit the metadata and output schema of the current
# job
store_output_job = store_inputs(replace.output)
new_config = JobConfig(
resolve_references=False, on_missing_references=OnMissing.NONE
)
store_output_job = Job(
store_inputs, function_args=(replace.output,), config=new_config
)
store_output_job.set_uuid(current_job.uuid)
store_output_job.index = current_job.index + 1
store_output_job.metadata = current_job.metadata
Expand Down Expand Up @@ -1370,17 +1374,15 @@ def pass_manager_config(
"""
Pass the manager config on to any jobs in the jobs array.

Merge with already specified manager config.

Parameters
----------
jobs
A job, flow, or list of jobs/flows.
manager_config
A manager config to pass on.
metadata
Metadata to pass on.
"""
from copy import deepcopy

all_jobs: list[Job] = []

def get_jobs(arg):
Expand All @@ -1400,4 +1402,4 @@ def get_jobs(arg):

# update manager config
for ajob in all_jobs:
ajob.config.manager_config = deepcopy(manager_config)
ajob.config.manager_config = manager_config | ajob.config.manager_config
4 changes: 2 additions & 2 deletions src/jobflow/core/reference.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -394,7 +394,7 @@ def find_and_get_references(arg: Any) -> tuple[OutputReference, ...]:
# argument is a primitive, we won't find a reference here
return ()

arg = jsanitize(arg, strict=True, enum_values=True, allow_bson=True)
arg = jsanitize(arg, strict=True, allow_bson=True)

# recursively find any reference classes
locations = find_key_value(arg, "@class", "OutputReference")
Expand Down Expand Up @@ -458,7 +458,7 @@ def find_and_resolve_references(
return arg

# serialize the argument to a dictionary
encoded_arg = jsanitize(arg, strict=True, enum_values=True, allow_bson=True)
encoded_arg = jsanitize(arg, strict=True, allow_bson=True)

# recursively find any reference classes
locations = find_key_value(encoded_arg, "@class", "OutputReference")
Expand Down
13 changes: 10 additions & 3 deletions tests/core/test_job.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -940,15 +940,22 @@ def add_schema_replace(a, b):


def test_store_inputs(memory_jobstore):
from jobflow.core.job import Job, OutputReference, store_inputs
from jobflow.core.job import (
Job,
OutputReference,
store_inputs,
JobConfig,
OnMissing,
)

test_job = store_inputs(1)
configs = JobConfig(resolve_references=False, on_missing_references=OnMissing.NONE)
test_job = Job(store_inputs, (1,), config=configs)
test_job.run(memory_jobstore)
output = memory_jobstore.query_one({"uuid": test_job.uuid}, ["output"])["output"]
assert output == 1

ref = OutputReference("abcd")
test_job = store_inputs(ref)
test_job = Job(store_inputs, (ref,), config=configs)
test_job.run(memory_jobstore)
output = memory_jobstore.query_one({"uuid": test_job.uuid}, ["output"])["output"]
assert OutputReference.from_dict(output) == ref
Expand Down