diff --git a/rootstock/cli.py b/rootstock/cli.py index 029f539..ba7e500 100644 --- a/rootstock/cli.py +++ b/rootstock/cli.py @@ -27,16 +27,19 @@ rootstock add --list [--root ] List every canonical checkpoint id that add accepts, grouped by env. - rootstock sync [] [--root | --cluster ] [--dry-run] + rootstock sync [] [--root | --cluster ] [--dry-run] Converge the install to its declared state: build missing/changed envs, download and verify missing/stale checkpoints, in parallel - phases. Idempotent — re-run to retry whatever failed. + phases. Idempotent — re-run to retry whatever failed. is a + local staging directory of env files, or a git spec + git+URL[@REF][#subdirectory=DIR] fetched to a temp checkout. rootstock sync --cluster delta --dry-run rootstock sync ./environments/ --jobs 8 + rootstock sync 'git+https://github.com/org/envs.git@main#subdirectory=delta' rootstock sync --rebuild # after a CLI version bump rootstock sync --phases build,download # login node (no GPU) - rootstock prune [] [--root | --cluster ] [--dry-run] [--yes] + rootstock prune [] [--root | --cluster ] [--dry-run] [--yes] The subtractive half of sync: remove built envs with no registered source, checkpoint records (and their unshared weight files) no source declares, and internal garbage (.build/.trash leftovers, @@ -318,9 +321,12 @@ def main(): sync_parser.add_argument( "source_dir", nargs="?", + metavar="SOURCE", help=( - "Optional directory of env source files (*.py) to register/update " - "before converging; defaults to the root's registered environments" + "Optional source of env definitions (*.py) to register/update " + "before converging: a local staging directory, or a git spec " + "'git+URL[@REF][#subdirectory=DIR]' shallow-fetched to a temp " + "checkout. Defaults to the root's registered environments" ), ) sync_parser.add_argument( @@ -434,13 +440,16 @@ def main(): prune_parser.add_argument( "source_dir", nargs="?", + metavar="SOURCE", help=( - "Optional directory declaring the *complete* desired set of env " - "sources (*.py): anything registered, built, or fetched beyond it " - "is pruned — including registered source files. An empty directory " - "declares zero environments. Defaults to the root's registered " - "environments (if that dir doesn't exist, nothing is declared and " - "only internal garbage is collected)." + "Optional source declaring the *complete* desired set of env " + "sources (*.py) — a local directory or a git spec " + "'git+URL[@REF][#subdirectory=DIR]': anything registered, built, " + "or fetched beyond it is pruned — including registered source " + "files. An empty directory declares zero environments. Defaults " + "to the root's registered environments (if that dir doesn't " + "exist, nothing is declared and only internal garbage is " + "collected)." ), ) prune_parser.add_argument( diff --git a/rootstock/commands/common.py b/rootstock/commands/common.py index 68c6d1a..6fa1456 100644 --- a/rootstock/commands/common.py +++ b/rootstock/commands/common.py @@ -39,6 +39,26 @@ def warn_on_permissions(root: Path, cache_root: Path) -> None: ) +def resolve_source_arg(spec: str) -> Path: + """The sync/prune source positional: a local directory of env sources, or + a ``git+URL[@REF][#subdirectory=DIR]`` spec (see :mod:`rootstock.gitsource`) + shallow-fetched to a temp checkout that lives until process exit. + + Raises OperationError for anything unusable (missing dir, malformed spec, + failed fetch) — callers treat that as a usage error. + """ + from ..gitsource import is_git_source, materialize_git_source + + if is_git_source(spec): + return materialize_git_source(spec) + source_dir = Path(spec) + if not source_dir.is_dir(): + from ..operations import OperationError + + raise OperationError(f"{source_dir} is not a directory") + return source_dir + + def resolve_root(args) -> Path: """``--root`` (or env/config fallback), with ``--cluster`` as a registry bootstrap for admins driving a known cluster by name.""" diff --git a/rootstock/commands/prune.py b/rootstock/commands/prune.py index 8640b0d..250b71e 100644 --- a/rootstock/commands/prune.py +++ b/rootstock/commands/prune.py @@ -26,7 +26,7 @@ from ..layout import ensure_layout_compatible, write_layout_marker from ..manifest import ManifestError from ..operations import OperationError -from .common import resolve_cache_root, resolve_root, warn_on_permissions +from .common import resolve_cache_root, resolve_root, resolve_source_arg, warn_on_permissions def cmd_prune(args) -> int: @@ -44,9 +44,10 @@ def cmd_prune(args) -> int: source_dir: Path | None = None if args.source_dir: - source_dir = Path(args.source_dir) - if not source_dir.is_dir(): - print(f"Error: {source_dir} is not a directory", file=sys.stderr) + try: + source_dir = resolve_source_arg(args.source_dir) + except OperationError as exc: + print(f"Error: {exc}", file=sys.stderr) return 2 if args.min_age < 0: diff --git a/rootstock/commands/sync.py b/rootstock/commands/sync.py index 38a8e05..0f6b65d 100644 --- a/rootstock/commands/sync.py +++ b/rootstock/commands/sync.py @@ -16,7 +16,7 @@ from ..environment import CheckpointNotFoundError from ..layout import ensure_layout_compatible, write_layout_marker from ..operations import OperationError, resolve_current_cluster -from .common import resolve_cache_root, resolve_root, warn_on_permissions +from .common import resolve_cache_root, resolve_root, resolve_source_arg, warn_on_permissions def _parse_phases(spec: str) -> tuple[str, ...]: @@ -55,9 +55,10 @@ def cmd_sync(args) -> int: source_dir: Path | None = None if args.source_dir: - source_dir = Path(args.source_dir) - if not source_dir.is_dir(): - print(f"Error: {source_dir} is not a directory", file=sys.stderr) + try: + source_dir = resolve_source_arg(args.source_dir) + except OperationError as exc: + print(f"Error: {exc}", file=sys.stderr) return 2 root = resolve_root(args) diff --git a/rootstock/gitsource.py b/rootstock/gitsource.py new file mode 100644 index 0000000..cd8e6d4 --- /dev/null +++ b/rootstock/gitsource.py @@ -0,0 +1,140 @@ +"""Fetch env-source directories from git repositories. + +``sync`` and ``prune`` take an optional positional source of env definitions. +Besides a local staging directory, that positional accepts a pip-style git +spec:: + + git+URL[@REF][#subdirectory=DIR] + + git+https://github.com/Garden-AI/rootstock.git#subdirectory=sample_model_configurations/nvidia_configs + git+https://github.com/Garden-AI/rootstock.git@v1.3.0#subdirectory=environments + git+ssh://git@github.com/Garden-AI/rootstock.git@main + +The spec is shallow-fetched (one commit, no history) into a temp directory +that lives until process exit — long enough for the planner to hash the +staged files and the executor to register/build from them, after which the +install root holds its own copies and the checkout is disposable. + +REF may be a branch, tag, or full commit SHA (whatever the server allows +``fetch`` by name — GitHub allows all three). Omitted REF means the remote's +default branch. Omitted subdirectory means the repository root. +""" + +from __future__ import annotations + +import atexit +import shutil +import subprocess +import tempfile +from dataclasses import dataclass +from pathlib import Path +from urllib.parse import parse_qs, urlsplit + +from .operations import OperationError + +GIT_SPEC_PREFIX = "git+" + + +def is_git_source(spec: str) -> bool: + """Whether a source positional names a git spec rather than a local dir.""" + return spec.startswith(GIT_SPEC_PREFIX) + + +@dataclass(frozen=True) +class GitSource: + url: str # clone URL, scheme intact, ref/fragment stripped + ref: str | None # branch, tag, or full SHA; None = remote default branch + subdirectory: str | None # repo-relative dir holding the *.py sources + + +def parse_git_source(spec: str) -> GitSource: + """Parse ``git+URL[@REF][#subdirectory=DIR]`` into its parts. + + The ``@REF`` separator is an ``@`` in the URL *path* (so the user info in + ``ssh://git@github.com/...`` never trips it); a ref containing ``/`` + (feature branches) parses fine since we split on the last ``@``. + """ + if not is_git_source(spec): + raise OperationError(f"Not a git source spec (expected git+URL): {spec}") + + parts = urlsplit(spec[len(GIT_SPEC_PREFIX) :]) + if not parts.scheme or not (parts.netloc or parts.path): + raise OperationError( + f"Malformed git source {spec!r}: expected git+URL[@REF][#subdirectory=DIR] " + "(scp-style addresses need the ssh:// form, e.g. git+ssh://git@host/org/repo.git)" + ) + + path, ref = parts.path, None + if "@" in path: + path, _, ref = path.rpartition("@") + if not ref: + raise OperationError(f"Malformed git source {spec!r}: empty ref after '@'") + + subdirectory = None + if parts.fragment: + fragment = parse_qs(parts.fragment) + unknown = sorted(set(fragment) - {"subdirectory"}) + if unknown: + raise OperationError( + f"Unknown fragment option(s) in git source: {', '.join(unknown)} " + "(only 'subdirectory' is supported)" + ) + subdirectory = fragment.get("subdirectory", [None])[0] + + url = parts._replace(path=path, fragment="").geturl() + return GitSource(url=url, ref=ref, subdirectory=subdirectory) + + +def _git(args: list[str], cwd: Path) -> None: + proc = subprocess.run( + ["git", *args], + cwd=cwd, + capture_output=True, + text=True, + ) + if proc.returncode != 0: + detail = proc.stderr.strip() or proc.stdout.strip() or "(no output)" + raise OperationError(f"git {args[0]} failed: {detail}") + + +def fetch_git_source(source: GitSource, dest: Path) -> Path: + """Shallow-fetch ``source`` into ``dest``; return the env-source dir. + + One code path covers branch, tag, and SHA refs: init + ``fetch --depth 1 + origin `` + detached checkout of FETCH_HEAD (a plain shallow clone + can't target a SHA). + """ + if shutil.which("git") is None: + raise OperationError("git not found in PATH (required for git+ sources)") + + dest.mkdir(parents=True, exist_ok=True) + _git(["init", "--quiet", "."], cwd=dest) + _git(["remote", "add", "origin", source.url], cwd=dest) + _git(["fetch", "--quiet", "--depth", "1", "origin", source.ref or "HEAD"], cwd=dest) + _git(["checkout", "--quiet", "--detach", "FETCH_HEAD"], cwd=dest) + + if source.subdirectory is None: + return dest + + subdir = (dest / source.subdirectory).resolve() + if dest.resolve() not in subdir.parents and subdir != dest.resolve(): + raise OperationError(f"subdirectory escapes the repository: {source.subdirectory}") + if not subdir.is_dir(): + raise OperationError( + f"subdirectory {source.subdirectory!r} not found in {source.url}" + f"{f' @ {source.ref}' if source.ref else ''}" + ) + return subdir + + +def materialize_git_source(spec: str) -> Path: + """Resolve a ``git+`` spec to a local directory of env sources. + + The checkout lands in a temp directory removed at process exit — callers + hold the returned path for at most the life of one CLI command, and + everything durable is copied into the install root during the run. + """ + source = parse_git_source(spec) + checkout = Path(tempfile.mkdtemp(prefix="rootstock-git-src-")) + atexit.register(shutil.rmtree, checkout, ignore_errors=True) + return fetch_git_source(source, checkout) diff --git a/tests/cli/test_sync.py b/tests/cli/test_sync.py index 589eb83..2372e94 100644 --- a/tests/cli/test_sync.py +++ b/tests/cli/test_sync.py @@ -140,6 +140,42 @@ def test_missing_source_dir_is_a_usage_error(tmp_path, stubbed): assert stubbed["plan_calls"] == [] +def test_git_source_spec_is_materialized_before_planning(tmp_path, stubbed, monkeypatch): + """A git+ positional resolves to a checkout dir, which is what the + planner sees as source_dir (fetch mechanics have their own tests in + tests/test_gitsource.py).""" + checkout = tmp_path / "checkout" / "environments" + checkout.mkdir(parents=True) + spec = "git+https://example.com/envs.git@main#subdirectory=environments" + + def fake_resolve(arg): + assert arg == spec + return checkout + + monkeypatch.setattr("rootstock.commands.sync.resolve_source_arg", fake_resolve) + + rc = cmd_sync(_make_args(tmp_path, source_dir=spec, dry_run=True)) + + assert rc == 0 + ((_, kwargs),) = stubbed["plan_calls"] + assert kwargs["source_dir"] == checkout + + +def test_failed_git_fetch_is_a_usage_error(tmp_path, stubbed, monkeypatch, capsys): + from rootstock.operations import OperationError + + def exploding_resolve(arg): + raise OperationError("git fetch failed: repository not found") + + monkeypatch.setattr("rootstock.commands.sync.resolve_source_arg", exploding_resolve) + + rc = cmd_sync(_make_args(tmp_path, source_dir="git+https://example.com/nope.git")) + + assert rc == 2 + assert stubbed["plan_calls"] == [] + assert "git fetch failed" in capsys.readouterr().err + + def test_json_dry_run_emits_the_plan_on_stdout(tmp_path, stubbed, capsys): stubbed["plan"] = ONE_BUILD diff --git a/tests/test_gitsource.py b/tests/test_gitsource.py new file mode 100644 index 0000000..5d4c5f5 --- /dev/null +++ b/tests/test_gitsource.py @@ -0,0 +1,190 @@ +"""Tests for git+ env-source specs (rootstock.gitsource). + +Parsing is pure; fetching is exercised against a throwaway local repository +over the file:// transport (which, unlike a plain path, supports the same +shallow-fetch path used against real remotes). +""" + +from __future__ import annotations + +import shutil +import subprocess +from pathlib import Path + +import pytest + +from rootstock.commands.common import resolve_source_arg +from rootstock.gitsource import ( + GitSource, + fetch_git_source, + is_git_source, + parse_git_source, +) +from rootstock.operations import OperationError + +# --------------------------------------------------------------------------- +# parsing + + +def test_is_git_source_only_matches_the_prefix(): + assert is_git_source("git+https://github.com/org/repo.git") + assert not is_git_source("./staging/delta-environments") + assert not is_git_source("/abs/path") + + +def test_parse_url_only(): + parsed = parse_git_source("git+https://github.com/org/repo.git") + assert parsed == GitSource("https://github.com/org/repo.git", None, None) + + +def test_parse_ref_and_subdirectory(): + parsed = parse_git_source( + "git+https://github.com/org/repo.git@v1.3.0#subdirectory=configs/delta" + ) + assert parsed == GitSource("https://github.com/org/repo.git", "v1.3.0", "configs/delta") + + +def test_parse_branch_ref_with_slash(): + parsed = parse_git_source("git+https://github.com/org/repo.git@feature/new-envs") + assert parsed.ref == "feature/new-envs" + assert parsed.url == "https://github.com/org/repo.git" + + +def test_parse_ssh_userinfo_is_not_a_ref_separator(): + parsed = parse_git_source("git+ssh://git@github.com/org/repo.git") + assert parsed == GitSource("ssh://git@github.com/org/repo.git", None, None) + + parsed = parse_git_source("git+ssh://git@github.com/org/repo.git@main") + assert parsed == GitSource("ssh://git@github.com/org/repo.git", "main", None) + + +@pytest.mark.parametrize( + "spec", + [ + "git+https://github.com/org/repo.git@", # empty ref + "git+github.com/org/repo", # no scheme + "git+https://github.com/org/repo.git#egg=whatever", # unknown fragment + ], +) +def test_parse_rejects_malformed_specs(spec): + with pytest.raises(OperationError): + parse_git_source(spec) + + +def test_parse_rejects_non_git_specs(): + with pytest.raises(OperationError, match="git\\+URL"): + parse_git_source("./local/dir") + + +# --------------------------------------------------------------------------- +# fetching (local file:// repository) + +requires_git = pytest.mark.skipif(shutil.which("git") is None, reason="git not on PATH") + + +def _run(args: list[str], cwd: Path) -> str: + return subprocess.run( + ["git", "-c", "user.email=t@t", "-c", "user.name=t", *args], + cwd=cwd, + check=True, + capture_output=True, + text=True, + ).stdout.strip() + + +@pytest.fixture +def upstream(tmp_path: Path) -> dict: + """A local repo with an environments/ subdir, two commits, and a tag.""" + repo = tmp_path / "upstream" + (repo / "environments").mkdir(parents=True) + _run(["init", "--quiet", "-b", "main"], cwd=repo) + # Serve arbitrary SHAs over the file:// transport, like GitHub does. + _run(["config", "uploadpack.allowAnySHA1InWant", "true"], cwd=repo) + + (repo / "environments" / "mace.py").write_text("CHECKPOINTS = {}\n") + (repo / "README.md").write_text("not an env source\n") + _run(["add", "-A"], cwd=repo) + _run(["commit", "--quiet", "-m", "first"], cwd=repo) + _run(["tag", "v1"], cwd=repo) + first_sha = _run(["rev-parse", "HEAD"], cwd=repo) + + (repo / "environments" / "uma.py").write_text("CHECKPOINTS = {}\n") + _run(["add", "-A"], cwd=repo) + _run(["commit", "--quiet", "-m", "second"], cwd=repo) + + return {"url": f"file://{repo}", "first_sha": first_sha} + + +@requires_git +def test_fetch_default_branch_repo_root(upstream, tmp_path): + dest = tmp_path / "checkout" + result = fetch_git_source(GitSource(upstream["url"], None, None), dest) + + assert result == dest + assert (result / "environments" / "uma.py").exists() + + +@requires_git +def test_fetch_subdirectory(upstream, tmp_path): + source = GitSource(upstream["url"], None, "environments") + result = fetch_git_source(source, tmp_path / "checkout") + + assert result.name == "environments" + assert sorted(p.name for p in result.glob("*.py")) == ["mace.py", "uma.py"] + + +@requires_git +def test_fetch_tag_ref(upstream, tmp_path): + source = GitSource(upstream["url"], "v1", "environments") + result = fetch_git_source(source, tmp_path / "checkout") + + assert (result / "mace.py").exists() + assert not (result / "uma.py").exists(), "tag predates the second commit" + + +@requires_git +def test_fetch_sha_ref(upstream, tmp_path): + source = GitSource(upstream["url"], upstream["first_sha"], None) + result = fetch_git_source(source, tmp_path / "checkout") + + assert not (result / "environments" / "uma.py").exists() + + +@requires_git +def test_fetch_missing_subdirectory_errors(upstream, tmp_path): + source = GitSource(upstream["url"], None, "no-such-dir") + with pytest.raises(OperationError, match="no-such-dir"): + fetch_git_source(source, tmp_path / "checkout") + + +@requires_git +def test_fetch_escaping_subdirectory_errors(upstream, tmp_path): + source = GitSource(upstream["url"], None, "../outside") + with pytest.raises(OperationError, match="escapes"): + fetch_git_source(source, tmp_path / "checkout") + + +@requires_git +def test_fetch_bad_ref_surfaces_git_stderr(upstream, tmp_path): + source = GitSource(upstream["url"], "no-such-ref", None) + with pytest.raises(OperationError, match="git fetch failed"): + fetch_git_source(source, tmp_path / "checkout") + + +# --------------------------------------------------------------------------- +# the shared CLI resolver + + +def test_resolve_source_arg_local_dir(tmp_path): + assert resolve_source_arg(str(tmp_path)) == tmp_path + + +def test_resolve_source_arg_missing_dir_raises(tmp_path): + with pytest.raises(OperationError, match="not a directory"): + resolve_source_arg(str(tmp_path / "nope")) + + +@requires_git +def test_resolve_source_arg_git_spec(upstream): + result = resolve_source_arg(f"git+{upstream['url']}#subdirectory=environments") + assert (result / "mace.py").exists()