Standardizing Python project setup using uv

.github/actions/python-uv-setup/action.yml

@@ -1,5 +1,6 @@
 name: "python-uv-setup"
 description: "Common Python + uv setup, deps install, and optional debug"
+
 inputs:
   python-version:
     description: "Python version"
@@ -8,6 +9,18 @@ inputs:
     description: "Enable debug output"
     required: false
     default: "false"
+  lockfile:
+    description: "Lockfile to use"
+    required: false
+    default: "uv.lock.cpu"
+  extras:
+    description: "Extra dependency set to install"
+    required: false
+    default: "cpu"
+  dependency-groups:
+    description: "Dependency group to install"
+    required: false
+    default: "dev"
 
 runs:
   using: "composite"
@@ -40,34 +53,22 @@ runs:
       uses: actions/cache@v4
       with:
         path: ~/.cache/uv
-        key: ${{ runner.os }}-uv-${{ hashFiles('pip/*.txt') }}
+        key: ${{ runner.os }}-py${{ inputs.python-version }}-uv-${{ hashFiles('pyproject.toml', 'uv.lock.cpu') }}
         restore-keys: |
-          ${{ runner.os }}-uv-
-
-    - name: Create virtual environment
-      shell: bash
-      run: |
-        uv venv --seed .venv
+          ${{ runner.os }}-py${{ inputs.python-version }}-uv-
 
     - name: Install dependencies
       shell: bash
       run: |
         set -euxo pipefail
+        cp "${{ inputs.lockfile }}" uv.lock
+        uv venv .venv
         . .venv/bin/activate
-
-        which python
-        python -V
-
-        uv pip install -e .
-        if [ -f pip/cpu_requirements.txt ]; then uv pip install -r pip/cpu_requirements.txt; fi
-        if [ -f pip/dev_requirements.txt ]; then uv pip install -r pip/dev_requirements.txt; fi
+        uv sync --extra "${{ inputs.extras }}" --group "${{ inputs.dependency-groups }}" --locked
 
     - name: Debug - what is installed
       if: ${{ inputs.debug == 'true' }}
       shell: bash
       run: |
         . .venv/bin/activate
-        python -m pip --version
-        python -V
-        uv --version
-        uv pip list
+        uv pip list

.github/workflows/pytest.yml

@@ -20,14 +20,24 @@ jobs:
       - uses: ./.github/actions/python-uv-setup
         with:
           python-version: ${{ matrix.python-version }}
+          lockfile: uv.lock.cpu
+          extras: cpu
+          dependency-groups: dev
           debug: "false"
 
-      - name: pytest
+      - name: Run modac unit tests
         shell: bash
+        working-directory: atomsci/modac/test/unit
         run: |
-          . .venv/bin/activate
-          cd atomsci/modac/test/unit && python -m pytest --capture=sys --capture=fd --cov=atomsci/ -vv
-          cd "${{ github.workspace }}/atomsci/ddm/test/unit" && python -m pytest -m "not moe_required" -n 2 --capture=fd --cov=atomsci -vv --durations=4 --durations-min=0.05
+          . "${{ github.workspace }}/.venv/bin/activate"
+          python -m pytest --capture=sys --capture=fd --cov=atomsci/ -vv
+
+      - name: Run ddm unit tests
+        shell: bash
+        working-directory: atomsci/ddm/test/unit
+        run: |
+          . "${{ github.workspace }}/.venv/bin/activate"
+          python -m pytest -m "not moe_required" -n 2 --capture=fd --cov=atomsci -vv --durations=4 --durations-min=0.05
 
       - name: Save coverage
         uses: actions/upload-artifact@v4
@@ -56,13 +66,16 @@ jobs:
       - uses: ./.github/actions/python-uv-setup
         with:
           python-version: ${{ matrix.python-version }}
+          lockfile: uv.lock.cpu
+          extras: cpu
+          dependency-groups: dev
           debug: "false"
 
       - name: Debug - dgl + torchdata + shadowing
         shell: bash
         if: ${{ env.DEBUG_DGL == 'true' }}
         run: |
-          . .venv/bin/activate
+          . "${{ github.workspace }}/.venv/bin/activate"
           set -x
 
           which python
@@ -76,11 +89,11 @@ jobs:
           python -m pip show torch torchdata dgl dgllife || true
           python -m pip freeze | grep -E '^(torch|torchdata|dgl|dgllife|numpy)=' || true
 
-      - name: pytest
+      - name: Run integrative tests
         shell: bash
+        working-directory: atomsci/ddm/test/integrative
         run: |
-          . .venv/bin/activate
-          cd atomsci/ddm/test/integrative
+          . "${{ github.workspace }}/.venv/bin/activate"
           ./integrative_batch_chunk_tests ${{ matrix.chunk }}
 
       - name: Save coverage

