deepmodeling / deepmodeling/deepmd-kit

[Feature Request] Improve pip installation path to handle CUDA Toolkit discovery for PyTorch/Inductor workflows

Open
#5,731 1 comment 0 reactions 0 assignees View on GitHub
Docs enhancement
Dominant language
Python
Stars
2k
Forks
649
Avg merge
6d 18h
Merged PRs (30d)
15

Description

### Summary

## Summary

I would like to request an improvement to the pip/source installation workflow for DeepMD-kit with the PyTorch backend and CUDA enabled.

Currently, after a successful installation, running a PyTorch backend workflow may fail at runtime when PyTorch Inductor attempts to compile CUDA-related code and cannot find `CUDA_HOME`:

```text
torch._inductor.exc.InductorError: OSError: CUDA_HOME environment variable is not set. Please set it to your CUDA install root.
```

In my environment, CUDA Toolkit is installed and PyTorch can use the GPU correctly, but the runtime still fails because the required CUDA Toolkit location is not propagated to the final DeepMD/PyTorch execution environment.

This is not a request to silently install the NVIDIA driver or system CUDA Toolkit. Instead, this is a feature request to make the pip/source installation path detect, record, validate, or clearly expose the required CUDA environment for PyTorch/Inductor workflows.

### Detailed Description

## Environment

```text
DeepMD-kit: source build
Python: 3.12.3
PyTorch: 2.11.0+cu128
torch.version.cuda: 12.8
CUDA Toolkit: 12.8.93
CUDA Toolkit path: /opt/cuda_128
GPU backend: PyTorch
DP_ENABLE_PYTORCH=1
DP_ENABLE_TENSORFLOW=0
DP_VARIANT=cuda
```

PyTorch itself can see CUDA:

```text
torch: 2.11.0+cu128
torch cuda: 12.8
cuda available: True
```

CUDA Toolkit is also installed:

```text
/opt/cuda_128/bin/nvcc
Cuda compilation tools, release 12.8, V12.8.93
```

However, during runtime, PyTorch Inductor fails because `CUDA_HOME` is not set in the environment used by the `dp` process.

## Error

```text
File "/opt/deepmd-kit/.venv/lib/python3.12/site-packages/torch/utils/cpp_extension.py", line 1633, in include_paths
cuda_home_include = _join_cuda_home('include')
File "/opt/deepmd-kit/.venv/lib/python3.12/site-packages/torch/utils/cpp_extension.py", line 3182, in _join_cuda_home
raise OSError('CUDA_HOME environment variable is not set. '
torch._inductor.exc.InductorError: OSError: CUDA_HOME environment variable is not set. Please set it to your CUDA install root.
```

## Why this matters

For PyTorch/DPA-4 workflows, users are increasingly likely to rely on PyTorch compilation paths such as TorchDynamo/Inductor, or other CUDA extension mechanisms. These runtime paths may require access to the CUDA Toolkit, not only the CUDA runtime libraries bundled with PyTorch wheels.

A user may have all of the following working correctly:

```text
nvidia-smi works
nvcc exists
torch.cuda.is_available() is True
DeepMD-kit installation completes
```

but still fail later because the `dp` runtime environment does not know where the CUDA Toolkit is.

This creates a fragile user experience: the installation appears successful, but the first real PyTorch backend run fails with a low-level `CUDA_HOME` error.

## Current manual workaround

Users currently need to manually set something like:

```bash
export CUDA_HOME=/opt/cuda_128
export CUDA_ROOT=/opt/cuda_128
export CUDA_PATH=/opt/cuda_128
export CUDAToolkit_ROOT=/opt/cuda_128
export CUDA_TOOLKIT_ROOT_DIR=/opt/cuda_128

export PATH=/opt/cuda_128/bin:$PATH
export LD_LIBRARY_PATH=/opt/cuda_128/lib64:${LD_LIBRARY_PATH:-}
```

This works, but it is not obvious from the pip/source installation flow. It is especially easy to lose these variables when using `nohup`, batch scripts, containers, or job schedulers.

## Requested feature

Please consider adding one or more of the following installation-time improvements.

### 1. Installation-time CUDA Toolkit detection

During `pip install .` or source build, if `DP_VARIANT=cuda` and `DP_ENABLE_PYTORCH=1`, detect common CUDA Toolkit locations, such as:

```text
$CUDA_HOME
$CUDAToolkit_ROOT
/usr/local/cuda
/usr/local/cuda-12.*
/opt/cuda*
```

If a valid Toolkit is found, record or expose the path for later use.

### 2. Generate an environment activation script

The installer could generate a script such as:

```text
/bin/deepmd-env.sh
```

or:

```text
/deepmd/env.sh
```

containing the required variables:

```bash
export CUDA_HOME=...
export CUDAToolkit_ROOT=...
export PATH=...
export LD_LIBRARY_PATH=...
```

Then the documentation can instruct users to run:

```bash
source deepmd-env.sh
```

before `dp`, LAMMPS, or batch jobs.

### 3. Provide a runtime diagnostic command

Add a command such as:

```bash
dp doctor
```

or:

```bash
dp --check-env
```

that validates:

```text
CUDA_HOME is set
nvcc is visible
CUDA include path exists
CUDA libcudart exists
torch.cuda.is_available() is True
torch.utils.cpp_extension.CUDA_HOME is valid
DeepMD C++ libraries are discoverable
LAMMPS plugin path is discoverable, if installed
```

This would turn low-level runtime failures into actionable diagnostics.

### 4. Improve runtime error message

If `DP_ENABLE_PYTORCH=1` and the PyTorch backend triggers a compilation path requiring CUDA Toolkit, DeepMD-kit could catch or pre-check this condition and report a clearer message:

```text
CUDA Toolkit was not found by PyTorch Inductor.

Please set CUDA_HOME to your CUDA Toolkit root, for example:

export CUDA_HOME=/usr/local/cuda

or source the DeepMD environment script generated during installation.
```

This would be much clearer than the current deep PyTorch stack trace.

### 5. Document pip/source CUDA environment requirements

The documentation should explicitly distinguish:

```text
1. CUDA runtime bundled with PyTorch wheels
2. System CUDA Toolkit required by nvcc / CMake / PyTorch Inductor / C++ extensions
3. NVIDIA driver
```

Many users reasonably assume that if `torch.cuda.is_available()` is true, the CUDA environment is complete. That is not sufficient for workflows involving PyTorch Inductor or CUDA extension compilation.

## Expected behavior

For a PyTorch + CUDA installation, users should either:

1. Get an installation-time warning that `CUDA_HOME` is missing, or
2. Get a generated environment script that sets the required CUDA variables, or
3. Get a clear `dp doctor` diagnostic before runtime, or
4. Get a DeepMD-specific actionable error message when CUDA Toolkit discovery fails.

The user should not need to discover this only after a long runtime stack trace inside `torch._inductor`.

## Actual behavior

The installation can complete successfully, but runtime later fails with:

```text
torch._inductor.exc.InductorError: OSError: CUDA_HOME environment variable is not set.
```

This makes the pip/source installation path feel incomplete for PyTorch + CUDA workflows.

## Request

Please consider improving the pip/source installation workflow so that CUDA Toolkit discovery for PyTorch/Inductor workflows is handled or validated during installation, rather than discovered only at runtime.

This will become increasingly important as DPA-4 and PyTorch-based workflows become more common among DeepMD-kit users.

### Further Information, Files, and Links

_No response_

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.