Skip to content

refactor: consolidate model dispatch into a single registry - #256

Open
Andrewxu313 wants to merge 1 commit into
mainfrom
refactor/model-dispatch-registry
Open

refactor: consolidate model dispatch into a single registry#256
Andrewxu313 wants to merge 1 commit into
mainfrom
refactor/model-dispatch-registry

Conversation

@Andrewxu313

Copy link
Copy Markdown
Contributor

Description

Collapse the duplicated if/elif model-name dispatch in get_initializer.py and get_parallel_strategy_manager.py into one MODEL_REGISTRY + resolve_model() in batchgen/model_dispatch.py. The two get_* functions become thin delegators. Behavior-preserving: identical substring/exact matching, same order, same classes; per-model imports stay lazy. Adds tests/test_model_dispatch.py (14 cases). No behavior change (pure refactor).

Motivation

The two get_* files duplicated the same name→class if/elif chain (incl. a copy-pasted KIMI_K25_BACKEND_NAME_PATTERNS), which is fragile (ordering/collision) and a maintenance hazard. Consolidating into a single registry makes model→implementation dispatch live in exactly one place, so the runtime core never branches on model names. First step of the modularization effort.

Type of Change

  • Bug fix
  • New feature
  • Breaking change
  • Documentation update
  • Refactor (no behavior change)

Checklist

  • I have read the CONTRIBUTION guide.
  • I have updated the tests (added tests/test_model_dispatch.py, 14 cases; verified e2e via GLM launch).
  • I have updated the documentation (design repo: model_architecture_spec.md S2.1).

Replace the duplicated if/elif name-matching in get_initializer.py and get_parallel_strategy_manager.py with one MODEL_REGISTRY + resolve_model() in batchgen/model_dispatch.py. Behavior-preserving: identical substring/exact matching, same order, same classes; per-model imports stay lazy. Adds tests/test_model_dispatch.py (14 cases).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Signed-off-by: Andrewxu313 <Andrewxu313@users.noreply.github.com>
@github-actions github-actions Bot added the ci:run Trigger build + GPU regression on H20 label Jun 11, 2026
@drunkcoding

Copy link
Copy Markdown
Collaborator

Review finding (import-time regression): model_dispatch.py adds a module-level from batchgen.config.model_name_utils import KIMI_K25_BACKEND_NAME_PATTERNS. That import chain transitively triggers the core_engine JIT build, so import batchgen.get_initializer (now a thin delegator through this module) hard-fails at import time in any env where the JIT build isn't warm/possible — verified in a fresh worktree (python 3.13 / torch 2.12+cu130): RuntimeError: Error building extension 'core_engine' during tests/test_model_dispatch.py collection. On main, get_initializer.py's module level is light (model imports are lazy inside branches), so this is a behavior change beyond the stated 'behavior-preserving' scope, and it makes the new test unrunnable outside the curated env.

Suggested fix: keep the KIMI patterns out of the module-level import path (inline the literal tuple, or make the import lazy inside the matcher), or guarantee batchgen.config.model_name_utils imports without touching op-builder chains.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

ci:run Trigger build + GPU regression on H20

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants