roboflow--rf-detr/tests/export/test_tensorrt_export.py

# ------------------------------------------------------------------------
# RF-DETR
# Copyright (c) 2025 Roboflow. All Rights Reserved.
# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
# ------------------------------------------------------------------------
"""Tests for TensorRT export helpers."""

import argparse
import subprocess

import pytest

from rfdetr.export import _tensorrt as tensorrt_export


def test_run_command_shell_dry_run_handles_missing_cuda_visible_devices(monkeypatch) -> None:
    """Dry-run logging should not crash when CUDA_VISIBLE_DEVICES is unset."""
    monkeypatch.delenv("CUDA_VISIBLE_DEVICES", raising=False)

    logged_messages = []
    monkeypatch.setattr(tensorrt_export.logger, "info", logged_messages.append)

    result = tensorrt_export.run_command_shell(["trtexec", "--help"], dry_run=True)

    assert result.returncode == 0
    assert any("CUDA_VISIBLE_DEVICES=" in message for message in logged_messages)


def test_run_command_shell_uses_list_not_string(monkeypatch) -> None:
    """subprocess.run must be called with a list (shell=False) to prevent injection."""
    captured = {}

    def _fake_run(command, shell, capture_output, text, check):
        captured["command"] = command
        captured["shell"] = shell
        return subprocess.CompletedProcess(command, 0, stdout="", stderr="")

    monkeypatch.setattr(tensorrt_export.subprocess, "run", _fake_run)

    tensorrt_export.run_command_shell(["trtexec", "--onnx=/some/model.onnx"], dry_run=False)

    assert isinstance(captured["command"], list), "command must be a list, not a string"
    assert captured["shell"] is False, "shell=False is required to prevent injection"


def test_run_command_shell_dry_run_does_not_invoke_subprocess(monkeypatch) -> None:
    """Dry-run must return early without calling subprocess.run."""
    was_called = []

    def _should_not_run(*args, **kwargs):
        was_called.append(True)
        return subprocess.CompletedProcess([], 0)

    monkeypatch.setattr(tensorrt_export.subprocess, "run", _should_not_run)
    monkeypatch.setattr(tensorrt_export.logger, "info", lambda _: None)

    result = tensorrt_export.run_command_shell(["trtexec", "--help"], dry_run=True)

    assert not was_called, "subprocess.run must not be called during dry_run"
    assert result.returncode == 0


def test_trtexec_returns_engine_path(monkeypatch) -> None:
    """Trtexec() must return the .engine path, not None."""
    captured_argv = []

    def _fake_run(command, **kwargs):
        captured_argv.extend(command)
        return subprocess.CompletedProcess(command, 0, stdout="", stderr="")

    monkeypatch.setattr(tensorrt_export.subprocess, "run", _fake_run)
    monkeypatch.setattr(tensorrt_export, "parse_trtexec_output", lambda _: {})

    args = argparse.Namespace(profile=False, verbose=False, dry_run=False)
    result = tensorrt_export.trtexec("/tmp/model.onnx", args)

    assert result == "/tmp/model.engine"


def test_trtexec_dry_run_returns_engine_path(monkeypatch) -> None:
    """Trtexec() with dry_run=True must still return the engine path."""
    monkeypatch.setattr(tensorrt_export.logger, "info", lambda _: None)
    monkeypatch.setattr(tensorrt_export, "parse_trtexec_output", lambda _: {})

    args = argparse.Namespace(profile=False, verbose=False, dry_run=True)
    result = tensorrt_export.trtexec("/tmp/model.onnx", args)

    assert result == "/tmp/model.engine"


@pytest.mark.parametrize(
    ("onnx_path", "expected_engine"),
    [
        pytest.param("/output/rfdetr.onnx", "/output/rfdetr.engine", id="plain-path"),
        pytest.param("/path with spaces/model.onnx", "/path with spaces/model.engine", id="path-with-spaces"),
        pytest.param("/model;rm -rf /.onnx", "/model;rm -rf /.engine", id="shell-metachar"),
    ],
)
def test_trtexec_argv_contains_no_shell_string(monkeypatch, onnx_path: str, expected_engine: str) -> None:
    """Trtexec builds an argv list; no shell string concatenation of user paths."""
    captured = {}

    def _fake_run(command, shell, **kwargs):
        captured["command"] = command
        captured["shell"] = shell
        return subprocess.CompletedProcess(command, 0, stdout="", stderr="")

    monkeypatch.setattr(tensorrt_export.subprocess, "run", _fake_run)
    monkeypatch.setattr(tensorrt_export, "parse_trtexec_output", lambda _: {})

    args = argparse.Namespace(profile=False, verbose=False, dry_run=False)
    result = tensorrt_export.trtexec(onnx_path, args)

    assert result == expected_engine
    assert isinstance(captured["command"], list), "argv must be a list"
    assert captured["shell"] is False, "shell=False required"
    # Verify the ONNX path appears as a standalone argument element (not shell-expanded)
    assert any(onnx_path in arg for arg in captured["command"])


