anyscale 0.24.86__py3-none-any.whl

anyscale/_private/workload/workload_config.py

@@ -0,0 +1,195 @@
+from dataclasses import dataclass, field
+from typing import Dict, List, Optional, Union
+
+from anyscale._private.models import ImageURI, ModelBase
+from anyscale.compute_config.models import ComputeConfig
+
+
+@dataclass(frozen=True)
+class WorkloadConfig(ModelBase):
+    name: Optional[str] = field(
+        default=None, metadata={"docstring": "Should be overwritten by subclass."}
+    )
+
+    def _validate_name(self, name: Optional[str]):
+        if name is not None and not isinstance(name, str):
+            raise TypeError("'name' must be a string.")
+
+    image_uri: Optional[str] = field(
+        default=None,
+        metadata={
+            "docstring": "URI of an existing image. Exclusive with `containerfile`."
+        },
+    )
+
+    def _validate_image_uri(self, image_uri: Optional[str]):
+        if image_uri is not None and not isinstance(image_uri, str):
+            raise TypeError(f"'image_uri' must be an str but it is {type(image_uri)}.")
+        if image_uri is not None:
+            ImageURI.from_str(image_uri)
+
+    containerfile: Optional[str] = field(
+        default=None,
+        repr=False,
+        metadata={
+            "docstring": "The file path to a containerfile that will be built into an image before running the workload. Exclusive with `image_uri`."
+        },
+    )
+
+    def _validate_containerfile(self, containerfile: Optional[str]):
+        if containerfile is not None and not isinstance(containerfile, str):
+            raise TypeError("'containerfile' must be a string.")
+
+    compute_config: Union[ComputeConfig, Dict, str, None] = field(
+        default=None,
+        metadata={
+            "docstring": "The name of an existing registered compute config or an inlined ComputeConfig object."
+        },
+    )
+
+    def _validate_compute_config(
+        self, compute_config: Union[ComputeConfig, Dict, str, None]
+    ) -> Union[None, str, ComputeConfig]:
+        if compute_config is None or isinstance(compute_config, str):
+            return compute_config
+
+        if isinstance(compute_config, dict):
+            compute_config = ComputeConfig.from_dict(compute_config)
+        if not isinstance(compute_config, ComputeConfig):
+            raise TypeError(
+                "'compute_config' must be a string, ComputeConfig, or corresponding dict"
+            )
+
+        return compute_config
+
+    working_dir: Optional[str] = field(
+        default=None,
+        repr=False,
+        metadata={
+            "docstring": "Directory that will be used as the working directory for the application. If a local directory is provided, it will be uploaded to cloud storage automatically. When running inside a workspace, this defaults to the current working directory ('.')."
+        },
+    )
+
+    def _validate_working_dir(self, working_dir: Optional[str]):
+        if working_dir is not None and not isinstance(working_dir, str):
+            raise TypeError("'working_dir' must be a string.")
+
+    excludes: Optional[List[str]] = field(
+        default=None,
+        repr=False,
+        metadata={
+            "docstring": "A list of file path globs that will be excluded when uploading local files for `working_dir`."
+        },
+    )
+
+    def _validate_excludes(self, excludes: Optional[List[str]]):
+        if excludes is not None and (
+            not isinstance(excludes, list)
+            or not all(isinstance(e, str) for e in excludes)
+        ):
+            raise TypeError("'excludes' must be a list of strings.")
+
+    requirements: Optional[Union[str, List[str]]] = field(
+        default=None,
+        repr=False,
+        metadata={
+            "docstring": "A list of pip requirements or a path to a `requirements.txt` file for the workload. When running inside a workspace, this defaults to the workspace-tracked requirements."
+        },
+    )
+
+    def _validate_requirements(self, requirements: Optional[Union[str, List[str]]]):
+        if requirements is None or isinstance(requirements, str):
+            return
+
+        if not isinstance(requirements, list) or not all(
+            isinstance(r, str) for r in requirements
+        ):
+            raise TypeError(
+                "'requirements' must be a string (file path) or list of strings."
+            )
+
+    env_vars: Optional[Dict[str, str]] = field(
+        default=None,
+        repr=True,
+        metadata={
+            "docstring": "A dictionary of environment variables that will be set for the workload."
+        },
+    )
+
+    def _validate_env_vars(self, env_vars: Optional[Dict[str, str]]):
+        if env_vars is not None and (
+            not isinstance(env_vars, dict)
+            or not all(
+                isinstance(k, str) and isinstance(v, str) for k, v in env_vars.items()
+            )
+        ):
+            raise TypeError("'env_vars' must be a Dict[str, str].")
+
+    py_modules: Optional[List[str]] = field(
+        default=None,
+        repr=True,
+        metadata={
+            "docstring": "A list of local directories or remote URIs that will be uploaded and added to the Python path."
+        },
+    )
+
+    def _validate_py_modules(self, py_modules: Optional[List[str]]):
+        if py_modules is not None and (
+            not isinstance(py_modules, list)
+            or not all(isinstance(m, str) for m in py_modules)
+        ):
+            raise TypeError("'py_modules' must be a list of strings.")
+
+    cloud: Optional[str] = field(
+        default=None,
+        metadata={
+            "docstring": "The Anyscale Cloud to run this workload on. If not provided, the organization default will be used (or, if running in a workspace, the cloud of the workspace)."
+        },
+    )
+
+    def _validate_cloud(self, cloud: Optional[str]):
+        if cloud is not None and not isinstance(cloud, str):
+            raise TypeError("'cloud' must be a string.")
+
+    project: Optional[str] = field(
+        default=None,
+        metadata={
+            "docstring": "The project for the workload. If not provided, the default project for the cloud will be used (or, if running in a workspace, the project of the workspace)."
+        },
+    )
+
+    def _validate_project(self, project: Optional[str]):
+        if project is not None and not isinstance(project, str):
+            raise TypeError("'project' must be a string.")
+
+    registry_login_secret: Optional[str] = field(
+        default=None,
+        repr=False,
+        metadata={
+            "docstring": "A name or identifier of the secret containing credentials to authenticate to the docker registry hosting the image. "
+            "This can only be used when 'image_uri' is specified and the image is not hosted on Anyscale."
+        },
+    )
+
+    def _validate_registry_login_secret(self, registry_login_secret: Optional[str]):
+        if registry_login_secret is not None and not isinstance(
+            registry_login_secret, str
+        ):
+            raise TypeError("'registry_login_secret' must be a string.")
+
+    ray_version: Optional[str] = field(
+        default=None,
+        metadata={
+            "docstring": "The Ray version (X.Y.Z) specified for this image specified by either an image URI or a containerfile. If not provided, the latest Ray version will be used."
+        },
+    )
+
+    def _validate_ray_version(self, ray_version: Optional[str]):
+        if ray_version is not None and not isinstance(ray_version, str):
+            raise TypeError("'ray_version' must be a string.")
+        if ray_version:
+            vs = ray_version.split(".")
+            if len(vs) != 3:
+                raise ValueError(
+                    f"Invalid Ray version format: {ray_version}. Must be in the format 'X.Y.Z'."
+                )

anyscale/_private/workload/workload_sdk.py

@@ -0,0 +1,324 @@
+import copy
+from dataclasses import replace
+import pathlib
+from typing import Any, Dict, List, Optional, Tuple, Union
+
+from anyscale._private.anyscale_client import (
+    AnyscaleClientInterface,
+    WORKSPACE_CLUSTER_NAME_PREFIX,
+)
+from anyscale._private.sdk.base_sdk import BaseSDK, Timer
+from anyscale.cli_logger import BlockLogger
+from anyscale.compute_config._private.compute_config_sdk import PrivateComputeConfigSDK
+from anyscale.compute_config.models import ComputeConfig
+from anyscale.image._private.image_sdk import PrivateImageSDK
+from anyscale.utils.runtime_env import is_dir_remote_uri, parse_requirements_file
+
+
+class WorkloadSDK(BaseSDK):
+    """Shared parent class for job and service SDKs."""
+
+    def __init__(
+        self,
+        *,
+        logger: Optional[BlockLogger] = None,
+        client: Optional[AnyscaleClientInterface] = None,
+        timer: Optional[Timer] = None,
+    ):
+        super().__init__(logger=logger, client=client, timer=timer)
+        self._compute_config_sdk = PrivateComputeConfigSDK(
+            logger=self.logger, client=self.client,
+        )
+        self._image_sdk = PrivateImageSDK(logger=self.logger, client=self.client)
+
+    @property
+    def image_sdk(self) -> PrivateImageSDK:
+        return self._image_sdk
+
+    def update_env_vars(
+        self,
+        runtime_envs: List[Dict[str, Any]],
+        *,
+        env_vars_updates: Optional[Dict[str, str]],
+    ) -> List[Dict[str, Any]]:
+        """Replaces 'env_vars' fields in runtime_envs with the override."""
+        new_runtime_envs = copy.deepcopy(runtime_envs)
+
+        if env_vars_updates:
+            for runtime_env in new_runtime_envs:
+                if "env_vars" in runtime_env:
+                    # the precedence should be config > runtime_env
+                    runtime_env["env_vars"].update(env_vars_updates)
+                else:
+                    runtime_env["env_vars"] = env_vars_updates
+
+        return new_runtime_envs
+
+    def override_and_load_requirements_files(
+        self,
+        runtime_envs: List[Dict[str, Any]],
+        *,
+        requirements_override: Union[None, str, List[str]],
+        workspace_requirements_path: Optional[str] = None,
+    ) -> List[Dict[str, Any]]:
+        """Replaces 'pip' fields in runtime_envs with their parsed file contents.
+
+        The precedence for overrides is: explicit overrides passed in > fields in the existing
+        runtime_envs > workspace defaults (if `autopopulate_in_workspace == True`).
+        """
+        new_runtime_envs = copy.deepcopy(runtime_envs)
+
+        local_path_to_parsed_requirements: Dict[str, List[str]] = {}
+
+        def _load_requirements_file_memoized(
+            target: Union[str, List[str]]
+        ) -> List[str]:
+            if isinstance(target, list):
+                return target
+            elif target in local_path_to_parsed_requirements:
+                return local_path_to_parsed_requirements[target]
+            elif isinstance(target, str):
+                parsed_requirements = parse_requirements_file(target)
+                if parsed_requirements is None:
+                    raise FileNotFoundError(
+                        f"Requirements file {target} does not exist."
+                    )
+                local_path_to_parsed_requirements[target] = parsed_requirements
+                return parsed_requirements
+            else:
+                raise TypeError("pip field in runtime_env must be a list or string.")
+
+        for runtime_env in new_runtime_envs:
+            if requirements_override is not None:
+                # Explicitly-specified override from the user.
+                runtime_env["pip"] = requirements_override
+            elif (
+                workspace_requirements_path is not None
+                and "pip" not in runtime_env
+                and "conda" not in runtime_env
+            ):
+                self.logger.info("Including workspace-managed pip dependencies.")
+                runtime_env["pip"] = workspace_requirements_path
+
+            if runtime_env.get("pip", None) is not None:
+                # Load requirements from the file if necessary.
+                runtime_env["pip"] = _load_requirements_file_memoized(
+                    runtime_env["pip"]
+                )
+
+        return new_runtime_envs
+
+    def override_and_upload_local_dirs(
+        self,
+        runtime_envs: List[Dict[str, Any]],
+        *,
+        working_dir_override: Optional[str],
+        excludes_override: Optional[List[str]],
+        cloud_id: Optional[str] = None,
+        autopopulate_in_workspace: bool = True,
+        additional_py_modules: Optional[List[str]] = None,
+    ) -> List[Dict[str, Any]]:
+        """Returns modified runtime_envs with all local dirs converted to remote URIs.
+
+        The precedence for overrides is: explicit overrides passed in > fields in the existing
+        runtime_envs > workspace defaults (if `autopopulate_in_workspace == True`).
+
+        Each unique local directory across these fields will be uploaded once to cloud storage,
+        then all occurrences of it in the config will be replaced with the corresponding remote URI.
+        """
+        new_runtime_envs = copy.deepcopy(runtime_envs)
+
+        local_path_to_uri: Dict[str, str] = {}
+
+        def _upload_dir_memoized(target: str, *, excludes: Optional[List[str]]) -> str:
+            if is_dir_remote_uri(target):
+                return target
+            if target in local_path_to_uri:
+                return local_path_to_uri[target]
+
+            self.logger.info(f"Uploading local dir '{target}' to cloud storage.")
+            uri = self._client.upload_local_dir_to_cloud_storage(
+                target, cloud_id=cloud_id, excludes=excludes,
+            )
+            local_path_to_uri[target] = uri
+            return uri
+
+        for runtime_env in new_runtime_envs:
+            # Extend, don't overwrite, excludes if it's provided.
+            if excludes_override is not None:
+                existing_excludes = runtime_env.get("excludes", None) or []
+                runtime_env["excludes"] = existing_excludes + excludes_override
+
+            final_excludes = runtime_env.get("excludes", [])
+
+            new_working_dir = None
+            if working_dir_override is not None:
+                new_working_dir = working_dir_override
+            elif "working_dir" in runtime_env:
+                new_working_dir = runtime_env["working_dir"]
+            elif autopopulate_in_workspace and self._client.inside_workspace():
+                new_working_dir = "."
+
+            if new_working_dir is not None:
+                runtime_env["working_dir"] = _upload_dir_memoized(
+                    new_working_dir, excludes=final_excludes
+                )
+
+            if additional_py_modules:
+                existing_py_modules = runtime_env.get("py_modules", [])
+                runtime_env["py_modules"] = existing_py_modules + additional_py_modules
+
+            final_py_modules = runtime_env.get("py_modules", None)
+            if final_py_modules is not None:
+                runtime_env["py_modules"] = [
+                    _upload_dir_memoized(py_module, excludes=final_excludes)
+                    for py_module in final_py_modules
+                ]
+
+        return new_runtime_envs
+
+    def _resolve_compute_config_id(
+        self, compute_config: Union[str, ComputeConfig]
+    ) -> str:
+        """Resolve the passed compute config to its ID.
+
+        Accepts either:
+            - A string of the form: '<name>[:<version>]'.
+            - A dictionary from which an anonymous compute config will be built.
+
+        Returns compute_config_id.
+        """
+        if isinstance(compute_config, str):
+            compute_config_id = self._client.get_compute_config_id(
+                compute_config_name=compute_config,
+            )
+            if compute_config_id is None:
+                raise ValueError(
+                    f"The compute config '{compute_config}' does not exist."
+                )
+        else:
+            _, compute_config_id = self._compute_config_sdk.create_compute_config(
+                compute_config
+            )
+
+        return compute_config_id
+
+    def resolve_compute_config_and_cloud_id(
+        self,
+        *,
+        compute_config: Union[None, str, ComputeConfig],
+        cloud: Union[None, str],
+    ) -> Tuple[str, str]:
+        """Resolve the passed compute config to its ID and corresponding cloud ID.
+
+        Accepts either:
+            - A string of the form: '<name>[:<version>]'.
+            - A dictionary from which an anonymous compute config will be built.
+
+        Returns (compute_config_id, cloud_id).
+        """
+        if compute_config is None and cloud is None:
+            compute_config_id = self._client.get_compute_config_id()
+            assert compute_config_id is not None
+            return (
+                compute_config_id,
+                self.client.get_cloud_id(compute_config_id=compute_config_id),
+            )
+        elif compute_config is None:
+            cloud_id = self.client.get_cloud_id(cloud_name=cloud)
+            compute_config_id = self._client.get_default_compute_config(
+                cloud_id=cloud_id
+            ).id
+            return (compute_config_id, cloud_id)
+        elif cloud is None:
+            compute_config_id = self._resolve_compute_config_id(
+                compute_config=compute_config
+            )
+            return (
+                compute_config_id,
+                self.client.get_cloud_id(compute_config_id=compute_config_id),
+            )
+        else:
+            # If we are creating a new compute config, ensure cloud is set accordingly
+            if (
+                isinstance(compute_config, ComputeConfig)
+                and compute_config.cloud is None
+            ):
+                compute_config = replace(compute_config, cloud=cloud)
+
+            compute_config_id = self._resolve_compute_config_id(
+                compute_config=compute_config
+            )
+            cloud_id_from_cc = self.client.get_cloud_id(
+                compute_config_id=compute_config_id
+            )
+            cloud_id_from_cloud = self.client.get_cloud_id(cloud_name=cloud)
+            if cloud_id_from_cc != cloud_id_from_cloud:
+                raise ValueError(
+                    "Cloud does not match from provided `cloud` and `compute_config`. Either pass one of these or ensure they match."
+                )
+            else:
+                return (compute_config_id, cloud_id_from_cc)
+
+    def get_current_workspace_name(self) -> Optional[str]:
+        """Get the name of the curernt workspace if running inside one."""
+        if not self._client.inside_workspace():
+            return None
+
+        name = self._client.get_current_workspace_cluster().name
+        # Defensively default to the workspace cluster name as-is if it doesn't
+        # start with the expected prefix.
+        if name.startswith(WORKSPACE_CLUSTER_NAME_PREFIX):
+            name = name[len(WORKSPACE_CLUSTER_NAME_PREFIX) :]
+
+        return name
+
+    def get_containerfile_contents(self, path: str) -> str:
+        """Get the full content of the containerfile as a string."""
+        containerfile_path = pathlib.Path(path)
+        if not containerfile_path.exists():
+            raise FileNotFoundError(
+                f"Containerfile '{containerfile_path}' does not exist."
+            )
+        if not containerfile_path.is_file():
+            raise ValueError(f"Containerfile '{containerfile_path}' must be a file.")
+
+        return containerfile_path.read_text()
+
+    def _strip_none_values_from_dict(self, d: Dict) -> Dict:
+        """Return a copy of the dictionary without any keys whose values are None.
+
+        Recursively calls into any dictionary values.
+        """
+        result = {}
+        for k, v in d.items():
+            if isinstance(v, dict):
+                result[k] = self._strip_none_values_from_dict(v)
+            elif v is not None:
+                result[k] = v
+
+        return result
+
+    def get_user_facing_compute_config(
+        self, compute_config_id: str,
+    ) -> Union[str, ComputeConfig]:
+        """Get the compute config in a format to be displayed in a user-facing status.
+
+        If the compute config refers to an anonymous compute config, its config
+        object will be returned. Else the name of the compute config will be
+        returned in the form: '<name>:<version>'.
+        """
+        compute_config = self._client.get_compute_config(compute_config_id)
+        if compute_config is None:
+            raise RuntimeError(
+                f"Failed to get compute config for ID {compute_config_id}."
+            )
+
+        compute_config_name = compute_config.name
+        if compute_config.version is not None:
+            compute_config_name += f":{compute_config.version}"
+
+        if not compute_config.anonymous:
+            return compute_config_name
+
+        return self._compute_config_sdk.get_compute_config(id=compute_config_id).config

