Question: localcolabfold on Nvidia T4 and COLAB_TPU_ADDR

[darthmodel]

What is your question? I'm new to localcolabfold. I installed it on an AWS instance with an Nvidia T4 GPU. The installation went smoothly. However, it complains about not being able to find TPU. I'm confused that why the code would look for a TPU, instead of the GPU there. Is there a way to force it to use the GPU? Thanks!

2023-06-15 19:45:01,406 Running colabfold 1.5.2 (affc6e0e89ee723795ab009f51f2600d843dafb0)                                                            Traceback (most recent call last):
  File "/scratch/colabfold/localcolabfold/colabfold-conda/lib/python3.10/site-packages/colabfold/batch.py", line 1197, in run
    jax.tools.colab_tpu.setup_tpu()                                                                                                                     File "/scratch/colabfold/localcolabfold/colabfold-conda/lib/python3.10/site-packages/jax/tools/colab_tpu.py", line 38, in setup_tpu
    colab_tpu_addr = os.environ['COLAB_TPU_ADDR'].split(':')[0]
  File "/scratch/colabfold/localcolabfold/colabfold-conda/lib/python3.10/os.py", line 680, in __getitem__
    raise KeyError(key) from None
KeyError: 'COLAB_TPU_ADDR'

Computational environment

• OS: AlmaLinux 8.7
• CUDA version

nvcc: NVIDIA (R) Cuda compiler driver
Copyright (c) 2005-2023 NVIDIA Corporation
Built on Mon_Apr__3_17:16:06_PDT_2023
Cuda compilation tools, release 12.1, V12.1.105
Build cuda_12.1.r12.1/compiler.32688072_0

G Cui

[vprobon]

Hi @YoshitakaMo, thank you for making localcolabfold available!

I also get the same error on our HPC system.

Computational environment

• OS: CentOS Linux release 7.8.2003 (Core)
• CUDA version

nvcc: NVIDIA (R) Cuda compiler driver Copyright (c) 2005-2021 NVIDIA Corporation Built on Mon_May__3_19:15:13_PDT_2021 Cuda compilation tools, release 11.3, V11.3.109 Build cuda_11.3.r11.3/compiler.29920130_0

Many thanks, Vasilis

[darthmodel]

Hi @vprobon,

Not sure if this helps, but I figured out the reason I was getting this is how my pip was configured. It wasn't pulling the packages from PyPi.

Cheers, Cui

[vprobon]

Hi @darthmodel and thank you for the tip. I do not know if this applies to my case but I will try to find out ...

Best, Vasilis