torchx-nightly 2025.8.5__py3-none-any.whl → 2026.1.11__py3-none-any.whl

torchx/runner/api.py

@@ -1,4 +1,3 @@
-#!/usr/bin/env python3
 # Copyright (c) Meta Platforms, Inc. and affiliates.
 # All rights reserved.
 #
@@ -25,6 +24,7 @@ from typing import (
     Type,
     TYPE_CHECKING,
     TypeVar,
+    Union,
 )
 
 from torchx.runner.events import log_event
@@ -42,6 +42,7 @@ from torchx.specs import (
     parse_app_handle,
     runopts,
     UnknownAppException,
+    Workspace,
 )
 from torchx.specs.finder import get_component
 from torchx.tracker.api import (
@@ -53,7 +54,7 @@ from torchx.tracker.api import (
 from torchx.util.session import get_session_id_or_create_new, TORCHX_INTERNAL_SESSION_ID
 
 from torchx.util.types import none_throws
-from torchx.workspace.api import PkgInfo, WorkspaceBuilder, WorkspaceMixin
+from torchx.workspace import WorkspaceMixin
 
 if TYPE_CHECKING:
     from typing_extensions import Self
@@ -129,9 +130,9 @@ class Runner:
     def _get_scheduler_params_from_env(self) -> Dict[str, str]:
         scheduler_params = {}
         for key, value in os.environ.items():
-            lower_case_key = key.lower()
-            if lower_case_key.startswith("torchx_"):
-                scheduler_params[lower_case_key.strip("torchx_")] = value
+            key = key.lower()
+            if key.startswith("torchx_"):
+                scheduler_params[key.removeprefix("torchx_")] = value
         return scheduler_params
 
     def __enter__(self) -> "Self":
@@ -164,25 +165,13 @@ class Runner:
         for scheduler in self._scheduler_instances.values():
             scheduler.close()
 
-    def build_standalone_workspace(
-        self,
-        workspace_builder: WorkspaceBuilder[S, T],
-        sync: bool = True,
-    ) -> PkgInfo[S]:
-        """
-        Build a standalone workspace for the given role.
-        This method is used to build a workspace for a role independent of the scheduler and
-        also enables asynchronous workspace building using the Role overrides.
-        """
-        return workspace_builder.build_workspace(sync)
-
     def run_component(
         self,
         component: str,
-        component_args: List[str],
+        component_args: Union[list[str], dict[str, Any]],
         scheduler: str,
         cfg: Optional[Mapping[str, CfgVal]] = None,
-        workspace: Optional[str] = None,
+        workspace: Optional[Union[Workspace, str]] = None,
         parent_run_id: Optional[str] = None,
     ) -> AppHandle:
         """
@@ -217,7 +206,7 @@ class Runner:
             ComponentNotFoundException: if the ``component_path`` is failed to resolve.
         """
 
-        with log_event("run_component", workspace=workspace) as ctx:
+        with log_event("run_component") as ctx:
             dryrun_info = self.dryrun_component(
                 component,
                 component_args,
@@ -228,7 +217,8 @@ class Runner:
             )
             handle = self.schedule(dryrun_info)
             app = none_throws(dryrun_info._app)
-            ctx._torchx_event.workspace = workspace
+
+            ctx._torchx_event.workspace = str(workspace)
             ctx._torchx_event.scheduler = none_throws(dryrun_info._scheduler)
             ctx._torchx_event.app_image = app.roles[0].image
             ctx._torchx_event.app_id = parse_app_handle(handle)[2]
@@ -238,10 +228,10 @@ class Runner:
     def dryrun_component(
         self,
         component: str,
-        component_args: List[str],
+        component_args: Union[list[str], dict[str, Any]],
         scheduler: str,
         cfg: Optional[Mapping[str, CfgVal]] = None,
-        workspace: Optional[str] = None,
+        workspace: Optional[Union[Workspace, str]] = None,
         parent_run_id: Optional[str] = None,
     ) -> AppDryRunInfo:
         """
@@ -249,10 +239,13 @@ class Runner:
         component, but just returns what "would" have run.
         """
         component_def = get_component(component)
+        args_from_cli = component_args if isinstance(component_args, list) else []
+        args_from_json = component_args if isinstance(component_args, dict) else {}
         app = materialize_appdef(
             component_def.fn,
-            component_args,
+            args_from_cli,
             self._component_defaults.get(component, None),
+            args_from_json,
         )
         return self.dryrun(
             app,
@@ -267,7 +260,7 @@ class Runner:
         app: AppDef,
         scheduler: str,
         cfg: Optional[Mapping[str, CfgVal]] = None,
-        workspace: Optional[str] = None,
+        workspace: Optional[Union[Workspace, str]] = None,
         parent_run_id: Optional[str] = None,
     ) -> AppHandle:
         """
@@ -280,9 +273,7 @@ class Runner:
             An application handle that is used to call other action APIs on the app.
         """
 
-        with log_event(
-            api="run", runcfg=json.dumps(cfg) if cfg else None, workspace=workspace
-        ) as ctx:
+        with log_event(api="run") as ctx:
             dryrun_info = self.dryrun(
                 app,
                 scheduler,
@@ -291,10 +282,15 @@ class Runner:
                 parent_run_id=parent_run_id,
             )
             handle = self.schedule(dryrun_info)
-            ctx._torchx_event.scheduler = none_throws(dryrun_info._scheduler)
-            ctx._torchx_event.app_image = none_throws(dryrun_info._app).roles[0].image
-            ctx._torchx_event.app_id = parse_app_handle(handle)[2]
-            ctx._torchx_event.app_metadata = app.metadata
+
+            event = ctx._torchx_event
+            event.scheduler = scheduler
+            event.runcfg = json.dumps(cfg) if cfg else None
+            event.workspace = str(workspace)
+            event.app_id = parse_app_handle(handle)[2]
+            event.app_image = none_throws(dryrun_info._app).roles[0].image
+            event.app_metadata = app.metadata
+
             return handle
 
     def schedule(self, dryrun_info: AppDryRunInfo) -> AppHandle:
@@ -328,21 +324,22 @@ class Runner:
 
         """
         scheduler = none_throws(dryrun_info._scheduler)
-        app_image = none_throws(dryrun_info._app).roles[0].image
         cfg = dryrun_info._cfg
-        with log_event(
-            "schedule",
-            scheduler,
-            app_image=app_image,
-            runcfg=json.dumps(cfg) if cfg else None,
-        ) as ctx:
+        with log_event("schedule") as ctx:
             sched = self._scheduler(scheduler)
             app_id = sched.schedule(dryrun_info)
             app_handle = make_app_handle(scheduler, self._name, app_id)
+
             app = none_throws(dryrun_info._app)
             self._apps[app_handle] = app
-            _, _, app_id = parse_app_handle(app_handle)
-            ctx._torchx_event.app_id = app_id
+
+            event = ctx._torchx_event
+            event.scheduler = scheduler
+            event.runcfg = json.dumps(cfg) if cfg else None
+            event.app_id = app_id
+            event.app_image = none_throws(dryrun_info._app).roles[0].image
+            event.app_metadata = app.metadata
+
             return app_handle
 
     def name(self) -> str:
@@ -353,7 +350,7 @@ class Runner:
         app: AppDef,
         scheduler: str,
         cfg: Optional[Mapping[str, CfgVal]] = None,
-        workspace: Optional[str] = None,
+        workspace: Optional[Union[Workspace, str]] = None,
         parent_run_id: Optional[str] = None,
     ) -> AppDryRunInfo:
         """
@@ -422,52 +419,45 @@ class Runner:
             "dryrun",
             scheduler,
             runcfg=json.dumps(cfg) if cfg else None,
-            workspace=workspace,
-        ):
+            workspace=str(workspace),
+        ) as ctx:
             sched = self._scheduler(scheduler)
             resolved_cfg = sched.run_opts().resolve(cfg)
 
