124 lines
4.7 KiB
Python
124 lines
4.7 KiB
Python
"""Run a bounded, content-free native dependency probe for one runtime cell."""
|
|
|
|
import argparse
|
|
import importlib.metadata
|
|
import json
|
|
import platform
|
|
import sys
|
|
|
|
|
|
CELL_IDS = ("M1", "M2", "M3", "M4", "M5")
|
|
|
|
|
|
def validate_cell_request(cell, *, cuda_probe):
|
|
"""Reject device/probe combinations that could silently broaden test scope."""
|
|
device = cell.get("device")
|
|
if device not in {"cpu", "cuda"}:
|
|
raise ValueError("cell device must be cpu or cuda")
|
|
if cuda_probe and device != "cuda":
|
|
raise ValueError("CUDA probe is forbidden for a CPU cell")
|
|
if device == "cuda" and not cuda_probe:
|
|
raise ValueError("CUDA cell requires --cuda-probe")
|
|
|
|
|
|
def _run_onnx_identity(onnx, onnxruntime):
|
|
tensor = onnx.helper.make_tensor_value_info("value", onnx.TensorProto.FLOAT, [1])
|
|
graph = onnx.helper.make_graph(
|
|
[onnx.helper.make_node("Identity", ["value"], ["result"])],
|
|
"runtime_dependency_probe",
|
|
[tensor],
|
|
[onnx.helper.make_tensor_value_info("result", onnx.TensorProto.FLOAT, [1])],
|
|
)
|
|
model = onnx.helper.make_model(graph, opset_imports=[onnx.helper.make_opsetid("", 13)])
|
|
# IMPORTANT: cap the in-memory probe IR so older supported ONNX Runtime wheels can execute it.
|
|
model.ir_version = min(model.ir_version, 10)
|
|
session = onnxruntime.InferenceSession(
|
|
model.SerializeToString(),
|
|
providers=["CPUExecutionProvider"],
|
|
)
|
|
result = session.run(None, {"value": [1.25]})
|
|
return len(result) == 1 and float(result[0][0]) == 1.25
|
|
|
|
|
|
def _load_runtime_modules():
|
|
import onnx
|
|
import onnxruntime
|
|
import torch
|
|
from audio_separator.separator import Separator
|
|
|
|
return {
|
|
"torch": torch,
|
|
"onnx_version": onnx.__version__,
|
|
"onnxruntime_version": onnxruntime.__version__,
|
|
"onnx_identity": _run_onnx_identity(onnx, onnxruntime),
|
|
"onnx_providers": list(onnxruntime.get_available_providers()),
|
|
"audio_separator_version": importlib.metadata.version("audio-separator"),
|
|
"separator_class": isinstance(Separator, type),
|
|
}
|
|
|
|
|
|
def build_runtime_report(cell, *, cuda_probe, modules=None, python_version=None):
|
|
"""Build a path-free JSON-compatible report; dependency modules may be injected by tests."""
|
|
validate_cell_request(cell, cuda_probe=cuda_probe)
|
|
runtime = modules if modules is not None else _load_runtime_modules()
|
|
torch = runtime["torch"]
|
|
report = {
|
|
"schema_version": 1,
|
|
"ok": bool(
|
|
runtime["onnx_identity"]
|
|
and runtime["separator_class"]
|
|
and "CPUExecutionProvider" in runtime["onnx_providers"]
|
|
),
|
|
"cell": cell["id"],
|
|
"device": cell["device"],
|
|
"python": python_version or platform.python_version(),
|
|
"python_implementation": platform.python_implementation(),
|
|
"torch": str(torch.__version__),
|
|
"onnx": str(runtime["onnx_version"]),
|
|
"onnxruntime": str(runtime["onnxruntime_version"]),
|
|
"onnx_identity": bool(runtime["onnx_identity"]),
|
|
"onnx_providers": list(runtime["onnx_providers"]),
|
|
"audio_separator": str(runtime["audio_separator_version"]),
|
|
"separator_class": bool(runtime["separator_class"]),
|
|
}
|
|
if cuda_probe:
|
|
available = bool(torch.cuda.is_available())
|
|
report["cuda_available"] = available
|
|
report["torch_cuda"] = str(torch.version.cuda) if torch.version.cuda is not None else None
|
|
report["ok"] = report["ok"] and available
|
|
if available:
|
|
allocation = torch.ones(4, device="cuda")
|
|
report["cuda_device"] = str(torch.cuda.get_device_name(0))
|
|
report["cuda_capability"] = list(torch.cuda.get_device_capability(0))
|
|
report["cuda_allocation"] = float(allocation.sum().item()) == 4.0
|
|
report["ok"] = report["ok"] and report["cuda_allocation"]
|
|
del allocation
|
|
return report
|
|
|
|
|
|
def _parse_args(argv):
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
|
parser.add_argument("--cell", choices=CELL_IDS, required=True)
|
|
parser.add_argument("--device", choices=("cpu", "cuda"), required=True)
|
|
parser.add_argument("--cuda-probe", action="store_true")
|
|
return parser.parse_args(argv)
|
|
|
|
|
|
def main(argv=None):
|
|
args = _parse_args(argv)
|
|
try:
|
|
report = build_runtime_report(
|
|
{"id": args.cell, "device": args.device},
|
|
cuda_probe=args.cuda_probe,
|
|
)
|
|
# IMPORTANT: this is the CLI privacy boundary; dependency failures can embed private absolute paths.
|
|
except Exception as exc:
|
|
print(json.dumps({"schema_version": 1, "ok": False, "cell": args.cell, "error": type(exc).__name__}))
|
|
return 2
|
|
print(json.dumps(report, sort_keys=True))
|
|
return 0 if report["ok"] else 1
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(main())
|