Skip to content

refactor(engine): extract InfraStack — infra assembly & lifecycle container#164

Merged
zc277584121 merged 2 commits into
zilliztech:mainfrom
code2tan:rebuild_infra
Jul 7, 2026
Merged

refactor(engine): extract InfraStack — infra assembly & lifecycle container#164
zc277584121 merged 2 commits into
zilliztech:mainfrom
code2tan:rebuild_infra

Conversation

@code2tan

@code2tan code2tan commented Jul 6, 2026

Copy link
Copy Markdown
Contributor

Overview

Moves the code that "constructs the 8 process-level infra clients" and "connects / initializes / closes storage" in startup / shutdown out of Engine and into a new InfraStack container component. Engine holds self.infra, and all handle access goes through self.infra.<handle> (eng.infra.<handle> in tests) — there are no longer any Engine-level forwarding properties.

This is the Engine infrastructure-layer extraction. It is a prerequisite for subsequent phases (refactoring, optimization, new logic): all business components depend on this set of infra clients, so first consolidating them into an independently testable container lets later components be constructed cleanly with InfraStack as a dependency.


Why it's needed

Engine is a god class of ~60 methods with 9 coupled responsibility areas. On the infrastructure side it has three specific problems:

  1. Assembly scattered across three places: the 8 infra clients (meta / milvus / artifact_cache / tx_cache / embed / converter / vlm / summary) are constructed in __init__, connected/initialized in startup, and closed in shutdown — spread across three sites, with the construction-order dependency (tx_cache must precede embed/vlm/summary because they cache through it) expressed only in implicit code ordering.

  2. Ownership is not explicit: ~50 sites in Engine access these clients directly via self.meta.execute(...) / self.milvus.upsert(...), with nothing in the code showing who owns meta or manages its lifecycle.

  3. Tests can only go E2E: infra assembly is coupled to Engine, so InfraStack cannot be constructed in isolation to verify its connect sequence, preload path, or shutdown order — the only option is to spin up the whole Engine and run the full pipeline.

After extracting InfraStack: construction + lifecycle live in one place, ~85 lines that can be read independently; every call site self.infra.meta makes ownership explicit; InfraStack is independently unit-testable (mock the 8 factories, assert the call sequence).


Changes

New engine/infra.py (82 lines)

The InfraStack class, single responsibility: construct the 8 infra clients + own the startup / shutdown lifecycle.

  • __init__: constructs in dependency order (tx_cache before embed/vlm/summary).
  • startup(*, preload_local_models=False): load_builtin()meta.connectmeta.init_schematx_cache.connectmilvus.connectmilvus.ensure_collection → optional _preload_models.
  • _preload_models(): preloads the local ONNX embedding provider (only when the provider declares it should be preloaded), so the first embed call doesn't pay the model-load latency.
  • shutdown(): closes only meta + tx_cache; milvus / artifact_cache have no lifecycle methods and are left untouched.

Engine changes (engine/engine.py)

  • __init__: 8 lines of infra construction → self.infra = InfraStack(cfg).
  • startup: original steps 1–7 (load_builtin + storage connect + optional preload) folded into await self.infra.startup(...); steps 8–11 (_build_pipeline / _gc_orphan_chunks / _recover_job_lane / job watcher) stay inline, to be moved out with PipelineSupervisor in 4b.
  • shutdown: the trailing meta.close() + tx_cache.close() folded into await self.infra.shutdown(); the preceding watcher/lane/consumer shutdown stays inline.
  • Removed _preload_startup_models (logic moved into InfraStack._preload_models).
  • Removed 8 Engine-level properties (getter+setter for meta/milvus/artifact_cache/tx_cache/embed/converter/vlm/summary) — see design decisions below. 47 sites self.<handle>self.infra.<handle>.

api/app.py

1 site: eng.meta.backendeng.infra.meta.backend (the only place outside Engine that reads an infra handle).

Tests

  • New tests/test_infra_stack.py (233 lines, 5 unit tests): construction order + tx_cache sharing, startup sequence, preload path, shutdown closes only meta+tx_cache, Engine exposes handles via eng.infra.
  • Mechanical rewrite across 20 test files, 258 sites: eng.<handle>eng.infra.<handle> (both reads and writes). This is the monkeypatch contract change — tests go from eng.milvus = _FakeMilvus() to eng.infra.milvus = _FakeMilvus().

File list

File Change
engine/infra.py new (+82)
engine/engine.py removed property block + _preload_startup_models; 47 sites rewritten
api/app.py 1 site
tests/test_infra_stack.py new (+233)
20 tests/test_engine_*.py etc. 258 sites eng.Xeng.infra.X

Design decisions

  • Handle access goes directly through self.infra.<handle>, with no Engine-level properties. An earlier version added 8 read-write properties as a compatibility shim (so the ~47 self.X + 258 eng.X sites wouldn't need to change), but that was pure forwarding boilerplate (~32 lines of zero-logic code) that hid ownership and created two access paths. Going direct with self.infra.X makes ownership visible at every call site, with zero boilerplate and a single path. The cost is ~300 mechanical replacements, verified by tests.

  • load_builtin() lives in InfraStack.startup. It is process bootstrap (registers built-in connector plugin classes; idempotent), at the same "make the process ready" level as connecting storage. This leaves the remaining lines of Engine.startup purely pipeline semantics, paving the way for 4b to move those lines wholesale into PipelineSupervisor.

  • Did not change the constructor signatures of already-shipped components. ObjectRepository / ArtifactCacheService / ConnectorFactory still take specific dependencies (meta / cfg / artifact_cache) rather than the whole InfraStack — avoiding cross-file churn and risk with no behavioral benefit. Only the future PipelineSupervisor (4b) will take InfraStack, because it needs the seven-item set (embed/milvus/tx_cache/...).

  • _gc_orphan_chunks stays on Engine. It is a startup reconcile (spanning milvus + objects), semantically a pipeline health check rather than an infrastructure connection. It moves with PipelineSupervisor in 4b; this PR leaves it alone.

  • Did not add close() for milvus / artifact_cache. The current shutdown closes only meta + tx_cache; that is the existing behavior (MilvusStore has no close method, LocalArtifactCache has no lifecycle). Adding a close is a behavior change and belongs in a separate issue; this refactor introduces none.


Behavioral equivalence

InfraStack's methods replicate the lifecycle in exactly the same order as the original Engine:

  • startup steps 1–7 sequence: load_builtinmeta.connectmeta.init_schematx_cache.connectmilvus.connectmilvus.ensure_collection → optional preload, matching one-to-one.
  • shutdown steps 4–5: meta.closetx_cache.close; milvus / artifact_cache untouched (same as before).
  • _preload_models's should_preload_on_server_start short-circuit logic is line-for-line identical to the original _preload_startup_models.
  • The 8 handles are the same object instances InfraStack constructed, read by Engine via self.infra.<handle> — identity is preserved, behavior unchanged.
  • Engine.startup / shutdown external signatures and semantics are unchanged; call sites in api/app.py and __main__.py are unchanged (only one site in app.py changes eng.meta.backendeng.infra.meta.backend, reading the same object).

Test impact

New unit tests (tests/test_infra_stack.py, 5 tests):

  1. Construction order + tx_cache shared by embed/vlm/summary.
  2. startup call sequence is exactly [load_builtin, meta.connect, meta.init_schema, tx_cache.connect, milvus.connect, milvus.ensure_collection]; _preload_models is not called when preload_local_models=False.
  3. preload_local_models=True with a provider that opts in fires preload_provider; short-circuits when should_preload is False.
  4. shutdown is exactly [meta.close, tx_cache.close]; milvus.close / artifact_cache.close are not called.
  5. After Engine(cfg), the 8 handles live on eng.infra, not on eng; eng.infra.milvus = fake can override.

Existing tests mechanically rewritten: 258 sites across 20 files, eng.<handle>eng.infra.<handle>. This is the follow-on change from the monkeypatch contract change — tests inject fakes into eng.infra.<handle> instead of eng.<handle>. No logic changes; ruff format passes throughout.

Regression result: full uv run --extra dev pytest356 passed / 11 skipped, behavior-equivalent to before the refactor.

Pre-existing failures (unrelated to this PR, confirmed by reproducing on baseline 072849a with git stash):

  • tests/test_connector_factory.py::TestValidateConfig, 2 tests (config validation not raising ValueError) — fail on the baseline.
  • tests/test_feishu_oauth.py / tests/test_enumeration_since.py, 2 collection errors — the optional connector SDK lark_oapi is not installed (documented in CLAUDE.md).

Risk & rollback

  • Risk: pure structural extraction, with no schema / protocol / data changes. The startup/shutdown sequences are preserved step-for-step, and handle identity is preserved. The main risk is a missed site among the ~300 mechanical rewrites — covered by the full test suite (356 passed).
  • Rollback: git revert 9fae00f d91d314 fully reverts the change, with no side effects.

code2t added 2 commits July 6, 2026 19:39
…cture components

This commit updates the test cases and the `Engine` class to consistently use the `infra` property for accessing infrastructure components such as `meta`, `milvus`, `artifact_cache`, and others. This change ensures that all infrastructure-related operations are properly encapsulated within the `InfraStack` and aligns with the recent addition of the `InfraStack` and its integration with the `Engine` class.
@code2tan code2tan changed the title PR: refactor(engine): extract InfraStack — infra assembly & lifecycle container refactor(engine): extract InfraStack — infra assembly & lifecycle container Jul 6, 2026
@zc277584121 zc277584121 merged commit 5981838 into zilliztech:main Jul 7, 2026
8 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants