From 1d9b6d91c07ca12204ff5f3f7b574740e3eaab8a Mon Sep 17 00:00:00 2001
From: dajiaohuang <mikewushuwen@outlook.com>
Date: Thu, 24 Sep 2026 07:30:24 +0800
Subject: [PATCH] Include serving modules in staged releases

---
 scripts/stage_release.py    |  5 ++++-
 tests/test_stage_release.py | 44 +++++++++++++++++++++++++++++++++++++
 2 files changed, 48 insertions(+), 1 deletion(-)
 create mode 100644 tests/test_stage_release.py

diff --git a/scripts/stage_release.py b/scripts/stage_release.py
index ec3d84b..bedae5c 100644
--- a/scripts/stage_release.py
+++ b/scripts/stage_release.py
@@ -4,7 +4,10 @@
 src, dst = sys.argv[1], sys.argv[2]; os.makedirs(f"{dst}/decider", exist_ok=True)
 for f in os.listdir(src):
     if os.path.isfile(f"{src}/{f}"): shutil.copy2(f"{src}/{f}", f"{dst}/{f}")
-for m in ["__init__", "prompt", "model", "systemone", "infer", "engine", "schema_engine", "fp8", "serve", "metrics"]:      # everything needed to run and serve; no training code
+for m in (
+    "__init__", "prompt", "model", "systemone", "infer", "batching", "prompt_fast",
+    "engine", "engine_v2", "shared_prefix", "schema_engine", "fp8", "serve", "metrics",
+):  # all internal modules needed to run and serve; no training code
     shutil.copy2(f"decider/{m}.py", f"{dst}/decider/{m}.py")
 shutil.copy2("MODEL_CARD.md", f"{dst}/README.md")
 if len(sys.argv) > 3: shutil.copy2(sys.argv[3], f"{dst}/eval_results.json")
diff --git a/tests/test_stage_release.py b/tests/test_stage_release.py
new file mode 100644
index 0000000..e8ef2f5
--- /dev/null
+++ b/tests/test_stage_release.py
@@ -0,0 +1,44 @@
+import importlib.util
+import json
+import os
+from pathlib import Path
+import subprocess
+import sys
+
+import pytest
+
+
+def test_staged_release_includes_server_module_closure(tmp_path):
+    repo = Path(__file__).resolve().parents[1]
+    model = tmp_path / "model"
+    release = tmp_path / "release"
+    model.mkdir()
+    (model / "decider_config.json").write_text(json.dumps({"version": 1}), encoding="utf-8")
+
+    staged = subprocess.run(
+        [sys.executable, str(repo / "scripts" / "stage_release.py"), str(model), str(release)],
+        cwd=repo,
+        capture_output=True,
+        text=True,
+        check=False,
+    )
+    assert staged.returncode == 0, staged.stdout + staged.stderr
+
+    required_modules = {"batching.py", "prompt_fast.py", "engine_v2.py", "shared_prefix.py"}
+    copied_modules = {path.name for path in (release / "decider").glob("*.py")}
+    assert required_modules <= copied_modules
+
+    if any(importlib.util.find_spec(package) is None for package in ("fastapi", "torch", "transformers")):
+        pytest.skip("staged server imports require fastapi, torch, and transformers")
+
+    env = os.environ.copy()
+    env["PYTHONPATH"] = os.pathsep.join((str(release), env.get("PYTHONPATH", "")))
+    imported = subprocess.run(
+        [sys.executable, "-c", "import decider.serve; import decider.engine_v2; import decider.shared_prefix"],
+        cwd=release,
+        env=env,
+        capture_output=True,
+        text=True,
+        check=False,
+    )
+    assert imported.returncode == 0, imported.stdout + imported.stderr