Parallelize independent Frontier configs
This commit is contained in:
@@ -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")
|
||||||
|
|||||||
Reference in New Issue
Block a user