cli: defer loading schedulers until requested by name

Author: d4l3k

torchx/cli/cmd_cancel.py

@@ -10,7 +10,6 @@
 
 from torchx.cli.cmd_base import SubCommand
 from torchx.runner import get_runner
-from torchx.specs.api import parse_app_handle
 
 logger: logging.Logger = logging.getLogger(__name__)
 
@@ -25,6 +24,5 @@ def add_arguments(self, subparser: argparse.ArgumentParser) -> None:
 
     def run(self, args: argparse.Namespace) -> None:
         app_handle = args.app_handle
-        _, session_name, _ = parse_app_handle(app_handle)
-        runner = get_runner(name=session_name)
+        runner = get_runner()
         runner.cancel(app_handle)

torchx/cli/cmd_describe.py

@@ -28,15 +28,15 @@ def add_arguments(self, subparser: argparse.ArgumentParser) -> None:
 
     def run(self, args: argparse.Namespace) -> None:
         app_handle = args.app_handle
-        scheduler, session_name, app_id = parse_app_handle(app_handle)
-        runner = get_runner(name=session_name)
+        scheduler, _, app_id = parse_app_handle(app_handle)
+        runner = get_runner()
         app = runner.describe(app_handle)
 
         if app:
             pprint.pprint(dataclasses.asdict(app), indent=2, width=80)
         else:
             logger.error(
-                f"AppDef: {app_id} on session: {session_name},"
+                f"AppDef: {app_id},"
                 f" does not exist or has been removed from {scheduler}'s data plane"
             )
             sys.exit(1)

torchx/cli/cmd_log.py

@@ -92,7 +92,7 @@ def get_logs(
     role_name = path[2] if len(path) > 2 else None
 
     if not runner:
-        runner = get_runner(name=session_name)
+        runner = get_runner()
     app_handle = make_app_handle(scheduler_backend, session_name, app_id)
 
     if len(path) == 4:

torchx/cli/cmd_run.py

@@ -125,15 +125,14 @@ def add_arguments(self, subparser: argparse.ArgumentParser) -> None:
             default=get_default_scheduler_name(),
             choices=list(scheduler_names),
             action=torchxconfig_run,
-            help=f"Name of the scheduler to use. One of: [{','.join(scheduler_names)}]",
+            help="Name of the scheduler to use.",
         )
         subparser.add_argument(
             "-cfg",
             "--scheduler_args",
             type=str,
             help="Arguments to pass to the scheduler (Ex:`cluster=foo,user=bar`)."
-            " For a list of scheduler run options run: `torchx runopts`"
-            "",
+            " For a list of scheduler run options run: `torchx runopts`",
         )
         subparser.add_argument(
             "--dryrun",

torchx/cli/cmd_status.py

@@ -159,15 +159,15 @@ def add_arguments(self, subparser: argparse.ArgumentParser) -> None:
 
     def run(self, args: argparse.Namespace) -> None:
         app_handle = args.app_handle
-        scheduler, session_name, app_id = parse_app_handle(app_handle)
-        runner = get_runner(name=session_name)
+        scheduler, _, app_id = parse_app_handle(app_handle)
+        runner = get_runner()
         app_status = runner.status(app_handle)
         filter_roles = parse_list_arg(args.roles)
         if app_status:
             logger.info(format_app_status(app_status, filter_roles))
         else:
             logger.error(
-                f"AppDef: {app_id} on session: {session_name},"
+                f"AppDef: {app_id},"
                 f" does not exist or has been removed from {scheduler}'s data plane"
             )
             sys.exit(1)

torchx/cli/main.py

@@ -57,7 +57,6 @@ def get_sub_cmds() -> Dict[str, SubCommand]:
     override_sub_cmds = load_group(
         "torchx.cli.cmds",
         default={},
-        ignore_missing=True,
     )
     for cmd_name, cmd_cls in override_sub_cmds.items():
         sub_cmds[cmd_name] = cmd_cls()

torchx/cli/test/main_test.py

@@ -7,12 +7,13 @@
 
 import argparse
 import os
+import sys
 import unittest
 from pathlib import Path
 from unittest.mock import MagicMock, patch
 
 from torchx.cli.cmd_base import SubCommand
-from torchx.cli.main import get_sub_cmds, main
+from torchx.cli.main import create_parser, get_sub_cmds, main
 
 
 _root: Path = Path(__file__).parent
@@ -68,6 +69,11 @@ def test_version(self) -> None:
                 ]
             )
 
