torchx-nightly 2024.12.17__py3-none-any.whl → 2024.12.19__py3-none-any.whl

torchx/runner/config.py

@@ -197,7 +197,15 @@ def _configparser() -> configparser.ConfigParser:
 
 
 def _get_scheduler(name: str) -> Scheduler:
-    schedulers = get_scheduler_factories()
+    schedulers = {
+        **get_scheduler_factories(),
+        **(
+            get_scheduler_factories(
+                group="torchx.schedulers.orchestrator", skip_defaults=True
+            )
+            or {}
+        ),
+    }
     if name not in schedulers:
         raise ValueError(
             f"`{name}` is not a registered scheduler. Valid scheduler names: {schedulers.keys()}"
@@ -241,7 +249,16 @@ def dump(
     if schedulers:
         scheds = schedulers
     else:
-        scheds = get_scheduler_factories().keys()
+        scheduler_factories = {
+            **get_scheduler_factories(),
+            **(
+                get_scheduler_factories(
+                    group="torchx.schedulers.orchestrator", skip_defaults=True
+                )
+                or {}
+            ),
+        }
+        scheds = scheduler_factories.keys()
 
     config = _configparser()
     for sched_name in scheds:

torchx/schedulers/__init__.py

@@ -42,9 +42,11 @@ def _defer_load_scheduler(path: str) -> SchedulerFactory:
     return run
 
 
-def get_scheduler_factories() -> Dict[str, SchedulerFactory]:
+def get_scheduler_factories(
+    group: str = "torchx.schedulers", skip_defaults: bool = False
+) -> Dict[str, SchedulerFactory]:
     """
-    get_scheduler_factories returns all the available schedulers names and the
+    get_scheduler_factories returns all the available schedulers names under `group` and the
     method to instantiate them.
 
     The first scheduler in the dictionary is used as the default scheduler.
@@ -55,8 +57,9 @@ def get_scheduler_factories() -> Dict[str, SchedulerFactory]:
         default_schedulers[scheduler] = _defer_load_scheduler(path)
 
     return load_group(
-        "torchx.schedulers",
+        group,
         default=default_schedulers,
+        skip_defaults=skip_defaults,
     )
 
 

torchx/schedulers/local_scheduler.py

@@ -696,12 +696,11 @@ class LocalScheduler(Scheduler[LocalOpts]):
         log.debug(f"Running {role_name} (replica {replica_id}):\n {args_pfmt}")
         env = self._get_replica_env(replica_params)
 
-        proc = subprocess.Popen(
+        proc = self.run_local_job(
             args=replica_params.args,
             env=env,
             stdout=stdout_,
             stderr=stderr_,
-            start_new_session=True,
             cwd=replica_params.cwd,
         )
         return _LocalReplica(
@@ -714,6 +713,23 @@ class LocalScheduler(Scheduler[LocalOpts]):
             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,

torchx/util/entrypoints.py

@@ -51,8 +51,7 @@ def _defer_load_ep(ep: EntryPoint) -> object:
 
 # pyre-ignore-all-errors[3, 2]
 def load_group(
-    group: str,
-    default: Optional[Dict[str, Any]] = None,
+    group: str, default: Optional[Dict[str, Any]] = None, skip_defaults: bool = False
 ):
     """
     Loads all the entry points specified by ``group`` and returns
@@ -72,6 +71,7 @@ def load_group(
     1. ``load_group("foo")["bar"]("baz")`` -> equivalent to calling ``this.is.a_fn("baz")``
     1. ``load_group("food")`` -> ``None``
     1. ``load_group("food", default={"hello": this.is.c_fn})["hello"]("world")`` -> equivalent to calling ``this.is.c_fn("world")``
+    1. ``load_group("food", default={"hello": this.is.c_fn}, skip_defaults=True)`` -> ``None``
 
 
     If the entrypoint is a module (versus a function as shown above), then calling the ``deferred_load_fn``
@@ -90,6 +90,8 @@ def load_group(
     entrypoints = metadata.entry_points().select(group=group)
 
     if len(entrypoints) == 0:
+        if skip_defaults:
+            return None
         return default
 
     eps = {}

{torchx_nightly-2024.12.17.dist-info → torchx_nightly-2024.12.19.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: torchx-nightly
-Version: 2024.12.17
+Version: 2024.12.19
 Summary: TorchX SDK and Components
 Home-page: https://github.com/pytorch/torchx
 Author: TorchX Devs

{torchx_nightly-2024.12.17.dist-info → torchx_nightly-2024.12.19.dist-info}/RECORD

@@ -57,14 +57,14 @@ torchx/pipelines/kfp/adapter.py,sha256=5GeHULjb1kxG6wJtYVLpNkgdzUi4iYEaR42VFOwT6
 torchx/pipelines/kfp/version.py,sha256=mYBxd6bm4MeR34D--xo-JLQ9wHeAl_ZQLwbItCf9tr0,539
 torchx/runner/__init__.py,sha256=x8Sz7s_tLxPgJgvWIhK4ju9BNZU61uBFywGwDY6CqJs,315
 torchx/runner/api.py,sha256=uI365qobKvTP6-fEAiGrVlBzipYDhy1L8RoOJ1BQ3mA,29338
-torchx/runner/config.py,sha256=fTdCcf-MKlBg6MzXopF4W0hYyDDoPAuvZs2v2bKzwG0,17849
+torchx/runner/config.py,sha256=qpq_RfUSUykquAsEKOhDT3xBsa-m3wc6F6O8oP2QJ7k,18303
 torchx/runner/events/__init__.py,sha256=1_y0bojXl3FL0zlAj7BI4Dg5cXKXUmaa2jZbVH0EDUA,5268
 torchx/runner/events/api.py,sha256=pPLfowWTXtN_XcrEDNI45pE6Ijvdc_Gdxq76RduqgGE,2664
 torchx/runner/events/handlers.py,sha256=ThHCIJW21BfBgB7b6ftyjASJmD1KdizpjuTtsyqnvJs,522
 torchx/runtime/__init__.py,sha256=Wxje2BryzeQneFu5r6P9JJiEKG-_C9W1CcZ_JNrKT6g,593
 torchx/runtime/tracking/__init__.py,sha256=dYnAPnrXYREfPXkpHhdOFkcYIODWEbA13PdD-wLQYBo,3055
 torchx/runtime/tracking/api.py,sha256=SmUQyUKZqG3KlAhT7CJOGqRz1O274E4m63wQeOVq3CU,5472
-torchx/schedulers/__init__.py,sha256=M9SBZiNdE3KI_yc1-BiRtAetfTgtX07uKkuvGUeZQLU,2230
+torchx/schedulers/__init__.py,sha256=gwy1opmKOPzQ_Lqh2GY0chYycLmdissLfd4846mPEMY,2334
 torchx/schedulers/api.py,sha256=jAjfTKNgw26Rc5-FsZCSFt1xQIBNzXdUjUeOlllCuIU,14163
 torchx/schedulers/aws_batch_scheduler.py,sha256=8jRTmi5gtqyKUQfhMRqCQqNVM0drXdu8RS0xUfZby48,28091
 torchx/schedulers/aws_sagemaker_scheduler.py,sha256=x33J_tQFpRr_AUv5dWf1qxqG0dbmIjaqFHyrSnHWijw,20921
@@ -74,7 +74,7 @@ torchx/schedulers/gcp_batch_scheduler.py,sha256=o0t70s7fpGgn8MqhkIPWV3ZuF0kc5KdL
 torchx/schedulers/ids.py,sha256=3E-_vwVYC-8Tv8kjuY9-W7TbOe_-Laqd8a65uIN3hQY,1798
 torchx/schedulers/kubernetes_mcad_scheduler.py,sha256=riUfzSBega0JcCbV9NspU_a5A46AD-74Jgf7JPvA2Is,42925
 torchx/schedulers/kubernetes_scheduler.py,sha256=W_vVvQVGjNUoasN-JIsOTU9f0qjpds2u3GteNblssGM,28193
-torchx/schedulers/local_scheduler.py,sha256=q_EkCjsf66X0QQXSynP-G5wT4rGS5HLk4F68KdsaE4g,41444
+torchx/schedulers/local_scheduler.py,sha256=wOL-Js1fzEJY7RugNpFKu6BjmdHBgRf3llONFdwg-xQ,41845
 torchx/schedulers/lsf_scheduler.py,sha256=BkXKyWoVrCZznjfRFQ8oW3vzM0sBGkGmGUcSX71D-nE,17667
 torchx/schedulers/ray_scheduler.py,sha256=d10PtA-sEfHmDxF-GKog_Qro7FpIDCjSZbkKBdyj72E,17462
 torchx/schedulers/slurm_scheduler.py,sha256=A11XvI0-Wa7n3qc1IdrYcDaQEZhj0GiGf-PjHj0Byn8,19386
@@ -103,7 +103,7 @@ torchx/tracker/backend/fsspec.py,sha256=528xKryBE27Rm_OHD7r2R6fmVAclknBtoy1s034N
 torchx/util/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 torchx/util/cuda.py,sha256=-ZTa1WCLnY2WtSWAdWufLQqZSDCZfZsloBuiS84LIkU,1099
 torchx/util/datetime.py,sha256=hV6Sg0u5KTBe68yrmy_RGCC5su0i4Tb_mAYphWamiXI,405
-torchx/util/entrypoints.py,sha256=4rqmA81XYLj4Kk7GboJi0z78h4NIQxSrcOzDuuTwCkw,2725
+torchx/util/entrypoints.py,sha256=J5WSFtCFVZ7ZjlAvzYcdWgVSUKC3aVo1eVmJAbWhpx8,2894
 torchx/util/io.py,sha256=HNpWLcFUX0WTAP3CsdamHz--FR5A4kSdLCPfNqa2UkA,1807
 torchx/util/log_tee_helpers.py,sha256=wPyozmh9BOt_2d3Gxa0iNogwnjzwFitIIMBJOJ1arIw,6330
 torchx/util/modules.py,sha256=LRTuZRH5bbRr0ZaCtCtvKbgwhMoPsTx-GokWbCLGPdk,1131
@@ -115,9 +115,9 @@ torchx/workspace/__init__.py,sha256=FqN8AN4VhR1C_SBY10MggQvNZmyanbbuPuE-JCjkyUY,
 torchx/workspace/api.py,sha256=PtDkGTC5lX03pRoYpuMz2KCmM1ZOycRP1UknqvNb97Y,6341
 torchx/workspace/dir_workspace.py,sha256=npNW_IjUZm_yS5r-8hrRkH46ndDd9a_eApT64m1S1T4,2268
 torchx/workspace/docker_workspace.py,sha256=PFu2KQNVC-0p2aKJ-W_BKA9ZOmXdCY2ABEkCExp3udQ,10269
-torchx_nightly-2024.12.17.dist-info/LICENSE,sha256=WVHfXhFC0Ia8LTKt_nJVYobdqTJVg_4J3Crrfm2A8KQ,1721
-torchx_nightly-2024.12.17.dist-info/METADATA,sha256=SdUrkacXCA9W4QDsLuIxmmhb7_iWsYHigv2t--fGacI,6170
-torchx_nightly-2024.12.17.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
-torchx_nightly-2024.12.17.dist-info/entry_points.txt,sha256=T328AMXeKI3JZnnxfkEew2ZcMN1oQDtkXjMz7lkV-P4,169
-torchx_nightly-2024.12.17.dist-info/top_level.txt,sha256=pxew3bc2gsiViS0zADs0jb6kC5v8o_Yy_85fhHj_J1A,7
-torchx_nightly-2024.12.17.dist-info/RECORD,,
+torchx_nightly-2024.12.19.dist-info/LICENSE,sha256=WVHfXhFC0Ia8LTKt_nJVYobdqTJVg_4J3Crrfm2A8KQ,1721
+torchx_nightly-2024.12.19.dist-info/METADATA,sha256=TTMZujCKPVefmghH92lcE9KgEISFAz_4WNx7lXSmvGs,6170
+torchx_nightly-2024.12.19.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
+torchx_nightly-2024.12.19.dist-info/entry_points.txt,sha256=T328AMXeKI3JZnnxfkEew2ZcMN1oQDtkXjMz7lkV-P4,169
+torchx_nightly-2024.12.19.dist-info/top_level.txt,sha256=pxew3bc2gsiViS0zADs0jb6kC5v8o_Yy_85fhHj_J1A,7
+torchx_nightly-2024.12.19.dist-info/RECORD,,