Merge main

Author: ryan-williams

.github/workflows/benchmarks.yml

@@ -25,9 +25,6 @@ jobs:
     needs: start-aws-runner
     runs-on: ${{ fromJSON(needs.start-aws-runner.outputs.instances) }}
 
-    env:
-      # see note about machine image above
-      CUDA_VERSION: 12.1
     steps:
       - name: Check out code
         uses: actions/checkout@v2

.github/workflows/test-gpu.yml

@@ -16,8 +16,8 @@ jobs:
   ec2:
     uses: Open-Athena/ec2-gha/.github/workflows/runner.yml@v2
     with:
-      ec2_instance_type: g4dn.xlarge
-      ec2_image_id: ami-00096836009b16a22  # Deep Learning OSS Nvidia Driver AMI GPU PyTorch
+      ec2_instance_type: g6.xlarge
+      ec2_image_id: ami-0365bff494b18bf93 # Deep Learning OSS Nvidia Driver AMI GPU PyTorch 2.7 (Ubuntu 22.04) in oa-ci-dev
     secrets:
       GH_SA_TOKEN: ${{ secrets.GH_SA_TOKEN }}
   run-tests:
@@ -36,7 +36,9 @@ jobs:
           key: ${{ runner.os }}-test-data-${{ hashFiles('tests/conftest.py', 'configs/*.test.yaml') }}
 
       - name: Install uv and set the python version
-        uses: astral-sh/setup-uv@v5
+        uses: astral-sh/setup-uv@v6
+        with:
+          enable-cache: true
 
       - name: Set up Python
         uses: actions/setup-python@v2

.github/workflows/test.yml

@@ -29,7 +29,9 @@ jobs:
           key: ${{ runner.os }}-test-data-${{ hashFiles('tests/conftest.py', 'configs/*.test.yaml') }}
 
       - name: Install uv and set the python version
-        uses: astral-sh/setup-uv@v5
+        uses: astral-sh/setup-uv@v6
+        with:
+          enable-cache: true
 
       - name: Set up Python
         uses: actions/setup-python@v2

AGENTS.md

@@ -121,5 +121,6 @@ scripts/                  # Data download and preprocessing
 2. **Distributed Training**: Supports multi-GPU via PyTorch DDP
 3. **Performance Mindset**: We include profiling tools (memray, py-spy, scalene) and aim to keep the code performant core train and eval loops.
 4. **Testing Philosophy**: Tests marked as `manual` or `cuda` for selective execution
-5. **Pre-commit Hooks**: Enforces code quality (ruff, mypy, detect-secrets)
-6. **Cloud Training**: Supports SkyPilot for remote job execution
+5. **Cloud Training**: Supports SkyPilot for remote job execution
+6. **Noisy Failure**: Do not swallow errors. If something goes wrong, let it fail loudly.
+7. **Avoid Hacks**: Don't accomodate bad designs by adding more cruft -- refactor separately first then make the nice change.

configs/eval_cm4.yaml

@@ -17,31 +17,31 @@ experiment:
     project: ocean-emulator
     entity: ai2cm
     group: samudra-inference-2in2out
-  network: Samudra
   prognostic_vars_key: thermo_dynamic_all
   boundary_vars_key: tau_hfds
 
 data: !include data/cm4.yaml
 
-samudra:
-  ch_width: [157, 200, 250, 300, 400]
-  n_out: 77
-  dilation: [1, 2, 4, 8]
-  n_layers: [1, 1, 1, 1]
+model:
   pred_residuals: false
   last_kernel_size: 3
   pad: "circular"
 
-  core_block:
-    block_type: "conv_next_block"
-    kernel_size: 3
-    activation: "capped_gelu"
-    upscale_factor: 4
-    norm: "batch"
+  unet:
+    ch_width: [200, 250, 300, 400]
+    dilation: [1, 2, 4, 8]
+    n_layers: [1, 1, 1, 1]
+
+    core_block:
+      block_type: "conv_next_block"
+      kernel_size: 3
+      activation: "capped_gelu"
+      upscale_factor: 4
+      norm: "batch"
+
+    down_sampling_block: "avg_pool"
+    up_sampling_block: "bilinear_upsample"
 
   corrector:
     non_negative_corrector_names:
       - "so"
-
-  down_sampling_block: "avg_pool"
-  up_sampling_block: "bilinear_upsample"

configs/eval_om4.yaml

@@ -6,7 +6,7 @@ disk_mode: true
 inference_time:
   start: "2014-10-10"
   end: "2022-12-24"
-num_model_steps_forward: -1
+num_model_steps_forward: 25
 
 experiment:
   name: local_om4_samudra
@@ -17,27 +17,27 @@ experiment:
     project: ocean-emulators
     entity: m2lines
     group: samudra-eval-om4
-  network: Samudra
   prognostic_vars_key: thermo_dynamic_all
   boundary_vars_key: tau_hfds_hfds_anom
 
 data: !include data/public.yaml
 
-samudra:
-  ch_width: [157, 200, 250, 300, 400]
-  n_out: 77
-  dilation: [1, 2, 4, 8]
-  n_layers: [1, 1, 1, 1]
+model:
   pred_residuals: false
   last_kernel_size: 3
   pad: "circular"
 
-  core_block:
-    block_type: "conv_next_block"
-    kernel_size: 3
-    activation: "capped_gelu"
-    upscale_factor: 4
-    norm: "batch"
+  unet:
+    ch_width: [200, 250, 300, 400]
+    dilation: [1, 2, 4, 8]
+    n_layers: [1, 1, 1, 1]
 
-  down_sampling_block: "avg_pool"
-  up_sampling_block: "bilinear_upsample"
+    core_block:
+      block_type: "conv_next_block"
+      kernel_size: 3
+      activation: "capped_gelu"
+      upscale_factor: 4
+      norm: "batch"
+
+    down_sampling_block: "avg_pool"
+    up_sampling_block: "bilinear_upsample"

configs/eval_om4_halfdeg.yaml

@@ -6,7 +6,7 @@ disk_mode: true
 inference_time:
   start: "2014-10-10"
   end: "2022-12-24"
-num_model_steps_forward: -1
+num_model_steps_forward: 25
 
 experiment:
   name: local_om4_samudra
@@ -17,32 +17,31 @@ experiment:
     project: ocean-emulators
     entity: m2lines
     group: samudra-eval-om4
-  network: Samudra
   prognostic_vars_key: thermo_dynamic_all
   boundary_vars_key: tau_hfds_hfds_anom
 
 data: !include data/public.yaml
 
-samudra:
+model:
   checkpointing: all
-  ch_width: [157, 280, 380, 480, 520]
-  n_out: 154
-  dilation: [1, 2, 4, 8]
-  n_layers: [1, 1, 1, 1]
   pred_residuals: false
   last_kernel_size: 3
   pad: "circular"
 
-  core_block:
-    block_type: "conv_next_block"
-    kernel_size: 3
-    activation: "capped_gelu"
-    upscale_factor: 2
-    norm: "batch"
+  unet:
+    ch_width: [280, 380, 480, 520]
+    dilation: [1, 2, 4, 8]
+    n_layers: [1, 1, 1, 1]
+    core_block:
+      block_type: "conv_next_block"
+      kernel_size: 3
+      activation: "capped_gelu"
+      upscale_factor: 2
+      norm: "batch"
 
-  down_sampling_block: "avg_pool"
-  up_sampling_block: "bilinear_upsample"
+    down_sampling_block: "avg_pool"
+    up_sampling_block: "bilinear_upsample"
 
   corrector:
     non_negative_corrector_names: null
-    ocean_heat_corrector: false
+    ocean_heat_corrector: false

configs/slurm_perlmutter_train_cm4.yaml

@@ -51,31 +51,31 @@ experiment:
     project: ocean-emulator
     entity: ai2cm
     group: samudra-train-2in2out
-  network: Samudra
   prognostic_vars_key: thermo_dynamic_all
   boundary_vars_key: tau_hfds
 
 data: !include data/cm4.yaml
 
