"""Distributed PGSO qualification primitives.""" from __future__ import annotations import dataclassesimport argparseimport jsonimport osimport pathlibimport reimport shutilimport sysimport uuidfrom collections.abc import Mapping, Sequence from scripts.pgso.corpus import (    Corpus,    CorpusRunError,    Scenario,    load_corpus,    run_behavior_corpus,    run_corpus,)from scripts.pgso.model import (    BuildIdentity,    PgsoError,    sha256_file,    verify_identity,)from scripts.pgso.pipeline import (    GENERATION_FLAGS,    ArtifactSpec,    PipelinePaths,    apply_profile,    emit_bitcode,    link_candidate,    merge_profile_batch,    verify_candidate,)from scripts.pgso.qualify import (    BENCHMARK_PLANS,    STARTUP_COMMANDS,    BenchmarkPair,    build_profile_linked_benchmarks,    measure_heavy_workloads,    measure_startup,    measurement_payload,    profile_linked_benchmark_evidence,    relink_profile_linked_benchmarks,    require_measurements_passed,)from scripts.pgso.runner import cancellation_guard, run_checkedfrom scripts.pgso.toolchain import SUPPORTED_TARGET, Toolchain  REPO_ROOT = pathlib.Path(__file__).resolve().parents[2]DEFAULT_CORPUS = REPO_ROOT / "scripts" / "pgso" / "corpus.json"REQUIRED_BUN_VERSION = "1.3.14"REQUIRED_HYPERFINE_VERSION = "1.20.0"  @dataclasses.dataclass(frozen=True)class ShardPlan:    shard_id: str    scenario_names: tuple[str, ...]    weight: float  def plan_shards(    scenarios: Sequence[Scenario],    *,    maximum_shards: int,    prefix: str,) -> tuple[ShardPlan, ...]:    if maximum_shards <= 0:        raise PgsoError("maximum shard count must be positive")    if re.fullmatch(r"[a-z][a-z0-9-]*", prefix) is None:        raise PgsoError(f"invalid shard prefix: {prefix!r}")    if not scenarios:        raise PgsoError("cannot shard an empty scenario collection")    names = [scenario.name for scenario in scenarios]    if len(names) != len(set(names)):        raise PgsoError("cannot shard duplicate scenario names")     shard_count = min(maximum_shards, len(scenarios))    assignments: list[list[Scenario]] = [[] for _ in range(shard_count)]    weights = [0.0] * shard_count    ordered = sorted(        scenarios,        key=lambda scenario: (-scenario.timeout_seconds, scenario.name),    )    for scenario in ordered:        shard_index = min(range(shard_count), key=lambda index: (weights[index], index))        assignments[shard_index].append(scenario)        weights[shard_index] += scenario.timeout_seconds     width = max(2, len(str(shard_count)))    return tuple(        ShardPlan(            shard_id=f"{prefix}-{index + 1:0{width}d}",            scenario_names=tuple(scenario.name for scenario in assignment),            weight=weights[index],        )        for index, assignment in enumerate(assignments)    )  def matrix_payload(shards: Sequence[ShardPlan]) -> dict[str, object]:    if not shards or any(not shard.scenario_names for shard in shards):        raise PgsoError("shard matrix cannot contain empty jobs")    return {        "include": [            {                "id": shard.shard_id,                "scenarios": ",".join(shard.scenario_names),            }            for shard in shards        ]    }  def select_corpus(    corpus: Corpus,    phase: str,    scenario_names: Sequence[str],) -> Corpus:    if phase == "training":        available = corpus.scenarios    elif phase == "behavior":        available = corpus.candidate_scenarios    else:        raise PgsoError(f"unsupported corpus phase: {phase}")    if not scenario_names:        raise PgsoError(f"{phase} shard has no assigned scenarios")    if len(scenario_names) != len(set(scenario_names)):        raise PgsoError(f"duplicate assigned scenario in {phase} shard")    by_name = {scenario.name: scenario for scenario in available}    unknown = sorted(set(scenario_names) - set(by_name))    if unknown:        raise PgsoError(            f"unknown {phase} scenario: " + ", ".join(unknown)        )    requested = set(scenario_names)    selected = tuple(        scenario for scenario in available if scenario.name in requested    )    return Corpus(        repo_root=corpus.repo_root,        manifest_path=corpus.manifest_path,        manifest_sha256=corpus.manifest_sha256,        scenarios=selected,        intentional_exclusions=corpus.intentional_exclusions,        verification_scenarios=(),    )  def _mapping(value: object, label: str) -> Mapping[str, object]:    if not isinstance(value, dict) or not all(        isinstance(key, str) for key in value    ):        raise PgsoError(f"{label} must be an object")    return value  def identity_from_mapping(value: object) -> BuildIdentity:    mapping = _mapping(value, "build identity")    flags = mapping.get("generation_flags")    if not isinstance(flags, (list, tuple)) or not all(        isinstance(flag, str) for flag in flags    ):        raise PgsoError("build identity generation_flags must be strings")    fields = (        "source_sha",        "target",        "host_arch",        "zig_version",        "llvm_version",        "bitcode_sha256",        "corpus_sha256",        "update_channel",    )    values: dict[str, str] = {}    for field in fields:        item = mapping.get(field)        if not isinstance(item, str) or not item:            raise PgsoError(f"build identity {field} must be a nonempty string")        values[field] = item    return BuildIdentity(        **values,        generation_flags=tuple(flags),    )  def validate_seed_identity(identity: BuildIdentity) -> None:    if identity.target != SUPPORTED_TARGET or identity.host_arch != "arm64":        raise PgsoError("seed identity is not native macOS arm64")    if identity.update_channel != "stable":        raise PgsoError("seed identity must use the stable update channel")    if identity.generation_flags != GENERATION_FLAGS:        raise PgsoError("seed identity has unexpected profile generation flags")  def aggregate_scenario_shards(    documents: Sequence[Mapping[str, object]],    *,    phase: str,    expected_names: Sequence[str],    expected_identity: BuildIdentity,    expected_artifact_sha256: str,) -> dict[str, object]:    if not documents:        raise PgsoError(f"no {phase} shard evidence found")    expected = tuple(expected_names)    expected_set = set(expected)    if len(expected) != len(expected_set):        raise PgsoError(f"expected {phase} scenarios contain duplicates")     shard_ids: set[str] = set()    combined: dict[str, Mapping[str, object]] = {}    merged_raw_profiles = 0    for raw_document in documents:        document = _mapping(raw_document, f"{phase} shard evidence")        if document.get("schema_version") != 1:            raise PgsoError(f"unsupported {phase} shard schema")        if document.get("phase") != phase:            raise PgsoError(f"unexpected shard phase: {document.get('phase')!r}")        if document.get("status") != "passed":            raise PgsoError(f"{phase} shard did not pass")        shard_id = document.get("shard_id")        if not isinstance(shard_id, str) or not shard_id:            raise PgsoError(f"{phase} shard ID must be a nonempty string")        if shard_id in shard_ids:            raise PgsoError(f"duplicate {phase} shard ID: {shard_id}")        shard_ids.add(shard_id)         verify_identity(expected_identity, identity_from_mapping(document.get("identity")))        if document.get("artifact_sha256") != expected_artifact_sha256:            raise PgsoError(f"{phase} artifact hash mismatch: {shard_id}")         assigned = document.get("assigned_scenarios")        if not isinstance(assigned, list) or not assigned or not all(            isinstance(name, str) and name for name in assigned        ):            raise PgsoError(f"{phase} shard assignment is invalid: {shard_id}")        if len(assigned) != len(set(assigned)):            raise PgsoError(f"duplicate {phase} scenario within shard: {shard_id}")         result = _mapping(document.get("result"), f"{phase} shard result")        raw_results = result.get("results")        if not isinstance(raw_results, list):            raise PgsoError(f"{phase} shard results must be a list: {shard_id}")        result_by_name: dict[str, Mapping[str, object]] = {}        for value in raw_results:            item = _mapping(value, f"{phase} scenario result")            name = item.get("name")            if not isinstance(name, str) or not name:                raise PgsoError(f"{phase} scenario result has an invalid name")            if name in result_by_name:                raise PgsoError(f"duplicate {phase} scenario result: {name}")            if item.get("status") != "passed":                raise PgsoError(f"{phase} scenario did not pass: {name}")            result_by_name[name] = item        if set(assigned) != set(result_by_name):            raise PgsoError(f"{phase} shard result does not match assignment: {shard_id}")        if (            result.get("passed") != len(assigned)            or result.get("failed") != 0            or result.get("skipped") != 0        ):            raise PgsoError(f"{phase} shard summary is inconsistent: {shard_id}")        for name, item in result_by_name.items():            if name in combined:                raise PgsoError(f"duplicate {phase} scenario: {name}")            combined[name] = item        profile_count = result.get("merged_raw_profiles")        if isinstance(profile_count, bool) or not isinstance(profile_count, int):            raise PgsoError(f"{phase} shard profile count is invalid: {shard_id}")        merged_raw_profiles += profile_count     missing = sorted(expected_set - set(combined))    if missing:        raise PgsoError(f"missing {phase} scenario: " + ", ".join(missing))    unexpected = sorted(set(combined) - expected_set)    if unexpected:        raise PgsoError(f"unexpected {phase} scenario: " + ", ".join(unexpected))    ordered_results = [dict(combined[name]) for name in expected]    return {        "passed": len(ordered_results),        "skipped": 0,        "failed": 0,        "merged_raw_profiles": merged_raw_profiles,        "results": ordered_results,    }  def aggregate_measurement_shards(    documents: Sequence[Mapping[str, object]],    *,    phase: str,    expected_names: Sequence[str],    expected_identity: BuildIdentity,    expected_artifact_sha256: str,    expected_profile_linkage: Mapping[str, Mapping[str, object]] | None = None,) -> list[dict[str, object]]:    if not documents:        raise PgsoError(f"no {phase} measurement evidence found")    expected = tuple(expected_names)    if len(expected) != len(set(expected)):        raise PgsoError(f"expected {phase} measurements contain duplicates")    if expected_profile_linkage is not None and set(expected_profile_linkage) != set(        expected    ):        raise PgsoError(f"expected {phase} profile linkage is incomplete")    measurements: dict[str, Mapping[str, object]] = {}    shard_ids: set[str] = set()    for raw_document in documents:        document = _mapping(raw_document, f"{phase} measurement evidence")        if document.get("schema_version") != 1:            raise PgsoError(f"unsupported {phase} measurement schema")        if document.get("phase") != phase:            raise PgsoError(f"unexpected measurement phase: {document.get('phase')!r}")        if document.get("status") != "passed":            raise PgsoError(f"{phase} measurement shard did not pass")        shard_id = document.get("shard_id")        if not isinstance(shard_id, str) or not shard_id:            raise PgsoError(f"{phase} measurement shard ID is invalid")        if shard_id in shard_ids:            raise PgsoError(f"duplicate {phase} measurement shard ID: {shard_id}")        shard_ids.add(shard_id)        verify_identity(expected_identity, identity_from_mapping(document.get("identity")))        if document.get("artifact_sha256") != expected_artifact_sha256:            raise PgsoError(f"{phase} measurement artifact hash mismatch: {shard_id}")        measurement = _mapping(            document.get("measurement"),            f"{phase} measurement",        )        name = measurement.get("name")        if not isinstance(name, str) or not name:            raise PgsoError(f"{phase} measurement name is invalid")        if name in measurements:            raise PgsoError(f"duplicate {phase} measurement: {name}")        if measurement.get("passed") is not True:            raise PgsoError(f"{phase} measurement failed: {name}")        required_samples = 100 if phase == "startup" else 50        if measurement.get("requested_samples") != required_samples:            raise PgsoError(                f"{phase} measurement requires exactly {required_samples} samples: {name}"            )        for sample_label in ("control_samples", "candidate_samples"):            samples = measurement.get(sample_label)            if not isinstance(samples, list) or len(samples) != required_samples:                raise PgsoError(                    f"{phase} measurement requires exactly {required_samples} samples: {name}"                )        if (            measurement.get("control_failures") != 0            or measurement.get("candidate_failures") != 0            or measurement.get("errors") != []        ):            raise PgsoError(f"{phase} measurement contains execution failures: {name}")        comparison = _mapping(            measurement.get("comparison"),            f"{phase} measurement comparison",        )        if comparison.get("passed") is not True:            raise PgsoError(f"{phase} measurement comparison failed: {name}")        if expected_profile_linkage is not None:            linkage = _mapping(                document.get("profile_linkage"),                f"{phase} profile linkage",            )            if dict(linkage) != dict(expected_profile_linkage[name]):                raise PgsoError(f"{phase} profile linkage mismatch: {name}")        measurements[name] = measurement     missing = sorted(set(expected) - set(measurements))    if missing:        raise PgsoError(f"missing {phase} measurement: " + ", ".join(missing))    unexpected = sorted(set(measurements) - set(expected))    if unexpected:        raise PgsoError(f"unexpected {phase} measurement: " + ", ".join(unexpected))    return [dict(measurements[name]) for name in expected]  def _write_json(path: pathlib.Path, payload: Mapping[str, object]) -> None:    path.parent.mkdir(parents=True, exist_ok=True)    temporary = path.with_name(f".{path.name}.{uuid.uuid4().hex}.tmp")    temporary.write_text(        json.dumps(payload, indent=2, sort_keys=True) + "\n",        encoding="utf-8",    )    os.replace(temporary, path)  def _read_json(path: pathlib.Path, label: str) -> Mapping[str, object]:    if not path.is_file():        raise PgsoError(f"{label} does not exist: {path}")    try:        payload = json.loads(path.read_text(encoding="utf-8"))    except (OSError, json.JSONDecodeError) as error:        raise PgsoError(f"could not read {label}: {error}") from error    return _mapping(payload, label)  def _fresh_directory(path: pathlib.Path) -> pathlib.Path:    path = path.resolve()    if path.exists():        if not path.is_dir():            raise PgsoError(f"output is not a directory: {path}")        if any(path.iterdir()):            raise PgsoError(f"output directory is not empty: {path}")    path.mkdir(parents=True, exist_ok=True)    return path  def _parse_names(value: str) -> tuple[str, ...]:    names = tuple(item for item in value.split(",") if item)    if not names or ",".join(names) != value:        raise PgsoError("scenario assignment must be comma-separated names")    if any(re.fullmatch(r"[a-z0-9][a-z0-9-]*", name) is None for name in names):        raise PgsoError("scenario assignment contains an invalid name")    return names  def _manifest_evidence(    manifest: Mapping[str, object],    *,    stage: str,    label: str,) -> Mapping[str, object]:    if manifest.get("schema_version") != 1:        raise PgsoError(f"unsupported {label} schema")    if manifest.get("status") != "passed" or manifest.get("stage") != stage:        raise PgsoError(f"{label} is not a passed {stage} artifact")    return _mapping(manifest.get("evidence"), f"{label} evidence")  def _require_hash(path: pathlib.Path, expected: object, label: str) -> str:    if not isinstance(expected, str) or len(expected) != 64:        raise PgsoError(f"{label} has an invalid expected hash")    if not path.is_file() or path.stat().st_size == 0:        raise PgsoError(f"{label} is missing or empty: {path}")    actual = sha256_file(path)    if actual != expected:        raise PgsoError(f"{label} hash mismatch: expected {expected}, got {actual}")    return actual  def _checkout_sha(log_dir: pathlib.Path) -> str:    result = run_checked(        ("git", "rev-parse", "HEAD"),        cwd=REPO_ROOT,        env=os.environ.copy(),        timeout_s=60,        log_path=log_dir / "source-sha.json",        require_empty_stderr=True,    )    return result.stdout.strip()  def _validate_checkout(    identity: BuildIdentity,    corpus: Corpus,    log_dir: pathlib.Path,) -> None:    if _checkout_sha(log_dir) != identity.source_sha:        raise PgsoError("distributed job source SHA does not match seed identity")    if corpus.manifest_sha256 != identity.corpus_sha256:        raise PgsoError("distributed job corpus hash does not match seed identity")    status = run_checked(        ("git", "status", "--porcelain", "--untracked-files=no"),        cwd=REPO_ROOT,        env=os.environ.copy(),        timeout_s=60,        log_path=log_dir / "source-status.json",        require_empty_stderr=True,    )    if status.stdout:        raise PgsoError("distributed job has tracked source changes")  def _validate_toolchain(identity: BuildIdentity, toolchain: Toolchain) -> None:    actual = dataclasses.replace(        identity,        target=toolchain.target,        host_arch=toolchain.host_arch,        zig_version=toolchain.zig_version,        llvm_version=toolchain.llvm_version,    )    verify_identity(identity, actual)  def _require_version(    command: str,    expected: str,    label: str,    log: pathlib.Path,) -> pathlib.Path:    resolved = shutil.which(command)    if resolved is None:

Lines 1–500 of 1282. Load more