Osdf images

.github/workflows/project-build-test.yaml

@@ -75,6 +75,10 @@ jobs:
         username: ${{ github.actor }}
         password: ${{ secrets.GITHUB_TOKEN }}
 
+    # scripts/env_hash.py runs `uv export`
+    -
+      name: Install uv
+      uses: astral-sh/setup-uv@v5
     -
       name: Install definition generation dependencies
       run: python -m pip install jsonargparse spython
@@ -101,6 +105,28 @@ jobs:
           quay.io/singularity/singularity:v3.8.1 \
           -c 'singularity build --sandbox /opt/aframe/sandbox apptainer.def'
 
+    # Every condor job fetches its image from OSDF, and IGWN asks for images
+    # under 5 GB, so check the compressed .sif that jobs transfer. Only for
+    # OSDF_PROJECTS in scripts/publish_images.py.
+    -
+      name: check image size
+      if: contains(fromJSON('["data", "infer", "plots"]'), matrix.project)
+      run: |
+        docker run \
+          --rm \
+          -v ${{ github.workspace }}:/opt/aframe \
+          --privileged \
+          --entrypoint /bin/bash \
+          quay.io/singularity/singularity:v3.8.1 \
+          -c 'singularity build /opt/aframe/image.sif /opt/aframe/sandbox'
+        size=$(stat -c %s image.sif)
+        sudo rm image.sif
+        echo "Image size: $((size / 1024**2)) MB"
+        if [ "$size" -gt $((5 * 1024**3)) ]; then
+          echo "::error::${{ matrix.project }} image exceeds IGWN's 5 GB limit"
+          exit 1
+        fi
+
     # run tests inside the sandbox;
     # if this is a push event or a tag, tar the sandbox
     # so that we can import it into docker

container_templates/micromamba.def

@@ -53,8 +53,11 @@ rm -rf /opt/build-tmp
 micromamba clean -ay
 
 # Record the environment hash (scripts/env_hash.py) so the pipeline can
-# tell whether this image matches the repo it's run from
+# tell whether this image matches the repo it's run from.
+# Also record commit it was built from, with "-dirty" if its environment 
+# files had uncommitted changes.
 echo @@ENV_HASH@@ > /opt/env_hash
+echo @@BUILD_COMMIT@@ > /opt/build_commit
 
 @@PURGE_BUILD_TOOLS@@
 

container_templates/uv.def

@@ -32,8 +32,11 @@ mkdir -p $TMPDIR
 rm -rf /opt/build-tmp
 
 # Record the environment hash (scripts/env_hash.py) so the pipeline can
-# tell whether this image matches the repo it's run from
+# tell whether this image matches the repo it's run from.
+# Also record commit it was built from, with "-dirty" if its environment 
+# files had uncommitted changes.
 echo @@ENV_HASH@@ > /opt/env_hash
+echo @@BUILD_COMMIT@@ > /opt/build_commit
 
 @@PURGE_BUILD_TOOLS@@
 

pipeline/config/config.yaml

@@ -90,6 +90,17 @@ max_num_samples: 3000   # max waveforms generated per rejection-sampling batch
 # Number of condor jobs that validation waveforms are split across
 num_validation_jobs: 200
 
+# --- Containers --------------------------------------------------------------
+# Where the data, infer and plots images come from (train and export are
+# always local):
+#   local: $AFRAME_CONTAINER_ROOT/<project>.sif, built with build-containers
+#   osdf:  the images published for this repo's environment with
+#          `python -m scripts.publish_images`, read through a CIT AP's /osdf
+container_source: local
+# The OSDF staging directory to read published images from. null for your
+# own (/igwn/cit/staging/<USER>). Anyone's can be read.
+osdf_staging_dir: null
+
 # --- Resources ---------------------------------------------------------------
 # Memory (MB) and walltime (minutes) for rules submitted as batch jobs,
 # under slurm or condor. Rules and keys not listed here use the profile's

pipeline/resources.smk

