Pass custom k8s container properties

[clumsy]

Description

k8s overlay does not allow modifying the container that TorchX itself creates (we can only append to the list of containers).

Motivation/Background

We would like to be able to set livenessProbe and readinessProbe properties as per: https://kubernetes.io/docs/tasks/configure-pod-container/configure-liveness-readiness-startup-probes

Detailed Proposal

We can maybe extend overlay support to overwrite/merge list elements or allow passing container kwargs as is.

Alternatives

Custom k8s scheduler that alters container/spec after torchx call.

Additional context/links

https://kubernetes.io/docs/tasks/configure-pod-container/configure-liveness-readiness-startup-probes

[clumsy]

Would appreciate your guidance, @kiukchung

[kiukchung]

can you give me an example of the k8s payload torchx creates and which one you want to edit?

[thomaslauer]

Hey @kiukchung, here's an example of what we're trying to add

What torchx creates right now:

spec:
  containers:
    - name: server-0
      image: vllm/vllm-openai:latest
      command: ["python", "-m", "vllm.entrypoints.openai.api_server", "--model", "..."]

What we'd like:

spec:
  containers:
    - name: server-0
      image: vllm/vllm-openai:latest
      command: ["python", "-m", "vllm.entrypoints.openai.api_server", "--model", "..."]
      livenessProbe:
        httpGet:
          path: /health
          port: 8000
        initialDelaySeconds: 120
        periodSeconds: 30
      readinessProbe:
        httpGet:
          path: /health
          port: 8000
        initialDelaySeconds: 60
        periodSeconds: 10

We're looking for a way to let k8s handle the health checks for vllm servers in RL rollout. We attempted to add this through through _apply_pod_overlay, but it tires to merge lists by concatenating them which won't let us add attributes to existing containers.

Any thoughts on how to best handle this? We're getting around it by injecting them via a custom scheduler, but it would be convenient to be supported natively.

[kiukchung]

spec:
containers:
- name: server-0
image: vllm/vllm-openai:latest
command: ["python", "-m", "vllm.entrypoints.openai.api_server", "--model", "..."]
livenessProbe:
httpGet:
path: /health
port: 8000
initialDelaySeconds: 120
periodSeconds: 30
readinessProbe:
httpGet:
path: /health
port: 8000
initialDelaySeconds: 60
periodSeconds: 10

thanks for the example. I see the issue now. This is really happening because we don't have a primitive that maps to "multi-container" Roles.

A quick (but dirty) workaround for not having to use a custom scheduler is to define the side-car container as the AppDef's Role and use the overlay to define the main container (vllm/vllm-openai:latest) this way you have total control over applying the healthcheck parameters.

I'd like to explore two approaches here:

1. Adding health-check attributes to Role
2. Re-defining how overlays work (e.g. allowing list overlay logic to update an existing element by matching the name attribute)