Add repo crawler and end-to-end ingest pipeline

.gitignore

@@ -7,3 +7,6 @@ __pycache__/
 dist/
 build/
 .DS_Store
+
+# local crawler output
+data/

ingest/README.md

@@ -2,6 +2,22 @@
 
 Clone CDS repos, filter, secret-scan, and chunk them into chunks.jsonl (steps 3-4).
 
+```bash
+uv run python -m ingest.pipeline --root data/repos --out data/chunks.jsonl
+```
+
+`pipeline.py` runs the crawler, then chunks every synced repo into one file. Add
+`--only FTM-Snake slurm-aware` or `--limit 5` to crawl a few repos; partial runs never delete
+clones. Set `GITHUB_TOKEN` for higher API rate limits.
+
+`crawler.py` lists the org's public, non-fork, non-archived, non-empty repos (minus `exclude`
+in `repos.toml`), keeps a shallow clone of each default branch under `--root`, and records
+every repo's commit and status (`ok`, `stale`, `skipped`, `failed`, `removed`) in
+`<root>/manifest.json`. A clone with local changes is left alone, a failed update keeps the
+last good copy, and a full run deletes clones of repos that are no longer listed.
+
+To chunk one checkout by hand:
+
 ```bash
 uv run python -m ingest.chunker path/to/repo --repo CornellDataScience/name --out chunks.jsonl
 ```

ingest/crawler.py

@@ -0,0 +1,131 @@
+"""Keep shallow clones of the org's repos under a root folder, recorded in root/manifest.json."""
+
+import json
+import os
+import shutil
+import subprocess
+import tomllib
+from concurrent.futures import ThreadPoolExecutor
+from dataclasses import asdict, dataclass
+from datetime import datetime, timezone
+from pathlib import Path
+
+CONFIG_PATH = Path(__file__).with_name("repos.toml")
+GIT_TIMEOUT_SECONDS = 600
+
+
+@dataclass
+class RepoInfo:
+    full_name: str
+    clone_url: str
+    default_branch: str
+    fork: bool = False
+    archived: bool = False
+    empty: bool = False
+
+
+@dataclass
+class ManifestEntry:
+    full_name: str
+    path: str  # relative to the root
+    default_branch: str
+    commit_sha: str | None
+    status: str  # ok | stale | skipped | failed | removed
+    synced_at: str
+    error: str | None = None
+
+
+def load_config(path: Path = CONFIG_PATH) -> dict:
+    with path.open("rb") as f:
+        return tomllib.load(f)
+
+
+def list_org_repos(org: str, token: str | None = None) -> list[RepoInfo]:
+    from github import Auth, Github
+
+    gh = Github(auth=Auth.Token(token)) if token else Github()
+    return [
+        RepoInfo(r.full_name, r.clone_url, r.default_branch, r.fork, r.archived, r.size == 0)
+        for r in gh.get_organization(org).get_repos(type="public")
+    ]
+
+
+def select_repos(repos: list[RepoInfo], exclude=(), only=(), limit: int | None = None) -> list[RepoInfo]:
+    names = lambda r: {r.full_name, r.full_name.split("/")[-1]}  # noqa: E731
+    selected = [
+        r for r in sorted(repos, key=lambda r: r.full_name.lower())
+        if not (r.fork or r.archived or r.empty) and not names(r) & set(exclude) and (not only or names(r) & set(only))
+    ]
+    return selected[:limit] if limit else selected
+
+
+def crawl(root: Path, repos: list[RepoInfo], jobs: int = 8, prune: bool = True) -> dict[str, ManifestEntry]:
+    """Clone or update each repo; with prune, delete clones of repos no longer listed."""
+    root.mkdir(parents=True, exist_ok=True)
+    previous = read_manifest(root)
+    with ThreadPoolExecutor(max_workers=jobs) as pool:
+        entries = {e.full_name: e for e in pool.map(lambda r: sync_repo(root, r), repos)}
+    for name, old in previous.items():
+        if name in entries:
+            continue
+        if prune and old.status != "removed":
+            shutil.rmtree(root / old.path, ignore_errors=True)
+            entries[name] = ManifestEntry(name, old.path, old.default_branch, None, "removed", _now())
+        elif not prune:
+            entries[name] = old
+    write_manifest(root, entries)
+    return entries
+
+
+def sync_repo(root: Path, repo: RepoInfo) -> ManifestEntry:
+    rel = repo.full_name
+    dest = root / rel
+    entry = lambda status, sha=None, error=None: ManifestEntry(  # noqa: E731
+        repo.full_name, rel, repo.default_branch, sha, status, _now(), error
+    )
+    if not (dest / ".git").is_dir():
+        shutil.rmtree(dest, ignore_errors=True)  # leftovers of an interrupted clone
+        try:
+            _git(None, "clone", "--quiet", "--depth", "1", "--single-branch", "--branch", repo.default_branch, repo.clone_url, str(dest))
+        except RuntimeError as e:
+            shutil.rmtree(dest, ignore_errors=True)
+            return entry("failed", error=str(e))
+        return entry("ok", _git(dest, "rev-parse", "HEAD"))
+
+    if _git(dest, "status", "--porcelain"):
+        return entry("skipped", _git(dest, "rev-parse", "HEAD"), "local changes in the clone")
+    try:
+        _git(dest, "fetch", "--quiet", "--depth", "1", "origin", repo.default_branch)
+        _git(dest, "reset", "--quiet", "--hard", "FETCH_HEAD")
+    except RuntimeError as e:
+        return entry("stale", _git(dest, "rev-parse", "HEAD"), str(e))  # keep indexing the last good copy
+    return entry("ok", _git(dest, "rev-parse", "HEAD"))
+
+
+def read_manifest(root: Path) -> dict[str, ManifestEntry]:
+    path = root / "manifest.json"
+    if not path.is_file():
+        return {}
+    return {name: ManifestEntry(**e) for name, e in json.loads(path.read_text())["repos"].items()}
+
+
+def write_manifest(root: Path, entries: dict[str, ManifestEntry]):
+    data = {"synced_at": _now(), "repos": {name: asdict(entries[name]) for name in sorted(entries)}}
+    tmp = root / "manifest.json.tmp"
+    tmp.write_text(json.dumps(data, indent=2) + "\n")
+    tmp.replace(root / "manifest.json")
+
+
+def _git(cwd: Path | None, *args: str) -> str:
+    env = {**os.environ, "GIT_TERMINAL_PROMPT": "0"}  # fail instead of prompting for credentials
+    try:
+        result = subprocess.run(["git", *args], cwd=cwd, env=env, capture_output=True, text=True, timeout=GIT_TIMEOUT_SECONDS)
+    except subprocess.TimeoutExpired:
+        raise RuntimeError(f"git {args[0]} timed out") from None
+    if result.returncode != 0:
+        raise RuntimeError(f"git {args[0]} failed: {result.stderr.strip().splitlines()[-1] if result.stderr.strip() else result.returncode}")
+    return result.stdout.strip()
+
+
+def _now() -> str:
+    return datetime.now(timezone.utc).isoformat(timespec="seconds")

ingest/pipeline.py

@@ -0,0 +1,73 @@
+"""Crawl the org, then chunk every synced repo into one chunks.jsonl.
+
+    uv run python -m ingest.pipeline --root data/repos --out data/chunks.jsonl --only FTM-Snake slurm-aware
+"""
+
+import argparse
+import json
+import os
+from collections import Counter
+from pathlib import Path
+
+from ingest import crawler
+from ingest.chunker import chunk_repo, llama_token_counter
+
+CHUNKABLE = {"ok", "stale"}
+
+
+def run(root: Path, out: Path, repos: list[crawler.RepoInfo], count_tokens, jobs=8, prune=True, scan_secrets=True) -> dict:
+    entries = crawler.crawl(root, repos, jobs=jobs, prune=prune)
+    listed = {r.full_name for r in repos}
+    skipped, redacted, chunks = Counter(), [], Counter()
+    out.parent.mkdir(parents=True, exist_ok=True)
+    tmp = out.with_name(out.name + ".tmp")
+    with tmp.open("w") as f:
+        for name in sorted(listed):
+            entry = entries[name]
+            if entry.status not in CHUNKABLE:
+                continue
+            for record in chunk_repo(root / entry.path, name, count_tokens, skipped, redacted, scan_secrets):
+                f.write(json.dumps(record) + "\n")
+                chunks[name] += 1
+    tmp.replace(out)  # a failed run never leaves a half-written file
+    return {"entries": entries, "chunks": chunks, "skipped": skipped, "redacted": redacted}
+
+
+def main():
+    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    parser.add_argument("--root", type=Path, default=Path("data/repos"), help="Where clones and manifest.json live")
+    parser.add_argument("--out", type=Path, default=Path("data/chunks.jsonl"))
+    parser.add_argument("--config", type=Path, default=crawler.CONFIG_PATH)
+    parser.add_argument("--only", nargs="+", default=(), help="Crawl just these repos (no pruning)")
+    parser.add_argument("--limit", type=int, help="Crawl the first N repos (no pruning)")
+    parser.add_argument("--jobs", type=int, default=8, help="Parallel clones")
+    parser.add_argument("--no-secret-scan", action="store_true", help="Skip gitleaks (never for real ingestion)")
+    args = parser.parse_args()
+
+    config = crawler.load_config(args.config)
+    listing = crawler.list_org_repos(config["org"], os.environ.get("GITHUB_TOKEN"))
+    repos = crawler.select_repos(listing, config.get("exclude", []), args.only, args.limit)
+    if args.only and len(repos) < len(set(args.only)):
+        found = {r.full_name.split("/")[-1] for r in repos}
+        print(f"Not found or not indexable: {sorted(set(args.only) - found)}")
+    partial = bool(args.only or args.limit)
+
+    result = run(args.root, args.out, repos, llama_token_counter(), args.jobs, prune=not partial, scan_secrets=not args.no_secret_scan)
+
+    statuses = Counter(e.status for e in result["entries"].values())
+    print(f"Repos: {dict(statuses)}")
+    for e in sorted(result["entries"].values(), key=lambda e: e.full_name):
+        if e.status in {"failed", "stale", "skipped"}:
+            print(f"  {e.status}: {e.full_name}: {e.error}")
+        elif e.status == "removed":
+            print(f"  removed: {e.full_name}")
+    print(f"Wrote {sum(result['chunks'].values())} chunks from {len(result['chunks'])} repos to {args.out}")
+    print(f"Skipped files: {dict(result['skipped']) or 'none'}")
+    if args.no_secret_scan:
+        print("WARNING: secret scan disabled")
+    for finding in result["redacted"]:
+        print(f"Redacted {finding.rule} at {finding.path}:{finding.line}")
+
+
+if __name__ == "__main__":
+    main()

ingest/repos.toml

@@ -0,0 +1,3 @@
+# Indexed: public, non-fork, non-archived, non-empty repos in `org`, minus `exclude` (names only).
+org = "CornellDataScience"
+exclude = []

pyproject.toml

@@ -9,6 +9,7 @@ dependencies = [
     "huggingface-hub>=0.30",
     "nbformat>=5.11.1",
     "pathspec>=1.1.1",
+    "pygithub>=2.10.0",
     "torch==2.14.1",
     "transformers>=4.50",
     "tree-sitter>=0.26.0",

tests/test_crawler.py

@@ -0,0 +1,143 @@
+import json
+import subprocess
+
+import pytest
+
+from ingest import crawler, pipeline
+from ingest.crawler import RepoInfo, crawl, read_manifest, select_repos
+
+
+def git(path, *args):
+    return subprocess.run(
+        ["git", "-C", str(path), "-c", "user.name=t", "-c", "user.email=t@t", *args],
+        check=True, capture_output=True, text=True,
+    ).stdout.strip()
+
+
+def make_remote(tmp_path, name, files, branch="main"):
+    path = tmp_path / "remotes" / name
+    path.mkdir(parents=True)
+    git(path, "init", "-q", "-b", branch)
+    commit(path, files)
+    return RepoInfo(f"org/{name}", path.as_uri(), branch)
+
+
+def commit(path, files):
+    for rel, text in files.items():
+        (path / rel).parent.mkdir(parents=True, exist_ok=True)
+        (path / rel).write_text(text)
+    git(path, "add", ".")
+    git(path, "commit", "-qm", "change")
+    return git(path, "rev-parse", "HEAD")
+
+
+def head(remote: RepoInfo):
+    return git(remote.clone_url.removeprefix("file://"), "rev-parse", "HEAD")
+
+
+@pytest.fixture
+def root(tmp_path):
+    return tmp_path / "root"
+
+
+def test_select_repos():
+    repos = [
+        RepoInfo("org/b", "", "main"),
+        RepoInfo("org/a", "", "main"),
+        RepoInfo("org/fork", "", "main", fork=True),
+        RepoInfo("org/old", "", "main", archived=True),
+        RepoInfo("org/empty", "", "main", empty=True),
+        RepoInfo("org/skip", "", "main"),
+    ]
+    assert [r.full_name for r in select_repos(repos, exclude=["skip"])] == ["org/a", "org/b"]
+    assert [r.full_name for r in select_repos(repos, only=["b"])] == ["org/b"]
+    assert [r.full_name for r in select_repos(repos, only=["fork"])] == []
+    assert [r.full_name for r in select_repos(repos, limit=1)] == ["org/a"]
+
+
+def test_clone_then_update(tmp_path, root):
+    remote = make_remote(tmp_path, "app", {"main.py": "x = 1\n"})
+    entries = crawl(root, [remote])
+    assert entries["org/app"].status == "ok"
+    assert entries["org/app"].commit_sha == head(remote)
+    assert git(root / "org/app", "rev-parse", "--is-shallow-repository") == "true"
+
+    new_sha = commit(tmp_path / "remotes/app", {"main.py": "x = 2\n"})
+    entries = crawl(root, [remote])
+    assert entries["org/app"].commit_sha == new_sha
+    assert (root / "org/app/main.py").read_text() == "x = 2\n"
+    assert read_manifest(root)["org/app"].commit_sha == new_sha
+
+
+def test_non_main_default_branch(tmp_path, root):
+    remote = make_remote(tmp_path, "legacy", {"a.py": "x = 1\n"}, branch="master")
+    assert crawl(root, [remote])["org/legacy"].status == "ok"
+
+
+def test_local_changes_are_left_alone(tmp_path, root):
+    remote = make_remote(tmp_path, "app", {"main.py": "x = 1\n"})
+    old_sha = crawl(root, [remote])["org/app"].commit_sha
+    (root / "org/app/main.py").write_text("edited locally\n")
+    commit(tmp_path / "remotes/app", {"main.py": "x = 2\n"})
+    entry = crawl(root, [remote])["org/app"]
+    assert entry.status == "skipped" and entry.commit_sha == old_sha
+    assert (root / "org/app/main.py").read_text() == "edited locally\n"
+
+
+def test_failures_are_isolated(tmp_path, root):
+    good = make_remote(tmp_path, "good", {"a.py": "x = 1\n"})
+    bad = RepoInfo("org/missing", (tmp_path / "nope").as_uri(), "main")
+    entries = crawl(root, [good, bad])
+    assert entries["org/good"].status == "ok"
+    assert entries["org/missing"].status == "failed" and entries["org/missing"].error
+    assert not (root / "org/missing").exists()
+
+
+def test_fetch_failure_keeps_last_good_copy(tmp_path, root):
+    remote = make_remote(tmp_path, "app", {"main.py": "x = 1\n"})
+    sha = crawl(root, [remote])["org/app"].commit_sha
+    moved = RepoInfo("org/app", (tmp_path / "gone").as_uri(), "main")
+    git(root / "org/app", "remote", "set-url", "origin", moved.clone_url)
+    entry = crawl(root, [moved])["org/app"]
+    assert entry.status == "stale" and entry.commit_sha == sha
+
+
+def test_prune_removes_only_repos_the_crawler_cloned(tmp_path, root):
+    a = make_remote(tmp_path, "a", {"a.py": "x = 1\n"})
+    b = make_remote(tmp_path, "b", {"b.py": "x = 1\n"})
+    crawl(root, [a, b])
+    (root / "someone-elses-dir").mkdir()
+
+    entries = crawl(root, [a], prune=False)  # partial runs (--only/--limit) never delete
+    assert (root / "org/b").exists() and entries["org/b"].status == "ok"
+
+    entries = crawl(root, [a])
+    assert entries["org/b"].status == "removed"
+    assert not (root / "org/b").exists()
+    assert (root / "someone-elses-dir").exists()
+    assert "org/b" not in crawl(root, [a])  # reported once, then dropped
+
+
+def test_pipeline_writes_combined_chunks(tmp_path, root):
+    a = make_remote(tmp_path, "a", {"src/app.py": "def main():\n    return 1\n", "data/x.csv": "a,b\n"})
+    b = make_remote(tmp_path, "b", {"README.md": "# B\n\nHello.\n"})
+    out = tmp_path / "chunks.jsonl"
+    result = pipeline.run(root, out, [a, b], lambda t: len(t.split()), jobs=2, scan_secrets=False)
+    records = [json.loads(line) for line in out.read_text().splitlines()]
+    assert {(r["repo"], r["path"]) for r in records} == {("org/a", "src/app.py"), ("org/b", "README.md")}
+    assert {r["commit_sha"] for r in records if r["repo"] == "org/a"} == {head(a)}
+    assert result["chunks"] == {"org/a": 1, "org/b": 1}
+    assert not out.with_name(out.name + ".tmp").exists()
+
+
+def test_pipeline_skips_repos_it_could_not_sync(tmp_path, root):
+    good = make_remote(tmp_path, "good", {"a.py": "x = 1\n"})
+    bad = RepoInfo("org/missing", (tmp_path / "nope").as_uri(), "main")
+    result = pipeline.run(root, tmp_path / "out.jsonl", [good, bad], lambda t: len(t.split()), scan_secrets=False)
+    assert result["chunks"] == {"org/good": 1}
+
+
+def test_default_config_is_valid():
+    config = crawler.load_config()
+    assert config["org"] == "CornellDataScience"
+    assert isinstance(config["exclude"], list)