CI: example job using EKS jumphost #1151
Conversation
.github/workflows/_ci.yaml
Outdated
| --output=/opt/output/rank%q{JOB_COMPLETION_INDEX}.zip \ | ||
| -- \ | ||
| jax-nccl-test \ | ||
| --coordinator-address \ |
There was a problem hiding this comment.
Why not use a JobSet? I've added automatic bootstrapping support for JAX in K8s jobset. A jobset provides a higher-level abstraction and can save the ClusterIP and bootstrapping boilerplates.
There was a problem hiding this comment.
For example of how to run a jobset, see https://github.com/jax-ml/jax/pull/24197/files
There was a problem hiding this comment.
Given that we are not currently exposing the code that configures the Kubenetes cluster itself, I was reluctant to go in a direction where the public code (CI config) does not work without non-trivial private logic (JobSet installation, service account configuration, ...)
Also, the bootstrapping doesn't work out of the box because it relies on Python code that is not shipped in our containers. Pursuing that felt like a tangent to the core of this PR.
There was a problem hiding this comment.
non-trivial private logic (JobSet installation, service account configuration, ...)
The service account logic is pretty straightforward --- it's not gonna rank in the top N most hairy K8s configs in the world. Plus, we already provided a config that works OOTB!
It's one line to install jobset:
kubectl apply --server-side -f https://github.com/kubernetes-sigs/jobset/releases/download/v0.7.1/manifests.yaml
That said, given that Jobs are 'built-in', maybe we should consider making jax.distributed.initialize automatic under it too.
There was a problem hiding this comment.
not currently exposing the code that configures the Kubenetes cluster
I wouldn't view the read-only service account and the installation of Jobset as the sensitive information of a cluster. It's the equivalence of one usermod -aG ... and one apt install for a SLURM cluster.
There was a problem hiding this comment.
I'm not saying we can't expose it because it's sensitive (it's not). But it's not currently exposed, and it won't be in this PR, so it can't be referred to. It would need to be checked what else is needed re: the AWS/EKS-specific parts of how the roles are managed, too.
Plus, we already provided a config that works OOTB!
You mean the minikube CI PR to JAX? That's not really discoverable?
The error message still includes a placeholder, BTW: https://github.com/yhtang/jax/blob/2f67710e8ce9cd6c872f85e57417aa9bba7aa270/jax/_src/clusters/k8s_cluster.py#L70
No description provided.