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

torchx/schedulers/ray_scheduler.py

@@ -1,453 +0,0 @@
-# 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.
-
-import dataclasses
-import json
-import logging
-import os
-import re
-import tempfile
-import time
-from dataclasses import dataclass, field
-from datetime import datetime
-from shutil import copy2, rmtree
-from typing import Any, cast, Dict, Final, Iterable, List, Optional, Tuple  # noqa
-
-import urllib3
-
-from torchx.schedulers.api import (
-    AppDryRunInfo,
-    AppState,
-    DescribeAppResponse,
-    filter_regex,
-    ListAppResponse,
-    Scheduler,
-    split_lines,
-    Stream,
-)
-from torchx.schedulers.ids import make_unique
-from torchx.schedulers.ray.ray_common import RayActor, TORCHX_RANK0_HOST
-from torchx.specs import AppDef, macros, NONE, ReplicaStatus, Role, RoleStatus, runopts
-from torchx.workspace.dir_workspace import TmpDirWorkspaceMixin
-from typing_extensions import TypedDict
-
-
-try:
-    from ray.autoscaler import sdk as ray_autoscaler_sdk
-    from ray.dashboard.modules.job.common import JobStatus
-    from ray.dashboard.modules.job.sdk import JobSubmissionClient
-
-    _has_ray = True
-
-except ImportError:
-    _has_ray = False
-
-
-def has_ray() -> bool:
-    """Indicates whether Ray is installed in the current Python environment."""
-    return _has_ray
-
-
-class RayOpts(TypedDict, total=False):
-    cluster_config_file: Optional[str]
-    cluster_name: Optional[str]
-    dashboard_address: Optional[str]
-    working_dir: Optional[str]
-    requirements: Optional[str]
-
-
-if _has_ray:
-    _logger: logging.Logger = logging.getLogger(__name__)
-
-    _ray_status_to_torchx_appstate: Dict[JobStatus, AppState] = {
-        JobStatus.PENDING: AppState.PENDING,
-        JobStatus.RUNNING: AppState.RUNNING,
-        JobStatus.SUCCEEDED: AppState.SUCCEEDED,
-        JobStatus.FAILED: AppState.FAILED,
-        JobStatus.STOPPED: AppState.CANCELLED,
-    }
-
-    class _EnhancedJSONEncoder(json.JSONEncoder):
-        def default(self, o: RayActor):  # pyre-ignore[3]
-            if dataclasses.is_dataclass(o):
-                return dataclasses.asdict(o)
-            return super().default(o)
-
-    def serialize(
-        actors: List[RayActor], dirpath: str, output_filename: str = "actors.json"
-    ) -> None:
-        actors_json = json.dumps(actors, cls=_EnhancedJSONEncoder)
-        with open(os.path.join(dirpath, output_filename), "w") as tmp:
-            json.dump(actors_json, tmp)
-
-    @dataclass
-    class RayJob:
-        """Represents a job that should be run on a Ray cluster.
-
-        Attributes:
-            app_id:
-                The unique ID of the application (a.k.a. job).
-            cluster_config_file:
-                The Ray cluster configuration file.
-            cluster_name:
-                The cluster name to use.
-            dashboard_address:
-                The existing dashboard IP address to connect to
-            working_dir:
-                The working directory to copy to the cluster
-            requirements:
-                The libraries to install on the cluster per requirements.txt
-            actors:
-                The Ray actors which represent the job to be run. This attribute is
-                dumped to a JSON file and copied to the cluster where `ray_main.py`
-                uses it to initiate the job.
-        """
-
-        app_id: str
-        working_dir: str
-        cluster_config_file: Optional[str] = None
-        cluster_name: Optional[str] = None
-        dashboard_address: Optional[str] = None
-        requirements: Optional[str] = None
-        actors: List[RayActor] = field(default_factory=list)
-
-    class RayScheduler(TmpDirWorkspaceMixin, Scheduler[RayOpts]):
-        """
-        RayScheduler is a TorchX scheduling interface to Ray. The job def
-        workers will be launched as Ray actors
-
-        The job environment is specified by the TorchX workspace. Any files in
-        the workspace will be present in the Ray job unless specified in
-        ``.torchxignore``. Python dependencies will be read from the
-        ``requirements.txt`` file located at the root of the workspace unless
-        it's overridden via ``-c ...,requirements=foo/requirements.txt``.
-
-        **Config Options**
-
-        .. runopts::
-            class: torchx.schedulers.ray_scheduler.create_scheduler
-
-        **Compatibility**
-
-        .. compatibility::
-            type: scheduler
-            features:
-                cancel: true
-                logs: |
-                    Partial support. Ray only supports a single log stream so
-                    only a dummy "ray/0" combined log role is supported.
-                    Tailing and time seeking are not supported.
-                distributed: true
-                describe: |
-                    Partial support. RayScheduler will return job status but
-                    does not provide the complete original AppSpec.
-                workspaces: true
-                mounts: false
-                elasticity: Partial support. Multi role jobs are not supported.
-
-        """
-
-        def __init__(
-            self, session_name: str, ray_client: Optional[JobSubmissionClient] = None
-        ) -> None:
-            # NOTE: make sure any new init options are supported in create_scheduler(...)
-            super().__init__("ray", session_name)
-
-            # w/o Final None check in _get_ray_client does not work as it pyre assumes mutability
-            self._ray_client: Final[Optional[JobSubmissionClient]] = ray_client
-
-        def _get_ray_client(
-            self, job_submission_netloc: Optional[str] = None
-        ) -> JobSubmissionClient:
-            if self._ray_client is not None:
-                client_netloc = urllib3.util.parse_url(
-                    self._ray_client.get_address()
-                ).netloc
-                if job_submission_netloc and job_submission_netloc != client_netloc:
-                    raise ValueError(
-                        f"client netloc ({client_netloc}) does not match job netloc ({job_submission_netloc})"
-                    )
-                return self._ray_client
-            elif os.getenv("RAY_ADDRESS"):
-                return JobSubmissionClient(os.getenv("RAY_ADDRESS"))
-            elif not job_submission_netloc:
-                raise Exception(
-                    "RAY_ADDRESS env variable or a scheduler with an attached Ray JobSubmissionClient is expected."
-                    " See https://docs.ray.io/en/latest/cluster/jobs-package-ref.html#job-submission-sdk for more info"
-                )
-            return JobSubmissionClient(f"http://{job_submission_netloc}")
-
-        # TODO: Add address as a potential CLI argument after writing ray.status() or passing in config file
-        def _run_opts(self) -> runopts:
-            opts = runopts()
-            opts.add(
-                "cluster_config_file",
-                type_=str,
-                required=False,
-                help="Use CLUSTER_CONFIG_FILE to access or create the Ray cluster.",
-            )
-            opts.add(
-                "cluster_name",
-                type_=str,
-                help="Override the configured cluster name.",
-            )
-            opts.add(
-                "dashboard_address",
-                type_=str,
-                required=False,
-                default="127.0.0.1:8265",
-                help="Use ray status to get the dashboard address you will submit jobs against",
-            )
-            opts.add("requirements", type_=str, help="Path to requirements.txt")
-            return opts
-
-        def schedule(self, dryrun_info: AppDryRunInfo[RayJob]) -> str:
-            cfg: RayJob = dryrun_info.request
-
-            # Create serialized actors for ray_driver.py
-            actors = cfg.actors
-            dirpath = cfg.working_dir
-            serialize(actors, dirpath)
-
-            job_submission_addr: str = ""
-            if cfg.cluster_config_file:
-                job_submission_addr = ray_autoscaler_sdk.get_head_node_ip(
-                    cfg.cluster_config_file
-                )  # pragma: no cover
-            elif cfg.dashboard_address:
-                job_submission_addr = cfg.dashboard_address
-            else:
-                raise RuntimeError(
-                    "Either `dashboard_address` or `cluster_config_file` must be specified"
-                )
-
-            # 0. Create Job Client
-            client = self._get_ray_client(job_submission_netloc=job_submission_addr)
-
-            # 1. Copy Ray driver utilities
-            current_directory = os.path.dirname(os.path.abspath(__file__))
-            copy2(os.path.join(current_directory, "ray", "ray_driver.py"), dirpath)
-            copy2(os.path.join(current_directory, "ray", "ray_common.py"), dirpath)
-            runtime_env = {"working_dir": dirpath}
-            if cfg.requirements:
-                runtime_env["pip"] = cfg.requirements
-
-            # 1. Submit Job via the Ray Job Submission API
-            try:
-                job_id: str = client.submit_job(
-                    submission_id=cfg.app_id,
-                    # we will pack, hash, zip, upload, register working_dir in GCS of ray cluster
-                    # and use it to configure your job execution.
-                    entrypoint="python3 ray_driver.py",
-                    runtime_env=runtime_env,
-                )
-
-            finally:
-                if dirpath.startswith(tempfile.gettempdir()):
-                    rmtree(dirpath)
-
-            # Encode job submission client in job_id
-            return f"{job_submission_addr}-{job_id}"
-
-        def _submit_dryrun(self, app: AppDef, cfg: RayOpts) -> AppDryRunInfo[RayJob]:
-            app_id = make_unique(app.name)
-
-            working_dir = app.roles[0].image
-            if not os.path.exists(working_dir):
-                raise RuntimeError(
-                    f"Role image must be a valid directory, got: {working_dir} "
-                )
-
-            requirements: Optional[str] = cfg.get("requirements")
-            if requirements is None:
-                workspace_reqs = os.path.join(working_dir, "requirements.txt")
-                if os.path.exists(workspace_reqs):
-                    requirements = workspace_reqs
-
-            cluster_cfg = cfg.get("cluster_config_file")
-            if cluster_cfg:
-                if not isinstance(cluster_cfg, str) or not os.path.isfile(cluster_cfg):
-                    raise ValueError(
-                        "The cluster configuration file must be a YAML file."
-                    )
-
-                job: RayJob = RayJob(
-                    app_id,
-                    cluster_config_file=cluster_cfg,
-                    requirements=requirements,
-                    working_dir=working_dir,
-                )
-
-            else:  # pragma: no cover
-                dashboard_address = cfg.get("dashboard_address")
-                job: RayJob = RayJob(
-                    app_id=app_id,
-                    dashboard_address=dashboard_address,
-                    requirements=requirements,
-                    working_dir=working_dir,
-                )
-            job.cluster_name = cfg.get("cluster_name")
-
-            for role in app.roles:
-                for replica_id in range(role.num_replicas):
-                    # Replace the ${img_root}, ${app_id}, and ${replica_id} placeholders
-                    # in arguments and environment variables.
-                    replica_role = macros.Values(
-                        img_root=role.image,
-                        app_id=app_id,
-                        replica_id=str(replica_id),
-                        rank0_env=TORCHX_RANK0_HOST,
-                    ).apply(role)
-
-                    actor = RayActor(
-                        name=role.name,
-                        min_replicas=role.min_replicas,
-                        command=[replica_role.entrypoint] + replica_role.args,
-                        env=replica_role.env,
-                        num_cpus=max(1, replica_role.resource.cpu),
-                        num_gpus=max(0, replica_role.resource.gpu),
-                    )
-
-                    job.actors.append(actor)
-
-            if len(app.roles) > 1 and app.roles[0].min_replicas is not None:
-                raise ValueError("min_replicas is only supported with single role jobs")
-
-            return AppDryRunInfo(job, repr)
-
-        def _validate(self, app: AppDef, scheduler: str) -> None:
-            if scheduler != "ray":
-                raise ValueError(
-                    f"An unknown scheduler backend '{scheduler}' has been passed to the Ray scheduler."
-                )
-
-            if app.metadata:
-                _logger.warning("The Ray scheduler does not use metadata information.")
-
-            for role in app.roles:
-                if role.resource.capabilities:
-                    _logger.warning(
-                        "The Ray scheduler does not support custom resource capabilities."
-                    )
-                    break
-
-            for role in app.roles:
-                if role.port_map:
-                    _logger.warning("The Ray scheduler does not support port mapping.")
-                    break
-
-        def wait_until_finish(self, app_id: str, timeout: int = 30) -> None:
-            """
-            ``wait_until_finish`` waits until the specified job has finished
-            with a given timeout. This is intended for testing. Programmatic
-            usage should use the runner wait method instead.
-            """
-
-            start = time.time()
-            while time.time() - start <= timeout:
-                status_info = self._get_job_status(app_id)
-                status = status_info
-                if status in {JobStatus.SUCCEEDED, JobStatus.STOPPED, JobStatus.FAILED}:
-                    break
-                time.sleep(1)
-
-        def _parse_app_id(self, app_id: str) -> Tuple[str, str]:
-            # find index of '-' in the first :\d+-
-            m = re.search(r":\d+-", app_id)
-            if m:
-                sep = m.span()[1]
-                addr = app_id[: sep - 1]
-                app_id = app_id[sep:]
-                return addr, app_id
-
-            addr, _, app_id = app_id.partition("-")
-            return addr, app_id
-
-        def _cancel_existing(self, app_id: str) -> None:  # pragma: no cover
-            addr, app_id = self._parse_app_id(app_id)
-            client = self._get_ray_client(job_submission_netloc=addr)
-            client.stop_job(app_id)
-
-        def _get_job_status(self, app_id: str) -> JobStatus:
-            addr, app_id = self._parse_app_id(app_id)
-            client = self._get_ray_client(job_submission_netloc=addr)
-            status = client.get_job_status(app_id)
-            if isinstance(status, str):
-                return cast(JobStatus, status)
-            return status.status
-
-        def describe(self, app_id: str) -> Optional[DescribeAppResponse]:
-            job_status_info = self._get_job_status(app_id)
-            state = _ray_status_to_torchx_appstate[job_status_info]
-            roles = [Role(name="ray", num_replicas=1, image="<N/A>")]
-
-            # get ip_address and put it in hostname
-
-            roles_statuses = [
-                RoleStatus(
-                    role="ray",
-                    replicas=[
-                        ReplicaStatus(
-                            id=0,
-                            role="ray",
-                            hostname=NONE,
-                            state=state,
-                        )
-                    ],
-                )
-            ]
-            return DescribeAppResponse(
-                app_id=app_id,
-                state=state,
-                msg=job_status_info,
-                roles_statuses=roles_statuses,
-                roles=roles,
-            )
-
-        def log_iter(
-            self,
-            app_id: str,
-            role_name: Optional[str] = None,
-            k: int = 0,
-            regex: Optional[str] = None,
-            since: Optional[datetime] = None,
-            until: Optional[datetime] = None,
-            should_tail: bool = False,
-            streams: Optional[Stream] = None,
-        ) -> Iterable[str]:
-            # TODO: support tailing, streams etc..
-            addr, app_id = self._parse_app_id(app_id)
-            client: JobSubmissionClient = self._get_ray_client(
-                job_submission_netloc=addr
-            )
-            logs: str = client.get_job_logs(app_id)
-            iterator = split_lines(logs)
-            if regex:
-                return filter_regex(regex, iterator)
-            return iterator
-
-        def list(self) -> List[ListAppResponse]:
-            client = self._get_ray_client()
-            jobs = client.list_jobs()
-            netloc = urllib3.util.parse_url(client.get_address()).netloc
-            return [
-                ListAppResponse(
-                    app_id=f"{netloc}-{details.submission_id}",
-                    state=_ray_status_to_torchx_appstate[details.status],
-                )
-                for details in jobs
-            ]
-
-
-def create_scheduler(
-    session_name: str, ray_client: Optional[JobSubmissionClient] = None, **kwargs: Any
-) -> "RayScheduler":
-    if not has_ray():  # pragma: no cover
-        raise ModuleNotFoundError(
-            "Ray is not installed in the current Python environment."
-        )
-
-    return RayScheduler(session_name=session_name, ray_client=ray_client)

