Skip to content

fix(docker): set TMS_CUDA_MAJOR for torch_memory_saver source build - #1774

Merged
Zhichenzzz merged 1 commit into
mainfrom
zhichen/fix-docker-tms-cuda-major
Jul 23, 2026
Merged

fix(docker): set TMS_CUDA_MAJOR for torch_memory_saver source build#1774
Zhichenzzz merged 1 commit into
mainfrom
zhichen/fix-docker-tms-cuda-major

Conversation

@Zhichenzzz

Copy link
Copy Markdown
Contributor

Problem

#1773 unpinned torch_memory_saver (@d64a639 → master), which crossed the multi-cuda-wheel change. Its setup.py now aborts a CUDA source build unless TMS_CUDA_MAJOR is set:

RuntimeError: TMS_CUDA_MAJOR env var must be set for CUDA builds

so the Docker Build & Push on main fails at the pip install git+...torch_memory_saver.git step (failing run).

Fix

Set TMS_CUDA_MAJOR for that build, derived from the image's own torch (torch.version.cuda), so the cu12 and cu13 variants each build the matching wheel with no hardcoding.

Notes

  • Verified the command form (TMS_CUDA_MAJOR=<n> pip install git+...torch_memory_saver.git) builds cleanly against current TMS master on an 8×H200 CUDA-13 box.
  • Dockerfile.rocm left unpinned as chore(docker): upgrade torch_memory_saver to 6d5bce48 #1773 set it — TMS_CUDA_MAJOR is a CUDA-only requirement; the ROCm build detects hip and does not need it.

@gemini-code-assist

Copy link
Copy Markdown
Contributor

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.

@Zhichenzzz
Zhichenzzz force-pushed the zhichen/fix-docker-tms-cuda-major branch from 4170030 to 4e1d79e Compare July 23, 2026 19:05
Unpinning torch_memory_saver (#1773) crossed the multi-cuda-wheel change, whose
setup.py now aborts a CUDA source build unless TMS_CUDA_MAJOR is set, breaking
the Docker Build. Derive it from the image's torch so cu12 and cu13 variants
both work.
@Zhichenzzz
Zhichenzzz force-pushed the zhichen/fix-docker-tms-cuda-major branch from 4e1d79e to 5658359 Compare July 23, 2026 19:06

@guapisolo guapisolo left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Good fix. LGTM.

@Zhichenzzz
Zhichenzzz merged commit 3b92142 into main Jul 23, 2026
38 checks passed
@Zhichenzzz
Zhichenzzz deleted the zhichen/fix-docker-tms-cuda-major branch July 23, 2026 19:29
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants