Spaces:
Running
Running
File size: 2,322 Bytes
89cb5bf | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 | from __future__ import annotations
import sys
import os
from pathlib import Path
ROOT = Path(__file__).resolve().parents[1]
WORKSPACE = ROOT.parent
os.environ.setdefault("MPLCONFIGDIR", "/tmp/matplotlib")
os.environ.setdefault("XDG_CACHE_HOME", "/tmp")
sys.path.insert(0, str(ROOT))
sys.path.insert(0, str(WORKSPACE / "pp-scheduler"))
from pp_scheduler import OpType, UnifiedBigMacVPP # noqa: E402
from pp_simulator import OpDurationSpec, PipelineSimulator # noqa: E402
def main() -> None:
scheduler = UnifiedBigMacVPP(
pp_size=4,
vpp_size=2,
num_microbatches=8,
ve_forward_limit=3,
)
scheduler.generate_schedule()
simulator = PipelineSimulator.from_scheduler(scheduler)
result = simulator.simulate(
{
OpType.FORWARD: OpDurationSpec(mean=1.0, variance=0),
OpType.BACKWARD: OpDurationSpec(mean=2.0, variance=0),
OpType.VEFORWARD: OpDurationSpec(mean=0.8, variance=0.01),
OpType.VEBACKWARD: OpDurationSpec(mean=0.9, variance=0.01),
OpType.GENFORWARD: OpDurationSpec(mean=0.6, variance=0.01),
OpType.GENBACKWARD: OpDurationSpec(mean=0.7, variance=0.01),
},
seed=7,
)
simulator.validate_result(result)
print("Simulation summary:")
for key, value in result.summary.items():
print(f" {key}: {value}")
print("\nFirst 12 simulated ops:")
for op in result.ops[:12]:
print(
f" rank={op.rank} {op.label:<8} "
f"start={op.start_time:6.2f} end={op.end_time:6.2f} "
f"wait={op.wait_time:5.2f} deps={len(op.deps)}"
)
print("\nFirst op with dependency reasons:")
for op in result.ops:
if op.dep_reasons:
print(f" {op.id}")
for dep_id, reasons in op.dep_reasons.items():
print(f" <- {dep_id}: {', '.join(reasons)}")
break
trace_path = ROOT / "simulate_demo_trace.json"
result.to_chrome_trace(trace_path)
print(f"\nChrome trace written to: {trace_path}")
perfetto_trace_path = ROOT / "simulate_demo_perfetto_trace.json"
result.to_chrome_trace(perfetto_trace_path, perfetto_compat=True)
print(f"Perfetto-compatible trace written to: {perfetto_trace_path}")
if __name__ == "__main__":
main()
|