zenml-nightly 0.66.0.dev20240927__py3-none-any.whl → 0.68.1.dev20241102__py3-none-any.whl

zenml/login/pro/tenant/client.py

@@ -0,0 +1,92 @@
+#  Copyright (c) ZenML GmbH 2024. 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.
+"""ZenML Pro tenant 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
+
+logger = get_logger(__name__)
+
+TENANTS_ROUTE = "/tenants"
+
+
+class TenantClient:
+    """Tenant management client."""
+
+    def __init__(
+        self,
+        client: ZenMLProClient,
+    ):
+        """Initialize the tenant client.
+
+        Args:
+            client: ZenML Pro client.
+        """
+        self.client = client
+
+    def get(self, id: UUID) -> TenantRead:
+        """Get a tenant by id.
+
+        Args:
+            id: Id. of the tenant to retrieve.
+
+        Returns:
+            A tenant.
+        """
+        return self.client._get_resource(
+            resource_id=id,
+            route=TENANTS_ROUTE,
+            response_model=TenantRead,
+        )
+
+    def list(
+        self,
+        offset: int = 0,
+        limit: int = 20,
+        tenant_name: Optional[str] = None,
+        url: Optional[str] = None,
+        organization_id: Optional[UUID] = None,
+        status: Optional[TenantStatus] = None,
+        member_only: bool = False,
+    ) -> List[TenantRead]:
+        """List tenants.
+
+        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.
+            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).
+
+        Returns:
+            List of tenants.
+        """
+        return self.client._list_resources(
+            route=TENANTS_ROUTE,
+            response_model=TenantRead,
+            offset=offset,
+            limit=limit,
+            tenant_name=tenant_name,
+            url=url,
+            organization_id=organization_id,
+            status=status,
+            member_only=member_only,
+        )

zenml/login/pro/tenant/models.py

@@ -0,0 +1,174 @@
+#  Copyright (c) ZenML GmbH 2024. 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.
+"""ZenML Pro tenant models."""
+
+from typing import Optional
+from uuid import UUID
+
+from pydantic import Field
+
+from zenml.login.pro.constants import ZENML_PRO_URL
+from zenml.login.pro.models import BaseRestAPIModel
+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.
+
+    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
+    are not valid values for `desired_state`)
+    * in the `status` field of a tenant object, to indicate the current state
+    of the tenant
+    """
+
+    # Tenant 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
+    # deployed, updated, deactivated, re-activated or deleted/cleaned up).
+    PENDING = "pending"
+    # Tenant is up and running
+    AVAILABLE = "available"
+    # Tenant is in a failure state (i.e. deployment, update or deletion failed)
+    FAILED = "failed"
+    # Tenant is deactivated
+    DEACTIVATED = "deactivated"
+    # Tenant resources have been deleted by the infrastructure provider but
+    # the tenant object still exists in the database
+    DELETED = "deleted"
+
+
+class ZenMLServiceConfiguration(BaseRestAPIModel):
+    """ZenML service configuration."""
+
+    version: str = Field(
+        description="The ZenML version.",
+    )
+
+
+class ZenMLServiceStatus(BaseRestAPIModel):
+    """ZenML service status."""
+
+    server_url: str = Field(
+        description="The ZenML server URL.",
+    )
+    version: Optional[str] = Field(
+        default=None,
+        description="The ZenML server version.",
+    )
+
+
+class ZenMLServiceRead(BaseRestAPIModel):
+    """Pydantic Model for viewing a ZenML service."""
+
+    configuration: ZenMLServiceConfiguration = Field(
+        description="The service configuration."
+    )
+
+    status: Optional[ZenMLServiceStatus] = Field(
+        default=None,
+        description="Information about the service status. Only set if the "
+        "service is deployed and active.",
+    )
+
+
+class TenantRead(BaseRestAPIModel):
+    """Pydantic Model for viewing a Tenant."""
+
+    id: UUID
+
+    name: str
+    description: Optional[str] = Field(
+        default=None, description="The description of the tenant."
+    )
+
+    organization: OrganizationRead
+
+    desired_state: str = Field(description="The desired state of the tenant.")
+    state_reason: str = Field(
+        description="The reason for the current tenant state.",
+    )
+    status: str = Field(
+        description="The current operational state of the tenant."
+    )
+    zenml_service: ZenMLServiceRead = Field(description="The ZenML service.")
+
+    @property
+    def organization_id(self) -> UUID:
+        """Get the organization id.
+
+        Returns:
+            The organization id.
+        """
+        return self.organization.id
+
+    @property
+    def organization_name(self) -> str:
+        """Get the organization name.
+
+        Returns:
+            The organization name.
+        """
+        return self.organization.name
+
+    @property
+    def version(self) -> Optional[str]:
+        """Get the ZenML service version.
+
+        Returns:
+            The ZenML service version.
+        """
+        version = self.zenml_service.configuration.version
+        if self.zenml_service.status and self.zenml_service.status.version:
+            version = self.zenml_service.status.version
+
+        return version
+
+    @property
+    def url(self) -> Optional[str]:
+        """Get the ZenML server URL.
+
+        Returns:
+            The ZenML server URL, if available.
+        """
+        return (
+            self.zenml_service.status.server_url
+            if self.zenml_service.status
+            else None
+        )
+
+    @property
+    def dashboard_url(self) -> str:
+        """Get the URL to the ZenML Pro dashboard for this tenant.
+
+        Returns:
+            The URL to the ZenML Pro dashboard for this tenant.
+        """
+        return (
+            ZENML_PRO_URL
+            + f"/organizations/{str(self.organization_id)}/tenants/{str(self.id)}"
+        )
+
+    @property
+    def dashboard_organization_url(self) -> str:
+        """Get the URL to the ZenML Pro dashboard for this tenant's organization.
+
+        Returns:
+            The URL to the ZenML Pro dashboard for this tenant's organization.
+        """
+        return ZENML_PRO_URL + f"/organizations/{str(self.organization_id)}"

zenml/login/pro/utils.py

@@ -0,0 +1,121 @@
+#  Copyright (c) ZenML GmbH 2024. 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.
+"""ZenML Pro login utils."""
+
+import re
+
+from zenml.logger import get_logger
+from zenml.login.credentials_store import get_credentials_store
+from zenml.login.pro.client import ZenMLProClient
+from zenml.login.pro.constants import ZENML_PRO_SERVER_SUBDOMAIN
+from zenml.login.pro.tenant.models import TenantStatus
+
+logger = get_logger(__name__)
+
+
+def is_zenml_pro_server_url(url: str) -> bool:
+    """Check if a given URL is a ZenML Pro server.
+
+    Args:
+        url: URL to check
+
+    Returns:
+        True if the URL is a ZenML Pro tenant, False otherwise
+    """
+    domain_regex = ZENML_PRO_SERVER_SUBDOMAIN.replace(".", r"\.")
+    return bool(
+        re.match(
+            r"^(https://)?[a-zA-Z0-9-\.]+\.{domain}/?$".format(
+                domain=domain_regex
+            ),
+            url,
+        )
+    )
+
+
+def get_troubleshooting_instructions(url: str) -> str:
+    """Get troubleshooting instructions for a given ZenML Pro server URL.
+
+    Args:
+        url: ZenML Pro server URL
+
+    Returns:
+        Troubleshooting instructions
+    """
+    credentials_store = get_credentials_store()
+    if credentials_store.has_valid_pro_authentication():
+        client = ZenMLProClient()
+
+        try:
+            servers = client.tenant.list(url=url, member_only=False)
+        except Exception as e:
+            logger.debug(f"Failed to list tenants: {e}")
+        else:
+            if servers:
+                server = servers[0]
+                if server.status == TenantStatus.AVAILABLE:
+                    return (
+                        f"The '{server.name}' ZenML Pro server that the client "
+                        "is connected to is currently running but you may not "
+                        "have the necessary permissions to access it. Please "
+                        "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}."
+                    )
+                if server.status == TenantStatus.DEACTIVATED:
+                    return (
+                        f"The '{server.name}' ZenML Pro server that the client "
+                        "is connected to has been deactivated. "
+                        "Please contact your ZenML Pro administrator for more "
+                        "information or to reactivate the server yourself if "
+                        "you have the necessary permissions by visiting the "
+                        f"ZenML Pro Organization page at {server.dashboard_organization_url}."
+                    )
+                if server.status == TenantStatus.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}."
+                    )
+                return (
+                    f"The '{server.name}' ZenML Pro server that the client "
+                    "is connected to is currently in a failed "
+                    "state. Please contact your ZenML Pro administrator for "
+                    "more information or try to re-deploy the server "
+                    "yourself if you have the necessary permissions by "
+                    "visiting the ZenML Pro Organization page at "
+                    f"{server.dashboard_organization_url}."
+                )
+
+            return (
+                f"The ZenML Pro server at URL '{url}' that the client is "
+                "connected to does not exist or you may not have access to it. "
+                "Please check the URL and your permissions and try again or "
+                "connect your client to a different server by running `zenml "
+                "login` or by using a service account API key."
+            )
+
+    return (
+        f"The ZenML Pro server at URL '{url}' that the client is connected to "
+        "does not exist, is not running, or you do not have permissions to "
+        "connect to it. Please check the URL and your permissions "
+        "and try again. The ZenML Pro server might have been deactivated or is "
+        "currently pending maintenance. Please contact your ZenML Pro "
+        "administrator for more information or try to manage the server "
+        "state by visiting the ZenML Pro dashboard."
+    )

zenml/{cli → login}/web_login.py

@@ -1,4 +1,4 @@
-#  Copyright (c) ZenML GmbH 2023. All Rights Reserved.
+#  Copyright (c) ZenML GmbH 2024. 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.
@@ -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.
-"""Web login CLI support."""
+"""ZenML OAuth2 device authorization grant client support."""
 
 import platform
 import time
@@ -24,19 +24,24 @@ from zenml import __version__
 from zenml.config.global_config import GlobalConfiguration
 from zenml.constants import (
     API,
+    AUTH,
     DEFAULT_HTTP_TIMEOUT,
     DEVICE_AUTHORIZATION,
     LOGIN,
     VERSION_1,
-    ZENML_PRO_CONNECTION_ISSUES_SUSPENDED_PAUSED_TENANT_HINT,
 )
 from zenml.exceptions import AuthorizationException, OAuthError
 from zenml.logger import get_logger
+from zenml.login.credentials import APIToken
+from zenml.login.pro.constants import ZENML_PRO_API_URL
+from zenml.login.pro.utils import is_zenml_pro_server_url
 
 logger = get_logger(__name__)
 
 
-def web_login(url: str, verify_ssl: Union[str, bool]) -> str:
+def web_login(
+    url: Optional[str] = None, verify_ssl: Optional[Union[str, bool]] = None
+) -> APIToken:
     """Implements the OAuth2 Device Authorization Grant flow.
 
     This function implements the client side of the OAuth2 Device Authorization
@@ -48,19 +53,23 @@ def web_login(url: str, verify_ssl: Union[str, bool]) -> str:
     * additional information is added to the user agent header to be used by
     users to identify the ZenML client
 
+    Upon completion of the flow, the access token is saved in the credentials store.
+
     Args:
-        url: The URL of the OAuth2 server.
+        url: The URL of the OAuth2 server. If not provided, the ZenML Pro API
+            server is used by default.
         verify_ssl: Whether to verify the SSL certificate of the OAuth2 server.
             If a string is passed, it is interpreted as the path to a CA bundle
             file.
 
     Returns:
-        The access token returned by the OAuth2 server.
+        The response returned by the OAuth2 server.
 
     Raises:
         AuthorizationException: If an error occurred during the authorization
             process.
     """
+    from zenml.login.credentials_store import get_credentials_store
     from zenml.models import (
         OAuthDeviceAuthorizationRequest,
         OAuthDeviceAuthorizationResponse,
@@ -69,9 +78,8 @@ def web_login(url: str, verify_ssl: Union[str, bool]) -> str:
         OAuthTokenResponse,
     )
 
-    auth_request = OAuthDeviceAuthorizationRequest(
-        client_id=GlobalConfiguration().user_id
-    )
+    credentials_store = get_credentials_store()
+
     # Make a request to the OAuth2 server to get the device code and user code.
     # The client ID used for the request is the unique ID of the ZenML client.
     response: Optional[requests.Response] = None
@@ -91,23 +99,48 @@ def web_login(url: str, verify_ssl: Union[str, bool]) -> str:
         os=platform.system(),
     )
 
-    # Get rid of any trailing slashes to prevent issues when having double
-    # slashes in the URL
-    url = url.rstrip("/")
-    zenml_pro_extra = ""
-    if ".zenml.io" in url:
-        zenml_pro_extra = (
-            ZENML_PRO_CONNECTION_ISSUES_SUSPENDED_PAUSED_TENANT_HINT
-        )
+    zenml_pro = False
+    if not url:
+        # If no URL is provided, we use the ZenML Pro API server by default
+        zenml_pro = True
+        url = base_url = ZENML_PRO_API_URL
+    else:
+        # Get rid of any trailing slashes to prevent issues when having double
+        # slashes in the URL
+        url = url.rstrip("/")
+        if is_zenml_pro_server_url(url):
+            # This is a ZenML Pro server. The device authentication is done
+            # through the ZenML Pro API.
+            zenml_pro = True
+            base_url = ZENML_PRO_API_URL
+        else:
+            base_url = url
+
+    auth_request = OAuthDeviceAuthorizationRequest(
+        client_id=GlobalConfiguration().user_id
+    )
+
+    # If an existing token is found in the credentials store, we reuse its
+    # device ID to avoid creating a new device ID for the same device.
+    existing_token = credentials_store.get_token(url)
+    if existing_token and existing_token.device_id:
+        auth_request.device_id = existing_token.device_id
+
+    if zenml_pro:
+        auth_url = base_url + AUTH + DEVICE_AUTHORIZATION
+        login_url = base_url + AUTH + LOGIN
+    else:
+        auth_url = base_url + API + VERSION_1 + DEVICE_AUTHORIZATION
+        login_url = base_url + API + VERSION_1 + LOGIN
+
     try:
-        auth_url = url + API + VERSION_1 + DEVICE_AUTHORIZATION
         response = requests.post(
             auth_url,
             headers={
                 "Content-Type": "application/x-www-form-urlencoded",
                 "User-Agent": user_agent_header.encode(),
             },
-            data=auth_request.model_dump(),
+            data=auth_request.model_dump(exclude_none=True),
             verify=verify_ssl,
             timeout=DEFAULT_HTTP_TIMEOUT,
         )
@@ -116,8 +149,7 @@ def web_login(url: str, verify_ssl: Union[str, bool]) -> str:
         else:
             logger.info(f"Error: {response.status_code} {response.text}")
             raise AuthorizationException(
-                "Could not connect to API server. Please check the URL."
-                + zenml_pro_extra
+                f"Could not connect to {base_url}. Please check the URL."
             )
     except (requests.exceptions.JSONDecodeError, ValueError, TypeError):
         logger.exception("Bad response received from API server.")
@@ -127,8 +159,7 @@ def web_login(url: str, verify_ssl: Union[str, bool]) -> str:
     except requests.exceptions.RequestException:
         logger.exception("Could not connect to API server.")
         raise AuthorizationException(
-            "Could not connect to API server. Please check the URL."
-            + zenml_pro_extra
+            f"Could not connect to {base_url}. Please check the URL."
         )
 
     # Open the verification URL in the user's browser
@@ -139,7 +170,7 @@ def web_login(url: str, verify_ssl: Union[str, bool]) -> str:
     if verification_uri.startswith("/"):
         # If the verification URI is a relative path, we need to add the base
         # URL to it
