cli: defer loading schedulers until used (#537)

Summary:
This updates the cli, schedulers and runners so that the schedulers are only imported when they need to be used. This means only the relevant scheduler is loaded which vastly improves responsiveness.

* `torchx --help`  900ms -> 300ms
* `torchx status local_cwd://` 1.38s to 840ms

Breaking changes:

* `schedulers.get_schedulers()` has been removed in favor of `schedulers.get_scheduler_factories()` since it's too dangerous
* `Runner.run_opts()` now takes a specific scheduler name instead of returning all scheduler runopts since that requires loading all schedulers.
*
get_schedulers is an internal interface since downstream users should be using the runner interface so changing this shouldn't be an issue

Runner.run_opts() is part of the user Runner interface but it's only really practical for the CLI so I doubt there's any OSS usage of it

Pull Request resolved: #537

Test Plan:
(torchx) tristanr@tristanr-arch2 ~/D/torchx (deferload)> time torchx --help
usage: torchx [-h] [--log_level LOG_LEVEL] [--version] {builtins,cancel,configure,describe,log,run,runopts,status} ...

torchx CLI

optional arguments:
  -h, --help            show this help message and exit
  --log_level LOG_LEVEL
                        Python logging log level
  --version             show program's version number and exit

sub-commands:
  Use the following commands to run operations, e.g.: torchx run ${JOB_NAME}

  {builtins,cancel,configure,describe,log,run,runopts,status}

________________________________________________________
Executed in  300.70 millis    fish           external
   usr time  280.99 millis  916.00 micros  280.07 millis
   sys time   16.46 millis    0.00 micros   16.46 millis

(torchx) tristanr@tristanr-arch2 ~/D/torchx (deferload)> time torchx status local_docker://torchx/sh-lsgcv92jm13ps
torchx 2022-06-24 15:55:58 INFO     AppDef:
  State: SUCCEEDED
  Num Restarts: -1
Roles:
 *sh[0]:SUCCEEDED
  sh[1]:SUCCEEDED
  sh[2]:SUCCEEDED

________________________________________________________
Executed in  507.45 millis    fish           external
   usr time  472.78 millis  926.00 micros  471.86 millis
   sys time   19.77 millis    0.00 micros   19.77 millis

Reviewed By: kurman

Differential Revision: D37432971

Pulled By: d4l3k

fbshipit-source-id: 52fd9bd9d8ec6accadac2128af877d0607b01176

Author: d4l3k

docs/source/schedulers.rst

@@ -32,7 +32,6 @@ Scheduler Functions
 .. automodule:: torchx.schedulers
 .. currentmodule:: torchx.schedulers
 
-.. autofunction:: get_schedulers
 .. autofunction:: get_scheduler_factories
 .. autofunction:: get_default_scheduler_name
 

docs/source/schedulers/local.rst

@@ -24,7 +24,7 @@ Image Providers
 Reference
 ~~~~~~~~~~~~
 
-.. autofunction:: create_cwd_scheduler
+.. autofunction:: create_scheduler
 
 .. autoclass:: LogIterator
    :members:

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_configure.py

@@ -11,7 +11,7 @@
 
 from torchx.cli.cmd_base import SubCommand
 from torchx.runner.config import dump
-from torchx.schedulers import get_schedulers
+from torchx.schedulers import get_scheduler_factories
 
 
 logger: logging.Logger = logging.getLogger(__name__)
@@ -42,7 +42,7 @@ def run(self, args: argparse.Namespace) -> None:
         if args.schedulers:
             schedulers = args.schedulers.split(",")
         else:
-            schedulers = get_schedulers(session_name="_").keys()
+            schedulers = get_scheduler_factories().keys()
 
         required_only = not args.all
 

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",
@@ -175,8 +174,7 @@ def _run(self, runner: Runner, args: argparse.Namespace) -> None:
                 " (e.g. `local_cwd`)"
             )
 
-        run_opts = runner.run_opts()
-        scheduler_opts = run_opts[args.scheduler]
+        scheduler_opts = runner.scheduler_run_opts(args.scheduler)
         cfg = scheduler_opts.cfg_from_str(args.scheduler_args)
         config.apply(scheduler=args.scheduler, cfg=cfg, dirs=CONFIG_DIRS)
 

torchx/cli/cmd_runopts.py

@@ -26,8 +26,12 @@ def add_arguments(self, subparser: argparse.ArgumentParser) -> None:
 
     def run(self, args: argparse.Namespace) -> None:
         filter = args.scheduler
-        run_opts = get_runner().run_opts()
-
-        for scheduler, opts in run_opts.items():
-            if not filter or scheduler == filter:
-                print(f"{GREEN}{scheduler}{ENDC}:\n{repr(opts)}\n")
+        with get_runner() as runner:
+            for scheduler in runner.scheduler_backends():
+                if filter and scheduler != filter:
+                    continue
+                try:
+                    opts = runner.scheduler_run_opts(scheduler)
+                    print(f"{GREEN}{scheduler}{ENDC}:\n{repr(opts)}\n")
+                except ModuleNotFoundError as e:
+                    print(f"{GREEN}{scheduler}{ENDC}: {e}\n")

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

@@ -59,7 +59,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()