Design: module splits, a per-role contract, and an import-cycle fix for the 2,000-line gate

[kaijunli-infr]

Summary

Frontier's layering is sound (events drive state, schedulers are stacked global → cluster → replica → stage, prediction and communication are pluggable). The strain is in three places: a handful of modules that grew past what one file can hold, a role-dispatch pattern that spreads every architecture decision across dozens of files, and an import cycle in the core package that the code routes around with hundreds of deferred imports.

This issue proposes one architectural change, a per-role contract that owns the lane, transfer, and metric identity currently described only in prose in AGENTS.md, and treats everything else as mechanical splitting behind a numeric safety net. Nothing here proposes changing simulation results.

AGENTS.md now caps critical modules at 2,000 lines and asks for a cleanup pass, then a documented split analysis, for anything still over the limit. This issue is intended to serve as that analysis for the nine modules that breach the cap. I am opening it as a design discussion first, per the naming guidance in the Development Gates, rather than sending PRs.

All numbers below were measured on main at 1f694f7 (2026-09-22), excluding the collective-sim submodule.

Findings

F1. Nine modules breach the 2,000-line gate, and cleanup alone will not close the gap

An AST scan for unreferenced private methods in the five largest classes finds under 150 removable lines in total, so each of these needs a functional split rather than a trim.

Module	Lines	What makes it large
execution_time_predictor/sklearn_execution_time_predictor.py	8,263	one class, 196 methods
config/config.py	5,720	38 dataclasses; ClusterConfig alone 2,593 lines
metrics/metrics_store.py	5,586	80 methods; _on_request_end 581 lines, _write_system_metrics 505
scheduler/replica_scheduler/vllm_v1_engine_replica_scheduler.py	5,138	110 methods; base for the SJ2Q and SGLang variants
execution_time_predictor/shared_prediction_model_manager.py	4,614	training functions of 466 and 437 lines
execution_time_predictor/sklearn_moe_execution_time_predictor.py	3,539	middle of a three-level inheritance chain
execution_time_predictor/sklearn_disaggregation_execution_time_predictor.py	2,985	one 1,595-line function with 18 role branches
profiling/attention/main.py	2,157	435-line main, 351-line parse_args
entities/request.py	2,125	148 methods, 185 instance fields

Package-wide: 99 functions exceed 150 lines, 24 exceed 300. The predictor subpackage holds 21,000 lines in 15 files.

F2. Cluster roles are dispatched by scattered branching, not by an abstraction

The "vLLM Parallel Semantics and Frontier Mapping" section of AGENTS.md specifies the lane identity contract per role. No object in the code owns it; each consumer re-derives it with an if on ClusterType.

Symbol	Files	Hits
ClusterType.DECODE_FFN	51	170
ClusterType.DECODE_ATTN	42	177
ClusterType.PREFILL	43	115
ClusterType.MONOLITHIC	41	133
literal "pd-af-disaggregation"	5	160
literal "pd-disaggregation"	3	105

Densest sites: the dense predictor (93 role checks), metrics_store.py (44), request.py (26), simulator.py (24), base_cluster_scheduler.py (24). That last file is the most-changed file in the repository, touched in 183 of 539 commits.

F3. The predictor stacks roles through inheritance

SklearnDisaggregationExecutionTimePredictor → SklearnMoEExecutionTimePredictor → SklearnExecutionTimePredictor. Role behaviour is layered by subclassing rather than composed, which produced the 1,595-line function above. The four predictor files are the second through fifth most-churned files in the history.

F4. The core package has a real import cycle

frontier.config/__init__ → quantization_manager → model_config → model_architectures → operators/families → config/parallel_semantics, and importing any config submodule runs config/__init__ again. Whether import frontier.operators.families succeeds depends on what was imported first. This surfaced twice during #29 when a test file imported the matrix harness before anything else. The code currently works around it with 308 function-level from frontier... imports across 104 files.

F5. Events carry scheduler logic

replica_stage_schedule_event.handle_event is 789 lines and cluster_batch_end_event.handle_event is 531. The stated boundary is that events drive state changes and schedulers own the decisions; these two invert it and cannot be unit-tested without the event loop.

F6. God objects in entities and metrics

Request carries 185 fields and 148 methods mixing lifecycle timestamps, PD-AF transfer state, spec-decode bookkeeping, and metrics hooks. ExecutionTime.__init__ is 417 lines. MetricsStore both collects and writes every metric family.

F7. Smaller items

• scheduler/utils/ is a 10,000-line catch-all (pdaf_transfer.py alone is 1,927 lines) holding transfer, expert-parallel, tracing, validation, and diagnostics code.
• config/global_vars.py holds simulation mode, architecture, CUDA-graph settings, and the MoE flag as module-level mutable state read by 8 modules. This makes instantiating the simulator twice in one process fragile, which matters for parity harnesses.
• 35 example scripts total 5,748 lines and define the same require_bool helper in 27 of them.
• No lint, type check, pre-commit, or import-layering enforcement exists, so none of the boundaries above are protected against regression.

Target structure

A strict downward dependency order with one new element: a leaf package that owns the role contract.

types + contracts (leaf)     ClusterType, TensorParallelMode, parallel semantics, RoleContract
  ↑ config                   per-domain modules, validators, SimulationContext
  ↑ entities, operators, model_architectures
  ↑ execution_time_predictor features · training+cache · per-family predictors · role composition
  ↑ scheduler                global / cluster / replica / stage · transfer/ · expert_parallel/
  ↑ events                   thin handlers
  ↑ metrics                  collectors · writers   (depends on entities, not on scheduler)
  ↑ simulator

The role contract: one object per ClusterType answering the questions the scattered branches ask today: which stages it executes; which transfer events it emits and consumes and to which peer roles; how a request lane is identified ((replica_id, dp_id), (replica_id, ep_id), or full-stage replica_local_id=None); which parallel field supplies the physical communication domain for each semantic TP mode; which utilization meter scope applies. It encodes what AGENTS.md already specifies. Migration is one consumer at a time, each gated on the golden matrix below.

Sequenced plan

Order matters. Steps 1 and 2 make everything after them safe and cheap; the riskiest change goes last.

1. Build the safety net (low risk, tests only). A checked-in golden matrix of system_metrics.json and request_metrics.csv across dense/MoE, offline/online, the three architectures, and the feature toggles. The dummy-mode examples run in seconds, so this comfortably meets the ≥50-scenario gate in CI. Add import-linter contracts for the target layering at the same time.
2. Break the import cycle (F4; low risk, zero fidelity exposure). Move ClusterType, TensorParallelMode, and parallel_semantics into a dependency-free leaf package. Stop config/__init__.py importing the quantization manager eagerly. Add a test that imports each core module first in a fresh interpreter. Then lift the deferred imports back to module top.
3. Split config.py and remove global state (F1, F7; low risk, mechanical). Per-domain modules inside config/ with config.py as a re-exporting facade; cross-field validation moved out of __post_init__; global_vars replaced by a SimulationContext passed explicitly.
4. Introduce the role contract (F2; architectural). Define it in the leaf package, then migrate consumers in order of branching density: predictor, metrics store, Request, simulator, cluster scheduler. Small PRs, each gated on the golden matrix. This is the step I would like agreement on before writing code.
5. Decompose the metrics store (F1, F6; medium). Collectors by concern (request, batch, utilization, transfer) separated from writers by format (CSV, JSON, plots, Chrome trace). Outputs stay byte-identical under the golden matrix.
6. Thin the event handlers and repackage scheduler utilities (F5, F7; medium). Move the two large handle_event bodies into scheduler methods; rename scheduler/utils/ contents into scheduler/transfer/ and scheduler/expert_parallel/ with shims at the old paths for one release.
7. Recompose the execution-time predictor (F1, F3; high, last). Separate feature extraction, training and cache, per-operator-family prediction, and role composition; replace the subclass chain with composition keyed on operator family plus the role contract. Needs the golden matrix at full width plus a per-operator predicted-time trace diff.

Risks and gate coverage

• Numeric drift: every step gated on the golden matrix; the predictor step additionally gets a per-operator trace so a change is caught at the operator that caused it.
• Import-path breakage: splits keep a re-exporting facade at the old path for at least one release; import-linter prevents new code depending on the facades.
• PD-AF parity harness: it pins source hashes of request.py, base_cluster_scheduler.py, and global_batch_end_event.py on the Reference side, so it is unaffected by refactoring the current branch, but splits touching those files should confirm the candidate hooks still resolve.
• Naming: all names here are proposals, per the Development Gates.

Open decisions

1. Role contract as a frozen dataclass per role, or a small class hierarchy with one subclass per role? I lean dataclass unless a role needs behaviour that cannot be expressed as data.
2. Leaf package location: extend the existing frontier/types/ (18 small files today) or create frontier/contracts/?
3. Keep the SJ2Q replica scheduler variants as subclasses of the vLLM V1 scheduler, or extract the 5,138-line base into admission, KV planning, and step scheduling? Not proposed here because the variants are research surfaces, but the base file is over the gate and needs its own analysis.
4. Check the golden matrix outputs into the repository, or generate on demand from a pinned commit?

Happy to start with steps 1 and 2 as PRs if the direction looks right; they carry no fidelity risk and make every later step verifiable.

Appendix: most-changed source files

Commits	File
183	scheduler/cluster_scheduler/base_cluster_scheduler.py
47	execution_time_predictor/sklearn_moe_execution_time_predictor.py
42	execution_time_predictor/sklearn_execution_time_predictor.py
34	execution_time_predictor/sklearn_disaggregation_execution_time_predictor.py
32	execution_time_predictor/shared_prediction_model_manager.py
26	metrics/metrics_store.py
26	config/config.py
20	events/replica_stage_schedule_event.py
19	scheduler/replica_scheduler/vllm_v1_engine_replica_scheduler.py