anyscale 0.24.86__py3-none-any.whl

anyscale/anyscale_pydantic/version.py

@@ -0,0 +1,38 @@
+__all__ = 'compiled', 'VERSION', 'version_info'
+
+VERSION = '1.10.12'
+
+try:
+    import cython  # type: ignore
+except ImportError:
+    compiled: bool = False
+else:  # pragma: no cover
+    try:
+        compiled = cython.compiled
+    except AttributeError:
+        compiled = False
+
+
+def version_info() -> str:
+    import platform
+    import sys
+    from importlib import import_module
+    from pathlib import Path
+
+    optional_deps = []
+    for p in ('devtools', 'dotenv', 'email-validator', 'typing-extensions'):
+        try:
+            import_module(p.replace('-', '_'))
+        except ImportError:
+            continue
+        optional_deps.append(p)
+
+    info = {
+        'pydantic version': VERSION,
+        'pydantic compiled': compiled,
+        'install path': Path(__file__).resolve().parent,
+        'python version': sys.version,
+        'platform': platform.platform(),
+        'optional deps. installed': optional_deps,
+    }
+    return '\n'.join('{:>30} {}'.format(k + ':', str(v).replace('\n', ' ')) for k, v in info.items())

anyscale/anyscale_schema.json

@@ -0,0 +1,9 @@
+{
+  "type": "object",
+  "properties": {
+    "output_files": {
+      "description": "Files and directories not in the project directory to save in a snapshot",
+      "type": "array"
+    }
+  }
+}

anyscale/api.py

