zenml-nightly 0.75.0.dev20250314__py3-none-any.whl → 0.75.0.dev20250315__py3-none-any.whl

zenml/integrations/gcp/vertex_custom_job_parameters.py

@@ -0,0 +1,50 @@
+#  Copyright (c) ZenML GmbH 2022. All Rights Reserved.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at:
+#
+#       https://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
+#  or implied. See the License for the specific language governing
+#  permissions and limitations under the License.
+"""Vertex custom job parameter model."""
+
+from typing import Optional
+
+from pydantic import BaseModel
+
+
+class VertexCustomJobParameters(BaseModel):
+    """Settings for the Vertex custom job parameters.
+
+    Attributes:
+        accelerator_type: Defines which accelerator (GPU, TPU) is used for the
+            job. Check out out this table to see which accelerator
+            type and count are compatible with your chosen machine type:
+            https://cloud.google.com/vertex-ai/docs/training/configure-compute#gpu-compatibility-table.
+        accelerator_count: Defines number of accelerators to be used for the
+            job. Check out out this table to see which accelerator
+            type and count are compatible with your chosen machine type:
+            https://cloud.google.com/vertex-ai/docs/training/configure-compute#gpu-compatibility-table.
+        machine_type: Machine type specified here
+            https://cloud.google.com/vertex-ai/docs/training/configure-compute#machine-types.
+        boot_disk_size_gb: Size of the boot disk in GB. (Default: 100)
+            https://cloud.google.com/vertex-ai/docs/training/configure-compute#boot_disk_options
+        boot_disk_type: Type of the boot disk. (Default: pd-ssd)
+            https://cloud.google.com/vertex-ai/docs/training/configure-compute#boot_disk_options
+        persistent_resource_id: The ID of the persistent resource to use for the job.
+            https://cloud.google.com/vertex-ai/docs/training/persistent-resource-overview
+        service_account: Specifies the service account to be used.
+    """
+
+    accelerator_type: Optional[str] = None
+    accelerator_count: int = 0
+    machine_type: str = "n1-standard-4"
+    boot_disk_size_gb: int = 100
+    boot_disk_type: str = "pd-ssd"
+    persistent_resource_id: Optional[str] = None
+    service_account: Optional[str] = None

zenml/login/credentials.py

@@ -21,7 +21,7 @@ from uuid import UUID
 from pydantic import BaseModel, ConfigDict
 
 from zenml.login.pro.constants import ZENML_PRO_API_URL, ZENML_PRO_URL
-from zenml.login.pro.tenant.models import TenantRead, TenantStatus
+from zenml.login.pro.workspace.models import WorkspaceRead, WorkspaceStatus
 from zenml.models import ServerModel
 from zenml.models.v2.misc.server_models import ServerDeploymentType
 from zenml.services.service_status import ServiceState
@@ -99,8 +99,8 @@ class ServerCredentials(BaseModel):
     # Pro server attributes
     organization_name: Optional[str] = None
     organization_id: Optional[UUID] = None
-    tenant_name: Optional[str] = None
-    tenant_id: Optional[UUID] = None
+    workspace_name: Optional[str] = None
+    workspace_id: Optional[UUID] = None
     pro_api_url: Optional[str] = None
     pro_dashboard_url: Optional[str] = None
 
@@ -128,7 +128,7 @@ class ServerCredentials(BaseModel):
             return ServerType.PRO_API
         if self.url == self.pro_api_url:
             return ServerType.PRO_API
-        if self.organization_id or self.tenant_id:
+        if self.organization_id or self.workspace_id:
             return ServerType.PRO
         if urlparse(self.url).hostname in [
             "localhost",
@@ -139,9 +139,9 @@ class ServerCredentials(BaseModel):
         return ServerType.REMOTE
 
     def update_server_info(
-        self, server_info: Union[ServerModel, TenantRead]
+        self, server_info: Union[ServerModel, WorkspaceRead]
     ) -> None:
-        """Update with server information received from the server itself or from a ZenML Pro tenant descriptor.
+        """Update with server information received from the server itself or from a ZenML Pro workspace descriptor.
 
         Args:
             server_info: The server information to update with.
@@ -152,20 +152,20 @@ class ServerCredentials(BaseModel):
             # All other attributes can change during the lifetime of the server
             self.deployment_type = server_info.deployment_type
             server_name = (
-                server_info.pro_tenant_name
-                or server_info.metadata.get("tenant_name")
+                server_info.pro_workspace_name
+                or server_info.metadata.get("workspace_name")
                 or server_info.name
             )
             if server_name:
                 self.server_name = server_name
             if server_info.pro_organization_id:
                 self.organization_id = server_info.pro_organization_id
-            if server_info.pro_tenant_id:
-                self.server_id = server_info.pro_tenant_id
+            if server_info.pro_workspace_id:
+                self.server_id = server_info.pro_workspace_id
             if server_info.pro_organization_name:
                 self.organization_name = server_info.pro_organization_name
-            if server_info.pro_tenant_name:
-                self.tenant_name = server_info.pro_tenant_name
+            if server_info.pro_workspace_name:
+                self.workspace_name = server_info.pro_workspace_name
             if server_info.pro_api_url:
                 self.pro_api_url = server_info.pro_api_url
             if server_info.pro_dashboard_url:
@@ -180,8 +180,8 @@ class ServerCredentials(BaseModel):
             self.server_name = server_info.name
             self.organization_name = server_info.organization_name
             self.organization_id = server_info.organization_id
-            self.tenant_name = server_info.name
-            self.tenant_id = server_info.id
+            self.workspace_name = server_info.name
+            self.workspace_id = server_info.id
             self.status = server_info.status
             self.version = server_info.version
 
@@ -192,7 +192,7 @@ class ServerCredentials(BaseModel):
         Returns:
             True if the server is available, False otherwise.
         """
-        if self.status not in [TenantStatus.AVAILABLE, ServiceState.ACTIVE]:
+        if self.status not in [WorkspaceStatus.AVAILABLE, ServiceState.ACTIVE]:
             return False
         if (
             self.api_key
@@ -270,20 +270,19 @@ class ServerCredentials(BaseModel):
         Returns:
             The URL to the ZenML dashboard for this server.
         """
-        if self.organization_id and self.server_id:
+        if self.pro_dashboard_url and self.workspace_name:
             return (
-                (self.pro_dashboard_url or ZENML_PRO_URL)
-                + f"/organizations/{str(self.organization_id)}/tenants/{str(self.server_id)}"
-            )
+                self.pro_dashboard_url or ZENML_PRO_URL
+            ) + f"/workspaces/{str(self.workspace_name)}"
 
         return self.url
 
     @property
     def dashboard_organization_url(self) -> str:
-        """Get the URL to the ZenML Pro dashboard for this tenant's organization.
+        """Get the URL to the ZenML Pro dashboard for this workspace's organization.
 
         Returns:
-            The URL to the ZenML Pro dashboard for this tenant's organization.
+            The URL to the ZenML Pro dashboard for this workspace's organization.
         """
         if self.organization_id:
             return (
@@ -293,19 +292,19 @@ class ServerCredentials(BaseModel):
 
     @property
     def dashboard_hyperlink(self) -> str:
-        """Get the hyperlink to the ZenML dashboard for this tenant.
+        """Get the hyperlink to the ZenML dashboard for this workspace.
 
         Returns:
-            The hyperlink to the ZenML dashboard for this tenant.
+            The hyperlink to the ZenML dashboard for this workspace.
         """
         return f"[link={self.dashboard_url}]{self.dashboard_url}[/link]"
 
     @property
     def api_hyperlink(self) -> str:
-        """Get the hyperlink to the ZenML OpenAPI dashboard for this tenant.
+        """Get the hyperlink to the ZenML OpenAPI dashboard for this workspace.
 
         Returns:
