Switch from pdm to poetry

.devcontainer/devcontainer.json

@@ -6,7 +6,7 @@
 	"image": "mcr.microsoft.com/devcontainers/python:1-3.12-bullseye",
 	// Features to add to the dev container. More info: https://containers.dev/features.
 	"features": {
-		"ghcr.io/devcontainers-contrib/features/pdm:2": {},
+		"ghcr.io/devcontainers-contrib/features/poetry:2": {"version": "1.8.3"},
 		"ghcr.io/devcontainers-contrib/features/pre-commit:2": {},
 		"ghcr.io/devcontainers-contrib/features/apt-get-packages": {
 			"packages": "vim"
@@ -52,9 +52,10 @@
 				"ms-python.vscode-pylance",
 				"charliermarsh.ruff",
 				"GitHub.copilot",
-				"knowsuchagency.pdm-task-provider",
 				"GitHub.copilot-chat",
-				"mutantdino.resourcemonitor"
+				"mutantdino.resourcemonitor",
+				"tamasfe.even-better-toml",
+				"zeshuaro.vscode-python-poetry"
 			]
 		}
 	},
@@ -65,22 +66,26 @@
 	},
 	"mounts": [
 		// https://code.visualstudio.com/remote/advancedcontainers/add-local-file-mount
-		// Mount a directory which will contain the pdm installation cache (shared with the host machine).
-		// This will use $SCRATCH/.cache/pdm, otherwise
+		// Mount a directory which will contain the poetry installation cache (shared with the host machine).
 		// Mount a "$SCRATCH" directory in the host to ~/scratch in the container.
 		"source=${localEnv:SCRATCH},target=/home/vscode/scratch,type=bind,consistency=cached",
-		"source=${localEnv:SCRATCH}/.cache/pdm,target=/home/vscode/.pdm_install_cache,type=bind,consistency=cached",
+		// NOTE: Mounting the hosts's ~/.cache/pypoetry dir to ~/.poetry_cache in the container because
+		// it causes issues to try to map it to ~/.cache/pypoetry directly (since ~/.cache doesn't
+		// exist initially, it seems to cause some permission issues?). Later below we
+		// `mkdir -p ~/.cache` and make a symlink to from ~/.cache/pypoetry to ~/.poetry_cache.
+		"source=${localEnv:SCRATCH}/.cache/pypoetry,target=/home/vscode/.poetry_cache,type=bind,consistency=cached",
 		// Mount a /network to match the /network directory on the host.
 		// FIXME: This assumes that either the NETWORK_DIR environment variable is set on the host, or
 		// that the /network directory exists.
 		"source=${localEnv:NETWORK_DIR:/network},target=/network,type=bind,readonly",
 		// Mount $SLURM_TMPDIR on the host machine to /tmp/slurm_tmpdir in the container.
 		// note: there's also a SLURM_TMPDIR env variable set to /tmp/slurm_tmpdir in the container.
-		// NOTE: this assumes that either $SLURM_TMPDIR is set on the host machine (e.g. a compute node)
+		// This assumes that either $SLURM_TMPDIR is set on the host machine (e.g. a compute node)
 		// or that `/tmp/slurm_tmpdir` exists on the host machine.
 		"source=${localEnv:SLURM_TMPDIR:/tmp/slurm_tmpdir},target=/tmp,type=bind,consistency=cached",
 		// Mount the ssh directory on the host machine to the container.
 		"source=${localEnv:HOME}/.ssh,target=/home/vscode/.ssh,type=bind,readonly"
+
 	],
 	"runArgs": [
 		"--gpus",
@@ -90,32 +95,31 @@
 	// create the pdm cache dir on the host machine if it doesn exist yet so the mount above
 	// doesn't fail.
 	"initializeCommand": {
-		"create pdm install cache": "mkdir -p ${SCRATCH?need the SCRATCH environment variable to be set.}/.cache/pdm",      // todo: put this on $SCRATCH on the host (e.g. compute node)
+		"create the poetry install cache dir if needed": "mkdir -p ${SCRATCH?need the SCRATCH environment variable to be set.}/.cache/pypoetry",      // todo: put this on $SCRATCH on the host (e.g. compute node)
+		// Create the dirs so the mount above doesn't fail.
+		// Note: here we use /tmp/slurm_tmpdir if $SLURM_TMPDIR is not set. This is fine for Mila compute nodes.
 		"create fake SLURM_TMPDIR": "mkdir -p ${SLURM_TMPDIR:-/tmp/slurm_tmpdir}",  // this is fine on compute nodes
+		// Create the ssh cache dir where ssh multiplexing sockets will be placed (to reuse connections).
 		"create ssh cache dir": "mkdir -p ~/.cache/ssh"
 	},
-	// NOTE: Getting some permission issues with the .cache dir if mounting .cache/pdm to
-	// .cache/pdm in the container. Therefore, here I'm making a symlink from ~/.cache/pdm to
-	// ~/.pdm_install_cache so the ~/.cache directory is writeable by the container.
+	// NOTE: Getting some permission issues with the .cache dir if mounting .cache/pypoetry to
+	// .cache/pypoetry in the container. Therefore, here I'm making a symlink from ~/.cache/pypoetry to
+	// ~/.poetry_cache so the ~/.cache directory is writeable by the container.
 	"onCreateCommand": {
-		"setup_pdm_install_cache": "mkdir -p ~/.cache && ln -s /home/vscode/.pdm_install_cache /home/vscode/.cache/pdm",
-		"pre-commit": "pre-commit install --install-hooks",
-		"setup_pdm_config": "pdm config install.cache true && pdm config venv.with_pip true && pdm config venv.in_project false"
+		"Setup poetry completions": "poetry completions bash >> ~/.bash_completion",
+		// todo: doesn't work, getting some permission issues: https://github.com/orgs/python-poetry/discussions/8182
+		// "setup dynamic versioning": "pipx run poetry self add poetry-dynamic-versioning[plugin] && pipx run poetry dynamic-versioning enable",
+		"Setup the poetry install cache": "mkdir -p ~/.cache && ln -s /home/vscode/.poetry_install_cache /home/vscode/.cache/pypoetry",
+		"Install pre-commit hooks": "pre-commit install --install-hooks"
 	},
 	"updateContentCommand": {
-		"pdm_install": "pdm install"
+		"pdm_install": "poetry install",
+		// "poetry dynamic versioning": "poetry dynamic-versioning"
 	},
 
 	// Use 'postCreateCommand' to run commands after the container is created.
-	// "postCreateCommand": "pip3 install --user -r requirements.txt",
 	"postCreateCommand": {
-		// NOTE: This sets up the python interpreter correctly, but the shell doesn't show the
-		// environment name, which isn't a big deal.
-		"setup_venv_shell": "pdm venv activate >> ~/.bash_aliases && chmod +x ~/.bash_aliases",
-		// "pdm-pep-582": "pdm --pep582 >> ~/.bash_profile",
-		// "pdm": "pdm config install.cache false && pdm config venv.in_project false && pdm install",
-		// "pre-commit": "pre-commit install --install-hooks"
-		// "activate_venv": "source .venv/bin/activate"
+		"Setup shell to use the venv": "echo \"source `poetry env info --path`/bin/activate\" >> ~/.bash_aliases && chmod +x ~/.bash_aliases",
 	}
 	// Uncomment to connect as root instead. More info: https://aka.ms/dev-containers-non-root.
 	// "remoteUser": "root"

actions-runner-job.sh → .github/actions-runner-job.sh

@@ -1,13 +1,18 @@
 #!/bin/bash
 #SBATCH --nodes=1
 #SBATCH --ntasks=1
-#SBATCH --cpus-per-task=1
+#SBATCH --cpus-per-task=4
 #SBATCH --mem=16G
-#SBATCH --gpus=rtx8000:1
+#SBATCH --gpus=1
 #SBATCH --time=00:30:00
 #SBATCH --dependency=singleton
 #SBATCH --output=logs/runner_%j.out
-
+#SBATCH --signal=B:TERM@60 # tells the controller to send SIGTERM to the job 1
+                            # min before its time ends to give it a chance for
+                            # better cleanup. If you cancel the job manually,
+                            # make sure that you specify the signal as TERM like
+                            # so `scancel --signal=TERM <jobid>`.
+                            # https://dhruveshp.com/blog/2021/signal-propagation-on-slurm/
 
 set -euo pipefail
 

.github/workflows/build.yml

@@ -44,21 +44,27 @@ jobs:
         python-version: ['3.12']
     steps:
     - uses: actions/checkout@v4
+    - name: Install poetry
+      # needs to be installed for the cache below to be used, which seems a bit weird (since the
+      # python version might be wrong?)
+      run: pip install poetry
+
     - name: Set up Python ${{ matrix.python-version }}
       uses: actions/setup-python@v5
       with:
         python-version: ${{ matrix.python-version }}
-    - run: pip install pdm
+        cache: poetry
+
     - name: Install dependencies
-      run: pdm install
+      run: poetry install
     - name: Test with pytest (very fast)
       env:
         JAX_PLATFORMS: cpu
-      run: pdm run pytest -v --shorter-than=1.0 --cov=project --cov-report=xml --cov-append
+      run: poetry run pytest -v --shorter-than=1.0 --cov=project --cov-report=xml --cov-append
     - name: Test with pytest (fast)
       env:
         JAX_PLATFORMS: cpu
-      run: pdm run pytest -v --cov=project --cov-report=xml --cov-append
+      run: poetry run pytest -v --cov=project --cov-report=xml --cov-append
 
     - name: Store coverage report as an artifact
       uses: actions/upload-artifact@v4
@@ -75,16 +81,19 @@ jobs:
         python-version: ['3.12']
     steps:
     - uses: actions/checkout@v4
+    - name: Install poetry
+      # See note above.
+      run: pip install poetry
     - name: Set up Python ${{ matrix.python-version }}
       uses: actions/setup-python@v5
       with:
         python-version: ${{ matrix.python-version }}
-    - run: pip install pdm
+        cache: poetry
     - name: Install dependencies
-      run: pdm config install.cache true && pdm install
+      run: poetry install
 
     - name: Test with pytest
-      run: pdm run pytest -v --cov=project --cov-report=xml --cov-append
+      run: poetry run pytest -v --cov=project --cov-report=xml --cov-append
 
     # TODO: this is taking too long to run, and is failing consistently. Need to debug this before making it part of the CI again.
     # - name: Test with pytest (only slow tests)
@@ -106,8 +115,10 @@ jobs:
     outputs:
       job_id: ${{ steps.sbatch.outputs.stdout }}
     steps:
+      - uses: actions/checkout@v4
+
       - name: Copy job script to the cluster
-        run: "scp actions-runner-job.sh ${{ matrix.cluster }}:actions-runner-job.sh"
+        run: "scp .github/actions-runner-job.sh ${{ matrix.cluster }}:"
 
       - name: Launch Slurm Actions Runner
         id: sbatch
@@ -132,16 +143,29 @@ jobs:
         cluster: ['mila']
     steps:
     - uses: actions/checkout@v4
+
+    - name: Set up Python 3.12 (to install poetry)
+      uses: actions/setup-python@v5
+      with:
+        python-version: 3.12
+
+    - name: Install poetry
+      # See note above.
+      run: pip install poetry
+
+    - name: Setup the installation cache in $SCRATCH
+      run: mkdir -p $SCRATCH/.cache && poetry config cache-dir $SCRATCH/.cache/pypoetry
+
     - name: Set up Python 3.12
       uses: actions/setup-python@v5
       with:
         python-version: 3.12
-    - run: pip install pdm
+        cache: poetry
     - name: Install dependencies
-      run: pdm install
+      run: poetry install
 
     - name: Test with pytest
-      run: pdm run pytest -v --cov=project --cov-report=xml --cov-append
+      run: poetry run pytest -v --cov=project --cov-report=xml --cov-append
 
     # TODO: Re-enable this later
     # - name: Test with pytest (only slow tests)

.github/workflows/docs.yml

@@ -11,17 +11,18 @@ jobs:
     steps:
       - name: Checkout
         uses: actions/checkout@v4
+      - name: Install poetry
+        run: pip install poetry
       - name: Set up Python 3.12
         uses: actions/setup-python@v5
         with:
           python-version: 3.12
-      - name: Install pdm
-        run: pip install pdm
+          cache: poetry
       - name: Install dependencies
-        run: pdm install
+        run: poetry install
 
       - name: Deploy docs
-        run: mkdocs gh-deploy --force
+        run: poetry run mkdocs gh-deploy --force
         # note: Checking if we really need the one below:
         # uses: mhausenblas/[email protected]
         # # Or use mhausenblas/mkdocs-deploy-gh-pages@nomaterial to build without the mkdocs-material theme

.gitignore

@@ -162,7 +162,6 @@ cython_debug/
 #.idea/
 # ADDED:
 .vscode
-.pdm-python
 wandb
 logs
 lightning_logs

.pre-commit-config.yaml

@@ -63,10 +63,10 @@ repos:
         require_serial: true
 
   # Dependency management
-  - repo: https://github.com/pdm-project/pdm
-    rev: 2.12.4
+  - repo: https://github.com/python-poetry/poetry
+    rev: 1.8.3
     hooks:
-      - id: pdm-lock-check
+      - id: poetry-check
         require_serial: true
 
 

docs/getting_started/install.md

@@ -40,11 +40,11 @@ You need to have [Conda](https://docs.conda.io/en/latest/) installed on your mac
     pip install -e .
     ```
 
-    Optionally, you can also install the package using [PDM](https://pdm-project.org/en/latest/). This makes it easier to add or change the dependencies later on:
+    Optionally, you can also install the package using [Poetry](https://python-poetry.org/). This makes it easier to add or change the package dependencies later on:
 
     ```bash
-    pip install pdm
-    pdm install
+    pip install poetry
+    poetry install
     ```
 
 ## Using a development container