Parallelize independent Frontier configs

This commit is contained in:
2026-07-15 22:11:44 +08:00
parent ae933f2256
commit 5dbcd9f51c

View File

@@ -21,6 +21,7 @@ import statistics
import subprocess import subprocess
import sys import sys
import time import time
from concurrent.futures import ThreadPoolExecutor, as_completed
from dataclasses import asdict, replace from dataclasses import asdict, replace
from pathlib import Path from pathlib import Path
from typing import Any, Iterable from typing import Any, Iterable
@@ -525,7 +526,7 @@ def run_frontier_load(
config=config, config=config,
probe_dir=load_dir, probe_dir=load_dir,
run_id=f"{config.name}_{rate_key(rate)}", run_id=f"{config.name}_{rate_key(rate)}",
cache_root=args.output_root / "cache", cache_root=args.output_root / "cache" / config.name,
) )
write_json(load_dir / "command.json", command) write_json(load_dir / "command.json", command)
environment = os.environ.copy() environment = os.environ.copy()
@@ -652,15 +653,20 @@ def run_frontier(args: argparse.Namespace) -> Path:
"end_to_end_action_specific_calibration": False, "end_to_end_action_specific_calibration": False,
"real_serving_results_observed_before_freeze": False, "real_serving_results_observed_before_freeze": False,
}, },
"orchestration": {
"max_parallel_configs": args.max_parallel_configs,
"cache_isolation": "one_cache_root_per_config",
"loads_within_config": "sequential",
},
"configs": [asdict(config) | {"name": config.name} for config in configs], "configs": [asdict(config) | {"name": config.name} for config in configs],
} }
write_json(args.output_root / "run_manifest.json", run_manifest) write_json(args.output_root / "run_manifest.json", run_manifest)
config_results = []
rate_records = sorted( rate_records = sorted(
protocol["rates"].values(), key=lambda item: item["offered_request_rate"] protocol["rates"].values(), key=lambda item: item["offered_request_rate"]
) )
for config in configs:
def run_config(config: Any) -> dict[str, Any]:
loads = [] loads = []
for record in rate_records: for record in rate_records:
result = run_frontier_load( result = run_frontier_load(
@@ -686,7 +692,15 @@ def run_frontier(args: argparse.Namespace) -> Path:
) )
summary = {"config": asdict(config) | {"name": config.name}, "loads": loads} summary = {"config": asdict(config) | {"name": config.name}, "loads": loads}
write_json(args.output_root / "results" / f"{config.name}.json", summary) write_json(args.output_root / "results" / f"{config.name}.json", summary)
config_results.append(summary) return summary
results_by_name = {}
with ThreadPoolExecutor(max_workers=args.max_parallel_configs) as pool:
futures = {pool.submit(run_config, config): config for config in configs}
for future in as_completed(futures):
config = futures[future]
results_by_name[config.name] = future.result()
config_results = [results_by_name[config.name] for config in configs]
freeze = { freeze = {
"schema": FRONTIER_SCHEMA, "schema": FRONTIER_SCHEMA,
"status": "frozen_before_community_run", "status": "frozen_before_community_run",
@@ -1370,6 +1384,7 @@ def parse_args() -> argparse.Namespace:
frontier.add_argument("--tp4-capacity-artifact", type=Path, required=True) frontier.add_argument("--tp4-capacity-artifact", type=Path, required=True)
frontier.add_argument("--tp8-capacity-artifact", type=Path, required=True) frontier.add_argument("--tp8-capacity-artifact", type=Path, required=True)
frontier.add_argument("--cache-equivalence-artifact", type=Path, required=True) frontier.add_argument("--cache-equivalence-artifact", type=Path, required=True)
frontier.add_argument("--max-parallel-configs", type=int, default=1)
frontier.add_argument("--config", action="append") frontier.add_argument("--config", action="append")
community = subparsers.add_parser("prepare-community") community = subparsers.add_parser("prepare-community")