fix dump and load behavior to load same set of schedulers

Differential Revision: D67341003 Pull Request resolved: #993
pytorch · Dec 17, 2024 · ff22758 · ff22758
1 parent 90884d7
commit ff22758
Showing 1 changed file with 10 additions and 1 deletion.
diff --git a/torchx/runner/config.py b/torchx/runner/config.py
@@ -249,7 +249,16 @@ def dump(
     if schedulers:
         scheds = schedulers
     else:
-        scheds = get_scheduler_factories().keys()
+        scheduler_factories = {
+            **get_scheduler_factories(),
+            **(
+                get_scheduler_factories(
+                    group="torchx.schedulers.orchestrator", skip_defaults=True
+                )
+                or {}
+            ),
+        }
+        scheds = scheduler_factories.keys()
 
     config = _configparser()
     for sched_name in scheds: