Push model to docker registry

e2e/engine-artifact.Dockerfile

@@ -0,0 +1,5 @@
+FROM alpine:3.20
+
+COPY . /models/
+
+CMD ["sh"]

e2e/scripts/publish_jetson_engine.sh

@@ -0,0 +1,43 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+COMPOSE="${COMPOSE:-docker compose -f docker-compose.test.yml --profile jetson}"
+REGISTRY_HOST="${REGISTRY_HOST:?REGISTRY_HOST is required}"
+ENGINE_REPOSITORY="${JETSON_ENGINE_REPOSITORY:-$REGISTRY_HOST/azaion/detections-jetson-engine}"
+BRANCH="${CI_COMMIT_BRANCH:-local}"
+ENGINE_TAG="${JETSON_ENGINE_TAG:-$(printf '%s' "$BRANCH" | tr -c 'A-Za-z0-9_.-' '-')}"
+OUT_DIR="${JETSON_ENGINE_OUT_DIR:-results/jetson-engine}"
+
+mkdir -p "$OUT_DIR/models"
+
+loader_id="$($COMPOSE ps -q mock-loader)"
+if [[ -z "$loader_id" ]]; then
+  echo "ERROR: mock-loader container is not running"
+  exit 1
+fi
+
+docker cp "$loader_id:/models/models/." "$OUT_DIR/models/"
+find "$OUT_DIR/models" -maxdepth 1 -type f ! -name 'azaion*.engine' -delete
+
+engine_count="$(find "$OUT_DIR/models" -maxdepth 1 -type f -name 'azaion*.engine' | wc -l | tr -d ' ')"
+if [[ "$engine_count" == "0" ]]; then
+  echo "ERROR: no converted TensorRT engine found in mock-loader /models/models"
+  find "$OUT_DIR/models" -maxdepth 2 -type f -print
+  exit 1
+fi
+
+echo "--- Converted TensorRT engine files:"
+find "$OUT_DIR/models" -maxdepth 1 -type f -name 'azaion*.engine' -print -exec ls -lh {} \;
+
+image="$ENGINE_REPOSITORY:$ENGINE_TAG"
+echo "--- Building Jetson engine artifact image: $image"
+docker build -f engine-artifact.Dockerfile -t "$image" "$OUT_DIR/models"
+docker push "$image"
+
+if [[ -n "${CI_COMMIT_SHA:-}" ]]; then
+  sha_tag="$(printf '%s' "$CI_COMMIT_SHA" | cut -c1-12)"
+  docker tag "$image" "$ENGINE_REPOSITORY:$sha_tag"
+  docker push "$ENGINE_REPOSITORY:$sha_tag"
+fi
+
+echo "--- Published Jetson engine artifact image: $image"

e2e/scripts/pull_jetson_engine.sh

@@ -0,0 +1,28 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+if [[ -z "${REGISTRY_HOST:-}" ]]; then
+  echo "--- REGISTRY_HOST is not set; skipping Jetson engine artifact pull"
+  exit 0
+fi
+
+ENGINE_REPOSITORY="${JETSON_ENGINE_REPOSITORY:-$REGISTRY_HOST/azaion/detections-jetson-engine}"
+BRANCH="${CI_COMMIT_BRANCH:-local}"
+ENGINE_TAG="${JETSON_ENGINE_TAG:-$(printf '%s' "$BRANCH" | tr -c 'A-Za-z0-9_.-' '-')}"
+TARGET_DIR="${JETSON_ENGINE_TARGET_DIR:-fixtures/models}"
+image="$ENGINE_REPOSITORY:$ENGINE_TAG"
+
+echo "--- Pulling Jetson engine artifact image: $image"
+if ! docker pull "$image"; then
+  echo "--- Jetson engine artifact image not found; smoke will use ONNX fallback"
+  exit 0
+fi
+
+cid="$(docker create "$image")"
+trap 'docker rm -f "$cid" >/dev/null 2>&1 || true' EXIT
+
+mkdir -p "$TARGET_DIR"
+docker cp "$cid:/models/." "$TARGET_DIR/"
+
+echo "--- Installed Jetson engine files:"
+find "$TARGET_DIR" -maxdepth 1 -type f -name 'azaion*.engine' -print -exec ls -lh {} \;

src/engines/tensorrt_engine.pxd

@@ -4,6 +4,8 @@ from engines.inference_engine cimport InferenceEngine
 cdef class TensorRTEngine(InferenceEngine):
 
     cdef public object context
+    cdef object cuda_context
+    cdef object cuda_lock
 
     cdef public object d_input
     cdef public object d_output

src/engines/tensorrt_engine.pyx

