squash linear history

Author: brianhenn

.github/pull_request_template.md

@@ -5,5 +5,6 @@ Changes:
 - Can group multiple related symbols on a single bullet
 
 - [ ] Tests added
+- [ ] If dependencies changed, "deps only" image rebuilt and "latest_deps_only_image.txt" file updated
 
 Resolves #<github issues> (delete if none)

.github/workflows/docs.yaml

@@ -8,6 +8,21 @@ on:
     branches:
       - main
 
+  # Add a button to manually trigger the workflow
+  workflow_dispatch:
+
+# Sets permissions of the GITHUB_TOKEN to allow deployment to GitHub Pages
+permissions:
+  contents: read
+  pages: write
+  id-token: write
+
+# Allow only one concurrent deployment, skipping runs queued between the run in-progress and latest queued.
+# However, do NOT cancel in-progress runs as we want to allow these production deployments to complete.
+concurrency:
+  group: "pages"
+  cancel-in-progress: false
+
 jobs:
   build:
     runs-on: ubuntu-latest
@@ -22,19 +37,27 @@ jobs:
       - uses: actions/cache@v4
         with:
           path: ${{ env.pythonLocation }}
-          key: v1-${{ env.CURRENT_WEEK }}-${{ env.pythonLocation }}-${{ hashFiles('requirements.txt') }}-${{ hashFiles('docs/requirements.txt') }}-${{ hashFiles('constraints.txt') }}
+          key: v1-${{ env.CURRENT_WEEK }}-${{ env.pythonLocation }}-${{ hashFiles('requirements.txt') }}-${{ hashFiles('docs/requirements.txt') }}-${{ hashFiles('requirements-healpix.txt') }}-${{ hashFiles('constraints.txt') }}
       - name: Install dependencies
         run: |
           python -m pip install uv
-          uv pip install --system -c constraints.txt -e .[docs]
+          # can't use --system here because it doesn't work with --no-build-isolation and a healpix dependency
+          uv venv --python 3.11
+          uv pip install -c constraints.txt -e .[docs]
+          uv pip install --no-build-isolation -c constraints.txt -e .[docs,healpix]
       - name: Build docs
         run: |
+          source .venv/bin/activate
           cd docs && make doctest html
-      - name: Deploy to GitHub Pages
-        uses: peaceiris/actions-gh-pages@v3
-        if: ${{ github.event_name == 'push' && github.ref == 'refs/heads/main' }}
+      - name: Upload to GitHub Pages
+        uses: actions/upload-pages-artifact@v3
         with:
-          publish_branch: gh-pages
-          github_token: ${{ secrets.GITHUB_TOKEN }}
-          publish_dir: docs/_build/
-          force_orphan: true
+          path: docs/_build/html
+  deploy:
+    runs-on: ubuntu-latest
+    # Only deploy on manual trigger or push to main
+    if: ${{ github.event_name == 'workflow_dispatch' || (github.event_name == 'push' && github.ref == 'refs/heads/main') }}
+    needs: build
+    steps:
+      - name: Deploy to GitHub Pages
+        uses: actions/deploy-pages@v4

.github/workflows/tests.yaml

@@ -22,15 +22,19 @@ jobs:
       - uses: actions/cache@v4
         with:
           path: ${{ env.pythonLocation }}
-          key: v1-${{ env.CURRENT_WEEK }}-${{ env.pythonLocation }}-${{ hashFiles('requirements.txt') }}-${{ hashFiles('requirements-dev.txt') }}-${{ hashFiles('constraints.txt') }}
+          key: v1-${{ env.CURRENT_WEEK }}-${{ env.pythonLocation }}-${{ hashFiles('requirements.txt') }}-${{ hashFiles('requirements-dev.txt') }}-${{ hashFiles('requirements-healpix.txt') }}-${{ hashFiles('constraints.txt') }}
       - name: Install dependencies
         run: |
           python -m pip install uv
-          uv pip install --system -c constraints.txt -e .[dev]
+          # can't use --system here because it doesn't work with --no-build-isolation and a healpix dependency
+          uv venv --python 3.11
+          uv pip install -c constraints.txt -e .[dev]
+          uv pip install --no-build-isolation -c constraints.txt -e .[dev,healpix]
       - name: Run pytest with code coverage
         run: |
