torchx-nightly 2025.10.16__py3-none-any.whl → 2025.12.2__py3-none-any.whl

torchx/_version.py

@@ -0,0 +1,8 @@
+# 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
+BASE_VERSION = "0.8.0dev0"

torchx/cli/cmd_delete.py

@@ -0,0 +1,30 @@
+#!/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 argparse
+import logging
+
+from torchx.cli.cmd_base import SubCommand
+from torchx.runner import get_runner
+
+logger: logging.Logger = logging.getLogger(__name__)
+
+
+class CmdDelete(SubCommand):
+    def add_arguments(self, subparser: argparse.ArgumentParser) -> None:
+        subparser.add_argument(
+            "app_handle",
+            type=str,
+            help="torchx app handle (e.g. local://session-name/app-id)",
+        )
+
+    def run(self, args: argparse.Namespace) -> None:
+        app_handle = args.app_handle
+        runner = get_runner()
+        runner.delete(app_handle)

torchx/cli/main.py

@@ -16,6 +16,7 @@ import torchx
 from torchx.cli.cmd_base import SubCommand
 from torchx.cli.cmd_cancel import CmdCancel
 from torchx.cli.cmd_configure import CmdConfigure
+from torchx.cli.cmd_delete import CmdDelete
 from torchx.cli.cmd_describe import CmdDescribe
 from torchx.cli.cmd_list import CmdList
 from torchx.cli.cmd_log import CmdLog
@@ -37,6 +38,7 @@ def get_default_sub_cmds() -> Dict[str, SubCommand]:
         "builtins": CmdBuiltins(),
         "cancel": CmdCancel(),
         "configure": CmdConfigure(),
+        "delete": CmdDelete(),
         "describe": CmdDescribe(),
         "list": CmdList(),
         "log": CmdLog(),

torchx/runner/api.py

@@ -420,52 +420,44 @@ class Runner:
             scheduler,
             runcfg=json.dumps(cfg) if cfg else None,
             workspace=str(workspace),
-        ):
+        ) as ctx:
             sched = self._scheduler(scheduler)
             resolved_cfg = sched.run_opts().resolve(cfg)
 
             sched._pre_build_validate(app, scheduler, resolved_cfg)
 
             if isinstance(sched, WorkspaceMixin):
-                for i, role in enumerate(app.roles):
-                    role_workspace = role.workspace
-
-                    if i == 0 and workspace:
-                        # NOTE: torchx originally took workspace as a runner arg and only applied the workspace to role[0]
-                        #   later, torchx added support for the workspace attr in Role
-                        #   for BC, give precedence to the workspace argument over the workspace attr for role[0]
-                        if role_workspace:
-                            logger.info(
-                                f"Using workspace={workspace} over role[{i}].workspace={role_workspace} for role[{i}]={role.name}."
-                                " To use the role's workspace attr pass: --workspace='' from CLI or workspace=None programmatically."  # noqa: B950
-                            )
-                        role_workspace = workspace
-
-                    if role_workspace:
-                        old_img = role.image
+                if workspace:
+                    # NOTE: torchx originally took workspace as a runner arg and only applied the workspace to role[0]
+                    # later, torchx added support for the workspace attr in Role
+                    # for BC, give precedence to the workspace argument over the workspace attr for role[0]
+                    if app.roles[0].workspace:
                         logger.info(
-                            f"Checking for changes in workspace `{role_workspace}` for role[{i}]={role.name}..."
-                        )
-                        # TODO kiuk@ once we deprecate the `workspace` argument in runner APIs we can simplify the signature of
-                        #   build_workspace_and_update_role2() to just taking the role and resolved_cfg
-                        sched.build_workspace_and_update_role2(
-                            role, role_workspace, resolved_cfg
+                            "Overriding role[%d] (%s) workspace to `%s`"
+                            "To use the role's workspace attr pass: --workspace='' from CLI or workspace=None programmatically.",
+                            0,
+                            role.name,
+                            str(app.roles[0].workspace),
                         )
+                    app.roles[0].workspace = (
+                        Workspace.from_str(workspace)
+                        if isinstance(workspace, str)
+                        else workspace
+                    )
 
-                        if old_img != role.image:
-                            logger.info(
-                                f"Built new image `{role.image}` based on original image `{old_img}`"
-                                f" and changes in workspace `{role_workspace}` for role[{i}]={role.name}."
-                            )
-                        else:
-                            logger.info(
-                                f"Reusing original image `{old_img}` for role[{i}]={role.name}."
-                                " Either a patch was built or no changes to workspace was detected."
-                            )
+                sched.build_workspaces(app.roles, resolved_cfg)
 
             sched._validate(app, scheduler, resolved_cfg)
             dryrun_info = sched.submit_dryrun(app, resolved_cfg)
             dryrun_info._scheduler = scheduler
+
+            event = ctx._torchx_event
+            event.scheduler = scheduler
+            event.runcfg = json.dumps(cfg) if cfg else None
+            event.app_id = app.name
+            event.app_image = none_throws(dryrun_info._app).roles[0].image
+            event.app_metadata = app.metadata
+
             return dryrun_info
 
     def scheduler_run_opts(self, scheduler: str) -> runopts:
@@ -595,6 +587,16 @@ class Runner:
             if status is not None and not status.is_terminal():
                 scheduler.cancel(app_id)
 
