Something went wrong. Try again.
My version of https://github.com/vercel-labs/fx
Something went wrong. Try again.
Python
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458from __future__ import annotations
import unittest
import pathlibimport dataclasses
from scripts.pgso.corpus import Corpus, Scenariofrom scripts.pgso.distributed import ( aggregate_measurement_shards, aggregate_scenario_shards, matrix_payload, plan_shards, run_plan, select_corpus, validate_seed_identity,)from scripts.pgso.model import BuildIdentity, PgsoErrorfrom scripts.pgso.pipeline import GENERATION_FLAGS
def scenario(name: str, timeout_seconds: float) -> Scenario: return Scenario( name=name, argv=("{binary}", "help"), cwd=".", env_set=(), env_unset=(), timeout_seconds=timeout_seconds, requires_tmux=False, allow_keychain=False, test_file=None, )
class WorkflowContractTests(unittest.TestCase): def test_behavior_workers_install_pinned_zig_without_llvm(self) -> None: repo_root = pathlib.Path(__file__).resolve().parents[3] action = (repo_root / ".github/actions/setup-pgso/action.yml").read_text() workflow = ( repo_root / ".github/workflows/pgso-macos-arm64.yml" ).read_text() behavior_job = workflow.split("\n behavior:\n", 1)[1].split( "\n startup:\n", 1 )[0]
self.assertIn(" zig:\n", action) self.assertIn( "if: inputs.zig == 'true' || inputs.llvm == 'true'", action, ) self.assertIn(' zig: "true"', behavior_job) self.assertNotIn(' llvm: "true"', behavior_job)
def test_heavy_workers_measure_candidate_built_profile_pairs(self) -> None: repo_root = pathlib.Path(__file__).resolve().parents[3] workflow = ( repo_root / ".github/workflows/pgso-macos-arm64.yml" ).read_text() candidate_job = workflow.split("\n candidate:\n", 1)[1].split( "\n behavior:\n", 1 )[0] heavy_job = workflow.split("\n heavy:\n", 1)[1].split( "\n aggregate:\n", 1 )[0]
self.assertIn("profiles/merged.profdata", candidate_job) self.assertIn("heavy/file_index/candidate/file-index-bench", candidate_job) self.assertIn( "heavy/approval_review/candidate/approval-review-bench", candidate_job, ) self.assertNotIn(' llvm: "true"', heavy_job) self.assertNotIn("--llvm-bin", heavy_job)
def test_measurement_matrices_are_generated_from_python_definitions(self) -> None: repo_root = pathlib.Path(__file__).resolve().parents[3] workflow = ( repo_root / ".github/workflows/pgso-macos-arm64.yml" ).read_text() startup_job = workflow.split("\n startup:\n", 1)[1].split( "\n heavy:\n", 1 )[0] heavy_job = workflow.split("\n heavy:\n", 1)[1].split( "\n aggregate:\n", 1 )[0]
self.assertIn("fromJSON(needs.seed.outputs.startup)", startup_job) self.assertIn("fromJSON(needs.seed.outputs.heavy)", heavy_job) self.assertNotIn("name: [help, version", startup_job) self.assertNotIn("- file-index-100k", heavy_job)
class ShardPlanningTests(unittest.TestCase): def test_measurement_plans_come_from_the_authoritative_definitions(self) -> None: missing_corpus = pathlib.Path("/does/not/need/to/exist.json")
self.assertEqual( { "include": [ {"name": name} for name in ("help", "version", "status", "background", "doctor", "sessions") ] }, run_plan("startup", missing_corpus, 20), ) self.assertEqual( { "include": [ {"name": name} for name in ( "file-index-100k", "ui-activity", "approval-transcript", "approval-diff", "approval-combined", "approval-payload", ) ] }, run_plan("heavy", missing_corpus, 20), )
def test_plan_assigns_every_scenario_exactly_once(self) -> None: scenarios = tuple( scenario(f"scenario-{index:02d}", float(index + 1)) for index in range(36) )
shards = plan_shards(scenarios, maximum_shards=20, prefix="training")
assigned = [name for shard in shards for name in shard.scenario_names] self.assertEqual(20, len(shards)) self.assertEqual(sorted(item.name for item in scenarios), sorted(assigned)) self.assertEqual(len(assigned), len(set(assigned))) self.assertTrue(all(shard.scenario_names for shard in shards))
def test_plan_is_deterministic_and_weight_balanced(self) -> None: scenarios = ( scenario("slow", 100.0), scenario("medium", 50.0), scenario("quick-a", 1.0), scenario("quick-b", 1.0), )
first = plan_shards(scenarios, maximum_shards=2, prefix="verify") second = plan_shards(scenarios, maximum_shards=2, prefix="verify")
self.assertEqual(first, second) self.assertEqual(("slow",), first[0].scenario_names) self.assertEqual( ("medium", "quick-a", "quick-b"), first[1].scenario_names, )
def test_plan_uses_one_nonempty_shard_per_scenario_when_below_limit(self) -> None: scenarios = (scenario("one", 1.0), scenario("two", 1.0))
shards = plan_shards(scenarios, maximum_shards=20, prefix="training")
self.assertEqual(2, len(shards)) self.assertEqual(("one",), shards[0].scenario_names) self.assertEqual(("two",), shards[1].scenario_names)
def test_matrix_payload_is_compact_and_contains_no_empty_jobs(self) -> None: shards = plan_shards( (scenario("one", 1.0), scenario("two", 2.0)), maximum_shards=20, prefix="training", )
payload = matrix_payload(shards)
self.assertEqual( { "include": [ {"id": "training-01", "scenarios": "two"}, {"id": "training-02", "scenarios": "one"}, ] }, payload, )
class CorpusSelectionTests(unittest.TestCase): def make_corpus(self) -> Corpus: return Corpus( repo_root=pathlib.Path("/tmp/repo"), manifest_path=pathlib.Path("/tmp/repo/scripts/pgso/corpus.json"), manifest_sha256="c" * 64, scenarios=(scenario("train-a", 1.0), scenario("train-b", 2.0)), intentional_exclusions=(), verification_scenarios=(scenario("verify-only", 3.0),), )
def test_select_training_rejects_names_outside_training_set(self) -> None: with self.assertRaisesRegex(PgsoError, "unknown training scenario"): select_corpus(self.make_corpus(), "training", ("verify-only",))
def test_select_behavior_can_mix_training_and_verification_scenarios(self) -> None: selected = select_corpus( self.make_corpus(), "behavior", ("verify-only", "train-a"), )
self.assertEqual( ("train-a", "verify-only"), tuple(item.name for item in selected.candidate_scenarios), )
def test_select_rejects_duplicate_assignments(self) -> None: with self.assertRaisesRegex(PgsoError, "duplicate assigned scenario"): select_corpus( self.make_corpus(), "training", ("train-a", "train-a"), )
class ShardAggregationTests(unittest.TestCase): def setUp(self) -> None: self.identity = BuildIdentity( source_sha="a" * 40, target="aarch64-macos", host_arch="arm64", zig_version="0.16.0", llvm_version="21.1.8", bitcode_sha256="b" * 64, corpus_sha256="c" * 64, update_channel="stable", generation_flags=GENERATION_FLAGS, )
def shard(self, shard_id: str, names: tuple[str, ...]) -> dict[str, object]: return { "schema_version": 1, "status": "passed", "phase": "training", "shard_id": shard_id, "identity": dataclasses.asdict(self.identity), "artifact_sha256": "d" * 64, "assigned_scenarios": list(names), "result": { "passed": len(names), "skipped": 0, "failed": 0, "merged_raw_profiles": len(names), "results": [ { "name": name, "status": "passed", "elapsed_seconds": 1.0, "raw_profiles": 1, "error": None, } for name in names ], }, }
def test_aggregate_requires_exactly_once_passing_coverage(self) -> None: result = aggregate_scenario_shards( (self.shard("training-01", ("one",)), self.shard("training-02", ("two",))), phase="training", expected_names=("one", "two"), expected_identity=self.identity, expected_artifact_sha256="d" * 64, )
self.assertEqual(2, result["passed"]) self.assertEqual(2, result["merged_raw_profiles"]) self.assertEqual(["one", "two"], [item["name"] for item in result["results"]])
def test_aggregate_rejects_duplicate_scenario_across_shards(self) -> None: with self.assertRaisesRegex(PgsoError, "duplicate training scenario"): aggregate_scenario_shards( (self.shard("training-01", ("one",)), self.shard("training-02", ("one",))), phase="training", expected_names=("one",), expected_identity=self.identity, expected_artifact_sha256="d" * 64, )
def test_aggregate_rejects_missing_scenario(self) -> None: with self.assertRaisesRegex(PgsoError, "missing training scenario: two"): aggregate_scenario_shards( (self.shard("training-01", ("one",)),), phase="training", expected_names=("one", "two"), expected_identity=self.identity, expected_artifact_sha256="d" * 64, )
def test_aggregate_rejects_identity_mismatch(self) -> None: shard = self.shard("training-01", ("one",)) identity = dict(shard["identity"]) identity["source_sha"] = "f" * 40 shard["identity"] = identity
with self.assertRaisesRegex(PgsoError, "identity mismatch: source_sha"): aggregate_scenario_shards( (shard,), phase="training", expected_names=("one",), expected_identity=self.identity, expected_artifact_sha256="d" * 64, )
def test_aggregate_rejects_artifact_hash_mismatch(self) -> None: shard = self.shard("training-01", ("one",)) shard["artifact_sha256"] = "e" * 64
with self.assertRaisesRegex(PgsoError, "training artifact hash mismatch"): aggregate_scenario_shards( (shard,), phase="training", expected_names=("one",), expected_identity=self.identity, expected_artifact_sha256="d" * 64, )
def test_aggregate_rejects_a_failed_shard_manifest(self) -> None: shard = self.shard("training-01", ("one",)) shard["status"] = "failed"
with self.assertRaisesRegex(PgsoError, "training shard did not pass"): aggregate_scenario_shards( (shard,), phase="training", expected_names=("one",), expected_identity=self.identity, expected_artifact_sha256="d" * 64, )
def test_seed_identity_requires_the_pinned_production_contract(self) -> None: validate_seed_identity(self.identity)
with self.assertRaisesRegex(PgsoError, "stable update channel"): validate_seed_identity(dataclasses.replace(self.identity, update_channel="dev")) with self.assertRaisesRegex(PgsoError, "profile generation flags"): validate_seed_identity( dataclasses.replace(self.identity, generation_flags=("--different",)) )
def measurement(self, name: str, passed: bool = True) -> dict[str, object]: return { "schema_version": 1, "status": "passed" if passed else "failed", "phase": "startup", "shard_id": name, "identity": dataclasses.asdict(self.identity), "artifact_sha256": "e" * 64, "measurement": { "name": name, "argv": ["help"], "requested_samples": 100, "control_samples": [1.0] * 100, "candidate_samples": [0.9] * 100, "control_failures": 0, "candidate_failures": 0, "errors": [], "comparison": { "control_samples": [1.0] * 100, "candidate_samples": [0.9] * 100, "control_p50": 1.0, "control_p95": 1.0, "candidate_p50": 0.9, "candidate_p95": 0.9, "p50_change": -0.1, "p95_change": -0.1, "passed": passed, }, "passed": passed, }, }
def test_measurement_aggregate_requires_all_passing_gates(self) -> None: measurements = aggregate_measurement_shards( (self.measurement("startup-help"), self.measurement("startup-doctor")), phase="startup", expected_names=("startup-help", "startup-doctor"), expected_identity=self.identity, expected_artifact_sha256="e" * 64, )
self.assertEqual( ["startup-help", "startup-doctor"], [item["name"] for item in measurements], )
def test_measurement_aggregate_rejects_failed_gate(self) -> None: with self.assertRaisesRegex(PgsoError, "startup measurement shard did not pass"): aggregate_measurement_shards( (self.measurement("startup-help", passed=False),), phase="startup", expected_names=("startup-help",), expected_identity=self.identity, expected_artifact_sha256="e" * 64, )
def test_heavy_aggregate_requires_production_profile_linkage(self) -> None: document = self.measurement("file-index-100k") document["phase"] = "heavy" measurement = dict(document["measurement"]) measurement["requested_samples"] = 50 measurement["control_samples"] = [1.0] * 50 measurement["candidate_samples"] = [0.9] * 50 document["measurement"] = measurement linkage = { "production_profile_sha256": "p" * 64, "benchmark_profile_sha256": "b" * 64, "supplement_sha256": "s" * 64, "selector": "file_index", } document["profile_linkage"] = linkage
result = aggregate_measurement_shards( (document,), phase="heavy", expected_names=("file-index-100k",), expected_identity=self.identity, expected_artifact_sha256="e" * 64, expected_profile_linkage={"file-index-100k": linkage}, )
self.assertEqual("file-index-100k", result[0]["name"])
document["profile_linkage"] = {**linkage, "supplement_sha256": "x" * 64} with self.assertRaisesRegex(PgsoError, "profile linkage mismatch"): aggregate_measurement_shards( (document,), phase="heavy", expected_names=("file-index-100k",), expected_identity=self.identity, expected_artifact_sha256="e" * 64, expected_profile_linkage={"file-index-100k": linkage}, )
def test_measurement_aggregate_rejects_truncated_samples(self) -> None: document = self.measurement("startup-help") measurement = dict(document["measurement"]) measurement["candidate_samples"] = [0.9] * 99 document["measurement"] = measurement
with self.assertRaisesRegex(PgsoError, "requires exactly 100 samples"): aggregate_measurement_shards( (document,), phase="startup", expected_names=("startup-help",), expected_identity=self.identity, expected_artifact_sha256="e" * 64, )
if __name__ == "__main__": unittest.main()