-            # early validation before build workspace
-            with log_event(
-                "pre_build_validate",
-                scheduler,
-            ):
-                sched._pre_build_validate(app, scheduler, resolved_cfg)
-
-            if workspace and isinstance(sched, WorkspaceMixin):
-                role = app.roles[0]
-                old_img = role.image
-
-                logger.info(f"Checking for changes in workspace `{workspace}`...")
-                logger.info(
-                    'To disable workspaces pass: --workspace="" from CLI or workspace=None programmatically.'
-                )
-                with log_event(
-                    "build_workspace_and_update_role",
-                    scheduler,
-                ) as ctx:
-                    sched.build_workspace_and_update_role(role, workspace, resolved_cfg)
-                    ctx._torchx_event.app_image = role.image
-                    ctx._torchx_event.workspace = 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 `{workspace}` for role[0]={role.name}."
-                    )
-                else:
-                    logger.info(
-                        f"Reusing original image `{old_img}` for role[0]={role.name}."
-                        " Either a patch was built or no changes to workspace was detected."
+            sched._pre_build_validate(app, scheduler, resolved_cfg)
+
+            if isinstance(sched, WorkspaceMixin):
+                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(
+                            "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
                     )
 
-            with log_event(
-                "validate",
-                scheduler,
-            ):
-                sched._validate(app, scheduler, resolved_cfg)
+                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:
@@ -486,6 +476,27 @@ class Runner:
         """
         return self._scheduler(scheduler).run_opts()
 
+    def cfg_from_str(self, scheduler: str, *cfg_literal: str) -> Mapping[str, CfgVal]:
+        """
+        Convenience function around the scheduler's ``runopts.cfg_from_str()`` method.
+
+        Usage:
+
+        .. doctest::
+
+            from torchx.runner import get_runner
+
+            runner = get_runner()
+            cfg = runner.cfg_from_str("local_cwd", "log_dir=/tmp/foobar", "prepend_cwd=True")
+            assert cfg == {"log_dir": "/tmp/foobar", "prepend_cwd": True, "auto_set_cuda_visible_devices": False}
+        """
+
+        opts = self._scheduler(scheduler).run_opts()
+        cfg = {}
+        for cfg_str in cfg_literal:
+            cfg.update(opts.cfg_from_str(cfg_str))
+        return cfg
+
     def scheduler_backends(self) -> List[str]:
         """
         Returns a list of all supported scheduler backends.
@@ -576,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/runner/config.py

@@ -73,7 +73,7 @@ CLI Usage
 
 #. In addition, it is possible to specify a different config other than .torchxconfig to
    load at runtime. Requirements are that the config path is specified by enviornment
-   variable TORCHX_CONFIG. It also disables hierarchy loading configs from multiple
+   variable TORCHXCONFIG. It also disables hierarchy loading configs from multiple
    directories as the cases otherwise.
 
 #. User level .torchxconfig
@@ -494,6 +494,8 @@ def find_configs(dirs: Optional[Iterable[str]] = None) -> List[str]:
 
     config = os.getenv(ENV_TORCHXCONFIG)
     if config is not None:
+        if not config:
+            return []
         configfile = Path(config)
         if not configfile.is_file():
             raise FileNotFoundError(

torchx/runner/events/__init__.py

@@ -33,8 +33,9 @@ from torchx.util.session import get_session_id_or_create_new
 
 from .api import SourceType, TorchxEvent  # noqa F401
 
-# pyre-fixme[9]: _events_logger is a global variable
-_events_logger: logging.Logger = None
+_events_logger: Optional[logging.Logger] = None
+
+log: logging.Logger = logging.getLogger(__name__)
 
 
 def _get_or_create_logger(destination: str = "null") -> logging.Logger:
@@ -51,19 +52,28 @@ def _get_or_create_logger(destination: str = "null") -> logging.Logger:
             a new logger if None provided.
     """
     global _events_logger
+
     if _events_logger:
         return _events_logger
-    logging_handler = get_logging_handler(destination)
-    logging_handler.setLevel(logging.DEBUG)
-    _events_logger = logging.getLogger(f"torchx-events-{destination}")
-    # Do not propagate message to the root logger
-    _events_logger.propagate = False
-    _events_logger.addHandler(logging_handler)
-    return _events_logger
+    else:
+        logging_handler = get_logging_handler(destination)
+        logging_handler.setLevel(logging.DEBUG)
+        _events_logger = logging.getLogger(f"torchx-events-{destination}")
+        # Do not propagate message to the root logger
+        _events_logger.propagate = False
+        _events_logger.addHandler(logging_handler)
+
+        assert _events_logger  # make type-checker happy
+        return _events_logger
 
 
 def record(event: TorchxEvent, destination: str = "null") -> None:
-    _get_or_create_logger(destination).info(event.serialize())
+    try:
+        serialized_event = event.serialize()
+    except Exception:
+        log.exception("failed to serialize event, will not record event")
+    else:
+        _get_or_create_logger(destination).info(serialized_event)
 
 
 class log_event:

torchx/runner/events/api.py

@@ -29,7 +29,7 @@ class TorchxEvent:
         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.

torchx/schedulers/__init__.py

@@ -21,8 +21,6 @@ DEFAULT_SCHEDULER_MODULES: Mapping[str, str] = {
     "kubernetes_mcad": "torchx.schedulers.kubernetes_mcad_scheduler",
     "aws_batch": "torchx.schedulers.aws_batch_scheduler",
     "aws_sagemaker": "torchx.schedulers.aws_sagemaker_scheduler",
-    "gcp_batch": "torchx.schedulers.gcp_batch_scheduler",
-    "ray": "torchx.schedulers.ray_scheduler",
     "lsf": "torchx.schedulers.lsf_scheduler",
 }
 
@@ -51,15 +49,14 @@ def get_scheduler_factories(
     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(
-        group,
-        default=default_schedulers,
-        skip_defaults=skip_defaults,
-    )
+    return load_group(group, default=default_schedulers)
 
 
 def get_default_scheduler_name() -> str:

torchx/schedulers/api.py

@@ -1,4 +1,3 @@
-#!/usr/bin/env python3
 # Copyright (c) Meta Platforms, Inc. and affiliates.
 # All rights reserved.
 #
@@ -16,14 +15,16 @@ from typing import Generic, Iterable, List, Optional, TypeVar
 
 from torchx.specs import (
     AppDef,
+    AppDryRunInfo,
     AppState,
     NONE,
     NULL_RESOURCE,
     Role,
     RoleStatus,
     runopts,
+    Workspace,
 )
-from torchx.workspace.api import WorkspaceMixin
+from torchx.workspace import WorkspaceMixin
 
 
 DAYS_IN_2_WEEKS = 14
@@ -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[str] = None,
+        workspace: str | Workspace | None = None,
     ) -> str:
         """
         Submits the application to be run by the scheduler.
@@ -144,16 +143,20 @@ class Scheduler(abc.ABC, Generic[T, A, D]):
         # 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)
 
     @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
@@ -169,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
@@ -183,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:
@@ -260,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,
@@ -350,19 +393,17 @@ 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 isinstance(app, AppDef):
-            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"
-                    )
+        for role in app.roles:
+            if role.resource == NULL_RESOURCE:
+                raise ValueError(
+                    f"No resource for role: {role.image}. Did you forget to attach resource to the role"
+                )
 
 
 def filter_regex(regex: str, data: Iterable[str]) -> Iterable[str]: