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

torchx/runner/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 json
 import logging
 import os
@@ -12,7 +13,19 @@ import time
 import warnings
 from datetime import datetime
 from types import TracebackType
-from typing import Any, Dict, Iterable, List, Mapping, Optional, Tuple, Type
+from typing import (
+    Any,
+    Dict,
+    Iterable,
+    List,
+    Mapping,
+    Optional,
+    Tuple,
+    Type,
+    TYPE_CHECKING,
+    TypeVar,
+    Union,
+)
 
 from torchx.runner.events import log_event
 from torchx.schedulers import get_scheduler_factories, SchedulerFactory
@@ -29,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 (
@@ -37,9 +51,13 @@ from torchx.tracker.api import (
     ENV_TORCHX_TRACKERS,
     tracker_config_env_var_name,
 )
+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 WorkspaceMixin
+from torchx.workspace import WorkspaceMixin
+
+if TYPE_CHECKING:
+    from typing_extensions import Self
 
 from .config import get_config, get_configs
 
@@ -47,6 +65,8 @@ logger: logging.Logger = logging.getLogger(__name__)
 
 
 NONE: str = "<NONE>"
+S = TypeVar("S")
+T = TypeVar("T")
 
 
 def get_configured_trackers() -> Dict[str, Optional[str]]:
@@ -96,15 +116,26 @@ class Runner:
         """
         self._name: str = name
         self._scheduler_factories = scheduler_factories
-        self._scheduler_params: Dict[str, object] = scheduler_params or {}
-        # pyre-ignore[24]: Scheduler opts
+        self._scheduler_params: Dict[str, Any] = {
+            **(self._get_scheduler_params_from_env()),
+            **(scheduler_params or {}),
+        }
+        # pyre-fixme[24]: SchedulerOpts is a generic, and we don't have access to the corresponding type
         self._scheduler_instances: Dict[str, Scheduler] = {}
         self._apps: Dict[AppHandle, AppDef] = {}
 
         # component_name -> map of component_fn_param_name -> user-specified default val encoded as str
         self._component_defaults: Dict[str, Dict[str, str]] = component_defaults or {}
 
-    def __enter__(self) -> "Runner":
+    def _get_scheduler_params_from_env(self) -> Dict[str, str]:
+        scheduler_params = {}
+        for key, value in os.environ.items():
+            key = key.lower()
+            if key.startswith("torchx_"):
+                scheduler_params[key.removeprefix("torchx_")] = value
+        return scheduler_params
+
+    def __enter__(self) -> "Self":
         return self
 
     def __exit__(
@@ -131,16 +162,16 @@ class Runner:
         It is ok to call this method multiple times on the same runner object.
         """
 
-        for name, scheduler in self._scheduler_instances.items():
+        for scheduler in self._scheduler_instances.values():
             scheduler.close()
 
     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:
         """
@@ -175,23 +206,32 @@ class Runner:
             ComponentNotFoundException: if the ``component_path`` is failed to resolve.
         """
 
-        dryrun_info = self.dryrun_component(
-            component,
-            component_args,
-            scheduler,
-            cfg=cfg,
-            workspace=workspace,
-            parent_run_id=parent_run_id,
-        )
-        return self.schedule(dryrun_info)
+        with log_event("run_component") as ctx:
+            dryrun_info = self.dryrun_component(
+                component,
+                component_args,
+                scheduler,
+                cfg=cfg,
+                workspace=workspace,
+                parent_run_id=parent_run_id,
+            )
+            handle = self.schedule(dryrun_info)
+            app = none_throws(dryrun_info._app)
+
+            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]
+            ctx._torchx_event.app_metadata = app.metadata
+            return handle
 
     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:
         """
@@ -199,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,
@@ -217,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:
         """
@@ -230,10 +273,25 @@ class Runner:
             An application handle that is used to call other action APIs on the app.
         """
 
-        dryrun_info = self.dryrun(
-            app, scheduler, cfg=cfg, workspace=workspace, parent_run_id=parent_run_id
-        )
-        return self.schedule(dryrun_info)
+        with log_event(api="run") as ctx:
+            dryrun_info = self.dryrun(
+                app,
+                scheduler,
+                cfg=cfg,
+                workspace=workspace,
+                parent_run_id=parent_run_id,
+            )
+            handle = self.schedule(dryrun_info)
+
+            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:
         """