@@ -1,10 +1,10 @@
 from engines.inference_engine cimport InferenceEngine
 import tensorrt as trt  # pyright: ignore[reportMissingImports]
 import pycuda.driver as cuda  # pyright: ignore[reportMissingImports]
-import pycuda.autoinit  # pyright: ignore[reportMissingImports]
 import pynvml
 import numpy as np
 import os
+import threading
 cimport constants_inf
 
 GPU_MEMORY_FRACTION = 0.8
@@ -32,6 +32,11 @@ class _CacheCalibrator(trt.IInt8EntropyCalibrator2):
 cdef class TensorRTEngine(InferenceEngine):
     def __init__(self, model_bytes: bytes, max_batch_size: int = 8, **kwargs):
         InferenceEngine.__init__(self, model_bytes, max_batch_size, engine_name="tensorrt")
+        self.cuda_context = TensorRTEngine.create_cuda_context()
+        self.cuda_lock = threading.Lock()
+        try:
+            with self.cuda_lock:
+                self.cuda_context.push()
                 try:
                     logger = trt.Logger(trt.Logger.WARNING)
                     runtime = trt.Runtime(logger)
@@ -70,10 +75,21 @@ cdef class TensorRTEngine(InferenceEngine):
                     self.d_output = cuda.mem_alloc(self.h_output.nbytes)
 
                     self.stream = cuda.Stream()
-
+                finally:
+                    try:
+                        self.cuda_context.pop()
+                    except Exception:
+                        pass
         except Exception as e:
             raise RuntimeError(f"Failed to initialize TensorRT engine: {str(e)}")
 
+    def __dealloc__(self):
+        try:
+            if self.cuda_context is not None:
+                self.cuda_context.detach()
+        except Exception:
+            pass
+
     @staticmethod
     def calculate_max_batch_size(gpu_memory_bytes, int input_h, int input_w):
         frame_input_bytes = 3 * input_h * input_w * 4
@@ -99,9 +115,18 @@ cdef class TensorRTEngine(InferenceEngine):
                 pass
         return 2 * 1024 * 1024 * 1024 if total_memory is None else total_memory
 
+    @staticmethod
+    def create_cuda_context():
+        cuda.init()
+        from engines import tensor_gpu_index
+        ctx = cuda.Device(max(tensor_gpu_index, 0)).make_context()
+        ctx.pop()
+        return ctx
+
     @staticmethod
     def get_engine_filename(str precision="fp16"):
         try:
+            cuda.init()
             from engines import tensor_gpu_index
             device = cuda.Device(max(tensor_gpu_index, 0))
             sm_count = device.multiprocessor_count
@@ -115,6 +140,8 @@ cdef class TensorRTEngine(InferenceEngine):
 
     @staticmethod
     def convert_from_source(bytes onnx_model, str calib_cache_path=None, bint force_static_input=False):
+        cuda_context = TensorRTEngine.create_cuda_context()
+        cuda_context.push()
         gpu_mem = TensorRTEngine.get_gpu_memory_bytes(0)
         workspace_bytes = int(gpu_mem * 0.9)
 
@@ -129,6 +156,7 @@ cdef class TensorRTEngine(InferenceEngine):
             except Exception as e:
                 constants_inf.logerror(<str>f'ONNX TensorRT compatibility preparation failed: {str(e)}')
 
+        try:
             with trt.Builder(trt_logger) as builder, \
                     builder.create_network(explicit_batch_flag) as network, \
                     trt.OnnxParser(network, trt_logger) as parser, \
@@ -180,11 +208,23 @@ cdef class TensorRTEngine(InferenceEngine):
                     return None
                 constants_inf.log('conversion done!')
                 return bytes(plan)
+        finally:
+            try:
+                cuda_context.pop()
+            except Exception:
+                pass
+            try:
+                cuda_context.detach()
+            except Exception:
+                pass
 
     cdef tuple get_input_shape(self):
         return <tuple>(self.input_shape[2], self.input_shape[3])
 
     cdef run(self, input_data):
+        try:
+            with self.cuda_lock:
+                self.cuda_context.push()
                 try:
                     actual_batch = input_data.shape[0]
                     if actual_batch != self.input_shape[0]:
@@ -202,6 +242,11 @@ cdef class TensorRTEngine(InferenceEngine):
                     output_shape = [actual_batch, self.output_shape[1], self.output_shape[2]]
                     output = self.h_output[:actual_batch].reshape(output_shape)
                     return [output]
+                finally:
+                    try:
+                        self.cuda_context.pop()
+                    except Exception:
+                        pass
 
         except Exception as e:
             raise RuntimeError(f"Failed to run TensorRT inference: {str(e)}")