(torchx/scheduler) make PenvImageProvider add fbpkg download dir to PATH and move PATH manipulation to _to_popen_request so that dryrun captures PATH set by the scheduler (#431)

Summary:
Pull Request resolved: #431

Resolves: T115068826 (enables DPP to be run as a separate reader role). Does two things:

1) (OSS + FB) moves PATH manipulation logic from `local_scheduler._popen()` into `local_scheduler._to_popen_request()` so that `dryrun` can better capture `PATH` env var that is manipulated by the local scheduler.

2) (FB-only) Makes PenvImageProvider add the fbpkg root directory to `PATH` of the role. This makes it possible to use entrypoints other than `python` (penv python) with PenvImageProvider, which is the case with DPP worker and master roles.

Reviewed By: aivanou

Differential Revision: D35066879

fbshipit-source-id: 259457a9fe2af7ca579099b91bbe75f5e2b51ce7

Author: Kiuk Chung

torchx/schedulers/local_scheduler.py

@@ -656,7 +656,6 @@ def _popen(
         role_name: RoleName,
         replica_id: int,
         replica_params: ReplicaParam,
-        prepend_cwd: bool,
     ) -> _LocalReplica:
         """
         Same as ``subprocess.Popen(**popen_kwargs)`` but is able to take ``stdout`` and ``stderr``
@@ -678,20 +677,8 @@ def _popen(
         # just make sure we override the parent's env vars with the user_defined ones
         env = os.environ.copy()
         env.update(replica_params.env)
-
-        # prepend extra_paths to PATH
-        env["PATH"] = _join_PATH(*self._extra_paths, env.get("PATH"))
-
-        cwd = replica_params.cwd
-        if cwd:
-            # if prepend_cwd is set, then prepend cwd to PATH
-            # making binaries in cwd take precedence to those in PATH
-            # otherwise append cwd to PATH so that the binaries in PATH
-            # precede over those in cwd
-            if prepend_cwd:
-                env["PATH"] = _join_PATH(cwd, env.get("PATH"))
-            else:
-                env["PATH"] = _join_PATH(env.get("PATH"), cwd)
+        # PATH is a special one, instead of overriding, append
+        env["PATH"] = _join_PATH(replica_params.env.get("PATH"), os.getenv("PATH"))
 
         # default to unbuffered python for faster responsiveness locally
         env.setdefault("PYTHONUNBUFFERED", "x")
@@ -763,8 +750,6 @@ def schedule(self, dryrun_info: AppDryRunInfo[PopenRequest]) -> str:
                     role_name,
                     replica_id,
                     replica_params,
-                    # pyre-ignore[6] cfg type checked by runopt.resolve()
-                    dryrun_info._cfg.get("prepend_cwd"),
                 )
                 local_app.add_replica(role_name, replica)
         self._apps[app_id] = local_app
@@ -789,7 +774,7 @@ def _get_gpu_device_count(self) -> int:
             gpus_info = [gpu_info for gpu_info in result.stdout.split("\n") if gpu_info]
             return len(gpus_info)
         except subprocess.CalledProcessError as e:
-            log.exception(f"Got exception while listing GPUs {e.stderr}")
+            log.exception(f"Got exception while listing GPUs: {e.stderr}")
             return 0
 
     def _set_cuda_visible_devices_for_role_replica(
@@ -813,18 +798,18 @@ def _update_env_cuda_visible_devices(
         app: AppDef,
         cfg: Mapping[str, CfgVal],
     ) -> None:
-        device_count = 0
         autoset = cfg.get("auto_set_cuda_visible_devices")
         if not autoset:
             return
+
         requested_gpus_total = sum(
             [role.resource.gpu * role.num_replicas for role in app.roles]
         )
         if requested_gpus_total <= 0:
             return
+
         device_count = self._get_gpu_device_count()
         if requested_gpus_total > device_count:
-            autoset = False
             log.warning(
                 "Cannot set `CUDA_VISIBLE_DEVICES` due to "
                 f"Available GPUs {device_count} less than requested {requested_gpus_total}"
@@ -860,6 +845,22 @@ def _to_popen_request(
 
             img_root = image_provider.fetch_role(role)
 
+            # prepend extra_paths to PATH
+            role.env["PATH"] = _join_PATH(*self._extra_paths, role.env.get("PATH"))
+            cwd = image_provider.get_cwd(role.image)
+
+            if cwd:
+                # if prepend_cwd is set, then prepend cwd to PATH
+                # making binaries in cwd take precedence to those in PATH
+                # otherwise append cwd to PATH so that the binaries in PATH
+                # precede over those in cwd
+                prepend_cwd = cfg.get("prepend_cwd")
+
+                if prepend_cwd:
+                    role.env["PATH"] = _join_PATH(cwd, role.env.get("PATH"))
+                else:
+                    role.env["PATH"] = _join_PATH(role.env.get("PATH"), cwd)
+
             for replica_id in range(role.num_replicas):
                 values = macros.Values(
                     img_root=img_root,

torchx/schedulers/test/local_scheduler_test.py

@@ -30,11 +30,11 @@
     LocalDirectoryImageProvider,
     LocalScheduler,
     ReplicaParam,
-    create_cwd_scheduler,
     _join_PATH,
+    create_cwd_scheduler,
     make_unique,
 )
-from torchx.specs.api import AppDef, AppState, Role, is_terminal, macros, Resource
+from torchx.specs.api import AppDef, AppState, Resource, Role, is_terminal, macros
 
 from .test_util import write_shell_script
 
@@ -336,24 +336,29 @@ def test_submit_inherit_parent_envs(self) -> None:
         os.environ, {"PATH": "/bin:/usr/bin", "TORCHELASTIC_ERROR_FILE": "ignored"}
     )
     def test_child_process_env_append_cwd(self, popen_mock: MagicMock) -> None:
-        ignored = 0
-        self.scheduler._popen(
-            role_name="ignored",
-            replica_id=ignored,
-            replica_params=ReplicaParam(
-                args=["a", "b"],
-                env={},
-                cwd="/home/bob",
-                stdout=None,
-                stderr=None,
-                combined=None,
+        popen_request = self.scheduler._to_popen_request(
+            app=AppDef(
+                name="_",
+                roles=[
+                    Role(name="foo", image=self.test_dir, env={"PATH": "/home/bob"}),
+                ],
             ),
-            prepend_cwd=False,
+            cfg={"prepend_cwd": False},
+        )
+
+        # the test scheduler is hooked up with LocalDirectoryImageProvider
+        # which treats the role's image as CWD
+        replica_params = popen_request.role_params["foo"][0]
+        self.assertEqual(f"/home/bob:{self.test_dir}", replica_params.env["PATH"])
+
+        # _popen actually adds in the os.environ[PATH]
+        self.scheduler._popen(
+            role_name="foo", replica_id=0, replica_params=replica_params
         )
 
         self.assertEqual(
             # for python 3.7 BC get call_args.kwargs by index
-            "/bin:/usr/bin:/home/bob",
+            f"/home/bob:{self.test_dir}:/bin:/usr/bin",
             popen_mock.call_args[1]["env"]["PATH"],
         )
 
@@ -362,61 +367,58 @@ def test_child_process_env_append_cwd(self, popen_mock: MagicMock) -> None:
         os.environ, {"PATH": "/bin:/usr/bin", "TORCHELASTIC_ERROR_FILE": "ignored"}
     )
     def test_child_process_env_prepend_cwd(self, popen_mock: MagicMock) -> None:
-        ignored = 0
-        self.scheduler._popen(
-            role_name="ignored",
-            replica_id=ignored,
-            replica_params=ReplicaParam(
-                args=["a", "b"],
-                env={},
-                cwd="/home/bob",
-                stdout=None,
-                stderr=None,
-                combined=None,
+        popen_request = self.scheduler._to_popen_request(
+            app=AppDef(
+                name="_",
+                roles=[
+                    Role(name="foo", image=self.test_dir, env={"PATH": "/home/bob"}),
+                ],
             ),
-            prepend_cwd=True,
+            cfg={"prepend_cwd": True},
+        )
+
+        # the test scheduler is hooked up with LocalDirectoryImageProvider
+        # which treats the role's image as CWD
+        replica_params = popen_request.role_params["foo"][0]
+        self.assertEqual(f"{self.test_dir}:/home/bob", replica_params.env["PATH"])
+
+        # _popen actually adds in the os.environ[PATH]
+        self.scheduler._popen(
+            role_name="foo", replica_id=0, replica_params=replica_params
         )
 
         self.assertEqual(
             # for python 3.7 BC get call_args.kwargs by index
-            "/home/bob:/bin:/usr/bin",
+            f"{self.test_dir}:/home/bob:/bin:/usr/bin",
             popen_mock.call_args[1]["env"]["PATH"],
         )
 
     @mock.patch("subprocess.Popen")
     @mock.patch.dict(os.environ, {"TORCHELASTIC_ERROR_FILE": "ignored"}, clear=True)
     def test_child_process_env_none(self, popen_mock: MagicMock) -> None:
-        ignored = 0
-        self.scheduler._popen(
-            role_name="ignored",
-            replica_id=ignored,
-            replica_params=ReplicaParam(
-                args=["a", "b"],
-                env={},
-                cwd="/home/bob",
-                stdout=None,
-                stderr=None,
-                combined=None,
+        popen_request = self.scheduler._to_popen_request(
+            app=AppDef(
+                name="_",
+                roles=[Role(name="foo", image=self.test_dir)],
             ),
-            prepend_cwd=True,
+            cfg={},
         )
-        self.assertEqual("/home/bob", popen_mock.call_args[1]["env"]["PATH"])
 
+        # the test scheduler is hooked up with LocalDirectoryImageProvider
+        # which treats the role's image as CWD
+        replica_params = popen_request.role_params["foo"][0]
+        self.assertEqual(f"{self.test_dir}", replica_params.env["PATH"])
+
+        # _popen actually adds in the os.environ[PATH]
         self.scheduler._popen(
-            role_name="ignored",
-            replica_id=ignored,
-            replica_params=ReplicaParam(
-                args=["a", "b"],
-                env={},
-                cwd=None,
-                stdout=None,
-                stderr=None,
-                combined=None,
-            ),
-            prepend_cwd=True,
+            role_name="foo", replica_id=0, replica_params=replica_params
+        )
+
+        self.assertEqual(
+            # for python 3.7 BC get call_args.kwargs by index
+            f"{self.test_dir}",
+            popen_mock.call_args[1]["env"]["PATH"],
         )
-        # for python 3.7 BC get call_args.kwargs by index
-        self.assertFalse(popen_mock.call_args[1]["env"]["PATH"])
 
     @mock.patch.dict(os.environ, {"FOO": "bar"})
     def test_submit_override_parent_env(self) -> None: