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

torchx/schedulers/docker_scheduler.py

@@ -4,18 +4,20 @@
 # 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
+
 import fnmatch
 import logging
 import os.path
+import re
 import tempfile
 from dataclasses import dataclass
 from datetime import datetime
-from typing import Any, Dict, Iterable, List, Optional, TYPE_CHECKING, Union
+from typing import Any, Dict, Iterable, List, Optional, TYPE_CHECKING, TypedDict, Union
 
 import torchx
 import yaml
 from torchx.schedulers.api import (
-    AppDryRunInfo,
     DescribeAppResponse,
     filter_regex,
     ListAppResponse,
@@ -27,6 +29,7 @@ from torchx.schedulers.devices import get_device_mounts
 from torchx.schedulers.ids import make_unique
 from torchx.specs.api import (
     AppDef,
+    AppDryRunInfo,
     AppState,
     BindMount,
     DeviceMount,
@@ -39,7 +42,6 @@ from torchx.specs.api import (
     VolumeMount,
 )
 from torchx.workspace.docker_workspace import DockerWorkspaceMixin
-from typing_extensions import TypedDict
 
 
 if TYPE_CHECKING:
@@ -82,6 +84,8 @@ LABEL_APP_ID: str = "torchx.pytorch.org/app-id"
 LABEL_ROLE_NAME: str = "torchx.pytorch.org/role-name"
 LABEL_REPLICA_ID: str = "torchx.pytorch.org/replica-id"
 
+ENV_TORCHX_IMAGE: str = "TORCHX_IMAGE"
+
 NETWORK = "torchx"
 
 
@@ -121,6 +125,8 @@ def ensure_network(client: Optional["DockerClient"] = None) -> None:
 
 class DockerOpts(TypedDict, total=False):
     copy_env: Optional[List[str]]
+    env: Optional[Dict[str, str]]
+    privileged: bool
 
 
 class DockerScheduler(DockerWorkspaceMixin, Scheduler[DockerOpts]):
@@ -215,9 +221,14 @@ class DockerScheduler(DockerWorkspaceMixin, Scheduler[DockerOpts]):
             for k in keys:
                 default_env[k] = os.environ[k]
 
+        env = cfg.get("env")
+        if env:
+            default_env.update(env)
+
         app_id = make_unique(app.name)
         req = DockerJob(app_id=app_id, containers=[])
-        rank0_name = f"{app_id}-{app.roles[0].name}-0"
+        # trim app_id and role name in case name is longer than 64 letters
+        rank0_name = f"{app_id[-30:]}-{app.roles[0].name[:30]}-0"
         for role in app.roles:
             mounts = []
             devices = []
@@ -256,14 +267,19 @@ class DockerScheduler(DockerWorkspaceMixin, Scheduler[DockerOpts]):
                     rank0_env="TORCHX_RANK0_HOST",
                 )
                 replica_role = values.apply(role)
-                name = f"{app_id}-{role.name}-{replica_id}"
-
+                # trim app_id and role name in case name is longer than 64 letters. Assume replica_id is less than 10_000. Remove invalid prefixes (https://github.com/moby/moby/blob/master/daemon/names/names.go#L6).
+                name = re.sub(
+                    r"^[^a-zA-Z0-9]+",
+                    "",
+                    f"{app_id[-30:]}-{role.name[:30]}-{replica_id}",
+                )
                 env = default_env.copy()
                 if replica_role.env:
                     env.update(replica_role.env)
 
                 # configure distributed host envs
                 env["TORCHX_RANK0_HOST"] = rank0_name
+                env[ENV_TORCHX_IMAGE] = replica_role.image
 
                 c = DockerContainer(
                     image=replica_role.image,
@@ -278,6 +294,7 @@ class DockerScheduler(DockerWorkspaceMixin, Scheduler[DockerOpts]):
                             LABEL_REPLICA_ID: str(replica_id),
                         },
                         "hostname": name,
+                        "privileged": cfg.get("privileged", False),
                         "network": NETWORK,
                         "mounts": mounts,
                         "devices": devices,
@@ -292,9 +309,9 @@ class DockerScheduler(DockerWorkspaceMixin, Scheduler[DockerOpts]):
                 if resource.memMB >= 0:
                     # To support PyTorch dataloaders we need to set /dev/shm to
                     # larger than the 64M default.
-                    c.kwargs["mem_limit"] = c.kwargs[
-                        "shm_size"
-                    ] = f"{int(resource.memMB)}m"
+                    c.kwargs["mem_limit"] = c.kwargs["shm_size"] = (
+                        f"{int(resource.memMB)}m"
+                    )
                 if resource.cpu >= 0:
                     c.kwargs["nano_cpus"] = int(resource.cpu * 1e9)
                 if resource.gpu > 0:
@@ -305,14 +322,14 @@ class DockerScheduler(DockerWorkspaceMixin, Scheduler[DockerOpts]):
                     c.kwargs["device_requests"] = [
                         DeviceRequest(
                             count=resource.gpu,
-                            capabilities=[["compute"]],
+                            capabilities=[["compute", "utility"]],
                         )
                     ]
                 req.containers.append(c)
 
         return AppDryRunInfo(req, repr)
 
-    def _validate(self, app: AppDef, scheduler: str) -> None:
+    def _validate(self, app: AppDef, scheduler: str, cfg: DockerOpts) -> None:
         # Skip validation step
         pass
 
@@ -357,6 +374,21 @@ class DockerScheduler(DockerWorkspaceMixin, Scheduler[DockerOpts]):
             default=None,
             help="list of glob patterns of environment variables to copy if not set in AppDef. Ex: FOO_*",
         )
