anyscale 0.24.86__py3-none-any.whl

anyscale/cluster.py

@@ -0,0 +1,138 @@
+import dataclasses
+from typing import Any, Dict, Optional
+
+import click
+
+from anyscale.authenticate import get_auth_api_client
+from anyscale.cli_logger import BlockLogger
+from anyscale.connect import ClientBuilder
+from anyscale.connect_utils.prepare_cluster import create_prepare_cluster_block
+from anyscale.connect_utils.project import create_project_block
+from anyscale.util import PROJECT_NAME_ENV_VAR
+
+
+log = BlockLogger()  # Anyscale CLI Logger
+
+
+@dataclasses.dataclass
+class ClusterInfo:
+    """
+    Synchronize with ray.dashboard.modules.job.sdk.ClusterInfo
+    """
+
+    address: str
+    cookies: Optional[Dict[str, Any]] = None
+    metadata: Optional[Dict[str, Any]] = None
+    headers: Optional[Dict[str, Any]] = None
+
+
+def get_job_submission_client_cluster_info(
+    address: str,
+    create_cluster_if_needed: Optional[bool] = None,
+    cookies: Optional[Dict[str, Any]] = None,  # noqa: ARG001
+    metadata: Optional[Dict[str, Any]] = None,
+    headers: Optional[Dict[str, Any]] = None,  # noqa: ARG001
+    **kwargs: Any,  # noqa: ARG001
+) -> ClusterInfo:
+    """
+    Get address and cookies used for ray JobSubmissionClient.
+
+    Args:
+        address (str): Address of same form as ray.init address without
+            anyscale:// prefix.
+        create_cluster_if_needed (bool): Indicates whether the cluster
+            of the address returned needs to be running. Raise an error
+            if cluster is not running and this is False. Create a cluster
+            if cluster is not running and this is True.
+
+    Returns:
+        ClusterInfo object consisting of address, cookies, and metadata for
+            JobSubmissionClient to use.
+    """
+
+    if create_cluster_if_needed:
+        # Use ClientBuilder to start cluster if needed because cluster needs to be active for
+        # the calling command.
+        api_client = get_auth_api_client().api_client
+        client_builder = ClientBuilder(
+            address=address, log=BlockLogger(log_output=False)
+        )
+        project_block = create_project_block(
+            client_builder._project_dir,  # noqa: SLF001
+            client_builder._project_name,  # noqa: SLF001
+            cloud_name=client_builder._cloud_name,  # noqa: SLF001
+            cluster_compute_name=client_builder._cluster_compute_name,  # noqa: SLF001
+            cluster_compute_dict=client_builder._cluster_compute_dict,  # noqa: SLF001
+        )
+        project_id = project_block.project_id
+        project_dir = project_block.project_dir
+        if project_dir and not client_builder._project_name:  # noqa: SLF001
+            # Warning when project dir provided or found in current directory and project name not provided
+            # as input.
+            # TODO(nikita): Remove after .anyscale.yaml is no longer supported
+            log.warning(
+                f"Project directory {project_dir} was detected. Using a project directory "
+                "to set the project for the cluster to use with `ray job submit` has been deprecated, and this "
+                "functionality will be removed in April 2022. To start an Anyscale cluster "
+                "in a particular project, please instead specify the project name as "
+                f'RAY_ADDRESS="anyscale://{project_block.project_name}/{client_builder._cluster_name if client_builder._cluster_name else ""}". '  # noqa: SLF001
+                "The project name can also be specified by setting the environment variable "
+                f'{PROJECT_NAME_ENV_VAR}="{project_block.project_name}". '
+                "Otherwise the cluster will not be grouped to a particular project.\n"
+            )
+        prepare_cluster_block = create_prepare_cluster_block(
+            project_id=project_id,
+            cluster_name=client_builder._cluster_name,  # noqa: SLF001
+            autosuspend_timeout=client_builder._autosuspend_timeout,  # noqa: SLF001
+            allow_public_internet_traffic=client_builder._allow_public_internet_traffic,  # noqa: SLF001
+            needs_update=client_builder._needs_update,  # noqa: SLF001
+            cluster_compute_name=client_builder._cluster_compute_name,  # noqa: SLF001
+            cluster_compute_dict=client_builder._cluster_compute_dict,  # noqa: SLF001
+            cloud_name=client_builder._cloud_name,  # noqa: SLF001
+            build_pr=client_builder._build_pr,  # noqa: SLF001
+            force_rebuild=client_builder._force_rebuild,  # noqa: SLF001
+            build_commit=client_builder._build_commit,  # noqa: SLF001
+            cluster_env_name=client_builder._cluster_env_name,  # noqa: SLF001
+            cluster_env_dict=client_builder._cluster_env_dict,  # noqa: SLF001
+            cluster_env_revision=client_builder._cluster_env_revision,  # noqa: SLF001
+            ray=client_builder._ray,  # noqa: SLF001
+        )
+        cluster_name = prepare_cluster_block.cluster_name
+    else:
+        # Calling ClientBuilder to parse address.
+        api_client = get_auth_api_client(log_output=False).api_client
+        client_builder = ClientBuilder(
+            address=address, log=BlockLogger(log_output=False)
+        )
+        cluster_name = client_builder._cluster_name  # noqa: SLF001
+        project_block = create_project_block(
+            client_builder._project_dir,  # noqa: SLF001
+            client_builder._project_name,  # noqa: SLF001
+            log_output=False,
+        )
+        project_id = project_block.project_id
+
+    user = api_client.get_user_info_api_v2_userinfo_get().result
+    metadata = {"creator_id": user.id}
+
+    cluster_list = api_client.list_sessions_api_v2_sessions_get(
+        project_id=project_id, active_only=True, name=cluster_name
+    ).results
+
+    if len(cluster_list) > 0:
+        cluster = cluster_list[0]
+        if cluster.host_name and cluster.access_token:
+            return ClusterInfo(
+                address=cluster.host_name,
+                cookies={"anyscale-token": cluster.access_token},
+                metadata=metadata,
+            )
+        else:
+            raise click.ClickException(
+                f"Host name or access token not found for cluster {cluster_name}. Please check the cluster is currently running."
+            )
+    else:
+        raise click.ClickException(
+            f"No running cluster found with name {cluster_name} in project {project_id}. Please start "
+            "the cluster, or change the project context if the wrong one is being used."
+        )

