diff --git a/adept/cloud.py b/adept/cloud.py new file mode 100644 index 00000000..e4727491 --- /dev/null +++ b/adept/cloud.py @@ -0,0 +1,483 @@ +"""Submit a working tree to AWS Batch as a content-addressed code bundle. + +The Batch sim images (`sim-cpu`, `sim-gpu`) carry the environment only -- no application code. +Code arrives at run time from an S3 tarball, so editing a line of physics costs an upload of a few +hundred KB instead of a docker build, an ECR push and a `cdk deploy`. This is the submit half of +that contract; the runtime half is `continuum-infra/sim-runner/bootstrap.sh`. + +adept is the natural home for it: every sim repo that submits (ml-for-lpi, kinetic-srs, +vp-turbulence) already depends on adept, and adept already depends on boto3, so a shared client +reaches all of them with no new dependency and no vendored copies. + + from adept.cloud import submit + + result = submit( + cmd="python lpi-scan.py --backend aws-worker --manifest scans/abc123", + queue="gpu", + array=48, + extras="gpu", + ) + print(result.job_id, result.adept_sha) + +Bundling uses `git ls-files -co --exclude-standard`: tracked files plus untracked ones that +.gitignore does not exclude. That carries uncommitted edits -- the thing iteration actually +consists of -- while keeping .venv, artifacts and mlflow.db out for free. + +`array=N` submits a Batch array job. The command can read `$AWS_BATCH_JOB_ARRAY_INDEX`, which +survives Batch's `Ref::` substitution unexpanded and is expanded by `bootstrap.sh` in the +container, so an N-point scan is one submission rather than N. +""" + +from __future__ import annotations + +import argparse +import gzip +import hashlib +import io +import os +import re +import subprocess +import sys +import tarfile +from collections.abc import Iterable, Mapping +from dataclasses import dataclass +from pathlib import Path + +# Bundled regardless of git status: uv.lock is what the runner installs from (`uv sync --frozen`), +# and it may have just been regenerated by relock=True without being committed. Some repos +# gitignore their lock outright -- kinetic-srs does -- so a bundler that merely respects +# .gitignore ships no lock and the runner has nothing to install from. +ALWAYS_INCLUDE = ("uv.lock", "pyproject.toml") + +EXCLUDE_DIRS = frozenset({".git", ".venv", "venv", "__pycache__", ".pytest_cache", ".ruff_cache", "node_modules"}) + +# .gitignore is not sufficient on its own, because these repos *commit* their outputs: kinetic-srs +# tracks ~42 MB of plots under sims/, vp-turbulence ~463 MB of .nc. Derived formats are therefore +# dropped by suffix as well -- but never silently: whatever is skipped is reported by count and +# size, and include_suffixes puts one back when a run genuinely needs it as input. +EXCLUDE_SUFFIXES = frozenset( + { + ".pyc", ".pyo", ".so", ".o", ".a", + ".png", ".jpg", ".jpeg", ".gif", ".svg", ".pdf", ".eps", ".mp4", ".webm", + ".h5", ".hdf5", ".nc", ".npz", ".npy", ".ckpt", ".db", ".sqlite", ".parquet", + ".zip", ".tar", ".tgz", ".whl", + } +) # fmt: skip + +# Report the biggest members once a bundle passes this, so 40 MB of accidental payload is noticed +# on the first submit rather than after a hundred of them. +LARGE_BUNDLE_BYTES = 20_000_000 + +DEFAULT_JOB_DEFINITIONS = {"gpu": "sim-gpu", "gpu-48g": "sim-gpu", "cpu": "sim-cpu", "cpu-hmem": "sim-cpu-hmem"} + +# Only the digits of a git sha, so a `?rev=main#` pin resolves to the commit actually locked +# rather than to the branch name. +_ADEPT_IN_LOCK = re.compile(r"github\.com/ergodicio/adept(?:\.git)?[^#\"']*#([0-9a-f]{7,40})") + + +@dataclass(frozen=True) +class Bundle: + """A deterministic tarball of a working tree, addressed by the digest of its own bytes.""" + + repo: Path + prefix: str + blob: bytes + digest: str + files: tuple[Path, ...] + dropped: dict[str, tuple[int, int]] + adept_sha: str | None + + @property + def key(self) -> str: + return f"{self.prefix}/{self.digest}.tar.gz" + + def report(self, out=sys.stdout) -> None: + """Print what was skipped, how big the bundle is, and the largest members if it is fat.""" + for reason, (count, total) in sorted(self.dropped.items(), key=lambda kv: -kv[1][1]): + print(f"[cloud] skipped {count} {reason} ({total / 1e6:.1f} MB)", file=out) + print( + f"[cloud] {len(self.files)} files, {len(self.blob) / 1e6:.1f} MB compressed, sha256:{self.digest[:12]}", + file=out, + ) + if self.adept_sha: + print(f"[cloud] adept pinned at {self.adept_sha}", file=out) + on_disk = sum((self.repo / f).stat().st_size for f in self.files if (self.repo / f).is_file()) + if on_disk > LARGE_BUNDLE_BYTES: + self._report_largest(out) + + def _report_largest(self, out, limit: int = 5) -> None: + sized = sorted(((self.repo / f).stat().st_size, f) for f in self.files if (self.repo / f).is_file())[-limit:] + print(f"[cloud] bundle is over {LARGE_BUNDLE_BYTES / 1e6:.0f} MB; largest members:", file=out) + for size, rel in reversed(sized): + print(f"[cloud] {size / 1e6:7.1f} MB {rel}", file=out) + top = sorted({str(f.parts[0]) for _, f in sized}) + if top: + print(f"[cloud] drop a directory with e.g. exclude=['{top[-1]}/*']", file=out) + + +@dataclass(frozen=True) +class SubmitResult: + """What a submission produced. `job_id` is None for a dry run.""" + + job_id: str | None + job_name: str + queue: str + job_definition: str + cmd: str + code_uri: str + digest: str + array_size: int + adept_sha: str | None + uploaded: bool + + +def _git(repo: Path, *args: str) -> str: + return subprocess.run(["git", "-C", str(repo), *args], check=True, capture_output=True, text=True).stdout + + +def repo_prefix(repo: Path) -> str: + """The S3 key prefix for a repo: its name, or the *main* worktree's name if this is a linked one. + + Iterating in a `git worktree` is normal here, and the directory a worktree happens to sit in is + a branch nickname, not a project. Keying bundles on it would scatter one repo's objects across + prefixes -- and the lifecycle rule and any manual audit both read by prefix. + """ + try: + common = Path(_git(repo, "rev-parse", "--path-format=absolute", "--git-common-dir").strip()) + except subprocess.CalledProcessError: + return repo.name + # /.git for a normal clone or a linked worktree; a bare repo has no parent name + # worth using, so fall back to the directory being bundled. + return common.parent.name if common.name == ".git" else repo.name + + +def adept_sha_from_lock(lock_text: str) -> str | None: + """The adept commit a `uv.lock` resolves to, or None if adept is not a git dependency there. + + Both dependent sim repos declare `adept @ git+...@main`, i.e. unpinned, so without this a run's + adept version is whatever was current whenever the environment was last built, with no record + of which. The lock inside the bundle is what `uv sync --frozen` installs, so this is the sha + the job actually runs -- worth recording next to the Batch job id. + """ + match = _ADEPT_IN_LOCK.search(lock_text) + return match.group(1) if match else None + + +def collect_files( + repo: Path, exclude: Iterable[str] = (), include_suffixes: Iterable[str] = () +) -> tuple[list[Path], dict[str, tuple[int, int]]]: + """Returns (files to bundle, {reason: (count, bytes)} for everything skipped).""" + keep_suffixes = {s if s.startswith(".") else f".{s}" for s in (x.lower() for x in include_suffixes)} + patterns = list(exclude) + listed = _git(repo, "ls-files", "-co", "--exclude-standard").splitlines() + dropped: dict[str, tuple[int, int]] = {} + + def drop(reason: str, path: Path) -> None: + count, total = dropped.get(reason, (0, 0)) + size = (repo / path).stat().st_size if (repo / path).is_file() else 0 + dropped[reason] = (count + 1, total + size) + + kept: list[Path] = [] + for rel in listed: + if not rel: + continue + path = Path(rel) + if not (repo / path).is_file(): + continue + if EXCLUDE_DIRS & set(path.parts): + drop("build/venv directories", path) + continue + suffix = path.suffix.lower() + if suffix in EXCLUDE_SUFFIXES and suffix not in keep_suffixes: + drop(f"derived {suffix} files", path) + continue + if any(path.match(pattern) or str(path).startswith(pattern.rstrip("*/")) for pattern in patterns): + drop("excluded by caller", path) + continue + kept.append(path) + + for extra in ALWAYS_INCLUDE: + if (repo / extra).is_file() and Path(extra) not in kept: + kept.append(Path(extra)) + return sorted(set(kept)), dropped + + +def build_blob(repo: Path, files: Iterable[Path], extra_files: Mapping[str, bytes | str] | None = None) -> bytes: + """Deterministic tar.gz: identical trees must hash identically or content addressing is moot. + + The tar is built uncompressed and gzipped separately with mtime=0. tarfile's own "w:gz" mode + stamps the gzip header with the current time and offers no way to override it, which makes + every bundle of the same tree hash differently -- so the cache would re-upload on every submit + and never hit. Member metadata is normalised for the same reason: the digest should depend on + content and layout only, not on checkout time, local uid, or umask. + """ + tar_buf = io.BytesIO() + with tarfile.open(fileobj=tar_buf, mode="w", format=tarfile.PAX_FORMAT) as tar: + for rel in files: + info = tar.gettarinfo(str(repo / rel), arcname=str(rel)) + _normalise(info) + with open(repo / rel, "rb") as handle: + tar.addfile(info, handle) + + for arcname, content in sorted((extra_files or {}).items()): + payload = content.encode() if isinstance(content, str) else content + info = tarfile.TarInfo(name=arcname) + info.size = len(payload) + info.mode = 0o644 + _normalise(info) + tar.addfile(info, io.BytesIO(payload)) + + out = io.BytesIO() + with gzip.GzipFile(fileobj=out, mode="wb", compresslevel=6, mtime=0) as gz: + gz.write(tar_buf.getvalue()) + return out.getvalue() + + +def _normalise(info: tarfile.TarInfo) -> None: + info.mtime = 0 + info.uid = info.gid = 0 + info.uname = info.gname = "" + info.mode = 0o755 if info.mode & 0o100 else 0o644 + + +def bundle( + repo: Path | str | None = None, + *, + exclude: Iterable[str] = (), + include_suffixes: Iterable[str] = (), + extra_files: Mapping[str, bytes | str] | None = None, + relock: bool = False, + prefix: str | None = None, +) -> Bundle: + """Bundle a working tree, plus any `extra_files` (arcname -> bytes) generated at submit time. + + `extra_files` is how a caller ships something that is not in the tree: a scan manifest, a + resolved config, a design pickle. They participate in the digest, so two submissions of the + same code with different manifests are different bundles. + """ + repo = Path(repo or Path.cwd()).expanduser().resolve() + # A file, not a directory, in a linked worktree -- so `.exists()`, not `.is_dir()`. + if not (repo / ".git").exists(): + raise ValueError(f"{repo} is not a git repository") + + if relock: + subprocess.run(["uv", "lock"], cwd=repo, check=True) + + files, dropped = collect_files(repo, exclude=exclude, include_suffixes=include_suffixes) + lock = repo / "uv.lock" + if not lock.is_file(): + raise FileNotFoundError( + f"{lock} not found -- the runner installs from it with `uv sync --frozen` (try relock=True)" + ) + + blob = build_blob(repo, files, extra_files) + return Bundle( + repo=repo, + prefix=prefix or repo_prefix(repo), + blob=blob, + digest=hashlib.sha256(blob).hexdigest(), + files=tuple(files), + dropped=dropped, + adept_sha=adept_sha_from_lock(lock.read_text()), + ) + + +def upload_bundle(bundle_: Bundle, bucket: str, s3=None) -> bool: + """Put the bundle at its content-addressed key. Returns False if it was already there.""" + from botocore.exceptions import ClientError + + if s3 is None: + import boto3 + + s3 = boto3.client("s3") + try: + s3.head_object(Bucket=bucket, Key=bundle_.key) + return False + except ClientError as err: + # 403 shows up instead of 404 on buckets that do not grant ListBucket, so treat it as + # "not known to be present" and let the put decide. + if err.response["Error"]["Code"] not in ("404", "NoSuchKey", "403"): + raise + s3.put_object(Bucket=bucket, Key=bundle_.key, Body=bundle_.blob) + return True + + +def submit( + cmd: str, + queue: str, + *, + job_definition: str | None = None, + repo: Path | str | None = None, + bucket: str | None = None, + extras: str = "", + array: int = 0, + name: str | None = None, + exclude: Iterable[str] = (), + include_suffixes: Iterable[str] = (), + extra_files: Mapping[str, bytes | str] | None = None, + relock: bool = False, + prefix: str | None = None, + dry_run: bool = False, + quiet: bool = False, + batch=None, + s3=None, +) -> SubmitResult: + """Bundle `repo`, upload it if new, and run `cmd` on the Batch `queue`. + + Args: + cmd: shell command run inside the project, e.g. `python run.py --cfg configs/foo`. A + command rather than an (entry, config) pair because the repos disagree about their + CLIs. `$AWS_BATCH_JOB_ARRAY_INDEX` is expanded in the container, not here. + queue: Batch job queue -- `gpu` (24 GB, spot), `gpu-48g` (48 GB, on-demand), `cpu`, + `cpu-hmem`. + job_definition: defaults to the generic definition for the queue's resource shape. + bucket: code bucket; defaults to $SIM_CODE_BUCKET. + extras: comma-separated uv extras passed to `uv sync`, e.g. "gpu". + array: array size. >1 submits an array job whose members differ only by + $AWS_BATCH_JOB_ARRAY_INDEX; 0 or 1 submits a single job, where that variable is + unset (so write `${AWS_BATCH_JOB_ARRAY_INDEX:-0}` if the command must work both ways). + prefix: S3 key prefix; defaults to the repo name (the main worktree's, from a linked one). + dry_run: bundle and report, upload nothing, submit nothing. + """ + job_definition = job_definition or DEFAULT_JOB_DEFINITIONS.get(queue) or "sim-gpu" + bucket = bucket or os.environ.get("SIM_CODE_BUCKET") + if not bucket and not dry_run: + raise ValueError("bucket is required (pass bucket= or set $SIM_CODE_BUCKET)") + + bundle_ = bundle( + repo, + exclude=exclude, + include_suffixes=include_suffixes, + extra_files=extra_files, + relock=relock, + prefix=prefix, + ) + if not quiet: + bundle_.report() + + code_uri = f"s3://{bucket or ''}/{bundle_.key}" + job_name = name or f"{bundle_.prefix}-{bundle_.digest[:12]}" + array_size = array if array > 1 else 0 + + if dry_run: + if not quiet: + print(f"[cloud] dry run; would upload {code_uri}") + print(f"[cloud] dry run; would submit {job_name} to {queue} / {job_definition}") + print(f"[cloud] dry run; array={array_size or 1} cmd={cmd}") + return SubmitResult( + job_id=None, + job_name=job_name, + queue=queue, + job_definition=job_definition, + cmd=cmd, + code_uri=code_uri, + digest=bundle_.digest, + array_size=array_size, + adept_sha=bundle_.adept_sha, + uploaded=False, + ) + + uploaded = upload_bundle(bundle_, bucket, s3=s3) + if not quiet: + print(f"[cloud] {'uploaded' if uploaded else 'unchanged, reusing'} {code_uri}") + + if batch is None: + import boto3 + + batch = boto3.client("batch") + + request = { + "jobName": job_name, + "jobQueue": queue, + "jobDefinition": job_definition, + "parameters": {"code_uri": code_uri, "cmd": cmd, "extras": extras}, + } + if array_size: + request["arrayProperties"] = {"size": array_size} + + response = batch.submit_job(**request) + if not quiet: + print( + f"[cloud] {response['jobName']} -> {response['jobId']}" + + (f" (array of {array_size})" if array_size else "") + ) + + return SubmitResult( + job_id=response["jobId"], + job_name=response["jobName"], + queue=queue, + job_definition=job_definition, + cmd=cmd, + code_uri=code_uri, + digest=bundle_.digest, + array_size=array_size, + adept_sha=bundle_.adept_sha, + uploaded=uploaded, + ) + + +def array_index(default: int = 0) -> int: + """This job's array index, or `default` when running outside an array job (or locally). + + Workers use this to pick their slice of a scan without caring whether they were launched as + one member of an array, as a single Batch job, or by hand on a laptop. + """ + value = os.environ.get("AWS_BATCH_JOB_ARRAY_INDEX") + return int(value) if value not in (None, "") else default + + +def main(argv: list[str] | None = None) -> int: + parser = argparse.ArgumentParser( + prog="python -m adept.cloud", + description="Bundle a working tree to S3 and submit it as an AWS Batch job.", + formatter_class=argparse.RawDescriptionHelpFormatter, + ) + parser.add_argument("--repo", default=None, help="repo to bundle (default: cwd)") + parser.add_argument("--cmd", required=True, help="command to run inside the project") + parser.add_argument("--queue", required=True, help="Batch job queue (gpu, gpu-48g, cpu, cpu-hmem)") + parser.add_argument("--job-definition", default=None, help="override the queue's default definition") + parser.add_argument("--bucket", default=None, help="code bucket [$SIM_CODE_BUCKET]") + parser.add_argument("--extras", default="", help="comma-separated uv extras, e.g. gpu") + parser.add_argument("--name", default=None, help="job name (default: -)") + parser.add_argument("--prefix", default=None, help="S3 key prefix (default: the repo name)") + parser.add_argument( + "--array", type=int, default=0, help="array size; the command can read $AWS_BATCH_JOB_ARRAY_INDEX" + ) + parser.add_argument("--relock", action="store_true", help="run `uv lock` before bundling") + parser.add_argument( + "--exclude", action="append", default=[], metavar="GLOB", help="path glob to leave out, repeatable" + ) + parser.add_argument( + "--include-suffix", + action="append", + default=[], + metavar="EXT", + help="keep a suffix that is dropped by default, repeatable (e.g. .npy)", + ) + parser.add_argument("--dry-run", action="store_true", help="bundle and report, upload nothing, submit nothing") + args = parser.parse_args(argv) + + try: + submit( + cmd=args.cmd, + queue=args.queue, + job_definition=args.job_definition, + repo=args.repo, + bucket=args.bucket, + extras=args.extras, + array=args.array, + name=args.name, + exclude=args.exclude, + include_suffixes=args.include_suffix, + relock=args.relock, + prefix=args.prefix, + dry_run=args.dry_run, + ) + except (ValueError, FileNotFoundError) as err: + print(f"error: {err}", file=sys.stderr) + return 2 + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/docs/source/usage.md b/docs/source/usage.md index 152c7447..93035ae1 100644 --- a/docs/source/usage.md +++ b/docs/source/usage.md @@ -49,3 +49,13 @@ usage/vlasov1d usage/vlasov1d2v usage/tf1d ``` + +--- + +## Running on AWS Batch + +```{toctree} +:maxdepth: 2 + +usage/cloud +``` diff --git a/docs/source/usage/cloud.md b/docs/source/usage/cloud.md new file mode 100644 index 00000000..0cac168b --- /dev/null +++ b/docs/source/usage/cloud.md @@ -0,0 +1,88 @@ +# Running on AWS Batch + +`adept.cloud` submits a working tree to AWS Batch. The Batch sim images (`sim-cpu`, `sim-gpu`) +contain the environment only — **no application code**. Code arrives at run time as a +content-addressed tarball in S3, so editing a line of physics costs an upload of a few hundred KB +instead of a `docker build`, an ECR push and a `cdk deploy`. + +The runtime half of the contract is `continuum-infra/sim-runner/bootstrap.sh`, which downloads the +bundle, runs `uv sync --frozen` against the lock inside it, and executes the submitted command. + +## Submitting + +```python +from adept.cloud import submit + +result = submit( + cmd="python run.py --cfg configs/vlasov-1d/epw", + queue="gpu", + extras="gpu", +) +print(result.job_id, result.adept_sha) +``` + +or from the command line: + +```bash +python -m adept.cloud --cmd 'python run.py --cfg configs/vlasov-1d/epw' --queue gpu --extras gpu +``` + +`--dry-run` bundles and reports without uploading or submitting — the cheapest way to see what a +bundle would contain. + +`cmd` is a command rather than an `(entry, config)` pair because the repos that submit disagree +about their CLIs. The bucket comes from `bucket=` or `$SIM_CODE_BUCKET`, and `job_definition` +defaults to the generic definition matching the queue's resource shape. + +## Queues + +| Queue | Shape | Pricing | +|---|---|---| +| `gpu` | 1 × 24 GB (L4 / A10G) | spot | +| `gpu-48g` | 1 × 48 GB (L40S) | on-demand | +| `cpu` | 2 vCPU | spot | +| `cpu-hmem` | 8 vCPU / 60 GB | spot | + +## Scans are array jobs + +Each sim is single-GPU, so an N-point scan is **one** array job of size N rather than N +submissions — with Batch's spot retries for free: + +```python +result = submit(cmd="python scan.py --index ${AWS_BATCH_JOB_ARRAY_INDEX:-0}", queue="gpu", array=48) +``` + +`$AWS_BATCH_JOB_ARRAY_INDEX` survives Batch's `Ref::` substitution unexpanded and is expanded by +`bootstrap.sh` inside the container. In the worker, read it with `adept.cloud.array_index()`, which +returns 0 when the process is not an array member, so the same entry point runs locally by hand. + +An `array` of 0 or 1 submits a plain job (Batch has no array of size 1), in which case the variable +is unset — hence the `${...:-0}` default above. + +## What gets bundled + +Bundling uses `git ls-files -co --exclude-standard`: tracked files plus untracked ones that +`.gitignore` does not exclude. That **carries uncommitted edits**, which is what iteration +consists of, while leaving `.venv`, artifacts and `mlflow.db` out for free. + +On top of that: + +- `uv.lock` and `pyproject.toml` ship regardless of git status. Some repos gitignore their lock, + and the runner installs from it. +- Derived formats (`.png`, `.nc`, `.h5`, `.npz`, …) are dropped by suffix, because these repos + commit their outputs. Nothing is dropped silently — the count and size are reported on every + submit, and `include_suffixes=[".npy"]` puts one back when a run needs it as *input*. +- `exclude=["sims/*"]` drops a path glob. +- `extra_files={"scans/abc/manifest.json": ...}` adds content generated at submit time that is not + in the tree at all. It participates in the digest. + +The archive is byte-reproducible: the tar is built uncompressed and gzipped with `mtime=0`, and +member metadata is normalised, so the same tree always yields the same digest and an unchanged +bundle is never re-uploaded. + +## Which adept actually ran + +Repos that depend on adept typically declare `adept @ git+https://github.com/ergodicio/adept@main`, +i.e. unpinned. `submit()` resolves the lock inside the bundle and returns the adept commit the job +will install as `result.adept_sha`. Record it alongside the Batch job id — it is the only record of +the physics that ran. diff --git a/tests/test_cloud/test_bundle.py b/tests/test_cloud/test_bundle.py new file mode 100644 index 00000000..0da48531 --- /dev/null +++ b/tests/test_cloud/test_bundle.py @@ -0,0 +1,253 @@ +"""Bundling and submission contract for adept.cloud. + +These are the properties the runtime depends on and that a rewrite is likely to break: the archive +must be reproducible or content addressing never hits, the lock must ship even when the repo +gitignores it, committed outputs must not ride along, and an array size must reach Batch. +""" + +import subprocess +import tarfile +from pathlib import Path + +import pytest + +from adept import cloud + + +def _repo(tmp_path: Path, *, lock: str = 'name = "x"\n', gitignore: str = "") -> Path: + repo = tmp_path / "demo-repo" + repo.mkdir() + subprocess.run(["git", "init", "-q"], cwd=repo, check=True) + (repo / "run.py").write_text("print('hello')\n") + (repo / "pyproject.toml").write_text('[project]\nname = "demo"\n') + (repo / "uv.lock").write_text(lock) + if gitignore: + (repo / ".gitignore").write_text(gitignore) + return repo + + +def _members(blob: bytes) -> set[str]: + import gzip + import io + + with tarfile.open(fileobj=io.BytesIO(gzip.decompress(blob))) as tar: + return set(tar.getnames()) + + +def test_bundle_is_reproducible(tmp_path): + """Same tree, same digest -- otherwise every submit re-uploads and the cache is pointless.""" + repo = _repo(tmp_path) + first = cloud.bundle(repo) + second = cloud.bundle(repo) + assert first.digest == second.digest + assert first.key == f"demo-repo/{first.digest}.tar.gz" + + +def test_lock_ships_even_when_gitignored(tmp_path): + """kinetic-srs gitignores uv.lock; a .gitignore-respecting bundler would ship no lock at all.""" + repo = _repo(tmp_path, gitignore="uv.lock\n") + assert "uv.lock" in _members(cloud.bundle(repo).blob) + + +def test_missing_lock_is_an_error(tmp_path): + repo = _repo(tmp_path) + (repo / "uv.lock").unlink() + with pytest.raises(FileNotFoundError, match="uv sync --frozen"): + cloud.bundle(repo) + + +def test_committed_outputs_are_dropped_and_reported(tmp_path): + """These repos commit their outputs, so .gitignore alone leaves bundles enormous.""" + repo = _repo(tmp_path) + (repo / "figure.png").write_bytes(b"\x89PNG" + b"0" * 4096) + (repo / "fields.nc").write_bytes(b"0" * 8192) + subprocess.run(["git", "add", "-A"], cwd=repo, check=True) + + bundle = cloud.bundle(repo) + assert "figure.png" not in _members(bundle.blob) + assert "fields.nc" not in _members(bundle.blob) + assert sum(count for count, _ in bundle.dropped.values()) == 2 + + kept = cloud.bundle(repo, include_suffixes=[".nc"]) + assert "fields.nc" in _members(kept.blob) + assert "figure.png" not in _members(kept.blob) + + +def test_exclude_glob(tmp_path): + repo = _repo(tmp_path) + (repo / "sims").mkdir() + (repo / "sims" / "big.txt").write_text("x" * 100) + assert "sims/big.txt" in _members(cloud.bundle(repo).blob) + assert "sims/big.txt" not in _members(cloud.bundle(repo, exclude=["sims/*"]).blob) + + +def test_extra_files_ride_along_and_change_the_digest(tmp_path): + """How a caller ships something generated at submit time -- a scan manifest, a resolved config.""" + repo = _repo(tmp_path) + plain = cloud.bundle(repo) + with_manifest = cloud.bundle(repo, extra_files={"scans/abc/manifest.json": '{"n": 3}'}) + + assert "scans/abc/manifest.json" in _members(with_manifest.blob) + assert with_manifest.digest != plain.digest + # Same manifest, same digest: a resubmission of identical work must not re-upload. + assert with_manifest.digest == cloud.bundle(repo, extra_files={"scans/abc/manifest.json": '{"n": 3}'}).digest + + +def test_untracked_edits_are_carried(tmp_path): + """Iteration consists of uncommitted edits; a git-archive-style bundler would miss them.""" + repo = _repo(tmp_path) + (repo / "scratch.py").write_text("x = 1\n") + assert "scratch.py" in _members(cloud.bundle(repo).blob) + + +def test_venv_is_left_out(tmp_path): + repo = _repo(tmp_path) + (repo / ".venv" / "lib").mkdir(parents=True) + (repo / ".venv" / "lib" / "thing.py").write_text("x = 1\n") + assert not any(name.startswith(".venv") for name in _members(cloud.bundle(repo).blob)) + + +def test_a_worktree_keys_under_the_main_repo(tmp_path): + """Iterating in a worktree is normal; its directory name is a branch nickname, not a project.""" + repo = _repo(tmp_path) + subprocess.run(["git", "add", "-A"], cwd=repo, check=True) + subprocess.run(["git", "-c", "user.email=t@t", "-c", "user.name=t", "commit", "-qm", "init"], cwd=repo, check=True) + tree = tmp_path / "wt-some-branch" + subprocess.run(["git", "worktree", "add", "-q", "-b", "topic", str(tree)], cwd=repo, check=True) + + assert cloud.bundle(tree).prefix == "demo-repo" + assert cloud.bundle(tree).key.startswith("demo-repo/") + assert cloud.bundle(tree, prefix="explicit").prefix == "explicit" + + +def test_not_a_git_repo(tmp_path): + (tmp_path / "plain").mkdir() + with pytest.raises(ValueError, match="not a git repository"): + cloud.bundle(tmp_path / "plain") + + +_ADEPT_SHA = "921d5766875134830ca7db24baf21ec33a0fe446" +_ADEPT_LOCK = f'source = {{ git = "https://github.com/ergodicio/adept.git?rev=main#{_ADEPT_SHA}" }}' + + +@pytest.mark.parametrize( + "lock,expected", + [ + (_ADEPT_LOCK, _ADEPT_SHA), + ('source = { registry = "https://pypi.org/simple" }', None), + ], +) +def test_adept_sha_from_lock(lock, expected): + """Both sim repos pin adept at @main, so the lock is the only record of what actually ran.""" + assert cloud.adept_sha_from_lock(lock) == expected + + +class _FakeS3: + def __init__(self, present=False): + self.present = present + self.puts = [] + + def head_object(self, **kwargs): + if self.present: + return {} + from botocore.exceptions import ClientError + + raise ClientError({"Error": {"Code": "404"}}, "HeadObject") + + def put_object(self, **kwargs): + self.puts.append(kwargs) + + +class _FakeBatch: + def __init__(self): + self.requests = [] + + def submit_job(self, **kwargs): + self.requests.append(kwargs) + return {"jobId": "fake-job-id", "jobName": kwargs["jobName"]} + + +def test_submit_passes_array_size_and_parameters(tmp_path): + repo = _repo(tmp_path, lock=_ADEPT_LOCK) + s3, batch = _FakeS3(), _FakeBatch() + result = cloud.submit( + cmd="python lpi-scan.py --array-index ${AWS_BATCH_JOB_ARRAY_INDEX:-0}", + queue="gpu", + repo=repo, + bucket="code-bucket", + extras="gpu", + array=48, + quiet=True, + s3=s3, + batch=batch, + ) + + (request,) = batch.requests + assert request["jobQueue"] == "gpu" + assert request["jobDefinition"] == "sim-gpu" # queue's default resource shape + assert request["arrayProperties"] == {"size": 48} + assert request["parameters"]["extras"] == "gpu" + assert request["parameters"]["code_uri"] == f"s3://code-bucket/demo-repo/{result.digest}.tar.gz" + assert result.job_id == "fake-job-id" + assert result.array_size == 48 + assert result.adept_sha == "921d5766875134830ca7db24baf21ec33a0fe446" + assert result.uploaded and len(s3.puts) == 1 + + +def test_submit_single_job_has_no_array_properties(tmp_path): + """Batch rejects an array of size 1, so array=1 must degrade to a plain job.""" + batch = _FakeBatch() + cloud.submit( + cmd="python run.py", + queue="gpu-48g", + repo=_repo(tmp_path), + bucket="b", + array=1, + quiet=True, + s3=_FakeS3(), + batch=batch, + ) + assert "arrayProperties" not in batch.requests[0] + + +def test_identical_bundle_is_not_reuploaded(tmp_path): + s3 = _FakeS3(present=True) + result = cloud.submit( + cmd="python run.py", + queue="cpu", + repo=_repo(tmp_path), + bucket="b", + quiet=True, + s3=s3, + batch=_FakeBatch(), + ) + assert not result.uploaded and s3.puts == [] + + +def test_dry_run_touches_nothing(tmp_path): + s3, batch = _FakeS3(), _FakeBatch() + result = cloud.submit( + cmd="python run.py", + queue="gpu", + repo=_repo(tmp_path), + bucket="b", + dry_run=True, + quiet=True, + s3=s3, + batch=batch, + ) + assert result.job_id is None + assert s3.puts == [] and batch.requests == [] + + +def test_bucket_is_required(tmp_path, monkeypatch): + monkeypatch.delenv("SIM_CODE_BUCKET", raising=False) + with pytest.raises(ValueError, match="SIM_CODE_BUCKET"): + cloud.submit(cmd="python run.py", queue="gpu", repo=_repo(tmp_path), quiet=True) + + +def test_array_index_defaults_outside_an_array_job(monkeypatch): + monkeypatch.delenv("AWS_BATCH_JOB_ARRAY_INDEX", raising=False) + assert cloud.array_index() == 0 + monkeypatch.setenv("AWS_BATCH_JOB_ARRAY_INDEX", "7") + assert cloud.array_index() == 7