+          source .venv/bin/activate
           make test_cov
-  cpu-very-fast:
+  cpu-very-fast-no-healpix:
     runs-on: ubuntu-latest
     steps:
       - uses: actions/checkout@v3
@@ -51,3 +55,35 @@ jobs:
       - name: Run pytest
         run: |
           make test_very_fast
+  gpu:
+    # These steps run on our self-hosted k8s GPU runners. See below for setup notes
+    # https://github.com/ai2cm/long-lived-infrastructure#installing-the-github-actions-runner-controller
+    runs-on: arc-runner-set
+    steps:
+      - uses: actions/checkout@v3
+      - name: Set environment variables
+        run: |
+          echo "CURRENT_WEEK=$(date +'%Y-%U')" >> $GITHUB_ENV
+      - uses: actions/setup-python@v4
+        with:
+          python-version: "3.11"
+      - uses: actions/cache@v4
+        with:
+          path: ${{ env.pythonLocation }}
+          key: ${{ env.CURRENT_WEEK }}-${{ env.pythonLocation }}-${{ hashFiles('requirements.txt') }}-${{ hashFiles('requirements-dev.txt') }}-${{ hashFiles('requirements-healpix.txt') }}-${{ hashFiles('constraints.txt') }}
+      - name: Install dependencies
+        run: |
+          sudo apt-get update
+          sudo apt-get install -y make git build-essential
+          python -m pip install uv
+          # can't use --system here because it doesn't work with --no-build-isolation and a healpix dependency
+          uv venv --python 3.11
+          uv pip install -c constraints.txt -e .[dev]
+          uv pip install --no-build-isolation -c constraints.txt -e .[dev,healpix]
+      - name: Run CUDA check and pytest
+        run: |
+            source .venv/bin/activate
+            export PATH=/usr/local/nvidia/bin:${PATH}
+            export LD_LIBRARY_PATH=/usr/local/nvidia/lib64:${LD_LIBRARY_PATH}
+            python3 fme/require_gpu.py
+            make test

.gitignore

@@ -1,5 +1,8 @@
 docs/available_modules.rst
+docs/available_steps.rst
 docs/_build
+filtered_repo/
+files_excluded_by_filter.txt
 
 .vscode
 

.pre-commit-config.yaml

@@ -10,6 +10,10 @@ repos:
     hooks:
     - id: check-added-large-files
       args: [--maxkb=250]
+      exclude: |
+        (?x)^(
+        fme/ace/aggregator/inference/testdata/.*-regression.pt
+        )$
     - id: trailing-whitespace
     - id: file-contents-sorter
       files: |
@@ -28,4 +32,5 @@ repos:
         (?x)^(
         .+/conf.py |
         .+/conftest.py |
+        fme/ace/models/makani_fcn3/.*
         )$

Makefile

@@ -1,7 +1,10 @@
 VERSION ?= $(shell git rev-parse --short HEAD)
 IMAGE ?= fme
+REGISTRY ?= registry.nersc.gov/m4492/ai2cm
 ENVIRONMENT_NAME ?= fme
+USERNAME ?= $(shell beaker account whoami --format=json | jq -r '.[0].name')
 DEPLOY_TARGET ?= pypi
+BEAKER_WORKSPACE = ai2/ace
 
 ifeq ($(shell uname), Linux)
 	CONDA_PACKAGES=gxx_linux-64 pip
@@ -12,15 +15,36 @@ endif
 build_docker_image:
 	DOCKER_BUILDKIT=1 docker build --platform=linux/amd64 -f docker/Dockerfile -t $(IMAGE):$(VERSION) --target production .
 
