Add validation for run.Partial based tasks in Experiment (#103)

* Add validation for run.Partial based tasks in Experiment

Signed-off-by: Hemil Desai <[email protected]>

* Add tests

Signed-off-by: Hemil Desai <[email protected]>

* fix

Signed-off-by: Hemil Desai <[email protected]>

---------

Signed-off-by: Hemil Desai <[email protected]>

Author: hemildesai

src/nemo_run/run/api.py

@@ -75,7 +75,7 @@ def run(
     )
     name = name or default_name
     with Experiment(title=name, executor=executor, log_level=log_level) as exp:
-        exp.add(fn_or_script, tail_logs=tail_logs, plugins=plugins)
+        exp.add(fn_or_script, tail_logs=tail_logs, plugins=plugins, name=name)
         if dryrun:
             exp.dryrun()
             return

src/nemo_run/run/experiment.py

@@ -19,6 +19,7 @@
 import inspect
 import json
 import os
+import pprint
 import shutil
 import sys
 import time
@@ -29,6 +30,7 @@
 
 import fiddle as fdl
 import networkx as nx
+from fiddle._src import daglish, diffing
 from rich.console import Group
 from rich.live import Live
 from rich.panel import Panel
@@ -263,6 +265,7 @@ def __init__(
         log_level: str = "INFO",
         _reconstruct: bool = False,
         jobs: list[Job | JobGroup] | None = None,
+        base_dir: str | None = None,
     ) -> None:
         """
         Initializes an experiment run by creating its metadata directory and saving the experiment config.
@@ -286,7 +289,8 @@ def __init__(
         self._title = title
         self._id = id or f"{title}_{int(time.time())}"
 
-        self._exp_dir = os.path.join(NEMORUN_HOME, "experiments", title, self._id)
+        base_dir = base_dir or NEMORUN_HOME
+        self._exp_dir = os.path.join(base_dir, "experiments", title, self._id)
 
         self.log_level = log_level
         self._runner = get_runner()
@@ -378,6 +382,8 @@ def _add_single_job(
         else:
             task_id = name
 
+        self._validate_task(task_info=task_id, task=task)
+
         executor = executor.clone()
         executor.assign(
             self._id,
@@ -418,6 +424,9 @@ def _add_job_group(
         else:
             task_id = name
 
+        for i, _task in enumerate(tasks):
+            self._validate_task(task_info=f"Job Group: {task_id}, job index: {i}", task=_task)
+
         executors = executor if isinstance(executor, list) else [executor]
         cloned_executors = []
         for executor in executors:
@@ -449,6 +458,30 @@ def _add_job_group(
         self._jobs.append(job_group)
         return job_group.id
 
+    def _validate_task(self, task_info: str, task: Union[Partial, Script]) -> None:
+        valid = True
+        message = ""
+        if isinstance(task, Partial):
+            serializer = ZlibJSONSerializer()
+            serialized = serializer.serialize(task)
+            deserialized = serializer.deserialize(serialized)
+            diff = diffing.build_diff(deserialized, task)
+            diff = {
+                daglish.path_str(d.target): (d.new_value if hasattr(d, "new_value") else None)  # type: ignore
+                for d in diff.changes
+            }
+            if deserialized != task:
+                valid = False
+                message += f"""
+Deserialized task does not match original task. The following paths in your task need to be wrapped in `run.Config` or `run.Partial`:
+
+{pprint.PrettyPrinter(indent=4).pformat(diff)}
+
+For more information about `run.Config` and `run.Partial`, please refer to https://github.com/NVIDIA/NeMo-Run/blob/main/docs/source/guides/configuration.md.
+"""
+        if not valid:
+            raise RuntimeError(f"Failed to validate task {task_info}.\n{message}")
+
     def add(
         self,
         task: Union[Partial, Script] | list[Union[Partial, Script]],

test/dummy_factory.py

@@ -25,6 +25,17 @@ class DummyModel:
     activation: str = "relu"
 
 
+class DummyTrainer:
+    def __init__(self, num_epochs: int = 10):
+        self.num_epochs = num_epochs
+
+    def __hash__(self) -> int:
+        return hash((self.num_epochs))
+
+    def __eq__(self, value: object) -> bool:
+        return isinstance(value, DummyTrainer) and self.num_epochs == value.num_epochs
+
+
 @dataclass
 class NestedModel:
     dummy: DummyModel
@@ -88,5 +99,8 @@ def plugin_list(arg: int = 20) -> List[run.Plugin]:
     ]
 
 
+def dummy_train(dummy_model: DummyModel, dummy_trainer: DummyTrainer): ...
+
+
 if __name__ == "__main__":
     run.cli.main(dummy_entrypoint)

test/run/test_experiment.py

@@ -0,0 +1,26 @@
+import pytest
+from fiddle._src.experimental.serialization import UnserializableValueError
+
+import nemo_run as run
+from test.dummy_factory import DummyModel, DummyTrainer, dummy_train
+
+
+@pytest.fixture
+def experiment(tmpdir):
+    return run.Experiment("dummy_experiment", base_dir=tmpdir)
+
+
+class TestValidateTask:
+    def test_validate_task(self, experiment: run.Experiment):
+        experiment._validate_task("valid_script", run.Script(inline="echo 'hello world'"))
+
+        valid_partial = run.Partial(
+            dummy_train, dummy_model=run.Config(DummyModel), dummy_trainer=run.Config(DummyTrainer)
+        )
+        experiment._validate_task("valid_partial", valid_partial)
+
+        invalid_partial = run.Partial(
+            dummy_train, dummy_model=DummyModel(), dummy_trainer=DummyTrainer()
+        )
+        with pytest.raises(UnserializableValueError):
+            experiment._validate_task("invalid_partial", invalid_partial)