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

torchx/runner/events/api.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 json
 from dataclasses import asdict, dataclass
 from enum import Enum
-from typing import Optional, Union
+from typing import Dict, Optional, Union
 
 
 class SourceType(str, Enum):
@@ -23,15 +25,18 @@ class TorchxEvent:
     The class represents the event produced by ``torchx.runner`` api calls.
 
     Arguments:
-        session: Session id that was used to execute request.
+        session: Session id of the current run
         scheduler: Scheduler that is used to execute request
         api: Api name
         app_id: Unique id that is set by the underlying scheduler
-        image: Image/container bundle that is used to execute request.
+        app_image: Image/container bundle that is used to execute request.
+        app_metadata: metadata to the app (treatment of metadata is scheduler dependent)
         runcfg: Run config that was used to schedule app.
         source: Type of source the event is generated.
         cpu_time_usec: CPU time spent in usec
         wall_time_usec: Wall time spent in usec
+        start_epoch_time_usec: Epoch time in usec when runner event starts
+        Workspace: Track how different workspaces/no workspace affects build and scheduler
     """
 
     session: str
@@ -39,11 +44,17 @@ class TorchxEvent:
     api: str
     app_id: Optional[str] = None
     app_image: Optional[str] = None
+    app_metadata: Optional[Dict[str, str]] = None
     runcfg: Optional[str] = None
     raw_exception: Optional[str] = None
     source: SourceType = SourceType.UNKNOWN
     cpu_time_usec: Optional[int] = None
     wall_time_usec: Optional[int] = None
+    start_epoch_time_usec: Optional[int] = None
+    workspace: Optional[str] = None
+    exception_type: Optional[str] = None
+    exception_message: Optional[str] = None
+    exception_source_location: Optional[str] = None
 
     def __str__(self) -> str:
         return self.serialize()

torchx/runner/events/handlers.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 logging
 from typing import Dict
 

torchx/runtime/tracking/__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
+
 """
 .. note:: EXPERIMENTAL, USE AT YOUR OWN RISK, APIs SUBJECT TO CHANGE
 

torchx/runtime/tracking/api.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 abc
 import json
 from typing import Dict, Union

torchx/schedulers/__init__.py

@@ -5,12 +5,13 @@
 # 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 importlib
-from typing import Dict, Mapping
+from typing import Mapping, Protocol
 
 from torchx.schedulers.api import Scheduler
 from torchx.util.entrypoints import load_group
