zenml-nightly 0.82.0.dev20250505__py3-none-any.whl → 0.82.0.dev20250506__py3-none-any.whl

zenml/VERSION

@@ -1 +1 @@
-0.82.0.dev20250505
+0.82.0.dev20250506

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/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/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/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.dist-info → zenml_nightly-0.82.0.dev20250506.dist-info}/METADATA

@@ -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.dist-info → zenml_nightly-0.82.0.dev20250506.dist-info}/RECORD

@@ -1,5 +1,5 @@
 zenml/README.md,sha256=827dekbOWAs1BpW7VF1a4d7EbwPbjwccX-2zdXBENZo,1777
-zenml/VERSION,sha256=xIFP7SS2Bk0B5oVwe5sVgUD_kvbgpml_jBENg4ew1rE,19
+zenml/VERSION,sha256=ppUfeGUhs-NasvRcl6HFrQq-3Bkc5Mvd-jli7V03Ufs,19
 zenml/__init__.py,sha256=CKEyepFK-7akXYiMrNVh92Nb01Cjs23w4_YyI6sgdc8,2242
 zenml/actions/__init__.py,sha256=mrt6wPo73iKRxK754_NqsGyJ3buW7RnVeIGXr1xEw8Y,681
 zenml/actions/base_action.py,sha256=UcaHev6BTuLDwuswnyaPjdA8AgUqB5xPZ-lRtuvf2FU,25553
@@ -64,7 +64,7 @@ zenml/code_repositories/git/local_git_repository_context.py,sha256=hfX7zVDQ27Le0
 zenml/code_repositories/local_repository_context.py,sha256=1VyiYkJBDVg0iGusgRQDToGRPJuu9lx7jTBDpplukDg,2816
 zenml/config/__init__.py,sha256=DZEic7euSbwI9Yb3FMRQhTgfhqz-C6OdAiYmOb0-opI,1519
 zenml/config/base_settings.py,sha256=itoLqc1cOwEYhgSGdZmSKSaBevQkvYH7NQh7PUamazc,1700
-zenml/config/build_configuration.py,sha256=vpea4cXTWfxRyijMFMo3U6jtHpC7wgL01NAIzp721RI,5263
+zenml/config/build_configuration.py,sha256=tEAIXwAGS8Flxo4mK5BGHsBk7LPWW16MOa6oRDymZ7k,6172
 zenml/config/compiler.py,sha256=jzyyTmM8nTxtpIVTSjZ-jrVOBJ1roVoJ3fs1A1nwv9M,24130
 zenml/config/constants.py,sha256=QvSgMwXWxtspcJ45CrFDP1ZY3w6gS3bIhXLOtIDAbZA,713
 zenml/config/docker_settings.py,sha256=w62SDW2Qo5jFxcikHr8p-Lk6ymC2ieh3MnEfkte9Lf0,13322
@@ -77,7 +77,7 @@ zenml/config/retry_config.py,sha256=4UH1xqw0G6fSEbXSNKfmiFEkwadxQef9BGMe3JBm6NI,
 zenml/config/schedule.py,sha256=qtMWa-mEo7jIKvDzQUstMwe57gdbvyWAQ7ggsoddbCA,5349
 zenml/config/secret_reference_mixin.py,sha256=YvY68MTd1gE23IVprf0BLkNn62hoxcvb5nqGgc8jMkU,5871
 zenml/config/secrets_store_config.py,sha256=y05zqyQhr_DGrs3IfBGa_FRoZ043hSYRT5wzrx-zHTU,2818
-zenml/config/server_config.py,sha256=Kns2ul12Zt4Y4ByKDJ4tgBrmUMvSrnxrc3bXRr5iQnw,31487
+zenml/config/server_config.py,sha256=xLBOs_fWUM_Imk2IDzLPC1aVUSpwYYRobYAYC9hd2qE,31558
 zenml/config/settings_resolver.py,sha256=PR9BRm_x1dy7nVKa9UqpeFdck8IEATSW6aWT8FKd-DI,4278
 zenml/config/source.py,sha256=RzUw8lin8QztUjz-AdoCzVM5Om_cSSPuroaPx-qAO4w,8226
 zenml/config/step_configurations.py,sha256=mngjobhHRj88f3klMdz6iw2mOj9wzYUPIV8Rp_2hV3g,10433
@@ -88,7 +88,7 @@ zenml/console.py,sha256=hj_KerPQKwnyKACj0ehSqUQX0mGVCJBKE1QvCt6ik3A,1160
 zenml/constants.py,sha256=ZVwWVxE2h54cNbJDyujJmhAc8w3J7I9MkWHOfP4HN9A,16431
 zenml/container_registries/__init__.py,sha256=ZSPbBIOnzhg88kQSpYgKe_POLuru14m629665-kAVAA,2200
 zenml/container_registries/azure_container_registry.py,sha256=t1sfDa94Vzbyqtb1iPFNutJ2EXV5_p9CUNITasoiQ70,2667
-zenml/container_registries/base_container_registry.py,sha256=6c2e32wuqxYHJXm5OV2LY1MtX9yopB7WZtes9fmTAz0,7625
+zenml/container_registries/base_container_registry.py,sha256=-9RIkD6oXNPaU59R3PB_PtyCqsFoLPLSn5xYZmEmzbc,8915
 zenml/container_registries/default_container_registry.py,sha256=_7lh2tQvoRgbQlSATc87Wzzl9dQzdTj8XvaTByTWKD4,1890
 zenml/container_registries/dockerhub_container_registry.py,sha256=RtdP-xqqTCX5-s9g3bYI_asTqs6BHBsTWeg3tS_Rg00,2684
 zenml/container_registries/gcp_container_registry.py,sha256=PNypXzqQ6RWDYP9_cPQUZDTf3NP6fxlsZWFD1Xu3jt4,2654
@@ -805,7 +805,7 @@ zenml/utils/yaml_utils.py,sha256=747M_BTf3lcHFgJDF8RJxnnGIbCU8e9YxBBMmoQ5O_U,583
 zenml/zen_server/__init__.py,sha256=WyltI9TzFW2mEHZVOs6alLWMCQrrZaFALtrQXs83STA,1355
 zenml/zen_server/auth.py,sha256=5EDqBb0HZJ8zr6OhwiTUcWkqoBFIEl3vfyit232UoOo,40259
 zenml/zen_server/cache.py,sha256=Tc4TSugmsU1bhThxlYfE8rv0KmltIX1CcVHgzrJ0Eus,6633
-zenml/zen_server/cloud_utils.py,sha256=nHCJjO4jRaiZiVIMbN339xG5_uLKzkRPk-FoId0hhx8,9307
+zenml/zen_server/cloud_utils.py,sha256=qKMtHjnkVfAk_yxF2K9iC0TkMCjp8s0gdz_cEsbB4DE,10908
 zenml/zen_server/csrf.py,sha256=Jsbi_IKriWCOuquSYTOEWqSXiGORJATIhR5Wrkm4XzQ,2684
 zenml/zen_server/dashboard/assets/404-D4aYbspS.js,sha256=w7_KToWud6ST-NPC7zVqs4lFHLr8_q95w3pUYc-W0xo,1033
 zenml/zen_server/dashboard/assets/@radix-C7hRs6Kx.js,sha256=cNA9UX8LGrKUQubGrls3E3Wq1fCrlK51W14yh22FE_U,296700
@@ -1318,8 +1318,8 @@ zenml/zen_stores/secrets_stores/sql_secrets_store.py,sha256=LPFW757WCJLP1S8vrvjs
 zenml/zen_stores/sql_zen_store.py,sha256=biOoDb2_zYmpsN-J-FSlKICYdwM9KDIe-_KN_yDf_mA,441414
 zenml/zen_stores/template_utils.py,sha256=GWBP5QEOyvhzndS_MLPmvh28sQaOPpPoZFXCIX9CRL4,9065
 zenml/zen_stores/zen_store_interface.py,sha256=fF_uL_FplnvGvM5o3jOQ8i1zHXhuhKLL2n4nvIKSR7E,92090
-zenml_nightly-0.82.0.dev20250505.dist-info/LICENSE,sha256=wbnfEnXnafPbqwANHkV6LUsPKOtdpsd-SNw37rogLtc,11359
-zenml_nightly-0.82.0.dev20250505.dist-info/METADATA,sha256=cnfux-P5M0Zllz1JYwFHgtBslGfFDkFSzvpzATb5wAQ,24304
-zenml_nightly-0.82.0.dev20250505.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
-zenml_nightly-0.82.0.dev20250505.dist-info/entry_points.txt,sha256=QK3ETQE0YswAM2mWypNMOv8TLtr7EjnqAFq1br_jEFE,43
-zenml_nightly-0.82.0.dev20250505.dist-info/RECORD,,
+zenml_nightly-0.82.0.dev20250506.dist-info/LICENSE,sha256=wbnfEnXnafPbqwANHkV6LUsPKOtdpsd-SNw37rogLtc,11359
+zenml_nightly-0.82.0.dev20250506.dist-info/METADATA,sha256=_sWp9WjyqeWPdnHJ3RUVzjNOID0vYkZhrN-yKVh7tuM,24315
+zenml_nightly-0.82.0.dev20250506.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
+zenml_nightly-0.82.0.dev20250506.dist-info/entry_points.txt,sha256=QK3ETQE0YswAM2mWypNMOv8TLtr7EjnqAFq1br_jEFE,43
+zenml_nightly-0.82.0.dev20250506.dist-info/RECORD,,