Cannot install JAX through a proxy server

[thibmonsel]

Description

Hi, I'm trying to install JAX on a cluster via a proxy server.

The following CLI done is :
python3 -m pip install --proxy=MYPROXY --upgrade "jax[cuda11_pip]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
I do get the error :

Looking in links: https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
Collecting jax[cuda11_local]
  Using cached jax-0.4.8.tar.gz (1.2 MB)
  Installing build dependencies ... error
  error: subprocess-exited-with-error
  
  × pip subprocess to install build dependencies did not run successfully.
  │ exit code: 1
  ╰─> [13 lines of output]
      Looking in links: https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
      WARNING: Retrying (Retry(total=4, connect=None, read=None, redirect=None, status=None)) after connection broken by 'ConnectTimeoutError(<pip._vendor.urllib3.connection.HTTPSConnection object at 0x7f7b27e6f040>, 'Connection to 134.158.74.101 timed out. (connect timeout=15)')': /jax-releases/jax_cuda_releases.html
      WARNING: Retrying (Retry(total=3, connect=None, read=None, redirect=None, status=None)) after connection broken by 'ConnectTimeoutError(<pip._vendor.urllib3.connection.HTTPSConnection object at 0x7f7b27e6f250>, 'Connection to 134.158.74.101 timed out. (connect timeout=15)')': /jax-releases/jax_cuda_releases.html
      WARNING: Retrying (Retry(total=2, connect=None, read=None, redirect=None, status=None)) after connection broken by 'ConnectTimeoutError(<pip._vendor.urllib3.connection.HTTPSConnection object at 0x7f7b27e6f4f0>, 'Connection to 134.158.74.101 timed out. (connect timeout=15)')': /jax-releases/jax_cuda_releases.html
      WARNING: Retrying (Retry(total=1, connect=None, read=None, redirect=None, status=None)) after connection broken by 'ConnectTimeoutError(<pip._vendor.urllib3.connection.HTTPSConnection object at 0x7f7b27e6f5e0>, 'Connection to 134.158.74.101 timed out. (connect timeout=15)')': /jax-releases/jax_cuda_releases.html
      WARNING: Retrying (Retry(total=0, connect=None, read=None, redirect=None, status=None)) after connection broken by 'ConnectTimeoutError(<pip._vendor.urllib3.connection.HTTPSConnection object at 0x7f7b27e6f6d0>, 'Connection to 134.158.74.101 timed out. (connect timeout=15)')': /jax-releases/jax_cuda_releases.html
      WARNING: Retrying (Retry(total=4, connect=None, read=None, redirect=None, status=None)) after connection broken by 'ConnectTimeoutError(<pip._vendor.urllib3.connection.HTTPSConnection object at 0x7f7b27e6fa90>, 'Connection to 134.158.74.101 timed out. (connect timeout=15)')': /simple/setuptools/
      WARNING: Retrying (Retry(total=3, connect=None, read=None, redirect=None, status=None)) after connection broken by 'ConnectTimeoutError(<pip._vendor.urllib3.connection.HTTPSConnection object at 0x7f7b27ef69d0>, 'Connection to 134.158.74.101 timed out. (connect timeout=15)')': /simple/setuptools/
      WARNING: Retrying (Retry(total=2, connect=None, read=None, redirect=None, status=None)) after connection broken by 'ConnectTimeoutError(<pip._vendor.urllib3.connection.HTTPSConnection object at 0x7f7b27ef6910>, 'Connection to 134.158.74.101 timed out. (connect timeout=15)')': /simple/setuptools/
      WARNING: Retrying (Retry(total=1, connect=None, read=None, redirect=None, status=None)) after connection broken by 'ConnectTimeoutError(<pip._vendor.urllib3.connection.HTTPSConnection object at 0x7f7b27ef69a0>, 'Connection to 134.158.74.101 timed out. (connect timeout=15)')': /simple/setuptools/
      WARNING: Retrying (Retry(total=0, connect=None, read=None, redirect=None, status=None)) after connection broken by 'ConnectTimeoutError(<pip._vendor.urllib3.connection.HTTPSConnection object at 0x7f7b27e6f280>, 'Connection to 134.158.74.101 timed out. (connect timeout=15)')': /simple/setuptools/
      ERROR: Could not find a version that satisfies the requirement setuptools (from versions: none)
      ERROR: No matching distribution found for setuptools
      [end of output]
  
  note: This error originates from a subprocess, and is likely not a problem with pip.
error: subprocess-exited-with-error

× pip subprocess to install build dependencies did not run successfully.
│ exit code: 1
╰─> See above for output.

note: This error originates from a subprocess, and is likely not a problem with pip.

When I try CPU version for JAX
python3 -m pip install --proxy=MYPROXY --upgrade "jax[cpu]"

I get the same problem.

I have tried installing other packages on the cluster and had no issue at all.

I have installed JAX on my local machine and it worked fine.
Both setuptools are the same on my cluster and local machine setuptools==67.6.7

I would imagine that some build dependencies need an internet connection and does not take into account the argument --proxy in the CLI for the install.

What jax/jaxlib version are you using?

jax==0.4.8

Which accelerator(s) are you using?

CPU/GPU

Additional system info

Python 3.9.16, Linux

NVIDIA GPU info

No response

[jakevdp]

I think you're running into this issue: pypa/pip#6018

It looks like a bug in pip, over which JAX has no control. There are conflicting reports there, but perhaps you could try updating to the latest pip version to see if it fixes the issue?

[thibmonsel]

Thanks for the reply, my current pip is the latest release pip 23.0.1. I'll try to investigate further with the issue you mentionned !