zenml-nightly 0.73.0.dev20250126__tar.gz → 0.73.0.dev20250128__tar.gz

{zenml_nightly-0.73.0.dev20250126 → zenml_nightly-0.73.0.dev20250128}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: zenml-nightly
-Version: 0.73.0.dev20250126
+Version: 0.73.0.dev20250128
 Summary: ZenML: Write production-ready ML code.
 License: Apache-2.0
 Keywords: machine learning,production,pipeline,mlops,devops

{zenml_nightly-0.73.0.dev20250126 → zenml_nightly-0.73.0.dev20250128}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "zenml-nightly"
-version = "0.73.0.dev20250126"
+version = "0.73.0.dev20250128"
 packages = [{ include = "zenml", from = "src" }]
 description = "ZenML: Write production-ready ML code."
 authors = ["ZenML GmbH <info@zenml.io>"]

zenml_nightly-0.73.0.dev20250128/src/zenml/VERSION

@@ -0,0 +1 @@
+0.73.0.dev20250128

{zenml_nightly-0.73.0.dev20250126 → zenml_nightly-0.73.0.dev20250128}/src/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_nightly-0.73.0.dev20250126 → zenml_nightly-0.73.0.dev20250128}/src/zenml/client.py

@@ -4952,25 +4952,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 +4973,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 +5103,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 +5113,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 +5127,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_nightly-0.73.0.dev20250126 → zenml_nightly-0.73.0.dev20250128}/src/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_nightly-0.73.0.dev20250126 → zenml_nightly-0.73.0.dev20250128}/src/zenml/entrypoints/base_entrypoint_configuration.py

@@ -214,12 +214,14 @@ class BaseEntrypointConfiguration(ABC):
         if not should_download_code:
             return
 
-        if code_reference := deployment.code_reference:
+        if code_path := deployment.code_path:
+            code_utils.download_code_from_artifact_store(code_path=code_path)
+        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."

{zenml_nightly-0.73.0.dev20250126 → zenml_nightly-0.73.0.dev20250128}/src/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_nightly-0.73.0.dev20250126 → zenml_nightly-0.73.0.dev20250128}/src/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_nightly-0.73.0.dev20250126 → zenml_nightly-0.73.0.dev20250128}/src/zenml/models/v2/core/schedule.py

@@ -13,11 +13,11 @@
 #  permissions and limitations under the License.
 """Models representing schedules."""
 
-import datetime
+from datetime import datetime, timedelta, timezone
 from typing import Dict, Optional, Union
 from uuid import UUID
 
-from pydantic import Field, model_validator
+from pydantic import Field, field_validator, model_validator
 
 from zenml.constants import STR_FIELD_MAX_LENGTH
 from zenml.logger import get_logger
@@ -46,15 +46,36 @@ class ScheduleRequest(WorkspaceScopedRequest):
     active: bool
 
     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
 
     orchestrator_id: Optional[UUID]
     pipeline_id: Optional[UUID]
 
+    @field_validator(
+        "start_time", "end_time", "run_once_start_time", mode="after"
+    )
+    @classmethod
+    def _ensure_tzunaware_utc(
+        cls, value: Optional[datetime]
+    ) -> Optional[datetime]:
+        """Ensures that all datetimes are timezone unaware and in UTC time.
+
+        Args:
+            value: The datetime.
+
+        Returns:
+            The datetime in UTC time without timezone.
+        """
+        if value and value.tzinfo:
+            value = value.astimezone(timezone.utc)
+            value = value.replace(tzinfo=None)
+
+        return value
+
     @model_validator(mode="after")
     def _ensure_cron_or_periodic_schedule_configured(
         self,
@@ -108,11 +129,11 @@ class ScheduleUpdate(BaseUpdate):
     name: Optional[str] = None
     active: Optional[bool] = 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: Optional[bool] = None
-    run_once_start_time: Optional[datetime.datetime] = None
+    run_once_start_time: Optional[datetime] = None
     orchestrator_id: Optional[UUID] = None
     pipeline_id: Optional[UUID] = None
 
@@ -125,11 +146,11 @@ class ScheduleResponseBody(WorkspaceScopedResponseBody):
 
     active: bool
     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
 
 
 class ScheduleResponseMetadata(WorkspaceScopedResponseMetadata):
@@ -217,7 +238,7 @@ class ScheduleResponse(
         return self.get_body().cron_expression
 
     @property
-    def start_time(self) -> Optional[datetime.datetime]:
+    def start_time(self) -> Optional[datetime]:
         """The `start_time` property.
 
         Returns:
@@ -226,7 +247,7 @@ class ScheduleResponse(
         return self.get_body().start_time
 
     @property
-    def end_time(self) -> Optional[datetime.datetime]:
+    def end_time(self) -> Optional[datetime]:
         """The `end_time` property.
 
         Returns:
@@ -235,7 +256,7 @@ class ScheduleResponse(
         return self.get_body().end_time
 
     @property
-    def run_once_start_time(self) -> Optional[datetime.datetime]:
+    def run_once_start_time(self) -> Optional[datetime]:
         """The `run_once_start_time` property.
 
         Returns:
@@ -244,7 +265,7 @@ class ScheduleResponse(
         return self.get_body().run_once_start_time
 
     @property
-    def interval_second(self) -> Optional[datetime.timedelta]:
+    def interval_second(self) -> Optional[timedelta]:
         """The `interval_second` property.
 
         Returns:
@@ -313,10 +334,10 @@ class ScheduleFilter(WorkspaceScopedFilter):
         default=None,
         description="The cron expression, describing the schedule",
     )
-    start_time: Optional[Union[datetime.datetime, str]] = Field(
+    start_time: Optional[Union[datetime, str]] = Field(
         default=None, description="Start time", union_mode="left_to_right"
     )
-    end_time: Optional[Union[datetime.datetime, str]] = Field(
+    end_time: Optional[Union[datetime, str]] = Field(
         default=None, description="End time", union_mode="left_to_right"
     )
     interval_second: Optional[Optional[float]] = Field(
@@ -332,7 +353,7 @@ class ScheduleFilter(WorkspaceScopedFilter):
         default=None,
         description="Name of the schedule",
     )
-    run_once_start_time: Optional[Union[datetime.datetime, str]] = Field(
+    run_once_start_time: Optional[Union[datetime, str]] = Field(
         default=None,
         description="The time at which the schedule should run once",
         union_mode="left_to_right",