@@ -0,0 +1,215 @@
+import asyncio
+import functools
+import os
+import socket
+import sys
+from typing import Any, Optional
+
+import click
+from urllib3.connection import HTTPConnection
+import wrapt
+
+from anyscale.client import openapi_client
+from anyscale.client.openapi_client.api.default_api import DefaultApi
+from anyscale.client.openapi_client.rest import ApiException as ApiExceptionInternal
+from anyscale.sdk import anyscale_client
+from anyscale.sdk.anyscale_client.api.default_api import DefaultApi as AnyscaleApi
+from anyscale.sdk.anyscale_client.rest import ApiException as ApiExceptionExternal
+from anyscale.shared_anyscale_utils.headers import RequestHeaders
+from anyscale.version import __version__ as version
+
+
+# NOTE: Some OSes don't implement all of these, so only include those that the OS supports.
+_SOCKET_OPTIONS = (
+    HTTPConnection.default_socket_options
+    + [
+        (getattr(socket, t[0]), getattr(socket, t[1]), t[2])
+        for t in [
+            ("SOL_SOCKET", "SO_KEEPALIVE", 1),
+            # Begin sending keepalives after 15 seconds of no activity:
+            ("SOL_TCP", "TCP_KEEPIDLE", 15),
+            # Send keepalives every 15 seconds:
+            ("SOL_TCP", "TCP_KEEPINTVL", 15),
+            # Number failed keepalives before marking the connection as dead:
+            ("SOL_TCP", "TCP_KEEPCNT", 6),
+        ]
+        if hasattr(socket, t[0]) and hasattr(socket, t[1])
+    ]
+    # Darwin uses TCP_KEEPALIVE[A] instead of TCP_KEEPIDLE. Unfortunatley, this was
+    # only added to the socket library recently (June 2021)[B], so we include the value
+    # of TCP_KEEPALIVE (16)[C] if the version of python is not super recent.
+    # [A] Bug: https://bugs.python.org/issue34932
+    # [B] Resolution: https://github.com/python/cpython/pull/25079
+    # [C] Darwin Source: https://github.com/apple/darwin-xnu/blob/main/bsd/netinet/tcp.h
+    + (
+        [(6, getattr(socket, "TCP_KEEPALIVE", 16), 15)]
+        if sys.platform == "darwin"
+        else []
+    )
+)
+
+
+# client is of type APIClient, which is auto-generated
+def configure_tcp_keepalive(client: Any) -> None:
+    assert hasattr(
+        client, "rest_client"
+    ), f"Incorrect object of type: {type(client)}\nThis object does not have a `rest_client` property."
+    client.rest_client.pool_manager.connection_pool_kw[
+        "socket_options"
+    ] = _SOCKET_OPTIONS
+
+
+# client is of type APIClient, which is auto-generated
+def configure_open_api_client_headers(client: Any, client_name: str) -> None:
+    client.set_default_header(RequestHeaders.CLIENT, client_name)
+    client.set_default_header(RequestHeaders.CLIENT_VERSION, version)
+
+
+class _ApiClient:
+    api_client: Optional[DefaultApi] = None
+    anyscale_client: Optional[AnyscaleApi] = None
+
+
+def format_api_exception(
+    e, method: str, resource_path: str, raise_structured_exception: bool = False,
+) -> None:
+    if os.environ.get("ANYSCALE_DEBUG") == "1" or raise_structured_exception:
+        raise e
+    else:
+        raise click.ClickException(
+            f"API Exception ({e.status}) from {method} {resource_path} \n"
+            f"Reason: {e.reason}\nHTTP response body: {e.body}\n"
+            f"Trace ID: {e.headers._container.get('x-trace-id', None)}"  # noqa: SLF001
+        )
+
+
+class ApiClientWrapperInternal(openapi_client.ApiClient):
+    def __init__(self, *args, raise_structured_exception: bool = False, **kwargs):
+        """
+        Arguments:
+            raise_structured_exception (bool): If True, API exceptions will be raised
+            as structured exceptions. If this and ANYSCALE_DEBUG are False, API
+            exceptions will be raised as user friendly but unstructured Click exceptions.
+            This arguement allows us to determine the type of raised error in code, but
+            users should ANYSCALE_DEBUG to configure this.
+        """
+        self.raise_structured_exception = raise_structured_exception
+        super().__init__(*args, **kwargs)
+
+    def call_api(  # noqa: PLR0913
+        self,
+        resource_path,
+        method,
+        path_params=None,
+        query_params=None,
+        header_params=None,
+        body=None,
+        post_params=None,
+        files=None,
+        response_type=None,
+        auth_settings=None,
+        async_req=None,
+        _return_http_data_only=None,
+        collection_formats=None,
+        _preload_content=True,
+        _request_timeout=None,
+        _host=None,
+    ):
+        try:
+            return openapi_client.ApiClient.call_api(
+                self,
+                resource_path,
+                method,
+                path_params,
+                query_params,
+                header_params,
+                body,
+                post_params,
+                files,
+                response_type,
+                auth_settings,
+                async_req,
+                _return_http_data_only,
+                collection_formats,
+                _preload_content,
+                _request_timeout,
+                _host,
+            )
+        except ApiExceptionInternal as e:
+            format_api_exception(
+                e, method, resource_path, self.raise_structured_exception
+            )
+
+
+class ApiClientWrapperExternal(anyscale_client.ApiClient):
+    def __init__(self, *args, raise_structured_exception: bool = False, **kwargs):
+        """
+        Arguments:
+            raise_structured_exception (bool): If True, API exceptions will be raised
+            as structured exceptions. If this and ANYSCALE_DEBUG are False, API
+            exceptions will be raised as user friendly but unstructured Click exceptions.
+            This arguement allows us to determine the type of raised error in code, but
+            users should ANYSCALE_DEBUG to configure this.
+        """
+        self.raise_structured_exception = raise_structured_exception
+        super().__init__(*args, **kwargs)
+
+    def call_api(  # noqa: PLR0913
+        self,
+        resource_path,
+        method,
+        path_params=None,
+        query_params=None,
+        header_params=None,
+        body=None,
+        post_params=None,
+        files=None,
+        response_type=None,
+        auth_settings=None,
+        async_req=None,
+        _return_http_data_only=None,
+        collection_formats=None,
+        _preload_content=True,
+        _request_timeout=None,
+        _host=None,
+    ):
+        try:
+            return anyscale_client.ApiClient.call_api(
+                self,
+                resource_path,
+                method,
+                path_params,
+                query_params,
+                header_params,
+                body,
+                post_params,
+                files,
+                response_type,
+                auth_settings,
+                async_req,
+                _return_http_data_only,
+                collection_formats,
+                _preload_content,
+                _request_timeout,
+                _host,
+            )
+        except ApiExceptionExternal as e:
+            format_api_exception(
+                e, method, resource_path, self.raise_structured_exception
+            )
+
+
+@wrapt.decorator
+def make_async(_func, instance, args, kwargs):  # noqa: ARG001
+    loop = asyncio.get_event_loop()
+    func = functools.partial(_func, *args, **kwargs)
+    return loop.run_in_executor(executor=None, func=func)
+
+
+class AsyncApiClientWrapperExternal(ApiClientWrapperExternal):
+    @make_async
+    def call_api(self, *args, **kwargs):
+        return super().call_api(*args, **kwargs)
+
+
+_api_client = _ApiClient()

