feat: list all registered schedulers (#1009)

Author: azzhipa

torchx/schedulers/__init__.py

@@ -48,15 +48,16 @@ def get_scheduler_factories(
 
     The first scheduler in the dictionary is used as the default scheduler.
     """
-
-    if skip_defaults:
-        default_schedulers = {}
-    else:
-        default_schedulers: dict[str, SchedulerFactory] = {}
-        for scheduler, path in DEFAULT_SCHEDULER_MODULES.items():
-            default_schedulers[scheduler] = _defer_load_scheduler(path)
-
-    return load_group(group, default=default_schedulers)
+    valid_schedulers = (
+        {}
+        if skip_defaults
+        else {
+            name: _defer_load_scheduler(path)
+            for name, path in DEFAULT_SCHEDULER_MODULES.items()
+        }
+    )
+    valid_schedulers.update(load_group(group, default={}))
+    return valid_schedulers
 
 
 def get_default_scheduler_name() -> str:

torchx/schedulers/test/registry_test.py

@@ -43,3 +43,37 @@ def test_get_local_schedulers(self, mock_load_group: MagicMock) -> None:
 
         for scheduler in schedulers.values():
             self.assertEqual("test_session", scheduler.session_name)
+
+    @patch("torchx.schedulers.load_group")
+    def test_custom_schedulers_merged(self, mock_load_group: MagicMock) -> None:
+        mock_scheduler = MagicMock()
+        mock_load_group.return_value = {"custom": mock_scheduler}
+
+        factories = get_scheduler_factories()
+
+        self.assertIn("custom", factories)
+        self.assertEqual(factories["custom"], mock_scheduler)
+        self.assertIn("local_docker", factories)
+
+    @patch("torchx.schedulers.load_group")
+    def test_custom_scheduler_overrides_default(
+        self, mock_load_group: MagicMock
+    ) -> None:
+        mock_scheduler = MagicMock()
+        mock_load_group.return_value = {"local_docker": mock_scheduler}
+
+        factories = get_scheduler_factories()
+
+        self.assertEqual(factories["local_docker"], mock_scheduler)
+
+    @patch("torchx.schedulers.load_group")
+    def test_skip_defaults_with_custom_schedulers(
+        self, mock_load_group: MagicMock
+    ) -> None:
+        mock_scheduler = MagicMock()
+        mock_load_group.return_value = {"custom": mock_scheduler}
+
+        factories = get_scheduler_factories(skip_defaults=True)
+
+        self.assertEqual(factories, {"custom": mock_scheduler})
+        self.assertNotIn("local_docker", factories)