-from typing_extensions import Protocol
 
 DEFAULT_SCHEDULER_MODULES: Mapping[str, str] = {
     "local_docker": "torchx.schedulers.docker_scheduler",
@@ -19,16 +20,14 @@ DEFAULT_SCHEDULER_MODULES: Mapping[str, str] = {
     "kubernetes": "torchx.schedulers.kubernetes_scheduler",
     "kubernetes_mcad": "torchx.schedulers.kubernetes_mcad_scheduler",
     "aws_batch": "torchx.schedulers.aws_batch_scheduler",
-    "gcp_batch": "torchx.schedulers.gcp_batch_scheduler",
-    "ray": "torchx.schedulers.ray_scheduler",
+    "aws_sagemaker": "torchx.schedulers.aws_sagemaker_scheduler",
     "lsf": "torchx.schedulers.lsf_scheduler",
 }
 
 
 class SchedulerFactory(Protocol):
     # pyre-fixme: Scheduler opts
-    def __call__(self, session_name: str, **kwargs: object) -> Scheduler:
-        ...
+    def __call__(self, session_name: str, **kwargs: object) -> Scheduler: ...
 
 
 def _defer_load_scheduler(path: str) -> SchedulerFactory:
@@ -40,22 +39,24 @@ 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.
     """
 
-    default_schedulers: Dict[str, SchedulerFactory] = {}
-    for scheduler, path in DEFAULT_SCHEDULER_MODULES.items():
-        default_schedulers[scheduler] = _defer_load_scheduler(path)
+    if skip_defaults:
+        default_schedulers = {}
+    else:
+        default_schedulers: dict[str, SchedulerFactory] = {}
+        for scheduler, path in DEFAULT_SCHEDULER_MODULES.items():
+            default_schedulers[scheduler] = _defer_load_scheduler(path)
 
-    return load_group(
-        "torchx.schedulers",
-        default=default_schedulers,
-    )
+    return load_group(group, default=default_schedulers)
 
 
 def get_default_scheduler_name() -> str:

torchx/schedulers/api.py

@@ -1,10 +1,11 @@
-#!/usr/bin/env python3
 # Copyright (c) Meta Platforms, Inc. and affiliates.
 # All rights reserved.
 #
 # 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 abc
 import re
 from dataclasses import dataclass, field
@@ -21,8 +22,9 @@ from torchx.specs import (
     Role,
     RoleStatus,
     runopts,
+    Workspace,
 )
-from torchx.workspace.api import WorkspaceMixin
+from torchx.workspace import WorkspaceMixin
 
 
 DAYS_IN_2_WEEKS = 14
@@ -41,7 +43,7 @@ class DescribeAppResponse:
     the status and description of the application as known by the scheduler.
     For some schedulers implementations this response object has necessary
     and sufficient information to recreate an ``AppDef`` object. For these types
-    of schedulers, the user can re-``run()`` the recreted application. Otherwise
+    of schedulers, the user can re-``run()`` the recreated application. Otherwise
     the user can only call non-creating methods (e.g. ``wait()``, ``status()``,
     etc).
 
@@ -59,6 +61,7 @@ class DescribeAppResponse:
     msg: str = NONE
     structured_error_msg: str = NONE
     ui_url: Optional[str] = None
+    metadata: dict[str, str] = field(default_factory=dict)
 
     roles_statuses: List[RoleStatus] = field(default_factory=list)
     roles: List[Role] = field(default_factory=list)
@@ -83,6 +86,7 @@ class ListAppResponse:
     app_id: str
     state: AppState
     app_handle: str = "<NOT_SET>"
+    name: str = ""
 
     # Implementing __hash__() makes ListAppResponse hashable which makes
     # it easier to check if a ListAppResponse object exists in a list of
@@ -126,7 +130,7 @@ class Scheduler(abc.ABC, Generic[T]):
         self,
         app: AppDef,
         cfg: T,
-        workspace: Optional[str] = None,
+        workspace: str | Workspace | None = None,
     ) -> str:
         """
         Submits the application to be run by the scheduler.
@@ -139,10 +143,14 @@ class Scheduler(abc.ABC, Generic[T]):
         # pyre-fixme: Generic cfg type passed to resolve
         resolved_cfg = self.run_opts().resolve(cfg)
         if workspace:
-            sched = self
-            assert isinstance(sched, WorkspaceMixin)
-            role = app.roles[0]
-            sched.build_workspace_and_update_role(role, workspace, resolved_cfg)
+            assert isinstance(self, WorkspaceMixin)
+
+            if isinstance(workspace, str):
+                workspace = Workspace.from_str(workspace)
+
+            app.roles[0].workspace = workspace
+            self.build_workspaces(app.roles, resolved_cfg)
+
         # pyre-fixme: submit_dryrun takes Generic type for resolved_cfg
         dryrun_info = self.submit_dryrun(app, resolved_cfg)
         return self.schedule(dryrun_info)
@@ -177,8 +185,10 @@ class Scheduler(abc.ABC, Generic[T]):
         resolved_cfg = self.run_opts().resolve(cfg)
         # pyre-fixme: _submit_dryrun takes Generic type for resolved_cfg
         dryrun_info = self._submit_dryrun(app, resolved_cfg)
+
         for role in app.roles:
             dryrun_info = role.pre_proc(self.backend, dryrun_info)
+
         dryrun_info._app = app
         dryrun_info._cfg = resolved_cfg
         return dryrun_info
@@ -253,6 +263,46 @@ class Scheduler(abc.ABC, Generic[T]):
             # do nothing if the app does not exist
             return
 
+    def delete(self, app_id: str) -> None:
+        """
+        Deletes the job information for the specified ``app_id`` from the
+        scheduler's data-plane. Basically "deep-purging" the job from the
+        scheduler's data-plane. Calling this API on a "live" job (e.g in a
+        non-terminal status such as PENDING or RUNNING) cancels the job.
+
+        Note that this API is only relevant for schedulers for which its
+        data-plane persistently stores the "JobDefinition" (which is often
+        versioned). AWS Batch and Kubernetes are examples of such schedulers.
+        On these schedulers, a finished job may fall out of the data-plane
+        (e.g. really old finished jobs get deleted) but the JobDefinition is
+        typically permanently stored. In this case, calling
+        :py:meth:`~cancel` would not delete the job definition.
+
+        In schedulers with no such feature (e.g. SLURM)
+        :py:meth:`~delete` is the same as :py:meth:`~cancel`, which is the
+        default implementation. Hence implementors of such schedulers need not
+        override this method.
+
+        .. warning::
+            Calling :py:meth:`~delete` on an ``app_id`` that has fallen out of
+            the scheduler's data-plane does nothing. The user is responsible for
+            manually tracking down and cleaning up any dangling resources related
+            to the job.
+        """
+        if self.exists(app_id):
+            self._delete_existing(app_id)
+
+    def _delete_existing(self, app_id: str) -> None:
+        """
+        Deletes the job information for the specified ``app_id`` from the
+        scheduler's data-plane. This method will only be called on an
+        application that exists.
+
+        The default implementation calls :py:meth:`~_cancel_existing` which is
+        appropriate for schedulers without persistent job definitions.
+        """
+        self._cancel_existing(app_id)
+
     def log_iter(
         self,
         app_id: str,
@@ -335,18 +385,24 @@ class Scheduler(abc.ABC, Generic[T]):
             f"{self.__class__.__qualname__} does not support application log iteration"
         )
 
-    def _validate(self, app: AppDef, scheduler: str) -> None:
+    def _pre_build_validate(self, app: AppDef, scheduler: str, cfg: T) -> None:
         """
-        Validates whether application is consistent with the scheduler.
+        validates before workspace build whether application is consistent with the scheduler.
 
-        Raises:
-            ValueError: if application is not compatible with scheduler
+        Raises error if application is not compatible with scheduler
+        """
+        pass
+
+    def _validate(self, app: AppDef, scheduler: str, cfg: T) -> None:
+        """
+        Validates after workspace build whether application is consistent with the scheduler.
+
+        Raises error if application is not compatible with scheduler
         """
         for role in app.roles:
             if role.resource == NULL_RESOURCE:
                 raise ValueError(
-                    f"No resource for role: {role.image}."
-                    f" Did you forget to attach resource to the role"
+                    f"No resource for role: {role.image}. Did you forget to attach resource to the role"
                 )
 
 

torchx/schedulers/aws_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 AWS Batch scheduler which can be used to run TorchX
@@ -51,13 +53,13 @@ from typing import (
     Optional,
     Tuple,
     TYPE_CHECKING,
+    TypedDict,
     TypeVar,
 )
 
 import torchx
 import yaml
 from torchx.schedulers.api import (
-    AppDryRunInfo,
     DescribeAppResponse,
     filter_regex,
     ListAppResponse,
@@ -69,6 +71,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,
     CfgVal,
@@ -84,12 +87,13 @@ from torchx.specs.api import (
 from torchx.specs.named_resources_aws import instance_type_from_resource
 from torchx.util.types import none_throws
 from torchx.workspace.docker_workspace import DockerWorkspaceMixin
-from typing_extensions import TypedDict
 
 ENV_TORCHX_ROLE_IDX = "TORCHX_ROLE_IDX"
 
 ENV_TORCHX_ROLE_NAME = "TORCHX_ROLE_NAME"
 
+ENV_TORCHX_IMAGE = "TORCHX_IMAGE"
+
 DEFAULT_ROLE_NAME = "node"
 
 TAG_TORCHX_VER = "torchx.pytorch.org/version"
@@ -97,6 +101,37 @@ TAG_TORCHX_APPNAME = "torchx.pytorch.org/app-name"
 TAG_TORCHX_USER = "torchx.pytorch.org/user"
 
 
+def parse_ulimits(ulimits_list: list[str]) -> List[Dict[str, Any]]:
+    """
+    Parse ulimit string in format: name:softLimit:hardLimit
+    Multiple ulimits separated by commas.
+    """
+    if not ulimits_list:
+        return []
+
+    ulimits = []
+    for ulimit_str in ulimits_list:
+        if not ulimit_str.strip():
+            continue
+
+        parts = ulimit_str.strip().split(":")
+        if len(parts) != 3:
+            raise ValueError(
+                f"ulimit must be in format name:softLimit:hardLimit, got: {ulimit_str}"
+            )
+
+        name, soft_limit, hard_limit = parts
+        ulimits.append(
+            {
+                "name": name,
+                "softLimit": int(soft_limit) if soft_limit != "-1" else -1,
+                "hardLimit": int(hard_limit) if hard_limit != "-1" else -1,
+            }
+        )
+
+    return ulimits
+
+
 if TYPE_CHECKING:
     from docker import DockerClient
 
@@ -170,8 +205,13 @@ def resource_from_resource_requirements(
 
 
 def _role_to_node_properties(
-    role: Role, start_idx: int, privileged: bool = False
-) -> Dict[str, object]:
+    role: Role,
+    start_idx: int,
+    privileged: bool = False,
+    job_role_arn: Optional[str] = None,
+    execution_role_arn: Optional[str] = None,
+    ulimits: Optional[List[Dict[str, Any]]] = None,
+) -> Dict[str, Any]:
     role.mounts += get_device_mounts(role.resource.devices)
 
     mount_points = []
@@ -233,6 +273,7 @@ def _role_to_node_properties(
         "environment": [{"name": k, "value": v} for k, v in role.env.items()],
         "privileged": privileged,
         "resourceRequirements": resource_requirements_from_resource(role.resource),
+        **({"ulimits": ulimits} if ulimits else {}),
         "linuxParameters": {
             # To support PyTorch dataloaders we need to set /dev/shm to larger
             # than the 64M default.
@@ -245,7 +286,11 @@ def _role_to_node_properties(
         "mountPoints": mount_points,
         "volumes": volumes,
     }
-    if role.num_replicas > 1:
+    if job_role_arn:
+        container["jobRoleArn"] = job_role_arn
+    if execution_role_arn:
+        container["executionRoleArn"] = execution_role_arn
+    if role.num_replicas > 0:
         instance_type = instance_type_from_resource(role.resource)
         if instance_type is not None:
             container["instanceType"] = instance_type
@@ -336,7 +381,7 @@ def _thread_local_cache(f: Callable[[], T]) -> Callable[[], T]:
 
 
 @_thread_local_cache
-def _local_session() -> "boto3.session.Session":
+def _local_session() -> "boto3.session.Session":  # noqa: F821
     import boto3.session
 
     return boto3.session.Session()
@@ -349,6 +394,9 @@ class AWSBatchOpts(TypedDict, total=False):
     privileged: bool
     share_id: Optional[str]
     priority: int
+    job_role_arn: Optional[str]
+    execution_role_arn: Optional[str]
+    ulimits: Optional[list[str]]
 
 
 class AWSBatchScheduler(DockerWorkspaceMixin, Scheduler[AWSBatchOpts]):
@@ -492,12 +540,16 @@ class AWSBatchScheduler(DockerWorkspaceMixin, Scheduler[AWSBatchOpts]):
             role = values.apply(role)
             role.env[ENV_TORCHX_ROLE_IDX] = str(role_idx)
             role.env[ENV_TORCHX_ROLE_NAME] = str(role.name)
+            role.env[ENV_TORCHX_IMAGE] = role.image
 
             nodes.append(
                 _role_to_node_properties(
                     role,
                     start_idx=node_idx,
                     privileged=cfg["privileged"],
+                    job_role_arn=cfg.get("job_role_arn"),
+                    execution_role_arn=cfg.get("execution_role_arn"),
+                    ulimits=parse_ulimits(cfg.get("ulimits") or []),
                 )
             )
             node_idx += role.num_replicas
@@ -573,6 +625,21 @@ class AWSBatchScheduler(DockerWorkspaceMixin, Scheduler[AWSBatchOpts]):
             "Higher number (between 0 and 9999) means higher priority. "
             "This will only take effect if the job queue has a scheduling policy.",
         )
+        opts.add(
+            "job_role_arn",
+            type_=str,
+            help="The Amazon Resource Name (ARN) of the IAM role that the container can assume for AWS permissions.",
+        )
+        opts.add(
+            "execution_role_arn",
+            type_=str,
+            help="The Amazon Resource Name (ARN) of the IAM role that the ECS agent can assume for AWS permissions.",
+        )
+        opts.add(
+            "ulimits",
+            type_=List[str],
+            help="Ulimit settings in format: name:softLimit:hardLimit (multiple separated by commas)",
+        )
         return opts
 
     def _get_job_id(self, app_id: str) -> Optional[str]:
@@ -785,6 +852,8 @@ class AWSBatchScheduler(DockerWorkspaceMixin, Scheduler[AWSBatchOpts]):
                     startFromHead=True,
                     **args,
                 )
+            # pyre-fixme[66]: Exception handler type annotation `unknown` must
+            #  extend BaseException.
             except self._log_client.exceptions.ResourceNotFoundException:
                 return []  # noqa: B901
             if response["nextForwardToken"] == next_token: