torchx-nightly 2024.2.11__py3-none-any.whl → 2025.1.14__py3-none-any.whl

torchx/schedulers/docker_scheduler.py

@@ -4,9 +4,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 fnmatch
 import logging
 import os.path
+import re
 import tempfile
 from dataclasses import dataclass
 from datetime import datetime
@@ -121,6 +124,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 +220,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,8 +266,12 @@ 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)
@@ -278,6 +292,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 +307,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 +320,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 +372,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/gcp_batch_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 GCP Batch scheduler which can be used to run TorchX
@@ -205,14 +207,12 @@ class GCPBatchScheduler(Scheduler[GCPBatchOpts]):
             if cpu <= 0:
                 cpu = 1
             MILLI = 1000
-            # pyre-fixme[8]: Attribute has type `Field`; used as `int`.
             res.cpu_milli = cpu * MILLI
             memMB = resource.memMB
             if memMB < 0:
                 raise ValueError(
                     f"memMB should to be set to a positive value, got {memMB}"
                 )
-            # pyre-fixme[8]: Attribute has type `Field`; used as `int`.
             res.memory_mib = memMB
 
             # TODO support named resources
@@ -360,13 +360,11 @@ class GCPBatchScheduler(Scheduler[GCPBatchOpts]):
             return None
 
         gpu = 0
-        # pyre-fixme[16]: `Field` has no attribute `instances`.
         if len(job.allocation_policy.instances) != 0:
             gpu_type = job.allocation_policy.instances[0].policy.machine_type
             gpu = GPU_TYPE_TO_COUNT[gpu_type]
 
         roles = {}
-        # pyre-fixme[16]: `RepeatedField` has no attribute `__iter__`.
         for tg in job.task_groups:
             env = tg.task_spec.environment.variables
             role = env["TORCHX_ROLE_NAME"]
@@ -390,7 +388,6 @@ class GCPBatchScheduler(Scheduler[GCPBatchOpts]):
         # TODO map role/replica status
         desc = DescribeAppResponse(
             app_id=app_id,
-            # pyre-fixme[16]: `Field` has no attribute `state`.
             state=JOB_STATE[job.status.state.name],
             roles=list(roles.values()),
         )
@@ -415,8 +412,10 @@ class GCPBatchScheduler(Scheduler[GCPBatchOpts]):
             raise ValueError(f"app not found: {app_id}")
 
         job_uid = job.uid
-        filters = [f"labels.job_uid={job_uid}"]
-        filters.append(f"resource.labels.task_id:task/{job_uid}-group0-{k}")
+        filters = [
+            f"labels.job_uid={job_uid}",
+            f"labels.task_id:{job_uid}-group0-{k}",
+        ]
 
         if since is not None:
             filters.append(f'timestamp>="{str(since.isoformat())}"')
@@ -437,7 +436,7 @@ class GCPBatchScheduler(Scheduler[GCPBatchOpts]):
 
         logger = logging.Client().logger(BATCH_LOGGER_NAME)
         for entry in logger.list_entries(filter_=filter):
-            yield entry.payload
+            yield entry.payload + "\n"
 
     def _job_full_name_to_app_id(self, job_full_name: str) -> str:
         """
@@ -465,7 +464,7 @@ class GCPBatchScheduler(Scheduler[GCPBatchOpts]):
             for job in all_jobs
         ]
 
-    def _validate(self, app: AppDef, scheduler: str) -> None:
+    def _validate(self, app: AppDef, scheduler: str, cfg: GCPBatchOpts) -> None:
         # Skip validation step
         pass
 

torchx/schedulers/ids.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
+
 import os
 import random
 import struct

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
@@ -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
 

torchx/schedulers/kubernetes_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 scheduler which can be used to run TorchX
@@ -23,7 +25,7 @@ Install Volcano:
     kubectl apply -f https://raw.githubusercontent.com/volcano-sh/volcano/v1.6.0/installer/volcano-development.yaml
 
 See the
-`Volcano Quickstart <https://github.com/volcano-sh/volcano#quick-start-guide>`_
+`Volcano Quickstart <https://github.com/volcano-sh/volcano>`_
 for more information.
 """
 