anyscale/api_utils/README.md

@@ -0,0 +1,2 @@
+## Overview
+For having logic that is both in CLI and the wrapped SDK (`AnyscaleSDK`).

anyscale/api_utils/common_utils.py

@@ -0,0 +1,81 @@
+import os
+from typing import Optional, Tuple
+
+from anyscale.client.openapi_client.api.default_api import DefaultApi as InternalApi
+from anyscale.client.openapi_client.models import (
+    Cloud,
+    ExperimentalWorkspace,
+)
+from anyscale.sdk.anyscale_client.api.default_api import DefaultApi as ExternalApi
+from anyscale.sdk.anyscale_client.models import Project
+
+
+WORKSPACE_ID_ENV_VAR = "ANYSCALE_EXPERIMENTAL_WORKSPACE_ID"
+
+# TODO (sumanthrh): Currently this same functionality of getting workspace id is repeated in many places, we should try to use this util everywhere
+def get_current_workspace_id() -> Optional[str]:
+    """Returns the ID of the workspace the client is running in (or `None`)."""
+    return os.environ.get(WORKSPACE_ID_ENV_VAR, None)
+
+
+def source_cloud_id_and_project_id(
+    *,
+    internal_api: InternalApi,
+    external_api: ExternalApi,
+    cloud_id: Optional[str],
+    project_id: Optional[str],
+) -> Tuple[str, str]:
+    """
+        Returns `Tuple[cloud_id, project_id]` based on optionally provided `cloud_id` and `project_id`.
+
+        - If neither `cloud_id` nor `project_id` are provided, the default cloud and its default project are used
+        (unless inside a workspace)
+        - If both `cloud_id` and `project_id` are provided, we verify that the project is under the cloud.
+
+        Note that the relationship hierarchy is:
+        ```
+        Cloud -----------<- Project
+        1                     Many
+        ```
+        - Each cloud will have at least one project (the default project).
+        - Each project will belong to exactly one cloud.
+        """
+    if not cloud_id:
+        if project_id:
+            # 1. No cloud, project
+            # (Use the cloud of the project)
+            project: Project = external_api.get_project(project_id).result
+            cloud_id = project.parent_cloud_id
+        else:
+            # 2. No cloud, no project
+            # (Use the default cloud and its default project)
+            # (unless we are in an Anyscale workspace, then use the workspace's cloud and project)
+            workspace_id = get_current_workspace_id()
+            is_anyscale_workspace = workspace_id is not None
+            if is_anyscale_workspace:
+                workspace: ExperimentalWorkspace = internal_api.get_workspace_api_v2_experimental_workspaces_workspace_id_get(
+                    workspace_id
+                ).result
+                cloud_id = workspace.cloud_id
+                project_id = workspace.project_id
+            else:
+                cloud: Cloud = external_api.get_default_cloud().result
+                cloud_id = cloud.id
+                project = external_api.get_default_project(
+                    parent_cloud_id=cloud_id
+                ).result
+                project_id = project.id
+    else:
+        if not project_id:
+            # 3. Cloud, no project
+            # (Use the default project of the cloud)
+            project = external_api.get_default_project(parent_cloud_id=cloud_id).result
+            project_id = project.id
+        else:
+            # 4. Cloud, project
+            # (Verify that the project is in the cloud)
+            project = external_api.get_project(project_id).result
+            assert (
+                project.parent_cloud_id == cloud_id
+            ), f"Project {project_id} is not in cloud {cloud_id}."
+    return (cloud_id, project_id)  # type: ignore

anyscale/api_utils/exceptions/job_errors.py

@@ -0,0 +1,2 @@
+class NoJobRunError(Exception):
+    """Raised if the job has no job runs (attempts) created yet"""

anyscale/api_utils/job_logs_util.py

@@ -0,0 +1,116 @@
+import asyncio
+from typing import List, Optional, Tuple
+
+from anyscale.api_utils.logs_util import (
+    _download_logs_concurrently,
+    _remove_ansi_escape_sequences,
+)
+from anyscale.cli_logger import LogsLogger
+from anyscale.client.openapi_client.api.default_api import DefaultApi
+from anyscale.client.openapi_client.models.base_job_status import BaseJobStatus
+from anyscale.controllers.logs_controller import DEFAULT_PARALLELISM
+from anyscale.sdk.anyscale_client.models.log_download_result import LogDownloadResult
+
+
+_JOB_RUN_TERMINAL_STATES = {
+    BaseJobStatus.SUCCEEDED,
+    BaseJobStatus.STOPPED,
+    BaseJobStatus.COMPLETED,
+    BaseJobStatus.FAILED,
+    BaseJobStatus.UNKNOWN,
+}
+
+
+async def _get_job_logs_from_storage_bucket_streaming(
+    api: DefaultApi,
+    log: LogsLogger,
+    *,
+    next_page_token: Optional[str] = None,
+    job_run_id: Optional[str] = None,
+    parallelism: int = DEFAULT_PARALLELISM,
+    remove_escape_chars: bool = True,
+    cluster_journal_events_start_line: int = 0,
+    no_cluster_journal_events: bool = False,
+) -> Tuple[str, Optional[str], int, bool]:
+    """Retrieves logs directly from the storage bucket.
+    Also fetches cluster journal logs and print them to stderr inline with job logs.
+
+    Args:
+        parallelism (int, optional): Defaults to 10 (AWS S3's default `max_concurrent_requests` value)
+            There is no documented upper limit. However, please set a reasonable value.
+        remove_escape_chars (bool, optional): Removes ANSI escape sequences from the logs, which are
+            commonly used to color terminal output. Defaults to True.
+
+    Returns (logs, next_page_token, cluster_journal_events_new_start_line, is_job_run_terminated):
+        Returns the logs output and the next page token for fetching additional logs.
+        cluster_journal_events_new_start_line is the new start line for fetching cluster journal events.
+        is_job_run_terminated is True if the job run is terminated, False otherwise.
+
+    Raises:
+        NoJobRunError: If a job run hasn't been created yet.
+    """
+    # First fetch both the job info and the log chunks in parallel
+
+    job_run_thread = api.get_decorated_job_api_v2_decorated_jobs_job_id_get(
+        job_id=job_run_id, async_req=True
+    )
+
+    log_download_result_thread = api.get_job_logs_download_v2_api_v2_logs_job_logs_download_v2_job_id_get(
+        job_id=job_run_id, next_page_token=next_page_token, async_req=True
+    )
+
+    # Once job info is fetched, fetch the cluster journal events
+    job_run = job_run_thread.get().result
+    cluster_id = job_run.cluster.id
+
+    if not no_cluster_journal_events:
+        cluster_journal_events_thread = api.get_startup_logs_api_v2_sessions_session_id_startup_logs_get(
+            cluster_id,
+            start_line=cluster_journal_events_start_line,
+            end_line=10000000,
+            async_req=True,
+        )
+
+    # Wait for the log chunks API and start downloading the chunks in parallel
+    log_download_result: LogDownloadResult = log_download_result_thread.get().result
+    all_log_chunk_urls: List[str] = [chunk.chunk_url for chunk in log_download_result.log_chunks]  # type: ignore
+    logs_task = asyncio.create_task(
+        _download_logs_concurrently(
+            all_log_chunk_urls,
+            parallelism,
+            bearer_token=log_download_result.bearer_token,
+        )
+    )
+
+    # Print out the cluster journal events
+    if not no_cluster_journal_events:
+        cluster_journal_events = cluster_journal_events_thread.get().result
+        lines = cluster_journal_events.lines.splitlines()
+        new_cluster_journal_events_start_line = (
+            cluster_journal_events.start_line + cluster_journal_events.num_lines
+        )
+        if cluster_journal_events_start_line == 0:
+            # When fetching the inital cluster journal events, we will just show the latest 10 events
+            # to avoid flooding the output with too many events. In particular, after a job is
+            # already running.
+            lines = lines[-10:]
+        for line in lines:
+            log.info(line)
+    else:
+        new_cluster_journal_events_start_line = cluster_journal_events_start_line
+
+    # Finally wait for all the log chunks to download and return all the logs
+    logs = await logs_task
+    if remove_escape_chars:
+        logs = _remove_ansi_escape_sequences(logs)
+    new_next_page_token = (
+        log_download_result.next_page_token
+        if len(log_download_result.log_chunks) > 0
+        else next_page_token
+    )
+    return (
+        logs,
+        new_next_page_token,
+        new_cluster_journal_events_start_line,
+        job_run.status in _JOB_RUN_TERMINAL_STATES,
+    )

anyscale/api_utils/job_util.py

@@ -0,0 +1,22 @@
+from typing import Optional
+
+from anyscale.api_utils.exceptions.job_errors import NoJobRunError
+from anyscale.sdk.anyscale_client.api.default_api import DefaultApi as BaseApi
+from anyscale.sdk.anyscale_client.models.production_job import ProductionJob
+
+
+def _get_job_run_id(
+    base_api: BaseApi,
+    *,
+    job_id: Optional[str] = None,
+    job_run_id: Optional[str] = None,
+) -> str:
+    assert bool(job_id) != bool(
+        job_run_id
+    ), "Exactly one of `job_id` or `job_run_id` must be provided."
+    if job_id:
+        prod_job: ProductionJob = base_api.get_production_job(job_id).result
+        if not prod_job.last_job_run_id:
+            raise NoJobRunError(f"Production job {job_id} has no job runs.")
+        job_run_id = prod_job.last_job_run_id
+    return job_run_id  # type: ignore

anyscale/api_utils/logs_util.py

@@ -0,0 +1,61 @@
+import asyncio
+import re
+from typing import List, Optional
+
+import aiohttp
+import requests
+
+from anyscale.shared_anyscale_utils.utils.asyncio import gather_in_batches
+
+
+CLUSTER_CONNECT_TIMEOUT = 30
+
+
+async def _download_logs_concurrently(
+    log_chunk_urls: List[str], parallelism: int, bearer_token: Optional[str] = None
+) -> str:
+    logs_across_chunks: List[str] = await gather_in_batches(  # type: ignore
+        parallelism,
+        *[
+            _download_log_from_s3_url(url, bearer_token=bearer_token)
+            for url in log_chunk_urls
+        ],
+    )
+    logs_across_chunks = [log.strip() for log in logs_across_chunks]
+    return "\n".join(logs_across_chunks)
+
+
+async def _download_log_from_ray_json_response(url: str) -> str:
+    async with aiohttp.ClientSession() as session:
+        response = await asyncio.wait_for(
+            session.get(url), timeout=CLUSTER_CONNECT_TIMEOUT
+        )
+        logs: str = (await response.json()).get("logs", "")
+        return logs
+
+
+async def _download_log_from_s3_url(
+    url: str, bearer_token: Optional[str] = None,
+) -> str:
+    # Note that the URL is presigned, so no token needs to be passed in the request
+    async with aiohttp.ClientSession() as session:
+        headers = {"Authorization": f"Bearer {bearer_token}"} if bearer_token else {}
+        async with session.get(url, headers=headers) as response:
+            return await response.text()
+
+
+def _download_log_from_s3_url_sync(
+    url: str, bearer_token: Optional[str] = None,
+) -> str:
+    # Note that the URL is presigned, so no token needs to be passed in the request
+    headers = {"Authorization": f"Bearer {bearer_token}"} if bearer_token else {}
+    response = requests.get(url, headers=headers)
+    response.raise_for_status()
+    return response.text
+
+
+def _remove_ansi_escape_sequences(s: str) -> str:
+    # Required as the log may contain ANSI escape sequeneces (e.g. for coloring in the terminal)
+    # Regex pattern from https://stackoverflow.com/a/14693789
+    ansi_escape = re.compile(r"\x1B(?:[@-Z\\-_]|\[[0-?]*[ -/]*[@-~])")
+    return ansi_escape.sub("", s)