Merge pull request #817 from gpetretto/jobconfig

Jobconfig

Author: utf

src/jobflow/core/flow.py

@@ -706,6 +706,7 @@ def update_config(
         function_filter: Callable = None,
         attributes: list[str] | str = None,
         dynamic: bool = True,
+        dict_mod: bool = False,
     ):
         """
         Update the job config of all Jobs in the Flow.
@@ -728,6 +729,9 @@ def update_config(
         dynamic
             The updates will be propagated to Jobs/Flows dynamically generated at
             runtime.
+        dict_mod
+            Use the dict mod language to apply updates. See :obj:`.DictMods` for more
+            details.
 
         Examples
         --------
@@ -767,6 +771,7 @@ def update_config(
                 function_filter=function_filter,
                 attributes=attributes,
                 dynamic=dynamic,
+                dict_mod=dict_mod,
             )
 
     def add_hosts_uuids(

src/jobflow/core/job.py

@@ -1022,6 +1022,7 @@ def update_config(
         function_filter: Callable = None,
         attributes: list[str] | str = None,
         dynamic: bool = True,
+        dict_mod: bool = False,
     ):
         """
         Update the job config.
@@ -1045,6 +1046,9 @@ def update_config(
         dynamic
             The updates will be propagated to Jobs/Flows dynamically generated at
             runtime.
+        dict_mod
+            Use the dict mod language to apply updates. See :obj:`.DictMods` for more
+            details.
 
         Examples
         --------
@@ -1096,12 +1100,23 @@ def update_config(
         At variance, if `dynamic` is set to `False` the `manager_config` option will
         only be set for the `test_job` and not for the generated Jobs.
         """
+        from jobflow.utils.dict_mods import apply_mod
+
+        if dict_mod and attributes:
+            raise ValueError("dict_mod and attributes options cannot be used together")
+
+        if dict_mod and isinstance(config, JobConfig):
+            raise ValueError(
+                "If dict_mod is selected the update config cannot be a JobConfig object"
+            )
+
         if dynamic:
             dict_input = {
                 "config": config,
                 "name_filter": name_filter,
                 "function_filter": function_filter,
                 "attributes": attributes,
+                "dict_mod": dict_mod,
             }
             self.config_updates.append(dict_input)
 
@@ -1119,29 +1134,34 @@ def update_config(
             return
 
         # if we get to here then we pass all the filters
-        if isinstance(config, dict):
-            # convert dict specification to a JobConfig but set the attributes
-            if attributes is None:
-                attributes = list(config.keys())
-
-            attributes = [attributes] if isinstance(attributes, str) else attributes
-            if not set(attributes).issubset(set(config.keys())):
-                raise ValueError(
-                    "Specified attributes include a key that is not present in the "
-                    "config dictionary."
-                )
-            config = JobConfig(**config)
-
-        if attributes is None:
-            # overwrite the whole config
-            self.config = config
+        if dict_mod:
+            conf_dict = self.config.as_dict()
+            apply_mod(config, conf_dict)
+            self.config = JobConfig.from_dict(conf_dict)
         else:
-            # only update the specified attributes
-            attributes = [attributes] if isinstance(attributes, str) else attributes
-            for attr in attributes:
-                if not hasattr(self.config, attr):
-                    raise ValueError(f"Unknown JobConfig attribute: {attr}")
-                setattr(self.config, attr, getattr(config, attr))
+            if isinstance(config, dict):
+                # convert dict specification to a JobConfig but set the attributes
+                if attributes is None:
+                    attributes = list(config.keys())
+
+                attributes = [attributes] if isinstance(attributes, str) else attributes
+                if not set(attributes).issubset(set(config.keys())):
+                    raise ValueError(
+                        "Specified attributes include a key that is not present in the "
+                        "config dictionary."
+                    )
+                config = JobConfig(**config)
+
+            if attributes is None:
+                # overwrite the whole config
+                self.config = config
+            else:
+                # only update the specified attributes
+                attributes = [attributes] if isinstance(attributes, str) else attributes
+                for attr in attributes:
+                    if not hasattr(self.config, attr):
+                        raise ValueError(f"Unknown JobConfig attribute: {attr}")
+                    setattr(self.config, attr, getattr(config, attr))
 
     def as_dict(self) -> dict:
         """Serialize the job as a dictionary."""

tests/core/test_job.py

@@ -1332,6 +1332,41 @@ def jsm_wrapped(a, b):
     ):
         test_job.update_config(new_config_dict, attributes="abc_xyz")
 
+    # test with dict_mod
+    test_job = Job(add)
+    with pytest.raises(
+        ValueError, match="dict_mod and attributes options cannot be used together"
+    ):
+        test_job.update_config({}, attributes="manager_config", dict_mod=True)
+
+    with pytest.raises(
+        ValueError,
+        match="If dict_mod is selected the update config cannot be a JobConfig object",
+    ):
+        test_job.update_config(JobConfig(), dict_mod=True)
+
+    dict_mod_config_1 = {
+        "_set": {
+            "manager_config->test->x": 1,
+            "resolve_references": True,
+        }
+    }
+    test_job.update_config(dict_mod_config_1, dict_mod=True)
+    assert test_job.config.manager_config["test"]["x"] == 1
+    assert test_job.config.resolve_references
+
+    dict_mod_config_2 = {
+        "_set": {
+            "manager_config->test->y": 2,
+            "manager_config->test2->z": 3,
+        }
+    }
+    test_job.update_config(dict_mod_config_2, dict_mod=True)
+    assert test_job.config.manager_config["test"]["x"] == 1
+    assert test_job.config.manager_config["test"]["y"] == 2
+    assert test_job.config.manager_config["test2"]["z"] == 3
+    assert test_job.config.resolve_references
+
     # test applied dynamic updates
     @dataclass
     class TestMaker(Maker):