# ---------------------------------------------------------------------------
# parse_trtexec_output (#1)
# ---------------------------------------------------------------------------

_FULL_TRTEXEC_STDOUT = """\
[I] GPU Compute Time: min = 1.23 ms, max = 4.56 ms, mean = 2.34 ms, median = 2.10 ms
[I] Host to Device Transfer Time: min = 0.10 ms, max = 0.20 ms, mean = 0.15 ms
[I] Device to Host Transfer Time: min = 0.05 ms, max = 0.08 ms, mean = 0.06 ms
[I] Latency: min = 1.40 ms, max = 4.80 ms, mean = 2.55 ms
[I] Throughput: 391.22 qps
"""

_PARTIAL_TRTEXEC_STDOUT = """\
[I] GPU Compute Time: min = 1.00 ms, max = 2.00 ms, mean = 1.50 ms, median = 1.45 ms
[I] Throughput: 100.00 qps
"""


@pytest.mark.parametrize(
    ("output_text", "expected"),
    [
        pytest.param(
            _FULL_TRTEXEC_STDOUT,
            {
                "compute_min_ms": 1.23,
                "compute_max_ms": 4.56,
                "compute_mean_ms": 2.34,
                "compute_median_ms": 2.10,
                "h2d_min_ms": 0.10,
                "h2d_max_ms": 0.20,
                "h2d_mean_ms": 0.15,
                "d2h_min_ms": 0.05,
                "d2h_max_ms": 0.08,
                "d2h_mean_ms": 0.06,
                "latency_min_ms": 1.40,
                "latency_max_ms": 4.80,
                "latency_mean_ms": 2.55,
                "throughput_qps": 391.22,
            },
            id="all-5-patterns",
        ),
        pytest.param(
            "",
            {},
            id="empty-stdout",
        ),
        pytest.param(
            _PARTIAL_TRTEXEC_STDOUT,
            {
                "compute_min_ms": 1.00,
                "compute_max_ms": 2.00,
                "compute_mean_ms": 1.50,
                "compute_median_ms": 1.45,
                "throughput_qps": 100.00,
            },
            id="partial-stdout",
        ),
    ],
)
def test_parse_trtexec_output(output_text: str, expected: dict) -> None:
    """parse_trtexec_output extracts timing statistics from trtexec stdout."""
    result = tensorrt_export.parse_trtexec_output(output_text)
    assert result == pytest.approx(expected, abs=1e-6)


# ---------------------------------------------------------------------------
# CalledProcessError logging path (#15)
# ---------------------------------------------------------------------------


def test_run_command_shell_called_process_error_is_reraised(monkeypatch) -> None:
    """CalledProcessError from subprocess.run is re-raised after logging."""
    error_messages = []

    def _fake_run(command, **kwargs):
        raise subprocess.CalledProcessError(returncode=1, cmd=["trtexec"], stderr="engine build failed")

    monkeypatch.setattr(tensorrt_export.subprocess, "run", _fake_run)
    monkeypatch.setattr(tensorrt_export.logger, "error", error_messages.append)

    with pytest.raises(subprocess.CalledProcessError):
        tensorrt_export.run_command_shell(["trtexec", "--onnx=/tmp/model.onnx"], dry_run=False)

    assert error_messages, "logger.error must be called when CalledProcessError is raised"


# ---------------------------------------------------------------------------
# profile=True argv path (#17)
# ---------------------------------------------------------------------------


def test_trtexec_profile_true_wraps_with_nsys(monkeypatch) -> None:
    """Profile=True wraps trtexec with 'nsys profile …' and the output flag is present."""
    captured_argv: list[str] = []

    def _fake_run(command, **kwargs):
        captured_argv.extend(command)
        return subprocess.CompletedProcess(command, 0, stdout="", stderr="")

    monkeypatch.setattr(tensorrt_export.subprocess, "run", _fake_run)
    monkeypatch.setattr(tensorrt_export, "parse_trtexec_output", lambda _: {})
    monkeypatch.setattr(tensorrt_export.logger, "info", lambda _: None)

    args = argparse.Namespace(profile=True, verbose=False, dry_run=False)
    tensorrt_export.trtexec("/tmp/model.onnx", args)

    assert captured_argv[0] == "nsys", "profile=True must wrap with nsys as argv[0]"
    argv_str = " ".join(captured_argv)
    assert "--output=" in argv_str, "nsys profile must include --output= flag"