@@ -266,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:
@@ -291,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:
         """
@@ -343,6 +402,7 @@ class Runner:
             role.env[ENV_TORCHX_JOB_ID] = make_app_handle(
                 scheduler, self._name, macros.app_id
             )
+            role.env[TORCHX_INTERNAL_SESSION_ID] = get_session_id_or_create_new()
 
             if parent_run_id:
                 role.env[ENV_TORCHX_PARENT_RUN_ID] = parent_run_id
@@ -355,33 +415,49 @@ class Runner:
                     role.env[tracker_config_env_var_name(name)] = config
 
         cfg = cfg or dict()
-        with log_event("dryrun", scheduler, runcfg=json.dumps(cfg) if cfg else None):
+        with log_event(
+            "dryrun",
+            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)
-            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.'
-                )
-                sched.build_workspace_and_update_role(role, workspace, resolved_cfg)
-
-                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
                     )
 
-            sched._validate(app, scheduler)
+                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:
@@ -400,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.
@@ -490,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``.
@@ -525,7 +632,7 @@ class Runner:
             if not app:
                 desc = scheduler.describe(app_id)
                 if desc:
-                    app = AppDef(name=app_id, roles=desc.roles)
+                    app = AppDef(name=app_id, roles=desc.roles, metadata=desc.metadata)
             return app
 
     def log_lines(
@@ -637,7 +744,7 @@ class Runner:
                 app.app_handle = make_app_handle(scheduler, self._name, app.app_id)
             return apps
 
-    # pyre-fixme: Scheduler opts
+    # pyre-fixme[24]: SchedulerOpts is a generic, and we don't have access to the corresponding type
     def _scheduler(self, scheduler: str) -> Scheduler:
         sched = self._scheduler_instances.get(scheduler)
         if not sched:
@@ -654,8 +761,8 @@ class Runner:
     def _scheduler_app_id(
         self,
         app_handle: AppHandle,
-        check_session: bool = True
-        # pyre-fixme: Scheduler opts
+        check_session: bool = True,
+        # pyre-fixme[24]: SchedulerOpts is a generic, and we don't have access to the corresponding type
     ) -> Tuple[Scheduler, str, str]:
         """
         Returns the scheduler and app_id from the app_handle.

torchx/runner/config.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
+
 """
 Status: Beta
 
@@ -71,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
@@ -195,7 +197,15 @@ def _configparser() -> configparser.ConfigParser:
 
 
 def _get_scheduler(name: str) -> Scheduler:
-    schedulers = get_scheduler_factories()
+    schedulers = {
+        **get_scheduler_factories(),
+        **(
+            get_scheduler_factories(
+                group="torchx.schedulers.orchestrator", skip_defaults=True
+            )
+            or {}
+        ),
+    }
     if name not in schedulers:
         raise ValueError(
             f"`{name}` is not a registered scheduler. Valid scheduler names: {schedulers.keys()}"
@@ -239,7 +249,16 @@ def dump(
     if schedulers:
         scheds = schedulers
     else:
-        scheds = get_scheduler_factories().keys()
+        scheduler_factories = {
+            **get_scheduler_factories(),
+            **(
+                get_scheduler_factories(
+                    group="torchx.schedulers.orchestrator", skip_defaults=True
+                )
+                or {}
+            ),
+        }
+        scheds = scheduler_factories.keys()
 
     config = _configparser()
     for sched_name in scheds:
@@ -259,13 +278,20 @@ def dump(
                     continue
 
                 # serialize list elements with `;` delimiter (consistent with torchx cli)
-                if opt.opt_type == List[str]:
+                if opt.is_type_list_of_str:
                     # deal with empty or None default lists
                     if opt.default:
                         # pyre-ignore[6] opt.default type checked already as List[str]
                         val = ";".join(opt.default)
                     else:
                         val = _NONE
+                elif opt.is_type_dict_of_str:
+                    # deal with empty or None default lists
+                    if opt.default:
+                        # pyre-ignore[16] opt.default type checked already as Dict[str, str]
+                        val = ";".join([f"{k}:{v}" for k, v in opt.default.items()])
+                    else:
+                        val = _NONE
                 else:
                     val = f"{opt.default}"
 
@@ -468,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(
@@ -480,7 +508,7 @@ def find_configs(dirs: Optional[Iterable[str]] = None) -> List[str]:
             dirs = DEFAULT_CONFIG_DIRS
         for d in dirs:
             configfile = Path(d) / CONFIG_FILE
-            if configfile.exists():
+            if os.access(configfile, os.R_OK):
                 config_files.append(str(configfile))
     return config_files
 
@@ -510,21 +538,26 @@ def load(scheduler: str, f: TextIO, cfg: Dict[str, CfgVal]) -> None:
                 # this also handles empty or None lists
                 cfg[name] = None
             else:
-                runopt = runopts.get(name)
+                opt = runopts.get(name)
 
-                if runopt is None:
+                if opt is None:
                     log.warning(
                         f"`{name} = {value}` was declared in the [{section}] section "
                         f" of the config file but is not a runopt of `{scheduler}` scheduler."
                         f" Remove the entry from the config file to no longer see this warning"
                     )
                 else:
-                    if runopt.opt_type is bool:
+                    if opt.opt_type is bool:
                         # need to handle bool specially since str -> bool is based on
                         # str emptiness not value (e.g. bool("False") == True)
                         cfg[name] = config.getboolean(section, name)
-                    elif runopt.opt_type is List[str]:
+                    elif opt.is_type_list_of_str:
                         cfg[name] = value.split(";")
+                    elif opt.is_type_dict_of_str:
+                        cfg[name] = {
+                            s.split(":", 1)[0]: s.split(":", 1)[1]
+                            for s in value.replace(",", ";").split(";")
+                        }
                     else:
                         # pyre-ignore[29]
-                        cfg[name] = runopt.opt_type(value)
+                        cfg[name] = opt.opt_type(value)

torchx/runner/events/__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
+
 """
 Module contains events processing mechanisms that are integrated with the standard python logging.
 
