amd-smi import error with rocm 6.1 + rocm/jax:latest image

[brettkoonce]

Description

docker pull rocm/jax
Using default tag: latest
latest: Pulling from rocm/jax
Digest: sha256:1d0238004c3049426a6a91e64fd0ef9f19c63e2418ceb34a4b1fcbc93e4f510e
Status: Image is up to date for rocm/jax:latest
docker.io/rocm/jax:latest

drun rocm/jax:latest
root@minerva:~# amd-smi
ModuleNotFoundError: No module named 'yaml'

See notes upstream: ROCm/amdsmi#19

Install:

python3 -m pip install 'PyYAML>=5.1'
root@minerva:~# amd-smi version
AMDSMI Tool: 23.4.2+3f96bd1 | AMDSMI Library version: 23.4.0.0

System info (python version, jaxlib version, accelerator, etc.)

ubuntu 22.04 lts, rocm 6.1, 7900 xtx

[superbobry]

Hi @brettkoonce, the ROCm repo looks like a better place for this issue. I don't think JAX directly depends on PyYAML.

[brettkoonce]

@superbobry Thx, forgot to close this!