zenml-nightly 0.84.0.dev20250718__tar.gz → 0.84.0.dev20250719__tar.gz

{zenml_nightly-0.84.0.dev20250718 → zenml_nightly-0.84.0.dev20250719}/PKG-INFO

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

{zenml_nightly-0.84.0.dev20250718 → zenml_nightly-0.84.0.dev20250719}/pyproject.toml

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

zenml_nightly-0.84.0.dev20250719/src/zenml/VERSION

@@ -0,0 +1 @@
+0.84.0.dev20250719

{zenml_nightly-0.84.0.dev20250718 → zenml_nightly-0.84.0.dev20250719}/src/zenml/cli/pipeline.py

@@ -451,6 +451,39 @@ def list_schedules(**kwargs: Any) -> None:
     )
 
 
+@schedule.command("update", help="Update a pipeline schedule.")
+@click.argument("schedule_name_or_id", type=str, required=True)
+@click.option(
+    "--cron-expression",
+    "-c",
+    type=str,
+    required=False,
+    help="The cron expression to update the schedule with.",
+)
+def update_schedule(
+    schedule_name_or_id: str, cron_expression: Optional[str] = None
+) -> None:
+    """Update a pipeline schedule.
+
+    Args:
+        schedule_name_or_id: The name or ID of the schedule to update.
+        cron_expression: The cron expression to update the schedule with.
+    """
+    if not cron_expression:
+        cli_utils.declare("No schedule update requested.")
+        return
+
+    try:
+        Client().update_schedule(
+            name_id_or_prefix=schedule_name_or_id,
+            cron_expression=cron_expression,
+        )
+    except Exception as e:
+        cli_utils.error(str(e))
+    else:
+        cli_utils.declare(f"Updated schedule '{schedule_name_or_id}'.")
+
+
 @schedule.command("delete", help="Delete a pipeline schedule.")
 @click.argument("schedule_name_or_id", type=str, required=True)
 @click.option(

{zenml_nightly-0.84.0.dev20250718 → zenml_nightly-0.84.0.dev20250719}/src/zenml/client.py

@@ -148,6 +148,7 @@ from zenml.models import (
     RunTemplateUpdate,
     ScheduleFilter,
     ScheduleResponse,
+    ScheduleUpdate,
     SecretFilter,
     SecretRequest,
     SecretResponse,
@@ -201,6 +202,7 @@ from zenml.utils.uuid_utils import is_valid_uuid
 
 if TYPE_CHECKING:
     from zenml.metadata.metadata_types import MetadataType, MetadataTypeEnum
+    from zenml.orchestrators import BaseOrchestrator
     from zenml.service_connectors.service_connector import ServiceConnector
     from zenml.services.service import ServiceConfig
     from zenml.stack import Stack
@@ -3806,6 +3808,74 @@ class Client(metaclass=ClientMetaClass):
             hydrate=hydrate,
         )
 
+    def _get_orchestrator_for_schedule(
+        self, schedule: ScheduleResponse
+    ) -> Optional["BaseOrchestrator"]:
+        """Get the orchestrator for a schedule.
+
+        Args:
+            schedule: The schedule to get the orchestrator for.
+
+        Returns:
+            The orchestrator for the schedule.
+        """
+        from zenml.orchestrators import BaseOrchestrator
+
+        if not schedule.orchestrator_id:
+            return None
+
+        try:
+            orchestrator_model = self.get_stack_component(
+                component_type=StackComponentType.ORCHESTRATOR,
+                name_id_or_prefix=schedule.orchestrator_id,
+            )
+        except KeyError:
+            return None
+
+        return cast(
+            BaseOrchestrator, BaseOrchestrator.from_model(orchestrator_model)
+        )
+
+    def update_schedule(
+        self,
+        name_id_or_prefix: Union[str, UUID],
+        cron_expression: Optional[str] = None,
+    ) -> ScheduleResponse:
+        """Update a schedule.
+
+        Args:
+            name_id_or_prefix: The name, id or prefix of the schedule to update.
+            cron_expression: The new cron expression for the schedule.
+
+        Returns:
+            The updated schedule.
+        """
+        schedule = self.get_schedule(
+            name_id_or_prefix=name_id_or_prefix,
+            allow_name_prefix_match=False,
+            project=self.active_project.id,
+        )
+
+        orchestrator = self._get_orchestrator_for_schedule(schedule)
+        if not orchestrator:
+            logger.warning(
+                "Unable to find orchestrator for schedule, skipping update."
+            )
+            return schedule
+        elif not orchestrator.supports_schedule_updates:
+            logger.warning(
+                "Orchestrator does not support schedule updates, skipping "
+                "update."
+            )
+            return schedule
+
+        update = ScheduleUpdate(cron_expression=cron_expression)
+        orchestrator.update_schedule(schedule, update)
+        return self.zen_store.update_schedule(
+            schedule_id=schedule.id,
+            schedule_update=update,
+        )
+
     def delete_schedule(
         self,
         name_id_or_prefix: Union[str, UUID],
@@ -3823,11 +3893,21 @@ class Client(metaclass=ClientMetaClass):
             allow_name_prefix_match=False,
             project=project,
         )
-        logger.warning(
-            f"Deleting schedule '{name_id_or_prefix}'... This will only delete "
-            "the reference of the schedule from ZenML. Please make sure to "
-            "manually stop/delete this schedule in your orchestrator as well!"
-        )
+
+        orchestrator = self._get_orchestrator_for_schedule(schedule)
+        if not orchestrator:
+            logger.warning(
+                "Unable to find orchestrator for schedule. Will only delete "
+                "the schedule reference from ZenML."
+            )
+        elif not orchestrator.supports_schedule_deletion:
+            logger.warning(
+                "Orchestrator does not support schedule deletion. Will only "
+                "delete the schedule reference from ZenML."
+            )
+        else:
+            orchestrator.delete_schedule(schedule)
+
         self.zen_store.delete_schedule(schedule_id=schedule.id)
 
     # ----------------------------- Pipeline runs ------------------------------

{zenml_nightly-0.84.0.dev20250718 → zenml_nightly-0.84.0.dev20250719}/src/zenml/integrations/gcp/step_operators/vertex_step_operator.py

@@ -306,6 +306,14 @@ class VertexStepOperator(BaseStepOperator, GoogleCredentialsMixin):
 
         while response.state not in VERTEX_JOB_STATES_COMPLETED:
             time.sleep(POLLING_INTERVAL_IN_SECONDS)
+            if self.connector_has_expired():
+                logger.warning("Connector has expired. Recreating client...")
+                # This call will refresh the credentials if they expired.
+                credentials, project_id = self._get_authentication()
+                # Recreate the Python API client.
+                client = aiplatform.gapic.JobServiceClient(
+                    credentials=credentials, client_options=client_options
+                )
             try:
                 response = client.get_custom_job(name=job_id)
                 retry_count = 0
@@ -318,12 +326,7 @@ class VertexStepOperator(BaseStepOperator, GoogleCredentialsMixin):
                         f"Error encountered when polling job "
                         f"{job_id}: {err}\nRetrying...",
                     )
-                    # This call will refresh the credentials if they expired.
-                    credentials, project_id = self._get_authentication()
-                    # Recreate the Python API client.
-                    client = aiplatform.gapic.JobServiceClient(
-                        credentials=credentials, client_options=client_options
-                    )
+                    continue
                 else:
                     logger.exception(
                         "Request failed after %s retries.",

{zenml_nightly-0.84.0.dev20250718 → zenml_nightly-0.84.0.dev20250719}/src/zenml/integrations/kubernetes/orchestrators/kubernetes_orchestrator.py

@@ -65,6 +65,7 @@ from zenml.integrations.kubernetes.orchestrators.manifest_utils import (
 from zenml.integrations.kubernetes.pod_settings import KubernetesPodSettings
 from zenml.logger import get_logger
 from zenml.metadata.metadata_types import MetadataType
+from zenml.models.v2.core.schedule import ScheduleUpdate
 from zenml.orchestrators import ContainerizedOrchestrator, SubmissionResult
 from zenml.orchestrators.utils import get_orchestrator_run_name
 from zenml.stack import StackValidator
@@ -74,6 +75,7 @@ if TYPE_CHECKING:
         PipelineDeploymentBase,
         PipelineDeploymentResponse,
         PipelineRunResponse,
+        ScheduleResponse,
     )
     from zenml.stack import Stack
 
@@ -81,6 +83,7 @@ logger = get_logger(__name__)
 
 ENV_ZENML_KUBERNETES_RUN_ID = "ZENML_KUBERNETES_RUN_ID"
 KUBERNETES_SECRET_TOKEN_KEY_NAME = "zenml_api_token"
+KUBERNETES_CRON_JOB_METADATA_KEY = "cron_job_name"
 
 
 class KubernetesOrchestrator(ContainerizedOrchestrator):
@@ -434,6 +437,7 @@ class KubernetesOrchestrator(ContainerizedOrchestrator):
 
         Raises:
             RuntimeError: If a schedule without cron expression is given.
+            Exception: If the orchestrator pod fails to start.
 
         Returns:
             Optional submission result.
@@ -576,7 +580,7 @@ class KubernetesOrchestrator(ContainerizedOrchestrator):
                 labels=orchestrator_pod_labels,
             )
 
-            self._k8s_batch_api.create_namespaced_cron_job(
+            cron_job = self._k8s_batch_api.create_namespaced_cron_job(
                 body=cron_job_manifest,
                 namespace=self.config.kubernetes_namespace,
             )
@@ -584,7 +588,11 @@ class KubernetesOrchestrator(ContainerizedOrchestrator):
                 f"Scheduling Kubernetes run `{pod_name}` with CRON expression "
                 f'`"{cron_expression}"`.'
             )
-            return None
+            return SubmissionResult(
+                metadata={
+                    KUBERNETES_CRON_JOB_METADATA_KEY: cron_job.metadata.name,
+                }
+            )
         else:
             # Create and run the orchestrator pod.
             pod_manifest = build_pod_manifest(
@@ -601,17 +609,34 @@ class KubernetesOrchestrator(ContainerizedOrchestrator):
                 termination_grace_period_seconds=settings.pod_stop_grace_period,
             )
 
-            kube_utils.create_and_wait_for_pod_to_start(
-                core_api=self._k8s_core_api,
-                pod_display_name="Kubernetes orchestrator pod",
-                pod_name=pod_name,
-                pod_manifest=pod_manifest,
-                namespace=self.config.kubernetes_namespace,
-                startup_max_retries=settings.pod_failure_max_retries,
-                startup_failure_delay=settings.pod_failure_retry_delay,
-                startup_failure_backoff=settings.pod_failure_backoff,
-                startup_timeout=settings.pod_startup_timeout,
-            )
+            try:
+                kube_utils.create_and_wait_for_pod_to_start(
+                    core_api=self._k8s_core_api,
+                    pod_display_name="Kubernetes orchestrator pod",
+                    pod_name=pod_name,
+                    pod_manifest=pod_manifest,
+                    namespace=self.config.kubernetes_namespace,
+                    startup_max_retries=settings.pod_failure_max_retries,
+                    startup_failure_delay=settings.pod_failure_retry_delay,
+                    startup_failure_backoff=settings.pod_failure_backoff,
+                    startup_timeout=settings.pod_startup_timeout,
+                )
+            except Exception as e:
+                if self.config.pass_zenml_token_as_secret:
+                    secret_name = self.get_token_secret_name(deployment.id)
+                    try:
+                        kube_utils.delete_secret(
+                            core_api=self._k8s_core_api,
+                            namespace=self.config.kubernetes_namespace,
+                            secret_name=secret_name,
+                        )
+                    except Exception as cleanup_error:
+                        logger.error(
+                            "Error cleaning up secret %s: %s",
+                            secret_name,
+                            cleanup_error,
+                        )
+                raise e
 
             metadata: Dict[str, MetadataType] = {
                 METADATA_ORCHESTRATOR_RUN_ID: pod_name,
@@ -980,3 +1005,48 @@ class KubernetesOrchestrator(ContainerizedOrchestrator):
         return {
             METADATA_ORCHESTRATOR_RUN_ID: self.get_orchestrator_run_id(),
         }
+
+    def update_schedule(
+        self, schedule: "ScheduleResponse", update: ScheduleUpdate
+    ) -> None:
+        """Updates a schedule.
+
+        Args:
+            schedule: The schedule to update.
+            update: The update to apply to the schedule.
+
+        Raises:
+            RuntimeError: If the cron job name is not found.
+        """
+        cron_job_name = schedule.run_metadata.get(
+            KUBERNETES_CRON_JOB_METADATA_KEY
+        )
+        if not cron_job_name:
+            raise RuntimeError("Unable to find cron job name for schedule.")
+
+        if update.cron_expression:
+            self._k8s_batch_api.patch_namespaced_cron_job(
+                name=cron_job_name,
+                namespace=self.config.kubernetes_namespace,
+                body={"spec": {"schedule": update.cron_expression}},
+            )
+
+    def delete_schedule(self, schedule: "ScheduleResponse") -> None:
+        """Deletes a schedule.
+
+        Args:
+            schedule: The schedule to delete.
+
+        Raises:
+            RuntimeError: If the cron job name is not found.
+        """
+        cron_job_name = schedule.run_metadata.get(
+            KUBERNETES_CRON_JOB_METADATA_KEY
+        )
+        if not cron_job_name:
+            raise RuntimeError("Unable to find cron job name for schedule.")
+
+        self._k8s_batch_api.delete_namespaced_cron_job(
+            name=cron_job_name,
+            namespace=self.config.kubernetes_namespace,
+        )

{zenml_nightly-0.84.0.dev20250718 → zenml_nightly-0.84.0.dev20250719}/src/zenml/models/v2/core/schedule.py

@@ -127,6 +127,7 @@ class ScheduleUpdate(BaseUpdate):
     """Update model for schedules."""
 
     name: Optional[str] = None
+    cron_expression: Optional[str] = None
 
 
 # ------------------ Response Model ------------------

{zenml_nightly-0.84.0.dev20250718 → zenml_nightly-0.84.0.dev20250719}/src/zenml/orchestrators/base_orchestrator.py

@@ -50,7 +50,12 @@ from zenml.utils.pydantic_utils import before_validator_handler
 
 if TYPE_CHECKING:
     from zenml.config.step_configurations import Step
-    from zenml.models import PipelineDeploymentResponse, PipelineRunResponse
+    from zenml.models import (
+        PipelineDeploymentResponse,
+        PipelineRunResponse,
+        ScheduleResponse,
+        ScheduleUpdate,
+    )
 
 logger = get_logger(__name__)
 
@@ -522,6 +527,61 @@ class BaseOrchestrator(StackComponent, ABC):
             f"'{self.__class__.__name__}' orchestrator."
         )
 
