zenml-nightly 0.73.0.dev20250127__py3-none-any.whl → 0.73.0.dev20250129__py3-none-any.whl

zenml/VERSION

@@ -1 +1 @@
-0.73.0.dev20250127
+0.73.0.dev20250129

zenml/cli/code_repository.py

@@ -42,7 +42,7 @@ def code_repository() -> None:
     context_settings={"ignore_unknown_options": True},
     help="Register a code repository.",
 )
-@click.argument("name", type=click.STRING)
+@click.argument("name", type=str)
 @click.option(
     "--type",
     "-t",
@@ -183,6 +183,74 @@ def list_code_repositories(**kwargs: Any) -> None:
         )
 
 
+@code_repository.command(
+    "update",
+    help="Update a code repository.",
+    context_settings={"ignore_unknown_options": True},
+)
+@click.argument("name_or_id", type=str, required=True)
+@click.option(
+    "--name",
+    "-n",
+    type=str,
+    required=False,
+    help="The new code repository name.",
+)
+@click.option(
+    "--description",
+    "-d",
+    type=str,
+    required=False,
+    help="The new code repository description.",
+)
+@click.option(
+    "--logo-url",
+    "-l",
+    type=str,
+    required=False,
+    help="New URL of a logo (png, jpg or svg) for the code repository.",
+)
+@click.argument(
+    "args",
+    nargs=-1,
+    type=click.UNPROCESSED,
+)
+def update_code_repository(
+    name_or_id: str,
+    name: Optional[str],
+    description: Optional[str],
+    logo_url: Optional[str],
+    args: List[str],
+) -> None:
+    """Update a code repository.
+
+    Args:
+        name_or_id: Name or ID of the code repository to update.
+        name: New name of the code repository.
+        description: New description of the code repository.
+        logo_url: New logo URL of the code repository.
+        args: Code repository configurations.
+    """
+    parsed_name_or_id, parsed_args = cli_utils.parse_name_and_extra_arguments(
+        list(args) + [name_or_id], expand_args=True, name_mandatory=True
+    )
+    assert parsed_name_or_id
+
+    with console.status(
+        f"Updating code repository '{parsed_name_or_id}'...\n"
+    ):
+        Client().update_code_repository(
+            name_id_or_prefix=parsed_name_or_id,
+            name=name,
+            description=description,
+            logo_url=logo_url,
+            config=parsed_args,
+        )
+        cli_utils.declare(
+            f"Successfully updated code repository `{parsed_name_or_id}`."
+        )
+
+
 @code_repository.command("delete")
 @click.argument("name_or_id", type=str, required=True)
 @click.option(

zenml/client.py

@@ -2453,7 +2453,9 @@ class Client(metaclass=ClientMetaClass):
     def trigger_pipeline(
         self,
         pipeline_name_or_id: Union[str, UUID, None] = None,
-        run_configuration: Optional[PipelineRunConfiguration] = None,
+        run_configuration: Union[
+            PipelineRunConfiguration, Dict[str, Any], None
+        ] = None,
         config_path: Optional[str] = None,
         template_id: Optional[UUID] = None,
         stack_name_or_id: Union[str, UUID, None] = None,
@@ -2524,6 +2526,11 @@ class Client(metaclass=ClientMetaClass):
         if config_path:
             run_configuration = PipelineRunConfiguration.from_yaml(config_path)
 
+        if isinstance(run_configuration, Dict):
+            run_configuration = PipelineRunConfiguration.model_validate(
+                run_configuration
+            )
+
         if run_configuration:
             validate_run_config_is_runnable_from_server(run_configuration)
 
@@ -4952,25 +4959,15 @@ class Client(metaclass=ClientMetaClass):
 
     # --------------------------- Code repositories ---------------------------
 
-    def create_code_repository(
-        self,
-        name: str,
-        config: Dict[str, Any],
-        source: Source,
-        description: Optional[str] = None,
-        logo_url: Optional[str] = None,
-    ) -> CodeRepositoryResponse:
-        """Create a new code repository.
+    @staticmethod
+    def _validate_code_repository_config(
+        source: Source, config: Dict[str, Any]
+    ) -> None:
+        """Validate a code repository config.
 
         Args:
-            name: Name of the code repository.
-            config: The configuration for the code repository.
-            source: The code repository implementation source.
-            description: The code repository description.
-            logo_url: URL of a logo (png, jpg or svg) for the code repository.
-
-        Returns:
-            The created code repository.
+            source: The code repository source.
+            config: The code repository config.
 
         Raises:
             RuntimeError: If the provided config is invalid.
@@ -4983,13 +4980,38 @@ class Client(metaclass=ClientMetaClass):
             )
         )
         try:
-            # Validate the repo config
+            # 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
         except Exception as e:
             raise RuntimeError(
                 "Failed to validate code repository config."
             ) from e
 
+    def create_code_repository(
+        self,
+        name: str,
+        config: Dict[str, Any],
+        source: Source,
+        description: Optional[str] = None,
+        logo_url: Optional[str] = None,
+    ) -> CodeRepositoryResponse:
+        """Create a new code repository.
+
+        Args:
+            name: Name of the code repository.
+            config: The configuration for the code repository.
+            source: The code repository implementation source.
+            description: The code repository description.
+            logo_url: URL of a logo (png, jpg or svg) for the code repository.
+
+        Returns:
+            The created code repository.
+        """
+        self._validate_code_repository_config(source=source, config=config)
         repo_request = CodeRepositoryRequest(
             user=self.active_user.id,
             workspace=self.active_workspace.id,
@@ -5088,6 +5110,7 @@ class Client(metaclass=ClientMetaClass):
         name: Optional[str] = None,
         description: Optional[str] = None,
         logo_url: Optional[str] = None,
+        config: Optional[Dict[str, Any]] = None,
     ) -> CodeRepositoryResponse:
         """Update a code repository.
 
@@ -5097,6 +5120,10 @@ class Client(metaclass=ClientMetaClass):
             name: New name of the code repository.
             description: New description of the code repository.
             logo_url: New logo URL of the code repository.
+            config: New configuration options for the code repository. Will
+                be used to update the existing configuration values. To remove
+                values from the existing configuration, set the value for that
+                key to `None`.
 
         Returns:
             The updated code repository.
@@ -5107,6 +5134,18 @@ class Client(metaclass=ClientMetaClass):
         update = CodeRepositoryUpdate(
             name=name, description=description, logo_url=logo_url
         )
+        if config is not None:
+            combined_config = repo.config
+            combined_config.update(config)
+            combined_config = {
+                k: v for k, v in combined_config.items() if v is not None
+            }
+
+            self._validate_code_repository_config(
+                source=repo.source, config=combined_config
+            )
+            update.config = combined_config
+
         return self.zen_store.update_code_repository(
             code_repository_id=repo.id, update=update
         )

zenml/config/schedule.py

@@ -13,10 +13,15 @@
 #  permissions and limitations under the License.
 """Class for defining a pipeline schedule."""
 
-import datetime
+from datetime import datetime, timedelta
 from typing import Optional
 
-from pydantic import BaseModel, model_validator
+from pydantic import (
+    BaseModel,
+    ValidationInfo,
+    field_validator,
+    model_validator,
+)
 
 from zenml.logger import get_logger
 
@@ -32,8 +37,12 @@ class Schedule(BaseModel):
             and time.
         cron_expression: Cron expression for the pipeline schedule. If a value
             for this is set it takes precedence over the start time + interval.
-        start_time: datetime object to indicate when to start the schedule.
-        end_time: datetime object to indicate when to end the schedule.
+        start_time: When the schedule should start. If this is a datetime object
+            without any timezone, it is treated as a datetime in the local
+            timezone.
+        end_time: When the schedule should end. If this is a datetime object
+            without any timezone, it is treated as a datetime in the local
+            timezone.
         interval_second: datetime timedelta indicating the seconds between two
             recurring runs for a periodic schedule.
         catchup: Whether the recurring run should catch up if behind schedule.
@@ -43,17 +52,45 @@ class Schedule(BaseModel):
             schedules the latest interval if more than one interval is ready to
             be scheduled. Usually, if your pipeline handles backfill
             internally, you should turn catchup off to avoid duplicate backfill.
-        run_once_start_time: datetime object to indicate when to run the
-            pipeline once. This is useful for one-off runs.
+        run_once_start_time: When to run the pipeline once. If this is a
+            datetime object without any timezone, it is treated as a datetime
+            in the local timezone.
     """
 
     name: Optional[str] = None
     cron_expression: Optional[str] = None
-    start_time: Optional[datetime.datetime] = None
-    end_time: Optional[datetime.datetime] = None
-    interval_second: Optional[datetime.timedelta] = None
+    start_time: Optional[datetime] = None
+    end_time: Optional[datetime] = None
+    interval_second: Optional[timedelta] = None
     catchup: bool = False
-    run_once_start_time: Optional[datetime.datetime] = None
+    run_once_start_time: Optional[datetime] = None
+
+    @field_validator(
+        "start_time", "end_time", "run_once_start_time", mode="after"
+    )
+    @classmethod
+    def _ensure_timezone(
+        cls, value: Optional[datetime], info: ValidationInfo
+    ) -> Optional[datetime]:
+        """Ensures that all datetimes are timezone aware.
+
+        Args:
+            value: The datetime.
+            info: The validation info.
+
+        Returns:
+            A timezone aware datetime or None.
+        """
+        if value and value.tzinfo is None:
+            assert info.field_name
+            logger.warning(
+                "Your schedule `%s` is missing a timezone. It will be treated "
+                "as a datetime in your local timezone.",
+                info.field_name,
+            )
+            value = value.astimezone()
+
+        return value
 
     @model_validator(mode="after")
     def _ensure_cron_or_periodic_schedule_configured(self) -> "Schedule":

zenml/constants.py

@@ -138,6 +138,7 @@ ENV_ZENML_USER_ID = "ZENML_USER_ID"
 ENV_ZENML_CONFIG_PATH = "ZENML_CONFIG_PATH"
 ENV_ZENML_DEBUG = "ZENML_DEBUG"
 ENV_ZENML_LOGGING_VERBOSITY = "ZENML_LOGGING_VERBOSITY"
+ENV_ZENML_LOGGING_FORMAT = "ZENML_LOGGING_FORMAT"
 ENV_ZENML_REPOSITORY_PATH = "ZENML_REPOSITORY_PATH"
 ENV_ZENML_PREVENT_PIPELINE_EXECUTION = "ZENML_PREVENT_PIPELINE_EXECUTION"
 ENV_ZENML_ENABLE_RICH_TRACEBACK = "ZENML_ENABLE_RICH_TRACEBACK"
@@ -154,6 +155,7 @@ ENV_ZENML_SECRETS_STORE_PREFIX = "ZENML_SECRETS_STORE_"
 ENV_ZENML_BACKUP_SECRETS_STORE_PREFIX = "ZENML_BACKUP_SECRETS_STORE_"
 ENV_ZENML_SKIP_PIPELINE_REGISTRATION = "ZENML_SKIP_PIPELINE_REGISTRATION"
 ENV_AUTO_OPEN_DASHBOARD = "AUTO_OPEN_DASHBOARD"
+ENV_ZENML_AUTO_OPEN_DASHBOARD = "ZENML_AUTO_OPEN_DASHBOARD"
 ENV_ZENML_DISABLE_DATABASE_MIGRATION = "DISABLE_DATABASE_MIGRATION"
 ENV_ZENML_LOCAL_STORES_PATH = "ZENML_LOCAL_STORES_PATH"
 ENV_ZENML_CONTAINER = "ZENML_CONTAINER"

zenml/entrypoints/base_entrypoint_configuration.py

@@ -22,6 +22,8 @@ from uuid import UUID
 
 from zenml.client import Client
 from zenml.code_repositories import BaseCodeRepository
+from zenml.enums import StackComponentType
+from zenml.exceptions import CustomFlavorImportError
 from zenml.logger import get_logger
 from zenml.utils import (
     code_repository_utils,
@@ -31,6 +33,7 @@ from zenml.utils import (
 )
 
 if TYPE_CHECKING:
+    from zenml.artifact_stores import BaseArtifactStore
     from zenml.models import CodeReferenceResponse, PipelineDeploymentResponse
 
 logger = get_logger(__name__)
@@ -204,6 +207,8 @@ class BaseEntrypointConfiguration(ABC):
                 decision instead.
 
         Raises:
+            CustomFlavorImportError: If the artifact store flavor can't be
+                imported.
             RuntimeError: If the current environment requires code download
                 but the deployment does not have a reference to any code.
         """
@@ -214,12 +219,34 @@ class BaseEntrypointConfiguration(ABC):
         if not should_download_code:
             return
 
-        if code_reference := deployment.code_reference:
+        if code_path := deployment.code_path:
+            # Load the artifact store not from the active stack but separately.
+            # This is required in case the stack has custom flavor components
+            # (other than the artifact store) for which the flavor
+            # implementations will only be available once the download finishes.
+            try:
+                artifact_store = self._load_active_artifact_store()
+            except CustomFlavorImportError as e:
+                raise CustomFlavorImportError(
+                    "Failed to import custom artifact store flavor. The "
+                    "artifact store flavor is needed to download your code, "
+                    "but it looks like it might be part of the files "
+                    "that we're trying to download. If this is the case, you "
+                    "should disable downloading code from the artifact store "
+                    "using `DockerSettings(allow_download_from_artifact_store=False)` "
+                    "or make sure the artifact flavor files are included in "
+                    "Docker image by using a custom parent image or installing "
+                    "them as part of a pip dependency."
+                ) from e
+            code_utils.download_code_from_artifact_store(
+                code_path=code_path, artifact_store=artifact_store
+            )
+        elif code_reference := deployment.code_reference:
+            # TODO: This might fail if the code repository had unpushed changes
+            # at the time the pipeline run was started.
             self.download_code_from_code_repository(
                 code_reference=code_reference
             )
-        elif code_path := deployment.code_path:
-            code_utils.download_code_from_artifact_store(code_path=code_path)
         else:
             raise RuntimeError(
                 "Code download required but no code reference or path provided."
@@ -298,6 +325,22 @@ class BaseEntrypointConfiguration(ABC):
 
         return False
 
+    def _load_active_artifact_store(self) -> "BaseArtifactStore":
+        """Load the active artifact store.
+
+        Returns:
+            The active artifact store.
+        """
+        from zenml.artifact_stores import BaseArtifactStore
+
+        artifact_store_model = Client().active_stack_model.components[
+            StackComponentType.ARTIFACT_STORE
+        ][0]
+        artifact_store = BaseArtifactStore.from_model(artifact_store_model)
+        assert isinstance(artifact_store, BaseArtifactStore)
+
+        return artifact_store
+
     @abstractmethod
     def run(self) -> None:
         """Runs the entrypoint configuration."""

zenml/exceptions.py

@@ -314,3 +314,7 @@ class SecretsStoreNotConfiguredError(NotImplementedError):
 
 class BackupSecretsStoreNotConfiguredError(NotImplementedError):
     """Raised when a backup secrets store is not configured."""
+
+
+class CustomFlavorImportError(ImportError):
+    """Raised when failing to import a custom flavor."""

zenml/integrations/github/code_repositories/github_code_repository.py

@@ -18,7 +18,7 @@ import re
 from typing import List, Optional
 
 import requests
-from github import Github, GithubException
+from github import Consts, Github, GithubException
 from github.Repository import Repository
 
 from zenml.code_repositories import (
@@ -30,6 +30,7 @@ from zenml.code_repositories.base_code_repository import (
 )
 from zenml.code_repositories.git import LocalGitRepositoryContext
 from zenml.logger import get_logger
+from zenml.utils import deprecation_utils
 from zenml.utils.secret_utils import SecretField
 
 logger = get_logger(__name__)
@@ -39,19 +40,24 @@ class GitHubCodeRepositoryConfig(BaseCodeRepositoryConfig):
     """Config for GitHub code repositories.
 
     Args:
-        url: The URL of the GitHub instance.
+        api_url: The GitHub API URL.
         owner: The owner of the repository.
         repository: The name of the repository.
         host: The host of the repository.
         token: The token to access the repository.
     """
 
-    url: Optional[str]
+    api_url: Optional[str] = None
     owner: str
     repository: str
     host: Optional[str] = "github.com"
     token: Optional[str] = SecretField(default=None)
 
+    url: Optional[str] = None
+    _deprecation_validator = deprecation_utils.deprecate_pydantic_attributes(
+        ("url", "api_url")
+    )
+
 
 class GitHubCodeRepository(BaseCodeRepository):
     """GitHub code repository."""
@@ -95,7 +101,7 @@ class GitHubCodeRepository(BaseCodeRepository):
         Raises:
             RuntimeError: If the repository is not public.
         """
-        url = f"https://api.github.com/repos/{owner}/{repo}"
+        url = f"{Consts.DEFAULT_BASE_URL}/repos/{owner}/{repo}"
         response = requests.get(url, timeout=7)
 
         try:
@@ -103,12 +109,15 @@ class GitHubCodeRepository(BaseCodeRepository):
                 pass
             else:
                 raise RuntimeError(
-                    "It is not possible to access this repository as it does not appear to be public."
-                    "Access to private repositories is only possible when a token is provided. Please provide a token and try again"
+                    "It is not possible to access this repository as it does "
+                    "not appear to be public. Access to private repositories "
+                    "is only possible when a token is provided. Please provide "
+                    "a token and try again"
                 )
         except Exception as e:
             raise RuntimeError(
-                f"An error occurred while checking if repository is public: {str(e)}"
+                "An error occurred while checking if repository is public: "
+                f"{str(e)}"
             )
 
     def login(
@@ -120,7 +129,10 @@ class GitHubCodeRepository(BaseCodeRepository):
             RuntimeError: If the login fails.
         """
         try:
-            self._github_session = Github(self.config.token)
+            self._github_session = Github(
+                login_or_token=self.config.token,
+                base_url=self.config.api_url or Consts.DEFAULT_BASE_URL,
+            )
             if self.config.token:
                 user = self._github_session.get_user().login
                 logger.debug(f"Logged in as {user}")

zenml/integrations/gitlab/code_repositories/gitlab_code_repository.py

@@ -31,6 +31,7 @@ from zenml.code_repositories.git.local_git_repository_context import (
     LocalGitRepositoryContext,
 )
 from zenml.logger import get_logger
+from zenml.utils import deprecation_utils
 from zenml.utils.secret_utils import SecretField
 
 logger = get_logger(__name__)
@@ -40,19 +41,24 @@ class GitLabCodeRepositoryConfig(BaseCodeRepositoryConfig):
     """Config for GitLab code repositories.
 
     Args:
-        url: The full URL of the GitLab project.
+        instance_url: The URL of the GitLab instance.
         group: The group of the project.
         project: The name of the GitLab project.
         host: The host of GitLab in case it is self-hosted instance.
         token: The token to access the repository.
     """
 
-    url: Optional[str]
+    instance_url: Optional[str] = None
     group: str
     project: str
     host: Optional[str] = "gitlab.com"
     token: str = SecretField()
 
+    url: Optional[str] = None
+    _deprecation_validator = deprecation_utils.deprecate_pydantic_attributes(
+        ("url", "instance_url")
+    )
+
 
 class GitLabCodeRepository(BaseCodeRepository):
     """GitLab code repository."""
@@ -85,7 +91,7 @@ class GitLabCodeRepository(BaseCodeRepository):
         """
         try:
             self._gitlab_session = Gitlab(
-                self.config.url, private_token=self.config.token
+                url=self.config.instance_url, private_token=self.config.token
             )
             self._gitlab_session.auth()
             user = self._gitlab_session.user or None

zenml/logger.py

@@ -24,6 +24,7 @@ from rich.traceback import install as rich_tb_install
 from zenml.constants import (
     ENABLE_RICH_TRACEBACK,
     ENV_ZENML_LOGGING_COLORS_DISABLED,
+    ENV_ZENML_LOGGING_FORMAT,
     ENV_ZENML_SUPPRESS_LOGS,
     ZENML_LOGGING_VERBOSITY,
     handle_bool_env_var,
@@ -144,6 +145,18 @@ def set_root_verbosity() -> None:
         get_logger(__name__).debug("Logging NOTSET")
 
 
+def get_formatter() -> logging.Formatter:
+    """Get a configured logging formatter.
+
+    Returns:
+        The formatter.
+    """
+    if log_format := os.environ.get(ENV_ZENML_LOGGING_FORMAT, None):
+        return logging.Formatter(fmt=log_format)
+    else:
+        return CustomFormatter()
+
+
 def get_console_handler() -> Any:
     """Get console handler for logging.
 
@@ -151,7 +164,7 @@ def get_console_handler() -> Any:
         A console handler.
     """
     console_handler = logging.StreamHandler(sys.stdout)
-    console_handler.setFormatter(CustomFormatter())
+    console_handler.setFormatter(get_formatter())
     return console_handler
 
 
@@ -179,7 +192,7 @@ def init_logging() -> None:
     set_root_verbosity()
 
     console_handler = logging.StreamHandler(sys.stdout)
-    console_handler.setFormatter(CustomFormatter())
+    console_handler.setFormatter(get_formatter())
     logging.root.addHandler(console_handler)
 
     # Enable logs if environment variable SUPPRESS_ZENML_LOGS is not set to True

zenml/models/v2/base/scoped.py

@@ -241,8 +241,9 @@ class UserScopedFilter(BaseFilter):
         if sort_by == "user":
             column = UserSchema.name
 
-            query = query.join(
-                UserSchema, getattr(table, "user_id") == UserSchema.id
+            query = query.outerjoin(
+                UserSchema,
+                getattr(table, "user_id") == UserSchema.id,
             )
 
             query = query.add_columns(UserSchema.name)

zenml/models/v2/core/pipeline_run.py

@@ -950,31 +950,31 @@ class PipelineRunFilter(WorkspaceScopedTaggableFilter):
         sort_by, operand = self.sorting_params
 
         if sort_by == "pipeline":
-            query = query.join(
+            query = query.outerjoin(
                 PipelineSchema,
                 PipelineRunSchema.pipeline_id == PipelineSchema.id,
             )
             column = PipelineSchema.name
         elif sort_by == "stack":
-            query = query.join(
+            query = query.outerjoin(
                 PipelineDeploymentSchema,
                 PipelineRunSchema.deployment_id == PipelineDeploymentSchema.id,
-            ).join(
+            ).outerjoin(
                 StackSchema,
                 PipelineDeploymentSchema.stack_id == StackSchema.id,
             )
             column = StackSchema.name
         elif sort_by == "model":
-            query = query.join(
+            query = query.outerjoin(
                 ModelVersionSchema,
                 PipelineRunSchema.model_version_id == ModelVersionSchema.id,
-            ).join(
+            ).outerjoin(
                 ModelSchema,
                 ModelVersionSchema.model_id == ModelSchema.id,
             )
             column = ModelSchema.name
         elif sort_by == "model_version":
-            query = query.join(
+            query = query.outerjoin(
                 ModelVersionSchema,
                 PipelineRunSchema.model_version_id == ModelVersionSchema.id,
             )