torchx-nightly 2024.1.6__py3-none-any.whl → 2025.12.24__py3-none-any.whl

torchx/schedulers/local_scheduler.py

@@ -5,6 +5,8 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
+# pyre-strict
+
 """
 This contains the TorchX local scheduler which can be used to run TorchX
 components locally via subprocesses.
@@ -27,10 +29,21 @@ import warnings
 from dataclasses import asdict, dataclass
 from datetime import datetime
 from types import FrameType
-from typing import Any, BinaryIO, Callable, Dict, Iterable, List, Optional, TextIO
+from typing import (
+    Any,
+    BinaryIO,
+    Callable,
+    Dict,
+    Iterable,
+    List,
+    Optional,
+    Protocol,
+    TextIO,
+    Tuple,
+    TypedDict,
+)
 
 from torchx.schedulers.api import (
-    AppDryRunInfo,
     DescribeAppResponse,
     filter_regex,
     ListAppResponse,
@@ -40,10 +53,10 @@ from torchx.schedulers.api import (
 )
 from torchx.schedulers.ids import make_unique
 from torchx.schedulers.streams import Tee
+from torchx.specs import AppDryRunInfo
 from torchx.specs.api import AppDef, AppState, is_terminal, macros, NONE, Role, runopts
 
 from torchx.util.types import none_throws
-from typing_extensions import TypedDict
 
 log: logging.Logger = logging.getLogger(__name__)
 
@@ -252,6 +265,26 @@ AppName = str
 RoleName = str
 
 
+class PopenProtocol(Protocol):
+    """
+    Protocol wrapper around python's ``subprocess.Popen``. Keeps track of
+    the a list of interface methods that the process scheduled by the `LocalScheduler`
+    must implement.
+    """
+
+    @property
+    def pid(self) -> int: ...
+
+    @property
+    def returncode(self) -> int: ...
+
+    def wait(self, timeout: Optional[float] = None) -> int: ...
+
+    def poll(self) -> Optional[int]: ...
+
+    def kill(self) -> None: ...
+
+
 @dataclass
 class _LocalReplica:
     """
@@ -260,8 +293,7 @@ class _LocalReplica:
 
     role_name: RoleName
     replica_id: int
-    # pyre-fixme[24]: Generic type `subprocess.Popen` expects 1 type parameter.
-    proc: subprocess.Popen
+    proc: PopenProtocol
 
     # IO streams:
     # None means no log_dir (out to console)
@@ -598,7 +630,7 @@ class LocalScheduler(Scheduler[LocalOpts]):
         )
         return opts
 
-    def _validate(self, app: AppDef, scheduler: str) -> None:
+    def _validate(self, app: AppDef, scheduler: str, cfg: LocalOpts) -> None:
         # Skip validation step for local application
         pass
 
@@ -658,6 +690,55 @@ class LocalScheduler(Scheduler[LocalOpts]):
         as file name ``str`` rather than a file-like obj.
         """
 
+        stdout_, stderr_, combined_ = self._get_replica_output_handles(replica_params)
+
+        args_pfmt = pprint.pformat(asdict(replica_params), indent=2, width=80)
+        log.debug(f"Running {role_name} (replica {replica_id}):\n {args_pfmt}")
+        env = self._get_replica_env(replica_params)
+
+        proc = self.run_local_job(
+            args=replica_params.args,
+            env=env,
+            stdout=stdout_,
+            stderr=stderr_,
+            cwd=replica_params.cwd,
+        )
+        return _LocalReplica(
+            role_name,
+            replica_id,
+            proc,
+            stdout=stdout_,
+            stderr=stderr_,
+            combined=combined_,
+            error_file=env.get("TORCHELASTIC_ERROR_FILE", "<N/A>"),
+        )
+
+    def run_local_job(
+        self,
+        args: List[str],
+        env: Dict[str, str],
+        stdout: Optional[io.FileIO],
+        stderr: Optional[io.FileIO],
+        cwd: Optional[str] = None,
+    ) -> "subprocess.Popen[bytes]":
+        return subprocess.Popen(
+            args=args,
+            env=env,
+            stdout=stdout,
+            stderr=stderr,
+            start_new_session=True,
+            cwd=cwd,
+        )
+
+    def _get_replica_output_handles(
+        self,
+        replica_params: ReplicaParam,
+    ) -> Tuple[Optional[io.FileIO], Optional[io.FileIO], Optional[Tee]]:
+        """
+        Returns the stdout, stderr, and combined outputs of the replica.
+        If the combined output file is not specified, then the combined output is ``None``.
+        """
+
         stdout_ = self._get_file_io(replica_params.stdout)
         stderr_ = self._get_file_io(replica_params.stderr)
         combined_: Optional[Tee] = None
@@ -668,6 +749,15 @@ class LocalScheduler(Scheduler[LocalOpts]):
                 none_throws(replica_params.stdout),
                 none_throws(replica_params.stderr),
             )
+        return stdout_, stderr_, combined_
+
+    def _get_replica_env(
+        self,
+        replica_params: ReplicaParam,
+    ) -> Dict[str, str]:
+        """
+        Returns environment variables for the ``_LocalReplica``
+        """
 
         # inherit parent's env vars since 99.9% of the time we want this behavior
         # just make sure we override the parent's env vars with the user_defined ones
@@ -679,26 +769,7 @@ class LocalScheduler(Scheduler[LocalOpts]):
         # default to unbuffered python for faster responsiveness locally
         env.setdefault("PYTHONUNBUFFERED", "x")
 
-        args_pfmt = pprint.pformat(asdict(replica_params), indent=2, width=80)
-        log.debug(f"Running {role_name} (replica {replica_id}):\n {args_pfmt}")
-
-        proc = subprocess.Popen(
-            args=replica_params.args,
-            env=env,
-            stdout=stdout_,
-            stderr=stderr_,
-            start_new_session=True,
-            cwd=replica_params.cwd,
-        )
-        return _LocalReplica(
-            role_name,
-            replica_id,
-            proc,
-            stdout=stdout_,
-            stderr=stderr_,
-            combined=combined_,
-            error_file=env.get("TORCHELASTIC_ERROR_FILE", "<N/A>"),
-        )
+        return env
 
     def _get_app_log_dir(self, app_id: str, cfg: LocalOpts) -> str:
         """
@@ -1088,6 +1159,7 @@ class LogIterator:
             self._check_finished()  # check to see if app has finished running
 
             if os.path.isfile(self._log_file):
+                time.sleep(0.1)  # fix timing issue
                 self._log_fp = open(
                     self._log_file,
                     mode="rt",
@@ -1129,11 +1201,12 @@ def create_scheduler(
     session_name: str,
     cache_size: int = 100,
     extra_paths: Optional[List[str]] = None,
+    image_provider_class: Callable[[LocalOpts], ImageProvider] = CWDImageProvider,
     **kwargs: Any,
 ) -> LocalScheduler:
     return LocalScheduler(
         session_name=session_name,
-        image_provider_class=CWDImageProvider,
+        image_provider_class=image_provider_class,
         cache_size=cache_size,
         extra_paths=extra_paths,
     )

torchx/schedulers/lsf_scheduler.py

@@ -5,6 +5,8 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
+# pyre-strict
+
 """
 This contains the TorchX LSF scheduler which can be used to run TorchX
 components on a LSF cluster.
@@ -27,11 +29,10 @@ import subprocess
 import tempfile
 from dataclasses import dataclass
 from datetime import datetime
-from typing import Any, Dict, Iterable, List, Optional
+from typing import Any, Dict, Iterable, List, Optional, TypedDict
 
 import torchx
 from torchx.schedulers.api import (
-    AppDryRunInfo,
     DescribeAppResponse,
     filter_regex,
     ListAppResponse,
@@ -43,6 +44,7 @@ from torchx.schedulers.ids import make_unique
 from torchx.schedulers.local_scheduler import LogIterator
 from torchx.specs import (
     AppDef,
+    AppDryRunInfo,
     AppState,
     BindMount,
     DeviceMount,
@@ -55,7 +57,6 @@ from torchx.specs import (
     VolumeMount,
 )
 from torchx.util import shlex
-from typing_extensions import TypedDict
 
 JOB_STATE: Dict[str, AppState] = {
     "DONE": AppState.SUCCEEDED,
@@ -486,7 +487,7 @@ class LsfScheduler(Scheduler[LsfOpts]):
             subprocess.run(req.cmd, stdout=subprocess.PIPE, check=True)
         return req.app_id
 
-    def _validate(self, app: AppDef, scheduler: str) -> None:
+    def _validate(self, app: AppDef, scheduler: str, cfg: LsfOpts) -> None:
         # Skip validation step for lsf
         pass