Update pixi configuration

[diegoferigo]

---

📚 Documentation preview 📚: https://jaxsim--281.org.readthedocs.build//281/

[diegoferigo]

Can you please check upon review if the current GPU configuration works on your setup? It seems that CUDA is detected and arrays can be allocated on GPUs, but then jax fails to jit-compile functions.

[flferretti]

Thanks Diego! I would keep the CI name of pixi e.g. drop 4d42133, as in the future we might think to use pixi for running the tests on CI and it may create confusion

[diegoferigo]

Thanks Diego! I would keep the CI name of pixi e.g. drop 4d42133, as in the future we might think to use pixi for running the tests on CI and it may create confusion

Reverting the name would be ok for me, just note that in this PR I also added a new step that runs the tests against the newly resolved pixi environment. For what regards CI/CD, I'd suggest to keep the functionality (regardless to the way dependencies are installed) in ci_cd.yml. What do you think?

[xela-95]

Thanks @diegoferigo ! LGTM!

[flferretti]

in this PR I also added a new step that runs the tests against the newly resolved pixi environment

This is great, thanks!

I'd suggest to keep the functionality (regardless to the way dependencies are installed) in ci_cd.yml. What do you think

Yes, it definitely makes sense to keep the two things in two separate workflows

[diegoferigo]

Great! I just need now someone that can confirm whether the resolved GPU environment fully works. I'm transitioning to a new setup and I don't have Nvidia GPUs available at the moment.

[flferretti]

I could solve the GPU environment with these modifications:

diff --git a/pyproject.toml b/pyproject.toml
index 1be8b1d..bb4fa88 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -258,13 +258,12 @@ robot_descriptions = "*"
 
 [tool.pixi.feature.gpu]
 platforms = ["linux-64"]
-system-requirements = { cuda = "12.1" }
+system-requirements = { cuda = "12"}
 
 [tool.pixi.feature.gpu.dependencies]
-cuda-version = "12.*"
 cuda-nvcc-tools = "*"
 cuda-nvvm-tools = "*"
-cudnn = "9.3.0.*"
+cudnn = "*"
 # Pinning a specific version awaiting the following to get addressed:
 # https://github.com/conda-forge/jaxlib-feedstock/issues/285
 jaxlib = { version = "<0.4.32", build = "*cuda*" }

[diegoferigo]

Does the following work on your setup? On mine it doesn't:

jaxsim on  update_pixi_confuguration via 🐍 v3.10.12
❯ pixi run --environment gpu python -c "import jax.numpy as jnp; jnp.linalg.norm(jnp.ones(3))"
Traceback (most recent call last):
  File "<string>", line 1, in <module>
jaxlib.xla_extension.XlaRuntimeError: INTERNAL: libdevice not found at ./libdevice.10.bc
--------------------
For simplicity, JAX has removed its internal frames from the traceback of the following exception. Set JAX_TRACEBACK_FILTERING=off to include these.

[flferretti]

Does the following work on your setup? On mine it doesn't:

jaxsim on  update_pixi_confuguration via 🐍 v3.10.12
❯ pixi run --environment gpu python -c "import jax.numpy as jnp; jnp.linalg.norm(jnp.ones(3))"
Traceback (most recent call last):
  File "<string>", line 1, in <module>
jaxlib.xla_extension.XlaRuntimeError: INTERNAL: libdevice not found at ./libdevice.10.bc
--------------------
For simplicity, JAX has removed its internal frames from the traceback of the following exception. Set JAX_TRACEBACK_FILTERING=off to include these.

Yes but I get some logs from PTXAS. It seems related to openxla/xla#16431 as I'm on CUDA 12.6:

$ pixi run --environment gpu python -c "import jax.numpy as jnp; jnp.linalg.norm(jnp.ones(3))"

 WARN The feature 'style' is defined but not used in any environment
 WARN The feature 'testing' is defined but not used in any environment
 WARN The feature 'viz' is defined but not used in any environment
 WARN The feature 'all' is defined but not used in any environment
'+ptx86' is not a recognized feature for this target (ignoring feature)
'+ptx86' is not a recognized feature for this target (ignoring feature)
'+ptx86' is not a recognized feature for this target (ignoring feature)
'+ptx86' is not a recognized feature for this target (ignoring feature)
'+ptx86' is not a recognized feature for this target (ignoring feature)
'+ptx86' is not a recognized feature for this target (ignoring feature)

[diegoferigo]

Can you please upload your pixi.lock? Also on my new setup, I get the same libdevice error.

[flferretti]

Can you please upload your pixi.lock? Also on my new setup, I get the same libdevice error.

You're right, it still does not work. I'm currently using jaxlib==0.4.28 in a working setup, maybe we could try to pin that as the upper limit

[diegoferigo]

Can you try now @flferretti? The resolved environment now seems working here after the last push.

[flferretti]

Can you try now @flferretti? The resolved environment now seems working here after the last push.

It works now, thanks!

[flferretti]

Thanks a lot Diego, LGTM!