[codex] Add Empire AI NVL72 training launcher

.github/workflows/container-physicsnemo.yml

@@ -5,7 +5,7 @@
 # Build the PhysicsNeMo-based container on PRs; on main/manual runs, also publish
 # x86_64 and arm64 images, run x86_64 CPU/GPU container tests, and run arm64
 # CPU-only container tests on a non-GPU ARM runner.
-name: Container PhysicsNeMo 25.11
+name: Container PhysicsNeMo 26.05
 
 on:
   push:
@@ -49,7 +49,7 @@ jobs:
     needs: ec2-build
     runs-on: ${{ needs.ec2-build.outputs.id }}
     env:
-      LOCAL_IMAGE_TAG: ocean-emulator:physicsnemo-25.11-${{ github.sha }}
+      LOCAL_IMAGE_TAG: ocean-emulator:physicsnemo-26.05-${{ github.sha }}
     steps:
       - name: Check out code
         uses: actions/checkout@v6
@@ -76,9 +76,9 @@ jobs:
         run: |
           BUILDKIT_PROGRESS=plain \
           IMAGE_TAG="${LOCAL_IMAGE_TAG}" \
-          DOCKERFILE=containers/Dockerfile.physicsnemo-25.11 \
+          DOCKERFILE=containers/Dockerfile.physicsnemo-26.05 \
           BUILD_APPTAINER=0 \
-          bash scripts/container/build_physicsnemo_25_11.sh
+          bash scripts/container/build_physicsnemo_26_05.sh
 
       - name: Log in to GHCR
         if: github.event_name == 'workflow_dispatch' || (github.event_name == 'push' && github.ref == 'refs/heads/main')
@@ -96,25 +96,25 @@ jobs:
           owner_lower="$(echo "${GITHUB_REPOSITORY_OWNER}" | tr '[:upper:]' '[:lower:]')"
           ghcr_image="ghcr.io/${owner_lower}/ocean-emulator-physicsnemo"
 
-          sha_tag="${ghcr_image}:25.11-${GITHUB_SHA}"
-          arch_sha_tag="${ghcr_image}:25.11-x86_64-${GITHUB_SHA}"
+          sha_tag="${ghcr_image}:26.05-${GITHUB_SHA}"
+          arch_sha_tag="${ghcr_image}:26.05-x86_64-${GITHUB_SHA}"
           docker tag "${LOCAL_IMAGE_TAG}" "${sha_tag}"
           docker tag "${LOCAL_IMAGE_TAG}" "${arch_sha_tag}"
           docker push "${sha_tag}"
           docker push "${arch_sha_tag}"
 
           if [ "${GITHUB_EVENT_NAME}" = "workflow_dispatch" ]; then
-            ref_tag="${ghcr_image}:25.11-manual-${GITHUB_REF_NAME//\//-}"
-            arch_ref_tag="${ghcr_image}:25.11-x86_64-manual-${GITHUB_REF_NAME//\//-}"
+            ref_tag="${ghcr_image}:26.05-manual-${GITHUB_REF_NAME//\//-}"
+            arch_ref_tag="${ghcr_image}:26.05-x86_64-manual-${GITHUB_REF_NAME//\//-}"
             docker tag "${LOCAL_IMAGE_TAG}" "${ref_tag}"
             docker tag "${LOCAL_IMAGE_TAG}" "${arch_ref_tag}"
             docker push "${ref_tag}"
             docker push "${arch_ref_tag}"
           fi
 
           if [ "${GITHUB_EVENT_NAME}" = "push" ] && [ "${GITHUB_REF}" = "refs/heads/main" ]; then
-            latest_tag="${ghcr_image}:25.11-latest"
-            arch_latest_tag="${ghcr_image}:25.11-x86_64-latest"
+            latest_tag="${ghcr_image}:26.05-latest"
+            arch_latest_tag="${ghcr_image}:26.05-x86_64-latest"
             docker tag "${LOCAL_IMAGE_TAG}" "${latest_tag}"
             docker tag "${LOCAL_IMAGE_TAG}" "${arch_latest_tag}"
             docker push "${latest_tag}"
@@ -184,7 +184,7 @@ jobs:
         run: |
           set -euo pipefail
           owner_lower="$(echo "${GITHUB_REPOSITORY_OWNER}" | tr '[:upper:]' '[:lower:]')"
-          image="ghcr.io/${owner_lower}/ocean-emulator-physicsnemo:25.11-${GITHUB_SHA}"
+          image="ghcr.io/${owner_lower}/ocean-emulator-physicsnemo:26.05-${GITHUB_SHA}"
           docker pull "${image}"
           echo "TEST_IMAGE=${image}" >> "${GITHUB_ENV}"
 
@@ -221,7 +221,7 @@ jobs:
     needs: ec2-arm
     runs-on: ${{ needs.ec2-arm.outputs.id }}
     env:
-      LOCAL_ARM64_IMAGE_TAG: ocean-emulator:physicsnemo-25.11-arm64-${{ github.sha }}
+      LOCAL_ARM64_IMAGE_TAG: ocean-emulator:physicsnemo-26.05-arm64-${{ github.sha }}
     steps:
       - name: Check out code
         uses: actions/checkout@v6
@@ -261,9 +261,9 @@ jobs:
         run: |
           BUILDKIT_PROGRESS=plain \
           IMAGE_TAG="${LOCAL_ARM64_IMAGE_TAG}" \
-          DOCKERFILE=containers/Dockerfile.physicsnemo-25.11 \
+          DOCKERFILE=containers/Dockerfile.physicsnemo-26.05 \
           BUILD_APPTAINER=0 \
-          bash scripts/container/build_physicsnemo_25_11.sh
+          bash scripts/container/build_physicsnemo_26_05.sh
 
       - name: Log in to GHCR
         uses: docker/login-action@v3
@@ -279,19 +279,19 @@ jobs:
           owner_lower="$(echo "${GITHUB_REPOSITORY_OWNER}" | tr '[:upper:]' '[:lower:]')"
           ghcr_image="ghcr.io/${owner_lower}/ocean-emulator-physicsnemo"
 
-          sha_tag="${ghcr_image}:25.11-arm64-${GITHUB_SHA}"
+          sha_tag="${ghcr_image}:26.05-arm64-${GITHUB_SHA}"
           docker tag "${LOCAL_ARM64_IMAGE_TAG}" "${sha_tag}"
           docker push "${sha_tag}"
           echo "TEST_ARM64_IMAGE=${sha_tag}" >> "${GITHUB_ENV}"
 
           if [ "${GITHUB_EVENT_NAME}" = "workflow_dispatch" ]; then
-            ref_tag="${ghcr_image}:25.11-arm64-manual-${GITHUB_REF_NAME//\//-}"
+            ref_tag="${ghcr_image}:26.05-arm64-manual-${GITHUB_REF_NAME//\//-}"
             docker tag "${LOCAL_ARM64_IMAGE_TAG}" "${ref_tag}"
             docker push "${ref_tag}"
           fi
 
           if [ "${GITHUB_EVENT_NAME}" = "push" ] && [ "${GITHUB_REF}" = "refs/heads/main" ]; then
-            latest_tag="${ghcr_image}:25.11-arm64-latest"
+            latest_tag="${ghcr_image}:26.05-arm64-latest"
             docker tag "${LOCAL_ARM64_IMAGE_TAG}" "${latest_tag}"
             docker push "${latest_tag}"
           fi

containers/Dockerfile.physicsnemo-25.11 → containers/Dockerfile.physicsnemo-26.05

@@ -3,7 +3,7 @@
 # SPDX-License-Identifier: Apache-2.0
 
 # syntax=docker/dockerfile:1.7
-FROM nvcr.io/nvidia/physicsnemo/physicsnemo:25.11
+FROM nvcr.io/nvidia/physicsnemo/physicsnemo:26.05
 
 WORKDIR /workspace
 
@@ -77,6 +77,21 @@ RUN . .venv/bin/activate \
         --no-install-package nvidia-nvjitlink-cu12 \
         --no-install-package nvidia-nvshmem-cu12 \
         --no-install-package nvidia-nvtx-cu12 \
+        --no-install-package nvidia-cublas \
+        --no-install-package nvidia-cuda-cupti \
+        --no-install-package nvidia-cuda-nvrtc \
+        --no-install-package nvidia-cuda-runtime \
+        --no-install-package nvidia-cudnn-cu13 \
+        --no-install-package nvidia-cufft \
+        --no-install-package nvidia-cufile \
+        --no-install-package nvidia-curand \
+        --no-install-package nvidia-cusolver \
+        --no-install-package nvidia-cusparse \
+        --no-install-package nvidia-cusparselt-cu13 \
+        --no-install-package nvidia-nccl-cu13 \
+        --no-install-package nvidia-nvjitlink \
+        --no-install-package nvidia-nvshmem-cu13 \
+        --no-install-package nvidia-nvtx \
         --no-install-package triton
 
 RUN . .venv/bin/activate \

containers/README.md

@@ -4,13 +4,13 @@ SPDX-FileCopyrightText: 2026 Samudra Authors
 SPDX-License-Identifier: CC-BY-4.0
 -->
 
-# PhysicsNeMo 25.11 Container Flow
+# PhysicsNeMo 26.05 Container Flow
 
-Build a project image on top of NVIDIA PhysicsNeMo `25.11`, verify compatibility of
+Build a project image on top of NVIDIA PhysicsNeMo `26.05`, verify compatibility of
 image-provided heavy packages, and run a smoke test:
 
 ```bash
-scripts/container/build_physicsnemo_25_11.sh
+scripts/container/build_physicsnemo_26_05.sh
 ```
 
 Run CUDA tests in the built image:
