zenml-nightly 0.73.0.dev20250131__py3-none-any.whl → 0.73.0.dev20250201__py3-none-any.whl

zenml/VERSION

@@ -1 +1 @@
-0.73.0.dev20250131
+0.73.0.dev20250201

zenml/integrations/azure/artifact_stores/azure_artifact_store.py

@@ -64,7 +64,10 @@ class AzureArtifactStore(BaseArtifactStore, AuthenticationMixin):
         """
         connector = self.get_connector()
         if connector:
-            from azure.identity import ClientSecretCredential
+            from azure.identity import (
+                ClientSecretCredential,
+                DefaultAzureCredential,
+            )
             from azure.storage.blob import BlobServiceClient
 
             client = connector.connect()
@@ -77,18 +80,24 @@ class AzureArtifactStore(BaseArtifactStore, AuthenticationMixin):
                 )
             # Get the credentials from the client
             credentials = client.credential
-            if not isinstance(credentials, ClientSecretCredential):
+
+            if isinstance(credentials, ClientSecretCredential):
+                return AzureSecretSchema(
+                    client_id=credentials._client_id,
+                    client_secret=credentials._client_credential,
+                    tenant_id=credentials._tenant_id,
+                    account_name=client.account_name,
+                )
+
+            elif isinstance(credentials, DefaultAzureCredential):
+                return AzureSecretSchema(account_name=client.account_name)
+
+            else:
                 raise RuntimeError(
                     "The Azure Artifact Store connector can only be used "
                     "with a service connector that is configured with "
-                    "Azure service principal credentials."
+                    "Azure service principal credentials or implicit authentication"
                 )
-            return AzureSecretSchema(
-                client_id=credentials._client_id,
-                client_secret=credentials._client_credential,
-                tenant_id=credentials._tenant_id,
-                account_name=client.account_name,
-            )
 
         secret = self.get_typed_authentication_secret(
             expected_schema_type=AzureSecretSchema

zenml/integrations/azure/service_connectors/azure_service_connector.py

@@ -14,12 +14,14 @@
 """Azure Service Connector."""
 
 import datetime
+import json
 import logging
 import re
 import subprocess
 from typing import Any, Dict, List, Optional, Tuple
 from uuid import UUID
 
+import requests
 import yaml
 from azure.core.credentials import AccessToken, TokenCredential
 from azure.core.exceptions import AzureError
@@ -69,6 +71,7 @@ logger = get_logger(__name__)
 AZURE_MANAGEMENT_TOKEN_SCOPE = "https://management.azure.com/.default"
 AZURE_SESSION_TOKEN_DEFAULT_EXPIRATION_TIME = 60 * 60  # 1 hour
 AZURE_SESSION_EXPIRATION_BUFFER = 15  # 15 minutes
+AZURE_ACR_OAUTH_SCOPE = "repository:*:*"
 
 
 class AzureBaseConfig(AuthenticationConfig):
@@ -362,8 +365,8 @@ neither a storage account nor a resource group is configured in the connector,
 all blob storage containers in all accessible storage accounts will be
 accessible.
 
-The only Azure authentication method that works with Azure blob storage
-resources is the service principal authentication method.
+The only Azure authentication methods that work with Azure blob storage resources are the implicit 
+authentication and the service principal authentication method.
 """,
             auth_methods=AzureAuthenticationMethods.values(),
             # Request a blob container to be configured in the
@@ -435,12 +438,10 @@ following formats:
 If a resource group is configured in the connector, only ACR registries in that
 resource group will be accessible.
 
-If an authentication method other than the Azure service principal is used for
-authentication, the admin account must be enabled for the registry, otherwise
-clients will not be able to authenticate to the registry. See the official Azure
-[documentation on the admin account](https://docs.microsoft.com/en-us/azure/container-registry/container-registry-authentication#admin-account
-)
-for more information.
+If an authentication method other than the Azure service principal is used, Entra ID authentication is used.
+This requires the configured identity to have the `AcrPush` role to be configured.
+If this fails, admin account authentication is tried. For this the admin account must be enabled for the registry.
+See the official Azure[documentation on the admin account](https://docs.microsoft.com/en-us/azure/container-registry/container-registry-authentication#admin-account) for more information. 
 """,
             auth_methods=AzureAuthenticationMethods.values(),
             supports_instances=True,
@@ -1718,10 +1719,14 @@ class AzureServiceConnector(ServiceConnector):
             resource_type=resource_type,
             resource_id=resource_id,
         )
-        resource_group: Optional[str] = None
+
+        resource_group: Optional[str]
+        registry_name: str
+        cluster_name: str
 
         if resource_type == DOCKER_REGISTRY_RESOURCE_TYPE:
             registry_name = self._parse_acr_resource_id(resource_id)
+            registry_domain = f"{registry_name}.azurecr.io"
 
             # If a service principal is used for authentication, the client ID
             # and client secret can be used to authenticate to the registry, if
@@ -1733,13 +1738,12 @@ class AzureServiceConnector(ServiceConnector):
             ):
                 assert isinstance(self.config, AzureServicePrincipalConfig)
                 username = str(self.config.client_id)
-                password = self.config.client_secret
+                password = str(self.config.client_secret)
 
-            # Without a service principal, this only works if the admin account
-            # is enabled for the registry, but this is not recommended and
-            # disabled by default.
+            # Without a service principal, we try to use the AzureDefaultCredentials to authenticate against the ACR.
+            # If this fails, we try to use the admin account.
+            # This has to be enabled for the registry, but this is not recommended and disabled by default.
             # https://docs.microsoft.com/en-us/azure/container-registry/container-registry-authentication#admin-account
-
             else:
                 registries = self._list_acr_registries(
                     credential,
@@ -1748,22 +1752,35 @@ class AzureServiceConnector(ServiceConnector):
                 registry_name, resource_group = registries.popitem()
 
                 try:
-                    client = ContainerRegistryManagementClient(
-                        credential, subscription_id
+                    username = "00000000-0000-0000-0000-000000000000"
+                    password = _ACRTokenExchangeClient(
+                        credential
+                    ).get_acr_access_token(
+                        registry_domain, AZURE_ACR_OAUTH_SCOPE
                     )
-
-                    registry_credentials = client.registries.list_credentials(
-                        resource_group, registry_name
+                except AuthorizationException:
+                    logger.warning(
+                        "Falling back to admin credentials for ACR authentication. Be sure to assign AcrPush role to the configured identity."
                     )
-                    username = registry_credentials.username
-                    password = registry_credentials.passwords[0].value
-                except AzureError as e:
-                    raise AuthorizationException(
-                        f"failed to list admin credentials for Azure Container "
-                        f"Registry '{registry_name}' in resource group "
-                        f"'{resource_group}'. Make sure the registry is "
-                        f"configured with an admin account: {e}"
-                    ) from e
+                    try:
+                        client = ContainerRegistryManagementClient(
+                            credential, subscription_id
+                        )
+
+                        registry_credentials = (
+                            client.registries.list_credentials(
+                                resource_group, registry_name
+                            )
+                        )
+                        username = registry_credentials.username
+                        password = registry_credentials.passwords[0].value
+                    except AzureError as e:
+                        raise AuthorizationException(
+                            f"failed to list admin credentials for Azure Container "
+                            f"Registry '{registry_name}' in resource group "
+                            f"'{resource_group}'. Make sure the registry is "
+                            f"configured with an admin account: {e}"
+                        ) from e
 
             # Create a client-side Docker connector instance with the temporary
             # Docker credentials
@@ -1775,7 +1792,7 @@ class AzureServiceConnector(ServiceConnector):
                 config=DockerConfiguration(
                     username=username,
                     password=password,
-                    registry=f"{registry_name}.azurecr.io",
+                    registry=registry_domain,
                 ),
                 expires_at=expires_at,
             )
@@ -1847,3 +1864,103 @@ class AzureServiceConnector(ServiceConnector):
             )
 
         raise ValueError(f"Unsupported resource type: {resource_type}")
+
+
+class _ACRTokenExchangeClient:
+    def __init__(self, credential: TokenCredential):
+        self._credential = credential
+
+    def _get_aad_access_token(self) -> str:
+        aad_access_token: str = self._credential.get_token(
+            AZURE_MANAGEMENT_TOKEN_SCOPE
+        ).token
+        return aad_access_token
+
+    # https://github.com/Azure/acr/blob/main/docs/AAD-OAuth.md#authenticating-docker-with-an-acr-refresh-token
+    def get_acr_refresh_token(self, acr_url: str) -> str:
+        try:
+            aad_access_token = self._get_aad_access_token()
+
+            headers = {
+                "Content-Type": "application/x-www-form-urlencoded",
+            }
+
+            data = {
+                "grant_type": "access_token",
+                "service": acr_url,
+                "access_token": aad_access_token,
+            }
+
+            response = requests.post(
+                f"https://{acr_url}/oauth2/exchange",
+                headers=headers,
+                data=data,
+                timeout=5,
+            )
+
+            if response.status_code != 200:
+                raise AuthorizationException(
+                    f"failed to get refresh token for Azure Container "
+                    f"Registry '{acr_url}' in resource group. "
+                    f"Be sure to assign AcrPush to the configured principal. "
+                    f"The token exchange returned status {response.status_code} "
+                    f"with body '{response.content.decode()}'"
+                )
+
+            acr_refresh_token_response = json.loads(response.content)
+            acr_refresh_token: str = acr_refresh_token_response[
+                "refresh_token"
+            ]
+            return acr_refresh_token
+
+        except (AzureError, requests.exceptions.RequestException) as e:
+            raise AuthorizationException(
+                f"failed to get refresh token for Azure Container "
+                f"Registry '{acr_url}' in resource group. "
+                f"Make sure the implicit authentication identity "
+                f"has access to the configured registry: {e}"
+            ) from e
+
+    # https://github.com/Azure/acr/blob/main/docs/AAD-OAuth.md#calling-post-oauth2token-to-get-an-acr-access-token
+    def get_acr_access_token(self, acr_url: str, scope: str) -> str:
+        acr_refresh_token = self.get_acr_refresh_token(acr_url)
+
+        headers = {
+            "Content-Type": "application/x-www-form-urlencoded",
+        }
+
+        data = {
+            "grant_type": "refresh_token",
+            "service": acr_url,
+            "scope": scope,
+            "refresh_token": acr_refresh_token,
+        }
+
+        try:
+            response = requests.post(
+                f"https://{acr_url}/oauth2/token",
+                headers=headers,
+                data=data,
+                timeout=5,
+            )
+
+            if response.status_code != 200:
+                raise AuthorizationException(
+                    f"failed to get access token for Azure Container "
+                    f"Registry '{acr_url}' in resource group. "
+                    f"Be sure to assign AcrPush to the configured principal. "
+                    f"The token exchange returned status {response.status_code} "
+                    f"with body '{response.content.decode()}'"
+                )
+
+            acr_access_token_response = json.loads(response.content)
+            acr_access_token: str = acr_access_token_response["access_token"]
+            return acr_access_token
+
+        except (AzureError, requests.exceptions.RequestException) as e:
+            raise AuthorizationException(
+                f"failed to get access token for Azure Container "
+                f"Registry '{acr_url}' in resource group. "
+                f"Make sure the implicit authentication identity "
+                f"has access to the configured registry: {e}"
+            ) from e

{zenml_nightly-0.73.0.dev20250131.dist-info → zenml_nightly-0.73.0.dev20250201.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: zenml-nightly
-Version: 0.73.0.dev20250131
+Version: 0.73.0.dev20250201
 Summary: ZenML: Write production-ready ML code.
 License: Apache-2.0
 Keywords: machine learning,production,pipeline,mlops,devops
@@ -104,6 +104,7 @@ Requires-Dist: python-dateutil (>=2.8.1,<3.0.0)
 Requires-Dist: python-multipart (>=0.0.9,<0.1.0) ; extra == "server"
 Requires-Dist: pyyaml (>=6.0.1)
 Requires-Dist: pyyaml-include (<2.0) ; extra == "templates"
+Requires-Dist: requests (>=2.27.11,<3.0.0) ; extra == "connectors-azure"
 Requires-Dist: rich[jupyter] (>=12.0.0)
 Requires-Dist: ruff (>=0.1.7) ; extra == "templates" or extra == "dev"
 Requires-Dist: s3fs (>=2022.11.0) ; extra == "s3fs"

{zenml_nightly-0.73.0.dev20250131.dist-info → zenml_nightly-0.73.0.dev20250201.dist-info}/RECORD

@@ -1,5 +1,5 @@
 zenml/README.md,sha256=827dekbOWAs1BpW7VF1a4d7EbwPbjwccX-2zdXBENZo,1777
-zenml/VERSION,sha256=yZ41gV2RoySfmEaBFCUmGXHDN81pnq8bzattD9wYAn4,19
+zenml/VERSION,sha256=cT39b7vWwaaT2iOZDbZ7n1WpXJzY1cAQW-F6qehm6kE,19
 zenml/__init__.py,sha256=SkMObQA41ajqdZqGErN00S1Vf3KAxpLvbZ-OBy5uYoo,2130
 zenml/actions/__init__.py,sha256=mrt6wPo73iKRxK754_NqsGyJ3buW7RnVeIGXr1xEw8Y,681
 zenml/actions/base_action.py,sha256=UcaHev6BTuLDwuswnyaPjdA8AgUqB5xPZ-lRtuvf2FU,25553
@@ -155,7 +155,7 @@ zenml/integrations/aws/step_operators/sagemaker_step_operator.py,sha256=-y1zqfID
 zenml/integrations/aws/step_operators/sagemaker_step_operator_entrypoint_config.py,sha256=2cXroe6-bXyHVkO5yPnZagNpqx6MrMDcvuvNr8J2j-A,1581
 zenml/integrations/azure/__init__.py,sha256=99cefcfBlEXl5lYCVUzWN0uvYn_TFGpBsEATvn1d3c4,2938
 zenml/integrations/azure/artifact_stores/__init__.py,sha256=dlIwbpgjE0Hy4rhMbelNJHVKm4t8tj_hRu9mQ_cEIAg,820
-zenml/integrations/azure/artifact_stores/azure_artifact_store.py,sha256=uVMpsWDIyeytn8aHxBAma5pdPwiWuK5fwew2Tvxsk3w,11955
+zenml/integrations/azure/artifact_stores/azure_artifact_store.py,sha256=wgXoZFdK1tT4XneOJTOr3uYd5e6dS5ruyiNo8xldCDk,12236
 zenml/integrations/azure/azureml_utils.py,sha256=kl_ld7jgWQgVBmagO46ADvfIKvBK3qaGboZom2C5UIA,7561
 zenml/integrations/azure/flavors/__init__.py,sha256=CRdRbAnd28j2SFETYz-6PlWwhyTzr3dMwW9-LQL2G7w,1411
 zenml/integrations/azure/flavors/azure_artifact_store_flavor.py,sha256=b3gwcqsY1sYI8MevIyz5pu5THo5aZHgmz8MCqiOfQ7c,3589
@@ -166,7 +166,7 @@ zenml/integrations/azure/orchestrators/__init__.py,sha256=q4rBPIJHcuUr6dLUBdrTkQ
 zenml/integrations/azure/orchestrators/azureml_orchestrator.py,sha256=mkbjg9j-s1e2G191Q0dy4LKOeVeXV4pYr1qTuxIbMFQ,19912
 zenml/integrations/azure/orchestrators/azureml_orchestrator_entrypoint_config.py,sha256=lacOyorsHa-HuD_kxN9M6BUiZDvs5jL9AnJWwrFCtp4,3104
 zenml/integrations/azure/service_connectors/__init__.py,sha256=yMz6bTCtIZqZwfEM6h7-PSWsd_DB8l0OD9z_bdzomZw,793
-zenml/integrations/azure/service_connectors/azure_service_connector.py,sha256=oJItaOh0GQX0UAQG7jVMHBb9T_LXMCjG8bzfGkB_JRc,76370
+zenml/integrations/azure/service_connectors/azure_service_connector.py,sha256=8kAnH6NK9J6SfJPPxlmxIc2oH73YtDI3qElgKOZSRuI,81330
 zenml/integrations/azure/step_operators/__init__.py,sha256=fV4_nAO0cH53x6_-F8-CbDEZwb_Vv64oq1r0-vtigEU,819
 zenml/integrations/azure/step_operators/azureml_step_operator.py,sha256=cmDLcxVd4lw_C2hz-_z50i9tcGgGbNtrBeL5FBQ3HTE,7394
 zenml/integrations/bentoml/__init__.py,sha256=q7Vo8DL4j5N3p2Gu0_vio5VUUlW8wQ6N-ZcunK_ibiY,1876
@@ -1293,8 +1293,8 @@ zenml/zen_stores/secrets_stores/sql_secrets_store.py,sha256=nEO0bAPlULBLxLVk-UTR
 zenml/zen_stores/sql_zen_store.py,sha256=GEBQDPhm52-YyxLBJcebNviwtr-VK_dnaHrg21fzJOw,417086
 zenml/zen_stores/template_utils.py,sha256=EKYBgmDLTS_PSMWaIO5yvHPLiQvMqHcsAe6NUCrv-i4,9068
 zenml/zen_stores/zen_store_interface.py,sha256=vf2gKBWfUUPtcGZC35oQB6pPNVzWVyQC8nWxVLjfrxM,92692
-zenml_nightly-0.73.0.dev20250131.dist-info/LICENSE,sha256=wbnfEnXnafPbqwANHkV6LUsPKOtdpsd-SNw37rogLtc,11359
-zenml_nightly-0.73.0.dev20250131.dist-info/METADATA,sha256=n0WxU1eiS-WyTdH9oHEG_eQRzqr7ZMGMLWPJnA6rU4o,21355
-zenml_nightly-0.73.0.dev20250131.dist-info/WHEEL,sha256=IYZQI976HJqqOpQU6PHkJ8fb3tMNBFjg-Cn-pwAbaFM,88
-zenml_nightly-0.73.0.dev20250131.dist-info/entry_points.txt,sha256=QK3ETQE0YswAM2mWypNMOv8TLtr7EjnqAFq1br_jEFE,43
-zenml_nightly-0.73.0.dev20250131.dist-info/RECORD,,
+zenml_nightly-0.73.0.dev20250201.dist-info/LICENSE,sha256=wbnfEnXnafPbqwANHkV6LUsPKOtdpsd-SNw37rogLtc,11359
+zenml_nightly-0.73.0.dev20250201.dist-info/METADATA,sha256=x5kwq8JoXQFnocvMyIGWJHVE7YfX6ZXpDI3MrVq53YA,21428
+zenml_nightly-0.73.0.dev20250201.dist-info/WHEEL,sha256=IYZQI976HJqqOpQU6PHkJ8fb3tMNBFjg-Cn-pwAbaFM,88
+zenml_nightly-0.73.0.dev20250201.dist-info/entry_points.txt,sha256=QK3ETQE0YswAM2mWypNMOv8TLtr7EjnqAFq1br_jEFE,43
+zenml_nightly-0.73.0.dev20250201.dist-info/RECORD,,