arch(probe): add build-time CoreX API discovery — stop guessing interfaces
probe_corex_api.py runs during docker build BEFORE qwen3_5.py deployment: 1. Lists ALL .py files in base image's vllm/model_executor/models/ 2. For each corex_gdn/corex_moe/corex_fa2: import → inspect signatures 3. If import fails: AST parse the .py file directly for class/method defs 4. Checks native qwen3_5.py for corex references before we overwrite it 5. Checks .so files exist (libcorex_gdn.so etc) 6. Dumps everything to /workspace/corex_probe_result.json Next deploy's build log will show EXACTLY what the corex API looks like. Then we write real dispatch code against real signatures, not guesses.
This commit is contained in:
Claude
@@ -66,6 +66,14 @@ else
|
||||
echo "[patch_ops] WARNING: transformers/models not found"
|
||||
fi
|
||||
|
||||
# 1b. CoreX API probe — MUST run BEFORE deploying our qwen3_5.py
|
||||
# Discovers corex_gdn.py, corex_moe.py, corex_fa2.py interfaces from base image.
|
||||
# Also inspects native qwen3_5.py before we overwrite it.
|
||||
# Results go to /workspace/corex_probe_result.json + build log.
|
||||
echo "[patch_ops] Running CoreX API probe..."
|
||||
python3 ./probe_corex_api.py 2>&1
|
||||
echo "[patch_ops] CoreX probe done (see above for results)"
|
||||
|
||||
# 2. Model module — qwen3_5.py with CoreX dispatch (CCCL env_dispatch pattern).
|
||||
# Our version tries to import corex_gdn/corex_moe from the base image.
|
||||
# If they exist → uses fused CUDA kernels (10x faster).
|
||||
|
||||
206
qwen3_6_scripts/probe_corex_api.py
Normal file
206
qwen3_6_scripts/probe_corex_api.py
Normal file
@@ -0,0 +1,206 @@
|
||||
"""
|
||||
CoreX API probe — runs at Docker build time (in patch_ops.sh).
|
||||
|
||||
Discovers the real interfaces of corex_gdn.py, corex_moe.py, corex_fa2.py
|
||||
from the base Docker image. Outputs:
|
||||
1. /workspace/corex_probe_result.json — machine-readable API map
|
||||
2. stdout — human-readable summary for build log
|
||||
|
||||
This is NOT runtime code. It runs once during `docker build`.
|
||||
"""
|
||||
|
||||
import importlib
|
||||
import inspect
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
|
||||
PROBE_RESULT = {}
|
||||
|
||||
def probe_module(module_path, name):
|
||||
"""Try to import a module and extract its public API."""
|
||||
result = {"available": False, "classes": {}, "functions": {}, "error": None}
|
||||
|
||||
try:
|
||||
# Try direct import first
|
||||
mod = importlib.import_module(module_path)
|
||||
result["available"] = True
|
||||
result["file"] = getattr(mod, "__file__", "unknown")
|
||||
|
||||
for attr_name in dir(mod):
|
||||
if attr_name.startswith("_"):
|
||||
continue
|
||||
obj = getattr(mod, attr_name)
|
||||
|
||||
if inspect.isclass(obj):
|
||||
cls_info = {
|
||||
"bases": [b.__name__ for b in obj.__bases__],
|
||||
"methods": {},
|
||||
}
|
||||
for method_name in dir(obj):
|
||||
if method_name.startswith("_") and method_name != "__init__":
|
||||
continue
|
||||
method = getattr(obj, method_name, None)
|
||||
if callable(method):
|
||||
try:
|
||||
sig = str(inspect.signature(method))
|
||||
cls_info["methods"][method_name] = sig
|
||||
except (ValueError, TypeError):
|
||||
cls_info["methods"][method_name] = "(unknown)"
|
||||
result["classes"][attr_name] = cls_info
|
||||
|
||||
elif callable(obj):
|
||||
try:
|
||||
sig = str(inspect.signature(obj))
|
||||
result["functions"][attr_name] = sig
|
||||
except (ValueError, TypeError):
|
||||
result["functions"][attr_name] = "(unknown)"
|
||||
|
||||
except ImportError as e:
|
||||
result["error"] = f"ImportError: {e}"
|
||||
except Exception as e:
|
||||
result["error"] = f"{type(e).__name__}: {e}"
|
||||
|
||||
return result
|
||||
|
||||
|
||||
def probe_file_directly(filepath, name):
|
||||
"""If import fails, try to read the file and extract class/function defs."""
|
||||
result = {"available": False, "classes": {}, "functions": {}, "error": None}
|
||||
|
||||
if not os.path.exists(filepath):
|
||||
result["error"] = f"File not found: {filepath}"
|
||||
return result
|
||||
|
||||
result["available"] = True
|
||||
result["file"] = filepath
|
||||
result["size"] = os.path.getsize(filepath)
|
||||
|
||||
try:
|
||||
import ast
|
||||
with open(filepath) as f:
|
||||
tree = ast.parse(f.read())
|
||||
|
||||
for node in ast.walk(tree):
|
||||
if isinstance(node, ast.ClassDef):
|
||||
methods = {}
|
||||
for item in node.body:
|
||||
if isinstance(item, (ast.FunctionDef, ast.AsyncFunctionDef)):
|
||||
args = []
|
||||
for arg in item.args.args:
|
||||
args.append(arg.arg)
|
||||
methods[item.name] = f"({', '.join(args)})"
|
||||
result["classes"][node.name] = {
|
||||
"bases": [ast.dump(b) for b in node.bases],
|
||||
"methods": methods,
|
||||
"lineno": node.lineno,
|
||||
}
|
||||
elif isinstance(node, ast.FunctionDef) and node.col_offset == 0:
|
||||
args = [arg.arg for arg in node.args.args]
|
||||
result["functions"][node.name] = {
|
||||
"signature": f"({', '.join(args)})",
|
||||
"lineno": node.lineno,
|
||||
}
|
||||
except Exception as e:
|
||||
result["error"] = f"AST parse error: {e}"
|
||||
|
||||
return result
|
||||
|
||||
|
||||
# Probe paths
|
||||
VLLM_MODELS = None
|
||||
for p in [
|
||||
"/usr/local/corex/lib/python3/dist-packages/vllm/model_executor/models",
|
||||
"/usr/local/corex/lib64/python3/dist-packages/vllm/model_executor/models",
|
||||
]:
|
||||
if os.path.isdir(p):
|
||||
VLLM_MODELS = p
|
||||
break
|
||||
|
||||
print("=" * 70)
|
||||
print("[corex_probe] CoreX API Discovery — Build Time")
|
||||
print("=" * 70)
|
||||
|
||||
if VLLM_MODELS:
|
||||
print(f"[corex_probe] vllm models dir: {VLLM_MODELS}")
|
||||
|
||||
# List ALL .py files in models dir to find corex modules
|
||||
all_files = sorted(os.listdir(VLLM_MODELS))
|
||||
corex_files = [f for f in all_files if "corex" in f.lower()]
|
||||
print(f"[corex_probe] CoreX files found: {corex_files}")
|
||||
print(f"[corex_probe] All model files: {[f for f in all_files if f.endswith('.py')]}")
|
||||
|
||||
# Probe each corex module
|
||||
for target in ["corex_gdn", "corex_moe", "corex_fa2"]:
|
||||
filepath = os.path.join(VLLM_MODELS, f"{target}.py")
|
||||
|
||||
# Try import first
|
||||
result = probe_module(f"vllm.model_executor.models.{target}", target)
|
||||
|
||||
# If import failed, try AST parse
|
||||
if not result["available"]:
|
||||
print(f"[corex_probe] {target}: import failed ({result['error']}), trying AST...")
|
||||
result = probe_file_directly(filepath, target)
|
||||
|
||||
PROBE_RESULT[target] = result
|
||||
|
||||
if result["available"]:
|
||||
print(f"[corex_probe] {target}: FOUND at {result.get('file', filepath)}")
|
||||
if result.get("size"):
|
||||
print(f"[corex_probe] size: {result['size']} bytes")
|
||||
for cls_name, cls_info in result.get("classes", {}).items():
|
||||
print(f"[corex_probe] class {cls_name}:")
|
||||
for method_name, sig in cls_info.get("methods", {}).items():
|
||||
print(f"[corex_probe] {method_name}{sig}")
|
||||
for func_name, func_info in result.get("functions", {}).items():
|
||||
if isinstance(func_info, dict):
|
||||
print(f"[corex_probe] def {func_name}{func_info['signature']} (line {func_info['lineno']})")
|
||||
else:
|
||||
print(f"[corex_probe] def {func_name}{func_info}")
|
||||
else:
|
||||
print(f"[corex_probe] {target}: NOT AVAILABLE — {result.get('error', 'unknown')}")
|
||||
|
||||
# Also probe the native qwen3_5.py BEFORE we overwrite it
|
||||
native_qw = os.path.join(VLLM_MODELS, "qwen3_5.py")
|
||||
if os.path.exists(native_qw):
|
||||
sz = os.path.getsize(native_qw)
|
||||
print(f"[corex_probe] Native qwen3_5.py: {sz} bytes")
|
||||
# Check if it imports corex
|
||||
with open(native_qw) as f:
|
||||
content = f.read()
|
||||
for keyword in ["corex_gdn", "corex_moe", "corex_fa2", "CoreXGDN", "CoreXMoE"]:
|
||||
if keyword in content:
|
||||
print(f"[corex_probe] → references '{keyword}'")
|
||||
PROBE_RESULT["native_qwen3_5"] = {
|
||||
"size": sz,
|
||||
"line_count": content.count("\n"),
|
||||
"has_corex_gdn": "corex_gdn" in content,
|
||||
"has_corex_moe": "corex_moe" in content,
|
||||
"has_corex_fa2": "corex_fa2" in content,
|
||||
}
|
||||
else:
|
||||
print(f"[corex_probe] Native qwen3_5.py: NOT FOUND (will deploy ours)")
|
||||
PROBE_RESULT["native_qwen3_5"] = {"size": 0, "exists": False}
|
||||
|
||||
else:
|
||||
print("[corex_probe] ERROR: vllm models directory not found")
|
||||
PROBE_RESULT["error"] = "vllm models dir not found"
|
||||
|
||||
# Also check .so files
|
||||
for so_name, env_var in [
|
||||
("libcorex_gdn.so", "VLLM_COREX_GDN_LIBRARY"),
|
||||
("libcorex_moe.so", "VLLM_COREX_MOE_LIBRARY"),
|
||||
("libcorex_fa2.so", "VLLM_COREX_FA2_LIBRARY"),
|
||||
]:
|
||||
path = os.environ.get(env_var, f"/usr/local/corex/lib64/{so_name}")
|
||||
exists = os.path.exists(path)
|
||||
size = os.path.getsize(path) if exists else 0
|
||||
print(f"[corex_probe] {so_name}: {'EXISTS' if exists else 'MISSING'} ({size} bytes) at {path}")
|
||||
PROBE_RESULT[so_name] = {"exists": exists, "size": size, "path": path}
|
||||
|
||||
# Write results
|
||||
output_path = "/workspace/corex_probe_result.json"
|
||||
with open(output_path, "w") as f:
|
||||
json.dump(PROBE_RESULT, f, indent=2, default=str)
|
||||
print(f"\n[corex_probe] Results written to {output_path}")
|
||||
print("=" * 70)
|
||||
Reference in New Issue
Block a user