cuda-nvcc missing again

[dhruvbalwada]

It seems that the problem detected and solved in issue #387
has resurfaced again. I think this happened after #435 was merged.

The problem:

There is a ptxas based error that shows up. Can be easily reproduced as:

from jax import random
random.PRNGkey(0)

gives the error that

2023-02-01 19:08:39.849007: W external/org_tensorflow/tensorflow/compiler/xla/stream_executor/gpu/asm_compiler.cc:85] Couldn't get ptxas version string: INTERNAL: Couldn't invoke ptxas --version
2023-02-01 19:08:39.849939: F external/org_tensorflow/tensorflow/compiler/xla/service/gpu/nvptx_compiler.cc:454] ptxas returned an error during compilation of ptx to sass: 'INTERNAL: Failed to launch ptxas'  If the error message indicates that a file could not be written, please verify that sufficient filesystem space is provided.
Aborted

During the last discussion, @ngam had asked to check what version of cuda-nvcc existed. When I check this

conda list | grep cuda-nvcc

This returns nothing, showing that there is no cuda-nvcc in the tensorflow/jax based ml-notebook.

Installing cuda-nvcc by using mamba install cuda-nvcc==11.6.* -c nvidia solves the problem.

However, it would be good if the user did not have to manually do this installation, and the docker image was properly setup.

[scottyhq]

@dhruvbalwada I thought it was removed intentionally b/c no longer needed? See conversation here #398 ...

[dhruvbalwada]

Maybe @yuvipanda or @ngam or @weiji14 can chip in about why the problem has resurfaced?

[ngam]

It’s a complicated issue with all sorts of stuff. I think for now the best thing is to keep it out and let the user find a resolution. This is generally a tricky problem with, and mismatches are bound to happen.

The good news is that cuda-nvcc is coming to conda-forge soon; the bad news is that it’ll be a while before the lengthy migration effort concludes.

Xref:

• Build order for CUDA SDK conda-forge/staged-recipes#21382
• Adding cuda-nvcc recipe conda-forge/staged-recipes#21350 (comment)

[ngam]

Btw, thanks @dhruvbalwada for keeping an eye on this, and for the detailed report :)

[ngam]

Small update: This is finally getting resolved... hopefully very soon! xref #450

[weiji14]

Looks like cuda-nvcc is now on conda-forge - https://github.com/conda-forge/cuda-nvcc-feedstock. Is it better to install in directly in the ml-notebook image, or wait for the ML libraries like Tensorflow/Jax to depend on cuda-nvcc directly first? I see some mention of it e.g. at conda-forge/tensorflow-feedstock#296 (comment).

[ngam]

We should likely wait. I am still trying to assess how best to migrate Jax and TensorFlow to the new packaging format. We in a bit of a bind here... with volunteer maintainers occupied with other tasks... but tensorflow 2.12 is very close and I am making small progress on jaxlib.

[weiji14]

Someone reported on the forum at https://discourse.pangeo.io/t/how-to-run-code-using-gpu-on-pangeo-saying-libdevice-not-found-at-libdevice-10-bc/3672 about missing cuda-nvcc and XLA_FLAGS causing issues. Can we revisit adding cuda-nvcc to the docker image again, if the matter is resolved on conda-forge @ngam? @yuvipanda mentioned that 2i2c doesn't use the old K80 GPUs anymore, so we don't need to worry about backward compatibility if it helps.

[weiji14]

Quick note to say that jaxlib-0.4.23-cuda120py* actually has an explicit runtime dependency on cuda-nvcc now (see conda-forge/jaxlib-feedstock#241), but we'll need some more updates on tensorflow to resolve an incompatibility with libabseil versions. See #549 (comment), and keep an eye on conda-forge/tensorflow-feedstock#385.

Once those PRs are merged, users shouldn't have to install cuda-nvcc manually anymore, as they should be installed directly with jaxlib.