Update to Pytorch 2.8 + cu129, update docker containers to 24.04 (#362)

* Pytorch 2.8 + cu129

* Recombine jax and torch env

* py3.12_torch2.8+cu129 (#364)

Co-authored-by: pierre.delaunay <[email protected]>

* Do not init process group on prepare step

* Add Distributed env variable on prepare when required

* Udpate purejaxrl to use the latest distrax + tfp-nightly

* ignore tensorflow-probability

* Pin Dependencies (#366)

* ignore tensorflow-probability

* Pin Dependencies

---------

Co-authored-by: pierre.delaunay <[email protected]>

* replace tree_map by tree.map

* update container to match current cuda version

* Update dockerfile

* Add a timer for rsync and add concept of job runner pipeline

* Handle rerun of jobs with dependencies

* make the server try its best even if the cluster is down

---------

Co-authored-by: pierre.delaunay <[email protected]>

Author: Delaunay

.pin/constraints-cuda-torch.txt

@@ -393,8 +393,8 @@ def convert_data(x: jnp.ndarray):
                 x = jnp.reshape(x, (args.num_minibatches, -1) + x.shape[1:])
                 return x
 
-            flatten_storage = jax.tree_map(flatten, storage)
-            shuffled_storage = jax.tree_map(convert_data, flatten_storage)
+            flatten_storage = jax.tree.map(flatten, storage)
+            shuffled_storage = jax.tree.map(convert_data, flatten_storage)
 
             def update_minibatch(agent_state, minibatch):
                 (loss, (pg_loss, v_loss, entropy_loss, approx_kl)), grads = ppo_loss_grad_fn(