+    def delete(self, app_handle: AppHandle) -> None:
+        """
+        Deletes the application from the scheduler.
+        """
+        scheduler, scheduler_backend, app_id = self._scheduler_app_id(app_handle)
+        with log_event("delete", scheduler_backend, app_id):
+            status = self.status(app_handle)
+            if status is not None:
+                scheduler.delete(app_id)
+
     def stop(self, app_handle: AppHandle) -> None:
         """
         See method ``cancel``.

torchx/schedulers/api.py

@@ -11,10 +11,11 @@ import re
 from dataclasses import dataclass, field
 from datetime import datetime
 from enum import Enum
-from typing import Generic, Iterable, List, Optional, TypeVar, Union
+from typing import Generic, Iterable, List, Optional, TypeVar
 
 from torchx.specs import (
     AppDef,
+    AppDryRunInfo,
     AppState,
     NONE,
     NULL_RESOURCE,
@@ -95,11 +96,9 @@ class ListAppResponse:
 
 
 T = TypeVar("T")
-A = TypeVar("A")
-D = TypeVar("D")
 
 
-class Scheduler(abc.ABC, Generic[T, A, D]):
+class Scheduler(abc.ABC, Generic[T]):
     """
     An interface abstracting functionalities of a scheduler.
     Implementers need only implement those methods annotated with
@@ -129,9 +128,9 @@ class Scheduler(abc.ABC, Generic[T, A, D]):
 
     def submit(
         self,
-        app: A,
+        app: AppDef,
         cfg: T,
-        workspace: Optional[Union[Workspace, str]] = None,
+        workspace: str | Workspace | None = None,
     ) -> str:
         """
         Submits the application to be run by the scheduler.
@@ -145,14 +144,19 @@ class Scheduler(abc.ABC, Generic[T, A, D]):
         resolved_cfg = self.run_opts().resolve(cfg)
         if workspace:
             assert isinstance(self, WorkspaceMixin)
-            self.build_workspace_and_update_role2(app.roles[0], workspace, resolved_cfg)
+
+            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)
 
     @abc.abstractmethod
-    def schedule(self, dryrun_info: D) -> str:
+    def schedule(self, dryrun_info: AppDryRunInfo) -> str:
         """
         Same as ``submit`` except that it takes an ``AppDryRunInfo``.
         Implementers are encouraged to implement this method rather than
@@ -168,7 +172,7 @@ class Scheduler(abc.ABC, Generic[T, A, D]):
 
         raise NotImplementedError()
 
-    def submit_dryrun(self, app: A, cfg: T) -> D:
+    def submit_dryrun(self, app: AppDef, cfg: T) -> AppDryRunInfo:
         """
         Rather than submitting the request to run the app, returns the
         request object that would have been submitted to the underlying
@@ -182,15 +186,15 @@ class Scheduler(abc.ABC, Generic[T, A, D]):
         # pyre-fixme: _submit_dryrun takes Generic type for resolved_cfg
         dryrun_info = self._submit_dryrun(app, resolved_cfg)
 
-        if isinstance(app, AppDef):
-            for role in app.roles:
-                dryrun_info = role.pre_proc(self.backend, dryrun_info)
+        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
 
     @abc.abstractmethod
-    def _submit_dryrun(self, app: A, cfg: T) -> D:
+    def _submit_dryrun(self, app: AppDef, cfg: T) -> AppDryRunInfo:
         raise NotImplementedError()
 
     def run_opts(self) -> runopts:
@@ -259,6 +263,46 @@ class Scheduler(abc.ABC, Generic[T, A, D]):
             # 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,
@@ -349,15 +393,12 @@ class Scheduler(abc.ABC, Generic[T, A, D]):
         """
         pass
 
-    def _validate(self, app: A, scheduler: str, cfg: T) -> None:
+    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
         """
-        if not isinstance(app, AppDef):
-            return
-
         for role in app.roles:
             if role.resource == NULL_RESOURCE:
                 raise ValueError(

torchx/schedulers/aws_batch_scheduler.py

@@ -381,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()
@@ -399,9 +399,7 @@ class AWSBatchOpts(TypedDict, total=False):
     ulimits: Optional[list[str]]
 
 
-class AWSBatchScheduler(
-    DockerWorkspaceMixin, Scheduler[AWSBatchOpts, AppDef, AppDryRunInfo[BatchJob]]
-):
+class AWSBatchScheduler(DockerWorkspaceMixin, Scheduler[AWSBatchOpts]):
     """
     AWSBatchScheduler is a TorchX scheduling interface to AWS Batch.
 

torchx/schedulers/aws_sagemaker_scheduler.py

@@ -157,7 +157,7 @@ def _merge_ordered(
 
 class AWSSageMakerScheduler(
     DockerWorkspaceMixin,
-    Scheduler[AWSSageMakerOpts, AppDef, AppDryRunInfo[AWSSageMakerJob]],
+    Scheduler[AWSSageMakerOpts],
 ):
     """
     AWSSageMakerScheduler is a TorchX scheduling interface to AWS SageMaker.

torchx/schedulers/docker_scheduler.py

@@ -129,9 +129,7 @@ class DockerOpts(TypedDict, total=False):
     privileged: bool
 
 
-class DockerScheduler(
-    DockerWorkspaceMixin, Scheduler[DockerOpts, AppDef, AppDryRunInfo[DockerJob]]
-):
+class DockerScheduler(DockerWorkspaceMixin, Scheduler[DockerOpts]):
     """
     DockerScheduler is a TorchX scheduling interface to Docker.
 

torchx/schedulers/kubernetes_mcad_scheduler.py

@@ -796,10 +796,7 @@ class KubernetesMCADOpts(TypedDict, total=False):
     network: Optional[str]
 
 
-class KubernetesMCADScheduler(
-    DockerWorkspaceMixin,
-    Scheduler[KubernetesMCADOpts, AppDef, AppDryRunInfo[KubernetesMCADJob]],
-):
+class KubernetesMCADScheduler(DockerWorkspaceMixin, Scheduler[KubernetesMCADOpts]):
     """
     KubernetesMCADScheduler is a TorchX scheduling interface to Kubernetes.