Repository: PrefectHQ/prefect

Issue:
Fail job if image can't be pulled (k8 agent)
## Current behavior

If the k8 agent tries to submit a job where the image is not accessible it the job gets "ImagePullBackOff". However the Job keeps being marked as "Submitted" in Prefect cloud.


## Proposed behavior
Agent detects that the Job newer started en reports the error.



## Example
This would enable the user to easier detect that the error is related to image pull.



Discussion:

Resolving patch:
diff --git a/src/prefect/agent/agent.py b/src/prefect/agent/agent.py
--- a/src/prefect/agent/agent.py
+++ b/src/prefect/agent/agent.py
@@ -126,6 +126,7 @@ def __init__(
         self._api_server = None  # type: ignore
         self._api_server_loop = None  # type: Optional[IOLoop]
         self._api_server_thread = None  # type: Optional[threading.Thread]
+        self._heartbeat_thread = None  # type: Optional[threading.Thread]
 
         logger = logging.getLogger(self.name)
         logger.setLevel(config.cloud.agent.get("level"))
@@ -199,6 +200,7 @@ def start(self, _loop_intervals: dict = None) -> None:
 
         try:
             self.setup()
+            self.run_heartbeat_thread()
 
             with exit_handler(self) as exit_event:
 
@@ -226,8 +228,6 @@ def start(self, _loop_intervals: dict = None) -> None:
                         # Reset the event in case it was set by poke handler.
                         AGENT_WAKE_EVENT.clear()
 
-                        self.heartbeat()
-
                         if self.agent_process(executor):
                             index = 0
                         elif index < max(loop_intervals.keys()):
@@ -299,6 +299,23 @@ def stop_server() -> None:
             # will terminate on exit anyway since it's a daemon thread.
             self._api_server_thread.join(timeout=1)
 
+        if self._heartbeat_thread is not None:
+            self.logger.debug("Stopping heartbeat thread")
+            self._heartbeat_thread.join(timeout=1)
+
+    def run_heartbeat_thread(self) -> None:
+        def run() -> None:
+            while True:
+                self.logger.debug("Running agent heartbeat...")
+                self.heartbeat()
+                self.logger.debug("Sleeping heartbeat for 60 seconds")
+                time.sleep(60)
+
+        self._heartbeat_thread = threading.Thread(
+            name="heartbeat", target=run, daemon=True
+        )
+        self._heartbeat_thread.start()
+
     def on_shutdown(self) -> None:
         """
         Invoked when the event loop is exiting and the agent is shutting down. Intended
diff --git a/src/prefect/agent/kubernetes/agent.py b/src/prefect/agent/kubernetes/agent.py
--- a/src/prefect/agent/kubernetes/agent.py
+++ b/src/prefect/agent/kubernetes/agent.py
@@ -1,6 +1,6 @@
 import os
-import uuid
 from os import path
+import uuid
 from typing import Iterable, List
 
 import yaml
@@ -8,6 +8,7 @@
 import prefect
 from prefect import config
 from prefect.agent import Agent
+from prefect.engine.state import Failed
 from prefect.utilities.agent import get_flow_image, get_flow_run_command
 from prefect.utilities.graphql import GraphQLResult
 
@@ -86,7 +87,7 @@ def __init__(
             no_cloud_logs=no_cloud_logs,
         )
 
-        self.namespace = namespace
+        self.namespace = namespace or os.getenv("NAMESPACE", "default")
         self.volume_mounts = volume_mounts
         self.volumes = volumes
 
@@ -103,9 +104,100 @@ def __init__(
             config.load_kube_config()
 
         self.batch_client = client.BatchV1Api()
+        self.core_client = client.CoreV1Api()
+        self.k8s_client = client
 
         self.logger.debug(f"Namespace: {self.namespace}")
 
+    def manage_jobs(self) -> None:
+        """
+        This function checks if jobs are `Failed` or `Succeeded` and if they are then the jobs are
+        deleted from the namespace. If one of the job's pods happen to run into image pulling errors
+        then the flow run is failed and the job is still deleted.
+        """
+        self.logger.debug(
+            "Retrieving information of jobs that are currently in the cluster..."
+        )
+
+        more = True
+        _continue = ""
+        while more:
+            try:
+                jobs = self.batch_client.list_namespaced_job(
+                    namespace=self.namespace,
+                    label_selector="prefect.io/identifier",
+                    limit=20,
+                    _continue=_continue,
+                )
+                _continue = jobs.metadata._continue
+                more = bool(_continue)
+
+                for job in jobs.items:
+                    delete_job = job.status.failed or job.status.succeeded
+                    job_name = job.metadata.name
+                    flow_run_id = job.metadata.labels.get("prefect.io/flow_run_id")
+
+                    if not delete_job:
+                        pods = self.core_client.list_namespaced_pod(
+                            namespace=self.namespace,
+                            label_selector="prefect.io/identifier={}".format(
+                                job.metadata.labels.get("prefect.io/identifier")
+                            ),
+                        )
+
+                        for pod in pods.items:
+                            if pod.status.container_statuses:
+                                for container_status in pod.status.container_statuses:
+                                    waiting = container_status.state.waiting
+                                    if waiting and (
+                                        waiting.reason == "ErrImagePull"
+                                        or waiting.reason == "ImagePullBackOff"
+                                    ):
+                                        self.logger.debug(
+                                            f"Failing flow run {flow_run_id} due to pod {waiting.reason}"
+                                        )
+                                        self.client.set_flow_run_state(
+                                            flow_run_id=flow_run_id,
+                                            state=Failed(
+                                                message="Kubernetes Error: {}".format(
+                                                    container_status.state.waiting.message
+                                                )
+                                            ),
+                                        )
+
+                                        delete_job = True
+                                        break
+
+                    if delete_job:
+                        self.logger.debug(f"Deleting job {job_name}")
+                        try:
+                            self.batch_client.delete_namespaced_job(
+                                name=job_name,
+                                namespace=self.namespace,
+                                body=self.k8s_client.V1DeleteOptions(
+                                    propagation_policy="Foreground"
+                                ),
+                            )
+                        except self.k8s_client.rest.ApiException as exc:
+                            if exc.status != 404:
+                                self.logger.error(
+                                    f"{exc.status} error attempting to delete job {job_name}"
+                                )
+            except self.k8s_client.rest.ApiException as exc:
+                if exc.status == 410:
+                    self.logger.debug("Refreshing job listing token...")
+                    _continue = ""
+                    continue
+                else:
+                    self.logger.debug(exc)
+
+    def heartbeat(self) -> None:
+        """
+        Check status of jobs created by this agent, delete completed jobs and failed containers.
+        """
+        self.manage_jobs()
+        super().heartbeat()
+
     def deploy_flow(self, flow_run: GraphQLResult) -> str:
         """
         Deploy flow runs on to a k8s cluster as jobs
@@ -120,20 +212,25 @@ def deploy_flow(self, flow_run: GraphQLResult) -> str:
 
         image = get_flow_image(flow_run=flow_run)
 
-        job_spec = self.replace_job_spec_yaml(flow_run, image)
+        identifier = str(uuid.uuid4())[:8]
+        job_spec = self.replace_job_spec_yaml(
+            flow_run=flow_run, image=image, identifier=identifier
+        )
 
         self.logger.debug(
             "Creating namespaced job {}".format(job_spec["metadata"]["name"])
         )
         job = self.batch_client.create_namespaced_job(
-            namespace=self.namespace or os.getenv("NAMESPACE", "default"), body=job_spec
+            namespace=self.namespace, body=job_spec
         )
 
         self.logger.debug("Job {} created".format(job.metadata.name))
 
         return "Job {}".format(job.metadata.name)
 
-    def replace_job_spec_yaml(self, flow_run: GraphQLResult, image: str) -> dict:
+    def replace_job_spec_yaml(
+        self, flow_run: GraphQLResult, image: str, identifier: str
+    ) -> dict:
         """
         Populate a k8s job spec. This spec defines a k8s job that handles
         executing a flow. This method runs each time the agent receives
@@ -161,6 +258,7 @@ def replace_job_spec_yaml(self, flow_run: GraphQLResult, image: str) -> dict:
         Args:
             - flow_run (GraphQLResult): A flow run object
             - image (str): The full name of an image to use for the job
+            - identifier (str): A unique identifier to identify this job
 
         Returns:
             - dict: a dictionary representation of a k8s job for flow execution
@@ -171,7 +269,6 @@ def replace_job_spec_yaml(self, flow_run: GraphQLResult, image: str) -> dict:
         with open(yaml_path, "r") as job_file:
             job = yaml.safe_load(job_file)
 
-        identifier = str(uuid.uuid4())[:8]
         job_name = "prefect-job-{}".format(identifier)
 
         # Populate job metadata for identification
diff --git a/src/prefect/agent/kubernetes/resource_manager.py b/src/prefect/agent/kubernetes/resource_manager.py
--- a/src/prefect/agent/kubernetes/resource_manager.py
+++ b/src/prefect/agent/kubernetes/resource_manager.py
@@ -19,6 +19,9 @@ class ResourceManager:
     The resource manager is responsible for cleaning up old completed/failed k8s jobs
     and pods from the cluster. This is optional and does not need to me used for the agent
     to work.
+
+    DEPRECATED: The resource manager is deprecated and it's main functionality is now present in the
+    Kubernetes agent.
     """
 
     def __init__(self) -> None:
@@ -55,6 +58,7 @@ def start(self) -> None:
         """
         Main loop which waits on a `LOOP_INTERVAL` and looks for finished jobs to clean
         """
+        self.logger.warning("DEPRECATED: The resource manager is deprecated")
         self.logger.info("Starting {}".format(type(self).__name__))
         while True:
             try:


Test patch:
diff --git a/tests/agent/test_agent.py b/tests/agent/test_agent.py
--- a/tests/agent/test_agent.py
+++ b/tests/agent/test_agent.py
@@ -474,7 +474,7 @@ def test_agent_start_max_polls_count(monkeypatch, runner_token, cloud_api):
 
     assert on_shutdown.call_count == 1
     assert agent_process.call_count == 2
-    assert heartbeat.call_count == 2
+    assert heartbeat.call_count == 1
 
 
 def test_agent_start_max_polls_zero(monkeypatch, runner_token, cloud_api):
@@ -495,7 +495,7 @@ def test_agent_start_max_polls_zero(monkeypatch, runner_token, cloud_api):
 
     assert on_shutdown.call_count == 1
     assert agent_process.call_count == 0
-    assert heartbeat.call_count == 0
+    assert heartbeat.call_count == 1
 
 
 def test_agent_registration_and_id(monkeypatch, cloud_api):
diff --git a/tests/agent/test_docker_agent.py b/tests/agent/test_docker_agent.py
--- a/tests/agent/test_docker_agent.py
+++ b/tests/agent/test_docker_agent.py
@@ -957,7 +957,7 @@ def test_docker_agent_start_max_polls_count(monkeypatch, runner_token, cloud_api
 
     assert on_shutdown.call_count == 1
     assert agent_process.call_count == 2
-    assert heartbeat.call_count == 2
+    assert heartbeat.call_count == 1
 
 
 def test_docker_agent_start_max_polls_zero(monkeypatch, runner_token, cloud_api):
@@ -986,7 +986,7 @@ def test_docker_agent_start_max_polls_zero(monkeypatch, runner_token, cloud_api)
 
     assert on_shutdown.call_count == 1
     assert agent_process.call_count == 0
-    assert heartbeat.call_count == 0
+    assert heartbeat.call_count == 1
 
 
 def test_docker_agent_network(monkeypatch, cloud_api):
diff --git a/tests/agent/test_fargate_agent.py b/tests/agent/test_fargate_agent.py
--- a/tests/agent/test_fargate_agent.py
+++ b/tests/agent/test_fargate_agent.py
@@ -1971,7 +1971,7 @@ def test_fargate_agent_start_max_polls_count(monkeypatch, runner_token, cloud_ap
 
     assert on_shutdown.call_count == 1
     assert agent_process.call_count == 2
-    assert heartbeat.call_count == 2
+    assert heartbeat.call_count == 1
 
 
 def test_fargate_agent_start_max_polls_zero(monkeypatch, runner_token, cloud_api):
@@ -1997,7 +1997,7 @@ def test_fargate_agent_start_max_polls_zero(monkeypatch, runner_token, cloud_api
 
     assert on_shutdown.call_count == 1
     assert agent_process.call_count == 0
-    assert heartbeat.call_count == 0
+    assert heartbeat.call_count == 1
 
 
 def test_agent_configuration_utility(monkeypatch, cloud_api):
diff --git a/tests/agent/test_k8s_agent.py b/tests/agent/test_k8s_agent.py
--- a/tests/agent/test_k8s_agent.py
+++ b/tests/agent/test_k8s_agent.py
@@ -1,4 +1,5 @@
-from unittest.mock import MagicMock
+import json
+from unittest.mock import MagicMock, PropertyMock
 
 import pytest
 
@@ -18,6 +19,12 @@ def test_k8s_agent_init(monkeypatch, cloud_api):
     k8s_config = MagicMock()
     monkeypatch.setattr("kubernetes.config", k8s_config)
 
+    get_jobs = MagicMock(return_value=[])
+    monkeypatch.setattr(
+        "prefect.agent.kubernetes.agent.KubernetesAgent.manage_jobs",
+        get_jobs,
+    )
+
     agent = KubernetesAgent()
     assert agent
     assert agent.labels == []
@@ -29,6 +36,15 @@ def test_k8s_agent_config_options(monkeypatch, cloud_api):
     k8s_config = MagicMock()
     monkeypatch.setattr("kubernetes.config", k8s_config)
 
+    k8s_client = MagicMock()
+    monkeypatch.setattr("kubernetes.client", k8s_client)
+
+    get_jobs = MagicMock(return_value=[])
+    monkeypatch.setattr(
+        "prefect.agent.kubernetes.agent.KubernetesAgent.manage_jobs",
+        get_jobs,
+    )
+
     with set_temporary_config({"cloud.agent.auth_token": "TEST_TOKEN"}):
         agent = KubernetesAgent(name="test", labels=["test"], namespace="namespace")
         assert agent
@@ -54,9 +70,20 @@ def test_k8s_agent_deploy_flow(core_version, command, monkeypatch, cloud_api):
     monkeypatch.setattr("kubernetes.config", k8s_config)
 
     batch_client = MagicMock()
-    batch_client.create_namespaced_job.return_value = {}
     monkeypatch.setattr(
-        "kubernetes.client.BatchV1Api", MagicMock(retrurn_value=batch_client)
+        "kubernetes.client.BatchV1Api", MagicMock(return_value=batch_client)
+    )
+
+    core_client = MagicMock()
+    core_client.list_namespaced_pod.return_value = MagicMock(items=[])
+    monkeypatch.setattr(
+        "kubernetes.client.CoreV1Api", MagicMock(return_value=core_client)
+    )
+
+    get_jobs = MagicMock(return_value=[])
+    monkeypatch.setattr(
+        "prefect.agent.kubernetes.agent.KubernetesAgent.manage_jobs",
+        get_jobs,
     )
 
     agent = KubernetesAgent()
@@ -96,9 +123,20 @@ def test_k8s_agent_deploy_flow_uses_environment_metadata(monkeypatch, cloud_api)
     monkeypatch.setattr("kubernetes.config", k8s_config)
 
     batch_client = MagicMock()
-    batch_client.create_namespaced_job.return_value = {}
     monkeypatch.setattr(
-        "kubernetes.client.BatchV1Api", MagicMock(retrurn_value=batch_client)
+        "kubernetes.client.BatchV1Api", MagicMock(return_value=batch_client)
+    )
+
+    core_client = MagicMock()
+    core_client.list_namespaced_pod.return_value = MagicMock(items=[])
+    monkeypatch.setattr(
+        "kubernetes.client.CoreV1Api", MagicMock(return_value=core_client)
+    )
+
+    get_jobs = MagicMock(return_value=[])
+    monkeypatch.setattr(
+        "prefect.agent.kubernetes.agent.KubernetesAgent.manage_jobs",
+        get_jobs,
     )
 
     agent = KubernetesAgent()
@@ -134,9 +172,20 @@ def test_k8s_agent_deploy_flow_raises(monkeypatch, cloud_api):
     monkeypatch.setattr("kubernetes.config", k8s_config)
 
     batch_client = MagicMock()
-    batch_client.create_namespaced_job.return_value = {}
     monkeypatch.setattr(
-        "kubernetes.client.BatchV1Api", MagicMock(retrurn_value=batch_client)
+        "kubernetes.client.BatchV1Api", MagicMock(return_value=batch_client)
+    )
+
+    core_client = MagicMock()
+    core_client.list_namespaced_pod.return_value = MagicMock(items=[])
+    monkeypatch.setattr(
+        "kubernetes.client.CoreV1Api", MagicMock(return_value=core_client)
+    )
+
+    get_jobs = MagicMock(return_value=[])
+    monkeypatch.setattr(
+        "prefect.agent.kubernetes.agent.KubernetesAgent.manage_jobs",
+        get_jobs,
     )
 
     agent = KubernetesAgent()
@@ -164,6 +213,12 @@ def test_k8s_agent_replace_yaml_uses_user_env_vars(monkeypatch, cloud_api):
     k8s_config = MagicMock()
     monkeypatch.setattr("kubernetes.config", k8s_config)
 
+    get_jobs = MagicMock(return_value=[])
+    monkeypatch.setattr(
+        "prefect.agent.kubernetes.agent.KubernetesAgent.manage_jobs",
+        get_jobs,
+    )
+
     monkeypatch.setenv("IMAGE_PULL_SECRETS", "my-secret")
     monkeypatch.setenv("JOB_MEM_REQUEST", "mr")
     monkeypatch.setenv("JOB_MEM_LIMIT", "ml")
@@ -192,7 +247,9 @@ def test_k8s_agent_replace_yaml_uses_user_env_vars(monkeypatch, cloud_api):
         {"cloud.agent.auth_token": "token", "logging.log_to_cloud": True}
     ):
         agent = KubernetesAgent(env_vars=dict(AUTH_THING="foo", PKG_SETTING="bar"))
-        job = agent.replace_job_spec_yaml(flow_run, image="test/name:tag")
+        job = agent.replace_job_spec_yaml(
+            flow_run, image="test/name:tag", identifier="identifier"
+        )
 
         assert job["metadata"]["labels"]["prefect.io/flow_run_id"] == "id"
         assert job["metadata"]["labels"]["prefect.io/flow_id"] == "new_id"
@@ -232,6 +289,12 @@ def test_k8s_agent_replace_yaml(monkeypatch, cloud_api):
     k8s_config = MagicMock()
     monkeypatch.setattr("kubernetes.config", k8s_config)
 
+    get_jobs = MagicMock(return_value=[])
+    monkeypatch.setattr(
+        "prefect.agent.kubernetes.agent.KubernetesAgent.manage_jobs",
+        get_jobs,
+    )
+
     monkeypatch.setenv("IMAGE_PULL_SECRETS", "my-secret")
     monkeypatch.setenv("JOB_MEM_REQUEST", "mr")
     monkeypatch.setenv("JOB_MEM_LIMIT", "ml")
@@ -260,7 +323,9 @@ def test_k8s_agent_replace_yaml(monkeypatch, cloud_api):
         volume_mounts = [{"name": "my-vol", "mountPath": "/mnt/my-mount"}]
         volumes = [{"name": "my-vol", "hostPath": "/host/folder"}]
         agent = KubernetesAgent(volume_mounts=volume_mounts, volumes=volumes)
-        job = agent.replace_job_spec_yaml(flow_run, image="test/name:tag")
+        job = agent.replace_job_spec_yaml(
+            flow_run, image="test/name:tag", identifier="identifier"
+        )
 
         assert job["metadata"]["labels"]["prefect.io/flow_run_id"] == "id"
         assert job["metadata"]["labels"]["prefect.io/flow_id"] == "new_id"
@@ -316,6 +381,12 @@ def test_k8s_agent_replace_yaml_responds_to_logging_config(
     k8s_config = MagicMock()
     monkeypatch.setattr("kubernetes.config", k8s_config)
 
+    get_jobs = MagicMock(return_value=[])
+    monkeypatch.setattr(
+        "prefect.agent.kubernetes.agent.KubernetesAgent.manage_jobs",
+        get_jobs,
+    )
+
     flow_run = GraphQLResult(
         {
             "flow": GraphQLResult(
@@ -334,7 +405,9 @@ def test_k8s_agent_replace_yaml_responds_to_logging_config(
     )
 
     agent = KubernetesAgent(no_cloud_logs=flag)
-    job = agent.replace_job_spec_yaml(flow_run, image="test/name:tag")
+    job = agent.replace_job_spec_yaml(
+        flow_run, image="test/name:tag", identifier="identifier"
+    )
     env = job["spec"]["template"]["spec"]["containers"][0]["env"]
     assert env[6]["value"] == str(not flag).lower()
 
@@ -343,6 +416,12 @@ def test_k8s_agent_replace_yaml_no_pull_secrets(monkeypatch, cloud_api):
     k8s_config = MagicMock()
     monkeypatch.setattr("kubernetes.config", k8s_config)
 
+    get_jobs = MagicMock(return_value=[])
+    monkeypatch.setattr(
+        "prefect.agent.kubernetes.agent.KubernetesAgent.manage_jobs",
+        get_jobs,
+    )
+
     flow_run = GraphQLResult(
         {
             "flow": GraphQLResult(
@@ -360,7 +439,9 @@ def test_k8s_agent_replace_yaml_no_pull_secrets(monkeypatch, cloud_api):
     )
 
     agent = KubernetesAgent()
-    job = agent.replace_job_spec_yaml(flow_run, image="test/name:tag")
+    job = agent.replace_job_spec_yaml(
+        flow_run, image="test/name:tag", identifier="identifier"
+    )
 
     assert not job["spec"]["template"]["spec"].get("imagePullSecrets", None)
 
@@ -386,7 +467,7 @@ def test_k8s_agent_removes_yaml_no_volume(monkeypatch, cloud_api):
     )
 
     agent = KubernetesAgent()
-    job = agent.replace_job_spec_yaml(flow_run, image="test/name:tag")
+    job = agent.replace_job_spec_yaml(flow_run, image="test/name:tag", identifier="id")
 
     assert not job["spec"]["template"]["spec"].get("volumes", None)
     assert not job["spec"]["template"]["spec"]["containers"][0].get(
@@ -398,6 +479,12 @@ def test_k8s_agent_includes_agent_labels_in_job(monkeypatch, cloud_api):
     k8s_config = MagicMock()
     monkeypatch.setattr("kubernetes.config", k8s_config)
 
+    get_jobs = MagicMock(return_value=[])
+    monkeypatch.setattr(
+        "prefect.agent.kubernetes.agent.KubernetesAgent.manage_jobs",
+        get_jobs,
+    )
+
     flow_run = GraphQLResult(
         {
             "flow": GraphQLResult(
@@ -415,7 +502,9 @@ def test_k8s_agent_includes_agent_labels_in_job(monkeypatch, cloud_api):
     )
 
     agent = KubernetesAgent(labels=["foo", "bar"])
-    job = agent.replace_job_spec_yaml(flow_run, image="test/name:tag")
+    job = agent.replace_job_spec_yaml(
+        flow_run, image="test/name:tag", identifier="identifier"
+    )
     env = job["spec"]["template"]["spec"]["containers"][0]["env"]
     assert env[5]["value"] == "['foo', 'bar']"
 
@@ -424,6 +513,12 @@ def test_k8s_agent_generate_deployment_yaml(monkeypatch, cloud_api):
     k8s_config = MagicMock()
     monkeypatch.setattr("kubernetes.config", k8s_config)
 
+    get_jobs = MagicMock(return_value=[])
+    monkeypatch.setattr(
+        "prefect.agent.kubernetes.agent.KubernetesAgent.manage_jobs",
+        get_jobs,
+    )
+
     agent = KubernetesAgent()
     deployment = agent.generate_deployment_yaml(
         token="test_token",
@@ -454,6 +549,12 @@ def test_k8s_agent_generate_deployment_yaml_env_vars(monkeypatch, cloud_api):
     k8s_config = MagicMock()
     monkeypatch.setattr("kubernetes.config", k8s_config)
 
+    get_jobs = MagicMock(return_value=[])
+    monkeypatch.setattr(
+        "prefect.agent.kubernetes.agent.KubernetesAgent.manage_jobs",
+        get_jobs,
+    )
+
     agent = KubernetesAgent()
     deployment = agent.generate_deployment_yaml(
         env_vars={"test1": "test2", "test3": "test4"}
@@ -476,6 +577,12 @@ def test_k8s_agent_generate_deployment_yaml_backend_default(monkeypatch, server_
     c = MagicMock()
     monkeypatch.setattr("prefect.agent.agent.Client", c)
 
+    get_jobs = MagicMock(return_value=[])
+    monkeypatch.setattr(
+        "prefect.agent.kubernetes.agent.KubernetesAgent.manage_jobs",
+        get_jobs,
+    )
+
     agent = KubernetesAgent()
     deployment = agent.generate_deployment_yaml()
 
@@ -502,6 +609,12 @@ def test_k8s_agent_generate_deployment_yaml_local_version(
     k8s_config = MagicMock()
     monkeypatch.setattr("kubernetes.config", k8s_config)
 
+    get_jobs = MagicMock(return_value=[])
+    monkeypatch.setattr(
+        "prefect.agent.kubernetes.agent.KubernetesAgent.manage_jobs",
+        get_jobs,
+    )
+
     agent = KubernetesAgent()
     deployment = agent.generate_deployment_yaml(
         token="test_token",
@@ -523,6 +636,12 @@ def test_k8s_agent_generate_deployment_yaml_latest(monkeypatch, cloud_api):
     k8s_config = MagicMock()
     monkeypatch.setattr("kubernetes.config", k8s_config)
 
+    get_jobs = MagicMock(return_value=[])
+    monkeypatch.setattr(
+        "prefect.agent.kubernetes.agent.KubernetesAgent.manage_jobs",
+        get_jobs,
+    )
+
     agent = KubernetesAgent()
     deployment = agent.generate_deployment_yaml(
         token="test_token",
@@ -545,6 +664,12 @@ def test_k8s_agent_generate_deployment_yaml_no_resource_manager(monkeypatch, clo
     k8s_config = MagicMock()
     monkeypatch.setattr("kubernetes.config", k8s_config)
 
+    get_jobs = MagicMock(return_value=[])
+    monkeypatch.setattr(
+        "prefect.agent.kubernetes.agent.KubernetesAgent.manage_jobs",
+        get_jobs,
+    )
+
     agent = KubernetesAgent()
     deployment = agent.generate_deployment_yaml(
         token="test_token", api="test_api", namespace="test_namespace"
@@ -565,6 +690,12 @@ def test_k8s_agent_generate_deployment_yaml_labels(monkeypatch, cloud_api):
     k8s_config = MagicMock()
     monkeypatch.setattr("kubernetes.config", k8s_config)
 
+    get_jobs = MagicMock(return_value=[])
+    monkeypatch.setattr(
+        "prefect.agent.kubernetes.agent.KubernetesAgent.manage_jobs",
+        get_jobs,
+    )
+
     agent = KubernetesAgent()
     deployment = agent.generate_deployment_yaml(
         token="test_token",
@@ -591,6 +722,12 @@ def test_k8s_agent_generate_deployment_yaml_no_image_pull_secrets(
     k8s_config = MagicMock()
     monkeypatch.setattr("kubernetes.config", k8s_config)
 
+    get_jobs = MagicMock(return_value=[])
+    monkeypatch.setattr(
+        "prefect.agent.kubernetes.agent.KubernetesAgent.manage_jobs",
+        get_jobs,
+    )
+
     agent = KubernetesAgent()
     deployment = agent.generate_deployment_yaml(
         token="test_token", api="test_api", namespace="test_namespace"
@@ -609,6 +746,12 @@ def test_k8s_agent_generate_deployment_yaml_contains_image_pull_secrets(
     k8s_config = MagicMock()
     monkeypatch.setattr("kubernetes.config", k8s_config)
 
+    get_jobs = MagicMock(return_value=[])
+    monkeypatch.setattr(
+        "prefect.agent.kubernetes.agent.KubernetesAgent.manage_jobs",
+        get_jobs,
+    )
+
     agent = KubernetesAgent()
     deployment = agent.generate_deployment_yaml(
         token="test_token",
@@ -631,6 +774,12 @@ def test_k8s_agent_generate_deployment_yaml_contains_resources(monkeypatch, clou
     k8s_config = MagicMock()
     monkeypatch.setattr("kubernetes.config", k8s_config)
 
+    get_jobs = MagicMock(return_value=[])
+    monkeypatch.setattr(
+        "prefect.agent.kubernetes.agent.KubernetesAgent.manage_jobs",
+        get_jobs,
+    )
+
     agent = KubernetesAgent()
     deployment = agent.generate_deployment_yaml(
         token="test_token",
@@ -660,6 +809,12 @@ def test_k8s_agent_generate_deployment_yaml_rbac(monkeypatch, cloud_api):
     k8s_config = MagicMock()
     monkeypatch.setattr("kubernetes.config", k8s_config)
 
+    get_jobs = MagicMock(return_value=[])
+    monkeypatch.setattr(
+        "prefect.agent.kubernetes.agent.KubernetesAgent.manage_jobs",
+        get_jobs,
+    )
+
     agent = KubernetesAgent()
     deployment = agent.generate_deployment_yaml(
         token="test_token", api="test_api", namespace="test_namespace", rbac=True
@@ -678,6 +833,12 @@ def test_k8s_agent_start_max_polls(monkeypatch, runner_token, cloud_api):
     k8s_config = MagicMock()
     monkeypatch.setattr("kubernetes.config", k8s_config)
 
+    get_jobs = MagicMock(return_value=[])
+    monkeypatch.setattr(
+        "prefect.agent.kubernetes.agent.KubernetesAgent.manage_jobs",
+        get_jobs,
+    )
+
     on_shutdown = MagicMock()
     monkeypatch.setattr(
         "prefect.agent.kubernetes.agent.KubernetesAgent.on_shutdown", on_shutdown
@@ -705,6 +866,12 @@ def test_k8s_gent_start_max_polls_count(monkeypatch, runner_token, cloud_api):
     k8s_config = MagicMock()
     monkeypatch.setattr("kubernetes.config", k8s_config)
 
+    get_jobs = MagicMock(return_value=[])
+    monkeypatch.setattr(
+        "prefect.agent.kubernetes.agent.KubernetesAgent.manage_jobs",
+        get_jobs,
+    )
+
     on_shutdown = MagicMock()
     monkeypatch.setattr(
         "prefect.agent.kubernetes.agent.KubernetesAgent.on_shutdown", on_shutdown
@@ -726,13 +893,19 @@ def test_k8s_gent_start_max_polls_count(monkeypatch, runner_token, cloud_api):
 
     assert on_shutdown.call_count == 1
     assert agent_process.call_count == 2
-    assert heartbeat.call_count == 2
+    assert heartbeat.call_count == 1
 
 
 def test_k8s_agent_start_max_polls_zero(monkeypatch, runner_token, cloud_api):
     k8s_config = MagicMock()
     monkeypatch.setattr("kubernetes.config", k8s_config)
 
+    get_jobs = MagicMock(return_value=[])
+    monkeypatch.setattr(
+        "prefect.agent.kubernetes.agent.KubernetesAgent.manage_jobs",
+        get_jobs,
+    )
+
     on_shutdown = MagicMock()
     monkeypatch.setattr(
         "prefect.agent.kubernetes.agent.KubernetesAgent.on_shutdown", on_shutdown
@@ -754,4 +927,111 @@ def test_k8s_agent_start_max_polls_zero(monkeypatch, runner_token, cloud_api):
 
     assert on_shutdown.call_count == 1
     assert agent_process.call_count == 0
-    assert heartbeat.call_count == 0
+    assert heartbeat.call_count == 1
+
+
+def test_k8s_agent_manage_jobs_pass(monkeypatch, cloud_api):
+    k8s_config = MagicMock()
+    monkeypatch.setattr("kubernetes.config", k8s_config)
+
+    job_mock = MagicMock()
+    job_mock.metadata.labels = {
+        "prefect.io/identifier": "id",
+        "prefect.io/flow_run_id": "fr",
+    }
+    job_mock.metadata.name = "my_job"
+    batch_client = MagicMock()
+    list_job = MagicMock()
+    list_job.metadata._continue = 0
+    list_job.items = [job_mock]
+    batch_client.list_namespaced_job.return_value = list_job
+    monkeypatch.setattr(
+        "kubernetes.client.BatchV1Api", MagicMock(return_value=batch_client)
+    )
+
+    pod = MagicMock()
+    pod.metadata.name = "pod_name"
+    core_client = MagicMock()
+    list_pods = MagicMock()
+    list_pods.items = [pod]
+    core_client.list_namespaced_pod.return_value = list_pods
+    monkeypatch.setattr(
+        "kubernetes.client.CoreV1Api", MagicMock(return_value=core_client)
+    )
+
+    agent = KubernetesAgent()
+    agent.heartbeat()
+
+
+def test_k8s_agent_manage_jobs_client_call(monkeypatch, cloud_api):
+    gql_return = MagicMock(
+        return_value=MagicMock(data=MagicMock(set_flow_run_state=None))
+    )
+    client = MagicMock()
+    client.return_value.graphql = gql_return
+    monkeypatch.setattr("prefect.agent.agent.Client", client)
+
+    k8s_config = MagicMock()
+    monkeypatch.setattr("kubernetes.config", k8s_config)
+
+    job_mock = MagicMock()
+    job_mock.metadata.labels = {
+        "prefect.io/identifier": "id",
+        "prefect.io/flow_run_id": "fr",
+    }
+    job_mock.metadata.name = "my_job"
+    job_mock.status.failed = False
+    job_mock.status.succeeded = False
+    batch_client = MagicMock()
+    list_job = MagicMock()
+    list_job.metadata._continue = 0
+    list_job.items = [job_mock]
+    batch_client.list_namespaced_job.return_value = list_job
+    monkeypatch.setattr(
+        "kubernetes.client.BatchV1Api", MagicMock(return_value=batch_client)
+    )
+
+    pod = MagicMock()
+    pod.metadata.name = "pod_name"
+    c_status = MagicMock()
+    c_status.state.waiting.reason = "ErrImagePull"
+    pod.status.container_statuses = [c_status]
+    core_client = MagicMock()
+    list_pods = MagicMock()
+    list_pods.items = [pod]
+    core_client.list_namespaced_pod.return_value = list_pods
+    monkeypatch.setattr(
+        "kubernetes.client.CoreV1Api", MagicMock(return_value=core_client)
+    )
+
+    agent = KubernetesAgent()
+    agent.manage_jobs()
+
+
+def test_k8s_agent_manage_jobs_delete_jobs(monkeypatch, cloud_api):
+
+    k8s_config = MagicMock()
+    monkeypatch.setattr("kubernetes.config", k8s_config)
+
+    job_mock = MagicMock()
+    job_mock.metadata.labels = {
+        "prefect.io/identifier": "id",
+        "prefect.io/flow_run_id": "fr",
+    }
+    job_mock.metadata.name = "my_job"
+    job_mock.status.failed = True
+    job_mock.status.succeeded = True
+    batch_client = MagicMock()
+    list_job = MagicMock()
+    list_job.metadata._continue = 0
+    list_job.items = [job_mock]
+    batch_client.list_namespaced_job.return_value = list_job
+    batch_client.delete_namespaced_job.return_value = None
+    monkeypatch.setattr(
+        "kubernetes.client.BatchV1Api", MagicMock(return_value=batch_client)
+    )
+
+    agent = KubernetesAgent()
+    agent.manage_jobs()
+
+    assert batch_client.delete_namespaced_job.called
diff --git a/tests/agent/test_local_agent.py b/tests/agent/test_local_agent.py
--- a/tests/agent/test_local_agent.py
+++ b/tests/agent/test_local_agent.py
@@ -606,7 +606,7 @@ def test_local_gent_start_max_polls_count(monkeypatch, runner_token, cloud_api):
 
     assert on_shutdown.call_count == 1
     assert agent_process.call_count == 2
-    assert heartbeat.call_count == 2
+    assert heartbeat.call_count == 1
 
 
 def test_local_agent_start_max_polls_zero(monkeypatch, runner_token, cloud_api):
@@ -627,4 +627,4 @@ def test_local_agent_start_max_polls_zero(monkeypatch, runner_token, cloud_api):
 
     assert on_shutdown.call_count == 1
     assert agent_process.call_count == 0
-    assert heartbeat.call_count == 0
+    assert heartbeat.call_count == 1
diff --git a/tests/cli/test_agent.py b/tests/cli/test_agent.py
--- a/tests/cli/test_agent.py
+++ b/tests/cli/test_agent.py
@@ -106,6 +106,12 @@ def test_agent_start_kubernetes(monkeypatch, cloud_api):
     start = MagicMock()
     monkeypatch.setattr("prefect.agent.kubernetes.KubernetesAgent.start", start)
 
+    get_jobs = MagicMock(return_value=[])
+    monkeypatch.setattr(
+        "prefect.agent.kubernetes.agent.KubernetesAgent.manage_jobs",
+        get_jobs,
+    )
+
     k8s_config = MagicMock()
     monkeypatch.setattr("kubernetes.config", k8s_config)
 
@@ -118,6 +124,12 @@ def test_agent_start_kubernetes_namespace(monkeypatch, cloud_api):
     start = MagicMock()
     monkeypatch.setattr("prefect.agent.kubernetes.KubernetesAgent.start", start)
 
+    get_jobs = MagicMock(return_value=[])
+    monkeypatch.setattr(
+        "prefect.agent.kubernetes.agent.KubernetesAgent.manage_jobs",
+        get_jobs,
+    )
+
     k8s_config = MagicMock()
     monkeypatch.setattr("kubernetes.config", k8s_config)
 
@@ -130,6 +142,12 @@ def test_agent_start_kubernetes_kwargs_ignored(monkeypatch, cloud_api):
     start = MagicMock()
     monkeypatch.setattr("prefect.agent.kubernetes.KubernetesAgent.start", start)
 
+    get_jobs = MagicMock(return_value=[])
+    monkeypatch.setattr(
+        "prefect.agent.kubernetes.agent.KubernetesAgent.manage_jobs",
+        get_jobs,
+    )
+
     k8s_config = MagicMock()
     monkeypatch.setattr("kubernetes.config", k8s_config)