anyscale/aggregated_instance_usage/__init__.py

@@ -0,0 +1,36 @@
+from typing import Optional
+
+from anyscale._private.anyscale_client import AnyscaleClientInterface
+from anyscale._private.sdk import sdk_docs
+from anyscale._private.sdk.base_sdk import Timer
+from anyscale.aggregated_instance_usage._private.aggregated_instance_usage_sdk import (
+    PrivateAggregatedInstanceUsageSDK,
+)
+from anyscale.aggregated_instance_usage.commands import (
+    _DOWNLOAD_ARG_DOCSTRINGS,
+    _DOWNLOAD_CSV_EXAMPLE,
+    download_csv,
+)
+from anyscale.aggregated_instance_usage.models import DownloadCSVFilters
+from anyscale.cli_logger import BlockLogger
+
+
+class AggregatedInstanceUsageSDK:
+    def __init__(
+        self,
+        *,
+        client: Optional[AnyscaleClientInterface] = None,
+        logger: Optional[BlockLogger] = None,
+        timer: Optional[Timer] = None,
+    ):
+        self._private_sdk = PrivateAggregatedInstanceUsageSDK(
+            client=client, logger=logger, timer=timer
+        )
+
+    @sdk_docs(
+        doc_py_example=_DOWNLOAD_CSV_EXAMPLE, arg_docstrings=_DOWNLOAD_ARG_DOCSTRINGS,
+    )
+    def download_csv(self, filters: DownloadCSVFilters,) -> str:  # noqa: F811
+        """Downloads aggregated instance usage CSV.
+        """
+        return self._private_sdk.download_csv(filters=filters)

anyscale/aggregated_instance_usage/_private/aggregated_instance_usage_sdk.py

@@ -0,0 +1,30 @@
+from datetime import date
+
+from anyscale._private.sdk.base_sdk import BaseSDK
+from anyscale.aggregated_instance_usage.models import DownloadCSVFilters
+
+
+class PrivateAggregatedInstanceUsageSDK(BaseSDK):
+    def download_csv(self, filters: DownloadCSVFilters) -> str:
+        cloud_id = (
+            self.client.get_cloud_id(cloud_name=filters.cloud, compute_config_id=None)
+            if filters.cloud
+            else None
+        )
+
+        project_id = (
+            self.client.get_project_id(parent_cloud_id=cloud_id, name=filters.project)
+            if filters.project
+            else None
+        )
+
+        return self.client.download_aggregated_instance_usage_csv(
+            start_date=date.fromisoformat(filters.start_date),
+            end_date=date.fromisoformat(filters.end_date),
+            cloud_id=cloud_id,
+            project_id=project_id,
+            directory=filters.directory,
+            hide_progress_bar=filters.hide_progress_bar
+            if filters.hide_progress_bar
+            else False,
+        )

anyscale/aggregated_instance_usage/commands.py

@@ -0,0 +1,42 @@
+from anyscale._private.sdk import sdk_command
+from anyscale.aggregated_instance_usage._private.aggregated_instance_usage_sdk import (
+    PrivateAggregatedInstanceUsageSDK,
+)
+from anyscale.aggregated_instance_usage.models import DownloadCSVFilters
+
+
+_AGGREGATED_INSTANCE_USAGE_SDK_SINGLETON_KEY = "aggregated_instance_usage_sdk"
+
+_DOWNLOAD_CSV_EXAMPLE = """
+import anyscale
+from anyscale.aggregated_instance_usage.models import DownloadCSVFilters
+
+anyscale.aggregated_instance_usage.download_csv(
+    DownloadCSVFilters(
+        start_date="2024-10-01",
+        end_date="2024-10-02",
+        cloud="cloud_name",
+        project="project_name",
+        directory="/directory",
+        hide_progress_bar=False,
+    ),
+)
+"""
+
+_DOWNLOAD_ARG_DOCSTRINGS = {
+    "filters": "The filter of the instance usage to be downloaded."
+}
+
+
+@sdk_command(
+    _AGGREGATED_INSTANCE_USAGE_SDK_SINGLETON_KEY,
+    PrivateAggregatedInstanceUsageSDK,
+    doc_py_example=_DOWNLOAD_CSV_EXAMPLE,
+    arg_docstrings=_DOWNLOAD_ARG_DOCSTRINGS,
+)
+def download_csv(
+    filters: DownloadCSVFilters, *, _sdk: PrivateAggregatedInstanceUsageSDK
+) -> str:
+    """Downloads aggregated instance usage CSV.
+    """
+    return _sdk.download_csv(filters)