anyscale 0.24.86__py3-none-any.whl

anyscale/utils/runtime_env.py

@@ -0,0 +1,487 @@
+# pylint:disable=private-import
+from contextlib import contextmanager
+import copy
+import hashlib
+import logging
+import os
+from pathlib import Path
+import tempfile
+from typing import (
+    Any,
+    Callable,
+    Dict,
+    Generator,
+    List,
+    Optional,
+    Tuple,
+    TYPE_CHECKING,
+)
+from urllib.parse import urlparse
+
+import click
+
+from anyscale.cli_logger import BlockLogger
+from anyscale.client.openapi_client import ComputeTemplate
+from anyscale.client.openapi_client.api.default_api import DefaultApi
+from anyscale.client.openapi_client.models.cloud_providers import CloudProviders
+from anyscale.client.openapi_client.models.cloud_with_cloud_resource import (
+    CloudWithCloudResource,
+)
+from anyscale.client.openapi_client.models.cloud_with_cloud_resource_gcp import (
+    CloudWithCloudResourceGCP,
+)
+from anyscale.client.openapi_client.models.user_info import UserInfo
+from anyscale.shared_anyscale_utils.aws import bucket_name_from_maybe_bucket_arn
+from anyscale.util import is_anyscale_workspace
+from anyscale.utils.ray_utils import zip_directory  # type: ignore
+from anyscale.utils.workload_types import Workload
+
+
+if TYPE_CHECKING:
+    pass
+
+logger = logging.getLogger(__name__)
+
+# TODO(austin): refactor to read requirements.txt and .skip_packages_tracking from s3 or gcs directly.
+# Default cluster storage directory.
+CLUSTER_STORAGE_DIR = "/mnt/cluster_storage"
+
+# Location of the directory containing workspace configurations.
+WORKSPACE_CONF_DIR = os.environ.get(
+    "ANYSCALE_WORKSPACE_CONF_DIR", os.path.join(CLUSTER_STORAGE_DIR, ".anyscale")
+)
+
+# Location of the workspace-managed requirements.txt file that is automatically populated.
+# This is not always guaranteed to exist within a workspace.
+WORKSPACE_REQUIREMENTS_FILE_PATH = os.path.join(WORKSPACE_CONF_DIR, "requirements.txt")
+
+# Feature flags for pip dependency tracking.
+SKIP_PACKAGES_TRACKING_PATH = os.path.join(
+    WORKSPACE_CONF_DIR, ".skip_packages_tracking"
+)
+
+
+def is_workspace_dependency_tracking_disabled() -> bool:
+    """Returns True if the workspace dependency tracking is disabled.
+
+    Enabled iff:
+        - ANYSCALE_WORKSPACE_DYNAMIC_DEPENDENCY_TRACKING env var is set to 1.
+        - ANYSCALE_SKIP_PYTHON_DEPENDENCY_TRACKING env var is not set to 1.
+        - No file exists at SKIP_PACKAGES_TRACKING_PATH.
+    """
+    # NOTE(edoakes): The environment variable is evaluated here instead of in the global scope
+    # so it can easily be overwritten for testing.
+    enabled = (
+        os.environ.get("ANYSCALE_WORKSPACE_DYNAMIC_DEPENDENCY_TRACKING", "0") == "1"
+    )
+    return (
+        not enabled
+        or os.environ.get("ANYSCALE_SKIP_PYTHON_DEPENDENCY_TRACKING", "0") == "1"
+        or os.path.exists(SKIP_PACKAGES_TRACKING_PATH)
+    )
+
+
+def _upload_file_to_google_cloud_storage(file: str, bucket: str, object_name: str):
+    try:
+        from google.cloud import storage
+
+    except Exception:  # noqa: BLE001
+        raise click.ClickException(
+            "Could not upload file to Google Storage. Could not import the Google Storage Python API via `from google.cloud import storage`.  Please check your installation or try running `pip install --upgrade google-cloud-storage`."
+        )
+    try:
+        storage_client = storage.Client()
+        bucket_obj = storage_client.bucket(bucket)
+        blob = bucket_obj.blob(object_name)
+        blob.upload_from_filename(file)
+    except Exception as e:  # noqa: BLE001
+        raise click.ClickException(
+            f"Failed to upload the working directory to Google Cloud Storage. Error {e!r}"
+            "Please validate you have exported cloud credentials with the correct write permissions and the intended bucket exists in your Cloud Storage account. "
+            "If you do not desire to upload your working directory, please set your working directory to a public remote URI or remove the runtime_environment from your job or service yaml."
+        ) from e
+
+
+def _upload_file_to_s3(file: str, bucket: str, object_key: str):
+    try:
+        import boto3
+        import botocore.config
+    except Exception:  # noqa: BLE001
+        raise click.ClickException(
+            "Could not upload file to S3: Could not import the Amazon S3 Python API via `import boto3`.  Please check your installation or try running `pip install boto3`."
+        )
+    try:
+        s3_client = boto3.client(
+            "s3", config=botocore.config.Config(signature_version="s3v4")
+        )
+        s3_client.upload_file(file, bucket, object_key)
+    except Exception as e:  # noqa: BLE001
+        raise click.ClickException(
+            f"Failed to upload the working directory to S3. Error {e!r}"
+            "Please validate you have exported cloud credentials with the correct write permissions and the intended bucket exists in your S3 account. "
+            "If you do not desire to upload your working directory, please set your working directory to a public remote URI or remove the runtime_environment from your job or service yaml."
+        ) from e
+
+
+def _get_remote_storage_object_name(upload_path, upload_filename):
+    # Strip leading slash, otherwise bucket will create a new directory called "/".
+    object_name = os.path.join(urlparse(upload_path).path, upload_filename).lstrip("/")
+    return object_name
+
+
+def _upload_file_to_remote_storage(
+    source_file: str, upload_path: str, upload_filename: str
+):
+    parsed_upload_path = urlparse(upload_path)
+    service = parsed_upload_path.scheme
+    bucket = parsed_upload_path.netloc
+    object_name = _get_remote_storage_object_name(upload_path, upload_filename)
+    if service == "s3":
+        _upload_file_to_s3(source_file, bucket, object_key=object_name)
+    if service == "gs":
+        _upload_file_to_google_cloud_storage(
+            source_file, bucket, object_name=object_name
+        )
+
+    final_uploaded_filepath = os.path.join(upload_path, upload_filename)
+    try:
+        from smart_open import open
+
+        open(final_uploaded_filepath)
+    except Exception as e:  # noqa: BLE001
+        raise click.ClickException(
+            f"Could not open uploaded file, maybe something went wrong while uploading: {e}."
+        )
+
+    return final_uploaded_filepath
+
+
+def is_dir_remote_uri(target_dir: str) -> bool:
+    parsed = urlparse(target_dir)
+    if parsed.scheme:
+        return True
+    return False
+
+
+@contextmanager
+def zip_local_dir(
+    path: str, *, excludes: Optional[List[str]] = None
+) -> Generator[Tuple[str, bytes, str], None, None]:
+    """Packs the local directory into a temporary zip file.
+
+    After the context manager exits, the file is deleted.
+
+    Yields: (temp_path, file_contents, content_hash).
+    """
+    with tempfile.TemporaryDirectory() as temp_dir:
+        zip_file_path = os.path.join(temp_dir, "anyscale_generated_working_dir.zip")
+        zip_directory(
+            path,
+            excludes=excludes if excludes is not None else [],
+            output_path=zip_file_path,
+            # Ray requires remote Zip URIs to consist of a single top-level directory when unzipped.
+            include_parent_dir=True,
+        )
+        zip_file_bytes = Path(zip_file_path).read_bytes()
+        yield zip_file_path, zip_file_bytes, hashlib.md5(zip_file_bytes).hexdigest()
+
+
+def upload_and_rewrite_working_dir(
+    runtime_env_json: Dict[str, Any],
+    upload_file_to_remote_storage_fn: Callable[
+        [str, str, str], str
+    ] = _upload_file_to_remote_storage,
+) -> Dict[str, Any]:
+    """Upload a local working_dir and rewrite the working_dir field with the destination remote URI.
+
+    After uploading, deletes the "upload_path" field because it is no longer used and is not a valid
+    OSS runtime env field.
+    """
+    if runtime_env_json.get("working_dir") is None:
+        return runtime_env_json
+
+    working_dir = runtime_env_json["working_dir"]
+    if is_dir_remote_uri(working_dir):
+        # The working dir is a remote URI already
+        return runtime_env_json
+
+    upload_path = runtime_env_json["upload_path"]
+    excludes = runtime_env_json.get("excludes")
+    with zip_local_dir(working_dir, excludes=excludes) as (zip_file_path, _, hash_val):
+        uploaded_zip_file_name = f"_anyscale_pkg_{hash_val}.zip"
+        final_uploaded_filepath = upload_file_to_remote_storage_fn(
+            zip_file_path, upload_path, uploaded_zip_file_name,
+        )
+
+    final_runtime_env = runtime_env_json.copy()
+    final_runtime_env["working_dir"] = final_uploaded_filepath
+    del final_runtime_env["upload_path"]
+    return final_runtime_env
+
+
+def override_runtime_env_config(
+    runtime_env: Optional[Dict[str, Any]],
+    anyscale_api_client: DefaultApi,
+    api_client: DefaultApi,
+    workload_type: Workload,
+    compute_config_id: Optional[str],
+    log: BlockLogger,
+) -> Optional[Dict[str, Any]]:
+    """Override the working_dir, upload_path, and pip fields for a runtime_env.
+
+    When running inside a workspace:
+        1. Autopopulates the `working_dir` with the current directory.
+        2. Autopopulates the `pip` field with the contents of the workspace-managed
+           `requirements.txt` file.
+
+    If the working_dir is a local path, will upload the contents to cloud storage and
+    replace the field with the resulting remote URI. The upload_path can be specified
+    in the runtime_env, else it will be auto-populated for the cloud.
+    """
+
+    existing_runtime_env = autopopulate_runtime_env_for_workspace(
+        runtime_env=runtime_env, log=log
+    )
+
+    if not existing_runtime_env:
+        return {}
+    elif not existing_runtime_env.get("working_dir"):
+        return existing_runtime_env
+
+    working_dir = existing_runtime_env.get("working_dir", "")
+    upload_path = existing_runtime_env.get("upload_path")
+
+    if not is_dir_remote_uri(working_dir):
+        if upload_path is not None:
+            # If upload_path is specified
+            # we back up the current working dir to the specified path
+            new_runtime_env = upload_and_rewrite_working_dir(existing_runtime_env)
+        elif is_anyscale_workspace() and "ANYSCALE_SESSION_ID" in os.environ:
+            # If submitting job v2 from workspaces and no upload_path is specified,
+            # we back up the current workspace content into S3
+            cluster_id = os.environ["ANYSCALE_SESSION_ID"]
+
+            decorated_cluster = api_client.get_decorated_cluster_api_v2_decorated_sessions_cluster_id_get(
+                cluster_id
+            ).result
+            cloud_id = decorated_cluster.cloud.id
+
+            workspace_id = os.environ["ANYSCALE_EXPERIMENTAL_WORKSPACE_ID"]
+
+            new_runtime_env = infer_upload_path_and_rewrite_working_dir(
+                api_client=api_client,
+                existing_runtime_env=existing_runtime_env,
+                workload_type=workload_type,
+                cloud_id=cloud_id,
+                log=log,
+                workspace_id=workspace_id,
+            )
+        else:
+            compute_template: ComputeTemplate = anyscale_api_client.get_compute_template(
+                compute_config_id
+            ).result
+            cloud_id = compute_template.config.cloud_id
+            new_runtime_env = infer_upload_path_and_rewrite_working_dir(
+                api_client=api_client,
+                existing_runtime_env=existing_runtime_env,
+                workload_type=workload_type,
+                cloud_id=cloud_id,
+                log=log,
+            )
+
+        return new_runtime_env
+    else:
+        return existing_runtime_env
+
+
+def parse_dot_env_file(dot_env_bytes: bytes) -> Dict[str, str]:
+    """Parse a .env file and return a dictionary of key-value pairs."""
+    dot_env = dot_env_bytes.split(b"\x00")
+    ret = {}
+    for kv in dot_env:
+        if len(kv) == 0:
+            # skip empty lines
+            continue
+        segs = kv.split(b"=", 1)
+        if len(segs) == 2:
+            key, value = segs
+            try:
+                ret[key.decode().strip()] = value.decode().strip()
+            except UnicodeDecodeError:
+                logger.error(f"Failed to decode env var entry: {kv!r}")
+        else:
+            logger.error(f"Invalid env var entry: {kv!r}")
+    return ret
+
+
+def parse_requirements_file(path: str) -> Optional[List[str]]:
+    """TODO: add comment."""
+    requirements_file = Path(path)
+    if requirements_file.is_file():
+        parsed_requirements = []
+        lines = requirements_file.read_text().strip().split("\n")
+        for line in lines:
+            # Strip comments.
+            hash_idx = line.find("#")
+            if hash_idx != -1:
+                line = line[:hash_idx]
+
+            # Ignore empty lines (after stripping comments).
+            line = line.strip()
+            if line:
+                parsed_requirements.append(line)
+    else:
+        parsed_requirements = None
+
+    return parsed_requirements
+
+
+def autopopulate_runtime_env_for_workspace(
+    runtime_env: Optional[Dict[str, Any]],
+    log: BlockLogger,
+    *,
+    requirements_file_path: str = WORKSPACE_REQUIREMENTS_FILE_PATH,
+) -> Optional[Dict[str, Any]]:
+    """Autopopulates fields of the runtime_env for commands run in a workspace.
+
+    Fields populated (if not specified by the user):
+        - working_dir: set to ".".
+        - pip: set to the contents of the workspace-managed requirements.txt file.
+          If the file does not exist, this field will not be set.
+
+    No-op if called outside of a workspace.
+    """
+    if not is_anyscale_workspace():
+        return runtime_env
+
+    if not runtime_env:
+        runtime_env = {}
+
+    if not runtime_env.get("working_dir"):
+        runtime_env["working_dir"] = "."
+        log.info("working_dir is not specified, using the current local directory.")
+
+    # Workspaces maintains a `requirements.txt` file in a well-known location.
+    # By default, populate the runtime environment with the contents of this file.
+    # If the user passes any "pip" dependencies, do not overwrite them.
+    if not is_workspace_dependency_tracking_disabled():
+        parsed_requirements = parse_requirements_file(requirements_file_path)
+        if parsed_requirements:
+            if runtime_env.get("pip"):
+                log.info(
+                    "Not including workspace-tracked dependencies because "
+                    "'pip' field is specified in the runtime_env."
+                )
+            elif runtime_env.get("conda"):
+                log.info(
+                    "Not including workspace-tracked dependencies because "
+                    "'conda' field is specified in the runtime_env."
+                )
+            else:
+                runtime_env["pip"] = parsed_requirements
+                log.info("Including workspace pip dependencies.")
+
+    return runtime_env
+
+
+def infer_upload_path_and_rewrite_working_dir(
+    *,
+    api_client: DefaultApi,
+    existing_runtime_env: Dict[str, Any],
+    cloud_id: str,
+    workload_type: Workload,
+    log: BlockLogger,
+    workspace_id: Optional[str] = None,
+) -> Dict[str, Any]:
+    """
+    Push working_dir to remote bucket and rewrite the working_dir field with the destination uri
+
+    If the upload_path is not specified by the user, we will get the bucket name based on the cloud.
+    We then rewrite the working_dir to the remote uri path
+    so that the launched service will read from remote bucket directly.
+
+    For Workspaces:
+        The remote path: [s3, gs]://{bucket_name}/{org_id}/{cloud_id}/workspace_snapshots/{workspace_id}/{workload_type}/{backup_zip}
+    Otherwise:
+        The remote path: [s3, gs]://{bucket_name}/{org_id}/{cloud_id}/{workload_type}/{backup_zip}
+        workload_type=[jobs, scheduled_jobs, services]
+    """
+
+    cloud: CloudWithCloudResource = api_client.get_cloud_with_cloud_resource_api_v2_clouds_with_cloud_resource_router_cloud_id_get(
+        cloud_id
+    ).result
+    org_id = _get_organization_id(api_client)
+
+    if cloud.provider == CloudProviders.AWS:
+        bucket_name = _get_cloud_s3_bucket_from_cloud(cloud)
+        protocol = "s3"
+    elif cloud.provider == CloudProviders.GCP:
+        bucket_name = _get_cloud_gs_bucket_from_cloud(api_client, cloud)
+        protocol = "gs"
+    else:
+        raise click.ClickException(
+            f"Currently launching a service from workspaces in a {cloud.provider} cloud is not supported. "
+            "Please contact Anyscale support for more info."
+        )
+
+    new_runtime_env = copy.deepcopy(existing_runtime_env)
+    working_dir_path = Path(new_runtime_env["working_dir"]).absolute()
+    log.info(f"Uploading local working_dir from '{working_dir_path}'.")
+
+    if workspace_id:
+        new_runtime_env[
+            "upload_path"
+        ] = f"{protocol}://{bucket_name}/{org_id}/{cloud_id}/workspace_snapshots/{workspace_id}/{workload_type}"
+    else:
+        new_runtime_env[
+            "upload_path"
+        ] = f"{protocol}://{bucket_name}/{org_id}/{cloud_id}/{workload_type}"
+
+    new_runtime_env = upload_and_rewrite_working_dir(new_runtime_env)
+    return new_runtime_env
+
+
+def _get_organization_id(api_client: DefaultApi):
+    user_info: UserInfo = (api_client.get_user_info_api_v2_userinfo_get().result)
+    orgs = user_info.organizations
+    return orgs[0].id
+
+
+def _get_cloud_s3_bucket_from_cloud(cloud: CloudWithCloudResource) -> Optional[str]:
+    """
+    If the cloud has an associated aws s3 bucket, we return its name.
+
+    Please note that this is only for v2 clouds where customers have their
+    own S3 buckets.
+    """
+    assert cloud.provider == CloudProviders.AWS
+    if cloud and cloud.cloud_resource and cloud.cloud_resource.aws_s3_id:
+        return bucket_name_from_maybe_bucket_arn(cloud.cloud_resource.aws_s3_id)
+    else:
+        return None
+
+
+def _get_cloud_gs_bucket_from_cloud(
+    api_client: DefaultApi, cloud: CloudWithCloudResource
+) -> Optional[str]:
+    """
+    If the cloud has an associated Google Storage bucket, we return its name.
+
+    Please note that this is only for v2 clouds where customers have their
+    own Google Storage.
+    """
+    assert cloud.provider == CloudProviders.GCP
+
+    gcp_cloud: CloudWithCloudResourceGCP = api_client.get_cloud_with_cloud_resource_api_v2_clouds_with_cloud_resource_gcp_router_cloud_id_get(
+        cloud.id
+    ).result
+
+    if (
+        gcp_cloud
+        and gcp_cloud.cloud_resource
+        and gcp_cloud.cloud_resource.gcp_cloud_storage_bucket_id
+    ):
+        gs_bucket_name = gcp_cloud.cloud_resource.gcp_cloud_storage_bucket_id
+        return gs_bucket_name
+    else:
+        return None

anyscale/utils/s3.py

@@ -0,0 +1,92 @@
+import json
+from typing import Any
+
+from botocore.exceptions import ClientError
+
+from anyscale.cli_logger import BlockLogger
+from anyscale.util import (  # pylint:disable=private-import
+    _coerce_to_list,
+    filter_actions_associated_with_role,
+    filter_actions_from_policy_document,
+)
+
+
+REQUIRED_S3_PERMISSIONS = {"s3:PutObject", "s3:ListBucket", "s3:GetObject"}
+
+
+def verify_s3_access(
+    boto3_session: Any, s3_bucket: Any, iam_role: Any, logger: BlockLogger
+) -> bool:
+    """This function verifies that either the specified IAM Role has an identity-based policy
+    that grants permission to access the S3 bucket OR the S3 bucket has a resource-based policy
+    that grants the IAM Role access.
+
+    NOTE: If the necessary permissions are split between the Role's policy and the Bucket's policy (i.e.
+    one has head bucket, the other has get/list objects), this function will return `False`.
+
+    TODO: Check for explicit DENY policies
+    TODO: Ensure the bucket is in the same account as the role.
+
+    TODO(#16330): Use AWS Policy Simulator or a similar policy simulation library.
+    """
+    return (
+        # black: no
+        _verify_resource_based_s3_access(s3_bucket, iam_role, logger)
+        or _verify_identity_based_s3_access(boto3_session, s3_bucket, iam_role, logger)
+    )
+
+
+def _verify_resource_based_s3_access(
+    s3_bucket: Any, iam_role: Any, logger: BlockLogger
+) -> bool:
+    try:
+        policy = s3_bucket.Policy().policy
+        policy_document = json.loads(policy)
+        allow_actions = filter_actions_from_policy_document(
+            policy_document,
+            lambda statement: statement["Effect"] == "Allow"
+            and iam_role.name in str(statement["Principal"]),
+        )
+        # TODO(#16330) ListBucket must be on `BUCKET_ARN`, while GetObject must be on `BUCKET_ARN/*`
+        if len(allow_actions) > 0:
+            missing_actions = REQUIRED_S3_PERMISSIONS - allow_actions
+            if "s3:*" in allow_actions or len(missing_actions) == 0:
+                return True
+            logger.info(
+                f"Bucket {s3_bucket.name} grants {iam_role.name} some permissions, "
+                "but is missing the following permissions:\n[{}]".format(
+                    ",".join(missing_actions)
+                )
+            )
+    except ClientError as e:
+        if e.response.get("Error", {}).get("Code") != "NoSuchBucketPolicy":
+            raise e
+
+    return False
+
+
+def _verify_identity_based_s3_access(
+    boto3_session: Any, s3_bucket: Any, iam_role: Any, logger: BlockLogger
+) -> bool:
+    # TODO(#16330) ListBucket must be on `BUCKET_ARN`, while GetObject must be on `BUCKET_ARN/*`
+    allow_actions_on_role = filter_actions_associated_with_role(
+        boto3_session,
+        iam_role,
+        lambda statement: statement["Effect"] == "Allow"
+        and (
+            "*" in _coerce_to_list(statement["Resource"])
+            or s3_bucket.name in str(statement["Resource"])
+        ),
+    )
+    if len(allow_actions_on_role) == 0:
+        return False
+    missing_actions = REQUIRED_S3_PERMISSIONS - allow_actions_on_role
+    if "s3:*" in allow_actions_on_role or len(missing_actions) == 0:
+        return True
+    logger.info(
+        f"Role {iam_role.name} has some permissions to access {s3_bucket.name}, "
+        "but is missing the following permissions:\n[{}]".format(
+            ",".join(missing_actions)
+        )
+    )
+    return False

anyscale/utils/user_utils.py

@@ -0,0 +1,17 @@
+from click import ClickException
+
+from anyscale.client.openapi_client.api.default_api import DefaultApi
+
+
+def get_user_id_by_email(api_client: DefaultApi, email: str) -> str:
+    users = api_client.list_organization_collaborators_api_v2_organization_collaborators_get(
+        email=email
+    ).results
+    if len(users) == 0:
+        raise ClickException(f"No user found with email {email}.")
+
+    if len(users) > 1:
+        raise ClickException(
+            f"Multiple users found with email {email}. Please provide a unique email."
+        )
+    return users[0].user_id

anyscale/utils/workload_types.py

@@ -0,0 +1,7 @@
+import enum
+
+
+class Workload(str, enum.Enum):
+    JOBS = "jobs"
+    SCHEDULED_JOBS = "scheduled_jobs"
+    SERVICES = "services"

anyscale/utils/workspace_notification.py

@@ -0,0 +1,39 @@
+import logging
+
+from typing_extensions import Literal
+
+from anyscale.sdk.anyscale_client.api.default_api import DefaultApi as BaseApi
+from anyscale.util import get_cluster_model_for_current_workspace
+
+
+logger = logging.getLogger(__name__)
+from anyscale.anyscale_pydantic import BaseModel
+
+
+WORKSPACE_NOTIFICATION_ADDRESS = "http://localhost:8266/simplefileserver/notifications"
+
+
+class WorkspaceNotificationAction(BaseModel):
+    type: Literal["navigate-service", "navigate-workspace-tab", "navigate-external-url"]
+    title: str
+    value: str
+
+
+class WorkspaceNotification(BaseModel):
+    body: str
+    action: WorkspaceNotificationAction
+
+
+def send_workspace_notification(
+    anyscale_api_client: BaseApi, notification: WorkspaceNotification
+):
+
+    try:
+        if get_cluster_model_for_current_workspace(anyscale_api_client) is None:
+            # We are not in a workspace, so we don't want to send a notification
+            return
+        import requests
+
+        requests.post(WORKSPACE_NOTIFICATION_ADDRESS, json=notification.dict())
+    except Exception:
+        logger.exception("Failed to send notification to UI")