zenml-nightly 0.72.0.dev20250115__py3-none-any.whl → 0.72.0.dev20250117__py3-none-any.whl

zenml/constants.py

@@ -176,11 +176,9 @@ ENV_ZENML_PIPELINE_RUN_API_TOKEN_EXPIRATION = (
 
 # ZenML Server environment variables
 ENV_ZENML_SERVER_PREFIX = "ZENML_SERVER_"
+ENV_ZENML_SERVER_PRO_PREFIX = "ZENML_SERVER_PRO_"
 ENV_ZENML_SERVER_DEPLOYMENT_TYPE = f"{ENV_ZENML_SERVER_PREFIX}DEPLOYMENT_TYPE"
 ENV_ZENML_SERVER_AUTH_SCHEME = f"{ENV_ZENML_SERVER_PREFIX}AUTH_SCHEME"
-ENV_ZENML_SERVER_REPORTABLE_RESOURCES = (
-    f"{ENV_ZENML_SERVER_PREFIX}REPORTABLE_RESOURCES"
-)
 ENV_ZENML_SERVER_AUTO_ACTIVATE = f"{ENV_ZENML_SERVER_PREFIX}AUTO_ACTIVATE"
 ENV_ZENML_RUN_SINGLE_STEPS_WITHOUT_STACK = (
     "ZENML_RUN_SINGLE_STEPS_WITHOUT_STACK"
@@ -321,14 +319,7 @@ DEFAULT_ZENML_SERVER_SECURE_HEADERS_REPORT_TO = "default"
 DEFAULT_ZENML_SERVER_REPORT_USER_ACTIVITY_TO_DB_SECONDS = 30
 DEFAULT_ZENML_SERVER_MAX_REQUEST_BODY_SIZE_IN_BYTES = 256 * 1024 * 1024
 
-# Configurations to decide which resources report their usage and check for
-# entitlement in the case of a cloud deployment. Expected Format is this:
-# ENV_ZENML_REPORTABLE_RESOURCES='["Foo", "bar"]'
-REPORTABLE_RESOURCES: List[str] = handle_json_env_var(
-    ENV_ZENML_SERVER_REPORTABLE_RESOURCES,
-    expected_type=list,
-    default=["pipeline", "pipeline_run", "model"],
-)
+DEFAULT_REPORTABLE_RESOURCES = ["pipeline", "pipeline_run", "model"]
 REQUIRES_CUSTOM_RESOURCE_REPORTING = ["pipeline", "pipeline_run"]
 
 # API Endpoint paths:

zenml/login/credentials.py

@@ -23,6 +23,7 @@ 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.models import ServerModel
+from zenml.models.v2.misc.server_models import ServerDeploymentType
 from zenml.services.service_status import ServiceState
 from zenml.utils.enum_utils import StrEnum
 from zenml.utils.string_utils import get_human_readable_time
@@ -44,7 +45,6 @@ class APIToken(BaseModel):
     expires_in: Optional[int] = None
     expires_at: Optional[datetime] = None
     leeway: Optional[int] = None
-    cookie_name: Optional[str] = None
     device_id: Optional[UUID] = None
     device_metadata: Optional[Dict[str, Any]] = None
 
@@ -89,13 +89,20 @@ class ServerCredentials(BaseModel):
     password: Optional[str] = None
 
     # Extra server attributes
+    deployment_type: Optional[ServerDeploymentType] = None
     server_id: Optional[UUID] = None
     server_name: Optional[str] = None
-    organization_name: Optional[str] = None
-    organization_id: Optional[UUID] = None
     status: Optional[str] = None
     version: Optional[str] = None
 
+    # Pro server attributes
+    organization_name: Optional[str] = None
+    organization_id: Optional[UUID] = None
+    tenant_name: Optional[str] = None
+    tenant_id: Optional[UUID] = None
+    pro_api_url: Optional[str] = None
+    pro_dashboard_url: Optional[str] = None
+
     @property
     def id(self) -> str:
         """Get the server identifier.
@@ -114,11 +121,13 @@ class ServerCredentials(BaseModel):
         Returns:
             The server type.
         """
-        from zenml.login.pro.utils import is_zenml_pro_server_url
-
+        if self.deployment_type == ServerDeploymentType.CLOUD:
+            return ServerType.PRO
         if self.url == ZENML_PRO_API_URL:
             return ServerType.PRO_API
-        if self.organization_id or is_zenml_pro_server_url(self.url):
+        if self.url == self.pro_api_url:
+            return ServerType.PRO_API
+        if self.organization_id or self.tenant_id:
             return ServerType.PRO
         if urlparse(self.url).hostname in [
             "localhost",
@@ -139,25 +148,39 @@ class ServerCredentials(BaseModel):
         if isinstance(server_info, ServerModel):
             # The server ID doesn't change during the lifetime of the server
             self.server_id = self.server_id or server_info.id
-
             # All other attributes can change during the lifetime of the server
+            self.deployment_type = server_info.deployment_type
             server_name = (
-                server_info.metadata.get("tenant_name") or server_info.name
+                server_info.pro_tenant_name
+                or server_info.metadata.get("tenant_name")
+                or server_info.name
             )
             if server_name:
                 self.server_name = server_name
-            organization_id = server_info.metadata.get("organization_id")
-            if organization_id:
-                self.organization_id = UUID(organization_id)
+            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_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_api_url:
+                self.pro_api_url = server_info.pro_api_url
+            if server_info.pro_dashboard_url:
+                self.pro_dashboard_url = server_info.pro_dashboard_url
             self.version = server_info.version or self.version
             # The server information was retrieved from the server itself, so we
             # can assume that the server is available
             self.status = "available"
         else:
+            self.deployment_type = ServerDeploymentType.CLOUD
             self.server_id = server_info.id
             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.status = server_info.status
             self.version = server_info.version
 
@@ -248,9 +271,10 @@ class ServerCredentials(BaseModel):
         """
         if self.organization_id and self.server_id:
             return (
-                ZENML_PRO_URL
+                (self.pro_dashboard_url or ZENML_PRO_URL)
                 + f"/organizations/{str(self.organization_id)}/tenants/{str(self.server_id)}"
             )
+
         return self.url
 
     @property
@@ -262,8 +286,8 @@ class ServerCredentials(BaseModel):
         """
         if self.organization_id:
             return (
-                ZENML_PRO_URL + f"/organizations/{str(self.organization_id)}"
-            )
+                self.pro_dashboard_url or ZENML_PRO_URL
+            ) + f"/organizations/{str(self.organization_id)}"
         return ""
 
     @property

zenml/login/credentials_store.py

@@ -25,7 +25,6 @@ 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.constants import ZENML_PRO_API_URL
 from zenml.login.pro.tenant.models import TenantRead
 from zenml.models import OAuthTokenResponse, ServerModel
 from zenml.utils import yaml_utils
@@ -289,10 +288,13 @@ class CredentialsStore(metaclass=SingletonMetaClass):
         self.check_and_reload_from_file()
         return self.credentials.get(server_url)
 
-    def get_pro_token(self, allow_expired: bool = False) -> Optional[APIToken]:
-        """Retrieve a valid token from the credentials store for the ZenML Pro API server.
+    def get_pro_token(
+        self, pro_api_url: str, allow_expired: bool = False
+    ) -> Optional[APIToken]:
+        """Retrieve a valid token from the credentials store for a ZenML Pro API server.
 
         Args:
+            pro_api_url: The URL of the ZenML Pro API server.
             allow_expired: Whether to allow expired tokens to be returned. The
                 default behavior is to return None if a token does exist but is
                 expired.
@@ -300,14 +302,18 @@ class CredentialsStore(metaclass=SingletonMetaClass):
         Returns:
             The stored token if it exists and is not expired, None otherwise.
         """
-        return self.get_token(ZENML_PRO_API_URL, allow_expired)
+        credential = self.get_pro_credentials(pro_api_url, allow_expired)
+        if credential:
+            return credential.api_token
+        return None
 
     def get_pro_credentials(
-        self, allow_expired: bool = False
+        self, pro_api_url: str, allow_expired: bool = False
     ) -> Optional[ServerCredentials]:
-        """Retrieve a valid token from the credentials store for the ZenML Pro API server.
+        """Retrieve a valid token from the credentials store for a ZenML Pro API server.
 
         Args:
+            pro_api_url: The URL of the ZenML Pro API server.
             allow_expired: Whether to allow expired tokens to be returned. The
                 default behavior is to return None if a token does exist but is
                 expired.
@@ -315,26 +321,47 @@ class CredentialsStore(metaclass=SingletonMetaClass):
         Returns:
             The stored credentials if they exist and are not expired, None otherwise.
         """
-        credential = self.get_credentials(ZENML_PRO_API_URL)
+        credential = self.get_credentials(pro_api_url)
         if (
             credential
+            and credential.type == ServerType.PRO_API
             and credential.api_token
             and (not credential.api_token.expired or allow_expired)
         ):
             return credential
         return None
 
-    def clear_pro_credentials(self) -> None:
-        """Delete the token from the store for the ZenML Pro API server."""
-        self.clear_token(ZENML_PRO_API_URL)
+    def clear_pro_credentials(self, pro_api_url: str) -> None:
+        """Delete the token from the store for a ZenML Pro API server.
+
+        Args:
+            pro_api_url: The URL of the ZenML Pro API server.
+        """
+        self.clear_token(pro_api_url)
+
+    def clear_all_pro_tokens(
+        self, pro_api_url: str
+    ) -> List[ServerCredentials]:
+        """Delete all tokens from the store for ZenML Pro servers connected to a given API server.
+
+        Args:
+            pro_api_url: The URL of the ZenML Pro API server.
 
-    def clear_all_pro_tokens(self) -> None:
-        """Delete all tokens from the store for ZenML Pro API servers."""
+        Returns:
+            A list of the credentials that were cleared.
+        """
+        credentials_to_clear = []
         for server_url, server in self.credentials.copy().items():
-            if server.type == ServerType.PRO:
+            if (
+                server.type == ServerType.PRO
+                and server.pro_api_url
+                and server.pro_api_url == pro_api_url
+            ):
                 if server.api_key:
                     continue
                 self.clear_token(server_url)
+                credentials_to_clear.append(server)
+        return credentials_to_clear
 
     def has_valid_authentication(self, url: str) -> bool:
         """Check if a valid authentication credential for the given server URL is stored.
@@ -357,13 +384,16 @@ class CredentialsStore(metaclass=SingletonMetaClass):
         token = credential.api_token
         return token is not None and not token.expired
 
-    def has_valid_pro_authentication(self) -> bool:
-        """Check if a valid token for the ZenML Pro API server is stored.
+    def has_valid_pro_authentication(self, pro_api_url: str) -> bool:
+        """Check if a valid token for a ZenML Pro API server is stored.
+
+        Args:
+            pro_api_url: The URL of the ZenML Pro API server.
 
         Returns:
             bool: True if a valid token is stored, False otherwise.
         """
-        return self.get_token(ZENML_PRO_API_URL) is not None
+        return self.get_pro_token(pro_api_url) is not None
 
     def set_api_key(
         self,
@@ -433,12 +463,14 @@ class CredentialsStore(metaclass=SingletonMetaClass):
         self,
         server_url: str,
         token_response: OAuthTokenResponse,
+        is_zenml_pro: bool = False,
     ) -> APIToken:
         """Store an API token received from an OAuth2 server.
 
         Args:
             server_url: The server URL for which the token is to be stored.
             token_response: Token response received from an OAuth2 server.
+            is_zenml_pro: Whether the token is for a ZenML Pro server.
 
         Returns:
             APIToken: The stored token.
@@ -468,7 +500,6 @@ class CredentialsStore(metaclass=SingletonMetaClass):
             expires_in=token_response.expires_in,
             expires_at=expires_at,
             leeway=leeway,
-            cookie_name=token_response.cookie_name,
             device_id=token_response.device_id,
             device_metadata=token_response.device_metadata,
         )
@@ -477,10 +508,14 @@ class CredentialsStore(metaclass=SingletonMetaClass):
         if credential:
             credential.api_token = api_token
         else:
-            self.credentials[server_url] = ServerCredentials(
+            credential = self.credentials[server_url] = ServerCredentials(
                 url=server_url, api_token=api_token
             )
 
+        if is_zenml_pro:
+            # This is how we encode that the token is for a ZenML Pro server
+            credential.pro_api_url = server_url
+
         self._save_credentials()
 
         return api_token

zenml/login/pro/client.py

@@ -33,7 +33,6 @@ from zenml.exceptions import AuthorizationException
 from zenml.logger import get_logger
 from zenml.login.credentials import APIToken
 from zenml.login.credentials_store import get_credentials_store
-from zenml.login.pro.constants import ZENML_PRO_API_URL
 from zenml.login.pro.models import BaseRestAPIModel
 from zenml.utils.singleton import SingletonMetaClass
 from zenml.zen_server.exceptions import exception_from_response
@@ -60,14 +59,11 @@ class ZenMLProClient(metaclass=SingletonMetaClass):
     _tenant: Optional["TenantClient"] = None
     _organization: Optional["OrganizationClient"] = None
 
-    def __init__(
-        self, url: Optional[str] = None, api_token: Optional[APIToken] = None
-    ) -> None:
+    def __init__(self, url: str, api_token: Optional[APIToken] = None) -> None:
         """Initialize the ZenML Pro client.
 
         Args:
-            url: The URL of the ZenML Pro API server. If not provided, the
-                default ZenML Pro API server URL is used.
+            url: The URL of the ZenML Pro API server.
             api_token: The API token to use for authentication. If not provided,
                 the token is fetched from the credentials store.
 
@@ -75,7 +71,7 @@ class ZenMLProClient(metaclass=SingletonMetaClass):
             AuthorizationException: If no API token is provided and no token
                 is found in the credentials store.
         """
-        self._url = url or ZENML_PRO_API_URL
+        self._url = url
         if api_token is None:
             logger.debug(
                 "No ZenML Pro API token provided. Fetching from credentials "

zenml/login/pro/constants.py

@@ -26,9 +26,3 @@ ENV_ZENML_PRO_URL = "ZENML_PRO_URL"
 DEFAULT_ZENML_PRO_URL = "https://cloud.zenml.io"
 
 ZENML_PRO_URL = os.getenv(ENV_ZENML_PRO_URL, default=DEFAULT_ZENML_PRO_URL)
-
-ENV_ZENML_PRO_SERVER_SUBDOMAIN = "ZENML_PRO_SERVER_SUBDOMAIN"
-DEFAULT_ZENML_PRO_SERVER_SUBDOMAIN = "cloudinfra.zenml.io"
-ZENML_PRO_SERVER_SUBDOMAIN = os.getenv(
-    ENV_ZENML_PRO_SERVER_SUBDOMAIN, default=DEFAULT_ZENML_PRO_SERVER_SUBDOMAIN
-)

zenml/login/pro/tenant/models.py

@@ -76,7 +76,7 @@ class ZenMLServiceStatus(BaseRestAPIModel):
 class ZenMLServiceRead(BaseRestAPIModel):
     """Pydantic Model for viewing a ZenML service."""
 
-    configuration: ZenMLServiceConfiguration = Field(
+    configuration: Optional[ZenMLServiceConfiguration] = Field(
         description="The service configuration."
     )
 
@@ -133,7 +133,9 @@ class TenantRead(BaseRestAPIModel):
         Returns:
             The ZenML service version.
         """
-        version = self.zenml_service.configuration.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
 

zenml/login/pro/utils.py

@@ -13,37 +13,16 @@
 #  permissions and limitations under the License.
 """ZenML Pro login utils."""
 
-import re
-
 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_SERVER_SUBDOMAIN
+from zenml.login.pro.constants import ZENML_PRO_API_URL
 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.
 
@@ -54,8 +33,15 @@ def get_troubleshooting_instructions(url: str) -> str:
         Troubleshooting instructions
     """
     credentials_store = get_credentials_store()
-    if credentials_store.has_valid_pro_authentication():
-        client = ZenMLProClient()
+
+    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)

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/login/web_login.py

@@ -34,13 +34,14 @@ 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: Optional[str] = None, verify_ssl: Optional[Union[str, bool]] = None
+    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.
 
@@ -61,6 +62,8 @@ def web_login(
         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 response returned by the OAuth2 server.
@@ -103,16 +106,16 @@ def web_login(
     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
+        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 is_zenml_pro_server_url(url):
+        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 = ZENML_PRO_API_URL
+            base_url = pro_api_url
         else:
             base_url = url
 
@@ -240,4 +243,6 @@ def web_login(
             )
 
     # Save the token in the credentials store
-    return credentials_store.set_token(url, token_response)
+    return credentials_store.set_token(
+        url, token_response, is_zenml_pro=zenml_pro
+    )

zenml/models/v2/misc/auth_models.py

@@ -119,8 +119,8 @@ class OAuthTokenResponse(BaseModel):
     token_type: str
     expires_in: Optional[int] = None
     refresh_token: Optional[str] = None
+    csrf_token: Optional[str] = None
     scope: Optional[str] = None
-    cookie_name: Optional[str] = None
     device_id: Optional[UUID] = None
     device_metadata: Optional[Dict[str, Any]] = None
 

zenml/models/v2/misc/server_models.py

@@ -107,6 +107,42 @@ class ServerModel(BaseModel):
         title="Timestamp of latest user activity traced on the server.",
     )
 
+    pro_dashboard_url: Optional[str] = Field(
+        None,
+        title="The base URL of the ZenML Pro dashboard to which the server "
+        "is connected. Only set if the server is a ZenML Pro server.",
+    )
+
+    pro_api_url: Optional[str] = Field(
+        None,
+        title="The base URL of the ZenML Pro API to which the server is "
+        "connected. Only set if the server is a ZenML Pro server.",
+    )
+
+    pro_organization_id: Optional[UUID] = Field(
+        None,
+        title="The ID of the ZenML Pro organization to which the server is "
+        "connected. Only set if the server is a ZenML Pro server.",
+    )
+
+    pro_organization_name: Optional[str] = Field(
+        None,
+        title="The name of the ZenML Pro organization to which the server is "
+        "connected. Only set if the server is a ZenML Pro server.",
+    )
+
+    pro_tenant_id: Optional[UUID] = Field(
+        None,
+        title="The ID of the ZenML Pro tenant to which the server is connected. "
+        "Only set if the server is a ZenML Pro server.",
+    )
+
+    pro_tenant_name: Optional[str] = Field(
+        None,
+        title="The name of the ZenML Pro tenant to which the server is connected. "
+        "Only set if the server is a ZenML Pro server.",
+    )
+
     def is_local(self) -> bool:
         """Return whether the server is running locally.
 
@@ -119,6 +155,14 @@ class ServerModel(BaseModel):
         # server ID is the same as the local client (user) ID.
         return self.id == GlobalConfiguration().user_id
 
+    def is_pro_server(self) -> bool:
+        """Return whether the server is a ZenML Pro server.
+
+        Returns:
+            True if the server is a ZenML Pro server, False otherwise.
+        """
+        return self.deployment_type == ServerDeploymentType.CLOUD
+
 
 class ServerLoadInfo(BaseModel):
     """Domain model for ZenML server load information."""