Update WorkPool's Job Configuration

[Blarc]

When we create a work pool CRD the operator creates a work pool deployment. The worker that is created by the deployment creates a work pool on Prefect's server on startup. The work pool's job configuration uses default values which is problematic. The CRD should allow setting the job configuration.

If we update the job configuration via Prefect's server UI, the following request body is sent:

{
    "base_job_template": {
        "job_configuration": {
            "command": "{{ command }}",
            "env": "{{ env }}",
            "labels": "{{ labels }}",
            "name": "{{ name }}",
            "namespace": "{{ namespace }}",
            "job_manifest": {
                "apiVersion": "batch/v1",
                "kind": "Job",
                "metadata": {
                    "generateName": "{{ name }}-",
                    "labels": "{{ labels }}",
                    "namespace": "{{ namespace }}"
                },
                "spec": {
                    "backoffLimit": 0,
                    "template": {
                        "spec": {
                            "completions": 1,
                            "containers": [
                                {
                                    "args": "{{ command }}",
                                    "env": "{{ env }}",
                                    "image": "{{ image }}",
                                    "imagePullPolicy": "{{ image_pull_policy }}",
                                    "name": "prefect-job"
                                }
                            ],
                            "parallelism": 1,
                            "restartPolicy": "Never",
                            "serviceAccountName": "{{ service_account_name }}"
                        }
                    },
                    "ttlSecondsAfterFinished": "{{ finished_job_ttl }}"
                }
            },
            "cluster_config": "{{ cluster_config }}",
            "job_watch_timeout_seconds": "{{ job_watch_timeout_seconds }}",
            "pod_watch_timeout_seconds": "{{ pod_watch_timeout_seconds }}",
            "stream_output": "{{ stream_output }}"
        },
        "variables": {
            "description": "Default variables for the Kubernetes worker.\n\nThe schema for this class is used to populate the `variables` section of the default\nbase job template.",
            "properties": {
                "name": {
                    "anyOf": [
                        {
                            "type": "string"
                        },
                        {
                            "type": "null"
                        }
                    ],
                    "description": "Name given to infrastructure created by a worker.",
                    "title": "Name"
                },
                "env": {
                    "additionalProperties": {
                        "anyOf": [
                            {
                                "type": "string"
                            },
                            {
                                "type": "null"
                            }
                        ]
                    },
                    "description": "Environment variables to set when starting a flow run.",
                    "title": "Environment Variables",
                    "type": "object"
                },
                "labels": {
                    "additionalProperties": {
                        "type": "string"
                    },
                    "description": "Labels applied to infrastructure created by a worker.",
                    "title": "Labels",
                    "type": "object"
                },
                "command": {
                    "anyOf": [
                        {
                            "type": "string"
                        },
                        {
                            "type": "null"
                        }
                    ],
                    "description": "The command to use when starting a flow run. In most cases, this should be left blank and the command will be automatically generated by the worker.",
                    "title": "Command"
                },
                "namespace": {
                    "default": "not-default",
                    "description": "The Kubernetes namespace to create jobs within.",
                    "title": "Namespace",
                    "type": "string"
                },
                "image": {
                    "anyOf": [
                        {
                            "type": "string"
                        },
                        {
                            "type": "null"
                        }
                    ],
                    "description": "The image reference of a container image to use for created jobs. If not set, the latest Prefect image will be used.",
                    "examples": [
                        "docker.io/prefecthq/prefect:3-latest"
                    ],
                    "title": "Image"
                },
                "service_account_name": {
                    "anyOf": [
                        {
                            "type": "string"
                        },
                        {
                            "type": "null"
                        }
                    ],
                    "description": "The Kubernetes service account to use for job creation.",
                    "title": "Service Account Name"
                },
                "image_pull_policy": {
                    "default": "IfNotPresent",
                    "description": "The Kubernetes image pull policy to use for job containers.",
                    "enum": [
                        "IfNotPresent",
                        "Always",
                        "Never"
                    ],
                    "title": "Image Pull Policy",
                    "type": "string"
                },
                "finished_job_ttl": {
                    "anyOf": [
                        {
                            "type": "integer"
                        },
                        {
                            "type": "null"
                        }
                    ],
                    "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 not set, jobs will be retained indefinitely.",
                    "title": "Finished Job TTL"
                },
                "job_watch_timeout_seconds": {
                    "anyOf": [
                        {
                            "type": "integer"
                        },
                        {
                            "type": "null"
                        }
                    ],
                    "description": "Number of seconds to wait for each event emitted by a job before timing out. If not set, the worker will wait for each event indefinitely.",
                    "title": "Job Watch Timeout Seconds"
                },
                "pod_watch_timeout_seconds": {
                    "default": 60,
                    "description": "Number of seconds to watch for pod creation before timing out.",
                    "title": "Pod Watch Timeout Seconds",
                    "type": "integer"
                },
                "stream_output": {
                    "default": true,
                    "description": "If set, output will be streamed from the job to local standard output.",
                    "title": "Stream Output",
                    "type": "boolean"
                },
                "cluster_config": {
                    "anyOf": [
                        {
                            "$ref": "#/definitions/KubernetesClusterConfig"
                        },
                        {
                            "type": "null"
                        }
                    ],
                    "description": "The Kubernetes cluster config to use for job creation."
                }
            },
            "type": "object",
            "definitions": {
                "KubernetesClusterConfig": {
                    "block_schema_references": {},
                    "block_type_slug": "kubernetes-cluster-config",
                    "description": "Stores configuration for interaction with Kubernetes clusters.\n\nSee `from_file` for creation.",
                    "properties": {
                        "config": {
                            "description": "The entire contents of a kubectl config file.",
                            "title": "Config",
                            "type": "object"
                        },
                        "context_name": {
                            "description": "The name of the kubectl context to use.",
                            "title": "Context Name",
                            "type": "string"
                        }
                    },
                    "required": [
                        "config",
                        "context_name"
                    ],
                    "secret_fields": [],
                    "title": "KubernetesClusterConfig",
                    "type": "object"
                }
            }
        }
    }
}

