JAX-vLLM Offloading k8s AWS EKS #1798
Conversation
yhtang
left a comment
yhtang
left a comment
There was a problem hiding this comment.
Thanks for making this! Made some comments. Let me know what you think.
|
|
||
| # Configuration | ||
| NAMESPACE="${NAMESPACE:-default}" | ||
| JOBSET_NAME="jax-vllm-multinode" |
There was a problem hiding this comment.
We have these variables defined twice in separate files. Is there anyway that we can provide a single source of truth to avoid unintentional errors (e.g. edited in one place but not another)?
There was a problem hiding this comment.
Sure thing Yu-Hang. Consider this file just for a local development. I'll change this and have a clear definition in the CI
| effect: NoSchedule | ||
| containers: | ||
| - name: gateway-container | ||
| image: 941377147396.dkr.ecr.us-east-1.amazonaws.com/sbosisio/jio:jax-k8s |
There was a problem hiding this comment.
Just a note that we need to turn it into a placeholder and set it dynamically for each job in the final production workflow.
| echo "Gateway URL: ${GATEWAY_URL}" | ||
| echo "Ray Head IP: ${RAY_HEAD_IP}" | ||
|
|
||
| # 1. Wait for gateway to be ready |
There was a problem hiding this comment.
Noting that in the long run this could be integrated into the bridge.
@Steboss how long could the gap be, between the start of the gateway and the application (jax/vLLM) pods? Can we make the launch of the application pods dependent on the gateway pod?
| #NCCL | ||
| # - name: NCCL_DEBUG | ||
| # value: "INFO" # Change to WARN after debugging | ||
| - name: NCCL_PROTO |
There was a problem hiding this comment.
NCCL's official guidelines are to avoid setting this variable explicitly whenever possible. Is this mandated by AWS?
yhtang
left a comment
yhtang
left a comment
There was a problem hiding this comment.
Now that we have a working example, could we integrate this into jio.yaml?
Also, could we add a performance-monitoring step to this job so that if throughput drops below a certain baseline, the job reports a failure?
|
hey @yhtang
We can definitely make up the fully working example, just a caveat, we're still investigating why NCCL doesn't pick up EFA on EKS. If it's ok with you we can start with this approach, performance will be low, I'll give you some numbers by EOW at most. |
No description provided.