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

zenml/VERSION

@@ -1 +1 @@
-0.73.0.dev20250130
+0.73.0.dev20250201

zenml/cli/code_repository.py

@@ -162,6 +162,32 @@ def register_code_repository(
         cli_utils.declare(f"Successfully registered code repository `{name}`.")
 
 
+@code_repository.command("describe", help="Describe a code repository.")
+@click.argument(
+    "name_id_or_prefix",
+    type=str,
+    required=True,
+)
+def describe_code_repository(name_id_or_prefix: str) -> None:
+    """Describe a code repository.
+
+    Args:
+        name_id_or_prefix: Name, ID or prefix of the code repository.
+    """
+    client = Client()
+    try:
+        code_repository = client.get_code_repository(
+            name_id_or_prefix=name_id_or_prefix,
+        )
+    except KeyError as err:
+        cli_utils.error(str(err))
+    else:
+        cli_utils.print_pydantic_model(
+            title=f"Code repository '{code_repository.name}'",
+            model=code_repository,
+        )
+
+
 @code_repository.command("list", help="List all connected code repositories.")
 @list_options(CodeRepositoryFilter)
 def list_code_repositories(**kwargs: Any) -> None:

zenml/client.py

@@ -34,7 +34,7 @@ from typing import (
     Union,
     cast,
 )
-from uuid import UUID, uuid4
+from uuid import UUID
 
 from pydantic import ConfigDict, SecretStr
 
@@ -4980,12 +4980,7 @@ class Client(metaclass=ClientMetaClass):
             )
         )
         try:
-            # This does a login to verify the credentials
-            code_repo_class(id=uuid4(), config=config)
-
-            # Explicitly access the config for pydantic validation, in case
-            # the login for some reason did not do that.
-            _ = code_repo_class.config
+            code_repo_class.validate_config(config)
         except Exception as e:
             raise RuntimeError(
                 "Failed to validate code repository config."

zenml/code_repositories/base_code_repository.py

@@ -15,7 +15,7 @@
 
 from abc import ABC, abstractmethod
 from typing import TYPE_CHECKING, Any, Dict, Optional, Set, Type
-from uuid import UUID
+from uuid import UUID, uuid4
 
 from zenml.config.secret_reference_mixin import SecretReferenceMixin
 from zenml.logger import get_logger
@@ -44,15 +44,18 @@ class BaseCodeRepository(ABC):
     def __init__(
         self,
         id: UUID,
+        name: str,
         config: Dict[str, Any],
     ) -> None:
         """Initializes a code repository.
 
         Args:
             id: The ID of the code repository.
+            name: The name of the code repository.
             config: The config of the code repository.
         """
         self._id = id
+        self._name = name
         self._config = config
         self.login()
 
@@ -80,7 +83,23 @@ class BaseCodeRepository(ABC):
                 source=model.source, expected_class=BaseCodeRepository
             )
         )
-        return class_(id=model.id, config=model.config)
+        return class_(id=model.id, name=model.name, config=model.config)
+
+    @classmethod
+    def validate_config(cls, config: Dict[str, Any]) -> None:
+        """Validate the code repository config.
+
+        This method should check that the config/credentials are valid and
+        the configured repository exists.
+
+        Args:
+            config: The configuration.
+        """
+        # The initialization calls the login to verify the credentials
+        code_repo = cls(id=uuid4(), name="", config=config)
+
+        # Explicitly access the config for pydantic validation
+        _ = code_repo.config
 
     @property
     def id(self) -> UUID:
@@ -91,6 +110,15 @@ class BaseCodeRepository(ABC):
         """
         return self._id
 
+    @property
+    def name(self) -> str:
+        """Name of the code repository.
+
+        Returns:
+            The name of the code repository.
+        """
+        return self._name
+
     @property
     def requirements(self) -> Set[str]:
         """Set of PyPI requirements for the repository.

zenml/code_repositories/git/local_git_repository_context.py

@@ -14,11 +14,14 @@
 """Implementation of the Local git repository context."""
 
 from typing import TYPE_CHECKING, Callable, Optional, cast
-from uuid import UUID
 
 from zenml.code_repositories import (
     LocalRepositoryContext,
 )
+from zenml.constants import (
+    ENV_ZENML_CODE_REPOSITORY_IGNORE_UNTRACKED_FILES,
+    handle_bool_env_var,
+)
 from zenml.logger import get_logger
 
 if TYPE_CHECKING:
@@ -26,6 +29,8 @@ if TYPE_CHECKING:
     from git.remote import Remote
     from git.repo.base import Repo
 
+    from zenml.code_repositories import BaseCodeRepository
+
 logger = get_logger(__name__)
 
 
@@ -33,16 +38,19 @@ class LocalGitRepositoryContext(LocalRepositoryContext):
     """Local git repository context."""
 
     def __init__(
-        self, code_repository_id: UUID, git_repo: "Repo", remote_name: str
+        self,
+        code_repository: "BaseCodeRepository",
+        git_repo: "Repo",
+        remote_name: str,
     ):
         """Initializes a local git repository context.
 
         Args:
-            code_repository_id: The ID of the code repository.
+            code_repository: The code repository.
             git_repo: The git repo.
             remote_name: Name of the remote.
         """
-        super().__init__(code_repository_id=code_repository_id)
+        super().__init__(code_repository=code_repository)
         self._git_repo = git_repo
         self._remote = git_repo.remote(name=remote_name)
 
@@ -50,14 +58,14 @@ class LocalGitRepositoryContext(LocalRepositoryContext):
     def at(
         cls,
         path: str,
-        code_repository_id: UUID,
+        code_repository: "BaseCodeRepository",
         remote_url_validation_callback: Callable[[str], bool],
     ) -> Optional["LocalGitRepositoryContext"]:
         """Returns a local git repository at the given path.
 
         Args:
             path: The path to the local git repository.
-            code_repository_id: The ID of the code repository.
+            code_repository: The code repository.
             remote_url_validation_callback: A callback that validates the
                 remote URL of the git repository.
 
@@ -70,11 +78,13 @@ class LocalGitRepositoryContext(LocalRepositoryContext):
             from git.exc import InvalidGitRepositoryError
             from git.repo.base import Repo
         except ImportError:
+            logger.debug("Failed to import git library.")
             return None
 
         try:
             git_repo = Repo(path=path, search_parent_directories=True)
         except InvalidGitRepositoryError:
+            logger.debug("No git repository exists at path %s.", path)
             return None
 
         remote_name = None
@@ -87,7 +97,7 @@ class LocalGitRepositoryContext(LocalRepositoryContext):
             return None
 
         return LocalGitRepositoryContext(
-            code_repository_id=code_repository_id,
+            code_repository=code_repository,
             git_repo=git_repo,
             remote_name=remote_name,
         )
@@ -124,13 +134,19 @@ class LocalGitRepositoryContext(LocalRepositoryContext):
     def is_dirty(self) -> bool:
         """Whether the git repo is dirty.
 
-        A repository counts as dirty if it has any untracked or uncommitted
-        changes.
+        By default, a repository counts as dirty if it has any untracked or
+        uncommitted changes. Users can use an environment variable to ignore
+        untracked files.
 
         Returns:
             True if the git repo is dirty, False otherwise.
         """
-        return self.git_repo.is_dirty(untracked_files=True)
+        ignore_untracked_files = handle_bool_env_var(
+            ENV_ZENML_CODE_REPOSITORY_IGNORE_UNTRACKED_FILES, default=False
+        )
+        return self.git_repo.is_dirty(
+            untracked_files=not ignore_untracked_files
+        )
 
     @property
     def has_local_changes(self) -> bool:

zenml/code_repositories/local_repository_context.py

@@ -14,10 +14,13 @@
 """Base class for local code repository contexts."""
 
 from abc import ABC, abstractmethod
-from uuid import UUID
+from typing import TYPE_CHECKING
 
 from zenml.logger import get_logger
 
+if TYPE_CHECKING:
+    from zenml.code_repositories import BaseCodeRepository
+
 logger = get_logger(__name__)
 
 
@@ -30,22 +33,22 @@ class LocalRepositoryContext(ABC):
     commit, and whether the repository is dirty.
     """
 
-    def __init__(self, code_repository_id: UUID) -> None:
+    def __init__(self, code_repository: "BaseCodeRepository") -> None:
         """Initializes a local repository context.
 
         Args:
-            code_repository_id: The ID of the code repository.
+            code_repository: The code repository.
         """
-        self._code_repository_id = code_repository_id
+        self._code_repository = code_repository
 
     @property
-    def code_repository_id(self) -> UUID:
-        """Returns the ID of the code repository.
+    def code_repository(self) -> "BaseCodeRepository":
+        """Returns the code repository.
 
         Returns:
-            The ID of the code repository.
+            The code repository.
         """
-        return self._code_repository_id
+        return self._code_repository
 
     @property
     @abstractmethod

zenml/constants.py

@@ -175,6 +175,9 @@ ENV_ZENML_WHEEL_PACKAGE_NAME = "ZENML_WHEEL_PACKAGE_NAME"
 ENV_ZENML_PIPELINE_RUN_API_TOKEN_EXPIRATION = (
     "ZENML_PIPELINE_API_TOKEN_EXPIRATION"
 )
+ENV_ZENML_CODE_REPOSITORY_IGNORE_UNTRACKED_FILES = (
+    "ZENML_CODE_REPOSITORY_IGNORE_UNTRACKED_FILES"
+)
 
 # Materializer environment variables
 ENV_ZENML_MATERIALIZER_ALLOW_NON_ASCII_JSON_DUMPS = (

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/integrations/github/code_repositories/github_code_repository.py

@@ -15,7 +15,8 @@
 
 import os
 import re
-from typing import List, Optional
+from typing import Any, Dict, List, Optional
+from uuid import uuid4
 
 import requests
 from github import Consts, Github, GithubException
@@ -62,6 +63,20 @@ class GitHubCodeRepositoryConfig(BaseCodeRepositoryConfig):
 class GitHubCodeRepository(BaseCodeRepository):
     """GitHub code repository."""
 
+    @classmethod
+    def validate_config(cls, config: Dict[str, Any]) -> None:
+        """Validate the code repository config.
+
+        This method should check that the config/credentials are valid and
+        the configured repository exists.
+
+        Args:
+            config: The configuration.
+        """
+        code_repo = cls(id=uuid4(), name="", config=config)
+        # Try to access the project to make sure it exists
+        _ = code_repo.github_repo
+
     @property
     def config(self) -> GitHubCodeRepositoryConfig:
         """Returns the `GitHubCodeRepositoryConfig` config.
@@ -190,7 +205,7 @@ class GitHubCodeRepository(BaseCodeRepository):
         """
         return LocalGitRepositoryContext.at(
             path=path,
-            code_repository_id=self.id,
+            code_repository=self,
             remote_url_validation_callback=self.check_remote_url,
         )
 

zenml/integrations/gitlab/code_repositories/gitlab_code_repository.py

@@ -15,7 +15,8 @@
 
 import os
 import re
-from typing import Optional
+from typing import Any, Dict, Optional
+from uuid import uuid4
 
 from gitlab import Gitlab
 from gitlab.v4.objects import Project
@@ -63,6 +64,20 @@ class GitLabCodeRepositoryConfig(BaseCodeRepositoryConfig):
 class GitLabCodeRepository(BaseCodeRepository):
     """GitLab code repository."""
 
+    @classmethod
+    def validate_config(cls, config: Dict[str, Any]) -> None:
+        """Validate the code repository config.
+
+        This method should check that the config/credentials are valid and
+        the configured repository exists.
+
+        Args:
+            config: The configuration.
+        """
+        code_repo = cls(id=uuid4(), name="", config=config)
+        # Try to access the project to make sure it exists
+        _ = code_repo.gitlab_project
+
     @property
     def config(self) -> GitLabCodeRepositoryConfig:
         """Returns the `GitLabCodeRepositoryConfig` config.
@@ -147,7 +162,7 @@ class GitLabCodeRepository(BaseCodeRepository):
         """
         return LocalGitRepositoryContext.at(
             path=path,
-            code_repository_id=self.id,
+            code_repository=self,
             remote_url_validation_callback=self.check_remote_url,
         )
 

zenml/integrations/pytorch/materializers/base_pytorch_materializer.py

@@ -44,7 +44,7 @@ class BasePyTorchMaterializer(BaseMaterializer):
             # NOTE (security): The `torch.load` function uses `pickle` as
             # the default unpickler, which is NOT secure. This materializer
             # is intended for use with trusted data sources.
-            return torch.load(f)  # nosec
+            return torch.load(f, weights_only=False)  # nosec
 
     def save(self, obj: Any) -> None:
         """Uses `torch.save` to save a PyTorch object.