Repository: PrefectHQ/prefect

Issue:
Add `EXTRA_PIP_PACKAGES` to the `KubernetesJob` block
### 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 proposed behavior

Add it to K8s block

### Describe the current behavior

Atm it only works with the `DockerContainer` block

### Example Use

_No response_

### Additional context

_No response_


Discussion:
Can you please add more detail? `EXTRA_PIP_PACKAGES` is a feature of the Prefect base image (https://github.com/PrefectHQ/prefect/blob/main/scripts/entrypoint.sh) and shouldn't rely on any details of the infrastructure block except setting environment variables.
Sure! This is supposed to solve this problem for users that leverage the `KubernetesJob` block with a Prefect base image:

https://discourse.prefect.io/t/i-m-getting-an-error-file-system-could-not-be-created-you-are-likely-missing-a-python-module-required-to-use-the-given-storage-protocol-how-to-solve-that/1459/1

They want to make it possible to set e.g. s3fs into this environment variable to install this dependency at runtime so that the flow run pod can download the flow code from S3 even when using base Prefect images that don't have s3fs dependency baked in.

This should work just fine for a `KubernetesJob` as well, per https://docs.prefect.io/concepts/filesystems/#filesystem-specific-package-dependencies-in-docker-images it's just an environment variable that can be used in any infrastructure type that executes Docker images.
it doesn't work for KubernetesJob block somehow -- I can cross-check and post a traceback here if that's helpful
Yeah it doesn't make sense that it wouldn't work there unless it's bypassing the entrypoint somehow. If it's not working this should be re-opened as a bug report with the MRE.
ok, I'll try to get an MRE here. Thanks Michael!
Started an agent using:

```yaml
apiVersion: apps/v1
kind: Deployment
metadata:
  name: dataflowops
spec:
  selector:
    matchLabels:
      app: dataflowops
  replicas: 1
  template:
    metadata:
      labels:
        app: dataflowops
    spec:
      containers:
        - name: agent
          image: prefecthq/prefect:2-python3.9
          command: ["prefect", "agent", "start", "-q", "prod"]
          imagePullPolicy: "Always"
          env:
            - name: PREFECT_API_URL
              value: https://api.prefect.cloud/api/accounts/c5276cbb-62a2-4501-b64a-74d3d900d781/workspaces/aaeffa0e-13fa-460e-a1f9-79b53c05ab36
            - name: PREFECT_API_KEY
              value: xxxxx
---
apiVersion: rbac.authorization.k8s.io/v1
kind: Role
metadata:
  name: dataflowops-role
rules:
  - apiGroups: [""]
    resources: ["pods", "pods/log", "pods/status"]
    verbs: ["get", "watch", "list"]
  - apiGroups: ["batch"]
    resources: ["jobs"]
    verbs: [ "get", "list", "watch", "create", "update", "patch", "delete" ]
---
apiVersion: rbac.authorization.k8s.io/v1
kind: RoleBinding
metadata:
  name: dataflowops-role-binding
subjects:
  - kind: ServiceAccount
    name: default
roleRef:
  kind: Role
  name: dataflowops-role
  apiGroup: rbac.authorization.k8s.io
```

Created blocks:
```python
from prefect.infrastructure import KubernetesJob
from prefect.filesystems import S3

k8s = KubernetesJob(
    image="prefecthq/prefect:2-python3.9",
    image_pull_policy="Always",
    env={"EXTRA_PIP_PACKAGES": "s3fs"},
)
k8s.save("prod", overwrite=True)


s3 = S3(
    bucket_path="prefect-orion/prod",
    aws_access_key_id="xxx",
    aws_secret_access_key="xxx",
)
s3.save("prod", overwrite=True)
```

Then deployed my flow and flow run using:
```
prefect deployment build flows/healthcheck.py:healthcheck --name k8s -q prod -t poc -o deploy/k8s.yaml -ib kubernetes-job/prod -sb s3/prod
prefect deployment apply deploy/k8s.yaml
prefect deployment run healthcheck/k8s
```

Exception at runtime:
```
Flow could not be retrieved from deployment.
Traceback (most recent call last):
  File "/usr/local/lib/python3.9/site-packages/fsspec/registry.py", line 232, in get_filesystem_class
    register_implementation(protocol, _import_class(bit["class"]))
  File "/usr/local/lib/python3.9/site-packages/fsspec/registry.py", line 255, in _import_class
    mod = importlib.import_module(mod)
  File "/usr/local/lib/python3.9/importlib/__init__.py", line 127, in import_module
    return _bootstrap._gcd_import(name[level:], package, level)
  File "<frozen importlib._bootstrap>", line 1030, in _gcd_import
  File "<frozen importlib._bootstrap>", line 1007, in _find_and_load
  File "<frozen importlib._bootstrap>", line 984, in _find_and_load_unlocked
ModuleNotFoundError: No module named 's3fs'

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "/usr/local/lib/python3.9/site-packages/prefect/filesystems.py", line 328, in filesystem
    self._filesystem = fsspec.filesystem(scheme, **self.settings)
  File "/usr/local/lib/python3.9/site-packages/fsspec/registry.py", line 265, in filesystem
    cls = get_filesystem_class(protocol)
  File "/usr/local/lib/python3.9/site-packages/fsspec/registry.py", line 234, in get_filesystem_class
    raise ImportError(bit["err"]) from e
ImportError: Install s3fs to access S3

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "/usr/local/lib/python3.9/site-packages/prefect/engine.py", line 254, in retrieve_flow_then_begin_flow_run
    flow = await load_flow_from_flow_run(flow_run, client=client)
  File "/usr/local/lib/python3.9/site-packages/prefect/client.py", line 104, in with_injected_client
    return await fn(*args, **kwargs)
  File "/usr/local/lib/python3.9/site-packages/prefect/deployments.py", line 55, in load_flow_from_flow_run
    await storage_block.get_directory(from_path=None, local_path=".")
  File "/usr/local/lib/python3.9/site-packages/prefect/filesystems.py", line 399, in get_directory
    return await self.filesystem.get_directory(
  File "/usr/local/lib/python3.9/site-packages/prefect/filesystems.py", line 260, in get_directory
    return self.filesystem.get(from_path, local_path, recursive=True)
  File "/usr/local/lib/python3.9/site-packages/prefect/filesystems.py", line 331, in filesystem
    raise RuntimeError(
RuntimeError: File system created with scheme 's3' from base path 's3://prefect-orion/prod' could not be created. You are likely missing a Python module required to use the given storage protocol.
```

![image](https://user-images.githubusercontent.com/86264395/186538611-5977e8ca-b8bb-452c-8eb2-f092192984c2.png)
 
Prefect Version:
```
Version:             2.2.0
API version:         0.8.0
Python version:      3.9.12
Git commit:          e3651362
Built:               Tue, Aug 23, 2022 2:18 PM
OS/Arch:             darwin/arm64
Profile:             dev
Server type:         hosted
```
This was running on a local K8s cluster and could be related to architecture (ARM) perhaps?
I think the MRE could be as simple as:

```python
await KubernetesJob(
    command=["pip", "show", "s3fs"],
    env={"EXTRA_PIP_PACKAGES": "s3fs"},
).run()
```

Thanks for putting this together! This makes sense, I think. This appears to be because of how Kubernetes treats _commands_ vs _arguments_ which is in contrast to Docker's behavior: https://kubernetes.io/docs/tasks/inject-data-application/define-command-argument-container/ or https://stackoverflow.com/questions/44316361/difference-between-docker-entrypoint-and-kubernetes-container-spec-command

This seems relatively easy to fix in the short term by just switching from "command" to "args" in our template but I'm not sure of the full implications, especially for non-Prefect images.
See also: 
- https://github.com/PrefectHQ/prefect/issues/5195
- https://github.com/PrefectHQ/prefect/issues/5585

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
@@ -315,7 +315,7 @@ def _shortcut_customizations(self) -> JsonPatch:
             shortcuts.append(
                 {
                     "op": "add",
-                    "path": "/spec/template/spec/containers/0/command",
+                    "path": "/spec/template/spec/containers/0/args",
                     "value": self.command,
                 }
             )


Test patch:
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
@@ -637,7 +637,7 @@ def test_user_can_supply_a_sidecar_container_and_volume(self):
 
         # the prefect-job container is still populated
         assert pod["containers"][0]["name"] == "prefect-job"
-        assert pod["containers"][0]["command"] == ["echo", "hello"]
+        assert pod["containers"][0]["args"] == ["echo", "hello"]
 
         assert pod["containers"][1] == {
             "name": "my-sidecar",
@@ -767,23 +767,23 @@ def test_user_overriding_command_line(self):
             customizations=[
                 {
                     "op": "add",
-                    "path": "/spec/template/spec/containers/0/command/0",
+                    "path": "/spec/template/spec/containers/0/args/0",
                     "value": "opentelemetry-instrument",
                 },
                 {
                     "op": "add",
-                    "path": "/spec/template/spec/containers/0/command/1",
+                    "path": "/spec/template/spec/containers/0/args/1",
                     "value": "--resource_attributes",
                 },
                 {
                     "op": "add",
-                    "path": "/spec/template/spec/containers/0/command/2",
+                    "path": "/spec/template/spec/containers/0/args/2",
                     "value": "service.name=my-cool-job",
                 },
             ],
         ).build_job()
 
-        assert manifest["spec"]["template"]["spec"]["containers"][0]["command"] == [
+        assert manifest["spec"]["template"]["spec"]["containers"][0]["args"] == [
             "opentelemetry-instrument",
             "--resource_attributes",
             "service.name=my-cool-job",
@@ -791,6 +791,30 @@ def test_user_overriding_command_line(self):
             "hello",
         ]
 
+    def test_user_overriding_entrypoint_command(self):
+        """Users should be able to wrap the command-line with another command"""
+        manifest = KubernetesJob(
+            command=["echo", "hello"],
+            customizations=[
+                {
+                    "op": "add",
+                    "path": "/spec/template/spec/containers/0/command",
+                    "value": ["conda", "run", "-n", "foo"],
+                },
+            ],
+        ).build_job()
+
+        assert manifest["spec"]["template"]["spec"]["containers"][0]["command"] == [
+            "conda",
+            "run",
+            "-n",
+            "foo",
+        ]
+        assert manifest["spec"]["template"]["spec"]["containers"][0]["args"] == [
+            "echo",
+            "hello",
+        ]
+
 
 class TestLoadingManifestsFromFiles:
     @pytest.fixture