Skip to content

Commit 39cddd4

Browse files
njzjznjzjz
authored
feat(dev): setup devcontainer for developers (#4263)
See `.devcontainer/READMD.md` for details. <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit ## Release Notes - **New Features** - Introduced a development container setup with a new Dockerfile. - Added scripts for building C++ and Python components, downloading the LibTorch library, and setting up the environment for LAMMPS simulations. - New README documentation for the development environment setup. - New configuration file for the development container to streamline setup processes. - **Bug Fixes** - Expanded `.gitignore` to prevent unnecessary files from being tracked. - **Chores** - Enhanced dependency management in `pyproject.toml` for improved organization and clarity. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Signed-off-by: GitHub <noreply@github.com> Signed-off-by: Jinzhe Zeng <jinzhe.zeng@rutgers.edu>
1 parent aba932c commit 39cddd4

File tree

10 files changed

+135
-0
lines changed

10 files changed

+135
-0
lines changed

‎.devcontainer/Dockerfile‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,3 @@
1+
FROM mcr.microsoft.com/devcontainers/cpp:1-ubuntu-24.04
2+
3+
COPY --from=ghcr.io/astral-sh/uv:latest /uv /uvx /bin/

‎.devcontainer/READMD.md‎

Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,35 @@
1+
# DeePMD-kit devcontainer environment
2+
3+
This [devcontainer](https://vscode.js.cn/docs/devcontainers/devcontainer-cli) environment setups Python and C++ environment to develop DeePMD-kit.
4+
One can setup locally or use [GitHub Codespaces](https://docs.github.com/en/codespaces) by clicking the Code button on the DeePMD-kit repository page.
5+
The whole setup process requires about 10 minutes, so one needs to be patient.
6+
7+
## Python environment
8+
9+
The following packages are installed into the Python environment `.venv`:
10+
11+
- DeePMD-kit (in edit mode)
12+
- Backends including TensorFlow, PyTorch, JAX
13+
- LAMMPS
14+
- MPICH
15+
- CMake
16+
- pre-commit (including hooks)
17+
- Test packages including pytest
18+
- Doc packages including sphinx
19+
20+
## C++ interface
21+
22+
The C++ interface with TensorFlow and PyTorch support is installed into `dp` directory.
23+
24+
When calling and debuging LAMMPS with DeePMD-kit, use the following scripts instead of the regular `lmp`:
25+
26+
- `.devcontainer/lmp`
27+
- `.devcontainer/gdb_lmp`
28+
29+
## Rebuild
30+
31+
Usually the Python package does not need to reinstall.
32+
But when one wants to recompile the C++ code, the following scripts can be executed.
33+
34+
- `.devcontainer/build_cxx.sh`
35+
- `.devcontainer/build_py.sh`

‎.devcontainer/build_cxx.sh‎

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,21 @@
1+
#!/bin/bash
2+
set -ev
3+
4+
NPROC=$(nproc --all)
5+
SCRIPT_PATH=$(dirname $(realpath -s $0))
6+
7+
export CMAKE_PREFIX_PATH=${SCRIPT_PATH}/../libtorch
8+
TENSORFLOW_ROOT=$(python -c 'import importlib,pathlib;print(pathlib.Path(importlib.util.find_spec("tensorflow").origin).parent)')
9+
10+
mkdir -p ${SCRIPT_PATH}/../buildcxx/
11+
cd ${SCRIPT_PATH}/../buildcxx/
12+
cmake -D ENABLE_TENSORFLOW=ON \
13+
-D ENABLE_PYTORCH=ON \
14+
-D CMAKE_INSTALL_PREFIX=${SCRIPT_PATH}/../dp/ \
15+
-D LAMMPS_VERSION=stable_29Aug2024_update1 \
16+
-D CMAKE_BUILD_TYPE=Debug \
17+
-D BUILD_TESTING:BOOL=TRUE \
18+
-D TENSORFLOW_ROOT=${TENSORFLOW_ROOT} \
19+
${SCRIPT_PATH}/../source
20+
cmake --build . -j${NPROC}
21+
cmake --install .

‎.devcontainer/build_py.sh‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,8 @@
1+
#!/bin/bash
2+
set -ev
3+
4+
SCRIPT_PATH=$(dirname $(realpath -s $0))
5+
cd ${SCRIPT_PATH}/..
6+
7+
uv sync --dev --python 3.12 --extra cpu --extra torch --extra jax --extra lmp --extra test --extra docs
8+
pre-commit install

‎.devcontainer/devcontainer.json‎

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,17 @@
1+
{
2+
"name": "DeePMD-kit",
3+
"build": {
4+
"dockerfile": "Dockerfile"
5+
},
6+
"features": {
7+
"ghcr.io/devcontainers/features/github-cli:1": {}
8+
},
9+
"postCreateCommand": ".devcontainer/build_py.sh && .devcontainer/download_libtorch.sh && .devcontainer/build_cxx.sh && pre-commit install-hooks",
10+
"remoteEnv": {
11+
"PATH": "${containerEnv:PATH}:${containerWorkspaceFolder}/.venv/bin",
12+
"DP_ENABLE_PYTORCH": "1",
13+
"DP_VARIANT": "cpu",
14+
"LMP_CXX11_ABI_0": "1",
15+
"UV_EXTRA_INDEX_URL": "https://download.pytorch.org/whl/cpu"
16+
}
17+
}

‎.devcontainer/download_libtorch.sh‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,8 @@
1+
#!/bin/bash
2+
set -ev
3+
4+
SCRIPT_PATH=$(dirname $(realpath -s $0))
5+
cd ${SCRIPT_PATH}/..
6+
7+
wget https://download.pytorch.org/libtorch/cpu/libtorch-cxx11-abi-shared-with-deps-2.5.0%2Bcpu.zip -O ~/libtorch.zip
8+
unzip ~/libtorch.zip

‎.devcontainer/gdb_lmp‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,9 @@
1+
#!/bin/bash
2+
SCRIPT_PATH=$(dirname $(realpath -s $0))
3+
4+
export CMAKE_PREFIX_PATH=${SCRIPT_PATH}/../libtorch
5+
TENSORFLOW_ROOT=$(python -c 'import importlib,pathlib;print(pathlib.Path(importlib.util.find_spec("tensorflow").origin).parent)')
6+
7+
env LAMMPS_PLUGIN_PATH=${SCRIPT_PATH}/../dp/lib/deepmd_lmp \
8+
LD_LIBRARY_PATH=${SCRIPT_PATH}/../dp/lib:${CMAKE_PREFIX_PATH}/lib:${TENSORFLOW_ROOT} \
9+
gdb ${SCRIPT_PATH}/../.venv/lib/python3.12/site-packages/lammps/lmp "$@"

‎.devcontainer/lmp‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,9 @@
1+
#!/bin/bash
2+
SCRIPT_PATH=$(dirname $(realpath -s $0))
3+
4+
export CMAKE_PREFIX_PATH=${SCRIPT_PATH}/../libtorch
5+
TENSORFLOW_ROOT=$(python -c 'import importlib,pathlib;print(pathlib.Path(importlib.util.find_spec("tensorflow").origin).parent)')
6+
7+
env LAMMPS_PLUGIN_PATH=${SCRIPT_PATH}/../dp/lib/deepmd_lmp \
8+
LD_LIBRARY_PATH=${SCRIPT_PATH}/../dp/lib:${CMAKE_PREFIX_PATH}/lib:${TENSORFLOW_ROOT} \
9+
${SCRIPT_PATH}/../.venv/bin/lmp "$@"

‎.gitignore‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -45,3 +45,8 @@ build_c_tests
4545
build_c/
4646
libdeepmd_c/
4747
.uv/
48+
libtorch/
49+
uv.lock
50+
buildcxx/
51+
node_modules/
52+
*.bib.original

‎pyproject.toml‎

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -136,16 +136,28 @@ cu12 = [
136136
"nvidia-cuda-nvcc-cu12",
137137
]
138138
jax = [
139+
# below is a funny workaround for
140+
# https://github.com/astral-sh/uv/issues/8601
139141
'jax>=0.4.33;python_version>="3.10"',
142+
'jax>=0.4.33;python_version>="3.10"',
143+
'flax>=0.10.0;python_version>="3.10"',
140144
'flax>=0.10.0;python_version>="3.10"',
141145
'orbax-checkpoint;python_version>="3.10"',
146+
'orbax-checkpoint;python_version>="3.10"',
142147
# The pinning of ml_dtypes may conflict with TF
143148
# 'jax-ai-stack;python_version>="3.10"',
144149
]
145150

146151
[tool.deepmd_build_backend.scripts]
147152
dp = "deepmd.main:main"
148153

154+
[dependency-groups]
155+
dev = [
156+
"pre-commit",
157+
"cmake",
158+
"mpich",
159+
]
160+
149161
[tool.setuptools_scm]
150162

151163
[tool.scikit-build]
@@ -428,3 +440,11 @@ select = [
428440
"TOR1",
429441
"TOR2",
430442
]
443+
444+
[tool.uv.sources]
445+
mpich = { index = "mpi4py" }
446+
447+
[[tool.uv.index]]
448+
name = "mpi4py"
449+
url = "https://pypi.anaconda.org/mpi4py/simple"
450+
explicit = true

0 commit comments

Comments
 (0)