[tinker] Support using Ray to manage tinker API server, engine and Jax workers

[andrewsykim]

I recently spent some time trying to get SkyRL-tx working on a Ray cluster. I initially started with an approach similar to #955, but I don't think it's quite what we want because it still requires running the uv run command on Ray nodes with GPUs / TPUs. Instead, it should be possible to run an entrypoint script on the head node that will then distribute Ray actors, similar to how skyrl-train works today.

This work may require some larger refactoring in skyrl/tinker/api.py, due to the way the engine parameters are extracted from the parent uv command used for the tinker API server, and because the tinker API server / engine need to run on the same Ray node as worker0. We may need to decouple how the arguments for API server / engine are set and add some additional logic to ensure the tinker apiserver / engine and worker 0 are co-located. I think the end result should be something similar to the existing main entrypoint for skyrl-train but with some added logic for scheduling the tinker API server / engine.

I'm happy to drive this but looking for some guidance on the overall approach.

[andrewsykim]

@erictang000 @pcmoritz interested to get your feedback on how to approach this.

[pcmoritz]

Thanks for creating this issue @andrewsykim !

The way it currently works with the skyrl-train backend is that (API server, engine) run as a Ray driver, and the engine gets connected to the Ray Cluster and execute the training there. We can mirror this design in the Jax backend by introducing a new class like https://github.com/NovaSky-AI/SkyRL/blob/main/skyrl/backends/jax.py#L1070 (e.g. RayJaxBackend) which forwards the calls via Ray actor calls from the engine to rank 0, and the rank 0 would be the actor instantiating the JaxBackendImpl that does the actual work https://github.com/NovaSky-AI/SkyRL/blob/main/skyrl/backends/jax.py#L183). This would decouple the components that don't need accelerators from the components that do, and still allow API server and engine to run on the same node (which is required if we want to support the sqlite DB).

The downside is that then the Ray and no Ray option run slightly different code, but I feel like that is ok since the difference is quite small and it seems better than the alternative solution of running API server and engine on one of the worker nodes that has an accelerator. There might also be a performance hit, but we also take that hit for skyrl-train, so it is probably ok.

What do you think about this?

[andrewsykim]

@pcmoritz makes sense to me! I explored this option a bit but I stopped due to the diverging code paths you mentioned. I will continue working on this approach if we're okay with the Ray / non-Ray code paths diverging a little bit.