anyscale/cluster_compute.py

@@ -0,0 +1,167 @@
+import re
+from typing import Optional, Tuple
+
+from anyscale.authenticate import get_auth_api_client
+from anyscale.cli_logger import BlockLogger
+from anyscale.client.openapi_client.api.default_api import DefaultApi
+from anyscale.cloud import get_cloud_id_and_name, get_last_used_cloud
+from anyscale.sdk.anyscale_client import (
+    ArchiveStatus,
+    ComputeTemplateConfig,
+    CreateComputeTemplate,
+)
+from anyscale.sdk.anyscale_client.models.cluster_compute_config import (
+    ClusterComputeConfig,
+)
+from anyscale.sdk.anyscale_client.models.compute_template import ComputeTemplate
+from anyscale.sdk.anyscale_client.models.compute_template_query import (
+    ComputeTemplateQuery,
+)
+from anyscale.utils.cloud_utils import get_organization_default_cloud
+
+
+log = BlockLogger()
+
+
+def get_default_cluster_compute(
+    cloud_name: Optional[str],
+    project_id: Optional[str],
+    api_client: Optional[DefaultApi] = None,
+    anyscale_api_client: Optional[DefaultApi] = None,
+) -> ComputeTemplate:
+    if api_client is None:
+        api_client = get_auth_api_client().api_client
+    if anyscale_api_client is None:
+        anyscale_api_client = get_auth_api_client().anyscale_api_client
+
+    if cloud_name is None:
+        default_cloud_name = get_organization_default_cloud(api_client)
+        if default_cloud_name:
+            cloud_name = default_cloud_name
+        else:
+            cloud_name = get_last_used_cloud(project_id, anyscale_api_client)
+
+    cloud_id, _ = get_cloud_id_and_name(api_client, cloud_name=cloud_name)
+    config_object = anyscale_api_client.get_default_compute_config(cloud_id).result
+    compute_template = register_compute_template(config_object, api_client=api_client)
+    return compute_template
+
+
+def parse_cluster_compute_name_version(
+    cluster_compute_name_version: str,
+) -> Tuple[str, Optional[int]]:
+
+    # Regex pattern to validate the cluster compute name version string:
+    # - the name could not contain colon, leading or trailing spaces
+    # - the version should be a positive integer
+    # - the version is optional
+    # - the name and version should be separated by a colon if version is provided
+    COMPUTE_CONFIG_NAME_VERSION_VALIDATION_REGEX_PATTERN = (
+        r"^[^\s:]+(\s+[^\s:]+)*(:[1-9]+[0-9]*)?$"
+    )
+
+    if not re.match(
+        COMPUTE_CONFIG_NAME_VERSION_VALIDATION_REGEX_PATTERN,
+        cluster_compute_name_version,
+    ):
+        raise ValueError(
+            f"Invalid compute config name and version '{cluster_compute_name_version}'. Please provide name and version in format `my-compute-config:3` or provide compute configuration name `my-compute-config`."
+        )
+
+    cluster_compute_name = cluster_compute_name_version
+    version = None
+    if cluster_compute_name_version.count(":") == 1:
+        cluster_compute_name, parsed_version = cluster_compute_name_version.split(":")
+        version = int(parsed_version)
+
+    return cluster_compute_name, version
+
+
+def get_cluster_compute_from_name(
+    cluster_compute_name: str,
+    api_client: Optional[DefaultApi] = None,
+    include_archived: Optional[bool] = False,
+) -> ComputeTemplate:
+    if api_client is None:
+        api_client = get_auth_api_client().api_client
+
+    version = None
+
+    enable_compute_config_versioning = api_client.check_is_feature_flag_on_api_v2_userinfo_check_is_feature_flag_on_get(
+        "compute-config-versioning"
+    ).result.is_on
+
+    if enable_compute_config_versioning:
+        (cluster_compute_name, version,) = parse_cluster_compute_name_version(
+            cluster_compute_name
+        )
+
+    cluster_computes = api_client.search_compute_templates_api_v2_compute_templates_search_post(
+        ComputeTemplateQuery(
+            orgwide=True,
+            name={"equals": cluster_compute_name},
+            include_anonymous=True,
+            archive_status=ArchiveStatus.ALL,
+            version=version,
+        )
+    ).results
+
+    if len(cluster_computes) == 0:
+        raise ValueError(
+            f"The compute config {cluster_compute_name} does not exist, or you don't have sufficient permissions."
+        )
+    cluster_compute = cluster_computes[0]
+    if not include_archived and cluster_compute.archived_at is not None:
+        raise ValueError(f"The compute config {cluster_compute_name} is archived.")
+
+    return cluster_compute
+
+
+def register_compute_template(
+    config_object: ComputeTemplateConfig, api_client: Optional[DefaultApi] = None,
+) -> ComputeTemplate:
+    """
+    Register compute template with a default name and return the compute template id."""
+    if api_client is None:
+        api_client = get_auth_api_client().api_client
+    created_template = api_client.create_compute_template_api_v2_compute_templates_post(
+        create_compute_template=CreateComputeTemplate(
+            config=config_object, anonymous=True,
+        )
+    ).result
+    return created_template
+
+
+def get_selected_cloud_id_or_default(
+    api_client: Optional[DefaultApi] = None,
+    anyscale_api_client: Optional[DefaultApi] = None,
+    cluster_compute_id: Optional[str] = None,
+    cluster_compute_config: Optional[ClusterComputeConfig] = None,
+    cloud_id: Optional[str] = None,
+    cloud_name: Optional[str] = None,
+):
+    """
+    Gets cloud_id that is selected for the current command from the
+    arguments `cloud_id`, `cloud_name`, `cluster_compute_id`, or `cluster_compute_config`.
+    If the cloud_id is not selected through any of these arguments, get the default cloud
+    from the default cluster compute.
+    """
+    api_client = api_client or get_auth_api_client(log_output=False).api_client
+    anyscale_api_client = (
+        anyscale_api_client or get_auth_api_client(log_output=False).anyscale_api_client
+    )
+    if cloud_id or cloud_name:
+        parent_cloud_id, _ = get_cloud_id_and_name(
+            api_client=api_client, cloud_id=cloud_id, cloud_name=cloud_name,
+        )
+    elif cluster_compute_id:
+        parent_cloud_id = anyscale_api_client.get_cluster_compute(
+            cluster_compute_id
+        ).result.config.cloud_id
+    elif cluster_compute_config:
+        parent_cloud_id = cluster_compute_config.cloud_id
+    else:
+        parent_cloud_id = (
+            anyscale_api_client.get_default_cluster_compute().result.config.cloud_id
+        )
+    return parent_cloud_id

anyscale/cluster_env.py

@@ -0,0 +1,173 @@
+from typing import List, Optional
+
+import click
+
+from anyscale.authenticate import get_auth_api_client
+from anyscale.cli_logger import BlockLogger
+from anyscale.client.openapi_client.api.default_api import DefaultApi
+from anyscale.sdk.anyscale_client import ClusterEnvironmentBuild
+from anyscale.shared_anyscale_utils.utils.byod import is_byod_id
+from anyscale.util import get_endpoint, get_ray_and_py_version_for_default_cluster_env
+
+
+log = BlockLogger()
+
+
+def get_default_cluster_env_build(
+    api_client: Optional[DefaultApi] = None,
+    anyscale_api_client: Optional[DefaultApi] = None,
+) -> ClusterEnvironmentBuild:
+    ray_version, py_version = get_ray_and_py_version_for_default_cluster_env()
+
+    # TODO(nikita): Condense with anyscale.connect after unifying logging.
+    if api_client is None:
+        api_client = get_auth_api_client().api_client
+    if anyscale_api_client is None:
+        anyscale_api_client = get_auth_api_client().anyscale_api_client
+
+    try:
+        build_id = api_client.get_default_cluster_env_build_api_v2_builds_default_py_version_ray_version_get(
+            f"py{py_version}", ray_version
+        ).result.id
+        build = anyscale_api_client.get_cluster_environment_build(build_id).result
+        return build
+    except Exception:  # noqa: BLE001
+        raise RuntimeError(
+            f"Failed to get default cluster env for Ray: {ray_version} on Python: py{py_version}"
+        )
+
+
+def get_build_from_cluster_env_identifier(
+    cluster_env_identifier: str, anyscale_api_client: Optional[DefaultApi] = None,
+) -> ClusterEnvironmentBuild:
+    """
+    Get a build id from a cluster environment identifier of form `my_cluster_env:1` or
+    `my_cluster_env`. If no build revision is specified, return id of latest build
+    for this application config.
+
+    TODO(nikita): Move this to behind api endpoint and consolidate with anyscale.connect
+    """
+
+    if anyscale_api_client is None:
+        anyscale_api_client = get_auth_api_client().anyscale_api_client
+
+    cluster_env_revision: Optional[int] = None
+
+    if is_byod_id(cluster_env_identifier):
+        cluster_env_name = cluster_env_identifier
+        cluster_env_revision = 1
+        cluster_env_id = cluster_env_identifier
+    else:
+        try:
+            components = cluster_env_identifier.rsplit(":", 1)
+            cluster_env_name = components[0]
+            cluster_env_revision = int(components[1]) if len(components) > 1 else None
+
+        except ValueError:
+            raise click.ClickException(
+                "Invalid cluster-env-name provided. Please make sure the provided name is of "
+                "the form <cluster-env-name>:<revision>. For example, `my_cluster_env:1`."
+            )
+        # ID of cluster env and not build itself
+        cluster_env_id = get_cluster_env_from_name(
+            cluster_env_name, anyscale_api_client
+        ).id
+    builds = list_builds(cluster_env_id, anyscale_api_client)
+    if cluster_env_revision:
+        for build in builds:
+            if build.revision == cluster_env_revision:
+                return build
+
+        raise click.ClickException(
+            "Revision {} of cluster environment '{}' not found.".format(
+                cluster_env_revision, cluster_env_name
+            )
+        )
+    else:
+        latest_build_revision = -1
+        build_to_use = None
+        for build in builds:
+            if build.revision > latest_build_revision:
+                latest_build_revision = build.revision
+                build_to_use = build
+
+        if not build_to_use:
+            raise click.ClickException(
+                "Error finding latest build of cluster environment {}. Please manually "
+                "specify the build version in the cluster environment name with the format "
+                "<cluster-env-name>:<revision>. For example, `my_cluster_env:1`.".format(
+                    cluster_env_name
+                )
+            )
+        return build_to_use
+
+
+def get_cluster_env_from_name(
+    cluster_env_name: str, anyscale_api_client: Optional[DefaultApi] = None,
+) -> ClusterEnvironmentBuild:
+    """
+    Get id of the cluster env (not build) given the name.
+    """
+
+    if anyscale_api_client is None:
+        anyscale_api_client = get_auth_api_client().anyscale_api_client
+    cluster_envs = anyscale_api_client.search_cluster_environments(
+        {"name": {"equals": cluster_env_name}, "paging": {"count": 1}}
+    ).results
+    for cluster_env in cluster_envs:
+        if cluster_env.name == cluster_env_name:
+            return cluster_env
+
+    raise click.ClickException(f"Cluster environment '{cluster_env_name}' not found.")
+
+
+def list_builds(
+    cluster_env_id: str,
+    anyscale_api_client: Optional[DefaultApi] = None,
+    max_items: Optional[int] = None,
+) -> List[ClusterEnvironmentBuild]:
+    """
+    List all builds for a given cluster env id.
+    """
+
+    if anyscale_api_client is None:
+        anyscale_api_client = get_auth_api_client().anyscale_api_client
+    entities: List[ClusterEnvironmentBuild] = []
+    has_more = (len(entities) < max_items) if max_items else True
+    paging_token = None
+    i = 0
+    while has_more and i < 100:
+        resp = anyscale_api_client.list_cluster_environment_builds(
+            cluster_env_id, count=50, paging_token=paging_token
+        )
+        entities.extend(resp.results)
+        paging_token = resp.metadata.next_paging_token
+        has_more = paging_token is not None
+        if max_items:
+            has_more = has_more and (len(entities) < max_items)
+        i += 1
+    if max_items:
+        return entities[:max_items]
+    else:
+        return entities
+
+
+def validate_successful_build(
+    build_id: str, anyscale_api_client: Optional[DefaultApi] = None,
+) -> None:
+    """
+    Validate build_id provided is of a successfully completed build.
+    """
+
+    if anyscale_api_client is None:
+        anyscale_api_client = get_auth_api_client().anyscale_api_client
+    build = anyscale_api_client.get_cluster_environment_build(build_id).result
+    if build.status != "succeeded":
+        cluster_env = anyscale_api_client.get_cluster_environment(
+            build.cluster_environment_id
+        ).result
+        raise click.ClickException(
+            f"The cluster environment build {cluster_env.name}:{build.revision} currently is in state: {build.status}. "
+            f'More information about this build can be viewed at {get_endpoint(f"configurations/app-config-details/{build_id}")}. '
+            "Please provide a cluster environment that has already been built successfully."
+        )

