anyscale 0.24.87__py3-none-any.whl → 0.24.91__py3-none-any.whl

anyscale/__init__.py

@@ -20,19 +20,25 @@ path.append(os.path.join(anyscale_dir, "sdk"))
 import anyscale
 from anyscale import (
     aggregated_instance_usage,
+    cloud,
     compute_config,
     image,
     integrations,
     job,
     llm,
+    organization_invitation,
+    project,
+    resource_quota,
     schedule,
     service,
+    user,
 )
 from anyscale._private.anyscale_client import AnyscaleClient, AnyscaleClientInterface
 from anyscale._private.sdk.base_sdk import Timer
 from anyscale.aggregated_instance_usage import AggregatedInstanceUsageSDK
 from anyscale.authenticate import AuthenticationBlock
 from anyscale.cli_logger import BlockLogger
+from anyscale.cloud import CloudSDK
 from anyscale.cluster import get_job_submission_client_cluster_info
 from anyscale.cluster_compute import get_cluster_compute_from_name
 from anyscale.compute_config import ComputeConfigSDK
@@ -40,9 +46,13 @@ from anyscale.connect import ClientBuilder
 from anyscale.image import ImageSDK
 from anyscale.job import JobSDK
 from anyscale.llm import LLMSDK
+from anyscale.organization_invitation import OrganizationInvitationSDK
+from anyscale.project import ProjectSDK
+from anyscale.resource_quota import ResourceQuotaSDK
 from anyscale.schedule import ScheduleSDK
 from anyscale.sdk.anyscale_client.sdk import AnyscaleSDK
 from anyscale.service import ServiceSDK
+from anyscale.user import UserSDK
 from anyscale.workspace import WorkspaceSDK
 
 
@@ -56,6 +66,10 @@ for attr, _ in inspect.getmembers(ClientBuilder, inspect.isfunction):
     if attr.startswith("_"):
         continue
 
+    # This is exposed in anyscale/cloud/__init__.py since anyscale.cloud is used as the SDK module too.
+    if attr == "cloud":
+        continue
+
     def _new_builder(attr: str) -> Any:
         target = getattr(ClientBuilder, attr)
 
@@ -105,9 +119,16 @@ class Anyscale:
         self._job_sdk = JobSDK(client=self._anyscale_client)
         self._service_sdk = ServiceSDK(client=self._anyscale_client)
         self._compute_config_sdk = ComputeConfigSDK(client=self._anyscale_client)
+        self._cloud_sdk = CloudSDK(client=self._anyscale_client)
         self._schedule_sdk = ScheduleSDK(client=self._anyscale_client)
         self._image_sdk = ImageSDK(client=self._anyscale_client)
         self._llm_sdk = LLMSDK(client=self._anyscale_client)
+        self._organization_invitation_sdk = OrganizationInvitationSDK(
+            client=self._anyscale_client
+        )
+        self._project_sdk = ProjectSDK(client=self._anyscale_client)
+        self._resource_quota_sdk = ResourceQuotaSDK(client=self._anyscale_client)
+        self._user_sdk = UserSDK(client=self._anyscale_client)
         self._workspace_sdk = WorkspaceSDK(client=self._anyscale_client)
 
     @classmethod
@@ -128,10 +149,25 @@ class Anyscale:
         obj._compute_config_sdk = ComputeConfigSDK(  # noqa: SLF001
             client=client, logger=logger, timer=timer
         )
+        obj._cloud_sdk = CloudSDK(  # noqa: SLF001
+            client=client, logger=logger, timer=timer
+        )
         obj._schedule_sdk = ScheduleSDK(  # noqa: SLF001
             client=client, logger=logger, timer=timer,
         )
         obj._image_sdk = ImageSDK(client=client, logger=logger)  # noqa: SLF001
+        obj._organization_invitation_sdk = OrganizationInvitationSDK(  # noqa: SLF001
+            client=client, logger=logger
+        )
+        obj._project_sdk = ProjectSDK(  # noqa: SLF001
+            client=client, logger=logger, timer=timer
+        )
+        obj._user_sdk = UserSDK(  # noqa: SLF001
+            client=client, logger=logger, timer=timer
+        )
+        obj._resource_quota_sdk = ResourceQuotaSDK(  # noqa: SLF001
+            client=client, logger=logger, timer=timer
+        )
         obj._workspace_sdk = WorkspaceSDK(  # noqa: SLF001
             client=client, logger=logger, timer=timer,
         )
@@ -153,6 +189,10 @@ class Anyscale:
     def compute_config(self) -> ComputeConfigSDK:  # noqa: F811
         return self._compute_config_sdk
 
+    @property
+    def cloud(self) -> CloudSDK:  # noqa: F811
+        return self._cloud_sdk
+
     @property
     def schedule(self) -> ScheduleSDK:  # noqa: F811
         return self._schedule_sdk
@@ -165,6 +205,22 @@ class Anyscale:
     def llm(self) -> LLMSDK:  # noqa: F811
         return self._llm_sdk
 
+    @property
+    def organization_invitation(self) -> OrganizationInvitationSDK:  # noqa: F811
+        return self._organization_invitation_sdk
+
+    @property
+    def project(self) -> ProjectSDK:  # noqa: F811
+        return self._project_sdk
+
+    @property
+    def resource_quota(self) -> ResourceQuotaSDK:  # noqa: F811
+        return self._resource_quota_sdk
+
+    @property
+    def user(self) -> UserSDK:  # noqa: F811
+        return self._user_sdk
+
     @property
     def workspace(self) -> WorkspaceSDK:  # noqa: F811
         return self._workspace_sdk

anyscale/_private/anyscale_client/anyscale_client.py

@@ -34,6 +34,8 @@ from anyscale.authenticate import AuthenticationBlock, get_auth_api_client
 from anyscale.cli_logger import BlockLogger
 from anyscale.client.openapi_client.api.default_api import DefaultApi as InternalApi
 from anyscale.client.openapi_client.models import (
+    AdminCreatedUser,
+    AdminCreateUser,
     ArchiveStatus,
     Cloud,
     CloudDataBucketAccessMode,
@@ -45,10 +47,14 @@ from anyscale.client.openapi_client.models import (
     ComputeTemplate,
     ComputeTemplateConfig,
     ComputeTemplateQuery,
+    CreateCloudCollaborator,
     CreateComputeTemplate,
     CreateDataset,
     CreateExperimentalWorkspace,
     CreateInternalProductionJob,
+    CreateOrganizationInvitation,
+    CreateResourceQuota,
+    CreateUserProjectCollaborator,
     Dataset as InternalDataset,
     DatasetUpload,
     DecoratedComputeTemplate,
@@ -58,6 +64,11 @@ from anyscale.client.openapi_client.models import (
     FineTunedModel,
     FinetunedmodelListResponse,
     InternalProductionJob,
+    ListResourceQuotasQuery,
+    OrganizationCollaborator,
+    OrganizationInvitation,
+    ResourceQuota,
+    ResourceQuotaStatus,
     SessionSshKey,
     StartSessionOptions,
     StopSessionOptions,
@@ -542,6 +553,14 @@ class AnyscaleClient(AnyscaleClientInterface):
 
             raise e from None
 
+    @handle_api_exceptions
+    def add_cloud_collaborators(
+        self, cloud_id: str, collaborators: List[CreateCloudCollaborator]
+    ) -> None:
+        self._internal_api_client.batch_create_cloud_collaborators_api_v2_clouds_cloud_id_collaborators_users_batch_create_post(
+            cloud_id, collaborators
+        )
+
     @handle_api_exceptions
     def create_compute_config(
         self, config: ComputeTemplateConfig, *, name: Optional[str] = None
@@ -984,6 +1003,14 @@ class AnyscaleClient(AnyscaleClientInterface):
             project_id
         ).result
 
+    @handle_api_exceptions
+    def add_project_collaborators(
+        self, project_id: str, collaborators: List[CreateUserProjectCollaborator]
+    ) -> None:
+        self._internal_api_client.batch_create_project_collaborators_api_v2_projects_project_id_collaborators_users_batch_create_post(
+            project_id, collaborators
+        )
+
     @handle_api_exceptions
     def get_job(
         self,
@@ -1120,39 +1147,26 @@ class AnyscaleClient(AnyscaleClientInterface):
                 # If the presigned URL scheme is SMART_OPEN, upload to cloud storage using the provided bucket name, path, & environment, and the smart_open library.
                 bucket_name = info.bucket_name
                 bucket_path = info.bucket_path
-                file_uri = info.file_uri
-
-                if file_uri and file_uri.startswith("azure"):
-
-                    from anyscale.utils.imports.azure import (
-                        try_import_azure_storage_blob_BlobServiceClient,
-                    )
-
-                    # Smartopen needs transport_params to be passed in for Azure.
-                    blob_service_client = (
-                        try_import_azure_storage_blob_BlobServiceClient()
-                    )
 
-                    transport_params = {
-                        "client": blob_service_client.from_connection_string(info.url),
-                    }
-                    with smart_open.open(
-                        f"{file_uri}", "wb", transport_params=transport_params,
-                    ) as fout:
-                        fout.write(zip_file_bytes)
-                else:
-                    env_vars: Dict[str, str] = {
-                        "AWS_ENDPOINT_URL": info.url,
-                    }
-                    with set_env(**env_vars), smart_open.open(
-                        f"{bucket_name}/{bucket_path}", "wb",
-                    ) as fout:
-                        fout.write(zip_file_bytes)
+                env_vars: Dict[str, str] = {
+                    "AWS_ENDPOINT_URL": info.url,
+                }
+                with set_env(**env_vars), smart_open.open(
+                    f"{bucket_name}/{bucket_path}", "wb",
+                ) as fout:
+                    fout.write(zip_file_bytes)
 
             else:
                 # Default to HTTP PUT.
                 internal_logger.debug(f"Uploading file '{file_name}' to cloud storage.")
-                requests.put(info.url, data=zip_file_bytes).raise_for_status()
+                headers = (
+                    {"x-ms-blob-type": "BlockBlob"}
+                    if info.file_uri.startswith("azure")
+                    else None
+                )
+                requests.put(
+                    info.url, data=zip_file_bytes, headers=headers
+                ).raise_for_status()
 
         return info.file_uri
 
@@ -1845,3 +1859,119 @@ class AnyscaleClient(AnyscaleClientInterface):
             )
 
         return filepath
+
+    @handle_api_exceptions
+    def admin_batch_create_users(
+        self, admin_create_users: List[AdminCreateUser]
+    ) -> List[AdminCreatedUser]:
+        return self._internal_api_client.admin_batch_create_users_api_v2_users_admin_batch_create_post(
+            admin_create_users
+        ).results
+
+    @handle_api_exceptions
+    def create_organization_invitations(
+        self, emails: List[str]
+    ) -> Tuple[List[str], List[str]]:
+        results = self._internal_api_client.batch_create_invitations_api_v2_organization_invitations_batch_create_post(
+            [CreateOrganizationInvitation(email=email) for email in emails]
+        ).results
+
+        success_emails = []
+        error_messages = []
+
+        for idx, result in enumerate(results):
+            if result.data:
+                success_emails.append(emails[idx])
+            else:
+                error_messages.append(result.error.detail)
+
+        return success_emails, error_messages
+
+    @handle_api_exceptions
+    def list_organization_invitations(self) -> List[OrganizationInvitation]:
+        results = (
+            self._internal_api_client.list_invitations_api_v2_organization_invitations_get().results
+        )
+
+        return results
+
+    @handle_api_exceptions
+    def delete_organization_invitation(self, email: str) -> OrganizationInvitation:
+        invitation = self._internal_api_client.list_invitations_api_v2_organization_invitations_get(
+            email=email
+        ).results
+
+        if len(invitation) == 0:
+            raise ValueError(f"Invitation for email '{email}' not found.")
+        elif len(invitation) > 1:
+            raise ValueError(
+                f"Multiple invitations found for email '{email}'. Please contact Anyscale support."
+            )
+
+        invitation_id = invitation[0].id
+
+        return self._internal_api_client.invalidate_invitation_api_v2_organization_invitations_invitation_id_invalidate_post(
+            invitation_id
+        ).result
+
+    @handle_api_exceptions
+    def get_organization_collaborator(self, email: str) -> OrganizationCollaborator:
+        results = self._internal_api_client.list_organization_collaborators_api_v2_organization_collaborators_get(
+            email=email
+        ).results
+
+        if len(results) == 0:
+            raise ValueError(f"User with email '{email}' not found.")
+
+        if len(results) > 1:
+            raise ValueError(
+                f"Multiple users found for email '{email}'. Please contact Anyscale support."
+            )
+
+        return results[0]
+
+    @handle_api_exceptions
+    def create_resource_quota(
+        self, create_resource_quota: CreateResourceQuota
+    ) -> ResourceQuota:
+        return self._internal_api_client.create_resource_quota_api_v2_resource_quotas_post(
+            create_resource_quota
+        ).result
+
+    @handle_api_exceptions
+    def list_resource_quotas(
+        self,
+        name: Optional[str] = None,
+        cloud_id: Optional[str] = None,
+        creator_id: Optional[str] = None,
+        is_enabled: Optional[bool] = None,
+        max_items: int = 20,
+    ) -> List[ResourceQuota]:
+
+        query = ListResourceQuotasQuery(
+            name=TextQuery(equals=name) if name else None,
+            cloud_id=cloud_id,
+            creator_id=creator_id,
+            is_enabled=is_enabled,
+            paging=PageQuery(count=max_items),
+        )
+
+        resource_quotas = self._internal_api_client.search_resource_quotas_api_v2_resource_quotas_search_post(
+            query
+        ).results
+
+        return resource_quotas
+
+    @handle_api_exceptions
+    def delete_resource_quota(self, resource_quota_id: str) -> None:
+        self._internal_api_client.delete_resource_quota_api_v2_resource_quotas_resource_quota_id_delete(
+            resource_quota_id
+        )
+
+    @handle_api_exceptions
+    def set_resource_quota_status(
+        self, resource_quota_id: str, is_enabled: bool
+    ) -> None:
+        self._internal_api_client.set_resource_quota_status_api_v2_resource_quotas_resource_quota_id_status_patch(
+            resource_quota_id, ResourceQuotaStatus(is_enabled=is_enabled)
+        ).result

anyscale/_private/anyscale_client/common.py

@@ -4,15 +4,23 @@ from typing import Dict, Generator, List, Optional, Tuple
 
 from anyscale._private.models.image_uri import ImageURI
 from anyscale.client.openapi_client.models import (
+    AdminCreatedUser,
+    AdminCreateUser,
     Cloud,
     ComputeTemplateConfig,
+    CreateCloudCollaborator,
     CreateExperimentalWorkspace,
     CreateInternalProductionJob,
+    CreateResourceQuota,
+    CreateUserProjectCollaborator,
     DecoratedComputeTemplate,
     DeletedPlatformFineTunedModel,
     FineTunedModel,
     InternalProductionJob,
+    OrganizationCollaborator,
+    OrganizationInvitation,
     Project,
+    ResourceQuota,
     WorkspaceDataplaneProxiedArtifacts,
 )
 from anyscale.client.openapi_client.models.create_schedule import CreateSchedule
@@ -140,6 +148,13 @@ class AnyscaleClientInterface(ABC):
         """
         raise NotImplementedError
 
+    @abstractmethod
+    def add_cloud_collaborators(
+        self, cloud_id: str, collaborators: List[CreateCloudCollaborator]
+    ) -> None:
+        """Batch add collaborators to a cloud."""
+        raise NotImplementedError
+
     @abstractmethod
     def create_compute_config(
         self, config: ComputeTemplateConfig, *, name: Optional[str] = None
@@ -283,6 +298,13 @@ class AnyscaleClientInterface(ABC):
         """
         raise NotImplementedError
 
+    @abstractmethod
+    def add_project_collaborators(
+        self, project_id: str, collaborators: List[CreateUserProjectCollaborator]
+    ) -> None:
+        """Batch add collaborators to a project."""
+        raise NotImplementedError
+
     @abstractmethod
     def get_job(
         self,
@@ -591,3 +613,63 @@ class AnyscaleClientInterface(ABC):
     ) -> str:
         """Download the aggregated instance usage csv."""
         raise NotImplementedError
+
+    @abstractmethod
+    def admin_batch_create_users(
+        self, admin_create_users: List[AdminCreateUser]
+    ) -> List[AdminCreatedUser]:
+        """Batch create, as an admin, users without email verification."""
+        raise NotImplementedError
+
+    @abstractmethod
+    def create_organization_invitations(
+        self, emails: List[str]
+    ) -> Tuple[List[str], List[str]]:
+        """Create organization invitations."""
+        raise NotImplementedError
+
+    @abstractmethod
+    def list_organization_invitations(self) -> List[OrganizationInvitation]:
+        """List organization invitations."""
+        raise NotImplementedError
+
+    @abstractmethod
+    def delete_organization_invitation(self, email: str) -> OrganizationInvitation:
+        """Delete organization invitation."""
+        raise NotImplementedError
+
+    @abstractmethod
+    def get_organization_collaborator(self, email: str) -> OrganizationCollaborator:
+        """Get organization collaborator."""
+        raise NotImplementedError
+
+    @abstractmethod
+    def create_resource_quota(
+        self, create_resource_quota: CreateResourceQuota
+    ) -> ResourceQuota:
+        """Create a resource quota."""
+        raise NotImplementedError
+
+    @abstractmethod
+    def list_resource_quotas(
+        self,
+        name: Optional[str] = None,
+        cloud_id: Optional[str] = None,
+        creator_id: Optional[str] = None,
+        is_enabled: Optional[bool] = None,
+        max_items: int = 20,
+    ) -> List[ResourceQuota]:
+        """List resource quotas."""
+        raise NotImplementedError
+
+    @abstractmethod
+    def delete_resource_quota(self, resource_quota_id: str) -> None:
+        """Delete a resource quota."""
+        raise NotImplementedError
+
+    @abstractmethod
+    def set_resource_quota_status(
+        self, resource_quota_id: str, is_enabled: bool
+    ) -> None:
+        """Set the status of a resource quota."""
+        raise NotImplementedError