/
/
/
1"""
2Out-of-process probe that verifies the CPU can execute on-device ML inference.
3
4Run as a short-lived subprocess by ``helpers.util.verify_cpu_supports_ml_inference()``; the
5process exit code carries the verdict. Running the inference here, in a throwaway process,
6means a CPU that faults on the vector instructions torch emits (common on virtual machines
7that do not pass the host CPU through) crashes only this probe -- never the server.
8"""
9
10from __future__ import annotations
11
12# Exit codes interpreted by the parent. A clean "no AVX2" rejection is distinct from a
13# native crash (which the parent sees as a negative, signal-valued return code) and from
14# any other exit (which the parent treats as inconclusive and lets through).
15PROBE_CAPABLE = 0
16PROBE_NO_AVX2 = 10
17
18
19def run_probe() -> int:
20 """Run representative inference kernels and return the verdict exit code."""
21 import torch # noqa: PLC0415
22
23 if torch.backends.cpu.get_cpu_capability() in ("DEFAULT", "NO AVX"):
24 return PROBE_NO_AVX2
25
26 # Exercise the kernel families the analysis providers depend on -- a BLAS matmul, an
27 # FFT, a convolution and an fbgemm-quantized matmul (the path most likely to use an
28 # instruction a VM masks) -- on tiny tensors. The cost is the torch import, not the math.
29 torch.set_num_threads(1)
30 # Prefer fbgemm -- the x86 quantized backend the analysis models use and the one most
31 # likely to hit a masked instruction -- and fall back to qnnpack on ARM.
32 for engine in ("fbgemm", "qnnpack"):
33 if engine in torch.backends.quantized.supported_engines:
34 torch.backends.quantized.engine = engine
35 break
36 with torch.inference_mode():
37 matrix = torch.randn(64, 64)
38 torch.mm(matrix, matrix)
39 torch.stft(torch.randn(2048), n_fft=512, window=torch.hann_window(512), return_complex=True)
40 conv = torch.nn.Conv1d(1, 4, kernel_size=8)
41 conv(torch.randn(1, 1, 2048))
42 linear = torch.nn.Sequential(torch.nn.Linear(64, 64))
43 quantized = torch.ao.quantization.quantize_dynamic( # type: ignore[no-untyped-call]
44 linear, {torch.nn.Linear}, dtype=torch.qint8
45 )
46 quantized(torch.randn(8, 64))
47 return PROBE_CAPABLE
48
49
50if __name__ == "__main__":
51 raise SystemExit(run_probe())
52