Repository: PrefectHQ/prefect

Issue:
Add `ttlSecondsAfterFinished` support to `KubernetesJob`
### First check

- [X] I added a descriptive title to this issue.
- [X] I used the GitHub search to find a similar request and didn't find it.
- [X] I searched the Prefect documentation for this feature.

### Prefect Version

2.x

### Describe the current behavior

No jobs created by the `KubernetesJob` infrastructure type are cleaned up yet.

### Describe the proposed behavior

Infrastructure blocks should generally allow automated cleanup of created resources.

Adding `ttlSecondsAfterFinished` to the job `spec` will facilitate automatic removal of completed jobs. This is especially nice as we do not need to manage removal ourselves.

Users can already add this via a `customizations`. Separately, this feels important enough to add a top-level shortcut for. Perhaps as `finished_job_ttl: int`.

### Example Use

```python
KubernetesJob(finished_job_ttl=300)
```

### Additional context

Replaces https://github.com/PrefectHQ/prefect/issues/5857


Discussion:
@madkinsz can I take this up?
Yeah definitely! Let me know if you have any questions. One minor concern is that I'm not sure which Kubernetes versions support this.; it'd be good to include that as a note in the pull request and docstring if it's not the current lts version.
@madkinsz Kubernetes v1.23 [stable] according this doc: https://kubernetes.io/docs/concepts/workloads/controllers/ttlafterfinished/


@hallenmaia great thanks! Looks like this is the ideal hook for us then.
@madkinsz can you guide me through what would be a good step to start
Sure, you'll want to:

- Add a `job_ttl_finished_seconds` field like https://github.com/PrefectHQ/prefect/blob/c2a76aa57fee8326de0883f1ed708878072eeba9/src/prefect/infrastructure/kubernetes.py#L102-L104 
    - We'll want it to be `Optional[int]` to allow it to be null. When null, the field should be excluded from the template.
- Add a handler to set it in the job specification at https://github.com/PrefectHQ/prefect/blob/c2a76aa57fee8326de0883f1ed708878072eeba9/src/prefect/infrastructure/kubernetes.py#L262-L271
- You may need to add a value for it in the base template to allow it to be patched, but probably not https://github.com/PrefectHQ/prefect/blob/c2a76aa57fee8326de0883f1ed708878072eeba9/src/prefect/infrastructure/kubernetes.py#L165-L172
- Add a test demonstrating that the value is set, similar to https://github.com/PrefectHQ/prefect/blob/c2a76aa57fee8326de0883f1ed708878072eeba9/tests/infrastructure/test_kubernetes_job.py#L317-L331
- Add a test demonstrating that the field is gone when `None`.
@meehawk I have this in my dev env. Do you mind if I send the PR?
@madkinsz Can I contribute to this?

Resolving patch:
diff --git a/src/prefect/infrastructure/kubernetes.py b/src/prefect/infrastructure/kubernetes.py
--- a/src/prefect/infrastructure/kubernetes.py
+++ b/src/prefect/infrastructure/kubernetes.py
@@ -64,6 +64,9 @@ class KubernetesJob(Infrastructure):
         pod_watch_timeout_seconds: Number of seconds to watch for pod creation before timing out (default 5).
         service_account_name: An optional string specifying which Kubernetes service account to use.
         stream_output: If set, stream output from the job to local standard output.
+        finished_job_ttl: The number of seconds to retain jobs after completion. If set, finished jobs will
+            be cleaned up by Kubernetes after the given delay. If None (default), jobs will need to be
+            manually removed.
     """
 
     type: Literal["kubernetes-job"] = Field(
@@ -109,6 +112,10 @@ class KubernetesJob(Infrastructure):
         True,
         description="If set, output will be streamed from the job to local standard output.",
     )
+    finished_job_ttl: Optional[int] = Field(
+        None,
+        description="The number of seconds to retain jobs after completion. If set, finished jobs will be cleaned up by Kubernetes after the given delay. If None (default), jobs will need to be manually removed.",
+    )
 
     # internal-use only right now
     _api_dns_name: Optional[str] = None  # Replaces 'localhost' in API URL
@@ -315,6 +322,15 @@ def _shortcut_customizations(self) -> JsonPatch:
                 }
             )
 
+        if self.finished_job_ttl is not None:
+            shortcuts.append(
+                {
+                    "op": "add",
+                    "path": "/spec/ttlSecondsAfterFinished",
+                    "value": self.finished_job_ttl,
+                }
+            )
+
         if self.command:
             shortcuts.append(
                 {


Test patch:
diff --git a/tests/blocks/checksums.json b/tests/blocks/checksums.json
--- a/tests/blocks/checksums.json
+++ b/tests/blocks/checksums.json
@@ -14,6 +14,6 @@
     "kubernetes-cluster-config": "sha256:90d421e948bfbe4cdc98b124995f0edd0f84b0837549ad1390423bad8e31cf3b",
     "docker-registry": "sha256:356410398f9957e9c96af04ae28875b71ca6174ac0978950b5f3641b21dc8544",
     "docker-container": "sha256:77cf9424bd42cd96851b7363cc21279546a73e55a61cd815c54e279e57206a2a",
-    "kubernetes-job": "sha256:0456ecdb37ef054437ab2c8383f4284c5b58a7161e7dd1707747d108b256cfda",
+    "kubernetes-job": "sha256:686f931093d8fa3a80dee6eb66516be7b022bf29c44a38766da8571f25fede8b",
     "process": "sha256:227b03037cf895d9a70603554b792f573e7e59b81a585d51a43ec1f826e81da3"
 }
diff --git a/tests/infrastructure/test_kubernetes_job.py b/tests/infrastructure/test_kubernetes_job.py
--- a/tests/infrastructure/test_kubernetes_job.py
+++ b/tests/infrastructure/test_kubernetes_job.py
@@ -372,6 +372,21 @@ def test_uses_service_account_name_setting(
     assert service_account_name == "foo"
 
 
+def test_uses_finished_job_ttl_setting(
+    mock_k8s_client,
+    mock_watch,
+    mock_k8s_batch_client,
+):
+    mock_watch.stream = _mock_pods_stream_that_returns_running_pod
+
+    KubernetesJob(command=["echo", "hello"], finished_job_ttl=123).run(MagicMock())
+    mock_k8s_batch_client.create_namespaced_job.assert_called_once()
+    finished_job_ttl = mock_k8s_batch_client.create_namespaced_job.call_args[0][1][
+        "spec"
+    ]["ttlSecondsAfterFinished"]
+    assert finished_job_ttl == 123
+
+
 def test_defaults_to_unspecified_image_pull_policy(
     mock_k8s_client,
     mock_watch,