chore: Improve core v2 init API [MD-441] (#9560)

determined-ai · Jul 15, 2024 · 02da2a2 · 02da2a2
1 parent 0494cdf
commit 02da2a2
Show file tree

Hide file tree

Showing 18 changed files with 140 additions and 92 deletions.
diff --git a/docs/release-notes/detached-mode-config-improvement.rst b/docs/release-notes/detached-mode-config-improvement.rst
@@ -0,0 +1,7 @@
+:orphan:
+
+**Deprecations**
+
+-  Detached mode: The ``defaults`` and ``unmanaged`` parameters of the ``init`` function for
+   unmanaged experiment have been deprecated and will be removed in a future version. Please use
+   ``config`` instead.
diff --git a/docs/tutorials/detached-mode/distributed-training-checkpointing.rst b/docs/tutorials/detached-mode/distributed-training-checkpointing.rst
@@ -67,7 +67,7 @@ context:
        dist.init_process_group("gloo")
        distributed = core_v2.DistributedContext.from_torch_distributed()
        core_v2.init(
-           defaults=core_v2.DefaultConfig(
+           config=core_v2.Config(
                name="unmanaged-3-torch-distributed",
            ),
            distributed=distributed,

diff --git a/docs/tutorials/detached-mode/save-load-checkpoints.rst b/docs/tutorials/detached-mode/save-load-checkpoints.rst
@@ -55,11 +55,9 @@ metrics to:
 
    def main():
        core_v2.init(
-           defaults=core_v2.DefaultConfig(
+           config=core_v2.Config(
                name="unmanaged-2-checkpoints",
                checkpoint_storage="/path/to/checkpoint",
-           ),
-           unmanaged=core_v2.UnmanagedConfig(
                external_experiment_id="my-existing-experiment",
                external_trial_id="my-existing-trial",
            ),

diff --git a/docs/tutorials/detached-mode/simple-metrics-reporting.rst b/docs/tutorials/detached-mode/simple-metrics-reporting.rst
@@ -68,7 +68,7 @@ function:
 
    def main():
        core_v2.init(
-           defaults=core_v2.DefaultConfig(
+           config=core_v2.Config(
                name="detached_mode_example",
            ),
        )

diff --git a/e2e_tests/tests/fixtures/core_api/pytorch_profiler_sync.py b/e2e_tests/tests/fixtures/core_api/pytorch_profiler_sync.py
@@ -7,7 +7,7 @@
 
 def main() -> None:
     core_v2.init(
-        defaults=core_v2.DefaultConfig(
+        config=core_v2.Config(
             name="pytorch-profiler-sync-test",
         ),
     )

diff --git a/e2e_tests/tests/fixtures/unmanaged/checkpoint_distributed.py b/e2e_tests/tests/fixtures/unmanaged/checkpoint_distributed.py
@@ -25,15 +25,13 @@ def main():
     logging.getLogger("determined").setLevel(logging.INFO)
     distributed = core_v2.DistributedContext.from_torch_distributed()
     core_v2.init(
-        defaults=core_v2.DefaultConfig(
+        config=core_v2.Config(
             name=name,
-            checkpoint_storage="/tmp/determined-cp",
-        ),
-        distributed=distributed,
-        unmanaged=core_v2.UnmanagedConfig(
             external_experiment_id=name,
             external_trial_id=name,
+            checkpoint_storage="/tmp/determined-cp",
         ),
+        distributed=distributed,
     )
 
     # Use framework-native dtrain utilities, as normal.

diff --git a/e2e_tests/tests/fixtures/unmanaged/checkpointing.py b/e2e_tests/tests/fixtures/unmanaged/checkpointing.py
@@ -10,13 +10,11 @@ def main():
     assert "DET_TEST_EXTERNAL_EXP_ID" in os.environ
     name = os.environ["DET_TEST_EXTERNAL_EXP_ID"]
     core_v2.init(
-        defaults=core_v2.DefaultConfig(
+        config=core_v2.Config(
             name=name,
-            checkpoint_storage="/tmp/determined-cp",
-        ),
-        unmanaged=core_v2.UnmanagedConfig(
             external_experiment_id=name,
             external_trial_id=name,
+            checkpoint_storage="/tmp/determined-cp",
         ),
     )
 

diff --git a/e2e_tests/tests/fixtures/unmanaged/error_termination.py b/e2e_tests/tests/fixtures/unmanaged/error_termination.py
@@ -14,13 +14,11 @@ def main():
 
     checkpoint_storage = tempfile.mkdtemp()
     core_v2.init(
-        defaults=core_v2.DefaultConfig(
+        config=core_v2.Config(
             name=name,
-        ),
-        checkpoint_storage=checkpoint_storage,
-        unmanaged=core_v2.UnmanagedConfig(
             external_experiment_id=name,
             external_trial_id=name,
+            checkpoint_storage=checkpoint_storage,
         ),
     )
 

diff --git a/examples/features/unmanaged/1_singleton.py b/examples/features/unmanaged/1_singleton.py
@@ -8,17 +8,13 @@
 def main():
     core_v2.init(
         # For managed experiments, will be overridden by the yaml config.
-        # Future: merge this and yaml configs field-by-field at runtime.
-        defaults=core_v2.DefaultConfig(
+        config=core_v2.Config(
             name="unmanaged-1-singleton",
             # labels=["some", "set", "of", "labels"],
             # description="some description",
+            # workspace="...",
+            # project="...",
         ),
-        # `UnmanagedConfig` values will not get merged, and will only be used in the unmanaged mode.
-        # unmanaged=core_v2.UnmanagedConfig(
-        #   workspace="...",
-        #   project="...",
-        # )
     )
 
     for i in range(100):

diff --git a/examples/features/unmanaged/2_checkpoints.py b/examples/features/unmanaged/2_checkpoints.py
@@ -7,12 +7,10 @@
 
 def main():
     core_v2.init(
-        defaults=core_v2.DefaultConfig(
+        config=core_v2.Config(
             name="unmanaged-2-checkpoints",
             # We allow configuring the local checkpoint storage directory.
             # checkpoint_storage="/tmp/determined-cp",
-        ),
-        unmanaged=core_v2.UnmanagedConfig(
             external_experiment_id="test-unmanaged-2-checkpoints",
             external_trial_id="test-unmanaged-2-checkpoints",
             # e.g., requeued jobs on slurm:

diff --git a/examples/features/unmanaged/3_torch_distributed.py b/examples/features/unmanaged/3_torch_distributed.py
@@ -21,7 +21,7 @@ def main():
     logging.getLogger("determined").setLevel(logging.INFO)
     distributed = core_v2.DistributedContext.from_torch_distributed()
     core_v2.init(
-        defaults=core_v2.DefaultConfig(
+        config=core_v2.Config(
             name="unmanaged-3-torch-distributed",
         ),
         distributed=distributed,

diff --git a/examples/features/unmanaged/advanced/checkpoints_advanced.py b/examples/features/unmanaged/advanced/checkpoints_advanced.py
@@ -21,16 +21,14 @@ def __getitem__(self, index: int) -> Tuple:
 
 def main():
     core_v2.init(
-        defaults=core_v2.DefaultConfig(
+        config=core_v2.Config(
             name="unmanaged-checkpoints-advanced",
             hparams={
                 "lr": 1e-5,
                 "max_epochs": 10,
             },
             labels=["some", "set", "of", "labels"],
             description="torch identity example",
-        ),
-        unmanaged=core_v2.UnmanagedConfig(
             external_experiment_id="unmanaged-checkpoints-advanced",
             external_trial_id="unmanaged-checkpoints-advanced",
         ),

diff --git a/examples/features/unmanaged/advanced/context_manager.py b/examples/features/unmanaged/advanced/context_manager.py
@@ -7,7 +7,7 @@
 
 def main():
     with core_v2.init_context(
-        defaults=core_v2.DefaultConfig(
+        config=core_v2.Config(
             name="unmanaged-context-manager",
         ),
     ) as core_context:

diff --git a/examples/features/unmanaged/ray/ray_hp_search.py b/examples/features/unmanaged/ray/ray_hp_search.py
@@ -17,7 +17,7 @@ def objective(config):
     print(f"experiment name: {experiment_name} trial name: {trial_name}")
 
     core_v2.init(
-        defaults=core_v2.DefaultConfig(
+        config=core_v2.Config(
             name=experiment_name,
             hparams={
                 "hp": config["hp"],
@@ -29,8 +29,6 @@ def objective(config):
                 "metric": "loss",
                 "smaller_is_better": True,
             },
-        ),
-        unmanaged=core_v2.UnmanagedConfig(
             external_experiment_id=experiment_name,
             external_trial_id=trial_name,
         ),

diff --git a/examples/features/unmanaged/ray/ray_hp_search_cleanup.py b/examples/features/unmanaged/ray/ray_hp_search_cleanup.py
@@ -19,7 +19,7 @@ def setup(self, config):
         print(f"experiment name: {experiment_name} trial name: {trial_name}")
 
         core_v2.init(
-            defaults=core_v2.DefaultConfig(
+            config=core_v2.Config(
                 name=experiment_name,
                 hparams={
                     "hp": config["hp"],
@@ -29,8 +29,6 @@ def setup(self, config):
                     "metric": "loss",
                     "smaller_is_better": True,
                 },
-            ),
-            unmanaged=core_v2.UnmanagedConfig(
                 external_experiment_id=experiment_name,
                 external_trial_id=trial_name,
             ),

diff --git a/harness/determined/experimental/core_v2/__init__.py b/harness/determined/experimental/core_v2/__init__.py
@@ -5,6 +5,7 @@
 # Import these directly to make core_v2 a complete package.
 from determined.core import DistributedContext, PreemptMode, TensorboardMode
 from determined.experimental.core_v2._core_v2 import (
+    Config,
     DefaultConfig,
     UnmanagedConfig,
     init_context,