airoh.provenance

  1# src/airoh/provenance.py
  2"""📜 Lightweight provenance records for a pipeline's inputs and outputs.
  3
  4Two records, written by two tasks:
  5
  6* ``record_sources`` → ``source_data/MANIFEST.json``, written at the end of
  7  ``fetch``: what each data asset actually resolved to (a URL, a real path
  8  behind a symlink), how big it was, its checksum, and — when the asset lives
  9  in a git or datalad repository — the commit it sat at.
 10* ``record_run`` → ``output_data/PROVENANCE.json``, written at the end of
 11  ``run``: the project commit, the environment, the manifest it consumed, and
 12  a checksum of every file the run produced.
 13
 14Together they answer "which inputs produced these outputs, on which code, in
 15which environment" without requiring datalad. When a project *is* tracked with
 16datalad, these records are redundant but harmless — datalad remains the only
 17thing that can actually *retrieve* a past state. See CLAUDE.md.
 18
 19Everything here is tolerant by design: a provenance record is documentation,
 20never a precondition, so a missing file, an absent git binary or an unreadable
 21asset is recorded as ``null`` with a warning. Nothing in this module raises.
 22"""
 23
 24import hashlib
 25import json
 26import os
 27import platform
 28import re
 29import subprocess
 30import sys
 31from datetime import datetime, timezone
 32from pathlib import Path
 33
 34from invoke import task
 35
 36# Hashing a multi-gigabyte asset would make every fetch crawl for a number
 37# nobody reads, so files above this size record their metadata only. Override
 38# with `provenance_hash_max_bytes` in invoke.yaml.
 39DEFAULT_HASH_MAX_BYTES = 100 * 1024 * 1024
 40
 41# git is invoked for metadata only, so it should answer immediately. A bounded
 42# timeout keeps an unreachable remote or a credential prompt from stalling a
 43# fetch (git is also run with prompts disabled, below).
 44_GIT_TIMEOUT_SECONDS = 15
 45
 46_GIT_ENV = dict(os.environ, GIT_TERMINAL_PROMPT="0")
 47
 48
 49def _utc_now():
 50    """Current UTC time as an ISO-8601 string."""
 51    return datetime.now(timezone.utc).isoformat(timespec="seconds")
 52
 53
 54def _git(args, cwd):
 55    """Return the stripped stdout of a git command, or None if it fails.
 56
 57    Tolerant: a missing git binary, a path outside any repository, or a
 58    timeout all yield None rather than an exception.
 59    """
 60    try:
 61        result = subprocess.run(
 62            ["git"] + args, cwd=str(cwd), capture_output=True, text=True,
 63            stdin=subprocess.DEVNULL, env=_GIT_ENV, timeout=_GIT_TIMEOUT_SECONDS,
 64        )
 65    except (OSError, subprocess.SubprocessError):
 66        return None
 67    if result.returncode != 0:
 68        return None
 69    return result.stdout.strip()
 70
 71
 72def git_info(path):
 73    """Describe the git state of the repository containing ``path``.
 74
 75    Returns a dict with ``commit``, ``branch``, ``remote`` and ``dirty``, or
 76    None when ``path`` is not inside a git repository (or git is unavailable).
 77    Used both for data assets — a symlinked datalad superdataset is a git
 78    repository, so its commit pins the input state — and for the project repo
 79    itself.
 80    """
 81    path = Path(path)
 82    cwd = path if path.is_dir() else path.parent
 83    if not cwd.exists():
 84        return None
 85    if _git(["rev-parse", "--is-inside-work-tree"], cwd) != "true":
 86        return None
 87    status = _git(["status", "--porcelain"], cwd)
 88    return {
 89        "toplevel": _git(["rev-parse", "--show-toplevel"], cwd),
 90        "commit": _git(["rev-parse", "HEAD"], cwd),
 91        "branch": _git(["rev-parse", "--abbrev-ref", "HEAD"], cwd),
 92        "remote": _git(["config", "--get", "remote.origin.url"], cwd),
 93        "dirty": bool(status) if status is not None else None,
 94    }
 95
 96
 97def _datalad_id(path):
 98    """The datalad dataset id at ``path``, or None if it is not one."""
 99    config = Path(path) / ".datalad" / "config"
100    if not config.is_file():
101        return None
102    try:
103        text = config.read_text()
104    except OSError:
105        return None
106    match = re.search(r"^\s*id\s*=\s*(\S+)", text, re.MULTILINE)
107    return match.group(1) if match else None
108
109
110def sha256_file(path, max_bytes=DEFAULT_HASH_MAX_BYTES):
111    """Hex sha256 of ``path``, or None if it is too large or unreadable.
112
113    Files over ``max_bytes`` are deliberately skipped — see
114    ``DEFAULT_HASH_MAX_BYTES``.
115    """
116    path = Path(path)
117    try:
118        if path.stat().st_size > max_bytes:
119            return None
120        digest = hashlib.sha256()
121        with open(path, "rb") as handle:
122            for block in iter(lambda: handle.read(1024 * 1024), b""):
123                digest.update(block)
124        return digest.hexdigest()
125    except OSError:
126        return None
127
128
129def describe_path(path, max_bytes=DEFAULT_HASH_MAX_BYTES):
130    """Describe one file or directory for a provenance record.
131
132    Records existence, size, mtime and (for files) a checksum, plus the git and
133    datalad state of whatever the path really points at. A symlink is resolved
134    first, so an asset linked to a checkout elsewhere on disk is attributed to
135    that checkout rather than to the link.
136    """
137    path = Path(path)
138    record = {
139        "path": str(path),
140        "exists": path.exists(),
141        "is_symlink": path.is_symlink(),
142        "resolved_path": None,
143        "size_bytes": None,
144        "mtime": None,
145        "sha256": None,
146        "git": None,
147        "datalad_id": None,
148    }
149    if not path.exists():
150        return record
151
152    resolved = path.resolve()
153    record["resolved_path"] = str(resolved)
154    try:
155        stat = resolved.stat()
156        record["mtime"] = datetime.fromtimestamp(
157            stat.st_mtime, timezone.utc).isoformat(timespec="seconds")
158        if resolved.is_file():
159            record["size_bytes"] = stat.st_size
160            record["sha256"] = sha256_file(resolved, max_bytes)
161    except OSError:
162        pass
163
164    record["git"] = git_info(resolved)
165    record["datalad_id"] = _datalad_id(resolved)
166    return record
167
168
169def _asset_entries(config):
170    """Yield ``(name, target_path, declared_source)`` for every data asset.
171
172    Covers both the ``files:`` section (assets with an ``output_file``) and the
173    ``datasets:`` section, which projects write either as a plain path string
174    or as a mapping with ``output_dir``/``output_file``.
175    """
176    for name, entry in (config.get("files") or {}).items():
177        if isinstance(entry, dict):
178            target = entry.get("output_file") or entry.get("output_dir")
179            yield name, target, entry.get("source") or entry.get("url")
180
181    for name, entry in (config.get("datasets") or {}).items():
182        if isinstance(entry, dict):
183            target = entry.get("output_dir") or entry.get("output_file")
184            yield name, target, entry.get("source") or entry.get("url")
185        else:
186            yield name, entry, None
187
188
189def _asset_mode(record, declared_source):
190    """How the asset got here: symlink, datalad, download, copy or local."""
191    if record["is_symlink"]:
192        return "symlink"
193    if record["datalad_id"]:
194        return "datalad"
195    if isinstance(declared_source, str) and declared_source.startswith(
196            ("http://", "https://", "ftp://")):
197        return "download"
198    return "copy" if declared_source else "local"
199
200
201def _drop_self_attribution(record, project_toplevel):
202    """Forget git state that just describes the project repo itself.
203
204    An asset downloaded into ``source_data/`` sits inside the project's own
205    repository, and reporting the project's commit as that asset's version
206    would be worse than reporting nothing — it looks like real input
207    provenance. Only an asset in a repository of its own keeps its git block.
208    """
209    git = record.get("git")
210    if git and project_toplevel and git.get("toplevel") == project_toplevel:
211        record["git"] = None
212    return record
213
214
215@task(help={"output": "Where to write the manifest (default: the "
216                      "`manifest_file` key in invoke.yaml)."})
217def record_sources(c, output=None):
218    """📜 Record what every data asset resolved to, into source_data/MANIFEST.json.
219
220    Call at the end of ``fetch``. For each asset declared under ``files:`` or
221    ``datasets:`` in invoke.yaml, records the path it landed at, what it really
222    points at, its size and checksum, and the git commit / datalad id of the
223    repository it belongs to — the part datalad would otherwise be needed for.
224
225    Tolerant: an asset that was never fetched is recorded as absent, and no
226    failure here can break a fetch.
227
228    Parameters
229    ----------
230    c : invoke.Context
231        The Invoke context.
232    output : str, optional
233        Manifest path. Defaults to the `manifest_file` key in invoke.yaml,
234        falling back to `source_data/MANIFEST.json`.
235
236    Examples
237    --------
238    ```bash
239    inv provenance.record-sources
240    ```
241    """
242    output_path = Path(output or c.config.get(
243        "manifest_file", "source_data/MANIFEST.json"))
244    max_bytes = c.config.get("provenance_hash_max_bytes", DEFAULT_HASH_MAX_BYTES)
245
246    project_toplevel = (git_info(Path.cwd()) or {}).get("toplevel")
247
248    assets = {}
249    for name, target, declared_source in _asset_entries(c.config):
250        if not target:
251            print(f"⚠️  Asset '{name}' declares no output path — skipping")
252            continue
253        record = describe_path(target, max_bytes)
254        record["declared_source"] = declared_source
255        record["mode"] = _asset_mode(record, declared_source)
256        assets[name] = _drop_self_attribution(record, project_toplevel)
257        if not record["exists"]:
258            print(f"⚠️  Asset '{name}' not present at {target} (recorded as absent)")
259
260    manifest = {
261        "schema": "airoh/manifest/1",
262        "recorded_at": _utc_now(),
263        "assets": assets,
264    }
265    output_path.parent.mkdir(parents=True, exist_ok=True)
266    output_path.write_text(json.dumps(manifest, indent=2, sort_keys=True) + "\n")
267    print(f"📜 Recorded {len(assets)} asset(s) → {output_path}")
268    return output_path
269
270
271def _declared_dependencies():
272    """Package names this project declares, from pyproject.toml or requirements.txt.
273
274    Recording the version of every installed package would bury the handful
275    that matter in a few hundred transitive ones, so only declared
276    dependencies are versioned.
277    """
278    names = []
279    pyproject = Path("pyproject.toml")
280    if pyproject.is_file():
281        names = _parse_pyproject_dependencies(pyproject)
282    if not names:
283        requirements = Path("requirements.txt")
284        if requirements.is_file():
285            names = parse_requirements(requirements)
286    return names
287
288
289def _parse_pyproject_dependencies(path):
290    """Package names in ``[project].dependencies``.
291
292    Uses tomllib/tomli when available (tomllib is stdlib from 3.11) and falls
293    back to reading the ``dependencies = [...]`` array directly, so airoh keeps
294    its single-dependency footprint on older interpreters.
295    """
296    try:
297        try:
298            import tomllib
299        except ImportError:
300            import tomli as tomllib
301        with open(path, "rb") as handle:
302            data = tomllib.load(handle)
303        return [requirement_name(spec)
304                for spec in data.get("project", {}).get("dependencies", [])]
305    except Exception:
306        pass
307
308    try:
309        text = path.read_text()
310    except OSError:
311        return []
312    match = re.search(r"^dependencies\s*=\s*\[(.*?)\]", text, re.MULTILINE | re.DOTALL)
313    if not match:
314        return []
315    return [requirement_name(spec)
316            for spec in re.findall(r"[\"']([^\"']+)[\"']", match.group(1))]
317
318
319def requirement_name(spec):
320    """Bare package name from a requirement spec ('numpy>=1.3' → 'numpy')."""
321    return re.split(r"[\[<>=!~;\s]", spec.strip(), maxsplit=1)[0].lower()
322
323
324def parse_requirements(path):
325    """Package names in a requirements file, ignoring comments and flags."""
326    names = []
327    try:
328        lines = Path(path).read_text().splitlines()
329    except OSError:
330        return names
331    for line in lines:
332        line = line.split("#", 1)[0].strip()
333        if not line or line.startswith("-"):
334            continue
335        names.append(requirement_name(line))
336    return names
337
338
339def _environment():
340    """Python, platform and declared-dependency versions."""
341    from importlib.metadata import PackageNotFoundError, version
342
343    packages = {}
344    for name in sorted(set(_declared_dependencies())):
345        try:
346            packages[name] = version(name)
347        except PackageNotFoundError:
348            packages[name] = None
349    try:
350        packages.setdefault("airoh", version("airoh"))
351    except PackageNotFoundError:
352        pass
353
354    lockfiles = {}
355    for name in ("uv.lock", "requirements.txt", "environment.yml"):
356        if Path(name).is_file():
357            lockfiles[name] = sha256_file(name)
358
359    return {
360        "python": sys.version.split()[0],
361        "platform": platform.platform(),
362        "packages": packages,
363        "lockfiles": lockfiles,
364    }
365
366
367def _output_files(output_dir, max_bytes, exclude):
368    """Checksum every file the pipeline produced under ``output_dir``.
369
370    The record itself, dotfiles and CONTENT.md are repository bookkeeping that
371    lives in the output folder without being an output of the run.
372    """
373    outputs = {}
374    root = Path(output_dir)
375    if not root.is_dir():
376        return outputs
377    for path in sorted(root.rglob("*")):
378        if not path.is_file() or path == exclude:
379            continue
380        if path.name == "CONTENT.md" or any(
381                part.startswith(".") for part in path.relative_to(root).parts):
382            continue
383        try:
384            size = path.stat().st_size
385        except OSError:
386            continue
387        outputs[str(path.relative_to(root))] = {
388            "size_bytes": size,
389            "sha256": sha256_file(path, max_bytes),
390        }
391    return outputs
392
393
394@task(help={"output": "Where to write the record (default: the "
395                      "`provenance_file` key in invoke.yaml).",
396            "tasks": "Comma-separated names of the tasks that ran."})
397def record_run(c, output=None, tasks=None):
398    """📜 Record what produced the current outputs, into output_data/PROVENANCE.json.
399
400    Call at the end of ``run``. Records the project's own git commit and dirty
401    flag, the environment, a checksum of the input manifest (plus an inlined
402    copy of each asset's commit and checksum, so the record still means
403    something if the manifest is lost), and the size and checksum of every
404    output file.
405
406    This file changes on every run — that is the point of it. Do not try to
407    make it stable; if the churn is unwanted, the answer is to run less often,
408    not to record less.
409
410    Tolerant: nothing here can break a run.
411
412    Parameters
413    ----------
414    c : invoke.Context
415        The Invoke context.
416    output : str, optional
417        Record path. Defaults to the `provenance_file` key in invoke.yaml,
418        falling back to `output_data/PROVENANCE.json`.
419    tasks : str, optional
420        Comma-separated names of the tasks that ran, recorded as-is.
421
422    Examples
423    --------
424    ```bash
425    inv provenance.record-run --tasks run-qc-measures,run-notebooks
426    ```
427    """
428    output_path = Path(output or c.config.get(
429        "provenance_file", "output_data/PROVENANCE.json"))
430    max_bytes = c.config.get("provenance_hash_max_bytes", DEFAULT_HASH_MAX_BYTES)
431    manifest_path = Path(c.config.get("manifest_file", "source_data/MANIFEST.json"))
432    output_dir = c.config.get("output_data_dir", "output_data")
433
434    inputs = {"manifest_file": str(manifest_path), "manifest_sha256": None,
435              "assets": {}}
436    if manifest_path.is_file():
437        inputs["manifest_sha256"] = sha256_file(manifest_path, max_bytes)
438        try:
439            manifest = json.loads(manifest_path.read_text())
440            for name, record in (manifest.get("assets") or {}).items():
441                inputs["assets"][name] = {
442                    "resolved_path": record.get("resolved_path"),
443                    "sha256": record.get("sha256"),
444                    "commit": (record.get("git") or {}).get("commit"),
445                }
446        except (OSError, ValueError):
447            print(f"⚠️  Could not read {manifest_path} — recording its hash only")
448    else:
449        print(f"⚠️  No manifest at {manifest_path} — run `invoke fetch` to create it")
450
451    record = {
452        "schema": "airoh/provenance/1",
453        "recorded_at": _utc_now(),
454        "tasks": [name.strip() for name in tasks.split(",")] if tasks else None,
455        "command": " ".join(sys.argv),
456        "repository": git_info(Path.cwd()),
457        "environment": _environment(),
458        "inputs": inputs,
459        "outputs": _output_files(output_dir, max_bytes, output_path),
460    }
461    output_path.parent.mkdir(parents=True, exist_ok=True)
462    output_path.write_text(json.dumps(record, indent=2, sort_keys=True) + "\n")
463    print(f"📜 Recorded {len(record['outputs'])} output(s) → {output_path}")
464    return output_path