+    def test_imports(self) -> None:
+        parser = create_parser(get_sub_cmds())
+        for scheduler in sys.modules:
+            self.assertNotIn("local_scheduler", scheduler)
+
     @patch("torchx.cli.main.load_group")
     def test_get_sub_cmds(self, load_group_mock: MagicMock) -> None:
         load_group_mock.return_value = {"run": _TestCmd}

torchx/runner/api.py

@@ -15,7 +15,7 @@
 
 from pyre_extensions import none_throws
 from torchx.runner.events import log_event
-from torchx.schedulers import get_schedulers
+from torchx.schedulers import get_scheduler_factories, SchedulerFactory
 from torchx.schedulers.api import Scheduler, Stream
 from torchx.specs import (
     AppDef,
@@ -50,9 +50,9 @@ class Runner:
     def __init__(
         self,
         name: str,
-        # pyre-fixme: Scheduler opts
-        schedulers: Dict[str, Scheduler],
+        schedulers_factories: Dict[str, SchedulerFactory],
         component_defaults: Optional[Dict[str, Dict[str, str]]] = None,
+        scheduler_params: Optional[Dict[str, object]] = None,
     ) -> None:
         """
         Creates a new runner instance.
@@ -63,7 +63,10 @@ def __init__(
             schedulers: a list of schedulers the runner can use.
         """
         self._name: str = name
-        self._schedulers = schedulers
+        self._scheduler_factories = schedulers_factories
+        self._scheduler_params: Dict[str, object] = scheduler_params or {}
+        # pyre-ignore[24]: Scheduler opts
+        self._scheduler_instances: Dict[str, Scheduler] = {}
         self._apps: Dict[AppHandle, AppDef] = {}
 
         # component_name -> map of component_fn_param_name -> user-specified default val encoded as str
@@ -96,7 +99,7 @@ def close(self) -> None:
         It is ok to call this method multiple times on the same runner object.
         """
 
-        for name, scheduler in self._schedulers.items():
+        for name, scheduler in self._scheduler_instances.items():
             scheduler.close()
 
     def run_component(
@@ -319,15 +322,15 @@ def run_opts(self) -> Dict[str, runopts]:
             A map of scheduler backend to its ``runopts``
         """
         return {
-            scheduler_backend: scheduler.run_opts()
-            for scheduler_backend, scheduler in self._schedulers.items()
+            name: self._scheduler(name).run_opts()
+            for name in self._scheduler_factories
         }
 
     def scheduler_backends(self) -> List[str]:
         """
         Returns a list of all supported scheduler backends.
         """
-        return list(self._schedulers.keys())
+        return list(self._scheduler_factories.keys())
 
     def status(self, app_handle: AppHandle) -> Optional[AppStatus]:
         """
@@ -557,10 +560,15 @@ def log_lines(
 
     # pyre-fixme: Scheduler opts
     def _scheduler(self, scheduler: str) -> Scheduler:
-        sched = self._schedulers.get(scheduler)
+        sched = self._scheduler_instances.get(scheduler)
+        if not sched:
+            factory = self._scheduler_factories.get(scheduler)
+            if factory:
+                sched = factory(self._name, **self._scheduler_params)
+                self._scheduler_instances[scheduler] = sched
         if not sched:
             raise KeyError(
-                f"Undefined scheduler backend: {scheduler}. Use one of: {self._schedulers.keys()}"
+                f"Undefined scheduler backend: {scheduler}. Use one of: {self._scheduler_factories.keys()}"
             )
         return sched
 
@@ -586,7 +594,7 @@ def _scheduler_app_id(
         return scheduler, scheduler_backend, app_id
 
     def __repr__(self) -> str:
-        return f"Runner(name={self._name}, schedulers={self._schedulers}, apps={self._apps})"
+        return f"Runner(name={self._name}, schedulers={self._scheduler_factories}, apps={self._apps})"
 
 
 def get_runner(
@@ -633,5 +641,5 @@ def get_runner(
     if not name:
         name = "torchx"
 
-    schedulers = get_schedulers(session_name=name, **scheduler_params)
-    return Runner(name, schedulers, component_defaults)
+    schedulers = get_scheduler_factories()
+    return Runner(name, schedulers, component_defaults, scheduler_params=scheduler_params)

torchx/schedulers/__init__.py

@@ -5,35 +5,36 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
-from typing import Dict, Optional
+import importlib
+from typing import Dict, Mapping
 
-import torchx.schedulers.aws_batch_scheduler as aws_batch_scheduler
-import torchx.schedulers.docker_scheduler as docker_scheduler
-import torchx.schedulers.kubernetes_scheduler as kubernetes_scheduler
-import torchx.schedulers.local_scheduler as local_scheduler
-import torchx.schedulers.slurm_scheduler as slurm_scheduler
 from torchx.schedulers.api import Scheduler
 from torchx.util.entrypoints import load_group
 from typing_extensions import Protocol
 
+DEFAULT_SCHEDULER_MODULES: Mapping[str, str] = {
+    "local_docker": "torchx.schedulers.docker_scheduler",
+    "local_cwd": "torchx.schedulers.local_scheduler",
+    "slurm": "torchx.schedulers.slurm_scheduler",
+    "kubernetes": "torchx.schedulers.kubernetes_scheduler",
+    "aws_batch": "torchx.schedulers.aws_batch_scheduler",
+    "ray": "torchx.schedulers.ray_scheduler",
+}
+
 
 class SchedulerFactory(Protocol):
     # pyre-fixme: Scheduler opts
     def __call__(self, session_name: str, **kwargs: object) -> Scheduler:
         ...
 
 
-def _try_get_ray_scheduler() -> Optional[SchedulerFactory]:
-    try:
-        from torchx.schedulers.ray_scheduler import _has_ray  # @manual
-
-        if _has_ray:
-            import torchx.schedulers.ray_scheduler as ray_scheduler  # @manual
-
-            return ray_scheduler.create_scheduler
+def _defer_load_scheduler(path: str) -> SchedulerFactory:
+    # pyre-ignore[24]: Scheduler opts
+    def run(*args: object, **kwargs: object) -> Scheduler:
+        module = importlib.import_module(path)
+        return module.create_scheduler(*args, **kwargs)
 
-    except ImportError:  # pragma: no cover
-        return None
+    return run
 
 
 def get_scheduler_factories() -> Dict[str, SchedulerFactory]:
@@ -44,22 +45,13 @@ def get_scheduler_factories() -> Dict[str, SchedulerFactory]:
     The first scheduler in the dictionary is used as the default scheduler.
     """
 
-    default_schedulers: Dict[str, SchedulerFactory] = {
-        "local_docker": docker_scheduler.create_scheduler,
-        "local_cwd": local_scheduler.create_cwd_scheduler,
-        "slurm": slurm_scheduler.create_scheduler,
-        "kubernetes": kubernetes_scheduler.create_scheduler,
-        "aws_batch": aws_batch_scheduler.create_scheduler,
-    }
-
-    ray_scheduler_creator = _try_get_ray_scheduler()
-    if ray_scheduler_creator:
-        default_schedulers["ray"] = ray_scheduler_creator
+    default_schedulers: Dict[str, SchedulerFactory] = {}
+    for scheduler, path in DEFAULT_SCHEDULER_MODULES.items():
+        default_schedulers[scheduler] = _defer_load_scheduler(path)
 
     return load_group(
         "torchx.schedulers",
         default=default_schedulers,
-        ignore_missing=True,
     )
 
 

torchx/schedulers/local_scheduler.py

@@ -1063,7 +1063,7 @@ def __next__(self) -> str:
                     return line
 
 
-def create_cwd_scheduler(session_name: str, **kwargs: Any) -> LocalScheduler:
+def create_scheduler(session_name: str, **kwargs: Any) -> LocalScheduler:
     return LocalScheduler(
         session_name=session_name,
         cache_size=kwargs.get("cache_size", 100),