.gitignore

@@ -123,6 +123,8 @@ celerybeat.pid
 # Environments
 .env
 .venv
+.venv-*
+uv.lock
 env/
 venv/
 ENV/

Makefile

@@ -46,7 +46,8 @@ VENV ?= atomsci-env
 # Work Directory
 WORK_DIR ?= work
 
-.PHONY: build-docker install install-dev install-system install-venv jupyter-notebook jupyter-lab \
+.PHONY: sync-cpu sync-cuda sync-rocm sync-mps test-cpu test-gpu test-rocm test-mps \
+        build-docker install install-dev install-system install-venv jupyter-notebook jupyter-lab \
 	pytest ruff ruff-fix setup uninstall uninstall-dev uninstall-system uninstall-venv
 
 # Load Docker image
@@ -157,6 +158,31 @@ setup:
 
 uninstall: uninstall-system
 
+# for uv env
+sync-cpu:
+	UV_PROJECT_ENVIRONMENT=.venv-cpu uv sync --extra cpu
+
+sync-cuda:
+	UV_PROJECT_ENVIRONMENT=.venv-cuda uv sync --extra cuda
+
+sync-rocm:
+	UV_PROJECT_ENVIRONMENT=.venv-rocm uv sync --extra rocm
+
+sync-mps:
+	UV_PROJECT_ENVIRONMENT=.venv-mps uv sync --extra mps
+
+test-cpu:
+	UV_PROJECT_ENVIRONMENT=.venv-cpu uv run pytest
+
+test-gpu:
+	UV_PROJECT_ENVIRONMENT=.venv-cuda uv run pytest
+
+test-rocm:
+	UV_PROJECT_ENVIRONMENT=.venv-rocm uv run pytest
+
+test-mps:
+	UV_PROJECT_ENVIRONMENT=.venv-mps uv run pytest
+
 # Uninstall atomsci-ampl from user space
 uninstall-dev:
 	@echo "Uninstalling atomsci-ampl for user"

atomsci/ddm/__init__.py

@@ -4,10 +4,8 @@
 #except TypeError:
 #    pass
 try:
-    # Python 3.8+
     from importlib.metadata import version, PackageNotFoundError
 except ImportError:
-    # For very old Pythons, use backport if available
     try:
         from importlib_metadata import version, PackageNotFoundError
     except ImportError:

atomsci/ddm/pipeline/model_wrapper.py

