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

torchx/pipelines/kfp/adapter.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 json
 import os
 import os.path
@@ -48,7 +50,9 @@ def component_spec_from_app(app: api.AppDef) -> Tuple[str, api.Role]:
 
     role = app.roles[0]
     assert (
-        role.num_replicas == 1
+        role.num_replicas
+        == 1
+        # pyre-fixme[16]: `AppDef` has no attribute `num_replicas`.
     ), f"KFP adapter only supports one replica, got {app.num_replicas}"
 
     command = [role.entrypoint, *role.args]
@@ -74,8 +78,7 @@ class ContainerFactory(Protocol):
     kfp.dsl.ContainerOp.
     """
 
-    def __call__(self, *args: object, **kwargs: object) -> dsl.ContainerOp:
-        ...
+    def __call__(self, *args: object, **kwargs: object) -> dsl.ContainerOp: ...
 
 
 class KFPContainerFactory(ContainerFactory, Protocol):
@@ -104,7 +107,7 @@ def component_from_app(
         app: The AppDef to generate a KFP container factory for.
         ui_metadata: KFP UI Metadata to output so you can have model results show
             up in the UI. See
-            https://www.kubeflow.org/docs/components/pipelines/sdk/output-viewer/
+            https://www.kubeflow.org/docs/components/pipelines/legacy-v1/sdk/output-viewer/
             for more info on the format.
 
     >>> from torchx import specs

torchx/pipelines/kfp/version.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
+
 # Follows PEP-0440 version scheme guidelines
 # https://www.python.org/dev/peps/pep-0440/#version-scheme
 #

torchx/runner/__init__.py

@@ -5,4 +5,6 @@
 # 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
+
 from torchx.runner.api import get_runner, Runner  # noqa: F401 F403

torchx/runner/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 json
 import logging
 import os
@@ -12,7 +14,7 @@ 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, TypeVar
 
 from torchx.runner.events import log_event
 from torchx.schedulers import get_scheduler_factories, SchedulerFactory
@@ -37,9 +39,10 @@ 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.api import PkgInfo, WorkspaceBuilder, WorkspaceMixin
 
 from .config import get_config, get_configs
 
@@ -47,6 +50,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,7 +101,10 @@ class Runner:
         """
         self._name: str = name
         self._scheduler_factories = scheduler_factories
-        self._scheduler_params: Dict[str, object] = scheduler_params or {}
+        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] = {}
@@ -104,6 +112,14 @@ class Runner:
         # 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 _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
+        return scheduler_params
+
     def __enter__(self) -> "Runner":
         return self
 
@@ -131,9 +147,21 @@ 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 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,
@@ -175,15 +203,23 @@ 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", workspace=workspace) 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 = 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,
@@ -230,10 +266,22 @@ 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", runcfg=json.dumps(cfg) if cfg else None, workspace=workspace
+        ) as ctx:
+            dryrun_info = self.dryrun(
+                app,
+                scheduler,
+                cfg=cfg,
+                workspace=workspace,
+                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
+            return handle
 
     def schedule(self, dryrun_info: AppDryRunInfo) -> AppHandle:
         """
@@ -343,6 +391,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,7 +404,12 @@ 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=workspace,
+        ):
             sched = self._scheduler(scheduler)
             resolved_cfg = sched.run_opts().resolve(cfg)
             if workspace and isinstance(sched, WorkspaceMixin):
@@ -379,7 +433,11 @@ class Runner:
                         " Either a patch was built or no changes to workspace was detected."
                     )
 
-            sched._validate(app, scheduler)
+            with log_event(
+                "validate",
+                scheduler,
+            ):
+                sched._validate(app, scheduler, resolved_cfg)
             dryrun_info = sched.submit_dryrun(app, resolved_cfg)
             dryrun_info._scheduler = scheduler
             return dryrun_info
@@ -654,7 +712,7 @@ class Runner:
     def _scheduler_app_id(
         self,
         app_handle: AppHandle,
-        check_session: bool = True
+        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]:
         """

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
 
@@ -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:
@@ -266,6 +285,13 @@ def dump(
                         val = ";".join(opt.default)
                     else:
                         val = _NONE
+                elif opt.opt_type == Dict[str, 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}"
 
@@ -480,7 +506,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
 
@@ -525,6 +551,11 @@ def load(scheduler: str, f: TextIO, cfg: Dict[str, CfgVal]) -> None:
                         cfg[name] = config.getboolean(section, name)
                     elif runopt.opt_type is List[str]:
                         cfg[name] = value.split(";")
+                    elif runopt.opt_type is Dict[str, 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)

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,13 +20,16 @@ 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
 
@@ -82,17 +87,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 +125,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 +147,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,
         )

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_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,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 importlib
 from typing import Dict, Mapping
 
@@ -19,6 +21,7 @@ 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",
+    "aws_sagemaker": "torchx.schedulers.aws_sagemaker_scheduler",
     "gcp_batch": "torchx.schedulers.gcp_batch_scheduler",
     "ray": "torchx.schedulers.ray_scheduler",
     "lsf": "torchx.schedulers.lsf_scheduler",
@@ -27,8 +30,7 @@ DEFAULT_SCHEDULER_MODULES: Mapping[str, str] = {
 
 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,9 +42,11 @@ 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.
@@ -53,8 +57,9 @@ def get_scheduler_factories() -> Dict[str, SchedulerFactory]:
         default_schedulers[scheduler] = _defer_load_scheduler(path)
 
     return load_group(
-        "torchx.schedulers",
+        group,
         default=default_schedulers,
+        skip_defaults=skip_defaults,
     )
 
 

torchx/schedulers/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 re
 from dataclasses import dataclass, field
@@ -335,7 +337,7 @@ 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 _validate(self, app: AppDef, scheduler: str, cfg: T) -> None:
         """
         Validates whether application is consistent with the scheduler.
 

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
@@ -807,6 +809,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: