refactor(perfmodel): unify estimator construction and CLI policies [AIC-1770]

[jasonqinzhou]

Summary

The takeover implementation is published on hzhou-codex/pr13-estimator-api at 2f72e025179508341777b7b4b299f42dc50f32bd. Review the implementation against main.

This PR's code head has not advanced. GitHub still exposes b88ebcb7fc07cdf86a87dcc26737fba5d1ba22bc, and the current maintainer account receives 404 for the private source repository jasonqinzhou/aisimulate. The previously described source commit 4ad74c5dcf9ad1c544edffecc74f866cfa5e4d93 could not be read or updated. The takeover branch integrates the accessible PR changes with main through 920453e6a6cfb958eae2b2f725db9b767788b4c7; the implementation and validation below refer to that published branch, not this PR's stale Files changed view.

The change gives Rust, Python, prediction, recommendation, and native replay one complete estimator-construction contract:

• ForwardPassPerfModel::best_available(config) and the matching Python facade are the production construction entry points. The config includes immutable model/system/backend/topology identity, required worker_type, selection policy, data roots, and structured estimator_config.
• Defaults are estimation_mode: auto and fallback_policy: deny. Auto always searches op_level -> fpm_interpolation -> fpm_regression, including with deny. Deny constrains an explicitly requested estimator; allow tries that estimator first and then the remaining modes in the same global priority order.
• estimator_config preserves separate regression/correction sampling controls, rectangular bucket shapes, regression feature weights and singular-fit regularization, correction bounds/ranges, and an enable switch. Rust owns defaults and validation; outer layers pass the nested request intact. Unknown fields include their paths, and invalid native configuration does not silently become regression.
• Correction behavior remains unchanged at default settings, explicitly identified as legacy_workload. Shared role-based correction is deferred until accuracy validation covers aggregated pure-prefill, pure-decode, mixed batches, and phase transitions. The latest main regression workload-store routing from feat(fpm): split regression fits into workload buckets #230 is preserved.
• Exact per-role candidate construction happens after topology and block size are known. Core records requested/selected modes, effective backend version/policy, selected root, and resolved parameters. Saved recommendation YAML pins these values, and native replay preserves scheduler TP/attention-DP, selected-root KV sizing, CUDA-graph reservation, FPM decode limits, heterogeneous P/D hardware, and existing timing evidence.
• Unified CLI predict and recommend expose database mode, transfer policy, ordered system roots, estimator selection, and nested estimator parameters. Legacy saved selection and direct-regression fallback semantics have explicit migration adapters; newly authored auto requests remain auto after serialization.

Validation

Validated implementation: 2f72e025179508341777b7b4b299f42dc50f32bd.

• Rust library tests with Python support: 1,382 passed, 1 ignored.
• External Rust public-API fixture: 7 passed.
• Application/Sweeper/CLI/runner regression suite: 1,042 passed, 5 expected cross-repository Dynamo skips.
• Isolated installed-wheel SDK/public-facade/construction contracts: 99 passed; imports were verified to originate from the isolated wheel installation.
• Installed-wheel native FPM parity: 21 passed.
• Built and installed the native cp311-abi3-manylinux_2_34_x86_64 wheel (development profile, stripped).
• Native Qwen3-32B/H200/TP2 checks cover an empty first data root followed by the bundled root, HYBRID/conservative policy, independent estimator parameters, and both fixed and inferred capacity. Recommendation evaluated 2/2 candidates successfully; its generated prediction YAML completed 4 requests with matching virtual-time metrics. The installed-wheel round trip also passed.
• Changed Python files pass Ruff check/format; Rust formatting, diff whitespace, application test inventory, and documentation destinations pass. Documentation-link unit tests: 7 passed.
• Repository-wide pre-commit is not green: its existing root-file lint/format configuration and missing ESLint 9 flat config reproduce on untouched main ce9371d1 (the same 25 remaining lint diagnostics). Unrelated hook-generated edits were excluded; targeted checks above pass.

Boundaries and integration

• Selection is pinned at construction; queries do not silently switch estimators on runtime data-domain errors. Regression can construct before training, but offline replay/search reject a cold estimator instead of producing fabricated latency.
• Native engine construction still uses the existing Python model compiler. A fully Rust-native compiler remains an architectural follow-up; the low-level AicEngineBuilder retains its separate compiled-engine scope.
• AFD, analytical encoder pools, and fixed/polynomial timing retain their existing providers. Nondefault public estimator policies require regular language workers with default timing in every role.
• Attach the published takeover branch to the review by granting access to the original source branch or opening a replacement PR from the takeover branch. No current-head CI or approval claim is made for this PR.
• Companion ai-dynamo/dynamo#14065 is still open and conflicted at d86bb11d5f1c92e0516b9780b833eeefaa82197f. Its canonical path must migrate to the complete single-config constructor, and its released-wheel bridge must retain the old behavior.
• Merge/release order remains: refresh and merge the Dynamo companion; merge/release the canonical AISimulate API; remove the temporary compatibility bridge after Dynamo's minimum dependency includes this API. Do not merge this PR first.

Linear: https://linear.app/nvidia/issue/AIC-1770

[linear-code]

AIC-1770 [SWEEPER] Preserve estimator and performance-data controls

Review in Linear

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[tedzhouhk]

Requesting changes on the perf-model construction boundary before this CLI/Sweeper contract is merged.

AIC currently has several partially overlapping construction paths: Planner builds an ad-hoc Python config for RustForwardPassPerfModel.best_available, Replay assembles another timing/AIC config, Rust exposes AicEngineBuilder, and this PR adds a Sweeper-owned ForwardPassEstimatorSpec plus resolver. Landing another resolver now would make the CLI schema a fourth source of truth and lock in fields before the underlying model-selection contract is stable.

Please first consolidate AIC around one carefully designed, typed, public construction entry point:

ForwardPassPerfModel::best_available(
    config: ForwardPassPerfModelConfig,
    options: Option<ForwardPassPerfOptions>,
) -> Result<ForwardPassPerfModel>

This should be the only production entry point used by Replay, Sweeper, and Planner. Strict native/regression constructors can remain internal or test-oriented implementation details.

The boundary needs a deliberate split:

• Config: immutable model identity and selection policy — model/system/backend/version, parallelism, quantization, op_level versus whole-forward fpm, database source mode, transfer policy, systems roots, and explicit fallback policy.
• Options: runtime/tuning behavior — observation limits, regression buckets, correction bounds, and capacity limits.
• AIC interpolation should remain a bootstrap-data producer that calls tune_with_fpms; it should not become another peer constructor hidden behind CLI parsing.

After that interface exists, this PR should become a thin consumer:

1. Parse CLI/YAML exactly once into the canonical Config/Option types.
2. Call ForwardPassPerfModel::best_available(Config, Option) before search starts.
3. Pin the resolved identity/diagnostics/provenance into ReplaySpec and candidate results.
4. Make Replay and Planner consume the same types and entry point instead of rebuilding dictionaries or re-resolving defaults.
5. Add parity tests proving the same Config/Option selects the same model and provenance through Replay, Sweeper, and Planner.

I do agree with the goal of resolving estimator identity once before trials. The requested change is about ownership: AIC core should own model selection and validation; the CLI/Sweeper should only parse into that single core contract. Please complete that consolidation before merging this PR.

[jasonqinzhou]

Summarizing the discussion with @tedzhouhk: we intend to land this refactor, but agreed to defer merging until AISimulate nightlies are available. Making the AISimulate repository public is a prerequisite for those nightlies. We agreed on this timing on August 31, and Hongkuan reaffirmed on September 9 that this PR and the Dynamo companion were still pending and waiting for nightly.

The compatibility concern is concrete: this PR replaces estimator-construction APIs that Dynamo Planner and its wheel smoke test use (from_regression, the old config dictionary shape, and from_native). Releasing this change before migrating those callers would break Planner initialization and the wheel smoke test.

Once the nightly prerequisite is met, the required merge/release order remains:

1. Merge fix(planner): use canonical AIC estimator config [AIC-1770] dynamo#14065 first. Its temporary compatibility bridge supports the currently released AISimulate wheel.
2. Merge this PR and release the canonical AISimulate API.
3. Remove the temporary bridge after Dynamo's minimum AISimulate dependency includes the canonical API.

Nightly availability alone does not remove the dependency on the companion Dynamo PR; this PR must not merge first.

[jasonqinzhou]

The migration guide has been substantially rewritten on current main: installation, six-command mapping, runnable general/advanced examples, and explicit compatibility gaps now replace this layout. Please carry the backend-version guidance into section 6.1 of the current guide so this PR preserves that rewrite. Also correct the statement that version resolution happens once before the search: this head calls Core for each concrete candidate role after topology and block size are known, then carries the resolved config and provenance into Replay.

[jasonqinzhou]

Pushed the documentation refresh to this PR's source branch in 468fe9c346d80bfb05a9ed74263551bc32782455. The migration guide now matches main at 7e84227dc211c521066b3d72608c4467088a6de6 exactly, preserving the rewritten examples and compatibility gaps. The current guide already covers backend-version pinning; the stale pre-search-resolution paragraph is removed.

Validation: all 44 inherited code blocks are unchanged, 36 Bash blocks pass syntax checks, 5 YAML blocks parse, and 54 local Markdown links/anchors pass against current main. The merge-tree check confirms this guide no longer conflicts; 27 other files still do. Some linked newer docs also require that pending main integration, and no runtime tests were rerun for this documentation-only update.

The source branch is verified at the new commit, but GitHub is still reporting the old PR head (b88ebcb) while processing the push. Leaving this review thread open until the new head is reflected. The PR remains a draft; its existing requested-changes review and Dynamo #14065 dependency remain. The PR description now separates this validation from the earlier implementation results.

[jasonqinzhou]

The migration guide and implementation are now refreshed on the source branch at 4ad74c5dcf9ad1c544edffecc74f866cfa5e4d93. Section 5.7 has runnable predict/recommend policy examples; remaining AIC-only diagnostics and tuning gaps are explicit. Native examples and round-trip tests pass. GitHub still reports b88ebcb as the upstream PR head, so this discussion remains open pending head synchronization and current-head checks.

[jasonqinzhou]

The remaining migration work is in the unified CLI: SearchSpace and the canonical Core constructor accept database_mode, transfer_policy, and systems_paths, but the strict predict/recommend engine schemas do not expose them and recommendation lowering/artifact generation does not preserve them. A user cannot yet translate the guide’s AIC policy example into an equivalent AISimulate CLI request. I am wiring these three controls through prediction, search, and saved prediction YAML, with tests for policy preservation, and replacing the corresponding section 5.7 gap text with runnable examples. This also requires reconciling this branch with current main so the new migration guide and current runtime are tested together.

[jasonqinzhou]

Integration review while updating this branch: running the new predict/recommend policy examples exposes a missing replay conversion. Canonical timing_model.config.tp / attention_dp must also determine the scheduler topology; otherwise TP>1 is treated as TP=1 and rejected. Inferred capacity must use the system root selected by Core, rather than the first unqualified search path. I am fixing that conversion and adding a native CLI regression before publishing the update.

[jasonqinzhou]

Final integration review also identified a saved-YAML edge case: a request mixing fixed timing with an estimator-backed role could acquire global estimator controls when serializing the selected root, then fail its own CLI validation. Since these public controls deliberately require default timing for every language worker, mixed/custom timing requests must retain their existing timing-provider path without adding canonical policy pins. I am covering that compatibility case in the resolver tests.

[jasonqinzhou]

The external Rust API fixture needs one merge correction: its existing empirical-mode builder check still uses DatabaseMode, but the merged import list omitted that symbol. I am restoring the import and running the fixture through its separate Cargo manifest (it is excluded from the root workspace).

[jasonqinzhou]

Published the migration implementation on the existing source branch at 4ad74c5dcf9ad1c544edffecc74f866cfa5e4d93 (normal push; includes main through 79ceb657). The updated PR description records the final scope and validation.

The CLI policy gap, scheduler TP conversion, selected-root capacity handling, mixed/custom timing serialization, and external Rust API import findings are fixed. Local validation includes 976 Python regression tests, 1,349 Rust library tests, 21 focused FPM parity tests, and the native custom-root/search/saved-YAML/predict round trip. The full migration examples also succeeded (100 prediction requests; all 8 search candidates). Follow-up checks for the latest main trace-admission change passed.

GitHub's fork branch API confirms 4ad74c5d, but this upstream PR still reports b88ebcb. Its checks are therefore stale and I am leaving the self-review discussion open until the published head appears on the PR. No readiness/approval/merge claim is made.

Dynamo #14065 remains a separate integration blocker: it is open/conflicted and its canonical construction must carry worker_type in the config; its compatibility branch also needs current-signature validation. The dependency and merge/release order are preserved in the description.

[tedzhouhk]

Architectural suggestion: could we move model resolution and op-spec construction into Rust and expose the constructor to Python through a thin PyO3 binding? This new canonical entry point still calls compile_engine_from_request, which acquires the GIL and imports aiconfigurator_core.sdk.engine.compile_engine, so Rust callers still need CPython and the Python SDK to construct an estimator. A Rust-native construction path would give both languages one implementation and let Rust consumers run independently.

[tedzhouhk]

Is it better to rename this field to estimation_mode? With supported values auto (default), op_level, fpm_interpolation, and fpm_regression.

auto should try estimators in this priority order: op_level -> fpm_interpolation -> fpm_regression. Set fallback_policy to allow by default; when fallback is allowed, it should follow the same priority order. Interpolation and regression should be explicit estimation modes.

Please keep these names, defaults, and selection semantics consistent across the Rust config, Python facade, and serialized CLI/Replay configuration.

[tedzhouhk]

Requesting changes to the canonical API design before it is stabilized. Please use the role-bound API proposed in the inline comment: one complete construction config, explicit estimator selection/fallback, structured estimator parameters, and correction driven by the same immutable worker role as regression.

This review is anchored to GitHub's currently exposed PR head b88ebcb7fc07cdf86a87dcc26737fba5d1ba22bc; the design baseline was checked against main at 612fbcf45053f2c5bdef7a37d8fbecc1c732640a.

[tedzhouhk]

Please revise the public construction contract around the current worker-role-bound implementation before finalizing this schema.

I recommend the following API:

1. Keep one production constructor, best_available(config), in Rust and Python. The complete config should own model/system/backend/version/topology/quantization/data-selection identity, a required immutable worker_type (prefill, decode, or aggregated), and the estimator configuration below. Fold the old standalone options into this object.
2. Rename forward_model to estimation_mode, supporting auto (default), op_level, fpm_interpolation, and fpm_regression. Default fallback_policy to allow. Auto tries op_level -> fpm_interpolation -> fpm_regression; an explicit mode is attempted first, then remaining modes in that global priority order when fallback is allowed. With deny, only the first selected candidate is attempted. Invalid configuration or incompatible FPM input must surface as errors, and an untrained regression must report not-ready.
3. Add a structured estimator_config namespace. CLI/Python/Sweeper/Replay should pass it through completely; Rust should own its typed schema, defaults, and validation. Preserve the latest regression feature weights and expose supported sampling, fitting, and correction controls here. Unknown fields should report their full paths instead of being ignored.
4. Let worker_type determine the shared two-dimensional feature definition for both regression and correction: critical attention and global FFN/MoE. Keep the feature weights in one shared location. Each instance has one role-bound regression and one role-bound correction feature space; do not expose correction.by_workload. An aggregated instance handles pure prefill, pure decode, and mixed batches in that shared space.
5. Reuse sampling/bucketing infrastructure and config types, but keep regression and correction settings and sample state separate. Regression buckets balance retained samples for one fit; correction buckets supply local observed/base ratios. Expose per-axis bucket shapes, sample budgets, readiness controls, correction bounds, and applicable range settings. Correction state must also be separate for op-level versus interpolation because their baseline predictions differ.

Illustrative shape; omitted fields use Core defaults:

worker_type: aggregated
estimation_mode: auto
fallback_policy: allow
estimator_config:
  features:
    attention_kv_weight: 1.0
    prefill_attention_pair_weight: 1.0
    ffn_token_weight: 1.0
  op_level: {}
  fpm_interpolation: {}
  fpm_regression:
    sampling:
      bins_per_axis: [4, 4]
      max_observations: 64
    min_observations: 5
    fit:
      kind: standardized_nnls
      singular_ridge_scale: 1.0e-9
  correction:
    enabled: true
    sampling:
      bins_per_axis: [4, 4]
      max_observations: 64
    min_observations: 5
    factor_bounds: {min: 0.5, max: 2.0}

Keep estimates read-only. If runtime fallback is supported, explicit tuning calls should also accumulate samples for an eligible standby regression; otherwise the fallback can remain permanently cold. Return the fully resolved config and selection/readiness provenance so replay artifacts and cache identity preserve the actual settings.

The shared role-bound correction space is a proposed modeling change beyond current main, so validate aggregated pure-prefill, pure-decode, mixed, and phase-transition accuracy before replacing the existing correction stores. Please also cover Rust/Python config round trips, independent regression/correction settings, fallback ordering and cold-start behavior, and replay preservation. The Rust-owned construction implementation with a thin Python binding remains the architectural target discussed in the earlier comment.

[jasonqinzhou]

Thanks Hongkuan, I agree with the overall direction, especially the explicit estimator modes and separate regression/correction settings. Two details I'd like us to settle before stabilizing the API:

1. Fallback default: I would prefer fallback_policy: deny by default, with allow as an explicit opt-in. For simulation and recommendation, an unexpected estimator change can affect results even when construction succeeds. If we choose allow as the new default, let's treat that as an intentional behavior change and preserve the original fallback semantics when migrating existing saved configs.

2. Shared correction feature space: As you noted, this is a modeling change beyond the config refactor. Let's make accuracy validation against the existing correction stores across aggregated pure-prefill, pure-decode, mixed batches, and phase transitions a prerequisite for switching behavior. We can adopt the config structure while keeping the correction implementation unchanged until that evidence is available.

With those two points addressed, I think this gives us a stronger basis for a stable public API.

[jasonqinzhou]

Superseded by #242, which carries forward this implementation against current main with additional fixes. Closing this PR in favor of #242; remaining review, validation, and the Dynamo #14065 integration dependency continue there.