add termination grace period (#1379)

Author: cosmicBboy

examples/advanced/abort_callback.py

@@ -1,7 +1,7 @@
 # /// script
 # requires-python = "==3.12"
 # dependencies = [
-#    "kubernetes",
+#    "flyte",
 # ]
 # ///
 """
@@ -22,33 +22,26 @@
    task. (Plain ``def`` tasks run in a worker thread, where ``signal.signal`` is not
    allowed -- for those, register the handler at module import time instead.)
 
-2. Give the callback enough time to finish. The default grace period is 30s, which
-   may be too short to upload a large checkpoint, so we raise
-   ``termination_grace_period_seconds`` on the pod template.
+2. Give the callback enough time to finish. The default grace period is 30s, which may
+   be too short to upload a large checkpoint, so we raise it with the
+   ``flyte.PodTemplate().with_termination_grace_period(...)`` helper. It accepts an int
+   number of seconds or a ``timedelta`` and sets ``terminationGracePeriodSeconds`` on the
+   pod spec -- no need to import the ``kubernetes`` package.
 """
 
 from __future__ import annotations
 
 import asyncio
 import signal
-
-from kubernetes.client import V1Container, V1PodSpec
+from datetime import timedelta
 
 import flyte
 
-# Raise the pod termination grace period so the abort callback has time to flush its
-# side effects (e.g. upload a checkpoint) before Kubernetes escalates to SIGKILL.
-pod_template = flyte.PodTemplate(
-    primary_container_name="primary",
-    pod_spec=V1PodSpec(
-        containers=[V1Container(name="primary")],
-        termination_grace_period_seconds=600,  # 10 minutes; default is 30s
-    ),
-)
-
 env = flyte.TaskEnvironment(
     name="abort_callback",
-    pod_template=pod_template,
+    # Give the abort callback time to flush its side effects (e.g. upload a checkpoint)
+    # before Kubernetes escalates to SIGKILL. Default is 30s.
+    pod_template=flyte.PodTemplate().with_termination_grace_period(timedelta(minutes=10)),
     image=flyte.Image.from_uv_script(__file__, name="flyte"),
 )
 
@@ -62,7 +55,7 @@ async def train(n_steps: int = 1000, step_seconds: float = 2.0) -> str:
     # Resume from the previous attempt's checkpoint, if any.
     prev = await checkpoint.load()
     start = int(prev.read_bytes().decode()) if prev is not None else 0
-    print(f"Starting training at step {start}")
+    flyte.logger.info(f"Starting training at step {start}")
 
     aborted = asyncio.Event()
 
@@ -71,45 +64,46 @@ def on_abort() -> None:
         # Keep the callback itself lightweight: flip a flag and let the loop below
         # perform the actual side effects. (If you'd rather do the work directly in
         # the callback, schedule a coroutine with `asyncio.ensure_future(...)`.)
-        print("Abort received -- will checkpoint current progress and exit.")
+        flyte.logger.info("Abort received -- will checkpoint current progress and exit.")
         aborted.set()
 
     # An `async def` task runs on the main thread's event loop, so we can register the
     # handler from inside the task. Aborting the run deletes the pod, and Kubernetes
     # sends SIGTERM to this process, which triggers `on_abort`.
     try:
         asyncio.get_running_loop().add_signal_handler(signal.SIGTERM, on_abort)
+        flyte.logger.info("Registered SIGTERM handler for abort callback.")
     except (NotImplementedError, RuntimeError):
         # e.g. running locally off the main thread; abort handling only applies in-cluster.
-        print("Could not register a SIGTERM handler here; continuing without abort handling.")
+        flyte.logger.info("Could not register a SIGTERM handler here; continuing without abort handling.")
 
     step = start
     for step in range(start, n_steps):
         if aborted.is_set():
             # ---- extra code to run on abort; side effects live here ----
             await checkpoint.save(f"{step}".encode())
-            print(f"Checkpoint saved at step {step}. Exiting due to abort.")
+            flyte.logger.info(f"Checkpoint saved at step {step}. Exiting due to abort.")
             return f"aborted at step {step}"
         # ... real training work would go here ...
         await asyncio.sleep(step_seconds)
         await checkpoint.save(f"{step + 1}".encode())
-        print(f"Completed step {step + 1}/{n_steps}")
+        flyte.logger.info(f"Completed step {step + 1}/{n_steps}")
 
     return f"completed all {n_steps} steps"
 
 
 if __name__ == "__main__":
     import time
 
-    flyte.init_from_config()
-    run = flyte.run(train, n_steps=1000, step_seconds=2.0)
+    flyte.init_from_config(log_level="INFO")
+    run = flyte.run(train, n_steps=100, step_seconds=2.0)
     print(run.url)
 
     # Wait until the task is actually running, let it take a few steps, then abort it
     # to trigger the on-abort checkpoint callback. (You could also just click "Abort"
     # in the UI instead of the code below.)
     run.wait(wait_for="running")
-    time.sleep(20)
+    time.sleep(10)
     print("Aborting the run to trigger the on-abort checkpoint callback...")
     run.abort()
     print("Aborted. Inspect the task logs -- it should have saved a checkpoint before exiting.")

src/flyte/_environment.py

@@ -49,7 +49,9 @@ class Environment:
     :param env_vars: Environment variables as `dict[str, str]`.
     :param secrets: Secrets to inject into the environment.
     :param pod_template: Kubernetes pod template as a string reference to a
-        named template or a `PodTemplate` object.
+        named template or a `PodTemplate` object. To set a termination grace
+        period without depending on the `kubernetes` package, use
+        `flyte.PodTemplate().with_termination_grace_period(...)`.
     :param description: Human-readable description (max 255 characters).
     :param interruptible: Whether the environment can be scheduled on
         spot/preemptible instances.

src/flyte/_pod.py

@@ -2,13 +2,21 @@
 
 import copy
 from dataclasses import dataclass
-from typing import TYPE_CHECKING, Dict, Optional
+from datetime import timedelta
+from typing import TYPE_CHECKING, Dict, Optional, Union
 
 if TYPE_CHECKING:
     from flyteidl2.core.tasks_pb2 import K8sPod
     from kubernetes.client import V1Container, V1PodSpec
 
 
+# A task's Kubernetes termination grace period, accepted as an ``int`` number of seconds or a
+# ``timedelta``. This is the time Kubernetes waits after sending SIGTERM (e.g. when a run is
+# aborted, which deletes the pod) before escalating to SIGKILL, and maps to
+# ``V1PodSpec.termination_grace_period_seconds``.
+TerminationGracePeriod = Union[int, timedelta]
+
+
 _PRIMARY_CONTAINER_NAME_FIELD = "primary_container_name"
 _PRIMARY_CONTAINER_DEFAULT_NAME = "primary"
 
@@ -187,6 +195,37 @@ def allow_nested_sandboxing(self) -> PodTemplate:
         """
         return _apply_sandboxing(self)
 
+    def with_termination_grace_period(self, termination_grace_period: TerminationGracePeriod) -> PodTemplate:
+        """
+        Return a copy of this template with Kubernetes' ``terminationGracePeriodSeconds`` set.
+
+        This is the time Kubernetes waits after sending SIGTERM (e.g. when a run is aborted,
+        which deletes the pod) before escalating to SIGKILL — raise it to give a task time to
+        checkpoint or otherwise clean up on abort. Accepts an ``int`` number of seconds or a
+        ``timedelta``.
+
+        Because the primary container is synthesized when the template has no pod spec, you can
+        set a grace period without depending on the ``kubernetes`` package::
+
+            env = flyte.TaskEnvironment(
+                name="train",
+                pod_template=flyte.PodTemplate().with_termination_grace_period(timedelta(minutes=10)),
+            )
+
+        The original template is never mutated; existing containers, volumes, labels, and other
+        pod-spec fields are preserved. Re-applying overwrites the previously set value.
+
+        :param termination_grace_period: Grace period as an ``int`` (seconds) or ``timedelta``.
+            Must be non-negative.
+        """
+        seconds = _normalize_termination_grace_period(termination_grace_period)
+        if seconds is None:
+            raise TypeError("termination_grace_period is required and must be an int (seconds) or timedelta.")
+        pt = _clone_with_primary(self)
+        assert pt.pod_spec is not None  # _clone_with_primary guarantees a pod spec
+        pt.pod_spec.termination_grace_period_seconds = seconds
+        return pt
+
     def to_k8s_pod(self) -> "K8sPod":
         from flyteidl2.core.tasks_pb2 import K8sObjectMetadata, K8sPod
         from kubernetes.client import ApiClient, V1PodSpec
@@ -216,6 +255,29 @@ def _clone_with_primary(pt: PodTemplate) -> PodTemplate:
     return pt
 
 
+def _normalize_termination_grace_period(value: Optional[TerminationGracePeriod]) -> Optional[int]:
+    """Normalize an ``int`` (seconds) or ``timedelta`` grace period to whole seconds.
+
+    Returns ``None`` when ``value`` is ``None``. Raises ``TypeError`` for other types and
+    ``ValueError`` for negative durations. A ``timedelta`` is truncated to whole seconds, since
+    Kubernetes' ``terminationGracePeriodSeconds`` is integer-valued.
+    """
+    if value is None:
+        return None
+    # bool is an int subclass; reject it explicitly to avoid True -> 1 surprises.
+    if isinstance(value, bool):
+        raise TypeError("termination_grace_period must be an int (seconds) or timedelta, not bool.")
+    if isinstance(value, timedelta):
+        seconds = int(value.total_seconds())
+    elif isinstance(value, int):
+        seconds = value
+    else:
+        raise TypeError(f"termination_grace_period must be an int (seconds) or timedelta, got {type(value).__name__}.")
+    if seconds < 0:
+        raise ValueError(f"termination_grace_period must be non-negative, got {seconds} seconds.")
+    return seconds
+
+
 def _get_primary_container(pt: PodTemplate) -> "V1Container":
     """Return the primary container of a template prepared by ``_clone_with_primary``."""
     assert pt.pod_spec is not None  # _clone_with_primary guarantees it

tests/user_api/test_termination_grace_period.py

@@ -0,0 +1,140 @@
+"""Tests for ``PodTemplate.with_termination_grace_period`` — the helper that sets
+``terminationGracePeriodSeconds`` on a pod template without needing the ``kubernetes`` package.
+"""
+
+import copy
+import pathlib
+from datetime import timedelta
+
+import pytest
+from kubernetes.client import V1Container, V1PodSpec
+
+import flyte
+from flyte._internal.runtime.task_serde import get_proto_task
+from flyte._pod import PodTemplate
+from flyte.models import SerializationContext
+
+
+def _grace_seconds(task):
+    """Serialize a task and return the pod spec's terminationGracePeriodSeconds (or None)."""
+    ctx = SerializationContext(project="p", domain="d", org="o", version="v1", root_dir=pathlib.Path.cwd())
+    proto = get_proto_task(task, ctx)
+    if not proto.HasField("k8s_pod"):
+        return None
+    pod_spec = proto.k8s_pod.pod_spec  # google.protobuf.Struct: supports `in` / `[]`, not `.get`
+    return pod_spec["terminationGracePeriodSeconds"] if "terminationGracePeriodSeconds" in pod_spec else None
+
+
+# --------------------------------------------------------------------------- #
+# PodTemplate.with_termination_grace_period
+# --------------------------------------------------------------------------- #
+
+
+class TestWithTerminationGracePeriod:
+    def test_synthesizes_pod_spec_with_primary_no_kubernetes_needed(self):
+        # Starting from a bare PodTemplate(), the helper builds the pod spec + primary container
+        # internally, so the caller never touches kubernetes.client.
+        pt = PodTemplate().with_termination_grace_period(42)
+        assert pt.pod_spec.termination_grace_period_seconds == 42
+        assert any(c.name == "primary" for c in pt.pod_spec.containers)
+
+    def test_int_seconds(self):
+        assert PodTemplate().with_termination_grace_period(30).pod_spec.termination_grace_period_seconds == 30
+
+    def test_timedelta(self):
+        pt = PodTemplate().with_termination_grace_period(timedelta(minutes=2))
+        assert pt.pod_spec.termination_grace_period_seconds == 120
+
+    def test_zero_is_kept(self):
+        # 0 is a meaningful k8s value (immediate kill), not "unset".
+        assert PodTemplate().with_termination_grace_period(0).pod_spec.termination_grace_period_seconds == 0
+
+    def test_custom_primary_container_name(self):
+        pt = PodTemplate(primary_container_name="worker").with_termination_grace_period(10)
+        assert any(c.name == "worker" for c in pt.pod_spec.containers)
+        assert pt.pod_spec.termination_grace_period_seconds == 10
+
+    def test_preserves_existing_pod_spec_fields(self):
+        base = PodTemplate(pod_spec=V1PodSpec(containers=[V1Container(name="primary", image="img")], hostname="h"))
+        pt = base.with_termination_grace_period(99)
+        assert pt.pod_spec.termination_grace_period_seconds == 99
+        assert pt.pod_spec.containers[0].image == "img"
+        assert pt.pod_spec.hostname == "h"
+
+    def test_does_not_mutate_original(self):
+        base = PodTemplate(pod_spec=V1PodSpec(containers=[V1Container(name="primary")]))
+        snapshot = copy.deepcopy(base)
+        base.with_termination_grace_period(99)
+        assert base == snapshot
+
+    def test_reapplying_overwrites(self):
+        pt = PodTemplate().with_termination_grace_period(30).with_termination_grace_period(60)
+        assert pt.pod_spec.termination_grace_period_seconds == 60
+
+    def test_composes_with_capability_helpers(self):
+        pt = PodTemplate().allow_fuse().with_termination_grace_period(60)
+        primary = next(c for c in pt.pod_spec.containers if c.name == "primary")
+        assert pt.pod_spec.termination_grace_period_seconds == 60
+        assert primary.resources.requests["smarter-devices/fuse"] == "1"  # allow_fuse survived
+
+    def test_bool_rejected(self):
+        with pytest.raises(TypeError, match="not bool"):
+            PodTemplate().with_termination_grace_period(True)
+
+    def test_negative_rejected(self):
+        with pytest.raises(ValueError, match="non-negative"):
+            PodTemplate().with_termination_grace_period(-1)
+
+    def test_wrong_type_rejected(self):
+        with pytest.raises(TypeError):
+            PodTemplate().with_termination_grace_period("30")
+
+    def test_none_rejected(self):
+        with pytest.raises(TypeError):
+            PodTemplate().with_termination_grace_period(None)
+
+
+# --------------------------------------------------------------------------- #
+# End-to-end: the helper output flows through task serialization
+# --------------------------------------------------------------------------- #
+
+
+env = flyte.TaskEnvironment(
+    name="tgp_env",
+    pod_template=PodTemplate().with_termination_grace_period(timedelta(minutes=10)),
+)
+
+
+@env.task
+async def env_task() -> int:
+    return 1
+
+
+env_plain = flyte.TaskEnvironment(name="tgp_plain")
+
+
+@env_plain.task
+async def plain_task() -> int:
+    return 1
+
+
+@env_plain.task(pod_template=PodTemplate().with_termination_grace_period(300))
+async def decorated_task() -> int:
+    return 1
+
+
+class TestSerialization:
+    def test_environment_pod_template(self):
+        assert _grace_seconds(env_task) == 600
+
+    def test_no_pod_template_uses_container_path(self):
+        assert _grace_seconds(plain_task) is None
+        ctx = SerializationContext(project="p", domain="d", org="o", version="v1", root_dir=pathlib.Path.cwd())
+        assert get_proto_task(plain_task, ctx).HasField("container")
+
+    def test_decorator_pod_template(self):
+        assert _grace_seconds(decorated_task) == 300
+
+    def test_override_pod_template(self):
+        overridden = plain_task.override(pod_template=PodTemplate().with_termination_grace_period(15))
+        assert _grace_seconds(overridden) == 15