Stop & Resume schedule on Kubernetes

Author: Json-Andriopoulos

src/zenml/cli/pipeline.py

@@ -677,23 +677,64 @@ def list_schedules(
     required=False,
     help="The cron expression to update the schedule with.",
 )
+@click.option(
+    "--activate",
+    "-a",
+    type=bool,
+    required=False,
+    is_flag=True,
+    help="Activate a pipeline schedule.",
+)
+@click.option(
+    "--deactivate",
+    "-d",
+    type=bool,
+    required=False,
+    is_flag=True,
+    help="Deactivate a pipeline schedule.",
+)
 def update_schedule(
-    schedule_name_or_id: str, cron_expression: Optional[str] = None
+    schedule_name_or_id: str,
+    cron_expression: str | None = None,
+    activate: bool | None = None,
+    deactivate: bool | None = 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.
+        activate: Activate a pipeline schedule.
+        deactivate: Deactivate a pipeline schedule.
     """
-    if not cron_expression:
+    options = [
+        cron_expression,
+        activate,
+        deactivate,
+    ]
+
+    if not any(options):
         cli_utils.declare("No schedule update requested.")
         return
 
+    if activate and deactivate:
+        cli_utils.declare(
+            "You can not both activate and deactivate at the same time."
+        )
+        return
+
+    if activate:
+        is_active = True
+    elif deactivate:
+        is_active = False
+    else:
+        is_active = None
+
     try:
         Client().update_schedule(
             name_id_or_prefix=schedule_name_or_id,
             cron_expression=cron_expression,
+            active=is_active,
         )
     except Exception as e:
         cli_utils.exception(e)

src/zenml/client.py

@@ -4613,12 +4613,14 @@ def update_schedule(
         self,
         name_id_or_prefix: Union[str, UUID],
         cron_expression: Optional[str] = None,
+        active: bool | None = 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.
+            active: Active status flag for the schedule.
 
         Returns:
             The updated schedule.
@@ -4642,7 +4644,13 @@ def update_schedule(
             )
             return schedule
 
-        update = ScheduleUpdate(cron_expression=cron_expression)
+        if schedule.active == active:
+            logger.warning(
+                f"Schedule active value is already {active}, skipping update."
+            )
+            return schedule
+
+        update = ScheduleUpdate(cron_expression=cron_expression, active=active)
         orchestrator.update_schedule(schedule, update)
         return self.zen_store.update_schedule(
             schedule_id=schedule.id,

src/zenml/integrations/kubernetes/orchestrators/kubernetes_orchestrator.py

@@ -1217,11 +1217,19 @@ def update_schedule(
         if not cron_job_name:
             raise RuntimeError("Unable to find cron job name for schedule.")
 
+        spec = {}
+
+        if update.cron_expression:
+            spec["schedule"] = update.cron_expression
+
+        if update.active is not None:
+            spec["suspend"] = not update.active
+
         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}},
+                body={"spec": spec},
             )
 
     def delete_schedule(self, schedule: "ScheduleResponse") -> None:

src/zenml/models/v2/core/schedule.py

@@ -128,6 +128,7 @@ class ScheduleUpdate(BaseUpdate):
 
     name: Optional[str] = None
     cron_expression: Optional[str] = None
+    active: bool | None = None
 
 
 # ------------------ Response Model ------------------

src/zenml/zen_stores/schemas/schedule_schema.py

@@ -205,6 +205,9 @@ def update(self, schedule_update: ScheduleUpdate) -> "ScheduleSchema":
         if schedule_update.cron_expression:
             self.cron_expression = schedule_update.cron_expression
 
+        if schedule_update.active is not None:
+            self.active = schedule_update.active
+
         self.updated = utc_now()
         return self