@@ -2588,7 +2588,6 @@ def train_with_early_stopping(self, pipeline):
         """
         self.data = pipeline.data
         feature_names = self.data.featurization.get_feature_columns()
-        nfeatures = len(feature_names)
         self.feature_weights = dict(zip(feature_names, [[] for f in feature_names]))
 
         em = perf.EpochManager(self,
@@ -2616,10 +2615,8 @@ def make_pred(dset):
                           pipeline.metric_type, test_perf))
 
             layer1_weights = self.model.model.layers[0].weight
-            feature_weights = np.zeros(nfeatures, dtype=float)
-            for node_weights in layer1_weights:
-                node_feat_weights = torch.abs(node_weights).detach().numpy()
-                feature_weights += node_feat_weights
+            feature_weights = layer1_weights.detach().abs().mean(dim=0).cpu().numpy()
+
             for fnum, fname in enumerate(feature_names):
                 self.feature_weights[fname].append(feature_weights[fnum])
 

atomsci/ddm/pipeline/parameter_parser.py

@@ -7,6 +7,7 @@
 
 import deepchem.models as dcm
 import deepchem.models.torch_models as dcmt
+from deepchem.models import GraphConvModel
 import deepchem.feat as dcf
 import inspect
 
@@ -34,8 +35,8 @@
 
 model_wl = {'AttentiveFPModel':dcm.AttentiveFPModel,
             'GCNModel':dcm.GCNModel,
-            'MPNNModel':dcm.MPNNModel,
-            'GraphConvModel':dcm.GraphConvModel,
+            'MPNNModel':dcmt.MPNNModel,
+            'GraphConvModel': GraphConvModel,
             'PytorchMPNNModel':dcmt.MPNNModel}#, dcm.GCNModel, dcm.GATModel]
 
 # featurizer white list

pyproject.toml

@@ -0,0 +1,98 @@
+[project]
+name = "atomsci-ampl"
+dynamic = ["version"]
+description = "atomsci AMPL python package"
+readme = "README.md"
+requires-python = ">=3.10,<3.11"
+dependencies = [
+  "lightning==2.2.5",
+  "pydantic>=1.10,<2",
+  "numpy<2",
+  "scipy>=1.10",
+  "scikit-learn==1.2.2",
+  "hyperopt",
+  "xgboost",
+  "lightgbm",
+  "umap-learn",
+  "bokeh",
+  "matplotlib",
+  "matplotlib-venn",
+  "seaborn>=0.13.0",
+  "pyarrow",
+  "pyyaml>=6.0",
+  "bravado",
+  "imbalanced-learn",
+  "imblearn",
+  "maestrowf",
+  "deepchem==2.8",
+  "rdkit==2024.3.5",
+  "mordred",
+  "setuptools",
+]
+
+[dependency-groups]
+dev = [
+  "ipykernel",
+  "jupyter",
+  "jupyterlab",
+  "pytest",
+  "pytest-cov",
+  "pytest-mock",
+  "pytest-xdist",
+  "matplotcheck",
+  "ruff",
+]
+
+[project.optional-dependencies]
+cpu = [
+  "torch==2.1.2",
+  "torchdata==0.7.1",
+  "tensorflow-cpu~=2.14.0; platform_system == 'Linux'",
+  "tensorflow~=2.14.0; platform_system == 'Darwin'",
+  "dgl==2.1.0",
+  "dgllife==0.3.2",
+  "torch_geometric==2.7.0",
+]
+
+cuda = [
+  "torch==2.1.2",
+  "torchdata==0.7.1",
+  "tensorflow~=2.14.0",
+  "dgl==2.1.0",
+  "dgllife==0.3.2",
+  "torch-geometric==2.7.0",
+]
+
+rocm = [
+  "torch==2.1.2",
+  "torchdata==0.7.1",
+  "tensorflow~=2.14.0",
+  "dgl==2.1.0",
+  "dgllife==0.3.2",
+  "torch-geometric==2.7.0",
+]
+
+mchip = [
+  "torch==2.1.2",
+  "torchdata",
+  "tensorflow~=2.14.0",
+  "dgl==2.1.0",
+  "dgllife==0.3.2",
+  "torch_geometric==2.7.0",
+]
+
+[build-system]
+requires = ["hatchling>=1.24.0"]
+build-backend = "hatchling.build"
+
+[tool.hatch.version]
+path = "VERSION"
+pattern = "(?P<version>\\d+\\.\\d+\\.\\d+.*)"
+
+[tool.hatch.build.targets.wheel]
+packages = ["atomsci"]
+
+[tool.pytest.ini_options]
+testpaths = ["atomsci/ddm/test"]
+addopts = "-ra"
+

sync_uv_env.sh

@@ -0,0 +1,54 @@
+#!/usr/bin/env bash
+# Purpose:
+#    syncs the environment, does not save a platform lockfile
+#
+# Use when:
+#    want to create/update a local env from existing lock/deps. build an env to run.
+#
+# Usage:
+#    ./sync_uv_env.sh <cpu|cuda|rocm|mchip>
+
+set -euo pipefail
+
+platform="${1:?usage: $0 <cpu|cuda|rocm|mchip>}"
+
+lockfile="uv.lock.${platform}"
+venv_dir=".venv-${platform}"
+
+if [[ ! -f "$lockfile" ]]; then
+  echo "Missing lockfile: $lockfile"
+  echo "Run ./update_lock.sh $platform first"
+  exit 1
+fi
+
+cp "$lockfile" uv.lock
+rm -rf "$venv_dir"
+
+case "$platform" in
+  cpu)
+    uv venv --python 3.10 "$venv_dir"
+    uv pip install --python "$venv_dir/bin/python" --index-url https://download.pytorch.org/whl/cpu torch==2.1.2
+    VIRTUAL_ENV="$venv_dir" uv sync --python "$venv_dir/bin/python" --extra cpu --group dev --locked
+    ;;
+  cuda)
+    # to force the right torch wheel into the env before syncing, especially for GPU variants
+    uv venv --python 3.10 "$venv_dir"
+    uv pip install --python "$venv_dir/bin/python" --index-url https://download.pytorch.org/whl/cu121 torch==2.1.2
+    VIRTUAL_ENV="$venv_dir" uv sync --python "$venv_dir/bin/python" --extra cuda --group dev --locked
+    ;;
+  rocm)
+    uv venv --python 3.10 "$venv_dir"
+    uv pip install --python "$venv_dir/bin/python" --index-url https://download.pytorch.org/whl/rocm5.6 torch==2.1.2
+    VIRTUAL_ENV="$venv_dir" uv sync --python "$venv_dir/bin/python" --extra rocm --group dev --locked
+    ;;
+  mchip)
+    uv venv --python 3.10 "$venv_dir"
+    VIRTUAL_ENV="$venv_dir" uv sync --python "$venv_dir/bin/python" --extra mchip --group dev --locked
+    ;;
+  *)
+    echo "Invalid platform: $platform"
+    exit 1
+    ;;
+esac
+
+echo "Synced $venv_dir from $lockfile"