@@ -10,36 +10,54 @@ Also resolves each project's container image.
 """
 
 import os
-import shutil
 import subprocess
 
+from snakemake.exceptions import WorkflowError
 from snakemake.logging import logger
 
 from scripts.env_hash import env_hash
+from scripts.publish_images import OSDF_PROJECTS, image_name, staging_dir
 
 
 def container(project):
-    """The image to run `project`'s rules in:
-    `$AFRAME_CONTAINER_ROOT/<project>.sif`.
+    """The image that `project`'s rules run in.
+
+    By default, the locally built `$AFRAME_CONTAINER_ROOT/<project>.sif`.
+    With `container_source: osdf`, the data, infer and plots rules instead
+    use the image published for the local repo's environment, read from
+    `osdf_staging_dir` (your own staging directory by default) through the
+    AP's `/osdf` mount.
     """
+    if config.get("container_source", "local") == "osdf" and project in OSDF_PROJECTS:
+        name = image_name(project, env_hash(project))
+        source = config.get("osdf_staging_dir") or staging_dir()
+        return f"/osdf{source}/{name}"
     return os.path.join(os.getenv("AFRAME_CONTAINER_ROOT", ""), f"{project}.sif")
 
 
 def check_images(projects=("data", "train", "export", "infer", "plots")):
-    """Warn about images built for a different environment than the local
-    repo's.
+    """Fail on missing published images, and warn about local images built
+    for a different environment than the local repo's.
 
     Images record their environment hash (scripts/env_hash.py) at build time.
     """
-    exe = shutil.which("apptainer") or shutil.which("singularity")
-    if exe is None:
-        return
     for project in projects:
         image = container(project)
+        if image.startswith("/osdf/"):
+            # published images are named by their environment hash
+            if not os.path.exists(image):
+                raise WorkflowError(
+                    f"{image} doesn't exist. Either the local repo's "
+                    f"{project} environment hasn't been published (build the "
+                    "image, then run `python -m scripts.publish_images "
+                    f"{project}` on a CIT AP), or this isn't a CIT AP. "
+                    "Otherwise, set `container_source: local`."
+                )
+            continue
         if not os.path.exists(image):
             continue
         result = subprocess.run(
-            [exe, "exec", image, "cat", "/opt/env_hash"],
+            ["apptainer", "exec", image, "cat", "/opt/env_hash"],
             capture_output=True,
             text=True,
         )
@@ -90,5 +108,5 @@ def gpu_resources():
     else:
         res["gpus_minimum_capability"] = config["gpu_min_capability"]
         if config.get("gpu_min_memory_mb"):
-            res["gpus_minimum_memory"] = config["gpu_min_memory_mb"]
+            res["gpus_minimum_memory"] = f"{config['gpu_min_memory_mb']}M"
     return res

projects/infer/infer/local.py

@@ -22,8 +22,9 @@ def build_model(weights, backend, device, aoti_path=None):
         if aoti_path is None:
             raise ValueError("backend 'aoti' requires aoti_path")
         # aoti_load_package references torch._inductor.codecache without
-        # importing it
-        import torch._inductor.codecache  # noqa: F401
+        # importing it. `import torch._inductor.codecache` would make
+        # `torch` local to this function, breaking the other backends.
+        from torch._inductor import codecache  # noqa: F401
 
         runner = torch._inductor.aoti_load_package(str(aoti_path))
 

scripts/build_containers.py

@@ -6,7 +6,8 @@
 from jsonargparse import ArgumentParser
 from spython.main import Client
 
-from scripts.env_hash import env_hash, local_libs
+from scripts.env_hash import build_commit, env_hash, local_libs, uv_command
+from scripts.publish_images import OSDF_PROJECTS
 
 # Define the directory where the projects are located
 ROOT_DIR: Path = Path(__file__).resolve().parent.parent
@@ -17,11 +18,6 @@
 # List of all available project names
 PROJECTS: list[str] = [x.name for x in BASE_DIR.iterdir() if x.is_dir()]
 
-# Extras to install into each project's container.
-# Currently only needed for `data`, which uses extras to keep CUDA-torch
-# out of its container.
-EXTRAS: dict[str, list[str]] = {"data": ["cpu"]}
-
 # Clear out the tools used to build the environment once complete to shrink
 # container size. A project that needs a compiler at run time should install
 # it from its apptainer.post, which marks it manually installed and so
@@ -82,19 +78,14 @@ def _get_files_block(project_name: str) -> str:
 
 def _get_uv_command(project_name: str, subcommand: str) -> str:
     """