anyscale/commands/aggregated_instance_usage_commands.py

@@ -0,0 +1,86 @@
+from typing import Optional
+
+import click
+
+import anyscale
+from anyscale.aggregated_instance_usage.models import DownloadCSVFilters
+from anyscale.cli_logger import BlockLogger
+from anyscale.commands import command_examples
+from anyscale.commands.util import AnyscaleCommand
+
+
+log = BlockLogger()  # CLI Logger
+
+
+@click.group("aggregated-instance-usage", help="Access aggregated instance usage.")
+def aggregated_instance_usage_cli() -> None:
+    pass
+
+
+@aggregated_instance_usage_cli.command(
+    name="download-csv",
+    cls=AnyscaleCommand,
+    example=command_examples.AGGREGATED_INSTANCE_USAGE_DOWNLOAD_CSV_EXAMPLE,
+)
+@click.option(
+    "--start-date",
+    required=True,
+    type=str,
+    help="The start date of the aggregated instance usage report (inclusive). Format: YYYY-MM-DD",
+)
+@click.option(
+    "--end-date",
+    required=True,
+    type=str,
+    help="The end date of the aggregated instance usage report (inclusive). Format: YYYY-MM-DD",
+)
+@click.option(
+    "--cloud",
+    required=False,
+    type=str,
+    help="The name of the cloud to filter the report by.",
+)
+@click.option(
+    "--project",
+    required=False,
+    type=str,
+    help="The name of the project to filter the report by.",
+)
+@click.option(
+    "--directory",
+    required=False,
+    type=str,
+    help="The directory to save the CSV file to. Default is the current directory",
+)
+@click.option(
+    "--hide-progress-bar",
+    required=False,
+    type=bool,
+    help="Whether to hide the progress bar. Default is False.",
+)
+def download_csv(
+    start_date: str,
+    end_date: str,
+    cloud: Optional[str],
+    project: Optional[str],
+    directory: Optional[str],
+    hide_progress_bar: Optional[bool] = False,
+) -> None:
+    """
+    Download an aggregated instance usage report as a zipped csv to the provided directory.
+    """
+    log.info("Downloading aggregated instance usage CSV...")
+    try:
+        anyscale.aggregated_instance_usage.download_csv(
+            DownloadCSVFilters(
+                start_date=start_date,
+                end_date=end_date,
+                cloud=cloud,
+                project=project,
+                directory=directory,
+                hide_progress_bar=hide_progress_bar,
+            )
+        )
+    except ValueError as e:
+        log.error(f"{e}")
+        return

anyscale/commands/anyscale_api/api_commands.py

@@ -0,0 +1,23 @@
+"""
+Commands to interact with the Anyscale API
+"""
+
+import click
+
+from anyscale.commands.anyscale_api.session_commands_commands import session_commands
+from anyscale.commands.anyscale_api.session_operations_commands import (
+    session_operations_commands,
+)
+from anyscale.commands.anyscale_api.sessions_commands import sessions
+
+
+@click.group(
+    "api", hidden=True, help="Various commands to interact with the Anyscale API."
+)
+def anyscale_api() -> None:
+    pass
+
+
+anyscale_api.add_command(session_commands)
+anyscale_api.add_command(sessions)
+anyscale_api.add_command(session_operations_commands)

anyscale/commands/anyscale_api/session_commands_commands.py

@@ -0,0 +1,80 @@
+"""
+Commands to interact with the Anyscale Session Commands API
+"""
+
+from typing import Optional
+
+import click
+
+from anyscale.authenticate import get_auth_api_client
+from anyscale.formatters import common_formatter
+from anyscale.sdk.anyscale_client import CreateSessionCommand
+
+
+@click.group(
+    "commands", help="Commands to interact with the Session Commands API.",
+)
+def session_commands() -> None:
+    pass
+
+
+@session_commands.command(name="list", short_help="Lists Session Commands.")
+@click.argument("session-id", required=True)
+@click.option(
+    "--count", type=int, default=10, help="Number of commands to show. Defaults to 10"
+)
+@click.option(
+    "--paging-token",
+    required=False,
+    help="Paging token used to fetch subsequent pages of commands.",
+)
+def list_session_commands(
+    session_id: str, count: int, paging_token: Optional[str],
+) -> None:
+    """Lists commands that have been executed on Session with SESSION_ID"""
+
+    api_client = get_auth_api_client().anyscale_api_client
+    response = api_client.list_session_commands(
+        session_id, count=count, paging_token=paging_token
+    )
+
+    print(common_formatter.prettify_json(response.to_dict()))
+
+
+@session_commands.command(name="get", short_help="Get details about a Session Command.")
+@click.argument("session-command-id", required=True)
+def get_session_command(session_command_id: str) -> None:
+    """Get details about the Session Command with id SESSION_COMMAND_ID"""
+
+    api_client = get_auth_api_client().anyscale_api_client
+    response = api_client.get_session_command(session_command_id)
+
+    print(common_formatter.prettify_json(response.to_dict()))
+
+
+@session_commands.command(name="kill", short_help="Kills a Session Command.")
+@click.argument("session-command-id", required=True)
+def kill_session_command(session_command_id: str) -> None:
+    """Kills the Session Command with id SESSION_COMMAND_ID"""
+
+    api_client = get_auth_api_client().anyscale_api_client
+    response = api_client.kill_session_command(session_command_id)
+
+    print(common_formatter.prettify_json(response.to_dict()))
+
+
+@session_commands.command(
+    name="create", short_help="Creates and executes a Session Command."
+)
+@click.argument("session-id", required=True)
+@click.argument("shell-command", required=True)
+def create_session_command(session_id: str, shell_command: str) -> None:
+    """Creates and executes SHELL_COMMAND on the session with id SESSION_ID."""
+
+    api_client = get_auth_api_client().anyscale_api_client
+    create_data = CreateSessionCommand(
+        session_id=session_id, shell_command=shell_command
+    )
+    response = api_client.create_session_command(create_data)
+
+    print(common_formatter.prettify_json(response.to_dict()))