+        opts.add(
+            "env",
+            type_=Dict[str, str],
+            default=None,
+            help="""environment variables to be passed to the run. The separator sign can be eiher comma or semicolon
+            (e.g. ENV1:v1,ENV2:v2,ENV3:v3 or ENV1:V1;ENV2:V2). Environment variables from env will be applied on top
+            of the ones from copy_env""",
+        )
+        opts.add(
+            "privileged",
+            type_=bool,
+            default=False,
+            help="If true runs the container with elevated permissions."
+            " Equivalent to running with `docker run --privileged`.",
+        )
         return opts
 
     def _get_app_state(self, container: "Container") -> AppState:

torchx/schedulers/ids.py

@@ -5,10 +5,12 @@
 # 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
+
 import os
-import random
 import struct
 
+
 START_CANDIDATES: str = "bcdfghjklmnpqrstvwxz"
 END_CANDIDATES: str = START_CANDIDATES + "012345679"
 
@@ -17,14 +19,19 @@ def make_unique(name: str, string_length: int = 0) -> str:
     """
     Appends a unique 64-bit string to the input argument.
 
+    Note that the unique string pulls entropy from `/dev/urandom` hence is not
+    affected by `random.seed()`
+
+    Args:
+        name: the name string to unique-ify
+        string_length: max length of the unique 64-bit string to append to the ``name``.
+          Default is 0, which returns the length of a randomly generated 64-bit string (typically 11-14 characters long).
+
     Returns:
-        string in format $name-$unique_suffix
+        string in format ``{name}-{unique_suffix}`
     """
-    return (
-        f"{name}-{random_id()}"
-        if string_length == 0
-        else f"{name}-{get_len_random_id(string_length)}"
-    )
+    max_length = None if string_length == 0 else string_length
+    return f"{name}-{random_id(max_length)}"
 
 
 def random_uint64() -> int:
@@ -34,13 +41,24 @@ def random_uint64() -> int:
     return struct.unpack("!Q", os.urandom(8))[0]
 
 
-def random_id() -> str:
+def random_id(max_length: int | None = None) -> str:
     """
     Generates an alphanumeric string ID that matches the requirements from
     https://kubernetes.io/docs/concepts/overview/working-with-objects/names/
+
+    Note that the unique string pulls entropy from `/dev/urandom` hence is not
+    affected by `random.seed()`
+
+    If ``max_length`` is provided, the returned ID will be at most that many characters long.
+
     """