+    @property
+    def supports_schedule_updates(self) -> bool:
+        """Whether the orchestrator supports updating schedules.
+
+        Returns:
+            Whether the orchestrator supports updating schedules.
+        """
+        return (
+            getattr(self.update_schedule, "__func__", None)
+            is not BaseOrchestrator.update_schedule
+        )
+
+    @property
+    def supports_schedule_deletion(self) -> bool:
+        """Whether the orchestrator supports deleting schedules.
+
+        Returns:
+            Whether the orchestrator supports deleting schedules.
+        """
+        return (
+            getattr(self.delete_schedule, "__func__", None)
+            is not BaseOrchestrator.delete_schedule
+        )
+
+    def update_schedule(
+        self, schedule: "ScheduleResponse", update: "ScheduleUpdate"
+    ) -> None:
+        """Updates a schedule.
+
+        Args:
+            schedule: The schedule to update.
+            update: The update to apply to the schedule.
+
+        Raises:
+            NotImplementedError: If the functionality is not implemented.
+        """
+        raise NotImplementedError(
+            "Schedule updating is not implemented for the "
+            f"'{self.__class__.__name__}' orchestrator."
+        )
+
+    def delete_schedule(self, schedule: "ScheduleResponse") -> None:
+        """Deletes a schedule.
+
+        Args:
+            schedule: The schedule to delete.
+
+        Raises:
+            NotImplementedError: If the functionality is not implemented.
+        """
+        raise NotImplementedError(
+            "Schedule deletion is not implemented for the "
+            f"'{self.__class__.__name__}' orchestrator."
+        )
+
 
 class BaseOrchestratorFlavor(Flavor):
     """Base orchestrator flavor class."""

{zenml_nightly-0.84.0.dev20250718 → zenml_nightly-0.84.0.dev20250719}/src/zenml/zen_server/rbac/models.py

@@ -62,6 +62,7 @@ class ResourceType(StrEnum):
     PIPELINE_RUN = "pipeline_run"
     PIPELINE_DEPLOYMENT = "pipeline_deployment"
     PIPELINE_BUILD = "pipeline_build"
+    SCHEDULE = "schedule"
     RUN_TEMPLATE = "run_template"
     SERVICE = "service"
     RUN_METADATA = "run_metadata"

{zenml_nightly-0.84.0.dev20250718 → zenml_nightly-0.84.0.dev20250719}/src/zenml/zen_server/rbac/utils.py

@@ -461,6 +461,8 @@ def get_resource_type_for_model(
         RunMetadataRequest,
         RunTemplateRequest,
         RunTemplateResponse,
+        ScheduleRequest,
+        ScheduleResponse,
         SecretRequest,
         SecretResponse,
         ServiceAccountRequest,
@@ -512,6 +514,8 @@ def get_resource_type_for_model(
         RunMetadataRequest: ResourceType.RUN_METADATA,
         RunTemplateRequest: ResourceType.RUN_TEMPLATE,
         RunTemplateResponse: ResourceType.RUN_TEMPLATE,
+        ScheduleRequest: ResourceType.SCHEDULE,
+        ScheduleResponse: ResourceType.SCHEDULE,
         SecretRequest: ResourceType.SECRET,
         SecretResponse: ResourceType.SECRET,
         ServiceAccountRequest: ResourceType.SERVICE_ACCOUNT,
@@ -649,6 +653,7 @@ def get_schema_for_resource_type(
         PipelineSchema,
         RunMetadataSchema,
         RunTemplateSchema,
+        ScheduleSchema,
         SecretSchema,
         ServiceConnectorSchema,
         ServiceSchema,
@@ -681,6 +686,7 @@ def get_schema_for_resource_type(
         ResourceType.PIPELINE_BUILD: PipelineBuildSchema,
         ResourceType.RUN_TEMPLATE: RunTemplateSchema,
         ResourceType.RUN_METADATA: RunMetadataSchema,
+        ResourceType.SCHEDULE: ScheduleSchema,
         # ResourceType.USER: UserSchema,
         ResourceType.ACTION: ActionSchema,
         ResourceType.EVENT_SOURCE: EventSourceSchema,

{zenml_nightly-0.84.0.dev20250718 → zenml_nightly-0.84.0.dev20250719}/src/zenml/zen_server/routers/schedule_endpoints.py

@@ -30,7 +30,12 @@ from zenml.zen_server.auth import AuthContext, authorize
 from zenml.zen_server.exceptions import error_response
 from zenml.zen_server.rbac.endpoint_utils import (
     verify_permissions_and_create_entity,
+    verify_permissions_and_delete_entity,
+    verify_permissions_and_get_entity,
+    verify_permissions_and_list_entities,
+    verify_permissions_and_update_entity,
 )
+from zenml.zen_server.rbac.models import ResourceType
 from zenml.zen_server.routers.projects_endpoints import workspace_router
 from zenml.zen_server.utils import (
     async_fastapi_endpoint_wrapper,
@@ -77,8 +82,6 @@ def create_schedule(
         project = zen_store().get_project(project_name_or_id)
         schedule.project = project.id
 
-    # NOTE: no RBAC is enforced currently for schedules, but we're
-    # keeping the RBAC checks here for consistency
     return verify_permissions_and_create_entity(
         request_model=schedule,
         create_method=zen_store().create_schedule,
@@ -121,8 +124,10 @@ def list_schedules(
     if project_name_or_id:
         schedule_filter_model.project = project_name_or_id
 
-    return zen_store().list_schedules(
-        schedule_filter_model=schedule_filter_model,
+    return verify_permissions_and_list_entities(
+        filter_model=schedule_filter_model,
+        resource_type=ResourceType.SCHEDULE,
+        list_method=zen_store().list_schedules,
         hydrate=hydrate,
     )
 
@@ -147,8 +152,9 @@ def get_schedule(
     Returns:
         A specific schedule object.
     """
-    return zen_store().get_schedule(
-        schedule_id=schedule_id,
+    return verify_permissions_and_get_entity(
+        id=schedule_id,
+        get_method=zen_store().get_schedule,
         hydrate=hydrate,
     )
 
@@ -172,9 +178,11 @@ def update_schedule(
     Returns:
         The updated schedule object.
     """
-    return zen_store().update_schedule(
-        schedule_id=schedule_id,
-        schedule_update=schedule_update,
+    return verify_permissions_and_update_entity(
+        id=schedule_id,
+        update_model=schedule_update,
+        get_method=zen_store().get_schedule,
+        update_method=zen_store().update_schedule,
     )
 
 
@@ -192,4 +200,8 @@ def delete_schedule(
     Args:
         schedule_id: ID of the schedule to delete.
     """
-    zen_store().delete_schedule(schedule_id=schedule_id)
+    verify_permissions_and_delete_entity(
+        id=schedule_id,
+        get_method=zen_store().get_schedule,
+        delete_method=zen_store().delete_schedule,
+    )

{zenml_nightly-0.84.0.dev20250718 → zenml_nightly-0.84.0.dev20250719}/src/zenml/zen_stores/schemas/schedule_schema.py

@@ -202,6 +202,9 @@ class ScheduleSchema(NamedSchema, RunMetadataInterface, table=True):
         if schedule_update.name is not None:
             self.name = schedule_update.name
 
+        if schedule_update.cron_expression:
+            self.cron_expression = schedule_update.cron_expression
+
         self.updated = utc_now()
         return self
 

zenml_nightly-0.84.0.dev20250718/src/zenml/VERSION

@@ -1 +0,0 @@
-0.84.0.dev20250718