zenml-nightly 0.55.0.dev20240124__py3-none-any.whl → 0.72.0.dev20250116__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,176 @@
+#  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: Optional[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 = None
+        if self.zenml_service.configuration:
+            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,107 @@
+#  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."""
+
+from zenml.logger import get_logger
+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
+
+logger = get_logger(__name__)
+
+
+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()
+
+    credentials = credentials_store.get_credentials(url)
+    if credentials and credentials.type == ServerType.PRO:
+        pro_api_url = credentials.pro_api_url or ZENML_PRO_API_URL
+
+    if pro_api_url and credentials_store.has_valid_pro_authentication(
+        pro_api_url
+    ):
+        client = ZenMLProClient(pro_api_url)
+
+        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/login/server_info.py

@@ -0,0 +1,52 @@
+#  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 server information retrieval."""
+
+from typing import Optional
+
+from zenml.logger import get_logger
+from zenml.models import ServerModel
+from zenml.zen_stores.rest_zen_store import (
+    RestZenStore,
+    RestZenStoreConfiguration,
+)
+
+logger = get_logger(__name__)
+
+
+def get_server_info(url: str) -> Optional[ServerModel]:
+    """Retrieve server information from a remote ZenML server.
+
+    Args:
+        url: The URL of the ZenML server.
+
+    Returns:
+        The server information or None if the server info could not be fetched.
+    """
+    # Here we try to leverage the existing RestZenStore support to fetch the
+    # server info and only the server info, which doesn't actually need
+    # any authentication.
+    try:
+        store = RestZenStore(
+            config=RestZenStoreConfiguration(
+                url=url,
+            )
+        )
+        return store.server_info
+    except Exception as e:
+        logger.warning(
+            f"Failed to fetch server info from the server running at {url}: {e}"
+        )
+
+    return None

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,8 @@
 #  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
 import webbrowser
@@ -23,6 +24,7 @@ from zenml import __version__
 from zenml.config.global_config import GlobalConfiguration
 from zenml.constants import (
     API,
+    AUTH,
     DEFAULT_HTTP_TIMEOUT,
     DEVICE_AUTHORIZATION,
     LOGIN,
@@ -30,11 +32,17 @@ from zenml.constants import (
 )
 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
 
 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,
+    pro_api_url: Optional[str] = None,
+) -> APIToken:
     """Implements the OAuth2 Device Authorization Grant flow.
 
     This function implements the client side of the OAuth2 Device Authorization
@@ -46,19 +54,25 @@ 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.
+        pro_api_url: The URL of the ZenML Pro API server. If not provided, the
+            default ZenML Pro API server URL is used.
 
     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,
@@ -67,9 +81,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
@@ -89,18 +102,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 = False
+    if not url:
+        # If no URL is provided, we use the ZenML Pro API server by default
+        zenml_pro = True
+        url = base_url = pro_api_url or 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 pro_api_url:
+            # This is a ZenML Pro server. The device authentication is done
+            # through the ZenML Pro API.
+            zenml_pro = True
+            base_url = 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.dict(),
+            data=auth_request.model_dump(exclude_none=True),
             verify=verify_ssl,
             timeout=DEFAULT_HTTP_TIMEOUT,
         )
@@ -109,7 +152,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."
+                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.")
@@ -119,7 +162,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."
+            f"Could not connect to {base_url}. Please check the URL."
         )
 
     # Open the verification URL in the user's browser
@@ -130,7 +173,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 "
@@ -147,19 +190,21 @@ 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"},
-            data=token_request.dict(),
+            data=token_request.model_dump(),
             verify=verify_ssl,
             timeout=DEFAULT_HTTP_TIMEOUT,
         )
         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())
@@ -169,7 +214,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":
@@ -196,3 +241,8 @@ 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, is_zenml_pro=zenml_pro
+    )

zenml/materializers/__init__.py

@@ -29,10 +29,9 @@ from zenml.materializers.cloudpickle_materializer import (
 from zenml.materializers.structured_string_materializer import (
     StructuredStringMaterializer,
 )
-from zenml.materializers.numpy_materializer import NumpyMaterializer
-from zenml.materializers.pandas_materializer import PandasMaterializer
 from zenml.materializers.pydantic_materializer import PydanticMaterializer
 from zenml.materializers.service_materializer import ServiceMaterializer
+from zenml.materializers.uuid_materializer import UUIDMaterializer
 
 __all__ = [
     "BuiltInContainerMaterializer",
@@ -40,8 +39,7 @@ __all__ = [
     "BytesMaterializer",
     "CloudpickleMaterializer",
     "StructuredStringMaterializer",
-    "NumpyMaterializer",
-    "PandasMaterializer",
     "PydanticMaterializer",
     "ServiceMaterializer",
+    "UUIDMaterializer",
 ]