zenml-nightly 0.82.0.dev20250505__tar.gz → 0.82.0.dev20250506__tar.gz

{zenml_nightly-0.82.0.dev20250505 → zenml_nightly-0.82.0.dev20250506}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: zenml-nightly
-Version: 0.82.0.dev20250505
+Version: 0.82.0.dev20250506
 Summary: ZenML: Write production-ready ML code.
 License: Apache-2.0
 Keywords: machine learning,production,pipeline,mlops,devops
@@ -112,7 +112,7 @@ Requires-Dist: ruff (>=0.1.7) ; extra == "templates" or extra == "dev"
 Requires-Dist: s3fs (>=2022.11.0,!=2025.3.1) ; extra == "s3fs"
 Requires-Dist: sagemaker (>=2.237.3) ; extra == "sagemaker"
 Requires-Dist: secure (>=0.3.0,<0.4.0) ; extra == "server"
-Requires-Dist: setuptools
+Requires-Dist: setuptools (>=70.0.0)
 Requires-Dist: sqlalchemy (>=2.0.0,<3.0.0)
 Requires-Dist: sqlalchemy_utils
 Requires-Dist: sqlmodel (==0.0.18)

{zenml_nightly-0.82.0.dev20250505 → zenml_nightly-0.82.0.dev20250506}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "zenml-nightly"
-version = "0.82.0.dev20250505"
+version = "0.82.0.dev20250506"
 packages = [{ include = "zenml", from = "src" }]
 description = "ZenML: Write production-ready ML code."
 authors = ["ZenML GmbH <info@zenml.io>"]
@@ -58,7 +58,7 @@ python = ">=3.9,<3.13"
 python-dateutil = "^2.8.1"
 pyyaml = ">=6.0.1"
 rich = { extras = ["jupyter"], version = ">=12.0.0" }
-setuptools = "*"
+setuptools = ">=70.0.0"
 sqlalchemy = "^2.0.0"
 sqlalchemy_utils = "*"
 sqlmodel = "0.0.18"

zenml_nightly-0.82.0.dev20250506/src/zenml/VERSION

@@ -0,0 +1 @@
+0.82.0.dev20250506

{zenml_nightly-0.82.0.dev20250505 → zenml_nightly-0.82.0.dev20250506}/src/zenml/config/build_configuration.py

@@ -20,12 +20,15 @@ from typing import TYPE_CHECKING, Dict, Optional
 from pydantic import BaseModel
 
 from zenml.config.docker_settings import DockerSettings
+from zenml.logger import get_logger
 from zenml.utils import json_utils
 
 if TYPE_CHECKING:
     from zenml.code_repositories import BaseCodeRepository
     from zenml.stack import Stack
 
+logger = get_logger(__name__)
+
 
 class BuildConfiguration(BaseModel):
     """Configuration of Docker builds.
@@ -99,6 +102,24 @@ class BuildConfiguration(BaseModel):
             with open(self.settings.dockerfile, "rb") as f:
                 hash_.update(f.read())
 
+        if self.settings.parent_image and stack.container_registry:
+            digest = stack.container_registry.get_image_repo_digest(
+                self.settings.parent_image
+            )
+            if digest:
+                hash_.update(digest.encode())
+            else:
+                logger.warning(
+                    "Unable to fetch parent image digest for image `%s`. "
+                    "This may lead to ZenML reusing existing builds even "
+                    "though a new version of the parent image has been "
+                    "pushed. Most likely you can fix this error by making sure "
+                    "the parent image is pushed to the container registry of "
+                    "your active stack `%s`.",
+                    self.settings.parent_image,
+                    stack.name,
+                )
+
         return hash_.hexdigest()
 
     def should_include_files(

{zenml_nightly-0.82.0.dev20250505 → zenml_nightly-0.82.0.dev20250506}/src/zenml/config/server_config.py

@@ -29,6 +29,7 @@ from pydantic import (
 )
 
 from zenml.constants import (
+    DEFAULT_HTTP_TIMEOUT,
     DEFAULT_REPORTABLE_RESOURCES,
     DEFAULT_ZENML_JWT_TOKEN_ALGORITHM,
     DEFAULT_ZENML_JWT_TOKEN_LEEWAY,
@@ -679,6 +680,7 @@ class ServerProConfiguration(BaseModel):
     organization_name: Optional[str] = None
     workspace_id: UUID
     workspace_name: Optional[str] = None
+    http_timeout: int = DEFAULT_HTTP_TIMEOUT
 
     @field_validator("api_url", "dashboard_url")
     @classmethod

{zenml_nightly-0.82.0.dev20250505 → zenml_nightly-0.82.0.dev20250506}/src/zenml/container_registries/base_container_registry.py

@@ -160,6 +160,25 @@ class BaseContainerRegistry(AuthenticationMixin):
 
         return self._docker_client
 
+    def is_valid_image_name_for_registry(self, image_name: str) -> bool:
+        """Check if the image name is valid for the container registry.
+
+        Args:
+            image_name: The name of the image.
+
+        Returns:
+            `True` if the image name is valid for the container registry,
+            `False` otherwise.
+        """
+        # Remove prefixes to make sure this logic also works for DockerHub
+        image_name = image_name.removeprefix("index.docker.io/")
+        image_name = image_name.removeprefix("docker.io/")
+
+        registry_uri = self.config.uri.removeprefix("index.docker.io/")
+        registry_uri = registry_uri.removeprefix("docker.io/")
+
+        return image_name.startswith(registry_uri)
+
     def prepare_image_push(self, image_name: str) -> None:
         """Preparation before an image gets pushed.
 
@@ -183,7 +202,7 @@ class BaseContainerRegistry(AuthenticationMixin):
             ValueError: If the image name is not associated with this
                 container registry.
         """
-        if not image_name.startswith(self.config.uri):
+        if not self.is_valid_image_name_for_registry(image_name):
             raise ValueError(
                 f"Docker image `{image_name}` does not belong to container "
                 f"registry `{self.config.uri}`."
@@ -194,6 +213,25 @@ class BaseContainerRegistry(AuthenticationMixin):
             image_name, docker_client=self.docker_client
         )
 
+    def get_image_repo_digest(self, image_name: str) -> Optional[str]:
+        """Get the repository digest of an image.
+
+        Args:
+            image_name: The name of the image.
+
+        Returns:
+            The repository digest of the image.
+        """
+        if not self.is_valid_image_name_for_registry(image_name):
+            return None
+
+        try:
+            metadata = self.docker_client.images.get_registry_data(image_name)
+        except Exception:
+            return None
+
+        return cast(str, metadata.id.split(":")[-1])
+
 
 class BaseContainerRegistryFlavor(Flavor):
     """Base flavor for container registries."""

{zenml_nightly-0.82.0.dev20250505 → zenml_nightly-0.82.0.dev20250506}/src/zenml/zen_server/cloud_utils.py

@@ -1,6 +1,7 @@
 """Utils concerning anything concerning the cloud control plane backend."""
 
 from datetime import datetime, timedelta
+from threading import RLock
 from typing import Any, Dict, Optional
 
 import requests
@@ -26,6 +27,7 @@ class ZenMLCloudConnection:
         self._session: Optional[requests.Session] = None
         self._token: Optional[str] = None
         self._token_expires_at: Optional[datetime] = None
+        self._lock = RLock()
 
     def request(
         self,
@@ -55,13 +57,21 @@ class ZenMLCloudConnection:
         url = self._config.api_url + endpoint
 
         response = self.session.request(
-            method=method, url=url, params=params, json=data, timeout=7
+            method=method,
+            url=url,
+            params=params,
+            json=data,
+            timeout=self._config.http_timeout,
         )
         if response.status_code == 401:
             # Refresh the auth token and try again
-            self._clear_session()
+            self._reset_login()
             response = self.session.request(
-                method=method, url=url, params=params, json=data, timeout=7
+                method=method,
+                url=url,
+                params=params,
+                json=data,
+                timeout=self._config.http_timeout,
             )
 
         try:
@@ -164,57 +174,90 @@ class ZenMLCloudConnection:
         Returns:
             A requests session with the authentication token.
         """
-        if self._session is None:
-            # Set up the session's connection pool size to match the server's
-            # thread pool size. This allows the server to cache one connection
-            # per thread, which means we can keep connections open for longer
-            # and avoid the overhead of setting up a new connection for each
-            # request.
-            conn_pool_size = server_config().thread_pool_size
-
-            self._session = requests.Session()
-            token = self._fetch_auth_token()
-            self._session.headers.update({"Authorization": "Bearer " + token})
-            # Add the ZenML specific headers
-            self._session.headers.update(get_zenml_headers())
-
-            retries = Retry(
-                total=5, backoff_factor=0.1, status_forcelist=[502, 504]
-            )
-            self._session.mount(
-                "https://",
-                HTTPAdapter(
-                    max_retries=retries,
-                    # We only use one connection pool to be cached because we
-                    # only communicate with one remote server (the control
-                    # plane)
-                    pool_connections=1,
-                    pool_maxsize=conn_pool_size,
-                ),
+        with self._lock:
+            if self._session is None:
+                # Set up the session's connection pool size to match the server's
+                # thread pool size. This allows the server to cache one connection
+                # per thread, which means we can keep connections open for longer
+                # and avoid the overhead of setting up a new connection for each
+                # request.
+                conn_pool_size = server_config().thread_pool_size
+
+                self._session = requests.Session()
+                # Add the ZenML specific headers
+                self._session.headers.update(get_zenml_headers())
+
+                retries = Retry(
+                    connect=5,
+                    read=8,
+                    redirect=3,
+                    status=10,
+                    allowed_methods=[
+                        "HEAD",
+                        "GET",
+                        "PUT",
+                        "PATCH",
+                        "POST",
+                        "DELETE",
+                        "OPTIONS",
+                    ],
+                    status_forcelist=[
+                        408,  # Request Timeout
+                        429,  # Too Many Requests
+                        502,  # Bad Gateway
+                        503,  # Service Unavailable
+                        504,  # Gateway Timeout
+                    ],
+                    other=3,
+                    backoff_factor=0.5,
+                )
+
+                self._session.mount(
+                    "https://",
+                    HTTPAdapter(
+                        max_retries=retries,
+                        # We only use one connection pool to be cached because we
+                        # only communicate with one remote server (the control
+                        # plane)
+                        pool_connections=1,
+                        pool_maxsize=conn_pool_size,
+                    ),
+                )
+
+            # Login to the ZenML Pro Management Plane. Calling this will fetch
+            # the active session token. It will also refresh the token if it is
+            # going to expire soon or if it is already expired.
+            access_token = self._fetch_auth_token(session=self._session)
+            self._session.headers.update(
+                {"Authorization": "Bearer " + access_token}
             )
 
-        return self._session
+            return self._session
 
-    def _clear_session(self) -> None:
-        """Clear the authentication session."""
-        self._session = None
-        self._token = None
-        self._token_expires_at = None
+    def _reset_login(self) -> None:
+        """Force a new login to the ZenML Pro Management Plane."""
+        with self._lock:
+            self._token = None
+            self._token_expires_at = None
 
-    def _fetch_auth_token(self) -> str:
+    def _fetch_auth_token(self, session: requests.Session) -> str:
         """Fetch an auth token from the Cloud API.
 
-        Raises:
-            RuntimeError: If the auth token can't be fetched.
+        Args:
+            session: The session to use to fetch the auth token.
 
         Returns:
-            Auth token.
+            The auth token.
+
+        Raises:
+            RuntimeError: If the auth token can't be fetched.
         """
         if (
             self._token is not None
             and self._token_expires_at is not None
             and utc_now() + timedelta(minutes=5) < self._token_expires_at
         ):
+            # Already logged in and token is still valid
             return self._token
 
         # Get an auth token from the Cloud API
@@ -229,9 +272,13 @@ class ZenMLCloudConnection:
             "audience": self._config.oauth2_audience,
             "grant_type": "client_credentials",
         }
+
         try:
-            response = requests.post(
-                login_url, headers=headers, data=payload, timeout=7
+            response = session.post(
+                login_url,
+                headers=headers,
+                data=payload,
+                timeout=self._config.http_timeout,
             )
             response.raise_for_status()
         except Exception as e:
@@ -253,8 +300,8 @@ class ZenMLCloudConnection:
                 "Could not fetch auth token from the Cloud API."
             )
 
-        self._token = access_token
         self._token_expires_at = utc_now() + timedelta(seconds=expires_in)
+        self._token = access_token
 
         assert self._token is not None
         return self._token

zenml_nightly-0.82.0.dev20250505/src/zenml/VERSION

@@ -1 +0,0 @@
-0.82.0.dev20250505