From ff22758025f7c90b7a5f7e2f25f99f502ce95391 Mon Sep 17 00:00:00 2001
From: lgarg26 <154364140+lgarg26@users.noreply.github.com>
Date: Tue, 17 Dec 2024 15:28:41 -0500
Subject: [PATCH] fix dump and load behavior to load same set of schedulers

Differential Revision: D67341003

Pull Request resolved: https://github.com/pytorch/torchx/pull/993
---
 torchx/runner/config.py | 11 ++++++++++-
 1 file changed, 10 insertions(+), 1 deletion(-)

diff --git a/torchx/runner/config.py b/torchx/runner/config.py
index ba243c7b..788acb27 100644
--- 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: