Merge branch 'main' into patch-1

Author: omrishiv

.github/workflows/slurm-local-integration-tests.yaml

@@ -6,8 +6,11 @@ on:
       - main
   pull_request:
 
+
 env:
-  SLURM_VERSION: 21.08.6
+  # slurm tag should be one of https://github.com/SchedMD/slurm/tags
+  SLURM_TAG: slurm-23-11-11-1
+  SLURM_VERSION: 23.11.11
 
 jobs:
   slurm:

.pyre_configuration

@@ -18,5 +18,6 @@
     "stubs"
   ],
   "strict": true,
-  "version": "0.0.101732536891"
+  "enable_strict_any_check": true,
+  "version": "0.0.101749035478"
 }

dev-requirements.txt

@@ -30,6 +30,7 @@ torchmetrics==1.6.3
 torchserve>=0.10.0
 torchtext==0.18.0
 torchvision==0.22.0
+typing-extensions
 ts==0.5.1
 ray[default]
 wheel

requirements.txt

@@ -1,6 +1,4 @@
-pyre-extensions
 docstring-parser>=0.8.1
-importlib-metadata
 pyyaml
 docker
 filelock

torchx/cli/cmd_list.py

@@ -21,6 +21,7 @@
 
 HANDLE_HEADER = "APP HANDLE"
 STATUS_HEADER = "APP STATUS"
+NAME_HEADER = "APP NAME"
 
 
 class CmdList(SubCommand):
@@ -39,5 +40,7 @@ def add_arguments(self, subparser: argparse.ArgumentParser) -> None:
     def run(self, args: argparse.Namespace) -> None:
         with get_runner() as runner:
             apps = runner.list(args.scheduler)
-            apps_data = [[app.app_handle, str(app.state)] for app in apps]
-            print(tabulate(apps_data, headers=[HANDLE_HEADER, STATUS_HEADER]))
+            apps_data = [[app.app_handle, app.name, str(app.state)] for app in apps]
+            print(
+                tabulate(apps_data, headers=[HANDLE_HEADER, NAME_HEADER, STATUS_HEADER])
+            )

torchx/components/dist.py

@@ -132,9 +132,6 @@ def spmd(
         j: {nnodes}x{nproc_per_node}. For GPU hosts omitting nproc_per_node will infer it from the GPU count on the host
         env: environment variables to be passed to the run (e.g. ENV1=v1,ENV2=v2,ENV3=v3)
         max_retries: the number of scheduler retries allowed
-        rdzv_port: the port on rank0's host to use for hosting the c10d store used for rendezvous.
-                   Only takes effect when running multi-node. When running single node, this parameter
-                   is ignored and a random free port is chosen.
         mounts: (for docker based runs only) mounts to mount into the worker environment/container
                 (ex. type=<bind/volume>,src=/host,dst=/job[,readonly]).
         debug: whether to run with preset debug flags enabled
@@ -174,6 +171,7 @@ def ddp(
     max_retries: int = 0,
     rdzv_port: int = 29500,
     rdzv_backend: str = "c10d",
+    rdzv_conf: Optional[str] = None,
     mounts: Optional[List[str]] = None,
     debug: bool = False,
     tee: int = 3,
@@ -208,6 +206,7 @@ def ddp(
                    Only takes effect when running multi-node. When running single node, this parameter
                    is ignored and a random free port is chosen.
         rdzv_backend: the rendezvous backend to use. Only takes effect when running multi-node.
+        rdzv_conf: the additional rendezvous configuration to use (ex. join_timeout=600,close_timeout=600,timeout=600).
         mounts: mounts to mount into the worker environment/container (ex. type=<bind/volume>,src=/host,dst=/job[,readonly]).
                 See scheduler documentation for more info.
         debug: whether to run with preset debug flags enabled
@@ -258,6 +257,7 @@ def ddp(
         "torchrun",
         "--rdzv_backend",
         rdzv_backend,
+        *(["--rdzv_conf", rdzv_conf] if rdzv_conf is not None else []),
         "--rdzv_endpoint",
         rdzv_endpoint,
         "--rdzv_id",

torchx/components/structured_arg.py

@@ -30,8 +30,6 @@
 from pathlib import Path
 from typing import Optional
 
-from pyre_extensions import none_throws
-
 from torchx import specs
 
 
@@ -148,7 +146,8 @@ def parse_from(
             if m:  # use the last module name
                 run_name = m.rpartition(".")[2]
             else:  # use script name w/ no extension
-                run_name = Path(none_throws(script)).stem
+                assert script, "`script` can't be `None` here due checks above"
+                run_name = Path(script).stem
         return StructuredNameArgument(
             experiment_name or default_experiment_name, run_name
         )

torchx/components/test/dist_test.py

@@ -41,8 +41,10 @@ def test_ddp_debug(self) -> None:
             self.assertEqual(env[k], v)
 
     def test_ddp_rdzv_backend_static(self) -> None:
-        app = ddp(script="foo.py", rdzv_backend="static")
+        rdzv_conf = "join_timeout=600,close_timeout=600,timeout=600"
+        app = ddp(script="foo.py", rdzv_backend="static", rdzv_conf=rdzv_conf)
         cmd = app.roles[0].args[1]
+        self.assertTrue(f"--rdzv_conf {rdzv_conf}" in cmd)
         self.assertTrue("--rdzv_backend static" in cmd)
         self.assertTrue("--node_rank" in cmd)
 

torchx/runner/api.py

@@ -14,7 +14,18 @@
 import warnings
 from datetime import datetime
 from types import TracebackType
-from typing import Any, Dict, Iterable, List, Mapping, Optional, Tuple, Type, TypeVar
+from typing import (
+    Any,
+    Dict,
+    Iterable,
+    List,
+    Mapping,
+    Optional,
+    Tuple,
+    Type,
+    TYPE_CHECKING,
+    TypeVar,
+)
 
 from torchx.runner.events import log_event
 from torchx.schedulers import get_scheduler_factories, SchedulerFactory
@@ -44,6 +55,9 @@
 from torchx.util.types import none_throws
 from torchx.workspace.api import PkgInfo, WorkspaceBuilder, WorkspaceMixin
 
+if TYPE_CHECKING:
+    from typing_extensions import Self
+
 from .config import get_config, get_configs
 
 logger: logging.Logger = logging.getLogger(__name__)
@@ -120,7 +134,7 @@ def _get_scheduler_params_from_env(self) -> Dict[str, str]:
                 scheduler_params[lower_case_key.strip("torchx_")] = value
         return scheduler_params
 
-    def __enter__(self) -> "Runner":
+    def __enter__(self) -> "Self":
         return self
 
     def __exit__(

torchx/runner/test/config_test.py

@@ -470,20 +470,17 @@ def test_dump_and_load_all_registered_schedulers(self) -> None:
         sfile = StringIO()
         dump(sfile)
 
-        scheduler_factories = {
-            **get_scheduler_factories(),
-            **(
-                get_scheduler_factories(
-                    group="torchx.schedulers.orchestrator", skip_defaults=True
-                )
-                or {}
-            ),
-        }
+        scheduler_factories = get_scheduler_factories()
 
         for sched_name, sched in scheduler_factories.items():
             sfile.seek(0)  # reset the file pos
             cfg = {}
-            load(scheduler=sched_name, f=sfile, cfg=cfg)
+            try:
+                load(scheduler=sched_name, f=sfile, cfg=cfg)
+            except ModuleNotFoundError:
+                # just test the ones that have been installed
+                continue
+
             for opt_name, _ in sched("test").run_opts():
                 self.assertTrue(
                     opt_name in cfg,