-    Build a `uv sync`/`uv export` command for a project. The `test`
-    group is installed so that CI can run tests inside the container.
+    Build a `uv sync`/`uv export` command for a project with the same
+    arguments that scripts/env_hash.py hashes.
     """
-    cmd = (
-        f"uv {subcommand} --frozen --no-default-groups --group test"
-        f" --package {project_name}"
-    )
+    cmd = " ".join(uv_command(project_name, subcommand))
     if subcommand == "export":
         # Need to use the pylock format here rather than requirements.txt
         # so that the index each package was locked from gets recorded.
         cmd += " --format pylock.toml"
-    for extra in EXTRAS.get(project_name, []):
-        cmd += f" --extra {extra}"
     return cmd
 
 
@@ -130,6 +121,7 @@ def create_definition_file(project_name: str) -> Path:
         .replace("@@EXTRA_ENV@@", extra_env)
         .replace("@@PURGE_BUILD_TOOLS@@", PURGE_BUILD_TOOLS)
         .replace("@@ENV_HASH@@", env_hash(project_name))
+        .replace("@@BUILD_COMMIT@@", build_commit(project_name))
     )
 
     output_path = project_dir / "apptainer.def"
@@ -183,6 +175,20 @@ def build(projects: list[str], container_root: Path, max_workers: int) -> None:
         logging.info("Container root path is not set.")
         return
 
+    # publish_images refuses images built from uncommitted environment files
+    unpublishable = [
+        p
+        for p in projects
+        if p in OSDF_PROJECTS and build_commit(p).endswith("-dirty")
+    ]
+    if unpublishable:
+        warning = (
+            "Environment files for the following projects have uncommitted "
+            f"changes: {', '.join(unpublishable)}. Their images will build, "
+            "but can't be published until they're committed and rebuilt."
+        )
+        logging.warning(warning)
+
     failed_projects = []
     with ProcessPoolExecutor(max_workers=max_workers) as executor:
         futures = {
@@ -206,6 +212,8 @@ def build(projects: list[str], container_root: Path, max_workers: int) -> None:
         )
     else:
         logging.info("All containers built successfully")
+    if unpublishable:
+        logging.warning(warning)
 
 
 def main():

scripts/env_hash.py

@@ -4,9 +4,16 @@
 image only has to be updated when the environment changes. Images
 record this hash at build time and the pipeline compares it with
 the local repo's to warn about stale images.
+
+Only the packages the project's image installed are hashed so that
+a lock change in one project doesn't impact other projects.
+
+Standard library only so that the Snakefile can import it.
 """
 
 import hashlib
+import json
+import subprocess
 import sys
 import tomllib
 from pathlib import Path
@@ -15,6 +22,16 @@
 PROJECTS_DIR: Path = ROOT_DIR / "projects"
 LIBS_DIR: Path = ROOT_DIR / "libs"
 TEMPLATES_DIR: Path = ROOT_DIR / "container_templates"
+LOCK_FILE: Path = ROOT_DIR / "uv.lock"
+
+# Extras to install into each project's container.
+# Currently only needed for `data`, which uses extras to keep CUDA-torch
+# out of its container.
+EXTRAS: dict[str, list[str]] = {"data": ["cpu"]}
+
+# Dependency groups installed into every container, so that CI can run
+# tests inside it
+GROUPS: list[str] = ["test"]
 
 
 def local_libs(project: str) -> list[str]:
@@ -43,14 +60,45 @@ def local_libs(project: str) -> list[str]:
     return sorted(seen)
 
 
