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).


KubernetesAgent overrides the 'image' and 'args' implied by 'job_template' in KubernetesRun
## Description

I'm running a `KubernetesAgent` to run flows with the following configuration:

* `prefect` version for agent: 0.14.6
* `prefect` version for flow: 0.14.8
* storage: `Webhook`
* run config: `KubernetesRun`

I'm creating the run config object by passing in a dictionary describing a `k8s` job

```python
KubernetesRun(
    job_template=some_job_spec
)
```

## Expected Behavior

I expected that when a flow run was kicked off, the agent would create a single k8s job whose first container had the `image`, `command`, and `args` specified in the spec that I passed into the  `job_template` argument to `KubernetesRun`'s constructor:

```yaml
containers:
- command:
  - "/bin/bash"
  - "-ec"
  args:
  - |
    source some_setup_script.sh
    prefect execute flow-run
  image: my-custom-image:12345
```

Instead, the `image` on my flow run job is set to `prefecthq/prefect:0.14.8` and the `args` are overridden to just `prefect execute flow-run`.

## Reproduction

It's tough to provide a clean reproducible example given that a KubernetesAgent is involved, but I think I can get close. I can see exactly where the issue is happening.

* overriding `image` regardless of the content of `job_template`: https://github.com/PrefectHQ/prefect/blob/24ce11873005c4e6f4f4853b2ec3b107aef05039/src/prefect/agent/kubernetes/agent.py#L571
    - `get_flow_image()` looks at the `image` property of the run config, but not in its template: https://github.com/PrefectHQ/prefect/blob/24ce11873005c4e6f4f4853b2ec3b107aef05039/src/prefect/utilities/agent.py#L36-L37
* overriding `args` regardless of the content of `job_template`: https://github.com/PrefectHQ/prefect/blob/24ce11873005c4e6f4f4853b2ec3b107aef05039/src/prefect/agent/kubernetes/agent.py#L574
    - `get_flow_run_command()` has no knowledge of k8s stuff at all and just returns the prefect CLI command to run the flow: https://github.com/PrefectHQ/prefect/blob/24ce11873005c4e6f4f4853b2ec3b107aef05039/src/prefect/utilities/agent.py#L36-L37

I think I'm looking in the right place, because when I changed my run config definition like this:

```python
KubernetesRun(
    job_template=some_job_spec,
    image="my-custom-image:12345"
)
```

the first container in the job got that image I passed in.

## Environment

* `prefect` version for the flow: 0.14.8
* `prefect` version for the agent: 0.14.6
    - NOTE: I'm using a very very thin wrapper around `KubernetesAgent`, instrumented with some my-company-specific stuff. Based on the links I provided in `## Reproduction` above, I'm pretty confident that the fact that I'm using this other wrapper agent is not related. 
* I've observed this behavior using k8s `v1.16.15` in microk8s `v1.16.15` running on Ubuntu 18.04

## Notes for Reviewers

If you're open to the idea, I'd be happy to submit a pull request that changes https://github.com/PrefectHQ/prefect/blob/24ce11873005c4e6f4f4853b2ec3b107aef05039/src/prefect/agent/kubernetes/agent.py#L570-L574 to respect values provided already in the RunConfig's template, like this:

```python
# treating "" as a special type of "None", since this is the documented
# default in https://docs.prefect.io/orchestration/execution/k8s_job_environment.html#job-spec-configuration
image_from_spec = container.get("image", None)
if image_from_spec is None or image_from_spec == "":
    container["image"] = image = get_flow_image(flow_run)

# treating [] as a special type of "None", since this is the documented
# default in https://docs.prefect.io/orchestration/execution/k8s_job_environment.html#job-spec-configuration
args_from_spec = container.get("args", None)
if args_from_spec is None or args_from_spec == []:
    container["args"] = [get_flow_run_command(flow_run)]
```

Thanks for your time and consideration!


Discussion:
Providing a default `run_config` on the agent doesn't really make much sense. We do support providing a custom job template (via `--job-template`) which will be used by default for all flows deployed by the agent. This lets you configure anything available in the k8s job spec that we don't explicitly override in the agent. That last bit is important though, we currently always set the `image` based on the flow, rather than relying on a default in the template. I can see a good use case for defaulting to the image in the template if no override is set on the flow (see also #4128). I'll push a PR up that accomplishes this.
Thanks for the issue @jameslamb. Overriding of `args` and `image` is intentional (although we could revisit that decision on `image`). The exact command for running a flow in prefect isn't something we want users to have to mess with, and we also may change the command over time (so keeping the args in the container "private" is preferred). As such, I'm pretty against letting users override `args` here (we don't support this for other agents either). If your goal is to add a script on startup, you could add that as part of the image's entrypoint (or set the `cmd` in the job template), neither of which are overridden, and are the intended place for adding some setup steps. We currently do this as part of the provided prefect images, but you're free to extend or override this any way you want: https://github.com/PrefectHQ/prefect/blob/559330bb549b9c92c413948b615509b6c6d33620/Dockerfile#L35.

For `image`, we initially went with always setting it (matching the previous behavior). I could see a case for using the default value in the template if the image isn't provided via `Docker` storage or as part of the run config though, and would accept that change. Either way, we should better document this behavior.

Thoughts?
Thanks for the detailed comments!

**images**

> the image isn't provided via Docker storage or as part of the run config though

I thought of everything in the `job_template` I provided to `KubernetesRun`, including the `image` in the first container, as "part of the run config". I didn't understand that the other keyword arguments like `image` would take precedence over values in that template.

I think it would make sense to figure out the image for ` KubernetesRun` job based on the first matching item in this list

1. image on storage if using `Docker` storage
2. `run_config.image`  (value that was passed to `image=` keyword in `KubernetesRun`)
3. image in the first container in the spec provided by `job_template` / `job_template_path`

I think that's what you're saying above too, just restated it to be sure we're talking about the same thing. I'd be happy to submit a pull request that makes that change and adds language about it in the docs.

**args**

I totally understand why you'd be cautious about letting people mess with `args`. Let me explain a bit why the "add an entrypoint to the image" is not ideal for my use case.

I'm working on an integration (at Saturn Cloud) where users bring any image they want, and separately are able to specify a shell script with commands that run every time a pod for a resource like a Dask cluster, JupyterLab instance, or a prefect flow run starts. People do things in there like `apt-get install` packages or clone git repos. The contract is "as long as this is valid code that can be run with `bash`, we will just run this shell script for you inside the container at startup". This is just like the experience with [EC2 user data](https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/user-data.html) on AWS. If I changed the mechanism for running this to "every image you bring has to have this entrypoint defined that runs this script in this standard place", it would introduce some friction into the process.

In `prefect` 0.13.x world, with `KubernetesJobEnvironment`, I could hack around this in a client library by taking advantage of the fact that setting `args` happened at build time, not run time. I know this is a hack and that I was accessing stuff you consider private, so I'm not saying "wow you broke this in prefect 0.14.x ugh".

* https://github.com/PrefectHQ/prefect/pull/2950
* (this is bad, don't stare directly at it haha): https://github.com/saturncloud/prefect-saturn/blob/2fb73e0972fac4a7aeee67724f8e1e0cc5f982ae/prefect_saturn/core.py#L357-L368

But using `KubernetesRun`, that isn't possible because setting up `args` happens in `generate_job_spec_from_run_config()` , at run time https://github.com/PrefectHQ/prefect/blob/24ce11873005c4e6f4f4853b2ec3b107aef05039/src/prefect/agent/kubernetes/agent.py#L493.

All that said, I understand that my requested change for `args` might be open up the risk of other issues, so I'm ok with handling my use case in a custom extension of `KubernetesAgent`.

thanks for hearing me out!
hey @jameslamb  I actually don't think you need to modify `args` in template OR add `ENTRYPOINT` inside the image to achieve what you want...

> If your goal is to add a script on startup, you could add that as part of the image's entrypoint (or set the cmd in the job template), neither of which are overridden, and are the intended place for adding some setup steps.

you could set the `command` in job template since it is guaranteed to not be overriden by Prefect, something like:

```
containers:
- command:
  - "dumb-init"
  - "--"
  - "some_setup_script.sh"
```

then in `some_setup_script.sh` you would be able to handle the args provided to invoke prefect cmd (`prefect execute flow-run`).
ohhh got it! Did not understand that before, that's just my lack of k8s knowledge. Ok, then that totally works for my use case, thanks!

Resolving patch:
diff --git a/src/prefect/agent/ecs/agent.py b/src/prefect/agent/ecs/agent.py
--- a/src/prefect/agent/ecs/agent.py
+++ b/src/prefect/agent/ecs/agent.py
@@ -387,7 +387,9 @@ def generate_task_definition(
             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
--- a/src/prefect/agent/kubernetes/agent.py
+++ b/src/prefect/agent/kubernetes/agent.py
@@ -651,7 +651,9 @@ def generate_job_spec_from_run_config(
         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
--- a/src/prefect/utilities/agent.py
+++ b/src/prefect/utilities/agent.py
@@ -3,12 +3,15 @@
 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
--- a/tests/agent/test_ecs_agent.py
+++ b/tests/agent/test_ecs_agent.py
@@ -335,9 +335,20 @@ def test_generate_task_definition_family_and_tags(self):
                 "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
--- a/tests/agent/test_k8s_agent.py
+++ b/tests/agent/test_k8s_agent.py
@@ -1283,19 +1283,44 @@ def test_generate_job_spec_metadata(self, tmpdir):
         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
--- 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():