(1/2)(torchx/cli) make internal schedulers a workspace scheduler; implement jetter_workspace_scheduler (#380)

Summary:
Pull Request resolved: #380

~~Part 1/2 (need an additional diff to make changes in penv python)~~ (done with this diff)

Does the following:

1. Adds a `torchx.workspace.Workspace` MixIn for Schedulers. Move Docker workspace build from docker_scheduler to torchx.workspace.DockerWorkspace

1. Implements a `JetterWorkspace` that uses jetter to build patches. Falls back to torchx custom patch building for BC.

1. Removes the now unnecessary fb/cmd_run (since we moved everything to workspaces.

1. Adds the `JetterWorkspace` mixin to mast, flow, qf, local_fb schedulers

1. Takes out `DockerWorkspace` from docker_scheduler.build_workspace_and_update_role()

1. Adds the `DockerWorkspace` mixin to docker and k8s schedulers

1. Gets rid of `WorkspaceScheduler` (in favor of `Workspace` mixin) and `WorkspaceRunner` (just builds an optional workspace parameter Runner APIs)

1. (bugfix) read from the config_dirs (home + cwd) to get component default args

Reviewed By: d4l3k

Differential Revision: D33971673

fbshipit-source-id: c44622b0d01f2d3bbe192b9e0ebf683777c47ec5

Author: Kiuk Chung

torchx/cli/cmd_run.py

@@ -18,9 +18,8 @@
 from pyre_extensions import none_throws
 from torchx.cli.cmd_base import SubCommand
 from torchx.cli.cmd_log import get_logs
-from torchx.runner import Runner, config
+from torchx.runner import Runner, config, get_runner
 from torchx.runner.config import load_sections
-from torchx.runner.workspaces import WorkspaceRunner, get_workspace_runner
 from torchx.schedulers import get_default_scheduler_name, get_scheduler_factories
 from torchx.specs import CfgVal
 from torchx.specs.finder import (
@@ -37,6 +36,8 @@
     "missing component name, either provide it from the CLI or in .torchxconfig"
 )
 
+CONFIG_DIRS = [str(Path.home()), str(Path.cwd())]
+
 logger: logging.Logger = logging.getLogger(__name__)
 
 
@@ -180,6 +181,12 @@ def add_arguments(self, subparser: argparse.ArgumentParser) -> None:
             default=False,
             help="Stream logs while waiting for app to finish.",
         )
+        subparser.add_argument(
+            "--workspace",
+            "--buck-target",
+            default=f"file://{Path.cwd()}",
+            help="local workspace to build/patch (buck-target of main binary if using buck)",
+        )
         subparser.add_argument(
             "component_name_and_args",
             nargs=argparse.REMAINDER,
@@ -197,57 +204,37 @@ def _run(self, runner: Runner, args: argparse.Namespace) -> None:
         run_opts = runner.run_opts()
         scheduler_opts = run_opts[args.scheduler]
         cfg = _parse_run_config(args.scheduler_args, scheduler_opts)
-        config_dirs = [str(Path.home()), str(Path.cwd())]
 
-        config.apply(scheduler=args.scheduler, cfg=cfg, dirs=config_dirs)
-
-        config_files = config.find_configs(dirs=config_dirs)
-        workspace = (
-            "file://" + os.path.dirname(config_files[0]) if config_files else None
-        )
+        config.apply(scheduler=args.scheduler, cfg=cfg, dirs=CONFIG_DIRS)
 
         component, component_args = _parse_component_name_and_args(
             args.component_name_and_args,
             none_throws(self._subparser),
-            dirs=config_dirs,
+            dirs=CONFIG_DIRS,
         )
-
         try:
             if args.dryrun:
-                if isinstance(runner, WorkspaceRunner):
-                    dryrun_info = runner.dryrun_component(
-                        component,
-                        component_args,
-                        args.scheduler,
-                        workspace=workspace,
-                        cfg=cfg,
-                    )
-                else:
-                    dryrun_info = runner.dryrun_component(
-                        component, component_args, args.scheduler, cfg=cfg
-                    )
+                dryrun_info = runner.dryrun_component(
+                    component,
+                    component_args,
+                    args.scheduler,
+                    workspace=args.workspace,
+                    cfg=cfg,
+                )
                 logger.info(
                     "\n=== APPLICATION ===\n"
                     f"{pformat(asdict(dryrun_info._app), indent=2, width=80)}"
                 )
 
                 logger.info("\n=== SCHEDULER REQUEST ===\n" f"{dryrun_info}")
             else:
-                if isinstance(runner, WorkspaceRunner):
-                    app_handle = runner.run_component(
-                        component,
-                        component_args,
-                        args.scheduler,
-                        workspace=workspace,
-                        cfg=cfg,
-                    )
-                else:
-                    app_handle = runner.run_component(
-                        component,
-                        component_args,
-                        args.scheduler,
-                        cfg=cfg,
-                    )
+                app_handle = runner.run_component(
+                    component,
+                    component_args,
+                    args.scheduler,
+                    workspace=args.workspace,
+                    cfg=cfg,
+                )
                 # DO NOT delete this line. It is used by slurm tests to retrieve the app id
                 print(app_handle)
 
@@ -278,8 +265,8 @@ def _run(self, runner: Runner, args: argparse.Namespace) -> None:
 
     def run(self, args: argparse.Namespace) -> None:
         os.environ["TORCHX_CONTEXT_NAME"] = os.getenv("TORCHX_CONTEXT_NAME", "cli_run")
-        component_defaults = load_sections(prefix="component")
-        with get_workspace_runner(component_defaults=component_defaults) as runner:
+        component_defaults = load_sections(prefix="component", dirs=CONFIG_DIRS)
+        with get_runner(component_defaults=component_defaults) as runner:
             self._run(runner, args)
 
     def _wait_and_exit(self, runner: Runner, app_handle: str, log: bool) -> None:

torchx/cli/test/cmd_run_test.py

@@ -22,7 +22,6 @@
     CmdRun,
     _parse_component_name_and_args,
     _parse_run_config,
-    logger,
 )
 from torchx.schedulers.local_scheduler import SignalException
 from torchx.specs import runopts
@@ -160,26 +159,6 @@ def test_run_dryrun(self, mock_runner_run: MagicMock) -> None:
         self.cmd_run.run(args)
         mock_runner_run.assert_not_called()
 
-    @patch("torchx.runner.workspaces.WorkspaceRunner._patch_app")
-    def test_runopts_not_found(self, patch_app: MagicMock) -> None:
-        args = self.parser.parse_args(
-            [
-                "--dryrun",
-                "--scheduler",
-                "kubernetes",
-                "utils.echo",
-                "--image",
-                "/tmp",
-            ]
-        )
-        with patch.object(logger, "error") as log_error:
-            with self.assertRaises(SystemExit) as e:
-                self.cmd_run.run(args)
-            msg = log_error.call_args[0][0]
-            self.assertTrue(
-                "Scheduler arg is incorrect or missing required option" in msg
-            )
-
     def _get_test_runopts(self) -> runopts:
         opts = runopts()
         opts.add("foo", type_=str, default="", help="")

torchx/components/component_test_base.py

@@ -48,11 +48,11 @@ def run_appdef_on_scheduler(
 
         runner = get_runner("test-runner")
         if dryrun:
-            dryrun_info = runner.dryrun(app_def, scheduler, cfg)
+            dryrun_info = runner.dryrun(app_def, scheduler, cfg=cfg)
             print(f"Dryrun info: {dryrun_info}")
             return dryrun_info
         else:
-            app_handle = runner.run(app_def, scheduler, cfg)
+            app_handle = runner.run(app_def, scheduler, cfg=cfg)
             print(f"AppHandle: {app_handle}")
             app_status = runner.wait(app_handle)
             print(f"Final status: {app_status}")

torchx/runner/api.py

@@ -30,6 +30,7 @@
     runopts,
 )
 from torchx.specs.finder import get_component
+from torchx.workspace.api import Workspace
 
 
 logger: logging.Logger = logging.getLogger(__name__)
@@ -103,6 +104,7 @@ def run_component(
         component_args: List[str],
         scheduler: SchedulerBackend,
         cfg: Optional[Mapping[str, CfgVal]] = None,
+        workspace: Optional[str] = None,
     ) -> AppHandle:
         """
         Runs a component.
@@ -136,7 +138,13 @@ def run_component(
             ComponentNotFoundException: if the ``component_path`` is failed to resolve.
         """
 
-        dryrun_info = self.dryrun_component(component, component_args, scheduler, cfg)
+        dryrun_info = self.dryrun_component(
+            component,
+            component_args,
+            scheduler,
+            cfg=cfg,
+            workspace=workspace,
+        )
         return self.schedule(dryrun_info)
 
     def dryrun_component(
@@ -145,6 +153,7 @@ def dryrun_component(
         component_args: List[str],
         scheduler: SchedulerBackend,
         cfg: Optional[Mapping[str, CfgVal]] = None,
+        workspace: Optional[str] = None,
     ) -> AppDryRunInfo:
         """
         Dryrun version of :py:func:`run_component`. Will not actually run the
@@ -156,13 +165,14 @@ def dryrun_component(
             component_args,
             self._component_defaults.get(component, None),
         )
-        return self.dryrun(app, scheduler, cfg)
+        return self.dryrun(app, scheduler, cfg=cfg, workspace=workspace)
 
     def run(
         self,
         app: AppDef,
         scheduler: SchedulerBackend,
         cfg: Optional[Mapping[str, CfgVal]] = None,
+        workspace: Optional[str] = None,
     ) -> AppHandle:
         """
         Runs the given application in the specified mode.
@@ -174,7 +184,7 @@ def run(
             An application handle that is used to call other action APIs on the app.
         """
 
-        dryrun_info = self.dryrun(app, scheduler, cfg)
+        dryrun_info = self.dryrun(app, scheduler, cfg=cfg, workspace=workspace)
         return self.schedule(dryrun_info)
 
     def schedule(self, dryrun_info: AppDryRunInfo) -> AppHandle:
@@ -229,6 +239,7 @@ def dryrun(
         app: AppDef,
         scheduler: SchedulerBackend,
         cfg: Optional[Mapping[str, CfgVal]] = None,
+        workspace: Optional[str] = None,
     ) -> AppDryRunInfo:
         """
         Dry runs an app on the given scheduler with the provided run configs.
@@ -265,6 +276,23 @@ def dryrun(
         cfg = cfg or dict()
         with log_event("dryrun", scheduler, runcfg=json.dumps(cfg) if cfg else None):
             sched = self._scheduler(scheduler)
+
+            if workspace and isinstance(sched, Workspace):
+                role = app.roles[0]
+                old_img = role.image
+                logger.info(
+                    f"Building workspace: {workspace} for role[0]: {role.name}, image: {old_img}"
+                )
+                sched.build_workspace_and_update_role(role, workspace)
+                logger.info("Done building workspace")
+                if old_img != role.image:
+                    logger.info(f"New image: {role.image} built from workspace")
+                else:
+                    logger.info(
+                        f"Reusing original image: {old_img} for role[0]: {role.name}."
+                        " Either a patch was built or no changes to workspace was detected."
+                    )
+
             sched._validate(app, scheduler)
             dryrun_info = sched.submit_dryrun(app, cfg)
             dryrun_info._scheduler = scheduler

torchx/runner/test/api_test.py

@@ -10,24 +10,21 @@
 import shutil
 import tempfile
 import unittest
+from typing import Mapping, Optional
 from unittest.mock import MagicMock, patch
 
 from pyre_extensions import none_throws
 from torchx.runner import Runner, get_runner
-from torchx.schedulers.api import DescribeAppResponse
+from torchx.schedulers.api import DescribeAppResponse, Scheduler
 from torchx.schedulers.local_scheduler import (
     LocalDirectoryImageProvider,
     LocalScheduler,
 )
 from torchx.schedulers.test.test_util import write_shell_script
-from torchx.specs.api import (
-    AppDef,
-    AppState,
-    Resource,
-    Role,
-    UnknownAppException,
-)
+from torchx.specs import AppDryRunInfo, CfgVal
+from torchx.specs.api import AppDef, AppState, Resource, Role, UnknownAppException
 from torchx.specs.finder import ComponentNotFoundException
+from torchx.workspace import Workspace
 
 
 GET_SCHEDULERS = "torchx.runner.api.get_schedulers"
@@ -136,6 +133,67 @@ def test_dryrun(self, _) -> None:
             scheduler_mock.submit_dryrun.assert_called_once_with(app, self.cfg)
             scheduler_mock._validate.assert_called_once()
 
+    def test_dryrun_with_workspace(self, _) -> None:
+        class TestScheduler(Scheduler, Workspace):
+            def __init__(self, build_new_img: bool):
+                Scheduler.__init__(self, backend="ignored", session_name="ignored")
+                self.build_new_img = build_new_img
+
+            def schedule(self, dryrun_info: AppDryRunInfo) -> str:
+                pass
+
+            def _submit_dryrun(
+                self, app: AppDef, cfg: Mapping[str, CfgVal]
+            ) -> AppDryRunInfo[AppDef]:
+                return AppDryRunInfo(app, lambda s: s)
+
+            def describe(self, app_id: str) -> Optional[DescribeAppResponse]:
+                pass
+
+            def _cancel_existing(self, app_id: str) -> None:
+                pass
+
+            def build_workspace_and_update_role(
+                self, role: Role, workspace: str
+            ) -> None:
+                if self.build_new_img:
+                    role.image = f"{role.image}_new"
+
+        with Runner(
+            name=SESSION_NAME,
+            schedulers={
+                "no-build-img": TestScheduler(build_new_img=False),
+                "builds-img": TestScheduler(build_new_img=True),
+            },
+        ) as runner:
+            app = AppDef(
+                "ignored",
+                roles=[
+                    Role(
+                        name="sleep",
+                        image="foo",
+                        resource=resource.SMALL,
+                        entrypoint="sleep",
+                        args=["1"],
+                    ),
+                    Role(
+                        name="sleep",
+                        image="bar",
+                        resource=resource.SMALL,
+                        entrypoint="sleep",
+                        args=["1"],
+                    ),
+                ],
+            )
+            dryruninfo = runner.dryrun(app, "no-build-img", workspace="//workspace")
+            self.assertEqual("foo", dryruninfo.request.roles[0].image)
+            self.assertEqual("bar", dryruninfo.request.roles[1].image)
+
+            dryruninfo = runner.dryrun(app, "builds-img", workspace="//workspace")
+            # workspace is attached to role[0] by default
+            self.assertEqual("foo_new", dryruninfo.request.roles[0].image)
+            self.assertEqual("bar", dryruninfo.request.roles[1].image)
+
     def test_describe(self, _) -> None:
         with Runner(
             name=SESSION_NAME, schedulers={"local_dir": self.scheduler}