omnius 1.0.627 → 1.0.629

package/dist/scripts/audio-yamnet-tensorrt-worker.py

@@ -0,0 +1,335 @@
+#!/usr/bin/env python3
+"""Persistent CUDA/TensorRT YAMNet worker for Omnius audio classification.
+
+The process deliberately accepts only already-captured WAV paths.  It never
+opens an ALSA device and never downloads a model.  Omnius bootstrap stages and
+checksums the ONNX model and TensorRT plan before this worker is launched.
+
+Protocol: JSON Lines on stdin/stdout.  stdout is protocol-only; diagnostics go
+to stderr so callers cannot confuse logs with inference results.
+"""
+from __future__ import annotations
+
+import argparse
+import csv
+import hashlib
+import json
+import os
+import sys
+import time
+import traceback
+import wave
+
+
+def emit(payload: dict) -> None:
+    sys.stdout.write(json.dumps(payload, separators=(",", ":")) + "\n")
+    sys.stdout.flush()
+
+
+def fail(message: str) -> None:
+    emit({"type": "fatal", "error": message})
+    raise RuntimeError(message)
+
+
+def cuda_check(result, label: str):
+    """cuda-python returns either an error enum or (error enum, value)."""
+    if isinstance(result, tuple):
+        status, value = result[0], result[1:]
+    else:
+        status, value = result, ()
+    if int(status) != 0:
+        raise RuntimeError(f"CUDA {label} failed: {status}")
+    if len(value) == 0:
+        return None
+    return value[0] if len(value) == 1 else value
+
+
+class CudaPythonRuntime:
+    """Thin adapter over JetPack's cuda-python bindings."""
+    def __init__(self, cudart):
+        self.cudart = cudart
+
+    def stream_create(self):
+        return cuda_check(self.cudart.cudaStreamCreate(), "stream create")
+
+    def alloc(self, size, label):
+        return cuda_check(self.cudart.cudaMalloc(size), f"{label} allocation")
+
+    def h2d(self, device, host, stream, label):
+        cuda_check(
+            self.cudart.cudaMemcpyAsync(
+                device,
+                host.ctypes.data,
+                host.nbytes,
+                self.cudart.cudaMemcpyKind.cudaMemcpyHostToDevice,
+                stream,
+            ),
+            f"{label} copy",
+        )
+
+    def d2h(self, host, device, stream, label):
+        cuda_check(
+            self.cudart.cudaMemcpyAsync(
+                host.ctypes.data,
+                device,
+                host.nbytes,
+                self.cudart.cudaMemcpyKind.cudaMemcpyDeviceToHost,
+                stream,
+            ),
+            f"{label} copy",
+        )
+
+    def synchronize(self, stream):
+        cuda_check(self.cudart.cudaStreamSynchronize(stream), "stream synchronize")
+
+    def free(self, allocation):
+        self.cudart.cudaFree(allocation)
+
+    def device_metadata(self):
+        device = cuda_check(self.cudart.cudaGetDevice(), "get device")
+        props = cuda_check(self.cudart.cudaGetDeviceProperties(device), "get device properties")
+        name = getattr(props, "name", "CUDA device")
+        if isinstance(name, bytes):
+            name = name.split(b"\0", 1)[0].decode("utf8", "replace")
+        return str(name), f"{getattr(props, 'major', 0)}.{getattr(props, 'minor', 0)}"
+
+
+class PyCudaRuntime:
+    """JetPack also ships python3-pycuda on supported L4T images.
+
+    Keeping this fallback avoids a generic PyPI CUDA runtime wheel: both
+    backends call the CUDA libraries supplied by the installed JetPack image.
+    """
+    def __init__(self, cuda):
+        cuda.init()
+        self.cuda = cuda
+        self.device = cuda.Device(0)
+        self.context = self.device.make_context()
+
+    def stream_create(self):
+        return self.cuda.Stream()
+
+    def alloc(self, size, _label):
+        return self.cuda.mem_alloc(size)
+
+    def h2d(self, device, host, stream, _label):
+        self.cuda.memcpy_htod_async(device, host, stream)
+
+    def d2h(self, host, device, stream, _label):
+        self.cuda.memcpy_dtoh_async(host, device, stream)
+
+    def synchronize(self, stream):
+        stream.synchronize()
+
+    def free(self, allocation):
+        allocation.free()
+
+    def stream_handle(self, stream):
+        return stream.handle
+
+    def device_metadata(self):
+        major, minor = self.device.compute_capability()
+        return self.device.name(), f"{major}.{minor}"
+
+
+def read_wav(path: str, np):
+    started = time.perf_counter()
+    with wave.open(path, "rb") as reader:
+        channels = reader.getnchannels()
+        sample_width = reader.getsampwidth()
+        sample_rate = reader.getframerate()
+        frames = reader.getnframes()
+        compression = reader.getcomptype()
+        raw = reader.readframes(frames)
+    if compression != "NONE":
+        raise RuntimeError("Only uncompressed RIFF/WAV input is supported")
+    if channels != 1:
+        raise RuntimeError(f"Expected mono WAV from the caller, received {channels} channels")
+    if sample_width != 2:
+        raise RuntimeError(f"Expected PCM16 WAV from the caller, received {sample_width * 8}-bit samples")
+    if sample_rate != 16000:
+        raise RuntimeError(f"Expected 16 kHz WAV from the caller, received {sample_rate} Hz")
+    waveform = np.frombuffer(raw, dtype="<i2").astype(np.float32) / 32768.0
+    return waveform, sample_rate, (time.perf_counter() - started) * 1000
+
+
+class YAMNetTensorRT:
+    def __init__(self, engine_path: str, class_map_path: str):
+        started = time.perf_counter()
+        import numpy as np
+        import tensorrt as trt
+        try:
+            from cuda import cudart
+            cuda = CudaPythonRuntime(cudart)
+        except ImportError:
+            import pycuda.driver as pycuda
+            cuda = PyCudaRuntime(pycuda)
+
+        self.np = np
+        self.trt = trt
+        self.cuda = cuda
+        self.logger = trt.Logger(trt.Logger.ERROR)
+        with open(engine_path, "rb") as handle:
+            serialized = handle.read()
+        self.runtime = trt.Runtime(self.logger)
+        self.engine = self.runtime.deserialize_cuda_engine(serialized)
+        if self.engine is None:
+            raise RuntimeError("TensorRT could not deserialize the YAMNet engine")
+        self.context = self.engine.create_execution_context()
+        if self.context is None:
+            raise RuntimeError("TensorRT could not create a YAMNet execution context")
+        self.stream = cuda.stream_create()
+        self.input_name = None
+        self.output_names = []
+        for index in range(self.engine.num_io_tensors):
+            name = self.engine.get_tensor_name(index)
+            if self.engine.get_tensor_mode(name) == trt.TensorIOMode.INPUT:
+                self.input_name = name
+            else:
+                self.output_names.append(name)
+        if self.input_name is None or not self.output_names:
+            raise RuntimeError("YAMNet TensorRT engine did not expose input and output tensors")
+        with open(class_map_path, newline="", encoding="utf8") as handle:
+            self.classes = [row["display_name"] for row in csv.DictReader(handle)]
+        if len(self.classes) != 521:
+            raise RuntimeError(f"YAMNet class map must contain 521 classes; found {len(self.classes)}")
+        self.device_name, self.compute_capability = self._device_metadata()
+        self.model_load_ms = (time.perf_counter() - started) * 1000
+
+    def _device_metadata(self):
+        # TensorRT uses the CUDA-visible ordinal.  Omnius preflight limits that
+        # namespace to exactly the approved Jetson GPU before process launch.
+        return self.cuda.device_metadata()
+
+    def warm(self):
+        # YAMNet needs >=0.975 seconds.  This is a model-load probe, not a
+        # captured-audio classification, and it verifies the TensorRT path is
+        # actually executable before readiness is reported.
+        self.infer(self.np.zeros(15600, dtype=self.np.float32), 1)
+
+    def infer(self, waveform, top_k: int):
+        started = time.perf_counter()
+        np = self.np
+        if waveform.ndim != 1 or waveform.size < 15600:
+            # YAMNet's framing kernel requires 0.975 s; pad silence only to
+            # satisfy the model shape, without altering any supplied samples.
+            waveform = np.pad(waveform.reshape(-1), (0, max(0, 15600 - waveform.size)))
+        waveform = np.ascontiguousarray(waveform.astype(np.float32, copy=False))
+        if not self.context.set_input_shape(self.input_name, tuple(waveform.shape)):
+            raise RuntimeError(f"TensorRT rejected YAMNet input shape {tuple(waveform.shape)}")
+
+        allocations = []
+        host_outputs = {}
+        try:
+            input_bytes = waveform.nbytes
+            input_device = self.cuda.alloc(input_bytes, "input")
+            allocations.append(input_device)
+            self.context.set_tensor_address(self.input_name, int(input_device))
+            for name in self.output_names:
+                shape = tuple(self.context.get_tensor_shape(name))
+                if any(int(dim) < 0 for dim in shape):
+                    raise RuntimeError(f"TensorRT did not resolve output shape for {name}: {shape}")
+                dtype = self.trt.nptype(self.engine.get_tensor_dtype(name))
+                host = np.empty(shape, dtype=dtype)
+                device = self.cuda.alloc(host.nbytes, name)
+                allocations.append(device)
+                host_outputs[name] = (host, device)
+                self.context.set_tensor_address(name, int(device))
+            self.cuda.h2d(input_device, waveform, self.stream, "input")
+            stream_handle = self.cuda.stream_handle(self.stream) if hasattr(self.cuda, "stream_handle") else self.stream
+            if not self.context.execute_async_v3(stream_handle):
+                raise RuntimeError("TensorRT YAMNet execute_async_v3 returned false")
+            for name, (host, device) in host_outputs.items():
+                self.cuda.d2h(host, device, self.stream, name)
+            self.cuda.synchronize(self.stream)
+            # The converted Google model exposes output_0 with AudioSet scores.
+            scores = next((value[0] for name, value in host_outputs.items() if value[0].shape[-1] == 521), None)
+            if scores is None:
+                raise RuntimeError("TensorRT YAMNet engine did not produce a 521-class score tensor")
+            mean_scores = np.mean(scores, axis=0)
+            indices = np.argsort(mean_scores)[-top_k:][::-1]
+            return [
+                {"label": self.classes[int(index)], "confidence": float(mean_scores[int(index)])}
+                for index in indices
+            ], (time.perf_counter() - started) * 1000
+        finally:
+            for allocation in allocations:
+                try:
+                    self.cuda.free(allocation)
+                except Exception:
+                    pass
+
+
+def file_digest(path: str) -> str:
+    digest = hashlib.sha256()
+    with open(path, "rb") as handle:
+        while True:
+            chunk = handle.read(1024 * 1024)
+            if not chunk:
+                break
+            digest.update(chunk)
+    return "sha256:" + digest.hexdigest()
+
+
+def main() -> int:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--engine", required=True)
+    parser.add_argument("--class-map", required=True)
+    parser.add_argument("--model-digest", required=True)
+    args = parser.parse_args()
+    # CUDA_VISIBLE_DEVICES must be a single approved accelerator before any
+    # CUDA/TensorRT import. Jetson's integrated GPU is logical device 0.
+    visible = os.environ.get("CUDA_VISIBLE_DEVICES", "").strip()
+    if visible != "0":
+        fail("Audio TensorRT worker requires exactly CUDA_VISIBLE_DEVICES=0 on Jetson")
+    worker = YAMNetTensorRT(args.engine, args.class_map)
+    worker.warm()
+    emit({
+        "type": "ready",
+        "pid": os.getpid(),
+        "backend": "tensorrt-fp16",
+        "device": worker.device_name,
+        "compute_capability": worker.compute_capability,
+        "cuda_visible_devices": visible,
+        "model": "yamnet",
+        "model_digest": args.model_digest,
+        "taxonomy": "AudioSet-521",
+        "model_load_ms": round(worker.model_load_ms, 3),
+        "warmed": True,
+    })
+    for line in sys.stdin:
+        try:
+            request = json.loads(line)
+            if request.get("type") == "shutdown":
+                emit({"type": "stopped"})
+                return 0
+            if request.get("type") != "classify":
+                raise RuntimeError("unknown request type")
+            request_id = str(request.get("id") or "")
+            file_path = str(request.get("file") or "")
+            if not request_id or not file_path:
+                raise RuntimeError("classify requires id and file")
+            waveform, sample_rate, decode_ms = read_wav(file_path, worker.np)
+            classifications, inference_ms = worker.infer(waveform, max(1, min(int(request.get("top_k", 5)), 25)))
+            emit({
+                "type": "result",
+                "id": request_id,
+                "success": True,
+                "sample_rate_hz": sample_rate,
+                "duration_seconds": round(float(waveform.size) / sample_rate, 6),
+                "classifications": classifications,
+                "timings_ms": {"decode": round(decode_ms, 3), "inference": round(inference_ms, 3)},
+            })
+        except Exception as exc:  # keep the persistent worker alive per request
+            emit({
+                "type": "result",
+                "id": str(locals().get("request", {}).get("id", "")),
+                "success": False,
+                "error": str(exc),
+            })
+            print(traceback.format_exc(), file=sys.stderr, flush=True)
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())