Dagster Orchestration of Azure Batch Pipelines

.containerignore

@@ -1,3 +1,4 @@
 Containerfile
 nssp_demo/private_data
 notebooks
+mounts/

.dockerignore

@@ -0,0 +1,4 @@
+Containerfile
+nssp_demo/private_data
+notebooks
+mounts

.github/workflows/update-dagster-code.yaml

@@ -0,0 +1,27 @@
+name: Deploy Dagster Code
+
+on:
+  push:
+    branches:
+      - main
+  pull_request:
+    branches:
+      - main
+  workflow_dispatch:
+
+jobs:
+  run:
+    runs-on: ubuntu-latest
+    environment: production
+    name: Update Dagster Code Location
+    steps:
+      - name: Run update script with cfa runner-action
+        uses: CDCgov/cfa-actions/[email protected]
+        with:
+          github_app_id: ${{ secrets.REPO_CDCENT_ACTOR_APP_ID }}
+          github_app_pem: ${{ secrets.REPO_CDCENT_ACTOR_APP_PEM }}
+          script: |
+            echo "Running update script"
+            uv run \
+              https://raw.githubusercontent.com/CDCgov/cfa-dagster/refs/heads/main/scripts/update_code_location.py \
+              --registry_image cfaprdbatchcr.azurecr.io/pyrenew-hew:dagster_latest

.gitignore

@@ -29,6 +29,8 @@
 *.gif
 *.pdf
 *.png