+def uv_command(project: str, subcommand: str) -> list[str]:
+    """
+    The `uv sync`/`uv export` command that installs a project's environment.
+    """
+    args = ["uv", subcommand, "--frozen", "--no-default-groups"]
+    for group in GROUPS:
+        args += ["--group", group]
+    args += ["--package", project]
+    for extra in EXTRAS.get(project, []):
+        args += ["--extra", extra]
+    return args
+
+
+def locked_requirements(project: str) -> bytes:
+    """The packages the project's image installs, pinned with their hashes,
+    as `uv export` resolves them from `uv.lock`.
+    """
+    args = [*uv_command(project, "export"), "--no-header", "--no-annotate"]
+    return subprocess.check_output(args, cwd=ROOT_DIR)
+
+
+def uv_settings() -> str:
+    """The root `pyproject.toml`'s `[tool.uv]` table, which can change an
+    install without changing the lock (e.g. build settings). The rest of the
+    file doesn't affect images.
+    """
+    with open(ROOT_DIR / "pyproject.toml", "rb") as f:
+        settings = tomllib.load(f)["tool"]["uv"]
+    return json.dumps(settings, sort_keys=True)
+
+
 def env_files(project: str) -> list[Path]:
-    """Every file whose contents impact the project's environment."""
+    """Every file, besides `uv.lock` and the root `pyproject.toml`, whose
+    contents impact the project's environment.
+    """
     project_dir = PROJECTS_DIR / project
     conda_lock = project_dir / f"{project}.conda-lock.yml"
     template = "micromamba.def" if conda_lock.exists() else "uv.def"
     candidates = [
-        ROOT_DIR / "uv.lock",
-        ROOT_DIR / "pyproject.toml",
         ROOT_DIR / "scripts" / "build_containers.py",
         TEMPLATES_DIR / template,
         project_dir / "pyproject.toml",
@@ -63,22 +111,42 @@ def env_files(project: str) -> list[Path]:
 
 
 def manifest(project: str) -> str:
-    """`sha256sum`-style lines for each environment file, sorted by path."""
+    """`sha256sum`-style lines for each environment file, sorted by path,
+    then the project's packages from `uv.lock` and the root `[tool.uv]`.
+    """
+
+    def line(content: bytes, name: str) -> str:
+        return f"{hashlib.sha256(content).hexdigest()}  {name}\n"
+
     paths = sorted(str(p.relative_to(ROOT_DIR)) for p in env_files(project))
-    return "".join(
-        f"{hashlib.sha256((ROOT_DIR / p).read_bytes()).hexdigest()}  {p}\n"
-        for p in paths
-    )
+    lines = [line((ROOT_DIR / p).read_bytes(), p) for p in paths]
+    lines.append(line(locked_requirements(project), "uv.lock"))
+    lines.append(line(uv_settings().encode(), "pyproject.toml [tool.uv]"))
+    return "".join(lines)
 
 
 def env_hash(project: str) -> str:
-    """A short hash of the files that impact the project's environment.
-
-    Hashes the manifest, so from the repo root this equals the first 12
-    characters of `sha256sum <files in order> | sha256sum`.
+    """A short hash of the project's environment: the first 12 characters
+    of the sha256 of its manifest.
     """
     return hashlib.sha256(manifest(project).encode()).hexdigest()[:12]
 
 
+def build_commit(project: str) -> str:
+    """The local repo's commit, with `-dirty` if any of the project's
+    environment files, `uv.lock` or the root `pyproject.toml` differ from it.
+    """
+    root_files = [LOCK_FILE, ROOT_DIR / "pyproject.toml"]
+    paths = [str(p) for p in [*env_files(project), *root_files]]
+
+    def git(*args: str) -> str:
+        cmd = ["git", *args]
+        return subprocess.check_output(cmd, cwd=ROOT_DIR, text=True).strip()
+
+    commit = git("rev-parse", "HEAD")
+    dirty = git("status", "--porcelain", "--", *paths)
+    return f"{commit}-dirty" if dirty else commit
+
+
 if __name__ == "__main__":
     print(env_hash(sys.argv[1]))  # noqa: T201