+push_shifter_image: build_docker_image
+	docker tag $(IMAGE):$(VERSION) $(REGISTRY)/$(IMAGE):$(VERSION)
+	docker push $(REGISTRY)/$(IMAGE):$(VERSION)
+
+build_beaker_image: build_docker_image
+	beaker image create --name $(IMAGE)-$(VERSION) $(IMAGE):$(VERSION)
+
+build_podman_image:
+	podman-hpc build -f docker/Dockerfile -t $(IMAGE):$(VERSION) .
+
+migrate_podman_image: build_podman_image
+	podman-hpc migrate $(IMAGE):$(VERSION)
+
 enter_docker_image: build_docker_image
 	docker run -it --rm $(IMAGE):$(VERSION) bash
 
+launch_beaker_session:
+	./launch-beaker-session.sh $(USERNAME)/$(IMAGE)-$(VERSION)
+
+build_deps_only_image:
+	DOCKER_BUILDKIT=1 docker build --platform=linux/amd64 -f docker/Dockerfile -t $(IMAGE)-deps-only:$(VERSION) --target deps-only .
+	beaker image create $(IMAGE)-deps-only:$(VERSION) --name $(IMAGE)-deps-only-$(VERSION) --workspace ai2/ace-ci-tests
+
 
 # recommended to deactivate current conda environment before running this
 create_environment:
 	conda create -n $(ENVIRONMENT_NAME) python=3.11 $(CONDA_PACKAGES)
 	conda run --no-capture-output -n $(ENVIRONMENT_NAME) python -m pip install uv
 	conda run --no-capture-output -n $(ENVIRONMENT_NAME) uv pip install -c constraints.txt -e .[dev,docs]
+	conda run --no-capture-output -n $(ENVIRONMENT_NAME) uv pip install --no-build-isolation -c constraints.txt -r requirements-healpix.txt
 	conda run --no-capture-output -n $(ENVIRONMENT_NAME) uv pip install -r analysis-deps.txt
 
 test:

PACKAGE_README.md

