Coverage for orchestrant/smoke/checks.py: 67%

1"""Runtime smoke checks that exercise the installed ML wheels with real work."""

3from __future__ import annotations

5import base64

6import dataclasses

7from typing import TYPE_CHECKING

10if TYPE_CHECKING:

11 from collections.abc import Callable

14# Y = X + [10, 10] (opset 9), inline so the check needs no model file and no onnx package.

15_TINY_ONNX_ADD_B64 = (

16 "CAk6SwoOCgFYCgFDEgFZIgNBZGQSBHRpbnkqEQgCEAEiCAAAIEEAACBBQgFDWg8K"

17 "AVgSCgoICAESBAoCCAJiDwoBWRIKCggIARIECgIIAkIECgAQEg=="

18)

21@dataclasses.dataclass(frozen=True)

22class CheckResult:

23 """Outcome of a single wheel smoke check.

25 Attributes:

26 optional: ``True`` when a failure is a warning, not a gate failure.

27 """

29 name: str

30 ok: bool

31 detail: str

32 optional: bool = False

35def _ok(name: str, detail: str) -> CheckResult:

36 """Build a passing result."""

37 return CheckResult(name=name, ok=True, detail=detail)

40def _fail(name: str, detail: str) -> CheckResult:

41 """Build a failing result."""

42 return CheckResult(name=name, ok=False, detail=detail)

45def _err(name: str, exc: Exception) -> CheckResult:

46 """Build a failing result from an exception."""

47 return _fail(name, f"{type(exc).__name__}: {exc}")

50def _optional_fail(name: str, detail: str) -> CheckResult:

51 """Build a failing result for a secondary runtime (warning, not gate fail)."""

52 return CheckResult(name=name, ok=False, detail=detail, optional=True)

55def check_numpy() -> CheckResult:

56 """Exercise NumPy: array creation, matmul, and a reduction."""

57 name = "numpy"

58 try:

59 import numpy as np

61 a = np.arange(9, dtype=np.float64).reshape(3, 3)

62 product = a @ a

63 if product.shape != (3, 3):

64 return _fail(name, f"unexpected matmul shape {product.shape}")

65 return _ok(name, f"{np.__version__}: 3x3 matmul + reduce ok")

66 except Exception as exc:

67 return _err(name, exc)

70def check_torch() -> CheckResult:

71 """Exercise PyTorch: matmul, autograd, and a tiny module forward+backward."""

72 name = "torch"

73 try:

74 import torch # ty: ignore[unresolved-import]

75 from torch import nn # ty: ignore[unresolved-import]

77 torch.manual_seed(0)

78 x = torch.randn(4, 3, requires_grad=True)

79 layer = nn.Linear(3, 2)

80 layer(x).sum().backward()

81 if x.grad is None or tuple(x.grad.shape) != (4, 3):

82 return _fail(name, "autograd did not populate gradients")

84 ones = torch.ones(2, 2)

85 if not torch.equal(ones @ ones, torch.full((2, 2), 2.0)):

86 return _fail(name, "matmul result incorrect")

88 threads = torch.get_num_threads()

89 return _ok(

90 name, f"{torch.__version__}: linear+backward+matmul ok (threads={threads})"

91 )

92 except Exception as exc:

93 return _err(name, exc)

96def check_torch_numpy_bridge() -> CheckResult:

97 """Verify the torch<->numpy zero-copy ABI bridge round-trips."""

98 name = "torch<->numpy"

99 try:

100 import numpy as np

101 import torch # ty: ignore[unresolved-import]

102

103 original = np.array([1.0, 2.0, 3.0], dtype=np.float32)

104 restored = torch.from_numpy(original).numpy()

105 if not np.array_equal(original, restored):

106 return _fail(name, "from_numpy/.numpy round-trip mismatch")

107 return _ok(name, "from_numpy/.numpy round-trip ok")

108 except Exception as exc:

109 return _err(name, exc)

110

111

112def check_torchvision() -> CheckResult:

113 """Exercise torchvision's compiled ops via non-maximum suppression."""

114 name = "torchvision"

115 try:

116 import torch # ty: ignore[unresolved-import]

117 import torchvision # ty: ignore[unresolved-import]

118 from torchvision.ops import nms # ty: ignore[unresolved-import]

119

120 boxes = torch.tensor(

121 [[0, 0, 10, 10], [1, 1, 11, 11], [50, 50, 60, 60]],

122 dtype=torch.float32,

123 )