+# except when in img/ folder
+!img/*.png
 
 # Compressed archives
 *.gz
@@ -401,7 +403,16 @@ nssp-etl
 output
 params
 nwss-vintages
+mounts
+prod-param-estimates
+pyrenew-hew-config
+pyrenew-hew-prod-output
+pyrenew-test-output
+test-output
 
 # Azure configuration files
 azureconfig.env
 azureconfig.sh
+
+# blobfuse
+config.yaml

.pre-commit-config.yaml

@@ -19,6 +19,7 @@ repos:
       - id: ruff-check
         # Run the formatter
       - id: ruff-format
+  #####extra
   - repo: https://github.com/astral-sh/uv-pre-commit
     rev: 0.9.17
     hooks:

Containerfile

@@ -1,6 +1,6 @@
 #syntax=docker/dockerfile:1-labs
 
-FROM rocker/tidyverse
+FROM rocker/tidyverse:4.5.1
 
 ARG GIT_COMMIT_SHA
 ENV GIT_COMMIT_SHA=$GIT_COMMIT_SHA
@@ -11,23 +11,43 @@ ENV GIT_BRANCH_NAME=$GIT_BRANCH_NAME
 ENV XLA_FLAGS=--xla_force_host_platform_device_count=4
 
 COPY ./hewr /pyrenew-hew/hewr
-
 WORKDIR /pyrenew-hew
 
+# install hewr dependencies
 RUN Rscript -e "install.packages('pak')"
 RUN Rscript -e "pak::pkg_install('cmu-delphi/epiprocess@main')"
 RUN Rscript -e "pak::pkg_install('cmu-delphi/epipredict@main')"
 RUN Rscript -e "pak::local_install('hewr', upgrade = FALSE)"
 
-COPY --exclude=pipelines/priors . .
-COPY pipelines/priors pipelines/priors
-
+#
 # Python from https://docs.astral.sh/uv/guides/integration/docker/
 COPY --from=ghcr.io/astral-sh/uv:latest /uv /uvx /bin/
 # Some handy uv environment variables
 ENV UV_COMPILE_BYTECODE=1
 ENV UV_LINK_MODE=copy
 ENV UV_PYTHON_CACHE_DIR=/root/.cache/uv/python
 
-RUN --mount=type=cache,target=/root/.cache/uv \
-    uv sync
+# copy in the project files
+COPY ./pyrenew_hew ./pyrenew_hew
+COPY ./pipelines ./pipelines
+COPY ./tests ./tests
+COPY README.md ./
+COPY ./pyproject.toml ./
+COPY ./uv.lock ./
+COPY ./.python-version ./
+
+# VENV MANAGEMENT AND DEPENDENCY SYNCING
+
+# Create the Pyrenew-Hew venv and sync dependencies from pyproject.toml
+RUN uv venv .venv
+RUN --mount=type=cache,target=/root/.cache/uv
+# Set VIRTUAL_ENV variable at runtime to choose which venv to activate
+# By default we'll do the main pyrenew-hew venv
+ENV VIRTUAL_ENV=/pyrenew-hew/.venv
+RUN uv sync
+
+# Copy in the dagster defs
+COPY ./dagster_defs.py ./
+
+# Update PATH to use the selected venv
+ENV PATH="${VIRTUAL_ENV}/bin:$PATH"

Makefile

@@ -1,21 +1,29 @@
-.PHONY: help container_build container_tag ghcr_login container_push run_timeseries run_e_model run_h_models post_process
+.PHONY: help container_build CONTAINER_IMAGE_VERSION ghcr_login container_push run_timeseries run_e_model run_h_models post_process run_he_model run_hw_model acc mount unmount config dagster dagster_build dagster_push
 
 # Build parameters
 
 ifndef ENGINE
 ENGINE = docker
 endif
 
-ifndef CONTAINER_NAME
-CONTAINER_NAME = pyrenew-hew
+ifndef CONTAINER_REGISTRY
+CONTAINER_REGISTRY = ghcr.io/cdcgov
 endif
 
-ifndef CONTAINERFILE
-CONTAINERFILE = Containerfile
+ifndef CONTAINER_IMAGE_NAME
+CONTAINER_IMAGE_NAME = pyrenew-hew
+endif
+
+ifndef CONTAINER_IMAGE_VERSION
+CONTAINER_IMAGE_VERSION = latest
 endif
 
 ifndef CONTAINER_REMOTE_NAME
-CONTAINER_REMOTE_NAME = ghcr.io/cdcgov/$(CONTAINER_NAME):latest
+CONTAINER_REMOTE_NAME = $(CONTAINER_REGISTRY)/$(CONTAINER_IMAGE_NAME):$(CONTAINER_IMAGE_VERSION)
+endif
+
+ifndef CONTAINERFILE
+CONTAINERFILE = Containerfile
 endif
 
 # Model Fit Parameters
@@ -47,13 +55,22 @@ endif
 help:
 	@echo "Usage: make [target] [ARGS]"
 	@echo ""
+
+	@echo "Blobfuse Mount Targets: "
+	@echo "  mount              : Mount blob storage containers using blobfuse2"
+	@echo "  unmount            : Unmount blob storage containers and clean up"
+	@echo ""
 	@echo "Container Build Targets: "
 	@echo "  container_build     : Build the container image"
-	@echo "  container_tag       : Tag the container image"
+	@echo "  dagster             : Run dagster definitions locally"
+	@echo "  dagster_build       : Build the dagster container image"
+	@echo "  dagster_push        : Push the dagster container image to the Azure Container Registry and code location"
+	@echo "  CONTAINER_IMAGE_VERSION : Tag the container image"
 	@echo "  ghcr_login          : Log in to the Github Container Registry. Requires GH_USERNAME and GH_PAT env vars"
 	@echo "  container_push      : Push the container image to the Azure Container Registry"
 	@echo ""
 	@echo "Model Fit Targets: "
+	@echo "  acc                 : Run the Azure Command Center for routine production jobs"
 	@echo "  run_timeseries      : Run the timeseries model fit job"
 	@echo "  run_e_model         : Run an e model fit job"
 	@echo "  run_h_models        : Run an h model fit job"
@@ -73,104 +90,167 @@ help:
 	@echo ""
 	@echo "To run the pyrenew-hew model and output to pyrenew-test-output:"
 	@echo "  make run_hew_model TEST=True MODEL_LETTERS=hew"
+	@echo "To use a custom container registry, image, and tag:"
+	@echo "  make run_hew_model CONTAINER_REGISTRY=<custom_registry> CONTAINER_IMAGE_NAME=<custom_image> CONTAINER_IMAGE_VERSION=<custom_tag>"
 	@echo ""
 	@echo "Any additional flags can be passed with ARGS, for example:"
 	@echo "  make run_hew_model ARGS=\"--locations-include 'NY GA'\""
 	@echo ""
 	@echo "Passing a flag through ARGS will also override the flags set previously."
 
+#------------------------ #
+# Blobfuse Mount Targets
+# ----------------------- #
+
+mount:
+	sudo bash -c "source ./blobfuse/mount.sh"
+
+unmount:
+	sudo bash -c "source ./blobfuse/cleanup.sh"
+
 # ----------------------- #
 # Container Build Targets
 # ----------------------- #
 
 container_build: ghcr_login
-	$(ENGINE) build . -t $(CONTAINER_NAME) -f $(CONTAINERFILE)
+	$(ENGINE) build . -t $(CONTAINER_IMAGE_NAME) -f $(CONTAINERFILE)
+
+dagster_build:
+	docker build -t cfaprdbatchcr.azurecr.io/pyrenew-hew:dagster_latest -f Containerfile .
 
-container_tag:
-	$(ENGINE) tag $(CONTAINER_NAME) $(CONTAINER_REMOTE_NAME)
+dagster_push: dagster_build
+	az login --identity && \
+	az acr login -n cfaprdbatchcr && \
+	docker push "cfaprdbatchcr.azurecr.io/pyrenew-hew:dagster_latest" && \
+	uv run https://raw.githubusercontent.com/CDCgov/cfa-dagster/refs/heads/main/scripts/update_code_location.py \
+	--registry_image "cfaprdbatchcr.azurecr.io/pyrenew-hew:dagster_latest"
+
+dagster:
+	uv run dagster_defs.py
+
+CONTAINER_IMAGE_VERSION:
+	$(ENGINE) tag $(CONTAINER_IMAGE_NAME) $(CONTAINER_REMOTE_NAME)
 
 ghcr_login:
 	echo $(GH_PAT) | $(ENGINE) login ghcr.io -u $(GH_USERNAME) --password-stdin
 
-container_push: container_tag ghcr_login
+container_push: CONTAINER_IMAGE_VERSION ghcr_login
 	$(ENGINE) push $(CONTAINER_REMOTE_NAME)
 
+config:
+	bash -c "source ./azureconfig.sh"
+
 # ---------------- #
 # Model Fit Targets
 # ---------------- #
 
-run_timeseries:
+acc: mount config
+	uv run pipelines/azure_command_center.py
+
+# Auto-set TEST/ENVIRONMENT based on each other
+ifeq ($(shell echo $(ENVIRONMENT) | tr A-Z a-z),test)
+override TEST = True
+endif
+ifeq ($(shell echo $(TEST) | tr A-Z a-z),true)
+override ENVIRONMENT = test
+endif
+
+run_timeseries: config
 	uv run python pipelines/batch/setup_job.py \
 		--model-family timeseries \
 		--output-subdir "${FORECAST_DATE}_forecasts" \
 		--model-letters "e" \
-		--job-id "pyrenew-e-${ENVIRONMENT}_${FORECAST_DATE}_t" \
+		--locations-exclude "WY" \
+		--job-id "timeseries-e-${ENVIRONMENT}_${FORECAST_DATE}_makefile" \
 		--pool-id pyrenew-pool \
 		--test "$(TEST)" \
 		--dry-run "$(DRY_RUN)" \
+		--container-registry "$(CONTAINER_REGISTRY)" \
+		--container-image-name "$(CONTAINER_IMAGE_NAME)" \
+		--container-image-version "$(CONTAINER_IMAGE_VERSION)" \
 		$(ARGS)
 
-run_e_model:
+run_e_model: config
 	uv run python pipelines/batch/setup_job.py \
 		--model-family pyrenew \
 		--output-subdir "${FORECAST_DATE}_forecasts" \
 		--model-letters "e" \
-		--job-id "pyrenew-e-${ENVIRONMENT}_${FORECAST_DATE}" \
+		--locations-exclude "WY" \
+		--job-id "pyrenew-e-${ENVIRONMENT}_${FORECAST_DATE}_makefile" \
 		--pool-id pyrenew-pool \
 		--rng-key "$(RNG_KEY)" \
 		--test "$(TEST)" \
 		--dry-run "$(DRY_RUN)" \
+		--container-registry "$(CONTAINER_REGISTRY)" \
+		--container-image-name "$(CONTAINER_IMAGE_NAME)" \
+		--container-image-version "$(CONTAINER_IMAGE_VERSION)" \
 		$(ARGS)
 
-run_h_model:
+run_h_model: config
 	uv run python pipelines/batch/setup_job.py \
 		--model-family pyrenew \
 		--output-subdir "${FORECAST_DATE}_forecasts" \
 		--model-letters "h" \
-		--job-id "pyrenew-h-${ENVIRONMENT}_${FORECAST_DATE}" \
+		--job-id "pyrenew-h-${ENVIRONMENT}_${FORECAST_DATE}_makefile" \
 		--pool-id pyrenew-pool \
 		--rng-key "$(RNG_KEY)" \
 		--test "$(TEST)" \
 		--dry-run "$(DRY_RUN)" \
+		--container-registry "$(CONTAINER_REGISTRY)" \
+		--container-image-name "$(CONTAINER_IMAGE_NAME)" \
+		--container-image-version "$(CONTAINER_IMAGE_VERSION)" \
 		$(ARGS)
 
-run_he_model:
+run_he_model: config
 	uv run python pipelines/batch/setup_job.py \
 		--model-family pyrenew \
 		--output-subdir "${FORECAST_DATE}_forecasts" \
 		--model-letters "he" \
-		--job-id "pyrenew-he-${ENVIRONMENT}_${FORECAST_DATE}" \
+		--locations-exclude "WY" \
+		--job-id "pyrenew-he-${ENVIRONMENT}_${FORECAST_DATE}_makefile" \
 		--pool-id pyrenew-pool \
 		--rng-key "$(RNG_KEY)" \
 		--test "$(TEST)" \
 		--dry-run "$(DRY_RUN)" \
+		--container-registry "$(CONTAINER_REGISTRY)" \
+		--container-image-name "$(CONTAINER_IMAGE_NAME)" \
+		--container-image-version "$(CONTAINER_IMAGE_VERSION)" \
 		$(ARGS)
 
-run_hw_model:
+run_hw_model: config
 	uv run python pipelines/batch/setup_job.py \
 		--model-family pyrenew \
 		--output-subdir "${FORECAST_DATE}_forecasts" \
 		--model-letters "hw" \
-		--job-id "pyrenew-hw-${ENVIRONMENT}_${FORECAST_DATE}" \
-		--pool-id pyrenew-pool-32gb \
+		--diseases "COVID-19" \
+		--job-id "pyrenew-hw-${ENVIRONMENT}_${FORECAST_DATE}_makefile" \
+		--pool-id pyrenew-pool \
 		--rng-key "$(RNG_KEY)" \
 		--test "$(TEST)" \
 		--dry-run "$(DRY_RUN)" \
+		--container-registry "$(CONTAINER_REGISTRY)" \
+		--container-image-name "$(CONTAINER_IMAGE_NAME)" \
+		--container-image-version "$(CONTAINER_IMAGE_VERSION)" \
 		$(ARGS)
 
-run_hew_model:
+run_hew_model: config
 	uv run python pipelines/batch/setup_job.py \
 		--model-family pyrenew \
 		--output-subdir "${FORECAST_DATE}_forecasts" \
 		--model-letters "hew" \
-		--job-id "pyrenew-hew-${ENVIRONMENT}_${FORECAST_DATE}" \
-		--pool-id pyrenew-pool-32gb \
+		--locations-exclude "WY" \
+		--diseases "COVID-19" \
+		--job-id "pyrenew-hew-${ENVIRONMENT}_${FORECAST_DATE}_makefile" \
+		--pool-id pyrenew-pool \
 		--rng-key "$(RNG_KEY)" \
 		--test "$(TEST)" \
 		--dry-run "$(DRY_RUN)" \
+		--container-registry "$(CONTAINER_REGISTRY)" \
+		--container-image-name "$(CONTAINER_IMAGE_NAME)" \
+		--container-image-version "$(CONTAINER_IMAGE_VERSION)" \
 		$(ARGS)
 
-post_process:
+post_process: config
 	uv run python pipelines/postprocess_forecast_batches.py \
     	--input "./blobfuse/mounts/pyrenew-hew-prod-output/${FORECAST_DATE}_forecasts" \
     	--output "./blobfuse/mounts/nssp-etl/gold/${FORECAST_DATE}_forecasts.parquet" \