-            The hyperlink to the ZenML OpenAPI dashboard for this tenant.
+            The hyperlink to the ZenML OpenAPI dashboard for this workspace.
         """
         api_url = self.url + "/docs"
         return f"[link={api_url}]{self.url}[/link]"
@@ -360,10 +359,10 @@ class ServerCredentials(BaseModel):
 
     @property
     def organization_id_hyperlink(self) -> str:
-        """Get the hyperlink to the ZenML Pro dashboard for this tenant's organization using its ID.
+        """Get the hyperlink to the ZenML Pro dashboard for this workspace's organization using its ID.
 
         Returns:
-            The hyperlink to the ZenML Pro dashboard for this tenant's
+            The hyperlink to the ZenML Pro dashboard for this workspace's
             organization using its ID.
         """
         if self.organization_id is None:

zenml/login/credentials_store.py

@@ -25,7 +25,7 @@ from zenml.constants import (
 from zenml.io import fileio
 from zenml.logger import get_logger
 from zenml.login.credentials import APIToken, ServerCredentials, ServerType
-from zenml.login.pro.tenant.models import TenantRead
+from zenml.login.pro.workspace.models import WorkspaceRead
 from zenml.models import OAuthTokenResponse, ServerModel
 from zenml.utils import yaml_utils
 from zenml.utils.singleton import SingletonMetaClass
@@ -75,7 +75,7 @@ class CredentialsStore(metaclass=SingletonMetaClass):
 
     Alongside credentials, the Credentials Store is also used to store
     additional server information:
-        * ZenML Pro tenant information populated by the `zenml login` command
+        * ZenML Pro workspace information populated by the `zenml login` command
         * ZenML server information populated by the REST zen store by fetching
         the server's information endpoint after authenticating
 
@@ -341,7 +341,7 @@ class CredentialsStore(metaclass=SingletonMetaClass):
         self.clear_token(pro_api_url)
 
     def clear_all_pro_tokens(
-        self, pro_api_url: str
+        self, pro_api_url: Optional[str] = None
     ) -> List[ServerCredentials]:
         """Delete all tokens from the store for ZenML Pro servers connected to a given API server.
 
@@ -356,7 +356,7 @@ class CredentialsStore(metaclass=SingletonMetaClass):
             if (
                 server.type == ServerType.PRO
                 and server.pro_api_url
-                and server.pro_api_url == pro_api_url
+                and (pro_api_url is None or server.pro_api_url == pro_api_url)
             ):
                 if server.api_key:
                     continue
@@ -556,7 +556,7 @@ class CredentialsStore(metaclass=SingletonMetaClass):
     def update_server_info(
         self,
         server_url: str,
-        server_info: Union[ServerModel, TenantRead],
+        server_info: Union[ServerModel, WorkspaceRead],
     ) -> None:
         """Update the server information stored for a specific server URL.
 

zenml/login/pro/client.py

@@ -41,7 +41,7 @@ logger = get_logger(__name__)
 
 if TYPE_CHECKING:
     from zenml.login.pro.organization.client import OrganizationClient
-    from zenml.login.pro.tenant.client import TenantClient
+    from zenml.login.pro.workspace.client import WorkspaceClient
 
 # type alias for possible json payloads (the Anys are recursive Json instances)
 Json = Union[Dict[str, Any], List[Any], str, int, float, bool, None]
@@ -56,7 +56,7 @@ class ZenMLProClient(metaclass=SingletonMetaClass):
     _url: str
     _api_token: APIToken
     _session: Optional[requests.Session] = None
-    _tenant: Optional["TenantClient"] = None
+    _workspace: Optional["WorkspaceClient"] = None
     _organization: Optional["OrganizationClient"] = None
 
     def __init__(self, url: str, api_token: Optional[APIToken] = None) -> None:
@@ -89,17 +89,17 @@ class ZenMLProClient(metaclass=SingletonMetaClass):
         self._api_token = api_token
 
     @property
-    def tenant(self) -> "TenantClient":
-        """Get the tenant client.
+    def workspace(self) -> "WorkspaceClient":
+        """Get the workspace client.
 
         Returns:
-            The tenant client.
+            The workspace client.
         """
-        if self._tenant is None:
-            from zenml.login.pro.tenant.client import TenantClient
+        if self._workspace is None:
+            from zenml.login.pro.workspace.client import WorkspaceClient
 
-            self._tenant = TenantClient(client=self)
-        return self._tenant
+            self._workspace = WorkspaceClient(client=self)
+        return self._workspace
 
     @property
     def organization(self) -> "OrganizationClient":

zenml/login/pro/utils.py

@@ -18,7 +18,7 @@ from zenml.login.credentials import ServerType
 from zenml.login.credentials_store import get_credentials_store
 from zenml.login.pro.client import ZenMLProClient
 from zenml.login.pro.constants import ZENML_PRO_API_URL
-from zenml.login.pro.tenant.models import TenantStatus
+from zenml.login.pro.workspace.models import WorkspaceStatus
 
 logger = get_logger(__name__)
 
@@ -44,13 +44,13 @@ def get_troubleshooting_instructions(url: str) -> str:
         client = ZenMLProClient(pro_api_url)
 
         try:
-            servers = client.tenant.list(url=url, member_only=False)
+            servers = client.workspace.list(url=url, member_only=False)
         except Exception as e:
-            logger.debug(f"Failed to list tenants: {e}")
+            logger.debug(f"Failed to list workspaces: {e}")
         else:
             if servers:
                 server = servers[0]
-                if server.status == TenantStatus.AVAILABLE:
+                if server.status == WorkspaceStatus.AVAILABLE:
                     return (
                         f"The '{server.name}' ZenML Pro server that the client "
                         "is connected to is currently running but you may not "
@@ -58,9 +58,9 @@ def get_troubleshooting_instructions(url: str) -> str:
                         "contact your ZenML Pro administrator for more "
                         "information or try to manage the server members "
                         "yourself if you have the necessary permissions by "
-                        f"visiting the ZenML Pro tenant page at {server.dashboard_url}."
+                        f"visiting the ZenML Pro workspace page at {server.dashboard_url}."
                     )
-                if server.status == TenantStatus.DEACTIVATED:
+                if server.status == WorkspaceStatus.DEACTIVATED:
                     return (
                         f"The '{server.name}' ZenML Pro server that the client "
                         "is connected to has been deactivated. "
@@ -69,14 +69,14 @@ def get_troubleshooting_instructions(url: str) -> str:
                         "you have the necessary permissions by visiting the "
                         f"ZenML Pro Organization page at {server.dashboard_organization_url}."
                     )
-                if server.status == TenantStatus.PENDING:
+                if server.status == WorkspaceStatus.PENDING:
                     return (
                         f"The '{server.name}' ZenML Pro server that the client "
                         "is connected to is currently undergoing maintenance "
                         "(e.g. being deployed, upgraded or re-activated). "
                         "Please try again later or contact your ZenML Pro "
                         "administrator for more information. You can also "
-                        f"visit the ZenML Pro tenant page at {server.dashboard_url}."
+                        f"visit the ZenML Pro workspace page at {server.dashboard_url}."
                     )
                 return (
                     f"The '{server.name}' ZenML Pro server that the client "

zenml/login/pro/{tenant → workspace}/__init__.py

@@ -11,4 +11,4 @@
 #  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
 #  or implied. See the License for the specific language governing
 #  permissions and limitations under the License.
-"""ZenML Pro tenant client."""
+"""ZenML Pro workspace client."""

zenml/login/pro/{tenant → workspace}/client.py

@@ -11,80 +11,80 @@
 #  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
 #  or implied. See the License for the specific language governing
 #  permissions and limitations under the License.
-"""ZenML Pro tenant client."""
+"""ZenML Pro workspace client."""
 
 from typing import List, Optional
 from uuid import UUID
 
 from zenml.logger import get_logger
 from zenml.login.pro.client import ZenMLProClient
-from zenml.login.pro.tenant.models import TenantRead, TenantStatus
+from zenml.login.pro.workspace.models import WorkspaceRead, WorkspaceStatus
 
 logger = get_logger(__name__)
 
-TENANTS_ROUTE = "/tenants"
+WORKSPACES_ROUTE = "/workspaces"
 
 
-class TenantClient:
-    """Tenant management client."""
+class WorkspaceClient:
+    """Workspace management client."""
 
     def __init__(
         self,
         client: ZenMLProClient,
     ):
-        """Initialize the tenant client.
+        """Initialize the workspace client.
 
         Args:
             client: ZenML Pro client.
         """
         self.client = client
 
-    def get(self, id: UUID) -> TenantRead:
-        """Get a tenant by id.
+    def get(self, id: UUID) -> WorkspaceRead:
+        """Get a workspace by id.
 
         Args:
-            id: Id. of the tenant to retrieve.
+            id: Id. of the workspace to retrieve.
 
         Returns:
-            A tenant.
+            A workspace.
         """
         return self.client._get_resource(
             resource_id=id,
-            route=TENANTS_ROUTE,
-            response_model=TenantRead,
+            route=WORKSPACES_ROUTE,
+            response_model=WorkspaceRead,
         )
 
     def list(
         self,
         offset: int = 0,
         limit: int = 20,
-        tenant_name: Optional[str] = None,
+        workspace_name: Optional[str] = None,
         url: Optional[str] = None,
         organization_id: Optional[UUID] = None,
-        status: Optional[TenantStatus] = None,
+        status: Optional[WorkspaceStatus] = None,
         member_only: bool = False,
-    ) -> List[TenantRead]:
-        """List tenants.
+    ) -> List[WorkspaceRead]:
+        """List workspaces.
 
         Args:
             offset: Offset to use for filtering.
             limit: Limit used for filtering.
-            tenant_name: Tenant name to filter by.
-            url: Tenant service URL to filter by.
+            workspace_name: Workspace name to filter by.
+            url: Workspace service URL to filter by.
             organization_id: Organization ID to filter by.
-            status: Filter for only tenants with this status.
-            member_only: If True, only list tenants where the user is a member
-                (i.e. users that can connect to the tenant).
+            status: Filter for only workspaces with this status.
+            member_only: If True, only list workspaces where the user is a member
+                (i.e. users that can connect to the workspace).
 
         Returns:
-            List of tenants.
+            List of workspaces.
         """
         return self.client._list_resources(
-            route=TENANTS_ROUTE,
-            response_model=TenantRead,
+            route=WORKSPACES_ROUTE,
+            response_model=WorkspaceRead,
             offset=offset,
             limit=limit,
-            tenant_name=tenant_name,
+            workspace_name=workspace_name,
             url=url,
             organization_id=organization_id,
             status=status,

zenml/login/pro/{tenant → workspace}/models.py

@@ -11,7 +11,7 @@
 #  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
 #  or implied. See the License for the specific language governing
 #  permissions and limitations under the License.
-"""ZenML Pro tenant models."""
+"""ZenML Pro workspace models."""
 
 from typing import Optional
 from uuid import UUID
@@ -24,32 +24,32 @@ from zenml.login.pro.organization.models import OrganizationRead
 from zenml.utils.enum_utils import StrEnum
 
 
-class TenantStatus(StrEnum):
-    """Enum that represents the desired state or status of a tenant.
+class WorkspaceStatus(StrEnum):
+    """Enum that represents the desired state or status of a workspace.
 
     These values can be used in two places:
 
-    * in the `desired_state` field of a tenant object, to indicate the desired
-    state of the tenant (with the exception of `PENDING` and `FAILED` which
+    * in the `desired_state` field of a workspace object, to indicate the desired
+    state of the workspace (with the exception of `PENDING` and `FAILED` which
     are not valid values for `desired_state`)
-    * in the `status` field of a tenant object, to indicate the current state
-    of the tenant
+    * in the `status` field of a workspace object, to indicate the current state
+    of the workspace
     """
 
-    # Tenant hasn't been deployed yet (i.e. newly created) or has been fully
+    # Workspace hasn't been deployed yet (i.e. newly created) or has been fully
     # deleted by the infrastructure provider
     NOT_INITIALIZED = "not_initialized"
-    # Tenant is being processed by the infrastructure provider (is being
+    # Workspace is being processed by the infrastructure provider (is being
     # deployed, updated, deactivated, re-activated or deleted/cleaned up).
     PENDING = "pending"
-    # Tenant is up and running
+    # Workspace is up and running
     AVAILABLE = "available"
-    # Tenant is in a failure state (i.e. deployment, update or deletion failed)
+    # Workspace is in a failure state (i.e. deployment, update or deletion failed)
     FAILED = "failed"
-    # Tenant is deactivated
+    # Workspace is deactivated
     DEACTIVATED = "deactivated"
-    # Tenant resources have been deleted by the infrastructure provider but
-    # the tenant object still exists in the database
+    # Workspace resources have been deleted by the infrastructure provider but
+    # the workspace object still exists in the database
     DELETED = "deleted"
 
 
@@ -87,24 +87,26 @@ class ZenMLServiceRead(BaseRestAPIModel):
     )
 
 
-class TenantRead(BaseRestAPIModel):
-    """Pydantic Model for viewing a Tenant."""
+class WorkspaceRead(BaseRestAPIModel):
+    """Pydantic Model for viewing a Workspace."""
 
     id: UUID
 
     name: str
     description: Optional[str] = Field(
-        default=None, description="The description of the tenant."
+        default=None, description="The description of the workspace."
     )
 
     organization: OrganizationRead
 
-    desired_state: str = Field(description="The desired state of the tenant.")
+    desired_state: str = Field(
+        description="The desired state of the workspace."
+    )
     state_reason: str = Field(
-        description="The reason for the current tenant state.",
+        description="The reason for the current workspace state.",
     )
     status: str = Field(
-        description="The current operational state of the tenant."
+        description="The current operational state of the workspace."
     )
     zenml_service: ZenMLServiceRead = Field(description="The ZenML service.")
 
@@ -156,21 +158,18 @@ class TenantRead(BaseRestAPIModel):
 
     @property
     def dashboard_url(self) -> str:
-        """Get the URL to the ZenML Pro dashboard for this tenant.
+        """Get the URL to the ZenML Pro dashboard for this workspace.
 
         Returns:
-            The URL to the ZenML Pro dashboard for this tenant.
+            The URL to the ZenML Pro dashboard for this workspace.
         """
-        return (
-            ZENML_PRO_URL
-            + f"/organizations/{str(self.organization_id)}/tenants/{str(self.id)}"
-        )
+        return ZENML_PRO_URL + f"/workspaces/{str(self.name)}"
 
     @property
     def dashboard_organization_url(self) -> str:
-        """Get the URL to the ZenML Pro dashboard for this tenant's organization.
+        """Get the URL to the ZenML Pro dashboard for this workspace's organization.
 
         Returns:
-            The URL to the ZenML Pro dashboard for this tenant's organization.
+            The URL to the ZenML Pro dashboard for this workspace's organization.
         """
         return ZENML_PRO_URL + f"/organizations/{str(self.organization_id)}"

zenml/models/v2/core/model.py

@@ -13,7 +13,15 @@
 #  permissions and limitations under the License.
 """Models representing models."""
 
-from typing import TYPE_CHECKING, Any, ClassVar, List, Optional, Type, TypeVar
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    ClassVar,
+    List,
+    Optional,
+    Type,
+    TypeVar,
+)
 from uuid import UUID
 
 from pydantic import Field