[ROCm][CI] Add uv pip compile workflow for rocm-test.txt lockfile

[AndreasKaratzas]

Adds a proper uv pip compile workflow for generating rocm-test.txt from rocm-test.in, matching the existing CUDA test.in in to test.txt pattern.

Changes

• Added pip-compile rocm-test hook in .pre-commit-config.yaml that compiles rocm-test.in in to rocm-test.txt using -c requirements/rocm.txt as a constraint file (since --torch-backend rocm doesn't exist in uv). CUDA/NVIDIA packages are excluded via --no-emit-package since ROCm torch is installed separately via rocm.txt.
• Added a verification step that asserts torch.version.hip is not None to catch accidental CUDA torch installation during the test image build.
• Regenerated as a proper uv-compiled lockfile (346 pinned packages, zero CUDA/NVIDIA packages).

Key design decisions

• Uses -c requirements/rocm.txt constraint instead of --torch-backend (no ROCm backend in uv)
• Excludes torch/torchvision/torchaudio/triton and all nvidia-/cuda- packages from the lockfile output since they are managed by rocm.txt and Dockerfile.rocm_base

Will run a full build on ROCm CI to verify no new regressions are introduced.

cc @kenroche

[gemini-code-assist]

Code Review

This pull request introduces a uv pip compile workflow for ROCm test dependencies, aligning it with the existing CUDA setup. The changes include a new pre-commit hook, a verification step in the Dockerfile to ensure the correct PyTorch build is used, and regenerated lockfiles. My review focuses on improving the maintainability of the pre-commit configuration by simplifying the package exclusion list.

[gemini-code-assist]

The extensive list of --no-emit-package arguments for CUDA and NVIDIA packages can be significantly simplified. uv supports glob patterns for this option, which would make this configuration more concise and easier to maintain. You can replace the individual nvidia-*, cuda-*, and cupy-cuda* package exclusions with wildcard patterns.

        --no-emit-package, 'cuda-*',
        --no-emit-package, 'cupy-cuda*',
        --no-emit-package, 'nvidia-*',

[AndreasKaratzas]

Specificity is better I think.

[hmellor]

Wildcards would likely be more robust. As it is now every CUDA dependency change would break the ROCm build, which is annoying for CUDA developers and ROCm developers.

[AndreasKaratzas]

@hmellor So apparently uv does not support glob/wildcard patterns in --no-emit-package. I tested it and got:

error: invalid value 'nvidia-' for '--no-emit-package <NO_EMIT_PACKAGE>':
Not a valid package or extra name: "nvidia-". Names must start and end with
a letter or digit and may only contain -, _, ., and alphanumeric characters.
The --no-emit-package flag requires exact package names

But, I implemented a post processing routine with warning log to account for this concern. I moved the pattern check under there.

[mergify]

Hi @AndreasKaratzas, the pre-commit checks have failed. Please run:

uv pip install pre-commit>=4.5.1
pre-commit install
pre-commit run --all-files

Then, commit the changes and push to your branch.

For future commits, pre-commit will run automatically on changed files before each commit.

Tip

Is mypy failing?

mypy is run differently in CI. If the failure is related to this check, please use the following command to run it locally:

# For mypy (substitute "3.10" with the failing version if needed)
pre-commit run --hook-stage manual mypy-3.10

[mergify]

Hi @AndreasKaratzas, the pre-commit checks have failed. Please run:

uv pip install pre-commit>=4.5.1
pre-commit install
pre-commit run --all-files

Then, commit the changes and push to your branch.

For future commits, pre-commit will run automatically on changed files before each commit.

Tip

Is mypy failing?

mypy is run differently in CI. If the failure is related to this check, please use the following command to run it locally:

# For mypy (substitute "3.10" with the failing version if needed)
pre-commit run --hook-stage manual mypy-3.10

[hmellor]

not a fan of the custom script

[hmellor]

Updating the uv-pre-commit version allows us to specify rocm7.0 (which based on the image build logs is what we are using), so we could have:

- repo: https://github.com/astral-sh/uv-pre-commit
  rev: 0.11.1
  hooks:
    - id: pip-compile
      args: [requirements/test.in, -o, requirements/test.txt, --index-strategy, unsafe-best-match, --torch-backend, cu129, --python-platform, x86_64-manylinux_2_28, --python-version, "3.12"]
      files: ^requirements/test\.(in|txt)$
    - id: pip-compile
      name: pip-compile-rocm
      args: [requirements/rocm-test.in, -c, requirements/rocm.txt, -o, requirements/rocm-test.txt, --index-strategy, unsafe-best-match, --torch-backend, rocm7.0, --python-platform, x86_64-manylinux_2_28, --python-version, "3.12"]
      files: ^requirements/rocm-test\.(in|txt)$

This almost works but there are conflicts in the provided requirements:

pip-compile-rocm.........................................................Failed
- hook id: pip-compile
- exit code: 1

  × No solution found when resolving dependencies:
  ╰─▶ Because there is no version of triton-rocm{sys_platform
      == 'linux'}==3.6.0 and torch==2.10.0+rocm7.0 depends on
      triton-rocm{sys_platform == 'linux'}==3.6.0, we can conclude that
      torch==2.10.0+rocm7.0 cannot be used.
      And because bitsandbytes==0.49.2 depends on torch>=2.3 and only the
      following versions of torch are available:
          torch<2.8.0
          torch==2.10.0+rocm7.0
      we can conclude that bitsandbytes==0.49.2 depends on torch>=2.3,<2.8.0.
      And because instanttensor>=0.1.5 depends on torch>=2.8.0 and only the
      following versions of instanttensor are available:
          instanttensor<=0.1.5
          instanttensor==0.1.6
          instanttensor==0.1.7
      we can conclude that bitsandbytes==0.49.2 and instanttensor>=0.1.5 are
      incompatible.
      And because you require bitsandbytes==0.49.2 and instanttensor>=0.1.5,
      we can conclude that your requirements are unsatisfiable.

[AndreasKaratzas]

not a fan of the custom script

I reverted the custom script :)

[AndreasKaratzas]

Updating the uv-pre-commit version allows us to specify rocm7.0 (which based on the image build logs is what we are using), so we could have:

- repo: https://github.com/astral-sh/uv-pre-commit
  rev: 0.11.1
  hooks:
    - id: pip-compile
      args: [requirements/test.in, -o, requirements/test.txt, --index-strategy, unsafe-best-match, --torch-backend, cu129, --python-platform, x86_64-manylinux_2_28, --python-version, "3.12"]
      files: ^requirements/test\.(in|txt)$
    - id: pip-compile
      name: pip-compile-rocm
      args: [requirements/rocm-test.in, -c, requirements/rocm.txt, -o, requirements/rocm-test.txt, --index-strategy, unsafe-best-match, --torch-backend, rocm7.0, --python-platform, x86_64-manylinux_2_28, --python-version, "3.12"]
      files: ^requirements/rocm-test\.(in|txt)$

This almost works but there are conflicts in the provided requirements:

pip-compile-rocm.........................................................Failed
- hook id: pip-compile
- exit code: 1

  × No solution found when resolving dependencies:
  ╰─▶ Because there is no version of triton-rocm{sys_platform
      == 'linux'}==3.6.0 and torch==2.10.0+rocm7.0 depends on
      triton-rocm{sys_platform == 'linux'}==3.6.0, we can conclude that
      torch==2.10.0+rocm7.0 cannot be used.
      And because bitsandbytes==0.49.2 depends on torch>=2.3 and only the
      following versions of torch are available:
          torch<2.8.0
          torch==2.10.0+rocm7.0
      we can conclude that bitsandbytes==0.49.2 depends on torch>=2.3,<2.8.0.
      And because instanttensor>=0.1.5 depends on torch>=2.8.0 and only the
      following versions of instanttensor are available:
          instanttensor<=0.1.5
          instanttensor==0.1.6
          instanttensor==0.1.7
      we can conclude that bitsandbytes==0.49.2 and instanttensor>=0.1.5 are
      incompatible.
      And because you require bitsandbytes==0.49.2 and instanttensor>=0.1.5,
      we can conclude that your requirements are unsatisfiable.

Investigated using --torch-backend rocm7.0 as you suggested. It resolves cleanly when paired with an override for triton-rocm. However, --torch-backend requires prebuilt PyTorch wheels on download.pytorch.org/whl/, and our ROCm Docker builds compile torch from source. If we target a ROCm version that PyTorch hasn't published wheels for yet, the hook would fail to resolve. Keeping the explicit --no-emit-package list for now, but applied your naming feedback (pip-compile-rocm with alias). Happy to revisit once PyTorch ROCm wheels are reliably published.

[hmellor]

LGTM, it should be fairly easy to update the --no-emit-package rules and once proper ROCm wheels are distributed we should be able to use --torch-backend instead