Add environment artifact: export_and_publish.py

export_and_publish.py

@@ -0,0 +1,232 @@
+"""Export an environment-only image and upload an explicit allowlist to a private HF repo."""
+import argparse
+import gzip
+import hashlib
+import io
+import json
+import os
+import re
+import shutil
+import subprocess
+import tarfile
+from pathlib import Path
+
+os.environ.setdefault("HF_HUB_DISABLE_TELEMETRY", "1")
+
+
+class PartsWriter(io.RawIOBase):
+    def __init__(self, directory, prefix, limit=1024**3):
+        self.directory, self.prefix, self.limit = directory, prefix, limit
+        self.files = []
+        self.current = None
+        self.used = 0
+
+    def writable(self):
+        return True
+
+    def write(self, data):
+        length = len(data)
+        view = memoryview(data)
+        while view:
+            if self.current is None or self.used == self.limit:
+                if self.current:
+                    self.current.close()
+                path = self.directory / f"{self.prefix}.part-{len(self.files):03d}"
+                self.current = path.open("xb")
+                self.files.append(path)
+                self.used = 0
+            take = min(len(view), self.limit - self.used)
+            self.current.write(view[:take])
+            self.used += take
+            view = view[take:]
+        return length
+
+    def close(self):
+        if self.current:
+            self.current.close()
+        super().close()
+
+
+def sha256(path):
+    digest = hashlib.sha256()
+    with path.open("rb") as source:
+        for block in iter(lambda: source.read(1024 * 1024), b""):
+            digest.update(block)
+    return digest.hexdigest()
+
+
+METADATA_FILES = {
+    "Dockerfile", "smoke_test.py", "environment-inventory.json", "requirements.snapshot.txt",
+    "prepare_environment.py", "export_and_publish.py", "model_setup.md", "validation.json", "README.md",
+}
+PART_PATTERN = re.compile(r"codiffact-env_20261006-cu128\.tar\.gz\.part-([0-9]{3})")
+
+
+def validate_release(directory, release):
+    seen, numbers = set(), []
+    for item in release["files"]:
+        name = item["file"]
+        match = PART_PATTERN.fullmatch(name)
+        assert name in METADATA_FILES or match, "Unapproved upload filename"
+        assert name not in seen, "Duplicate artifact"
+        seen.add(name)
+        path = directory / name
+        assert path.is_file() and not path.is_symlink()
+        assert path.resolve().parent == directory.resolve()
+        assert path.stat().st_size == item["bytes"] and sha256(path) == item["sha256"]
+        if match:
+            numbers.append(int(match.group(1)))
+    assert METADATA_FILES <= seen and sorted(numbers) == list(range(len(numbers))) and numbers
+    actual_parts = {p.name for p in directory.glob("*.part-*")}
+    assert actual_parts == {name for name in seen if PART_PATTERN.fullmatch(name)}
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--image", required=True)
+    parser.add_argument("--repo", required=True)
+    parser.add_argument("--token-file", type=Path, required=True)
+    parser.add_argument("--context", type=Path, required=True)
+    parser.add_argument("--out", type=Path, required=True)
+    parser.add_argument("--resume-upload", action="store_true")
+    args = parser.parse_args()
+    tokens = set(re.findall(r"(?<![A-Za-z0-9_])hf_[A-Za-z0-9]+(?![A-Za-z0-9_])", args.token_file.read_text()))
+    if len(tokens) != 1:
+        raise ValueError("Expected one token in the supplied token file")
+    token = next(iter(tokens))
+    token_bytes = token.encode()
+    from huggingface_hub import HfApi
+
+    api = HfApi(token=token)
+    identity = api.whoami()
+    assert args.repo.split("/")[0] == identity["name"], "Only publish to the authenticated user's namespace"
+    root = Path(__file__).resolve().parents[2]
+    manifest_path = args.out / "environment-release.json"
+    if args.resume_upload:
+        release = json.loads(manifest_path.read_text())
+        assert release["repo"] == args.repo and release["image_tag"] == args.image
+        for item in release["files"]:
+            assert sha256(args.out / item["file"]) == item["sha256"]
+    else:
+        assert not args.out.exists(), "Use a fresh output directory or --resume-upload"
+        info = json.loads(subprocess.check_output(["docker", "image", "inspect", args.image], text=True))[0]
+        assert not any("TOKEN=" in value or "PROXY=" in value.upper() for value in info["Config"].get("Env", []))
+        inventory = json.loads((args.context / "environment-inventory.json").read_text())
+        assert not any(any(part in item["file"].split("/") for part in (".git", ".cache", "hfkey.txt")) for item in inventory["files"])
+        args.out.mkdir(parents=True)
+        validations = {}
+        for mode in ("cpu", "gpu"):
+            cmd = ["docker", "run", "--rm", "--network", "none"]
+            if mode == "gpu":
+                cmd.extend(["--gpus", "device=0"])
+            cmd.extend([args.image, "python", "/opt/codiffact/smoke_test.py"])
+            if mode == "gpu":
+                cmd.append("--gpu")
+            result = subprocess.run(cmd, check=True, capture_output=True, text=True)
+            validations[mode] = json.loads(result.stdout.strip().splitlines()[-1])
+            print("OFFLINE_VALIDATION", mode, "passed", flush=True)
+        (args.out / "validation.json").write_text(json.dumps(validations, indent=2))
+        prefix = "codiffact-env_20261006-cu128.tar.gz"
+        writer = PartsWriter(args.out, prefix)
+        process = subprocess.Popen(["docker", "save", args.image], stdout=subprocess.PIPE)
+        image_secret_scan_tail = b""
+        try:
+            with gzip.GzipFile(filename="", mode="wb", fileobj=writer, compresslevel=1, mtime=0) as compressed:
+                for block in iter(lambda: process.stdout.read(4 * 1024 * 1024), b""):
+                    scan = image_secret_scan_tail + block
+                    assert token_bytes not in scan, "Token detected; export aborted"
+                    image_secret_scan_tail = scan[-len(token_bytes):]
+                    compressed.write(block)
+            assert process.wait() == 0, "docker save failed"
+        finally:
+            if process.poll() is None:
+                process.terminate()
+                process.wait()
+            writer.close()
+        files = list(writer.files)
+        for name in ("Dockerfile", "smoke_test.py", "environment-inventory.json", "requirements.snapshot.txt"):
+            destination = args.out / name
+            shutil.copy2(args.context / name, destination)
+            files.append(destination)
+        for name in ("prepare_environment.py", "export_and_publish.py"):
+            destination = args.out / name
+            shutil.copy2(Path(__file__).parent / name, destination)
+            files.append(destination)
+        shutil.copy2(root / "docs/model_setup.md", args.out / "model_setup.md")
+        files.extend([args.out / "model_setup.md", args.out / "validation.json"])
+        body = (
+            "---\ntags:\n- docker\n- environment\n- codiffact\n---\n\n"
+            "# CoDiffAct environment-only image\n\n"
+            "Private Linux/amd64 Docker archive. Contains installed runtime dependencies; no project, datasets, pretrained weights or credentials. "
+            "Each dependency retains its own license; this card does not relicense them.\n\n"
+            "Download all parts, verify SHA256SUMS, then load:\n\n```bash\n"
+            "sha256sum -c SHA256SUMS\n"
+            "cat codiffact-env_20261006-cu128.tar.gz.part-* | gzip -dc | docker load\n"
+            f"docker run --rm --network none --gpus all {args.image} python /opt/codiffact/smoke_test.py --gpu\n"
+            "```\n\n"
+            "See model_setup.md for public model sources, paths and mounts. validation.json records offline dependency, OSMesa, audio/video and CUDA smoke checks, "
+            "not full model inference/training. A compatible host NVIDIA driver and NVIDIA Container Toolkit are required.\n\n"
+            "The exact installed-prefix snapshot is embedded in the archive. requirements.snapshot.txt is an inventory, not a universal pip installation lock. "
+            "Rebuilding with prepare_environment.py requires the recorded source environment; Dockerfile alone is not a complete from-PyPI recipe.\n"
+        )
+        (args.out / "README.md").write_text(body)
+        files.append(args.out / "README.md")
+        for path in files:
+            if ".part-" not in path.name:
+                assert token_bytes not in path.read_bytes(), "Token detected in upload material"
+        release = {
+            "repo": args.repo, "private": True, "image_tag": args.image,
+            "image_id": info["Id"], "uncompressed_image_bytes": info["Size"],
+            "base_image": "nvidia/cuda@sha256:001469ea0f3dec85a1ca929aeea3b58ae369d4c11228b10aec1f642bb6ca7a6f",
+            "files": [{"file": p.name, "bytes": p.stat().st_size, "sha256": sha256(p)} for p in files],
+            "no_user_token_in_export": True,
+        }
+        manifest_path.write_text(json.dumps(release, indent=2))
+        sums = [f"{item['sha256']}  {item['file']}" for item in release["files"]]
+        sums.append(f"{sha256(manifest_path)}  {manifest_path.name}")
+        (args.out / "SHA256SUMS").write_text("\n".join(sums) + "\n")
+        print("EXPORT_COMPLETE", sum(p.stat().st_size for p in writer.files), "bytes", flush=True)
+    validate_release(args.out, release)
+    api.create_repo(repo_id=args.repo, repo_type="model", private=True, exist_ok=True)
+    assert api.repo_info(args.repo).private, "Refusing to upload to a public repository"
+    allowed = [item["file"] for item in release["files"]] + ["SHA256SUMS", manifest_path.name]
+    for name in ("SHA256SUMS", manifest_path.name):
+        assert not (args.out / name).is_symlink()
+        assert token_bytes not in (args.out / name).read_bytes()
+    parts = sorted(name for name in allowed if ".part-" in name)
+    metadata = sorted(set(allowed) - set(parts))
+    remote = {s.rfilename: s for s in api.repo_info(args.repo, files_metadata=True).siblings}
+    assert not ({name for name in remote if ".part-" in name} - set(parts)), "Old remote parts present; use a new version repository"
+    for name in parts:
+        previous = getattr(remote.get(name), "lfs", None)
+        assert not previous or previous.sha256 == sha256(args.out / name), "Never overwrite another image release; use a new version repository"
+    for name in parts + metadata:
+        path = args.out / name
+        expected = sha256(path)
+        existing = remote.get(name)
+        lfs = getattr(existing, "lfs", None)
+        if lfs and lfs.sha256 == expected:
+            print("ALREADY_UPLOADED", name, flush=True)
+            continue
+        print("UPLOADING", name, path.stat().st_size, flush=True)
+        api.upload_file(path_or_fileobj=str(path), path_in_repo=name, repo_id=args.repo,
+                        repo_type="model", commit_message=f"Add environment artifact: {name}")
+    final = api.repo_info(args.repo, files_metadata=True)
+    remote = {s.rfilename: s for s in final.siblings}
+    for name in allowed:
+        assert name in remote, name
+        assert remote[name].size == (args.out / name).stat().st_size, name
+        lfs = getattr(remote[name], "lfs", None)
+        if lfs:
+            assert lfs.sha256 == sha256(args.out / name), name
+    receipt = {"repo": args.repo, "private": final.private, "commit": final.sha, "files": allowed, "complete": True}
+    (args.out.parent / "upload_receipt.json").write_text(json.dumps(receipt, indent=2))
+    print("HF_UPLOAD_VERIFIED", args.repo, final.sha, flush=True)
+
+
+if __name__ == "__main__":
+    try:
+        main()
+    except Exception as error:
+        print("FAILED", type(error).__name__, "See local build/validation state; secret values suppressed.", flush=True)
+        raise SystemExit(1)