diff --git a/.gitignore b/.gitignore index cb03964..9d79a0f 100644 --- a/.gitignore +++ b/.gitignore @@ -165,3 +165,5 @@ cline_mcp.json # Local scratch for removed/parked code (not tracked) scratch/ + +use_cases/*.tar.gz diff --git a/CLAUDE.md b/CLAUDE.md index a911517..8de4206 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -47,7 +47,8 @@ The codebase separates **commands** (entry points with argparse, launched as CLI - `agents/` — Per-agent-platform setup (`base.py` ABC + `claude.py` / `goose.py` / `cline.py` / `codex.py` / `opencode.py`). Each subclass owns its `write_static`, `write_dynamic`, `runtime_env`, `vscode_hint`. Shared helpers (`_mcp_env_block`, `_build_mcp_servers_dict`) in `base.py`. DSAGT sets no telemetry/OTel env, writes no launch shim, and never touches provider credentials — agents are expected pre-authenticated (shell env / their own auth flows) before dsagt is pointed at them; dsagt prints no credential hints and never troubleshoots auth. - `knowledge.py` — ChromaDB document retrieval, embedding backends, per-collection routing (the reference example of the house style). - `registry.py` — `CodeRegistry` (CLI codes) + `SkillRegistry` (agent instruction skills), KB indexing. -- `provenance.py` — Code execution records (`run_and_record`), execution-record indexing into ChromaDB (`CodeUseIndexer` → `code_use` collection), pipeline reconstruction (`reconstruct_pipeline`, dependency graph). +- `provenance.py` — Code execution records (`run_and_record`), execution-record indexing into ChromaDB (`CodeUseIndexer` → `code_use` collection), pipeline reconstruction (`reconstruct_pipeline`, dependency graph, `compute_pipeline_fingerprint`). +- `contract.py` — the sample-contract schema: `validate_contract` / `load_contract` / `save_contract` for `/dataset_contract.yaml`, the user-facing pivot artifact the dataset builder derives from (sample keys, collation, split policy, normalization ownership, pipeline fingerprint, producer/consumer reconciliation). See `docs/dataset-contract.md`. - `observability.py` — first-party span emission over the serverless sqlite store via MLflow's native `mlflow.start_span` (no OTel `TracerProvider`). `resolve_tracking_uri` (never-raise), `init_tracing`, `@traced`/`obs`/`child_span` + typed span helpers. Each internal trace's root is tagged `dsagt.source` with the MCP tool *category* (`memory`/`skill`/`knowledge`/`registry`, or `execution` for dsagt-run) so the MLflow UI can filter the debug view apart from agent traces. `MLflowSink` (a `traces.Trace` consumer) replays finished transcripts via `start_span_no_context`. - `memory.py` — Explicit memory (YAML, `ExplicitMemory`) + the episodic `MemoryExtractor` (a `traces.TraceCollector` consumer that mechanically chunks+tags+embeds every turn, no LLM). Turns carry `ts_epoch` for recency-weighted retrieval. `extract_session` is a no-op stub kept only for the deferred cross-session N+1 catch-up call site. - `skills.py` — External skill-catalog data plane (`SkillsCatalog`: clone/sync/index/install), the `SkillRouter` render facade, and the Genesis-derived keyword scorer (`rank_skills`). diff --git a/docs/dataset-contract.md b/docs/dataset-contract.md new file mode 100644 index 0000000..af54569 --- /dev/null +++ b/docs/dataset-contract.md @@ -0,0 +1,256 @@ +# Dataset Contract + +A **sample contract** is the pivot artifact for turning a data directory into a +training-ready PyTorch dataset package. It is a YAML file, persisted at +`/dataset_contract.yaml` (project root, not `.dsagt/`, because it is a +user-facing artifact the scientist reviews and signs off on) declaring what one +training sample looks like, how samples collate into a batch, how the data +splits, and which side owns normalization. + +This page documents the schema (`dsagt.contract`). The tooling that builds a +contract from a pipeline or a `Dataset` and consumes it (`check-dataset`, the +`dataset-builder` skill, the package scaffolder, the reference model) is +separate, forthcoming work; this schema is what they share. + +## Two contracts, reconciled + +A `Dataset`'s `__getitem__` is the adapter between two views of the same +sample: + +- **Producer contract**: what is actually on disk. +- **Consumer contract**: what `model.forward()` expects. + +The `reconciliation` section of the schema records where the two differ and +how the adapter resolves each difference: dtype casts, channel layout, +normalization ownership, padding and ragged handling, label encoding. + +## Schema + +```yaml +version: 1 # schema version, int +mode: pipeline # "pipeline" | "standalone" +pipeline_fingerprint: sha256:... # required iff mode == pipeline; absent iff standalone + +keys: + : + dtype: float32 # framework dtype name + shape: [N, 3] # ints, or symbolic dim names (e.g. "N" for a + # variable node count) + role: input # input | target | mask | metadata + value_range: [0.0, 1.0] # optional [min, max] + collation: stack # per-key collation rule (stack, graph_batch, list, ...) + +normalization: + owner: pipeline # pipeline | dataset — whether a pipeline step + # already normalized, or __getitem__ must + +split: + strategy: group # random | group | time + group_key: # required for group/time — the id samples + # sharing it must not be split across + seed: 42 + ratios: + train: 0.7 + val: 0.15 + test: 0.15 + +reconciliation: + - key: + producer: "what's actually on disk" + consumer: "what model.forward() expects" + resolution: "how __getitem__ bridges the two" +``` + +`dsagt.contract.validate_contract` enforces this shape; `load_contract` / +`save_contract` validate on the way in and out. + +## Field reference + +**`mode`** + +| Value | Meaning | +|---|---| +| `pipeline` | The sample's inputs are terminal outputs of a DSAgt-tracked pipeline. The contract carries `pipeline_fingerprint`. | +| `standalone` | No DSAgt pipeline; the data root was characterized directly (e.g. via the `scan-directory` code). No pipeline to fingerprint — the field must be absent. | + +**`keys..role`** — what `__getitem__` hands the training loop + +| Value | Meaning | +|---|---| +| `input` | Fed to `model.forward()`. | +| `target` | The ground truth compared against the model's output in the loss. | +| `mask` | A boolean/float mask consumed alongside an input or target key (padding mask, loss mask, node/edge validity mask). | +| `metadata` | Carried for bookkeeping, splitting, or debugging (ids, provenance, plot coordinates) — never passed to `forward()` or the loss. | + +**`keys..collation`** is a free-form string, not a closed enum — the collate rule a key needs is domain-specific, and mapping a name to generated `collate_fn` code is the dataset-builder skill's job, not this schema's. Common values seen in practice: `stack` (equal-shape tensors, `torch.stack`), `pad` (ragged sequences, pad to batch max), `list` (leave as a Python list, e.g. non-tensor metadata), `graph_batch` (PyG-style disjoint-union batching for graph keys). + +**`normalization.owner`** — who has already applied centering/scaling to a key's raw values + +| Value | Meaning | +|---|---| +| `pipeline` | An upstream registered code already normalized the data on disk; `__getitem__` passes the value through unchanged. | +| `dataset` | No upstream normalization; `__getitem__` computes and applies it itself (e.g. using statistics fit over the train split). | + +**`split.strategy`** — how `split.ratios` is turned into train/val/test membership + +| Value | `split.group_key` | Meaning | +|---|---|---| +| `random` | not used | i.i.d. row-level assignment, seeded by `split.seed`. Leaks whenever samples share an identity (a patient, a simulation case, augmented copies of one image); use only when samples are genuinely independent. | +| `group` | required | Every sample whose `group_key` takes the same value is assigned to the same split, so no group straddles a split boundary. | +| `time` | required | Samples are ordered by `group_key`'s value (a timestamp or step index) and cut chronologically, so later time periods are held out rather than interleaved with training data. | + +## The pipeline fingerprint + +In pipeline mode, the contract carries a hash over the upstream pipeline's +structural shape, computed by `provenance.compute_pipeline_fingerprint` from +`reconstruct_pipeline(..., fmt="json")`'s `dependency_graph` and +`terminal_outputs` (never `records` — timestamps and stdout vary rerun to +rerun even when the pipeline hasn't changed). A later staleness check +recomputes the fingerprint and flags the contract for review when it no +longer matches: a step was added or removed, or an output path changed. + +Standalone mode (no DSAgt pipeline; the data root was characterized directly +with codes like `scan-directory`) has no pipeline to fingerprint, so the +field is absent. + +## Worked example: tabular case (standalone) + +A CSV of patient records: 37 numeric features, an integer label, grouped by +patient so repeated visits from the same patient never split across +train/val/test. + +```yaml +version: 1 +mode: standalone + +keys: + features: + dtype: float32 + shape: [37] + role: input + value_range: [-3.0, 3.0] + collation: stack + label: + dtype: int64 + shape: [] + role: target + collation: stack + patient_id: + dtype: string + shape: [] + role: metadata + collation: list + +normalization: + owner: dataset + +split: + strategy: group + group_key: patient_id + seed: 42 + ratios: + train: 0.7 + val: 0.15 + test: 0.15 + +reconciliation: + - key: features + producer: float64 columns in the source CSV, unnormalized + consumer: float32 tensor, zero mean / unit variance expected by model.forward() + resolution: cast to float32 and standardize in __getitem__ using stats computed + over the train split + - key: label + producer: string category name + consumer: integer class index expected by the loss function + resolution: label encoding fit at contract-build time; mapping stored alongside + the split manifest +``` + +## Worked example: XGC graph case (pipeline) + +[`XGCGraphDataset`](https://github.com/AI-ModCon/dsagt/blob/main/use_cases/fusion-fm/skills/xgc-ai-training/scripts/xgc_dataset.py) +returns a PyG `Data` graph per `(phi_plane, start_timestep)` sample: node +features `x`, target field values `y`, node positions, mesh edges, and a few +scalar metadata fields. Splits are made by phi-plane group so every timestep +from one plane stays in one split. This confirms the schema handles a +variable node count (`N`), an edge dimension (`E`), and PyG's own graph +collation. + +```yaml +version: 1 +mode: pipeline +pipeline_fingerprint: sha256:6f1ea1e0c1a1c9e5c9b9f2e8f7d4a3b2c1d0e9f8a7b6c5d4e3f2a1b0c9d8e7f6 + +keys: + x: + dtype: float32 + shape: [N, n_steps, "7+F"] + role: input + collation: graph_batch + pos: + dtype: float32 + shape: [N, 2] + role: input + collation: graph_batch + edge_index: + dtype: int64 + shape: [2, E] + role: input + collation: graph_batch + edge_attr: + dtype: float32 + shape: [E, 3] + role: input + collation: graph_batch + leadtime: + dtype: float32 + shape: [1, 1] + role: input + collation: stack + y: + dtype: float32 + shape: [N, F] + role: target + collation: graph_batch + phi: + dtype: int64 + shape: [] + role: metadata + collation: list + step0: + dtype: int64 + shape: [] + role: metadata + collation: list + target_step: + dtype: int64 + shape: [] + role: metadata + collation: list + +normalization: + owner: pipeline + +split: + strategy: group + group_key: phi + seed: 7 + ratios: + train: 0.7 + val: 0.15 + test: 0.15 + +reconciliation: + - key: x + producer: per-field npz arrays, one file per simulation step + resolution: __getitem__ loads the requested step_*.npz files and concatenates + static_ctx with the field slice + consumer: single [N, n_steps, 7+F] tensor stacking static node context and + per-step field values + - key: edge_index + producer: triangular mesh connectivity in mesh.npz + consumer: undirected PyG edge_index expected by the message-passing layers + resolution: cells_to_edge_index(undirected=True), cached once as topology.pt +``` + +Both worked examples are validated as part of the test suite (`tests/test_contract.py`). diff --git a/mkdocs.yml b/mkdocs.yml index 26b41fe..05c3a17 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -57,6 +57,7 @@ nav: - Architecture: architecture.md - Capabilities: - Provenance: provenance.md + - Dataset Contract: dataset-contract.md - Knowledge Base: knowledge-base.md - Skills: skills.md - Memory: memory.md diff --git a/src/dsagt/contract.py b/src/dsagt/contract.py new file mode 100644 index 0000000..c6b4cc0 --- /dev/null +++ b/src/dsagt/contract.py @@ -0,0 +1,233 @@ +""" +Sample contract: the pivot artifact the dataset builder derives from. + +A **sample contract** is a YAML file, persisted at ``/dataset_contract.yaml``, +declaring what one training sample looks like: its keys (dtype, shape, value +range, semantic role), the collation rule per key, the split policy, and which +side (an upstream pipeline step, or ``__getitem__``) owns normalization. It is +written at the project root rather than under ``.dsagt/`` because it is a +user-facing artifact the scientist reviews and signs off on, not server-owned +state. + +The contract reconciles two views of the same sample: the **producer** +contract (what is actually on disk) and the **consumer** contract (what +``model.forward()`` expects). The ``reconciliation`` section records where +they differ and how the generated ``Dataset`` adapts one to the other (dtype +casts, channel layout, normalization ownership, padding/ragged handling, +label encoding). + +In **pipeline mode**, the contract also carries a ``pipeline_fingerprint``: a +hash (see ``provenance.compute_pipeline_fingerprint``) over the dependency +graph and terminal outputs of the upstream ``reconstruct_pipeline`` run that +produced the sample's inputs. A later staleness check recomputes the +fingerprint and flags the contract for review if the upstream pipeline has +changed. In **standalone mode** (no DSAgt pipeline; the data root was +characterized directly) there is nothing to fingerprint, so the field is +absent. + +This module only defines and validates the schema; nothing here builds a +contract from a pipeline or a ``Dataset`` — that is the dataset-builder +skill's job. +""" + +from __future__ import annotations + +from pathlib import Path + +import yaml + +#: Filename the contract is persisted under, at the project root. +CONTRACT_FILENAME = "dataset_contract.yaml" + +#: How the sample's inputs were produced. +#: pipeline - the inputs are terminal outputs of a DSAgt-tracked pipeline; +#: the contract carries a pipeline_fingerprint (see below) for +#: the staleness check. +#: standalone - no DSAgt pipeline; the data root was characterized directly +#: (e.g. via the scan-directory code). No pipeline to +#: fingerprint, so the field must be absent. +VALID_MODES = {"pipeline", "standalone"} + +#: What a sample key is for, i.e. what __getitem__ hands the training loop. +#: input - fed to model.forward(). +#: target - the ground truth compared against the model's output in the +#: loss. +#: mask - a boolean/float mask consumed alongside an input or target key +#: (padding mask, loss mask, node/edge validity mask). +#: metadata - carried on the sample for bookkeeping, splitting, or debugging +#: (ids, provenance, plot coordinates) but never passed to +#: forward() or the loss. +VALID_ROLES = {"input", "target", "mask", "metadata"} + +#: Who has already applied normalization (centering/scaling) to a key's raw +#: values. +#: pipeline - an upstream registered code already normalized the data on +#: disk; __getitem__ passes the value through unchanged. +#: dataset - no upstream normalization; __getitem__ computes and applies it +#: itself (e.g. using statistics fit over the train split). +VALID_NORMALIZATION_OWNERS = {"pipeline", "dataset"} + +#: How split.ratios is turned into train/val/test membership. +#: random - i.i.d. row-level assignment, seeded by split.seed. Leaks +#: whenever samples share an identity (a patient, a simulation +#: case, augmented copies of one image); use only when samples are +#: genuinely independent. +#: group - every sample whose split.group_key takes the same value is +#: assigned to the same split, so no group straddles a boundary. +#: time - samples are ordered by split.group_key's value (a timestamp or +#: step index) and cut chronologically, so later time periods are +#: held out rather than interleaved with training data. +VALID_SPLIT_STRATEGIES = {"random", "group", "time"} + +#: Split strategies where split.group_key is required: "group" clusters equal +#: values into one split, "time" sorts by the same field before cutting +#: chronologically. Excludes "random", which has no notion of a key to +#: leak across. +_GROUPED_SPLIT_STRATEGIES = {"group", "time"} + +_RATIO_TOLERANCE = 1e-6 + + +def validate_contract(contract: dict) -> None: + """Validate a sample contract dict against the schema. Raises ``ValueError``.""" + version = contract.get("version") + if not isinstance(version, int): + raise ValueError(f"'version' must be an int, got {version!r}") + + mode = contract.get("mode") + if mode not in VALID_MODES: + raise ValueError(f"'mode' must be one of {sorted(VALID_MODES)}, got {mode!r}") + + fingerprint = contract.get("pipeline_fingerprint") + if mode == "pipeline": + if not fingerprint or not isinstance(fingerprint, str): + raise ValueError( + "'pipeline_fingerprint' is required when mode is 'pipeline'" + ) + elif fingerprint is not None: + raise ValueError( + "'pipeline_fingerprint' must be absent when mode is 'standalone'" + ) + + keys = contract.get("keys") + if not isinstance(keys, dict) or not keys: + raise ValueError( + "'keys' must be a non-empty mapping of sample key name to spec" + ) + for name, spec in keys.items(): + _validate_key_spec(name, spec) + + _validate_normalization(contract.get("normalization")) + _validate_split(contract.get("split")) + + reconciliation = contract.get("reconciliation", []) + _validate_reconciliation(reconciliation, keys) + + +def _validate_key_spec(name: str, spec: dict) -> None: + if not isinstance(spec, dict): + raise ValueError(f"keys.{name} must be a mapping, got {type(spec).__name__}") + + dtype = spec.get("dtype") + if not dtype or not isinstance(dtype, str): + raise ValueError(f"keys.{name}.dtype is required and must be a string") + + shape = spec.get("shape") + if not isinstance(shape, list): + raise ValueError( + f"keys.{name}.shape must be a list (ints or symbolic dim names)" + ) + for dim in shape: + if not isinstance(dim, (int, str)) or isinstance(dim, bool): + raise ValueError( + f"keys.{name}.shape entries must be ints or symbolic dim names, got {dim!r}" + ) + + role = spec.get("role") + if role not in VALID_ROLES: + raise ValueError( + f"keys.{name}.role must be one of {sorted(VALID_ROLES)}, got {role!r}" + ) + + # Not a closed enum: the collate_fn a key needs is domain-specific (stack, + # pad, list, graph_batch, ...), and the dataset-builder skill (not this + # module) is what maps a collation name to generated code. Only presence + # is validated here. + collation = spec.get("collation") + if not collation or not isinstance(collation, str): + raise ValueError(f"keys.{name}.collation is required and must be a string") + + value_range = spec.get("value_range") + if value_range is not None: + if len(value_range) != 2 or value_range[0] > value_range[1]: + raise ValueError(f"keys.{name}.value_range must be a 2-element [min, max]") + + +def _validate_normalization(normalization: dict) -> None: + if not isinstance(normalization, dict): + raise ValueError("'normalization' must be a mapping with an 'owner' field") + owner = normalization.get("owner") + if owner not in VALID_NORMALIZATION_OWNERS: + raise ValueError( + f"normalization.owner must be one of {sorted(VALID_NORMALIZATION_OWNERS)}, " + f"got {owner!r}" + ) + + +def _validate_split(split: dict) -> None: + if not isinstance(split, dict): + raise ValueError("'split' must be a mapping") + + strategy = split.get("strategy") + if strategy not in VALID_SPLIT_STRATEGIES: + raise ValueError( + f"split.strategy must be one of {sorted(VALID_SPLIT_STRATEGIES)}, got {strategy!r}" + ) + + if not isinstance(split.get("seed"), int): + raise ValueError("split.seed is required and must be an int") + + if strategy in _GROUPED_SPLIT_STRATEGIES and not split.get("group_key"): + raise ValueError( + f"split.group_key is required when split.strategy is '{strategy}'" + ) + + ratios = split.get("ratios") + if not isinstance(ratios, dict) or not ratios: + raise ValueError( + "split.ratios must be a non-empty mapping of split name to fraction" + ) + total = sum(ratios.values()) + if abs(total - 1.0) > _RATIO_TOLERANCE: + raise ValueError(f"split.ratios must sum to 1.0, got {total}") + + +def _validate_reconciliation(reconciliation: list, keys: dict) -> None: + if not isinstance(reconciliation, list): + raise ValueError("'reconciliation' must be a list") + for i, entry in enumerate(reconciliation): + if not isinstance(entry, dict): + raise ValueError(f"reconciliation[{i}] must be a mapping") + key = entry.get("key") + if key not in keys: + raise ValueError( + f"reconciliation[{i}].key {key!r} is not a declared sample key" + ) + for field in ("producer", "consumer", "resolution"): + if not entry.get(field): + raise ValueError(f"reconciliation[{i}].{field} is required") + + +def load_contract(path: str | Path) -> dict: + """Read and validate a sample contract from disk.""" + contract = yaml.safe_load(Path(path).read_text()) or {} + validate_contract(contract) + return contract + + +def save_contract(path: str | Path, contract: dict) -> None: + """Validate and write a sample contract to disk.""" + validate_contract(contract) + Path(path).write_text( + yaml.dump(contract, default_flow_style=False, sort_keys=False) + ) diff --git a/src/dsagt/provenance.py b/src/dsagt/provenance.py index 6df8860..24c9668 100644 --- a/src/dsagt/provenance.py +++ b/src/dsagt/provenance.py @@ -12,12 +12,17 @@ **Pipeline reconstruction**: Reads execution records, builds a dependency graph from input/output - file overlap, and renders as a bash script or Snakemake workflow. + file overlap, and renders as a bash script, a Snakemake workflow, or + structured JSON (records, dependency graph, terminal outputs). + ``compute_pipeline_fingerprint`` hashes the JSON output's dependency + graph and terminal outputs into the pipeline fingerprint a dataset + contract (``contract.py``) uses for its staleness check. """ from __future__ import annotations import fcntl +import hashlib import json import logging import subprocess @@ -608,6 +613,26 @@ def render_json(records: list[dict], deps: dict[int, list[int]]) -> str: return json.dumps(payload, indent=2) +def compute_pipeline_fingerprint(structured_output: dict) -> str: + """Hash a pipeline's structural shape from a ``reconstruct_pipeline`` JSON payload. + + Hashes only ``dependency_graph`` and ``terminal_outputs`` — never + ``records``, whose ``stdout``/``stderr``/timestamps vary rerun to rerun + even when the pipeline itself hasn't changed. Consumed by the dataset + contract's staleness check: stable across reruns of an unchanged + pipeline, and changes when a step is added, removed, or has its output + paths altered. + """ + payload = { + "dependency_graph": structured_output.get("dependency_graph", {}), + "terminal_outputs": structured_output.get("terminal_outputs", []), + } + digest = hashlib.sha256( + json.dumps(payload, sort_keys=True, separators=(",", ":")).encode() + ).hexdigest() + return f"sha256:{digest}" + + def _shell_quote(s: str) -> str: """Quote a string for shell if it contains special characters.""" if not s: diff --git a/tests/test_contract.py b/tests/test_contract.py new file mode 100644 index 0000000..32c26c4 --- /dev/null +++ b/tests/test_contract.py @@ -0,0 +1,293 @@ +""" +Tests for the sample-contract schema (``dsagt.contract``). + +Includes the two worked examples from the schema design: a tabular case +(standalone mode) and the XGC graph case +(``use_cases/fusion-fm/skills/xgc-ai-training/scripts/xgc_dataset.py``, +pipeline mode), confirming the schema expresses both. +""" + +import copy + +import pytest + +from dsagt.contract import load_contract, save_contract, validate_contract + +# --------------------------------------------------------------------------- +# Worked example: tabular case (standalone mode) +# --------------------------------------------------------------------------- + +TABULAR_CONTRACT = { + "version": 1, + "mode": "standalone", + "keys": { + "features": { + "dtype": "float32", + "shape": [37], + "role": "input", + "value_range": [-3.0, 3.0], + "collation": "stack", + }, + "label": { + "dtype": "int64", + "shape": [], + "role": "target", + "collation": "stack", + }, + "patient_id": { + "dtype": "string", + "shape": [], + "role": "metadata", + "collation": "list", + }, + }, + "normalization": {"owner": "dataset"}, + "split": { + "strategy": "group", + "group_key": "patient_id", + "seed": 42, + "ratios": {"train": 0.7, "val": 0.15, "test": 0.15}, + }, + "reconciliation": [ + { + "key": "features", + "producer": "float64 columns in the source CSV, unnormalized", + "consumer": "float32 tensor, zero mean / unit variance expected by model.forward()", + "resolution": "cast to float32 and standardize in __getitem__ using stats " + "computed over the train split", + }, + { + "key": "label", + "producer": "string category name", + "consumer": "integer class index expected by the loss function", + "resolution": "label encoding fit at contract-build time; mapping stored " + "alongside the split manifest", + }, + ], +} + +# --------------------------------------------------------------------------- +# Worked example: XGC graph case (pipeline mode) +# --------------------------------------------------------------------------- +# +# Mirrors XGCGraphDataset.__getitem__ in +# use_cases/fusion-fm/skills/xgc-ai-training/scripts/xgc_dataset.py: a PyG +# Data graph (x, y, pos, edge_index, edge_attr) plus scalar metadata +# (leadtime, phi, step0, target_step), split by phi-plane group so every +# step from one plane stays in one split. + +XGC_CONTRACT = { + "version": 1, + "mode": "pipeline", + "pipeline_fingerprint": "sha256:6f1ea1e0c1a1c9e5c9b9f2e8f7d4a3b2c1d0e9f8a7b6c5d4e3f2a1b0c9d8e7f6", + "keys": { + "x": { + "dtype": "float32", + "shape": ["N", "n_steps", "7+F"], + "role": "input", + "collation": "graph_batch", + }, + "pos": { + "dtype": "float32", + "shape": ["N", 2], + "role": "input", + "collation": "graph_batch", + }, + "edge_index": { + "dtype": "int64", + "shape": [2, "E"], + "role": "input", + "collation": "graph_batch", + }, + "edge_attr": { + "dtype": "float32", + "shape": ["E", 3], + "role": "input", + "collation": "graph_batch", + }, + "leadtime": { + "dtype": "float32", + "shape": [1, 1], + "role": "input", + "collation": "stack", + }, + "y": { + "dtype": "float32", + "shape": ["N", "F"], + "role": "target", + "collation": "graph_batch", + }, + "phi": { + "dtype": "int64", + "shape": [], + "role": "metadata", + "collation": "list", + }, + "step0": { + "dtype": "int64", + "shape": [], + "role": "metadata", + "collation": "list", + }, + "target_step": { + "dtype": "int64", + "shape": [], + "role": "metadata", + "collation": "list", + }, + }, + "normalization": {"owner": "pipeline"}, + "split": { + "strategy": "group", + "group_key": "phi", + "seed": 7, + "ratios": {"train": 0.7, "val": 0.15, "test": 0.15}, + }, + "reconciliation": [ + { + "key": "x", + "producer": "per-field npz arrays, one file per simulation step", + "consumer": "single [N, n_steps, 7+F] tensor stacking static node " + "context and per-step field values", + "resolution": "__getitem__ loads the requested step_*.npz files and " + "concatenates static_ctx with the field slice", + }, + { + "key": "edge_index", + "producer": "triangular mesh connectivity in mesh.npz", + "consumer": "undirected PyG edge_index expected by the message-passing layers", + "resolution": "cells_to_edge_index(undirected=True), cached once as topology.pt", + }, + ], +} + + +# --------------------------------------------------------------------------- +# Worked examples validate +# --------------------------------------------------------------------------- + + +class TestWorkedExamples: + + def test_tabular_contract_is_valid(self): + validate_contract(TABULAR_CONTRACT) + + def test_xgc_contract_is_valid(self): + validate_contract(XGC_CONTRACT) + + def test_tabular_and_xgc_round_trip_through_disk(self, tmp_path): + for name, contract in (("tabular", TABULAR_CONTRACT), ("xgc", XGC_CONTRACT)): + path = tmp_path / f"{name}_dataset_contract.yaml" + save_contract(path, contract) + assert load_contract(path) == contract + + +# --------------------------------------------------------------------------- +# Schema validation failures +# --------------------------------------------------------------------------- + + +class TestValidateContract: + + def _tabular(self) -> dict: + return copy.deepcopy(TABULAR_CONTRACT) + + def test_missing_version(self): + contract = self._tabular() + del contract["version"] + with pytest.raises(ValueError, match="version"): + validate_contract(contract) + + def test_invalid_mode(self): + contract = self._tabular() + contract["mode"] = "bogus" + with pytest.raises(ValueError, match="mode"): + validate_contract(contract) + + def test_pipeline_mode_requires_fingerprint(self): + contract = self._tabular() + contract["mode"] = "pipeline" + with pytest.raises(ValueError, match="pipeline_fingerprint"): + validate_contract(contract) + + def test_standalone_mode_forbids_fingerprint(self): + contract = copy.deepcopy(XGC_CONTRACT) + contract["mode"] = "standalone" + with pytest.raises(ValueError, match="pipeline_fingerprint"): + validate_contract(contract) + + def test_empty_keys_rejected(self): + contract = self._tabular() + contract["keys"] = {} + with pytest.raises(ValueError, match="keys"): + validate_contract(contract) + + def test_key_missing_dtype(self): + contract = self._tabular() + del contract["keys"]["label"]["dtype"] + with pytest.raises(ValueError, match="dtype"): + validate_contract(contract) + + def test_key_invalid_role(self): + contract = self._tabular() + contract["keys"]["label"]["role"] = "bogus" + with pytest.raises(ValueError, match="role"): + validate_contract(contract) + + def test_key_shape_must_be_int_or_symbolic_name(self): + contract = self._tabular() + contract["keys"]["label"]["shape"] = [3.5] + with pytest.raises(ValueError, match="shape"): + validate_contract(contract) + + def test_key_symbolic_shape_accepted(self): + contract = self._tabular() + contract["keys"]["label"]["shape"] = ["N"] + validate_contract(contract) + + def test_invalid_normalization_owner(self): + contract = self._tabular() + contract["normalization"]["owner"] = "bogus" + with pytest.raises(ValueError, match="normalization"): + validate_contract(contract) + + def test_invalid_split_strategy(self): + contract = self._tabular() + contract["split"]["strategy"] = "bogus" + with pytest.raises(ValueError, match="strategy"): + validate_contract(contract) + + def test_group_strategy_requires_group_key(self): + contract = self._tabular() + del contract["split"]["group_key"] + with pytest.raises(ValueError, match="group_key"): + validate_contract(contract) + + def test_ratios_must_sum_to_one(self): + contract = self._tabular() + contract["split"]["ratios"] = {"train": 0.9, "val": 0.3} + with pytest.raises(ValueError, match="ratios"): + validate_contract(contract) + + def test_reconciliation_key_must_exist(self): + contract = self._tabular() + contract["reconciliation"][0]["key"] = "not_a_key" + with pytest.raises(ValueError, match="not_a_key"): + validate_contract(contract) + + def test_reconciliation_missing_field(self): + contract = self._tabular() + del contract["reconciliation"][0]["resolution"] + with pytest.raises(ValueError, match="resolution"): + validate_contract(contract) + + +class TestSaveContractRejectsInvalid: + + def test_save_contract_validates_before_writing(self, tmp_path): + contract = copy.deepcopy(TABULAR_CONTRACT) + del contract["version"] + path = tmp_path / "dataset_contract.yaml" + with pytest.raises(ValueError): + save_contract(path, contract) + assert not path.exists() diff --git a/tests/test_pipeline.py b/tests/test_pipeline.py index b9e9ded..f8c95f3 100644 --- a/tests/test_pipeline.py +++ b/tests/test_pipeline.py @@ -8,6 +8,7 @@ from dsagt.provenance import ( build_dependency_graph, + compute_pipeline_fingerprint, compute_terminal_outputs, load_pipeline_records, reconstruct_pipeline, @@ -362,6 +363,106 @@ def test_diamond_terminal_outputs(self): assert result["terminal_outputs"] == ["final.txt"] +# --------------------------------------------------------------------------- +# compute_pipeline_fingerprint +# --------------------------------------------------------------------------- + + +class TestComputePipelineFingerprint: + + def _structured(self, records): + deps = build_dependency_graph(records) + return json.loads(render_json(records, deps)) + + def test_stable_across_reruns(self): + # Same pipeline shape, different timestamps/stdout/exit code — as if + # rerun a second time — must fingerprint identically. + run1 = [ + _make_record( + "fastp", + ["fastp"], + output_files=["clean.fq"], + timestamp="2024-01-15T10:00:00Z", + ), + _make_record( + "align", + ["bwa"], + input_files=["clean.fq"], + output_files=["aligned.bam"], + record_id="r1", + timestamp="2024-01-15T10:05:00Z", + ), + ] + run2 = [ + _make_record( + "fastp", + ["fastp"], + output_files=["clean.fq"], + timestamp="2024-02-01T09:00:00Z", + return_code=0, + ), + _make_record( + "align", + ["bwa"], + input_files=["clean.fq"], + output_files=["aligned.bam"], + record_id="r1", + timestamp="2024-02-01T09:07:00Z", + ), + ] + fp1 = compute_pipeline_fingerprint(self._structured(run1)) + fp2 = compute_pipeline_fingerprint(self._structured(run2)) + assert fp1 == fp2 + assert fp1.startswith("sha256:") + + def test_changes_when_step_added(self): + base = [_make_record("fastp", ["fastp"], output_files=["clean.fq"])] + extended = base + [ + _make_record( + "align", + ["bwa"], + input_files=["clean.fq"], + output_files=["aligned.bam"], + record_id="r1", + ) + ] + fp_base = compute_pipeline_fingerprint(self._structured(base)) + fp_extended = compute_pipeline_fingerprint(self._structured(extended)) + assert fp_base != fp_extended + + def test_changes_when_step_removed(self): + full = [ + _make_record("fastp", ["fastp"], output_files=["clean.fq"]), + _make_record( + "align", + ["bwa"], + input_files=["clean.fq"], + output_files=["aligned.bam"], + record_id="r1", + ), + ] + reduced = full[:1] + fp_full = compute_pipeline_fingerprint(self._structured(full)) + fp_reduced = compute_pipeline_fingerprint(self._structured(reduced)) + assert fp_full != fp_reduced + + def test_changes_when_output_path_altered(self): + original = [_make_record("fastp", ["fastp"], output_files=["clean.fq"])] + renamed = [_make_record("fastp", ["fastp"], output_files=["clean_v2.fq"])] + fp_original = compute_pipeline_fingerprint(self._structured(original)) + fp_renamed = compute_pipeline_fingerprint(self._structured(renamed)) + assert fp_original != fp_renamed + + def test_ignores_stdout_and_exit_code(self): + records = [_make_record("fastp", ["fastp"], output_files=["clean.fq"])] + structured = self._structured(records) + structured["records"][0]["execution"]["stdout"] = "noisy rerun output" + structured["records"][0]["execution"]["return_code"] = 1 + fp_original = compute_pipeline_fingerprint(self._structured(records)) + fp_noisy = compute_pipeline_fingerprint(structured) + assert fp_original == fp_noisy + + # --------------------------------------------------------------------------- # reconstruct_pipeline (end-to-end) # ---------------------------------------------------------------------------