@@ -1,8 +1,11 @@
 [![Docs](https://readthedocs.org/projects/ai2-climate-emulator/badge/?version=latest)](https://ai2-climate-emulator.readthedocs.io/en/latest/)
-[![PyPI](https://img.shields.io/pypi/v/fme.svg)](https://pypi.org/project/fme/)
 
 # FME: Weather/Climate Model Emulation
-This package contains code to train and evaluate weather/climate model emulators.
+This package contains code to train and evaluate weather/climate model emulators as seen in
+"ACE: A fast, skillful learned global atmospheric model for climate prediction" ([arxiv:2310.02074](https://arxiv.org/abs/2310.02074))
+and "Application of the Ai2 Climate Emulator to E3SMv2's global atmosphere model, with a focus on precipitation fidelity"
+([JGR-ML](https://agupubs.onlinelibrary.wiley.com/doi/full/10.1029/2024JH000136)).
+
 
 ## Installation
 

README.md

@@ -11,8 +11,8 @@ This repo contains code accompanying five papers describing ACE models:
 - "ACE: A fast, skillful learned global atmospheric model for climate prediction" ([link](https://arxiv.org/abs/2310.02074))
 - "Application of the Ai2 Climate Emulator to E3SMv2's global atmosphere model, with a focus on precipitation fidelity" ([link](https://agupubs.onlinelibrary.wiley.com/doi/full/10.1029/2024JH000136))
 - "ACE2: Accurately learning subseasonal to decadal atmospheric variability and forced responses" ([link](https://www.nature.com/articles/s41612-025-01090-0))
-- "ACE2-SOM: Coupling to a slab ocean and learning the sensitivity of climate to changes in CO2" ([link](https://arxiv.org/abs/2412.04418))
-- "Applying the ACE2 Emulator to SST Green's Functions for the E3SMv3 Global Atmosphere Model" ([link](https://arxiv.org/abs/2505.08742))
+- "ACE2-SOM: Coupling an ML Atmospheric Emulator to a Slab Ocean and Learning the Sensitivity of Climate to Changed CO<sub>2</sub>" ([link](https://agupubs.onlinelibrary.wiley.com/doi/full/10.1029/2024JH000575))
+- "Applying the ACE2 Emulator to SST Green's Functions for the E3SMv3 Global Atmosphere Model" ([link](https://agupubs.onlinelibrary.wiley.com/doi/full/10.1029/2025JH000774))
 
 ## Installation
 
@@ -26,4 +26,25 @@ See complete documentation [here](https://ai2-climate-emulator.readthedocs.io/en
 
 ## Model checkpoints
 
-Pretrained model checkpoints, and datasets to run them, are available in the [ACE Hugging Face](https://huggingface.co/collections/allenai/ace-67327d822f0f0d8e0e5e6ca4) collection.
+Pretrained model checkpoints are available in the [ACE Hugging Face](https://huggingface.co/collections/allenai/ace-67327d822f0f0d8e0e5e6ca4) collection.
+
+## Available datasets
+Two versions of the complete dataset described in [arxiv:2310.02074](https://arxiv.org/abs/2310.02074)
+are available on a [requester pays](https://cloud.google.com/storage/docs/requester-pays) Google Cloud Storage bucket:
+```
+gs://ai2cm-public-requester-pays/2023-11-29-ai2-climate-emulator-v1/data/repeating-climSST-1deg-zarrs
+gs://ai2cm-public-requester-pays/2023-11-29-ai2-climate-emulator-v1/data/repeating-climSST-1deg-netCDFs
+```
+The `zarr` format is convenient for ad-hoc analysis. The netCDF version contains our
+train/validation split which was used for training and inference.
+
+The datasets used in the [ACE2 paper](https://arxiv.org/abs/2411.11268) are available at:
+```
+gs://ai2cm-public-requester-pays/2024-11-13-ai2-climate-emulator-v2-amip/data/c96-1deg-shield/
+gs://ai2cm-public-requester-pays/2024-11-13-ai2-climate-emulator-v2-amip/data/era5-1deg-1940-2022.zarr/
+```
+
+The dataset used in the [ACE2-SOM paper](https://arxiv.org/abs/2412.04418) is available at:
+```
+gs://ai2cm-public-requester-pays/2024-12-05-ai2-climate-emulator-v2-som/SHiELD-SOM-C96
+```

analysis-deps.txt

@@ -1,5 +1,7 @@
 # these are convenient to have installed for ad-hoc analysis but are not requirements of the "fme" package
 Bottleneck
+beaker-gantry==3.*
+beaker-py
 bokeh>=3.1.0
 dask[distributed]
 ipywidgets

docker/Dockerfile

@@ -21,6 +21,8 @@ ENV FORCE_CUDA_EXTENSION=1
 # install python deps
 COPY requirements.txt /tmp/requirements.txt
 RUN python3 -m pip install -r /tmp/requirements.txt
+COPY requirements-healpix.txt /tmp/requirements-healpix.txt
+RUN python3 -m pip install --no-build-isolation -r /tmp/requirements-healpix.txt
 
 FROM base AS production
 
@@ -29,4 +31,22 @@ COPY . ${FME_DIR}
 RUN cd $FME_DIR && pip install --no-deps -e .
 
 # copy after install so editing scripts does not trigger reinstall
-COPY scripts ${FME_DIR}/scripts
+COPY scripts ${FME_DIR}/scripts
+
+FROM base AS deps-only
+
+# install github cli and yq
+RUN apt-get update && apt-get install -y \
+    ca-certificates \
+    && curl -fsSL https://cli.github.com/packages/githubcli-archive-keyring.gpg | gpg --dearmor -o /usr/share/keyrings/githubcli-archive-keyring.gpg \
+    && echo "deb [signed-by=/usr/share/keyrings/githubcli-archive-keyring.gpg] https://cli.github.com/packages stable main" | tee /etc/apt/sources.list.d/github-cli.list > /dev/null \
+    && apt-get update && apt-get install -y gh \
+    && apt-get clean \
+    && rm -rf /var/lib/apt/lists/*
+
+RUN wget https://github.com/mikefarah/yq/releases/latest/download/yq_linux_amd64 -O /usr/bin/yq &&\
+chmod +x /usr/bin/yq
+
+COPY scripts/ci_beaker/entrypoint.sh /entrypoint.sh
+RUN chmod +x /entrypoint.sh
+ENTRYPOINT ["/entrypoint.sh"]