@@ -41,7 +41,7 @@ GitHub Actions workflow:
 ```
 
 That scrub runs on the local ARM GPU host, pulls
-`ghcr.io/<owner>/ocean-emulator-physicsnemo:25.11-arm64-latest` from `main`, and
+`ghcr.io/<owner>/ocean-emulator-physicsnemo:26.05-arm64-latest` from `main`, and
 runs the CUDA-marked tests inside the published ARM64 container with:
 
 ```bash
@@ -54,24 +54,24 @@ non-interactive `sudo docker`.
 Published image tags:
 
 ```text
-25.11-<sha>                  # existing x86_64 compatibility tag
-25.11-x86_64-<sha>
-25.11-arm64-<sha>
-25.11-latest                 # existing x86_64 compatibility tag from main
-25.11-x86_64-latest
-25.11-arm64-latest
-25.11-manual-<ref>           # existing x86_64 compatibility tag from workflow_dispatch
-25.11-x86_64-manual-<ref>
-25.11-arm64-manual-<ref>
+26.05-<sha>                  # existing x86_64 compatibility tag
+26.05-x86_64-<sha>
+26.05-arm64-<sha>
+26.05-latest                 # existing x86_64 compatibility tag from main
+26.05-x86_64-latest
+26.05-arm64-latest
+26.05-manual-<ref>           # existing x86_64 compatibility tag from workflow_dispatch
+26.05-x86_64-manual-<ref>
+26.05-arm64-manual-<ref>
 ```
 
 Useful environment variables:
 
 ```bash
-IMAGE_TAG=ocean-emulator:physicsnemo-25.11
-DOCKERFILE=containers/Dockerfile.physicsnemo-25.11
+IMAGE_TAG=ocean-emulator:physicsnemo-26.05
+DOCKERFILE=containers/Dockerfile.physicsnemo-26.05
 BUILD_APPTAINER=1
-SIF_PATH=dist/ocean-emulator_physicsnemo-25.11.sif
+SIF_PATH=dist/ocean-emulator_physicsnemo-26.05.sif
 PYTEST_MARK_EXPR="cuda and not manual"
 PYTEST_ARGS="-k test_trainer"
 ```

pyproject.toml

@@ -42,8 +42,9 @@ dependencies = [
   "scipy>=1.15.2",
   "skypilot[aws,lambda]>=0.10",
   "sqlalchemy>=2",                # skypilot requires sqlalchemy 2.0 but doesn't include that requirement
-  "torch>=2.9.0a0,<2.10",         # Matches PhysicsNeMo 25.11 for Docker build
+  "torch>=2.9.0a0,<2.13",         # Matches PhysicsNeMo container builds
   "torchinfo>=1.8",
+  "typing-extensions>=4.15",
   "wandb>=0.19.8",
   "xarray[io]>=2025.1.2",
   "xarray-einstats[einops]>=0.8",

scripts/container/build_physicsnemo_25_11.sh → scripts/container/build_physicsnemo_26_05.sh

@@ -6,11 +6,11 @@
 
 set -euo pipefail
 
-IMAGE_TAG="${IMAGE_TAG:-ocean-emulator:physicsnemo-25.11}"
-DOCKERFILE="${DOCKERFILE:-containers/Dockerfile.physicsnemo-25.11}"
+IMAGE_TAG="${IMAGE_TAG:-ocean-emulator:physicsnemo-26.05}"
+DOCKERFILE="${DOCKERFILE:-containers/Dockerfile.physicsnemo-26.05}"
 BUILD_APPTAINER="${BUILD_APPTAINER:-1}"
 RUN_SMOKE_TEST="${RUN_SMOKE_TEST:-1}"
-SIF_PATH="${SIF_PATH:-dist/ocean-emulator_physicsnemo-25.11.sif}"
+SIF_PATH="${SIF_PATH:-dist/ocean-emulator_physicsnemo-26.05.sif}"
 
 docker_cmd=(docker)
 if ! docker version >/dev/null 2>&1; then

scripts/container/run_cuda_tests_in_image.sh

@@ -6,7 +6,7 @@
 
 set -euo pipefail
 
-IMAGE_TAG="${IMAGE_TAG:-ocean-emulator:physicsnemo-25.11}"
+IMAGE_TAG="${IMAGE_TAG:-ocean-emulator:physicsnemo-26.05}"
 PYTEST_MARK_EXPR="${PYTEST_MARK_EXPR:-cuda and not manual}"
 PYTEST_ARGS="${PYTEST_ARGS:-}"
 

scripts/container/scrub_arm64_gpu_container_tests.sh

@@ -39,7 +39,7 @@ fi
 
 if [[ -z "${IMAGE_TAG}" ]]; then
   docker_repo_lower="$(echo "${DOCKER_REPO}" | tr '[:upper:]' '[:lower:]')"
-  IMAGE_TAG="${docker_repo_lower}/ocean-emulator-physicsnemo:25.11-arm64-latest"
+  IMAGE_TAG="${docker_repo_lower}/ocean-emulator-physicsnemo:26.05-arm64-latest"
 fi
 
 login_token="${GHCR_TOKEN:-${GITHUB_TOKEN:-${GH_TOKEN:-}}}"