124 scores = torch.tensor([0.9, 0.8, 0.7])

125 # Boxes 0 and 1 overlap (IoU ~0.68) so one is suppressed; box 2 stands.

126 keep = nms(boxes, scores, iou_threshold=0.5)

127 if int(keep.numel()) != 2:

128 return _fail(name, f"nms kept {int(keep.numel())} boxes, expected 2")

129 return _ok(

130 name, f"{torchvision.__version__}: ops.nms ok (kept {keep.tolist()})"

131 )

132 except Exception as exc:

133 return _err(name, exc)

134

135

136def check_onnxruntime() -> CheckResult:

137 """Run a tiny embedded ONNX Add model through the CPU execution provider."""

138 name = "onnxruntime"

139 try:

140 import numpy as np

141 import onnxruntime as ort

142

143 providers = ort.get_available_providers()

144 if "CPUExecutionProvider" not in providers:

145 return _fail(name, f"CPUExecutionProvider missing from {providers}")

146 model = base64.b64decode(_TINY_ONNX_ADD_B64)

147 session = ort.InferenceSession(model, providers=["CPUExecutionProvider"])

148 out = session.run(None, {"X": np.array([1.0, 2.0], dtype=np.float32)})[0]

149 # run() returns a union and only an ndarray carries numbers; ravel gives a flat shape.

150 if not isinstance(out, np.ndarray):

151 return _fail(name, f"output is {type(out).__name__}, expected ndarray")

152 values = [float(v) for v in np.ravel(out)]

153 if values != [11.0, 12.0]:

154 return _fail(name, f"inference output {values} != [11.0, 12.0]")

155 return _ok(name, f"{ort.__version__}: CPU-EP inference ok {values}")

156 except Exception as exc:

157 return _err(name, exc)

158

159

160def check_onnxruntime_genai() -> CheckResult:

161 """Confirm ONNX Runtime GenAI registered its entry points; generation needs a model dir."""

162 name = "onnxruntime-genai"

163 try:

164 import onnxruntime_genai as og

165 except Exception as exc:

166 return _optional_fail(name, f"{type(exc).__name__}: {exc}")

167 version = getattr(og, "__version__", "?")

168 missing = [

169 attr

170 for attr in ("Model", "Tokenizer", "GeneratorParams")

171 if not hasattr(og, attr)

172 ]

173 if missing:

174 return _fail(name, f"{version}: missing entry points {missing}")

175 return _ok(name, f"{version}: native binding + entry points present")

176

177

178def check_tvm() -> CheckResult:

179 """Exercise TVM's runtime: an NDArray round-trip through the FFI layer."""

180 name = "tvm"

181 try:

182 import numpy as np

183 import tvm # ty: ignore[unresolved-import]

184 except Exception as exc:

185 return _optional_fail(name, f"{type(exc).__name__}: {exc}")

186 try:

187 original = np.array([1.5, 2.5, 3.5], dtype=np.float32)

188 if hasattr(tvm.runtime, "tensor"):

189 # TVM >= 0.25: the FFI split replaced tvm.nd with runtime.tensor.

190 restored = tvm.runtime.tensor(original).numpy()

191 api = "runtime.tensor"

192 else:

193 restored = tvm.nd.array(original, device=tvm.cpu(0)).numpy()

194 api = "nd.array"

195 if not np.array_equal(original, restored):

196 return _fail(name, f"{api} round-trip mismatch")

197 return _ok(name, f"{tvm.__version__}: {api} round-trip ok")

198 except Exception as exc:

199 return _err(name, exc)

200

201

202def check_pyav() -> CheckResult:

203 """Exercise PyAV with mpeg4 by name: ``h264`` may resolve to a GPU-only encoder."""

204 name = "pyav"

205 try:

206 import av # ty: ignore[unresolved-import]

207 except Exception as exc:

208 return _optional_fail(name, f"{type(exc).__name__}: {exc}")

209 try:

210 import io

211

212 buf = io.BytesIO()

213 with av.open(buf, mode="w", format="mp4") as container:

214 stream = container.add_stream("mpeg4", rate=24)

215 stream.width, stream.height, stream.pix_fmt = 64, 64, "yuv420p"

216 frame = av.VideoFrame(64, 64, "yuv420p")

217 for packet in stream.encode(frame):

218 container.mux(packet)

219 for packet in stream.encode():

220 container.mux(packet)

221 if not buf.getvalue():

222 return _fail(name, "mpeg4 encode produced no bytes")

223 return _ok(name, f"{av.__version__}: in-memory mpeg4 encode ok")

224 except Exception as exc:

225 return _err(name, exc)

226

227

228def check_iree() -> CheckResult:

229 """Exercise IREE end-to-end: compile MLIR and execute it on local-task."""

230 name = "iree"

231 try:

232 import iree.compiler.tools as compiler_tools

233 import iree.runtime as iree_runtime

234 import numpy as np

235 except Exception as exc:

236 return _optional_fail(name, f"{type(exc).__name__}: {exc}")

237 try:

238 mlir = (

239 "func.func @abs(%input : tensor<f32>) -> (tensor<f32>) {"

240 " %result = math.absf %input : tensor<f32>"

241 " return %result : tensor<f32> }"

242 )

243 vmfb = compiler_tools.compile_str(mlir, target_backends=["llvm-cpu"])

244 module = iree_runtime.load_vm_flatbuffer(vmfb, driver="local-task")

245 # A bare float dies in the VM marshaling layer with FAILED_PRECONDITION.

246 value = float(module.abs(np.asarray(-5.0, dtype=np.float32)).to_host())

247 if value != 5.0:

248 return _fail(name, f"abs(-5) returned {value}, expected 5.0")

249 version = getattr(iree_runtime, "__version__", "n/a")

250 return _ok(name, f"{version}: MLIR compile + local-task run ok (abs(-5)=5)")

251 except Exception as exc:

252 return _err(name, exc)

253

254

255def check_opencv() -> CheckResult:

256 """Exercise OpenCV: PNG and JPEG (the MJPEG stream codec) round-trips, color conversion."""

257 name = "opencv"

258 try:

259 import cv2

260 import numpy as np

261

262 img = (np.random.default_rng(0).random((16, 16, 3)) * 255).astype(np.uint8)

263 decoded = None

264 for ext, params in ((".png", []), (".jpg", [cv2.IMWRITE_JPEG_QUALITY, 80])):

265 encoded, buffer = cv2.imencode(ext, img, params)

266 if not encoded:

267 return _fail(name, f"cv2.imencode({ext}) failed")

268 decoded = cv2.imdecode(buffer, cv2.IMREAD_COLOR)

269 if decoded is None or decoded.shape != img.shape:

270 shape = None if decoded is None else decoded.shape

271 return _fail(name, f"cv2.imdecode({ext}) gave shape {shape}")

272 gray = cv2.cvtColor(decoded, cv2.COLOR_BGR2GRAY)

273 if gray.shape != (16, 16):

274 return _fail(name, f"cvtColor gave shape {gray.shape}")

275 return _ok(name, f"{cv2.__version__}: imencode/imdecode (png+jpg)/cvtColor ok")

276 except Exception as exc:

277 return _err(name, exc)

278

279

280def check_opencv_dnn() -> CheckResult:

281 """Exercise the OpenCV DNN module (protobuf-linked) without a model file."""

282 name = "opencv-dnn"

283 try:

284 import cv2

285 import numpy as np

286

287 if not hasattr(cv2, "dnn"):

288 return _optional_fail(name, "cv2.dnn module not present")

289 img = (np.random.default_rng(0).random((32, 32, 3)) * 255).astype(np.uint8)

290 blob = cv2.dnn.blobFromImage(img, size=(32, 32))

291 if blob.shape != (1, 3, 32, 32):

292 return _optional_fail(name, f"blobFromImage shape {blob.shape}")

293 return _ok(name, f"{cv2.__version__}: cv2.dnn.blobFromImage ok")

294 except Exception as exc:

295 return _optional_fail(name, f"{type(exc).__name__}: {exc}")

296

297

298def check_opencv_codecs() -> CheckResult:

299 """Round-trip TIFF/WEBP/OpenEXR; optional, since per-arch builds can drop these codecs."""

300 name = "opencv-codecs"

301 try:

302 import cv2

303 import numpy as np

304

305 rng = np.random.default_rng(0)

306 img = (rng.random((16, 16, 3)) * 255).astype(np.uint8)

307 exr_img = rng.random((16, 16, 3)).astype(np.float32)

308 ok_exts: list[str] = []

309 bad_exts: list[str] = []

310 for ext, src in ((".tif", img), (".webp", img), (".exr", exr_img)):

311 try:

312 encoded, buffer = cv2.imencode(ext, src)

313 if encoded and cv2.imdecode(buffer, cv2.IMREAD_UNCHANGED) is not None:

314 ok_exts.append(ext)

315 else:

316 bad_exts.append(ext)

317 except cv2.error:

318 bad_exts.append(ext)

319 if bad_exts:

320 return _optional_fail(name, f"ok={ok_exts} unavailable={bad_exts}")

321 return _ok(name, f"{cv2.__version__}: tiff/webp/exr round-trip ok")

322 except Exception as exc:

323 return _optional_fail(name, f"{type(exc).__name__}: {exc}")

324

325

326def check_opencv_freetype() -> CheckResult:

327 """Render text via OpenCV's freetype module (source-built freetype on riscv64)."""

328 name = "opencv-freetype"

329 try:

330 from pathlib import Path

331

332 import cv2

333 import numpy as np

334

335 if not hasattr(cv2, "freetype"):

336 return _optional_fail(name, "cv2.freetype module not present")

337 # matplotlib (a core dep) bundles DejaVuSans.ttf; scanning /opt would be slow.

338 import matplotlib as mpl

339

340 fonts = list((Path(mpl.get_data_path()) / "fonts" / "ttf").glob("*.ttf"))

341 if not fonts:

342 fonts = list(Path("/usr/share/fonts").rglob("*.ttf"))

343 if not fonts:

344 return _optional_fail(name, "no .ttf font found in image to render")

345 ft = cv2.freetype.createFreeType2()

346 ft.loadFontData(str(fonts[0]), 0)

347 canvas = np.zeros((32, 64, 3), dtype=np.uint8)

348 bottom_left_origin = True

349 ft.putText(

350 canvas,

351 "Ok",

352 (2, 24),

353 16,

354 (255, 255, 255),

355 -1,

356 cv2.LINE_AA,

357 bottom_left_origin,

358 )

359 if int(canvas.sum()) == 0:

360 return _optional_fail(name, "putText drew no pixels")

361 return _ok(name, f"cv2.freetype rendered text ok ({fonts[0].name})")

362 except Exception as exc:

363 return _optional_fail(name, f"{type(exc).__name__}: {exc}")

364

365

366def check_pillow() -> CheckResult:

367 """Exercise Pillow: create, resize, and inspect an image."""

368 name = "pillow"

369 try:

370 import PIL

371 from PIL import Image

372

373 image = Image.new("RGB", (32, 16), (128, 64, 32))

374 resized = image.resize((8, 4))

375 if resized.size != (8, 4):

376 return _fail(name, f"resize produced {resized.size}")

377 return _ok(name, f"{PIL.__version__}: new/resize ok")

378 except Exception as exc:

379 return _err(name, exc)

380

381

382def check_litert() -> CheckResult:

383 """Confirm LiteRT imports as ``ai_edge_litert`` or (riscv64 build) ``tflite_runtime``."""

384 name = "ai-edge-litert"

385 import importlib

386 import importlib.metadata

387

388 try:

389 version = importlib.metadata.version("ai-edge-litert")

390 except importlib.metadata.PackageNotFoundError:

391 version = "?"

392

393 last_exc: Exception | None = None

394 for module_name in ("ai_edge_litert", "tflite_runtime"):

395 try:

396 interpreter = importlib.import_module(f"{module_name}.interpreter")

397 except Exception as exc:

398 last_exc = exc

399 continue

400 if hasattr(interpreter, "Interpreter"):

401 return _ok(name, f"{version}: {module_name}.Interpreter present")

402 return _optional_fail(name, f"{module_name}.interpreter has no Interpreter")

403

404 if last_exc is not None:

405 return _optional_fail(name, f"{type(last_exc).__name__}: {last_exc}")

406 return _optional_fail(name, "no LiteRT interpreter module found")

407

408

409ALL_CHECKS: tuple[Callable[[], CheckResult], ...] = (

410 check_numpy,

411 check_torch,

412 check_torch_numpy_bridge,

413 check_torchvision,

414 check_onnxruntime,

415 check_onnxruntime_genai,

416 check_opencv,

417 check_opencv_dnn,

418 check_opencv_codecs,

419 check_opencv_freetype,

420 check_pillow,

421 check_pyav,

422 check_tvm,

423 check_iree,

424 check_litert,

425)

426

427

428def run_all() -> list[CheckResult]:

429 """Run every wheel smoke check in order and collect the results."""

430 return [check() for check in ALL_CHECKS]