+    # If a max_length is provided and is non-positive, return empty string
+    if max_length is not None and max_length <= 0:
+        return ""
+
     out = ""
     v = random_uint64()
+
     while v > 0:
         if out == "":
             candidates = START_CANDIDATES
@@ -50,21 +68,9 @@ def random_id() -> str:
         char = v % len(candidates)
         v = v // len(candidates)
         out += candidates[char]
-    return out
-
-
-def get_len_random_id(string_length: int) -> str:
-    """
-    Generates an alphanumeric string ID that matches the requirements from
-    https://kubernetes.io/docs/concepts/overview/working-with-objects/names/
-    """
-    out = ""
-    for i in range(string_length):
-        if out == "":
-            candidates = START_CANDIDATES
-        else:
-            candidates = END_CANDIDATES
 
-        out += random.choice(candidates)
+        if max_length is not None and len(out) >= max_length:
+            break
 
+    # NOTE: statistically the length of `out` is typically between 12-14 characters long
     return out

torchx/schedulers/kubernetes_mcad_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 Kubernetes_MCAD scheduler which can be used to run TorchX
@@ -15,8 +17,8 @@ Prerequisites
 
 TorchX Kubernetes_MCAD scheduler depends on AppWrapper + MCAD.
 
-Install MCAD: 
-See deploying Multi-Cluster-Application-Dispatcher guide 
+Install MCAD:
+See deploying Multi-Cluster-Application-Dispatcher guide
 https://github.com/project-codeflare/multi-cluster-app-dispatcher/blob/main/doc/deploy/deployment.md
 
 This implementation requires MCAD v1.34.1 or higher.
@@ -44,12 +46,12 @@ from typing import (
     Optional,
     Tuple,
     TYPE_CHECKING,
+    TypedDict,
 )
 
 import torchx
 import yaml
 from torchx.schedulers.api import (
-    AppDryRunInfo,
     DescribeAppResponse,
     filter_regex,
     ListAppResponse,
@@ -60,6 +62,7 @@ from torchx.schedulers.api import (
 from torchx.schedulers.ids import make_unique
 from torchx.specs.api import (
     AppDef,
+    AppDryRunInfo,
     AppState,
     BindMount,
     CfgVal,
@@ -76,7 +79,6 @@ from torchx.specs.api import (
 )
 
 from torchx.workspace.docker_workspace import DockerWorkspaceMixin
-from typing_extensions import TypedDict
 
 if TYPE_CHECKING:
     from docker import DockerClient
@@ -436,7 +438,7 @@ def mcad_svc(
                     target_port=int(service_port),
                 )
             ],
-            selector={"appwrapper.workload.codeflare.dev": svc_name},
+            selector={LABEL_UNIQUE_NAME: svc_name},
             session_affinity="None",
             type="ClusterIP",
         ),
@@ -598,7 +600,7 @@ def app_to_resource(
 
     """
     Create Service:
-    The selector will have the key 'appwrapper.workload.codeflare.dev', and the value will be 
+    The selector will have the key 'appwrapper.workload.codeflare.dev', and the value will be
     the appwrapper name
     """
 
@@ -990,7 +992,7 @@ class KubernetesMCADScheduler(DockerWorkspaceMixin, Scheduler[KubernetesMCADOpts
         if image_secret is not None and service_account is not None:
             msg = """Service Account and Image Secret names are both provided.
  Depending on the Service Account configuration, an ImagePullSecret may be defined in your Service Account.
- If this is the case, check service account and image secret configurations to understand the expected behavior for 
+ If this is the case, check service account and image secret configurations to understand the expected behavior for
  patched image push access."""
             warnings.warn(msg)
         namespace = cfg.get("namespace")
@@ -1031,7 +1033,7 @@ class KubernetesMCADScheduler(DockerWorkspaceMixin, Scheduler[KubernetesMCADOpts
         info._cfg = cfg
         return info
 
-    def _validate(self, app: AppDef, scheduler: str) -> None:
+    def _validate(self, app: AppDef, scheduler: str, cfg: KubernetesMCADOpts) -> None:
         # Skip validation step
         pass