anyscale 0.24.63__tar.gz → 0.24.65__tar.gz

{anyscale-0.24.63 → anyscale-0.24.65}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: anyscale
-Version: 0.24.63
+Version: 0.24.65
 Summary: Command Line Interface for Anyscale
 Author: Anyscale Inc.
 License: AS License

{anyscale-0.24.63 → anyscale-0.24.65}/anyscale/_private/anyscale_client/anyscale_client.py

@@ -11,6 +11,7 @@ from urllib.parse import urlparse
 
 from openapi_client.exceptions import ApiException
 import requests
+from rich.style import Style
 
 from anyscale._private.anyscale_client.common import (
     AnyscaleClientInterface,
@@ -19,6 +20,10 @@ from anyscale._private.anyscale_client.common import (
     RUNTIME_ENV_PACKAGE_FORMAT,
 )
 from anyscale._private.models.image_uri import ImageURI
+from anyscale._private.utils.progress_util import (
+    FileDownloadProgress,
+    ProgressFileReader,
+)
 from anyscale.api_utils.common_utils import source_cloud_id_and_project_id
 from anyscale.api_utils.logs_util import _download_log_from_s3_url_sync
 from anyscale.authenticate import AuthenticationBlock, get_auth_api_client
@@ -35,8 +40,11 @@ from anyscale.client.openapi_client.models import (
     ComputeTemplateConfig,
     ComputeTemplateQuery,
     CreateComputeTemplate,
+    CreateDataset,
     CreateExperimentalWorkspace,
     CreateInternalProductionJob,
+    Dataset,
+    DatasetUpload,
     DecoratedComputeTemplate,
     DeletedPlatformFineTunedModel,
     ExperimentalWorkspace,
@@ -80,7 +88,9 @@ from anyscale.sdk.anyscale_client.models.sort_by_clause_jobs_sort_field import (
     SortByClauseJobsSortField,
 )
 from anyscale.sdk.anyscale_client.models.sort_order import SortOrder
+from anyscale.sdk.anyscale_client.models.update_cluster import UpdateCluster
 from anyscale.sdk.anyscale_client.rest import ApiException as ExternalApiException
+from anyscale.shared_anyscale_utils.bytes_util import Bytes
 from anyscale.shared_anyscale_utils.conf import ANYSCALE_HOST
 from anyscale.shared_anyscale_utils.latest_ray_version import LATEST_RAY_VERSION
 from anyscale.util import (
@@ -357,18 +367,29 @@ class AnyscaleClient(AnyscaleClientInterface):
     def _get_project_id_by_name(
         self, *, parent_cloud_id: Optional[str] = None, name: Optional[str] = None
     ) -> str:
+        """Resolves `name`, a project name, to a project ID.
+
+        Args:
+            parent_cloud_id: If specified, return the project that has this `parent_cloud_id`. \
+                Else (`None`), return the first project with the given `name`. \
+                Defaults to `None`.
+            name: The name of the project.
+        """
         # Find if project with name already exists
-        existing_projects = self._internal_api_client.find_project_by_project_name_api_v2_projects_find_by_name_get(
+        matching_projects = self._internal_api_client.find_project_by_project_name_api_v2_projects_find_by_name_get(
             name
         ).results
-        if len(existing_projects) == 0:
+        if len(matching_projects) == 0:
             raise ValueError(f"Project '{name}' was not found.")
         else:
-            for project in existing_projects:
-                if project.parent_cloud_id == parent_cloud_id:
+            for project in matching_projects:
+                if (
+                    parent_cloud_id is None
+                    or project.parent_cloud_id == parent_cloud_id
+                ):
                     return project.id
             raise ValueError(
-                f"{len(existing_projects)} project(s) found with name '{name}' and none matched cloud_id '{parent_cloud_id}'"
+                f"{len(matching_projects)} project(s) found with name '{name}' and none matched cloud_id '{parent_cloud_id}'"
             )
 
     def _get_project_id_by_cloud_id(
@@ -606,7 +627,7 @@ class AnyscaleClient(AnyscaleClientInterface):
         return self._external_api_client.get_cluster_environment_build(build_id).result
 
     def get_cluster_env_build_image_uri(
-        self, cluster_env_build_id: str
+        self, cluster_env_build_id: str, use_image_alias: bool = False
     ) -> Optional[ImageURI]:
         try:
             build: ClusterEnvironmentBuild = self._external_api_client.get_cluster_environment_build(
@@ -615,7 +636,7 @@ class AnyscaleClient(AnyscaleClientInterface):
             cluster_env = self._external_api_client.get_cluster_environment(
                 build.cluster_environment_id
             ).result
-            return ImageURI.from_cluster_env_build(cluster_env, build)
+            return ImageURI.from_cluster_env_build(cluster_env, build, use_image_alias)
         except ExternalApiException as e:
             if e.status == 404:
                 return None
@@ -720,10 +741,7 @@ class AnyscaleClient(AnyscaleClientInterface):
         cluster_env = self._find_or_create_cluster_env(
             cluster_env_name, anonymous=anonymous
         )
-        cluster_env_builds = self._external_api_client.list_cluster_environment_builds(
-            cluster_environment_id=cluster_env.id, count=self.LIST_ENDPOINT_COUNT
-        ).results
-        for build in cluster_env_builds:
+        for build in self.list_cluster_env_builds(cluster_env.id):
             if (
                 build.status == ClusterEnvironmentBuildStatus.SUCCEEDED
                 and build.containerfile == containerfile
@@ -760,6 +778,7 @@ class AnyscaleClient(AnyscaleClientInterface):
         image_uri: ImageURI,
         registry_login_secret: Optional[str] = None,
         ray_version: Optional[str] = None,
+        name: Optional[str] = None,
     ) -> str:
         if image_uri.is_cluster_env_image():
             identifier = image_uri.to_cluster_env_identifier()
@@ -779,22 +798,20 @@ class AnyscaleClient(AnyscaleClientInterface):
                         f"Legacy cluster environment '{identifier}' is not found."
                     )
 
-        cluster_env_name = image_uri.to_cluster_env_name()
-        cluster_env = self._find_or_create_cluster_env(cluster_env_name, anonymous=True)
+        cluster_env_name = name if name else image_uri.to_cluster_env_name()
+        cluster_env = self._find_or_create_cluster_env(
+            cluster_env_name, anonymous=not name
+        )
         image_uri_str = str(image_uri)
 
-        cluster_env_builds = self._external_api_client.list_cluster_environment_builds(
-            cluster_environment_id=cluster_env.id, count=1
-        ).results
-        build = cluster_env_builds[0] if cluster_env_builds else None
-        if (
-            build is not None
-            and build.docker_image_name == image_uri_str
-            and build.registry_login_secret == registry_login_secret
-            and build.status == ClusterEnvironmentBuildStatus.SUCCEEDED
-            and build.ray_version == ray_version
-        ):
-            return build.id
+        for build in self.list_cluster_env_builds(cluster_env.id):
+            if (
+                build.docker_image_name == image_uri_str
+                and build.registry_login_secret == registry_login_secret
+                and build.status == ClusterEnvironmentBuildStatus.SUCCEEDED
+                and build.ray_version == ray_version
+            ):
+                return build.id
 
         # Still create a new build if the cluster env already exists but the build does not match the image_uri.
         result = self._external_api_client.create_cluster_environment_build(
@@ -1107,34 +1124,113 @@ class AnyscaleClient(AnyscaleClientInterface):
             id
         )
 
+    def get_dataset(
+        self, name: str, version: Optional[int], project: Optional[str]
+    ) -> Dataset:
+        project_id = self._source_project_id(project)
+        dataset = self._internal_api_client.find_dataset_api_v2_datasets_find_get(
+            name=name, version=version, project_id=project_id
+        ).result
+        return dataset
+
     def upload_dataset(
         self,
         dataset_file: str,
-        existing_dataset_id_or_name: Optional[str],
+        name: Optional[str],
         description: Optional[str],
-        cloud_id: Optional[str],
-        project_id: Optional[str],
-    ):
+        cloud: Optional[str],
+        project: Optional[str],
+    ) -> Dataset:
+        # Resolve `~/.../file` to `/home/user/.../file`
+        dataset_file = os.path.expanduser(dataset_file)
+
         if not os.path.isfile(dataset_file):
             raise ValueError(f"Path '{dataset_file}' is not a valid file.")
+        dataset_file_size = os.path.getsize(dataset_file)
+        if dataset_file_size > 5 * Bytes.GB:
+            raise ValueError(
+                f"File '{dataset_file}' is too large to upload. The maximum size is 5 GB."
+            )
+        project_id = self._get_project_id_by_name(name=project) if project else None
+        cloud_id = self.get_cloud_id(cloud_name=cloud) if cloud else None
 
-        cloud_id, project_id = source_cloud_id_and_project_id(
-            internal_api=self._internal_api_client,
-            external_api=self._external_api_client,
-            cloud_id=cloud_id,
-            project_id=project_id,
-        )
-        print(existing_dataset_id_or_name, description)
+        with FileDownloadProgress() as progress:
+            task_id = progress.add_task(
+                description=f"Creating an upload request for '{dataset_file}'",
+                total=dataset_file_size,
+            )
+            _, project_id = source_cloud_id_and_project_id(
+                internal_api=self._internal_api_client,
+                external_api=self._external_api_client,
+                cloud_id=cloud_id,
+                project_id=project_id,
+            )
+            dataset_upload: DatasetUpload = self._internal_api_client.create_dataset_upload_api_v2_datasets_upload_post(
+                create_dataset=CreateDataset(
+                    filename=os.path.basename(dataset_file),
+                    description=description,
+                    name=name,
+                    project_id=project_id,
+                )
+            ).result
+
+            progress.update(task_id, description=f"Uploading '{dataset_file}'")
+
+            with open(dataset_file, "rb") as file_reader:
+                progress_reader = ProgressFileReader(file_reader, progress, task_id)
+                response = requests.put(
+                    dataset_upload.upload_url, data=progress_reader,
+                )
+                response.raise_for_status()
+
+            progress.update(task_id, completed=os.path.getsize(dataset_file))
+            progress.console.print(
+                "Upload complete!", style=Style(bold=True, color="green")
+            )
+        return dataset_upload.dataset
 
     def download_dataset(
-        self, dataset_id_or_name: str, version: Optional[int]  # noqa: ARG002
+        self, name: str, version: Optional[int], project: Optional[str]
     ) -> bytes:
-        # TODO: replace with self._external_api_client.get_dataset_download_url()
-        dataset_bytes_url = "https://example.com/dataset_bytes"
-        response = requests.get(dataset_bytes_url)
-        dataset_bytes = response.content
+        project_id = self._source_project_id(project)
+        with FileDownloadProgress() as progress:
+            task_id = progress.add_task(
+                description=f"Getting download info for '{name}'",
+            )
+            download_url: str = self._internal_api_client.get_dataset_download_url_api_v2_datasets_download_get(
+                name, version=version, project_id=project_id,
+            )
+            progress.update(task_id, description=f"Downloading '{name}'")
+            response = requests.get(download_url, stream=True)
+            total_size = int(response.headers.get("content-length", 0))
+            progress.update(task_id, total=total_size)
+
+            # For CLI, consider writing to disk instead of loading the entire file into memory.
+            dataset_bytes = b""
+            for data in response.iter_content(Bytes.MB):
+                dataset_bytes += data
+                progress.update(task_id, advance=len(data))
+
+            progress.update(task_id, completed=total_size)
+            progress.console.print(
+                "Download complete!", style=Style(bold=True, color="green")
+            )
+
         return dataset_bytes
 
+    def _source_project_id(self, project_name: Optional[str]) -> str:
+        """Sources a project ID from an optionally-provided project name."""
+        if project_name:
+            project_id = self._get_project_id_by_name(name=project_name)
+        else:
+            _, project_id = source_cloud_id_and_project_id(
+                internal_api=self._internal_api_client,
+                external_api=self._external_api_client,
+                cloud_id=None,
+                project_id=None,
+            )
+        return project_id
+
     def get_finetuned_model(
         self, model_id: Optional[str], job_id: Optional[str]
     ) -> FineTunedModel:
@@ -1193,6 +1289,40 @@ class AnyscaleClient(AnyscaleClientInterface):
             workspace = resp.results[0]
             return workspace
 
+    @handle_api_exceptions
+    def update_workspace(
+        self,
+        *,
+        workspace_id: str,
+        name: Optional[str] = None,
+        compute_config_id: Optional[str] = None,
+        cluster_environment_build_id: Optional[str] = None,
+        idle_timeout_minutes: Optional[int] = None,
+    ):
+        workspace = self.get_workspace(id=workspace_id)
+        if not workspace:
+            raise ValueError(f"Workspace '{workspace_id}' not found.")
+
+        if name:
+            # Update the workspace name with workspaces patch API
+            self._internal_api_client.patch_workspace_api_v2_experimental_workspaces_workspace_id_patch(
+                workspace_id=workspace_id,
+                json_patch_operation=[
+                    {"op": "replace", "path": "/name", "value": name,},
+                ],
+            )
+
+        if compute_config_id or cluster_environment_build_id or idle_timeout_minutes:
+            # Update cluster with external cluster API
+            self._external_api_client.update_cluster(
+                cluster_id=workspace.cluster_id,
+                update_cluster=UpdateCluster(
+                    idle_timeout_minutes=idle_timeout_minutes,
+                    cluster_environment_build_id=cluster_environment_build_id,
+                    cluster_compute_id=compute_config_id,
+                ),
+            )
+
     @handle_api_exceptions
     def update_workspace_dependencies_offline_only(
         self, workspace_id: str, requirements: List[str]

{anyscale-0.24.63 → anyscale-0.24.65}/anyscale/_private/anyscale_client/common.py

@@ -7,6 +7,7 @@ from anyscale.client.openapi_client.models import (
     ComputeTemplateConfig,
     CreateExperimentalWorkspace,
     CreateInternalProductionJob,
+    Dataset,
     DecoratedComputeTemplate,
     DeletedPlatformFineTunedModel,
     FineTunedModel,
@@ -36,8 +37,8 @@ from anyscale.utils.workspace_notification import WorkspaceNotification
 # Maybe just make it part of the release process to update it, or fetch the
 # default builds and get the latest one. The best thing to do is probably
 # to populate this in the backend.
-DEFAULT_RAY_VERSION = "2.10.0"
-DEFAULT_PYTHON_VERSION = "py310"
+DEFAULT_RAY_VERSION = "2.36.0"
+DEFAULT_PYTHON_VERSION = "py311"
 RUNTIME_ENV_PACKAGE_FORMAT = "pkg_{content_hash}.zip"
 
 # All workspace cluster names should start with this prefix.
@@ -236,10 +237,11 @@ class AnyscaleClientInterface(ABC):
         image_uri: ImageURI,
         registry_login_secret: Optional[str] = None,
         ray_version: Optional[str] = None,
+        name: Optional[str] = None,
     ):
         """Get the cluster environment build ID for the cluster environment with provided image_uri.
 
-        It maps a image_uri to a cluster env name and then use the name to get the cluster env.
+        It maps a image_uri to a cluster env name if a name is not provided, and then use the name to get the cluster env.
         If there exists a cluster environment for the image_uri, it will reuse the cluster env.
         Else it will create a new cluster environment.
         It will create a new cluster environment build with the provided image_uri or try to reuse one with the same image_uri if exists.
@@ -252,9 +254,10 @@ class AnyscaleClientInterface(ABC):
 
     @abstractmethod
     def get_cluster_env_build_image_uri(
-        self, cluster_env_build_id: str
+        self, cluster_env_build_id: str, use_image_alias: bool = False
     ) -> Optional[ImageURI]:
         """Get the image URI for the provided build ID.
+        If `use_image_alias` is True, the container image alias will be used if the build is BYOD.
 
         Returns None if not found.
         """
@@ -408,22 +411,30 @@ class AnyscaleClientInterface(ABC):
         """
         raise NotImplementedError
 
+    @abstractmethod
+    def get_dataset(
+        self, name: str, version: Optional[int], project: Optional[str]
+    ) -> Dataset:
+        """See docstring for `anyscale.llm.dataset.get()`."""
+        raise NotImplementedError
+
+    @abstractmethod
     def upload_dataset(
         self,
         dataset_file: str,
-        existing_dataset_id_or_name: Optional[str],
+        name: Optional[str],
         description: Optional[str],
-        cloud_id: Optional[str],
-        project_id: Optional[str],
-    ):
-        """See docstring for `anyscale.llm.dataset.upload`."""
+        cloud: Optional[str],
+        project: Optional[str],
+    ) -> Dataset:
+        """See docstring for `anyscale.llm.dataset.upload()`."""
         raise NotImplementedError
 
     @abstractmethod
     def download_dataset(
-        self, dataset_id_or_name: str, version: Optional[int]
+        self, name: str, version: Optional[int], project: Optional[str]
     ) -> bytes:
-        """See docstring for `anyscale.llm.dataset.download`."""
+        """See docstring for `anyscale.llm.dataset.download()`."""
 
     @abstractmethod
     def create_workspace(self, model: CreateExperimentalWorkspace) -> str:
@@ -509,3 +520,16 @@ class AnyscaleClientInterface(ABC):
         self, cloud_id: Optional[str], project_id: Optional[str], max_items: int,
     ) -> List[FineTunedModel]:
         raise NotImplementedError
+
+    @abstractmethod
+    def update_workspace(
+        self,
+        *,
+        workspace_id: str,
+        name: Optional[str] = None,
+        compute_config_id: Optional[str] = None,
+        cluster_environment_build_id: Optional[str] = None,
+        idle_timeout_minutes: Optional[int] = None,
+    ):
+        """Update a workspace."""
+        raise NotImplementedError

{anyscale-0.24.63 → anyscale-0.24.65}/anyscale/_private/anyscale_client/fake_anyscale_client.py

@@ -15,6 +15,7 @@ from anyscale.client.openapi_client.models import (
     ComputeTemplateConfig,
     CreateExperimentalWorkspace,
     CreateInternalProductionJob,
+    Dataset,
     DecoratedComputeTemplate,
     DeletedPlatformFineTunedModel,
     ExperimentalWorkspace,
@@ -495,6 +496,7 @@ class FakeAnyscaleClient(AnyscaleClientInterface):
         image_uri: ImageURI,
         registry_login_secret: Optional[str] = None,
         ray_version: Optional[str] = None,
+        name: Optional[str] = None,
     ) -> str:
         for build in self._builds.values():
             build_image_uri = self.get_cluster_env_build_image_uri(build.id)
@@ -505,7 +507,7 @@ class FakeAnyscaleClient(AnyscaleClientInterface):
             ):
                 return build.id  # type: ignore
         cluster_env_id = f"cluster-env-id-{uuid.uuid4()!s}"
-        cluster_env_name = image_uri.to_cluster_env_name()
+        cluster_env_name = name if name else image_uri.to_cluster_env_name()
         self._images[cluster_env_id] = ClusterEnvironment(
             id=cluster_env_id,
             name=cluster_env_name,
@@ -525,13 +527,13 @@ class FakeAnyscaleClient(AnyscaleClientInterface):
         return build_id
 
     def get_cluster_env_build_image_uri(
-        self, cluster_env_build_id: str
+        self, cluster_env_build_id: str, use_image_alias: bool = False
     ) -> Optional[ImageURI]:
         build = self._builds.get(cluster_env_build_id, None)
         if build is None:
             return None
         else:
-            if build.docker_image_name is not None:
+            if build.docker_image_name is not None and not use_image_alias:
                 return ImageURI.from_str(build.docker_image_name)
             else:
                 cluster_env = self._images[build.cluster_environment_id]
@@ -818,18 +820,23 @@ class FakeAnyscaleClient(AnyscaleClientInterface):
     def trigger_counts(self, id: str):  # noqa: A002
         return self._schedule_trigger_counts[id]
 
+    def get_dataset(
+        self, name: str, version: Optional[int], project: Optional[str]  # noqa: ARG002
+    ) -> Dataset:
+        pass
+
     def upload_dataset(
         self,
         dataset_file: str,  # noqa: ARG002
-        existing_dataset_id_or_name: Optional[str],  # noqa: ARG002
+        name: Optional[str],  # noqa: ARG002
         description: Optional[str],  # noqa: ARG002
-        cloud_id: Optional[str],  # noqa: ARG002
-        project_id: Optional[str],  # noqa: ARG002
-    ):
+        cloud: Optional[str],  # noqa: ARG002
+        project: Optional[str],  # noqa: ARG002
+    ) -> Dataset:
         pass
 
     def download_dataset(
-        self, dataset_id_or_name: str, version: Optional[int]  # noqa: ARG002
+        self, name: str, version: Optional[int], project: Optional[str]  # noqa: ARG002
     ) -> bytes:
         return b""
 
@@ -968,3 +975,33 @@ class FakeAnyscaleClient(AnyscaleClientInterface):
             )
             for _ in range(max_items)
         ]
+
+    def update_workspace(
+        self,
+        *,
+        workspace_id: str,
+        name: Optional[str] = None,
+        compute_config_id: Optional[str] = None,
+        cluster_environment_build_id: Optional[str] = None,
+        idle_timeout_minutes: Optional[int] = None,
+    ):
+        if workspace_id not in self._workspaces:
+            raise ValueError(f"Workspace '{workspace_id}' not found.")
+
+        workspace = self._workspaces[workspace_id]
+
+        if name:
+            workspace.name = name
+
+        if compute_config_id:
+            workspace.compute_config_id = compute_config_id
+
+        if cluster_environment_build_id:
+            workspace.environment_id = cluster_environment_build_id
+
+        if idle_timeout_minutes:
+            workspace.idle_timeout_minutes = idle_timeout_minutes
+
+            compute_config = self.get_compute_config(workspace.compute_config_id)
+            assert compute_config is not None
+            compute_config.idle_timeout_minutes = idle_timeout_minutes