Install on a Ubuntu server without admin rights

[astanziola]

Hello! This must be quite a recurring question, yet I can't find a way to solve it.

How can we install jax, with GPU support (hence, with appropriate cuda and cudnn) on a server without admin rights? In particular for an ubuntu server 😄

Of course, I can use tools like conda and similar (so a pytorch installation, for example, works fine and can run on the GPUs).

[hawkinsp]

I'd try the pip installation:
https://github.com/google/jax#pip-installation-gpu-cuda-installed-via-pip-easier

This does not need admin rights (beyond the ability to install a pip package, for which you can use a user-local installation).

The installation does require that you have the NVIDIA driver installed, but presumably your admin already did that, or you wouldn't be able to run anything on GPU. It also needs to be sufficiently new: if that's a problem the best option is to ask the admin to upgrade it.

[astanziola]

I do have cuda installed, as follows;

$ nvcc --version
nvcc: NVIDIA (R) Cuda compiler driver
Copyright (c) 2005-2022 NVIDIA Corporation
Built on Mon_Oct_24_19:12:58_PDT_2022
Cuda compilation tools, release 12.0, V12.0.76
Build cuda_12.0.r12.0/compiler.31968024_0

However, a fresh install of jax on a new environment gives me a DNN initialization error

$ pip install --upgrade "jax[cuda12_local]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html

... [ Install logs ] ...

$ python cudnn_test.py 
2023-08-21 14:48:23.456666: E external/xla/xla/stream_executor/cuda/cuda_dnn.cc:445] Could not create cudnn handle: CUDNN_STATUS_INTERNAL_ERROR
2023-08-21 14:48:23.457054: E external/xla/xla/stream_executor/cuda/cuda_dnn.cc:449] Memory usage: 6238306304 bytes free, 25633226752 bytes total.
Traceback (most recent call last):
  File "/home/astanziola/repos/holab/cudnn_test.py", line 5, in <module>
    m@m
    ~^~
  File "/home/astanziola/anaconda3/envs/hbfail/lib/python3.11/site-packages/jax/_src/numpy/array_methods.py", line 256, in deferring_binary_op
    return binary_op(*args)
           ^^^^^^^^^^^^^^^^
  File "/home/astanziola/anaconda3/envs/hbfail/lib/python3.11/site-packages/jax/_src/traceback_util.py", line 166, in reraise_with_filtered_traceback
    return fun(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^
  File "/home/astanziola/anaconda3/envs/hbfail/lib/python3.11/site-packages/jax/_src/pjit.py", line 253, in cache_miss
    outs, out_flat, out_tree, args_flat, jaxpr = _python_pjit_helper(
                                                 ^^^^^^^^^^^^^^^^^^^^
  File "/home/astanziola/anaconda3/envs/hbfail/lib/python3.11/site-packages/jax/_src/pjit.py", line 166, in _python_pjit_helper
    out_flat = pjit_p.bind(*args_flat, **params)
               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/astanziola/anaconda3/envs/hbfail/lib/python3.11/site-packages/jax/_src/core.py", line 2596, in bind
    return self.bind_with_trace(top_trace, args, params)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/astanziola/anaconda3/envs/hbfail/lib/python3.11/site-packages/jax/_src/core.py", line 389, in bind_with_trace
    out = trace.process_primitive(self, map(trace.full_raise, args), params)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/astanziola/anaconda3/envs/hbfail/lib/python3.11/site-packages/jax/_src/core.py", line 821, in process_primitive
    return primitive.impl(*tracers, **params)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/astanziola/anaconda3/envs/hbfail/lib/python3.11/site-packages/jax/_src/pjit.py", line 1209, in _pjit_call_impl
    return xc._xla.pjit(name, f, call_impl_cache_miss, [], [], donated_argnums,
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/astanziola/anaconda3/envs/hbfail/lib/python3.11/site-packages/jax/_src/pjit.py", line 1192, in call_impl_cache_miss
    out_flat, compiled = _pjit_call_impl_python(
                         ^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/astanziola/anaconda3/envs/hbfail/lib/python3.11/site-packages/jax/_src/pjit.py", line 1128, in _pjit_call_impl_python
    always_lower=False, lowering_platform=None).compile()
                                                ^^^^^^^^^
  File "/home/astanziola/anaconda3/envs/hbfail/lib/python3.11/site-packages/jax/_src/interpreters/pxla.py", line 2206, in compile
    executable = UnloadedMeshExecutable.from_hlo(
                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/astanziola/anaconda3/envs/hbfail/lib/python3.11/site-packages/jax/_src/interpreters/pxla.py", line 2544, in from_hlo
    xla_executable, compile_options = _cached_compilation(
                                      ^^^^^^^^^^^^^^^^^^^^
  File "/home/astanziola/anaconda3/envs/hbfail/lib/python3.11/site-packages/jax/_src/interpreters/pxla.py", line 2454, in _cached_compilation
    xla_executable = dispatch.compile_or_get_cached(
                     ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/astanziola/anaconda3/envs/hbfail/lib/python3.11/site-packages/jax/_src/dispatch.py", line 496, in compile_or_get_cached
    return backend_compile(backend, computation, compile_options,
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/astanziola/anaconda3/envs/hbfail/lib/python3.11/site-packages/jax/_src/profiler.py", line 314, in wrapper
    return func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/astanziola/anaconda3/envs/hbfail/lib/python3.11/site-packages/jax/_src/dispatch.py", line 464, in backend_compile
    return backend.compile(built_c, compile_options=options)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
jax._src.traceback_util.UnfilteredStackTrace: jaxlib.xla_extension.XlaRuntimeError: FAILED_PRECONDITION: DNN library initialization failed. Look at the errors above for more details.

The stack trace below excludes JAX-internal frames.
The preceding is the original exception that occurred, unmodified.

--------------------

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "/home/astanziola/repos/holab/cudnn_test.py", line 5, in <module>
    m@m
    ~^~
  File "/home/astanziola/anaconda3/envs/hbfail/lib/python3.11/site-packages/jax/_src/numpy/array_methods.py", line 256, in deferring_binary_op
    return binary_op(*args)
           ^^^^^^^^^^^^^^^^
jaxlib.xla_extension.XlaRuntimeError: FAILED_PRECONDITION: DNN library initialization failed. Look at the errors above for more details.

However, this method worked: #16547 (reply in thread)

[hawkinsp]

Exactly, that's why I suggested cuda12_pip, not cuda12_local. For a "local" install it's your responsibility to install all of the dependencies, and that's most easily done with admin rights.

[astanziola]

😳 oops, that's embarassing... I absolutely didn't see the _pip vs _local difference! And now I'm not even sure where I got the _local command from.

Anyhow, I've tested it an it works, thanks a lot for the patience!