mamba-ssm setup.py fails with InvalidVersion for vendor PyTorch/CUDA version strings instead of falling back to source build

[janbernloehr]

Summary

Installing mamba-ssm can fail during wheel build when setup.py:get_wheel_url() parses vendor-specific PyTorch/CUDA version metadata with packaging.version.parse. If the environment contains a non-PEP-440 vendor build string, get_wheel_url() raises packaging.version.InvalidVersion and aborts the install instead of falling back to a source build.

This is reproducible in vendor PyTorch containers where CUDA/PyTorch build metadata may be exported as values like gpgpu.<build-id>.

Observed failure

Package:

• mamba-ssm==2.3.2.post1
• Python 3.12
• Linux x86_64
• Vendor/internal PyTorch prerelease container

During installation as a transitive dependency:

Building mamba-ssm==2.3.2.post1
× Failed to build `mamba-ssm==2.3.2.post1`
...
File "<string>", line 340, in run
File "<string>", line 315, in get_wheel_url
File ".../packaging/version.py", line 103, in parse
  return Version(version)
File ".../packaging/version.py", line 361, in __init__
  raise InvalidVersion(f"Invalid version: {version!r}")
packaging.version.InvalidVersion: Invalid version: 'gpgpu.<build-id>'

[rehan243]

Oh yeah, we hit something super similar with a custom PyTorch build in a Docker container. The version string was like torch-1.12.0+cuda11.3.local-dev.12345, and packaging.version.parse just absolutely choked on it. Exact same InvalidVersion error you're seeing with gpgpu.<build-id>.

The thing is, packaging.version is pretty strict about PEP-440 compliance, so any extra metadata or non-standard formatting blows it up. What we ended up doing was wrapping the parse call in a try/except and logging a warning if it failed, then falling back to a source build path. Hacky, but it unblocked us.

Honestly, for something like get_wheel_url() in setup.py, feels like it should already be accounting for this — vendor strings aren’t that rare. Maybe worth adding a regex pre-check to strip out anything non-PEP-440 before calling parse?