I was about to propose to add new fields to CRD that would allow configuring the job configuration, however, this API is quite annoying to deal with. I do think this is quite important though, because, among other things (service account name etc.), this is also how we can set the namespace where the jobs will be created.

How should we tackle this? Do we need to create a struct for the whole base_job_template ? Is there a better way?

[tinkerborg]

I needed this too, will be opening a PR later today

[Blarc]

@tinkerborg How are you planning to implement this? It would maybe make sense to first generate Prefect API client from OpenAPI schema as mentioned in #188.

[tinkerborg]

@tinkerborg How are you planning to implement this? It would maybe make sense to first generate Prefect API client from OpenAPI schema as mentioned in #188.

I followed the existing pattern currently used in this codebase. I didn't want to introduce new patterns not directly related to the feature being implemented. PR is open as a draft now, will work on implementing unit tests today then convert to normal PR.

[Blarc]

@tinkerborg How are you planning to implement this? It would maybe make sense to first generate Prefect API client from OpenAPI schema as mentioned in #188.

I followed the existing pattern currently used in this codebase. I didn't want to introduce new patterns not directly related to the feature being implemented. PR is open as a draft now, will work on implementing unit tests today then convert to normal PR.

Aha, you'll take the existing base_job_template and merge it with the object that will be provided by the CRD?

[tinkerborg]

@tinkerborg How are you planning to implement this? It would maybe make sense to first generate Prefect API client from OpenAPI schema as mentioned in #188.

I followed the existing pattern currently used in this codebase. I didn't want to introduce new patterns not directly related to the feature being implemented. PR is open as a draft now, will work on implementing unit tests today then convert to normal PR.

Aha, you'll take the existing base_job_template and merge it with the object that will be provided by the CRD?

I went back and forth on that and ultimately went with the same behavior as the helm chart, which expects you to provide a new template in its entirety. Hmm, maybe I should add a baseJobTemplatePatches field to provide the option of patching the default one as well?

[Blarc]

I went back and forth on that and ultimately went with the same behavior as the helm chart, which expects you to provide a new template in its entirety. Hmm, maybe I should add a baseJobTemplatePatches field to provide the option of patching the default one as well?

That's the main issue I had with implementing this. The problem is, you need to send the whole base_job_template json that is mentioned above, which is too big to add whenever I want to configure few things like namespace and service account.

I think baseJobTemplatePatches is a much better option. However, I think it would be even better if the CRD specified all the fields that can also be set via the UI and then create the base_job_template based on that. This way you get auto-completion which I think is really important.

It just seems so wierd that the API accepts a generic JSON object. I would like the API to specify the JSON schema which we could use to generate the struct in Go. This would make things way better in my opinion.

I'd love to hear what @chrisguidry has to say about this.

[chrisguidry]

The challenge with the base_job_template for Prefect workers in general is that they don't all have the same job configurations. KubernetesWorker expects one set of fields, while the GCP Cloud Run expects another, etc. This is why the API is open-ended and doesn't have a fixed schema. But it leads to really challenging UX issues like this one you're pointing out.

@tinkerborg's approach is definitely the best "steam valve" here, where it can allow for any type of work pool (even if it's not a kubernetes one), but that approach isn't mutually exclusive with creating some kubernetes-specific shortcuts like you're describing. This could be a purely client-side feature in the operator itself, and could use a pre-defined schema based on the KubernetesWorker's job configuration, like:

apiVersion: prefect.io/v1
kind: PrefectWorkPool
metadata:
  namespace: where-my-worker-runs
  name: my-k8s-work-pool
spec:
  type: kubernetes
  server:
    namespace: where-my-server-runs
    name: prefect-ephemeral
  workers: 3
  kubernetesJobConfiguration:
    namespace: where-my-prefect-jobs-run
    command: "my-entry-point prefect flow-run execute"
    env:
      HELLO: "world"
      GOODBYE: "moon"
    labels:
      example.com/mything: one
      example.com/myotherthing: two
    ...

The base job template is carrying both the job configuration and the definition of the variables (in OpenAPI format), which is what drives the UI, but often all you want to do is tweak one of the job configuration values (like the namespace or environment variables), but we do kind of require you to provide the whole thing. With this option, we can query the server for the default job template after it's created, then stitch in these values to perform a subsequent update. It's a little complicated, but it could be a nice UX layer on top of also supporting a full replacement of the baseJobTemplate.

[Blarc]

Just thinking out loud here, but would it maybe make more sense to make separate CRDs for each type of work pool then? For example: PrefectKubernetesWorkPool, PrefectGCPWorkPool etc.?

[chrisguidry]

It's a good thought, for sure. The challenge there is similar to the general challenge we have with work pool configurations in Prefect: it's intentionally open-ended so that people can theoretically implement new workers just by subclassing our Worker base classes, and it might be in libraries that Prefect doesn't even know about. I do think adding typed work pool CRDs could definitely ease the pain, but we'll still need a way to deal with full base_job_templates.

Another idea occurred to me that could be something of a compromise. Building off of that last idea, I've often seen CRDs where the presence/absence of a key in the CRD indicates the type of a thing. We do something like that with the sqlite versus postgres configuration for PrefectServer:

apiVersion: prefect.io/v1
kind: PrefectWorkPool
metadata:
  namespace: where-my-worker-runs
  name: my-k8s-work-pool
spec:
  # no type: kubernetess here
  server:
    namespace: where-my-server-runs
    name: prefect-ephemeral
  workers: 3
  # but there is a kubernetes: key here that indicates that it's a kubernetes 
  # work pool and lets us have typed configuration for kubernetes
  kubernetes:
    namespace: where-my-prefect-jobs-run
    command: "my-entry-point prefect flow-run execute"
    env:
      HELLO: "world"
      GOODBYE: "moon"
    labels:
      example.com/mything: one
      example.com/myotherthing: two
    ...

Just spitballing here too….

[Blarc]

But even though the API is open-ended, you still need to specify the type (e.g kubernetes), so the correct image is used and the base_job_templates depends on the type, does it not? So if I understand correctly, the operator is currently not as "open-ended" as the API? Do we want it to be? Would in that case need to allow a type generic?

I do like your suggestion, although I still think it's cleaner to have multiple CRDs. What would happen if a user defined multiple types in the same CRD? Since the CRD schema would need to allow it, if we want to use one CRD for different types? Operator would probably need to throw an error, that only one type can be defined. Maybe it's not that bad.

[chrisguidry]

You're definitely right, the operator is currently not supporting all work pool types and that's been a long-standing issue I've been wanting to resolve, and you two are hitting on all the thorny issues that have slowed us from doing this 😀

Yeah I've seen this pattern with a few CRDs where there's kind of a wide variety of options for something, and a great example is the PersistentVolume that has a truly bewildering number of implementations (like nfs, iscsi, a bunch of cloud provider-specific types) that could be subclasses in other environments, but that doesn't fit the K8s CRD pattern very well.

I think no matter what, we can't support truly open-ended work pool types unless we can figure out the image/python package issue described in #30 as well as this issue of configuration. Honestly, I think @tinkerborg's approach is the closest thing we'll get, but we can design some shortcuts for specific work pool types that are very common (process and kubernetes are by far the most likely to be used with the operator, I'd think)

[tinkerborg]

@chrisguidry Do you think it makes more sense to have a shortcut specific to kubernetes workers, and strongly type the schema there to allow for strategic merge patches, or have a more generic baseJobTemplatePatches shortcut, using JSON templates? I think json patch isn't as pretty, but it saves a lot of complexity, lets users modify the variables definitions, and doesn't require changes for new worker types. JSON merge is generic too but would require redefining the entire list when adding items to an existing list, e.g. if someone wanted to add an additional container for whatever reason. That kind of defeats the purpose.

I fixed the failing tests in my PR, working on adding unit tests for the new functionality but won't have time to finish til this weekend.

[chrisguidry]

Much appreciated, the PR looks good so far!

We kind of have a precedent for using JSON 6902 patches for job configurations in the Kubernetes worker, the GCP Cloud Run workers v1 and v2, and the GCP Vertex AI worker, but it's been kind of obscure for most developers. It's fairly well standardized, but I don't think too many people are familiar with it.

It does show up fairly prominently in kustomize, so it's definitely more appropriate for a Kubernetes operator and folks familiar with the K8s ecosystem, I'd say.

I'm going to run this by Claude and have it produce alternatives for us to consider, will post them on this issue.