@@ -167,6 +169,17 @@ ANNOTATION_ISTIO_SIDECAR = "sidecar.istio.io/inject"
 
 LABEL_INSTANCE_TYPE = "node.kubernetes.io/instance-type"
 
+# role.env translates to static env variables in the yaml
+# {"FOO" : "bar"}               =====>      - name: FOO
+#                                             value: bar
+# unless this placeholder is present at the start of the role.env value then the env variable
+# in the yaml will be dynamically populated at runtime (placeholder is stripped out of the value)
+# {"FOO" : "[FIELD_PATH]bar"}   =====>      - name: FOO
+#                                             valueFrom:
+#                                               fieldRef:
+#                                                 fieldPath: bar
+PLACEHOLDER_FIELD_PATH = "[FIELD_PATH]"
+
 
 def sanitize_for_serialization(obj: object) -> object:
     from kubernetes import client
@@ -181,7 +194,9 @@ def role_to_pod(name: str, role: Role, service_account: Optional[str]) -> "V1Pod
         V1ContainerPort,
         V1EmptyDirVolumeSource,
         V1EnvVar,
+        V1EnvVarSource,
         V1HostPathVolumeSource,
+        V1ObjectFieldSelector,
         V1ObjectMeta,
         V1PersistentVolumeClaimVolumeSource,
         V1Pod,
@@ -301,9 +316,20 @@ def role_to_pod(name: str, role: Role, service_account: Optional[str]) -> "V1Pod
         image=role.image,
         name=name,
         env=[
-            V1EnvVar(
-                name=name,
-                value=value,
+            (
+                V1EnvVar(
+                    name=name,
+                    value_from=V1EnvVarSource(
+                        field_ref=V1ObjectFieldSelector(
+                            field_path=value.strip(PLACEHOLDER_FIELD_PATH)
+                        )
+                    ),
+                )
+                if value.startswith(PLACEHOLDER_FIELD_PATH)
+                else V1EnvVar(
+                    name=name,
+                    value=value,
+                )
             )
             for name, value in role.env.items()
         ],
@@ -635,7 +661,7 @@ class KubernetesScheduler(DockerWorkspaceMixin, Scheduler[KubernetesOpts]):
         )
         return AppDryRunInfo(req, repr)
 
-    def _validate(self, app: AppDef, scheduler: str) -> None:
+    def _validate(self, app: AppDef, scheduler: str, cfg: KubernetesOpts) -> None:
         # Skip validation step
         pass
 

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.
@@ -35,6 +37,7 @@ from typing import (
     Iterable,
     List,
     Optional,
+    Protocol,
     TextIO,
     Tuple,
 )
@@ -262,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:
     """
@@ -270,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)
@@ -608,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
 
@@ -674,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(
@@ -692,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,
@@ -1162,11 +1200,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.
@@ -486,7 +488,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
 

torchx/schedulers/ray/ray_driver.py

@@ -116,7 +116,7 @@ def load_actor_json(filename: str) -> List[RayActor]:
         return actors
 
 
-def create_placement_group_async(replicas: List[RayActor]) -> PlacementGroup:
+def create_placement_group_async(replicas: List[RayActor]) -> PlacementGroup:  # type: ignore
     """return a placement group reference, the corresponding placement group could be scheduled or pending"""
     bundles = []
     for replica in replicas:
@@ -148,12 +148,12 @@ class RayDriver:
         else:
             self.min_replicas = replicas[0].min_replicas  # pyre-ignore[8]
 
-        self.placement_groups: List[
-            PlacementGroup
-        ] = []  # all the placement groups, shall never change
-        self.actor_info_of_id: Dict[
-            str, ActorInfo
-        ] = {}  # store the info used to recover an actor
+        self.placement_groups: List[PlacementGroup] = (
+            []
+        )  # all the placement groups, shall never change
+        self.actor_info_of_id: Dict[str, ActorInfo] = (
+            {}
+        )  # store the info used to recover an actor
         self.active_tasks: List["ray.ObjectRef"] = []  # list of active tasks
 
         self.terminating: bool = False  # if the job has finished and being terminated

torchx/schedulers/ray_scheduler.py

@@ -318,7 +318,7 @@ if _has_ray:
 
             return AppDryRunInfo(job, repr)
 
-        def _validate(self, app: AppDef, scheduler: str) -> None:
+        def _validate(self, app: AppDef, scheduler: str, cfg: RayOpts) -> None:
             if scheduler != "ray":
                 raise ValueError(
                     f"An unknown scheduler backend '{scheduler}' has been passed to the Ray scheduler."

torchx/schedulers/slurm_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 Slurm scheduler which can be used to run TorchX
 components on a Slurm cluster.
@@ -470,7 +472,7 @@ class SlurmScheduler(DirWorkspaceMixin, Scheduler[SlurmOpts]):
 
         return AppDryRunInfo(req, repr)
 
-    def _validate(self, app: AppDef, scheduler: str) -> None:
+    def _validate(self, app: AppDef, scheduler: str, cfg: SlurmOpts) -> None:
         # Skip validation step for slurm
         pass
 

torchx/schedulers/streams.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
+
 import io
 import os
 import threading

torchx/specs/__init__.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 AppDef and related component definitions. These are
 used by components to define the apps which can then be launched via a TorchX
@@ -13,13 +15,7 @@ scheduler or pipeline adapter.
 import difflib
 from typing import Callable, Dict, Optional
 
-from torchx.specs.named_resources_aws import NAMED_RESOURCES as AWS_NAMED_RESOURCES
-from torchx.specs.named_resources_generic import (
-    NAMED_RESOURCES as GENERIC_NAMED_RESOURCES,
-)
-from torchx.util.entrypoints import load_group
-
-from .api import (  # noqa: F401 F403
+from torchx.specs.api import (
     ALL,
     AppDef,
     AppDryRunInfo,
@@ -50,7 +46,13 @@ from .api import (  # noqa: F401 F403
     UnknownSchedulerException,
     VolumeMount,
 )
-from .builders import make_app_handle, materialize_appdef, parse_mounts  # noqa
+from torchx.specs.builders import make_app_handle, materialize_appdef, parse_mounts
+
+from torchx.specs.named_resources_aws import NAMED_RESOURCES as AWS_NAMED_RESOURCES
+from torchx.specs.named_resources_generic import (
+    NAMED_RESOURCES as GENERIC_NAMED_RESOURCES,
+)
+from torchx.util.entrypoints import load_group
 
 GiB: int = 1024
 
@@ -181,3 +183,42 @@ def get_named_resources(res: str) -> Resource:
 
     """
     return named_resources[res]
+
+
+__all__ = [
+    "AppDef",
+    "AppDryRunInfo",
+    "AppHandle",
+    "AppState",
+    "AppStatus",
+    "BindMount",
+    "CfgVal",
+    "DeviceMount",
+    "get_type_name",
+    "is_terminal",
+    "macros",
+    "MISSING",
+    "NONE",
+    "NULL_RESOURCE",
+    "parse_app_handle",
+    "ReplicaState",
+    "ReplicaStatus",
+    "Resource",
+    "RetryPolicy",
+    "Role",
+    "RoleStatus",
+    "runopt",
+    "runopts",
+    "UnknownAppException",
+    "UnknownSchedulerException",
+    "InvalidRunConfigException",
+    "MalformedAppHandleException",
+    "VolumeMount",
+    "resource",
+    "get_named_resources",
+    "named_resources",
+    "make_app_handle",
+    "materialize_appdef",
+    "parse_mounts",
+    "ALL",
+]