@@ -18,18 +20,22 @@ Example of usage:
 
 """
 
+import json
 import logging
+import sys
 import time
 import traceback
 from types import TracebackType
-from typing import Optional, Type
+from typing import Dict, Optional, Type
 
 from torchx.runner.events.handlers import get_logging_handler
+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:
@@ -46,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)
-    _events_logger = logging.getLogger(f"torchx-events-{destination}")
-    _events_logger.setLevel(logging.DEBUG)
-    # 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:
@@ -82,17 +97,28 @@ class log_event:
         scheduler: Optional[str] = None,
         app_id: Optional[str] = None,
         app_image: Optional[str] = None,
+        app_metadata: Optional[Dict[str, str]] = None,
         runcfg: Optional[str] = None,
+        workspace: Optional[str] = None,
     ) -> None:
         self._torchx_event: TorchxEvent = self._generate_torchx_event(
-            api, scheduler or "", app_id, app_image=app_image, runcfg=runcfg
+            api,
+            scheduler or "",
+            app_id,
+            app_image=app_image,
+            app_metadata=app_metadata,
+            runcfg=runcfg,
+            workspace=workspace,
         )
         self._start_cpu_time_ns = 0
         self._start_wall_time_ns = 0
+        self._start_epoch_time_usec = 0
 
     def __enter__(self) -> "log_event":
         self._start_cpu_time_ns = time.process_time_ns()
         self._start_wall_time_ns = time.perf_counter_ns()
+        self._torchx_event.start_epoch_time_usec = int(time.time() * 1_000_000)
+
         return self
 
     def __exit__(
@@ -109,6 +135,20 @@ class log_event:
         ) // 1000
         if traceback_type:
             self._torchx_event.raw_exception = traceback.format_exc()
+            typ, value, tb = sys.exc_info()
+            if tb:
+                last_frame = traceback.extract_tb(tb)[-1]
+                self._torchx_event.exception_source_location = json.dumps(
+                    {
+                        "filename": last_frame.filename,
+                        "lineno": last_frame.lineno,
+                        "name": last_frame.name,
+                    }
+                )
+        if exec_type:
+            self._torchx_event.exception_type = exec_type.__name__
+        if exec_value:
+            self._torchx_event.exception_message = str(exec_value)
         record(self._torchx_event)
 
     def _generate_torchx_event(
@@ -117,15 +157,19 @@ class log_event:
         scheduler: str,
         app_id: Optional[str] = None,
         app_image: Optional[str] = None,
+        app_metadata: Optional[Dict[str, str]] = None,
         runcfg: Optional[str] = None,
         source: SourceType = SourceType.UNKNOWN,
+        workspace: Optional[str] = None,
     ) -> TorchxEvent:
         return TorchxEvent(
-            session=app_id or "",
+            session=get_session_id_or_create_new(),
             scheduler=scheduler,
             api=api,
             app_id=app_id,
             app_image=app_image,
+            app_metadata=app_metadata,
             runcfg=runcfg,
             source=source,
+            workspace=workspace,
         )