Remove pytorch channel from conda env file and upper bound deepspeed

.github/workflows/ci-integration-test-conda-reusable.yml

@@ -110,11 +110,24 @@ jobs:
           restore-keys: |
             shared-params_of3
 
-      - name: Install AWS CLI and Download OpenFold parameters
-        if: steps.cache-openfold_parameters.outputs.cache-hit != 'true'
+      - name: Download OpenFold parameters
+        # openfold3 owns the download: `setup_openfold` resolves the checkpoint file
+        # name from OPENFOLD_MODEL_CHECKPOINT_REGISTRY, writes `ckpt_root`, sets up the
+        # Biotite CCD, and skips any transfer that is already on disk
+        # (force_download=False), so this is a no-op on a good cache restore. It runs
+        # inside the image, so the host needs neither the AWS CLI nor credentials --
+        # the bucket is public and boto3 comes with the environment.
+        #
+        # Deliberately NOT gated on `cache-hit`: a restored archive is only useful if it
+        # lands in the directory that gets mounted, and checking for the file is the
+        # honest test of that. `--non-interactive` takes the OpenFoldSetupConfig
+        # defaults, which fetch only the default checkpoint and leave
+        # run_integration_tests False -- the tests are the next step's job.
         run: |
-          echo "Cache miss: Downloading OpenFold parameters..."
-          aws s3 cp s3://openfold3-data/openfold3-parameters/of3-ob-2025-06-30-174k.pt ~/.openfold3/ --no-sign-request
+          docker run --rm \
+            -v ~/.openfold3:/root/.openfold3 \
+            ${{ env.REGISTRY }}/${{ env.IMAGE_NAME }}:test-${{ inputs.cuda_base_image_tag }}-${{ github.sha }} \
+            setup_openfold --non-interactive
 
       - name: Run integration test
         run: |

.github/workflows/ci-integration-test-pixi-amd-reusable.yml

@@ -68,14 +68,23 @@ jobs:
             shared-params_of3
 
       - name: Download OpenFold parameters
-        if: steps.cache-openfold_parameters.outputs.cache-hit != 'true'
+        # openfold3 owns the download: `setup_openfold` resolves the checkpoint file
+        # name from OPENFOLD_MODEL_CHECKPOINT_REGISTRY, writes `ckpt_root`, sets up the
+        # Biotite CCD, and skips any transfer that is already on disk
+        # (force_download=False), so this is a no-op on a good cache restore. It runs
+        # inside the image, so the host needs neither the AWS CLI nor credentials --
+        # the bucket is public and boto3 comes with the environment.
+        #
+        # Deliberately NOT gated on `cache-hit`: a restored archive is only useful if it
+        # lands in the directory that gets mounted, and checking for the file is the
+        # honest test of that. `--non-interactive` takes the OpenFoldSetupConfig
+        # defaults, which fetch only the default checkpoint and leave
+        # run_integration_tests False -- the tests are the next step's job.
         run: |
-          echo "Cache miss: Downloading OpenFold parameters..."
-          # Public bucket — fetch over HTTPS so the self-hosted runner doesn't
-          # need the AWS CLI installed (equivalent to `aws s3 cp --no-sign-request`).
-          curl -fSL \
-            https://openfold3-data.s3.amazonaws.com/openfold3-parameters/of3-ob-2025-06-30-174k.pt \
-            -o ~/.openfold3/of3-ob-2025-06-30-174k.pt
+          docker run --rm \
+            -v ~/.openfold3:/root/.openfold3 \
+            ${{ env.REGISTRY }}/${{ env.IMAGE_NAME }}:test-pixi-${{ inputs.pixi_env }}-${{ github.sha }} \
+            setup_openfold --non-interactive
 
       - name: Run integration test
         run: |

.github/workflows/ci-integration-test-pixi-cuda-reusable.yml

@@ -111,11 +111,24 @@ jobs:
           restore-keys: |
             shared-params_of3
 
-      - name: Install AWS CLI and Download OpenFold parameters
-        if: steps.cache-openfold_parameters.outputs.cache-hit != 'true'
+      - name: Download OpenFold parameters
+        # openfold3 owns the download: `setup_openfold` resolves the checkpoint file
+        # name from OPENFOLD_MODEL_CHECKPOINT_REGISTRY, writes `ckpt_root`, sets up the
+        # Biotite CCD, and skips any transfer that is already on disk
+        # (force_download=False), so this is a no-op on a good cache restore. It runs
+        # inside the image, so the host needs neither the AWS CLI nor credentials --
+        # the bucket is public and boto3 comes with the environment.
+        #
+        # Deliberately NOT gated on `cache-hit`: a restored archive is only useful if it
+        # lands in the directory that gets mounted, and checking for the file is the
+        # honest test of that. `--non-interactive` takes the OpenFoldSetupConfig
+        # defaults, which fetch only the default checkpoint and leave
+        # run_integration_tests False -- the tests are the next step's job.
         run: |
-          echo "Cache miss: Downloading OpenFold parameters..."
-          aws s3 cp s3://openfold3-data/openfold3-parameters/of3-ob-2025-06-30-174k.pt ~/.openfold3/ --no-sign-request
+          docker run --rm \
+            -v ~/.openfold3:/root/.openfold3 \
+            ${{ env.REGISTRY }}/${{ env.IMAGE_NAME }}:test-pixi-${{ inputs.pixi_env }}-${{ github.sha }} \
+            setup_openfold --non-interactive
 
       - name: Run integration test
         run: |

docker/Dockerfile.conda

@@ -37,7 +37,12 @@ COPY environments/production-linux-64.yml /opt/openfold3/production-linux-64.yml
 # Install environment based on BUILD_MODE
 # - lock: uses conda-lock for exact reproducible builds (training/production)
 # - yaml: uses mamba env create for flexible version resolution (development/testing)
-RUN mamba install -n base -c conda-forge conda-lock --yes \
+# CONDA_OVERRIDE_CUDA bypasses GPU detection while solving: conda-forge's CUDA
+# builds of pytorch are gated on the `__cuda` virtual package, which is absent on a
+# GPU-less build machine. Advertise a CUDA 12 driver so the CUDA build is selectable
+# (same trick as docker/Dockerfile.pixi); it changes nothing at runtime.
+RUN export CONDA_OVERRIDE_CUDA=12.0 \
+    && mamba install -n base -c conda-forge conda-lock --yes \
     && if [ "$BUILD_MODE" = "lock" ]; then \
          conda-lock install --name openfold3 /opt/openfold3/production-linux-64.lock; \
        elif [ "$BUILD_MODE" = "yaml" ]; then \

docker/Dockerfile.update-reqs

@@ -19,8 +19,13 @@ COPY --chown=$MAMBA_USER:$MAMBA_USER environments/production-linux-64.yml /tmp/p
 
 # Generate explicit lock file for linux-64
 # The explicit format is directly consumable by mamba/conda
+# --with-cuda is required for conda-forge's CUDA builds of pytorch: conda-lock never
+# looks at the host's GPU, it solves against its bundled default virtual packages,
+# which declare `__cuda: 11.4` -- too low for the CUDA 12 stack (`__cuda >=12`).
+# It does not honour CONDA_OVERRIDE_CUDA, so this has to be the flag.
 RUN micromamba run -n base conda-lock lock \
     --mamba \
+    --with-cuda 12.0 \
     --platform linux-64 \
     --file /tmp/production-linux-64.yml \
     --kind explicit \