Repository: PrefectHQ/prefect

Issue:
Enable default run_config on an agent
## Opened from the [Prefect Public Slack Community](https://prefect-slackin.herokuapp.com/)

**zach.hodowanec**: Hi Prefect Team!

I am wondering if there is a way to set static `run_config` parameters on a Kubernetes Agent for subsequent flows to consume rather than duplicating similar `run_configs` across various flows. We currently make use of the `PREFECT__CLOUD__AGENT__ENV_VARS` to pass along `storage` configurations, but not having much success attempting to update the execution environment to use an internal custom image. I have tried passing the `IMAGE` and the `PREFECT__CONTEXT__IMAGE` environment variables to my job spec thus far to no avail.

**zach.hodowanec**: As a follow-up, when using the `PREFECT__CONTEXT__IMAGE` environment variable I am seeing the value set correctly in my configmap, however when the job kicks off the env variable is using the `prefecthq/prefect:0.14.6` as it's value. What's interesting about this use case is that if don't pass in a `PREFECT__CONTEXT__IMAGE` value at all the k8s job does *not* assign a value for it.

**jenny**: Hi <@U01NZ4MBCF5> - thanks for your message.  I want to make sure I'm understanding your request and use case correctly.  Am I right that what you'd like to do is re-use your run-configs for different flows?  You might find <https://github.com/PrefectHQ/prefect/discussions/4042|this discussion> on flow deployment useful but come back if you've got more questions.

**zach.hodowanec**: Hi <@ULXMV9SD7> - thanks for your response. Yes, instead of defining a `flow.run_config = KubernetesRun(image=os.getenv("MY_CUSTOM_IMAGE")` on every flow I was hoping to set something on the agent config itself to override the default execution environment image. Since we know each flow is going to execute on a Kubernetes agent it seems redundant to also define the `flow.run_config` as such.

Potentially <https://github.com/PrefectHQ/prefect/issues/4128|this issue> is more in line with what I'm trying to accomplish?

**jenny**: Thanks Zach.  I've checked with the team and I don't think that's possible right now.  But I'm going to open an issue to request it.  If you want to add any more information in that issue to help us understand your use case, please do!  <@ULVA73B9P> open "Enable default run_config on an agent"

Original thread can be found [here](https://prefect-community.slack.com/archives/CL09KU1K7/p1616081635141800?thread_ts=1616081635.141800&cid=CL09KU1K7).



Discussion:

Patch:
diff --git a/changes/pr4270.yaml b/changes/pr4270.yaml
new file mode 100644
index 0000000000..90b3069a4b
--- /dev/null
+++ b/changes/pr4270.yaml
@@ -0,0 +1,2 @@
+enhancement:
+  - "Support setting default `image` in `--job-template`/`--task-definition` in Kubernetes/ECS agents - [#4270](https://github.com/PrefectHQ/prefect/pull/4270)"
diff --git a/src/prefect/agent/ecs/agent.py b/src/prefect/agent/ecs/agent.py
index 1542685431..862c17c8fe 100644
--- a/src/prefect/agent/ecs/agent.py
+++ b/src/prefect/agent/ecs/agent.py
@@ -387,7 +387,9 @@ class ECSAgent(Agent):
             containers.append(container)
 
         # Set flow image
-        container["image"] = image = get_flow_image(flow_run)
+        container["image"] = image = get_flow_image(
+            flow_run, default=container.get("image")
+        )
 
         # Add `PREFECT__CONTEXT__IMAGE` environment variable
         env = {"PREFECT__CONTEXT__IMAGE": image}
diff --git a/src/prefect/agent/kubernetes/agent.py b/src/prefect/agent/kubernetes/agent.py
index ddf4c03506..2db9e88015 100644
--- a/src/prefect/agent/kubernetes/agent.py
+++ b/src/prefect/agent/kubernetes/agent.py
@@ -651,7 +651,9 @@ class KubernetesAgent(Agent):
         container = containers[0]
 
         # Set container image
-        container["image"] = image = get_flow_image(flow_run)
+        container["image"] = image = get_flow_image(
+            flow_run, default=container.get("image")
+        )
 
         # Set flow run command
         container["args"] = get_flow_run_command(flow_run).split()
diff --git a/src/prefect/utilities/agent.py b/src/prefect/utilities/agent.py
index 5ca5999d02..1438ffde40 100644
--- a/src/prefect/utilities/agent.py
+++ b/src/prefect/utilities/agent.py
@@ -3,12 +3,15 @@ from distutils.version import LooseVersion
 from prefect.utilities.graphql import GraphQLResult
 
 
-def get_flow_image(flow_run: GraphQLResult) -> str:
+def get_flow_image(flow_run: GraphQLResult, default: str = None) -> str:
     """
     Retrieve the image to use for this flow run deployment.
 
     Args:
         - flow_run (GraphQLResult): A GraphQLResult flow run object
+        - default (str, optional): A default image to use. If not specified,
+            The `prefecthq/prefect` image corresponding with the flow's prefect
+            version will be used.
 
     Returns:
         - str: a full image name to use for this flow run
@@ -29,13 +32,19 @@ def get_flow_image(flow_run: GraphQLResult) -> str:
     # Not having an environment implies run-config based flow, even if
     # run_config is None.
     if has_run_config or not has_environment:
+        # Precedence:
+        # - Image on docker storage
+        # - Image on run_config
+        # - Provided default
+        # - `prefecthq/prefect` for flow's core version
         if isinstance(storage, Docker):
             return storage.name
-        elif has_run_config:
+        if has_run_config:
             run_config = RunConfigSchema().load(flow_run.run_config)
             if getattr(run_config, "image", None) is not None:
                 return run_config.image
-        # No image found on run-config, and no environment present. Use default.
+        if default is not None:
+            return default
         # core_version should always be present, but just in case
         version = flow_run.flow.get("core_version") or "latest"
         cleaned_version = version.split("+")[0]


Test patch:
diff --git a/tests/agent/test_ecs_agent.py b/tests/agent/test_ecs_agent.py
index 629fb44998..d6d9e44f47 100644
--- a/tests/agent/test_ecs_agent.py
+++ b/tests/agent/test_ecs_agent.py
@@ -335,9 +335,20 @@ class TestGenerateTaskDefinition:
                 "test/name:tag",
             ),
             (ECSRun(image="myimage"), Local(), "myimage"),
+            (
+                ECSRun(
+                    task_definition={
+                        "containerDefinitions": [
+                            {"name": "flow", "image": "on-template"}
+                        ]
+                    }
+                ),
+                Local(),
+                "on-template",
+            ),
             (ECSRun(), Local(), "prefecthq/prefect:0.13.0"),
         ],
-        ids=["on-storage", "on-run_config", "default"],
+        ids=["on-storage", "on-run_config", "on-template", "default"],
     )
     def test_generate_task_definition_image(self, run_config, storage, expected):
         taskdef = self.generate_task_definition(run_config, storage)
diff --git a/tests/agent/test_k8s_agent.py b/tests/agent/test_k8s_agent.py
index 7bb1825bd1..b23a6703e9 100644
--- a/tests/agent/test_k8s_agent.py
+++ b/tests/agent/test_k8s_agent.py
@@ -1283,19 +1283,44 @@ class TestK8sAgentRunConfig:
         assert job["spec"]["template"]["spec"]["restartPolicy"] == "Never"
 
     @pytest.mark.parametrize(
-        "run_config, storage, expected",
+        "run_config, storage, on_template, expected",
         [
             (
                 KubernetesRun(),
                 Docker(registry_url="test", image_name="name", image_tag="tag"),
+                None,
                 "test/name:tag",
             ),
-            (KubernetesRun(image="myimage"), Local(), "myimage"),
-            (KubernetesRun(), Local(), "prefecthq/prefect:0.13.0"),
+            (
+                KubernetesRun(),
+                Docker(registry_url="test", image_name="name", image_tag="tag"),
+                "default-image",
+                "test/name:tag",
+            ),
+            (KubernetesRun(image="myimage"), Local(), None, "myimage"),
+            (KubernetesRun(image="myimage"), Local(), "default-image", "myimage"),
+            (KubernetesRun(), Local(), None, "prefecthq/prefect:0.13.0"),
+            (KubernetesRun(), Local(), "default-image", "default-image"),
+        ],
+        ids=[
+            "on-storage",
+            "on-storage-2",
+            "on-run_config",
+            "on-run_config-2",
+            "on-template",
+            "default",
         ],
-        ids=["on-storage", "on-run_config", "default"],
     )
-    def test_generate_job_spec_image(self, run_config, storage, expected):
+    def test_generate_job_spec_image(
+        self, tmpdir, run_config, storage, on_template, expected
+    ):
+        if on_template:
+            template_path = str(tmpdir.join("job.yaml"))
+            template = self.read_default_template()
+            template["spec"]["template"]["spec"]["containers"][0]["image"] = on_template
+            with open(template_path, "w") as f:
+                yaml.safe_dump(template, f)
+            self.agent.job_template_path = template_path
         flow_run = self.build_flow_run(run_config, storage)
         job = self.agent.generate_job_spec(flow_run)
         image = job["spec"]["template"]["spec"]["containers"][0]["image"]
diff --git a/tests/utilities/test_agent.py b/tests/utilities/test_agent.py
index d5750babfe..c998a77cbd 100644
--- a/tests/utilities/test_agent.py
+++ b/tests/utilities/test_agent.py
@@ -84,7 +84,8 @@ def test_get_flow_image_run_config_docker_storage(run_config):
 
 @pytest.mark.parametrize("run_config", [KubernetesRun(), LocalRun(), None])
 @pytest.mark.parametrize("version", ["0.13.0", "0.10.0+182.g385a32514.dirty", None])
-def test_get_flow_image_run_config_default_value_from_core_version(run_config, version):
+@pytest.mark.parametrize("default", [None, "default-value"])
+def test_get_flow_image_run_config_default(run_config, version, default):
     flow_run = GraphQLResult(
         {
             "flow": GraphQLResult(
@@ -98,9 +99,14 @@ def test_get_flow_image_run_config_default_value_from_core_version(run_config, v
             "id": "id",
         }
     )
-    image = get_flow_image(flow_run)
-    expected_version = version.split("+")[0] if version else "latest"
-    assert image == f"prefecthq/prefect:{expected_version}"
+    if default is None:
+        expected_version = version.split("+")[0] if version else "latest"
+        expected = f"prefecthq/prefect:{expected_version}"
+    else:
+        expected = default
+
+    image = get_flow_image(flow_run, default=default)
+    assert image == expected
 
 
 def test_get_flow_image_run_config_image_on_RunConfig():