-        verification_uri = url + verification_uri
+        verification_uri = base_url + verification_uri
     webbrowser.open(verification_uri)
     logger.info(
         f"If your browser did not open automatically, please open the "
@@ -156,7 +187,6 @@ def web_login(url: str, verify_ssl: Union[str, bool]) -> str:
     interval = auth_response.interval
     token_response: OAuthTokenResponse
     while True:
-        login_url = url + API + VERSION_1 + LOGIN
         response = requests.post(
             login_url,
             headers={"Content-Type": "application/x-www-form-urlencoded"},
@@ -167,8 +197,11 @@ def web_login(url: str, verify_ssl: Union[str, bool]) -> str:
         if response.status_code == 200:
             # The user has authorized the device, so we can extract the access token
             token_response = OAuthTokenResponse(**response.json())
-            logger.info("Successfully logged in.")
-            return token_response.access_token
+            if zenml_pro:
+                logger.info("Successfully logged in to ZenML Pro.")
+            else:
+                logger.info(f"Successfully logged in to {url}.")
+            break
         elif response.status_code == 400:
             try:
                 error_response = OAuthError(**response.json())
@@ -178,7 +211,7 @@ def web_login(url: str, verify_ssl: Union[str, bool]) -> str:
                 TypeError,
             ):
                 raise AuthorizationException(
-                    f"Error received from API server: {response.text}"
+                    f"Error received from {base_url}: {response.text}"
                 )
 
             if error_response.error == "authorization_pending":
@@ -205,3 +238,6 @@ def web_login(url: str, verify_ssl: Union[str, bool]) -> str:
             raise AuthorizationException(
                 f"Error: {response.status_code} {response.json()['error']}"
             )
+
+    # Save the token in the credentials store
+    return credentials_store.set_token(url, token_response)

zenml/materializers/__init__.py

@@ -31,6 +31,7 @@ from zenml.materializers.structured_string_materializer import (
 )
 from zenml.materializers.pydantic_materializer import PydanticMaterializer
 from zenml.materializers.service_materializer import ServiceMaterializer
+from zenml.materializers.uuid_materializer import UUIDMaterializer
 
 __all__ = [
     "BuiltInContainerMaterializer",
@@ -40,4 +41,5 @@ __all__ = [
     "StructuredStringMaterializer",
     "PydanticMaterializer",
     "ServiceMaterializer",
+    "UUIDMaterializer",
 ]

zenml/materializers/base_materializer.py

@@ -229,38 +229,72 @@ class BaseMaterializer(metaclass=BaseMaterializerMeta):
     # ================
     # Internal Methods
     # ================
+    def validate_save_type_compatibility(self, data_type: Type[Any]) -> None:
+        """Checks whether the materializer can save the given type.
 
-    def validate_type_compatibility(self, data_type: Type[Any]) -> None:
-        """Checks whether the materializer can read/write the given type.
+        Args:
+            data_type: The type to check.
+
+        Raises:
+            TypeError: If the materializer cannot save the given type.
+        """
+        if not self.can_save_type(data_type):
+            raise TypeError(
+                f"Unable to save type {data_type}. {self.__class__.__name__} "
+                f"can only save artifacts of the following types: "
+                f"{self.ASSOCIATED_TYPES}."
+            )
+
+    def validate_load_type_compatibility(self, data_type: Type[Any]) -> None:
+        """Checks whether the materializer can load the given type.
 
         Args:
             data_type: The type to check.
 
         Raises:
-            TypeError: If the materializer cannot read/write the given type.
+            TypeError: If the materializer cannot load the given type.
         """
-        if not self.can_handle_type(data_type):
+        if not self.can_load_type(data_type):
             raise TypeError(
-                f"Unable to handle type {data_type}. {self.__class__.__name__} "
-                f"can only read/write artifacts of the following types: "
+                f"Unable to load type {data_type}. {self.__class__.__name__} "
+                f"can only load artifacts of the following types: "
                 f"{self.ASSOCIATED_TYPES}."
             )
 
     @classmethod
-    def can_handle_type(cls, data_type: Type[Any]) -> bool:
-        """Whether the materializer can read/write a certain type.
+    def can_save_type(cls, data_type: Type[Any]) -> bool:
+        """Whether the materializer can save a certain type.
 
         Args:
             data_type: The type to check.
 
         Returns:
-            Whether the materializer can read/write the given type.
+            Whether the materializer can save the given type.
         """
         return any(
             issubclass(data_type, associated_type)
             for associated_type in cls.ASSOCIATED_TYPES
         )
 
+    @classmethod
+    def can_load_type(cls, data_type: Type[Any]) -> bool:
+        """Whether the materializer can load an artifact as the given type.
+
+        Args:
+            data_type: The type to check.
+
+        Returns:
+            Whether the materializer can load an artifact as the given type.
+        """
+        return any(
+            issubclass(associated_type, data_type)
+            # This next condition is not always correct, but better to have a
+            # false positive here instead of failing for cases where it would
+            # have worked.
+            or issubclass(data_type, associated_type)
+            for associated_type in cls.ASSOCIATED_TYPES
+        )
+
     def extract_full_metadata(self, data: Any) -> Dict[str, "MetadataType"]:
         """Extract both base and custom metadata from the given data.