diff --git a/.github/workflows/project-build-test.yaml b/.github/workflows/project-build-test.yaml index de0d21fba..e574ee0fc 100644 --- a/.github/workflows/project-build-test.yaml +++ b/.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 diff --git a/container_templates/micromamba.def b/container_templates/micromamba.def index d368ee53d..01d77b39e 100644 --- a/container_templates/micromamba.def +++ b/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@@ diff --git a/container_templates/uv.def b/container_templates/uv.def index dd6469d6b..2475d45d7 100644 --- a/container_templates/uv.def +++ b/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@@ diff --git a/pipeline/config/config.yaml b/pipeline/config/config.yaml index 720fbb4f7..ba8cb9716 100644 --- a/pipeline/config/config.yaml +++ b/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/.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/). 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 diff --git a/pipeline/resources.smk b/pipeline/resources.smk index 8c996a407..042050b32 100644 --- a/pipeline/resources.smk +++ b/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/.sif`. + """The image that `project`'s rules run in. + + By default, the locally built `$AFRAME_CONTAINER_ROOT/.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 diff --git a/projects/infer/infer/local.py b/projects/infer/infer/local.py index ea985b72f..aac45a314 100644 --- a/projects/infer/infer/local.py +++ b/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)) diff --git a/scripts/build_containers.py b/scripts/build_containers.py index e9b699e6b..07d41ac40 100644 --- a/scripts/build_containers.py +++ b/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(): diff --git a/scripts/env_hash.py b/scripts/env_hash.py index abc7b2949..ca6c16dae 100644 --- a/scripts/env_hash.py +++ b/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 | 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 diff --git a/scripts/publish_images.py b/scripts/publish_images.py new file mode 100644 index 000000000..a8fa6af31 --- /dev/null +++ b/scripts/publish_images.py @@ -0,0 +1,111 @@ +"""Publish container images to your OSDF staging directory, where condor +jobs fetch them and `container_source: osdf` finds them. + + python -m scripts.publish_images [data infer plots] + +Run on a CIT AP with /osdf mounted. Copies each +`$AFRAME_CONTAINER_ROOT/.sif` to `aframe--.sif` +in `/igwn/cit/staging/`, where `` is the environment hash +stamped into the image at build time. Only images built with their +environment files committed are published. + +Jobs read images through OSDF caches, which keep serving the copy they +already hold even if the file at CIT is replaced. This means a name must +never be reused for a different image. Any environment change gets a new +hash, and an image already published under its hash is skipped. + +Standard library only, so that the Snakefile can import it. +""" + +import argparse +import getpass +import os +import shutil +import subprocess +from pathlib import Path + +from scripts.env_hash import env_hash + +# The images condor jobs can run in. train and export only run on the AP. +# Must match the CI size check in .github/workflows/project-build-test.yaml. +OSDF_PROJECTS = ("data", "infer", "plots") + + +def staging_dir(user: str | None = None) -> str: + return f"/igwn/cit/staging/{user or getpass.getuser()}" + + +def image_name(project: str, env: str) -> str: + return f"aframe-{project}-{env}.sif" + + +def stamp(image: Path, name: str) -> str | None: + result = subprocess.run( + ["apptainer", "exec", str(image), "cat", f"/opt/{name}"], + capture_output=True, + text=True, + ) + return result.stdout.strip() if result.returncode == 0 else None + + +def publish(project: str, container_root: Path, dest_dir: Path) -> None: + image = container_root / f"{project}.sif" + if not image.exists(): + raise FileNotFoundError( + f"{image} doesn't exist. Use `build_containers` to build it." + ) + env, commit = stamp(image, "env_hash"), stamp(image, "build_commit") + if env is None or commit is None: + raise ValueError( + f"{image} predates build stamps. Rebuild with `build_containers`." + ) + if commit.endswith("-dirty"): + raise ValueError( + f"{image} was built with uncommitted changes to its environment " + "files. Commit them, then rebuild it." + ) + if env != env_hash(project): + print( # noqa: T201 + f"warning: {image} was built for environment {env}, not the " + f"local repo's {env_hash(project)}." + ) + + dest = dest_dir / image_name(project, env) + if dest.exists(): + print(f"{dest} already published") # noqa: T201 + return + print(f"publishing {image} (built from {commit}) to {dest}") # noqa: T201 + try: + shutil.copyfile(image, dest) + except BaseException: + # a partial copy would otherwise count as published next time + dest.unlink(missing_ok=True) + raise + + +def main() -> None: + parser = argparse.ArgumentParser( + description="Publish container images to OSDF staging" + ) + parser.add_argument( + "projects", nargs="*", default=OSDF_PROJECTS, choices=OSDF_PROJECTS + ) + parser.add_argument( + "--container_root", + type=Path, + default=os.getenv("AFRAME_CONTAINER_ROOT"), + ) + args = parser.parse_args() + if args.container_root is None: + raise ValueError("set AFRAME_CONTAINER_ROOT or pass --container_root") + dest_dir = Path("/osdf" + staging_dir()) + if not dest_dir.is_dir(): + raise FileNotFoundError( + f"{dest_dir} doesn't exist. Use an AP with OSDF mounted." + ) + for project in args.projects: + publish(project, args.container_root, dest_dir) + + +if __name__ == "__main__": + main()