-samudra:
-  ch_width: [157, 200, 250, 300, 400]
-  n_out: 77
-  dilation: [1, 2, 4, 8]
-  n_layers: [1, 1, 1, 1]
+model:
   pred_residuals: false
   last_kernel_size: 3
   pad: "circular"
 
-  core_block:
-    block_type: "conv_next_block"
-    kernel_size: 3
-    activation: "capped_gelu"
-    upscale_factor: 4
-    norm: "batch"
+  unet:
+    ch_width: [200, 250, 300, 400]
+    dilation: [1, 2, 4, 8]
+    n_layers: [1, 1, 1, 1]
+
+    core_block:
+      block_type: "conv_next_block"
+      kernel_size: 3
+      activation: "capped_gelu"
+      upscale_factor: 4
+      norm: "batch"
+
+    down_sampling_block: "avg_pool"
+    up_sampling_block: "bilinear_upsample"
 
   corrector:
     non_negative_corrector_names:
-      - "so"
-
-  down_sampling_block: "avg_pool"
-  up_sampling_block: "bilinear_upsample"
+      - "so"

configs/slurm_perlmutter_train_om4.yaml

@@ -43,30 +43,30 @@ experiment:
     project: ocean-emulators
     entity: m2lines
     group: samudra-train-om4
-  network: Samudra
   prognostic_vars_key: thermo_all
   boundary_vars_key: tau_hfds_hfds_anom
 
 data: !include data/om4_ohc.yaml
 
-samudra:
-  ch_width: [157, 200, 250, 300, 400]
-  n_out: 77
-  dilation: [1, 2, 4, 8]
-  n_layers: [1, 1, 1, 1]
+model:
   pred_residuals: false
   last_kernel_size: 3
   pad: "circular"
 
-  core_block:
-    block_type: "conv_next_block"
-    kernel_size: 3
-    activation: "capped_gelu"
-    upscale_factor: 4
-    norm: "batch"
+  unet:
+    ch_width: [200, 250, 300, 400]
+    dilation: [1, 2, 4, 8]
+    n_layers: [1, 1, 1, 1]
 
-  down_sampling_block: "avg_pool"
-  up_sampling_block: "bilinear_upsample"
+    core_block:
+      block_type: "conv_next_block"
+      kernel_size: 3
+      activation: "capped_gelu"
+      upscale_factor: 4
+      norm: "batch"
+
+    down_sampling_block: "avg_pool"
+    up_sampling_block: "bilinear_upsample"
 
   corrector:
     non_negative_corrector_names: null

configs/train_cm4.yaml

@@ -50,33 +50,33 @@ experiment:
     project: ocean-emulators
     entity: m2lines
     group: samudra-train-cm4
-  network: Samudra
   prognostic_vars_key: thermo_dynamic_all
   boundary_vars_key: tau_hfds
 
 data: !include data/cm4.yaml
 
 
-samudra:
-  ch_width: [157, 200, 250, 300, 400]
-  n_out: 77
-  dilation: [1, 2, 4, 8]
-  n_layers: [1, 1, 1, 1]
+model:
   pred_residuals: false
   last_kernel_size: 3
   pad: "circular"
 
-  core_block:
-    block_type: "conv_next_block"
-    kernel_size: 3
-    activation: "capped_gelu"
-    upscale_factor: 4
-    norm: "batch"
+  unet:
+    ch_width: [200, 250, 300, 400]
+    dilation: [1, 2, 4, 8]
+    n_layers: [1, 1, 1, 1]
+
+    core_block:
+      block_type: "conv_next_block"
+      kernel_size: 3
+      activation: "capped_gelu"
+      upscale_factor: 4
+      norm: "batch"
+
+    down_sampling_block: "avg_pool"
+    up_sampling_block: "bilinear_upsample"
 
   corrector:
     non_negative_corrector_names:
       - "so"
-    ocean_heat_corrector: true
-
-  down_sampling_block: "avg_pool"
-  up_sampling_block: "bilinear_upsample"
+    ocean_heat_corrector: true