torchx_nightly-2024.1.6.dist-info/METADATA

@@ -1,176 +0,0 @@
-Metadata-Version: 2.1
-Name: torchx-nightly
-Version: 2024.1.6
-Summary: TorchX SDK and Components
-Home-page: https://github.com/pytorch/torchx
-Author: TorchX Devs
-Author-email: torchx@fb.com
-License: BSD-3
-Keywords: pytorch,machine learning
-Platform: UNKNOWN
-Classifier: Development Status :: 4 - Beta
-Classifier: Intended Audience :: Developers
-Classifier: Intended Audience :: Science/Research
-Classifier: License :: OSI Approved :: BSD License
-Classifier: Programming Language :: Python :: 3
-Classifier: Programming Language :: Python :: 3.8
-Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
-Requires-Python: >=3.7
-Description-Content-Type: text/markdown
-License-File: LICENSE
-Requires-Dist: pyre-extensions
-Requires-Dist: docstring-parser >=0.8.1
-Requires-Dist: importlib-metadata
-Requires-Dist: pyyaml
-Requires-Dist: docker
-Requires-Dist: filelock
-Requires-Dist: fsspec ==2023.10.0
-Requires-Dist: urllib3 <1.27,>=1.21.1
-Requires-Dist: tabulate
-Provides-Extra: aws_batch
-Requires-Dist: boto3 ; extra == 'aws_batch'
-Provides-Extra: dev
-Requires-Dist: aiobotocore ; extra == 'dev'
-Requires-Dist: ax-platform[mysql] ==0.2.3 ; extra == 'dev'
-Requires-Dist: black ==23.3.0 ; extra == 'dev'
-Requires-Dist: boto3 ; extra == 'dev'
-Requires-Dist: captum >=0.4.0 ; extra == 'dev'
-Requires-Dist: docker ; extra == 'dev'
-Requires-Dist: flake8 ==3.9.0 ; extra == 'dev'
-Requires-Dist: fsspec[s3] ==2023.10.0 ; extra == 'dev'
-Requires-Dist: google-api-core ; extra == 'dev'
-Requires-Dist: google-cloud-batch >=0.5.0 ; extra == 'dev'
-Requires-Dist: google-cloud-logging >=3.0.0 ; extra == 'dev'
-Requires-Dist: google-cloud-runtimeconfig >=0.33.2 ; extra == 'dev'
-Requires-Dist: hydra-core ; extra == 'dev'
-Requires-Dist: ipython ; extra == 'dev'
-Requires-Dist: kfp ==1.8.22 ; extra == 'dev'
-Requires-Dist: mlflow-skinny ; extra == 'dev'
-Requires-Dist: moto ==4.1.6 ; extra == 'dev'
-Requires-Dist: protobuf ==3.20.3 ; extra == 'dev'
-Requires-Dist: pyre-extensions ; extra == 'dev'
-Requires-Dist: pyre-check ; extra == 'dev'
-Requires-Dist: pytest ; extra == 'dev'
-Requires-Dist: pytorch-lightning ==1.5.10 ; extra == 'dev'
-Requires-Dist: torch-model-archiver >=0.4.2 ; extra == 'dev'
-Requires-Dist: torch >=1.10.0 ; extra == 'dev'
-Requires-Dist: torchmetrics <0.11.0 ; extra == 'dev'
-Requires-Dist: torchserve >=0.4.2 ; extra == 'dev'
-Requires-Dist: torchtext >=0.11.0 ; extra == 'dev'
-Requires-Dist: torchvision >=0.11.1 ; extra == 'dev'
-Requires-Dist: ts ==0.5.1 ; extra == 'dev'
-Requires-Dist: usort ==1.0.2 ; extra == 'dev'
-Requires-Dist: ray[default] ; extra == 'dev'
-Provides-Extra: gcp_batch
-Requires-Dist: google-cloud-batch >=0.5.0 ; extra == 'gcp_batch'
-Requires-Dist: google-cloud-logging >=3.0.0 ; extra == 'gcp_batch'
-Requires-Dist: google-cloud-runtimeconfig >=0.33.2 ; extra == 'gcp_batch'
-Provides-Extra: kfp
-Requires-Dist: kfp ==1.6.2 ; extra == 'kfp'
-Provides-Extra: kubernetes
-Requires-Dist: kubernetes >=11 ; extra == 'kubernetes'
-Provides-Extra: ray
-Requires-Dist: ray >=1.12.1 ; extra == 'ray'
-
-[![PyPI](https://img.shields.io/pypi/v/torchx)](https://pypi.org/project/torchx/)
-[![License](https://img.shields.io/badge/License-BSD%203--Clause-blue.svg)](LICENSE)
-![Tests](https://github.com/pytorch/torchx/actions/workflows/python-unittests.yaml/badge.svg)
-![Lint](https://github.com/pytorch/torchx/actions/workflows/lint.yaml/badge.svg)
-[![codecov](https://codecov.io/gh/pytorch/torchx/branch/main/graph/badge.svg?token=ceHHIm0hXy)](https://codecov.io/gh/pytorch/torchx)
-
-
-# TorchX
-
-
-TorchX is a universal job launcher for PyTorch applications.
-TorchX is designed to have fast iteration time for training/research and support
-for E2E production ML pipelines when you're ready.
-
-TorchX currently supports:
-
-* Kubernetes (EKS, GKE, AKS, etc)
-* Slurm
-* AWS Batch
-* Docker
-* Local
-* Ray (prototype)
-* GCP Batch (prototype)
-
-Need a scheduler not listed? [Let us know!](https://github.com/pytorch/torchx/issues?q=is%3Aopen+is%3Aissue+label%3Ascheduler-request)
-
-## Quickstart
-
-See the [quickstart guide](https://pytorch.org/torchx/latest/quickstart.html).
-
-## Documentation
-
-* [Stable Documentation](https://pytorch.org/torchx/latest/)
-* [Nightly Documentation](https://pytorch.org/torchx/main/)
-
-## Requirements
-
-torchx:
-
-* python3 (3.8+)
-* [PyTorch](https://pytorch.org/get-started/locally/)
-* optional: [Docker](https://docs.docker.com/get-docker/) (needed for docker based schedulers)
-
-Certain schedulers may require scheduler specific requirements. See installation
-for info.
-
-## Installation
-
-### Stable
-
-```bash
-# install torchx sdk and CLI -- minimum dependencies
-pip install torchx
-
-# install torchx sdk and CLI -- all dependencies
-pip install "torchx[dev]"
-
-# install torchx kubeflow pipelines (kfp) support
-pip install "torchx[kfp]"
-
-# install torchx Kubernetes / Volcano support
-pip install "torchx[kubernetes]"
-
-# install torchx Ray support
-pip install "torchx[ray]"
-
-# install torchx GCP Batch support
-pip install "torchx[gcp_batch]"
-```
-
-### Nightly
-
-```bash
-# install torchx sdk and CLI
-pip install torchx-nightly[dev]
-```
-
-### Source
-
-```bash
-# install torchx sdk and CLI from source
-$ pip install -e git+https://github.com/pytorch/torchx.git#egg=torchx
-
-# install extra dependencies
-$ pip install -e git+https://github.com/pytorch/torchx.git#egg=torchx[dev]
-```
-
-### Docker
-
-TorchX provides a docker container for using as as part of a TorchX role.
-
-See: https://github.com/pytorch/torchx/pkgs/container/torchx
-
-## Contributing
-
-We welcome PRs! See the [CONTRIBUTING](CONTRIBUTING.md) file.
-
-## License
-
-TorchX is BSD licensed, as found in the [LICENSE](LICENSE) file.
-
-