allow configurable scheduler load group (#992)

Summary:

Allow configurable scheduler load group for clean scheduler splits

Reviewed By: jesszzzz

Differential Revision: D67290464

Author: lgarg26

torchx/runner/config.py

@@ -197,7 +197,12 @@ def _configparser() -> configparser.ConfigParser:
 
 
 def _get_scheduler(name: str) -> Scheduler:
-    schedulers = get_scheduler_factories()
+    schedulers = {
+        **get_scheduler_factories(
+            group="torchx.schedulers.orchestrator", skip_defaults=True
+        ),
+        **get_scheduler_factories(),
+    }
     if name not in schedulers:
         raise ValueError(
             f"`{name}` is not a registered scheduler. Valid scheduler names: {schedulers.keys()}"
@@ -248,6 +253,7 @@ def dump(
         try:
             sched = _get_scheduler(sched_name)
         except ModuleNotFoundError:
+            print(f"[LG] no schedulule load for {sched_name}")
             continue
 
         section = f"{sched_name}"

torchx/runner/test/config_test.py

@@ -470,10 +470,19 @@ def test_dump_and_load_all_registered_schedulers(self) -> None:
         sfile = StringIO()
         dump(sfile)
 
-        for sched_name, sched in get_scheduler_factories().items():
+        scheduler_factories = {
+            **get_scheduler_factories(),
+            **get_scheduler_factories(
+                group="torchx.schedulers.orchestrator", skip_defaults=True
+            ),
+        }
+
+        for sched_name, sched in scheduler_factories.items():
             sfile.seek(0)  # reset the file pos
             cfg = {}
             load(scheduler=sched_name, f=sfile, cfg=cfg)
-
             for opt_name, _ in sched("test").run_opts():
-                self.assertTrue(opt_name in cfg)
+                self.assertTrue(
+                    opt_name in cfg,
+                    f"missing {opt_name} in {sched} run opts with cfg {cfg}",
+                )

torchx/schedulers/__init__.py

@@ -42,9 +42,11 @@ def run(*args: object, **kwargs: object) -> Scheduler:
     return run
 
 
-def get_scheduler_factories() -> Dict[str, SchedulerFactory]:
+def get_scheduler_factories(
+    group: str = "torchx.schedulers", skip_defaults: bool = False
+) -> Dict[str, SchedulerFactory]:
     """
-    get_scheduler_factories returns all the available schedulers names and the
+    get_scheduler_factories returns all the available schedulers names under `group` and the
     method to instantiate them.
 
     The first scheduler in the dictionary is used as the default scheduler.
@@ -55,8 +57,9 @@ def get_scheduler_factories() -> Dict[str, SchedulerFactory]:
         default_schedulers[scheduler] = _defer_load_scheduler(path)
 
     return load_group(
-        "torchx.schedulers",
+        group,
         default=default_schedulers,
+        skip_defaults=skip_defaults,
     )
 
 

torchx/util/entrypoints.py

@@ -51,8 +51,7 @@ def run(*args: object, **kwargs: object) -> object:
 
 # pyre-ignore-all-errors[3, 2]
 def load_group(
-    group: str,
-    default: Optional[Dict[str, Any]] = None,
+    group: str, default: Optional[Dict[str, Any]] = None, skip_defaults: bool = False
 ):
     """
     Loads all the entry points specified by ``group`` and returns
@@ -90,6 +89,8 @@ def load_group(
     entrypoints = metadata.entry_points().select(group=group)
 
     if len(entrypoints) == 0:
+        if skip_defaults:
+            return {}
         return default
 
     eps = {}