From 34e87c339dad5401bd5a571e1b171f11cb6ab9b8 Mon Sep 17 00:00:00 2001 From: jason Date: Wed, 23 Sep 2026 17:12:06 -0400 Subject: [PATCH 1/3] add ci unit tests --- .github/workflows/tests.yml | 62 +++++++++++++++++++++ docker/Dockerfile.trt | 3 + scripts/export_ultralytics_image_encoder.py | 28 ++++------ scripts/export_ultralytics_mask_decoder.py | 19 ++++--- scripts/export_ultralytics_text_encoder.py | 4 +- 5 files changed, 89 insertions(+), 27 deletions(-) create mode 100644 .github/workflows/tests.yml diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml new file mode 100644 index 0000000..84481f1 --- /dev/null +++ b/.github/workflows/tests.yml @@ -0,0 +1,62 @@ +name: Unit Tests + +on: + pull_request: + branches: [main] + push: + branches: [main] + workflow_dispatch: + +permissions: + contents: read + +concurrency: + group: unit-tests-${{ github.ref }} + cancel-in-progress: true + +jobs: + unit-tests: + runs-on: [self-hosted, linux, gpu] + timeout-minutes: 60 + + if: >- + github.event_name != 'pull_request' || + github.event.pull_request.head.repo.full_name == github.repository + + steps: + - uses: actions/checkout@v4 + with: + clean: false + + - name: Check GPU and engines + run: | + nvidia-smi + echo "HOME=$HOME" + ls -lh "$HOME/models" + if ! ls "$HOME"/models/*.engine >/dev/null 2>&1; then + echo "::error::No .engine files found in $HOME/models on this runner" + exit 1 + fi + + # Cpp tests + - name: Configure (C++) + working-directory: sam3trt + run: cmake -S . -B build -DSAM3TRT_BUILD_TESTS=ON -DSAM3TRT_BUILD_ROS=OFF + + - name: Build (C++) + working-directory: sam3trt + run: cmake --build build -j"$(nproc)" + + - name: Run C++ tests + working-directory: sam3trt + run: ctest --test-dir build --output-on-failure -E Timing + + # Python tests + - name: Install and run Python tests + if: ${{ !cancelled() }} + run: | + python3 -m venv --system-site-packages .venv + . .venv/bin/activate + pip install ./sam3trtpy + cd sam3trtpy + python3 -m unittest discover -s tests -v diff --git a/docker/Dockerfile.trt b/docker/Dockerfile.trt index bc7490d..45460b0 100644 --- a/docker/Dockerfile.trt +++ b/docker/Dockerfile.trt @@ -35,6 +35,9 @@ WORKDIR /home/$USER RUN sudo apt install -y python3-pip RUN pip3 install --no-deps opencv-python-headless ultralytics +RUN pip install --no-cache-dir ultralytics==8.4.121 +ENV LD_LIBRARY_PATH=/opt/hpcx/ucx/lib:/opt/hpcx/ucc/lib:${LD_LIBRARY_PATH} + RUN sudo apt update \ && sudo apt install -y \ diff --git a/scripts/export_ultralytics_image_encoder.py b/scripts/export_ultralytics_image_encoder.py index 26a7f5c..2488f92 100644 --- a/scripts/export_ultralytics_image_encoder.py +++ b/scripts/export_ultralytics_image_encoder.py @@ -80,12 +80,6 @@ def patch_layernorm(model): def trace_and_export_image_encoder(model : torch.nn.Module, input : Any, engine_path : str, fp16 : bool) -> Tuple[torch.Tensor]: wrapper = ImageEncoderWrapper(model).to(DEVICE).eval() - predictor.setup_source(input) - - for batch in predictor.dataset: - input = predictor.preprocess(batch[1]) - break - print("[ImageEncoderExport] Tracing Image Encoder Model") torch_output = wrapper(input) # warm up exp_program = torch.export.export(wrapper, (input,), strict=False) @@ -223,16 +217,17 @@ def export_and_verify_image_encoder(predictor : SAM3SemanticPredictor, fp16 : bo precision = "fp16" if fp16 else "fp32" engine_path = os.path.join(os.environ["HOME"], "models", f"image_encoder_{precision}.engine") - #_ = trace_and_export_image_encoder( - # predictor.model, - # img, - # engine_path, - # fp16 - #) for batch in predictor.dataset: im = predictor.preprocess(batch[1]) break + _ = trace_and_export_image_encoder( + predictor.model, + im, + engine_path, + fp16 + ) + _verify_engine(engine_path, im, torch_output, min_cos=0.999 if fp16 else 0.9999) return torch_output @@ -255,15 +250,16 @@ def export_and_verify_image_encoder(predictor : SAM3SemanticPredictor, fp16 : bo precision = "fp16" if args.fp16 else "fp32" engine_path = os.path.join(os.environ["HOME"], "models", f"image_encoder_{precision}.engine") + + for batch in predictor.dataset: + im = predictor.preprocess(batch[1]) + break _ = trace_and_export_image_encoder( predictor.model, - img, + im, engine_path, args.fp16 ) - for batch in predictor.dataset: - im = predictor.preprocess(batch[1]) - break _verify_engine(engine_path, im, torch_output, min_cos=0.999 if args.fp16 else 0.9999) diff --git a/scripts/export_ultralytics_mask_decoder.py b/scripts/export_ultralytics_mask_decoder.py index 606c4ff..02836f7 100644 --- a/scripts/export_ultralytics_mask_decoder.py +++ b/scripts/export_ultralytics_mask_decoder.py @@ -238,7 +238,7 @@ def trace_and_export_mask_deocder(model : torch.nn.Module, input : Any, engine_p print("[MaskDecoderExport] Patching Mask Decoder Model") _, _, spatial_shapes, _ = fusion_wrapper(*input) H, W = int(spatial_shapes[0, 0]), int(spatial_shapes[0, 1]) - patch_rpb(predictor.model, H, W) + patch_rpb(model, H, W) torch_output = wrapper(*input) @@ -256,9 +256,10 @@ def trace_and_export_mask_deocder(model : torch.nn.Module, input : Any, engine_p offload_module_to_cpu=True, device=torch_tensorrt.Device("cuda:0"), ) + model.to(DEVICE) # offload_module_to_cpu leaves the weights on the CPU - #with open(engine_path, "wb") as f: - # f.write(engine_bytes) + with open(engine_path, "wb") as f: + f.write(engine_bytes) return torch_output @@ -376,12 +377,12 @@ def export_and_verify_mask_decoder(predictor : SAM3SemanticPredictor, input : Tu precision = "fp16" if fp16 else "fp32" engine_path = os.path.join(os.environ["HOME"], "models", f"mask_decoder_{precision}.engine") - #trace_and_export_mask_deocder( - # predictor.model, - # input, - # engine_path, - # fp16 - #) + trace_and_export_mask_deocder( + predictor.model, + input, + engine_path, + fp16 + ) predictor.set_image(img) ref = predictor(text=captions)[0] for batch in predictor.dataset: diff --git a/scripts/export_ultralytics_text_encoder.py b/scripts/export_ultralytics_text_encoder.py index ecc6dd7..b7a44b1 100644 --- a/scripts/export_ultralytics_text_encoder.py +++ b/scripts/export_ultralytics_text_encoder.py @@ -79,7 +79,7 @@ def trace_and_export_text_encoder(model : torch.nn.Module, input : Any, engine_p exp_program, arg_inputs=[input[0].to(DEVICE), input[1].to(DEVICE)], optimization_level=5, - use_explicit_typeing=True, + use_explicit_typing=True, device=torch_tensorrt.Device(f"cuda:0"), ) @@ -184,7 +184,7 @@ def export_and_verify_text_encoder(predictor : SAM3SemanticPredictor, fp16 : boo precision = "fp16" if fp16 else "fp32" engine_path = os.path.join(os.environ["HOME"], "models", f"text_encoder_{precision}.engine") - #trace_and_export_text_encoder(predictor.model, (input_ids, attention_mask), engine_path, args.fp16) + trace_and_export_text_encoder(predictor.model, (input_ids, attention_mask), engine_path, fp16) _verify_engine(engine_path, (input_ids, attention_mask), (torch_output[0], torch_output[2]), 0.999 if fp16 else 0.9999) From cca379c27d2582f2d68de997eb820c049c8b684b Mon Sep 17 00:00:00 2001 From: jason Date: Wed, 23 Sep 2026 20:39:43 -0400 Subject: [PATCH 2/3] update ci --- .github/workflows/tests.yml | 94 ++++++++++++++++++++++++------------- 1 file changed, 61 insertions(+), 33 deletions(-) diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml index 84481f1..f4835f7 100644 --- a/.github/workflows/tests.yml +++ b/.github/workflows/tests.yml @@ -1,3 +1,4 @@ +# Save as .github/workflows/tests.yml name: Unit Tests on: @@ -14,11 +15,17 @@ concurrency: group: unit-tests-${{ github.ref }} cancel-in-progress: true +env: + # Tag produced by docker/build-trt.bash + IMAGE: sam3-tensorrt:infer + jobs: + # This job name is what you select in the branch protection rule. unit-tests: - runs-on: [self-hosted, linux, gpu] - timeout-minutes: 60 + runs-on: [self-hosted, linux, gpu] # must match --labels used in config.sh + timeout-minutes: 90 + # Don't run fork PRs on your own hardware. if: >- github.event_name != 'pull_request' || github.event.pull_request.head.repo.full_name == github.repository @@ -26,37 +33,58 @@ jobs: steps: - uses: actions/checkout@v4 with: + # Keep sam3trt/build between runs for incremental C++ builds. clean: false - - name: Check GPU and engines - run: | - nvidia-smi - echo "HOME=$HOME" - ls -lh "$HOME/models" - if ! ls "$HOME"/models/*.engine >/dev/null 2>&1; then - echo "::error::No .engine files found in $HOME/models on this runner" - exit 1 - fi - - # Cpp tests - - name: Configure (C++) - working-directory: sam3trt - run: cmake -S . -B build -DSAM3TRT_BUILD_TESTS=ON -DSAM3TRT_BUILD_ROS=OFF - - - name: Build (C++) - working-directory: sam3trt - run: cmake --build build -j"$(nproc)" - - - name: Run C++ tests - working-directory: sam3trt - run: ctest --test-dir build --output-on-failure -E Timing - - # Python tests - - name: Install and run Python tests - if: ${{ !cancelled() }} + # Same script you use by hand. It must run from inside docker/ (it uses + # `-f Dockerfile.trt .`). Layers are cached, so only the first run is slow. + - name: Build docker image + working-directory: docker + run: ./build-trt.bash + + - name: Run unit tests in container + env: + # Engines live on the runner machine (they are not in git). + # Must contain image_encoder_fp16.engine, text_encoder_fp16.engine, + # mask_decoder_fp16.engine. + MODELS_DIR: /home/jason/models + TEST_SCRIPT: | + set -uo pipefail + rc=0 + + if ! ls "$HOME"/models/*.engine >/dev/null 2>&1; then + echo "::error::No .engine files found in the mounted models dir" + exit 1 + fi + nvidia-smi || true + + echo "::group::C++ build and tests" + ( cd "$HOME/sam3trt" \ + && cmake -S . -B build -DSAM3TRT_BUILD_TESTS=ON -DSAM3TRT_BUILD_ROS=OFF \ + && cmake --build build -j"$(nproc)" \ + && ctest --test-dir build --output-on-failure -E Timing ) || rc=1 + echo "::endgroup::" + + echo "::group::Python tests" + ( cd "$HOME" \ + && pip install --user --break-system-packages --no-cache-dir ./sam3trtpy \ + && cd sam3trtpy \ + && python3 -m unittest discover -s tests -v ) || rc=1 + echo "::endgroup::" + + exit $rc run: | - python3 -m venv --system-site-packages .venv - . .venv/bin/activate - pip install ./sam3trtpy - cd sam3trtpy - python3 -m unittest discover -s tests -v + U="$(whoami)" + # Same mounts as docker/run-trt.bash, minus the interactive/X11 parts. + docker run --rm \ + --gpus all \ + --ipc=host \ + --user "$(id -u):$(id -g)" \ + -e HOME="/home/$U" \ + -e TEST_SCRIPT \ + -v "$MODELS_DIR:/home/$U/models" \ + -v "$GITHUB_WORKSPACE/sam3trt:/home/$U/sam3trt" \ + -v "$GITHUB_WORKSPACE/sam3trtpy:/home/$U/sam3trtpy" \ + -w "/home/$U" \ + "$IMAGE" \ + bash -c "$TEST_SCRIPT" From 6b9b582dac03639f9b75b36e6da39fe9f76ba259 Mon Sep 17 00:00:00 2001 From: jason Date: Wed, 23 Sep 2026 20:59:11 -0400 Subject: [PATCH 3/3] fixing ci... again --- .github/workflows/tests.yml | 12 +----------- 1 file changed, 1 insertion(+), 11 deletions(-) diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml index f4835f7..d4ccf7b 100644 --- a/.github/workflows/tests.yml +++ b/.github/workflows/tests.yml @@ -1,4 +1,3 @@ -# Save as .github/workflows/tests.yml name: Unit Tests on: @@ -16,16 +15,13 @@ concurrency: cancel-in-progress: true env: - # Tag produced by docker/build-trt.bash IMAGE: sam3-tensorrt:infer jobs: - # This job name is what you select in the branch protection rule. unit-tests: runs-on: [self-hosted, linux, gpu] # must match --labels used in config.sh timeout-minutes: 90 - # Don't run fork PRs on your own hardware. if: >- github.event_name != 'pull_request' || github.event.pull_request.head.repo.full_name == github.repository @@ -33,21 +29,15 @@ jobs: steps: - uses: actions/checkout@v4 with: - # Keep sam3trt/build between runs for incremental C++ builds. clean: false - # Same script you use by hand. It must run from inside docker/ (it uses - # `-f Dockerfile.trt .`). Layers are cached, so only the first run is slow. - name: Build docker image working-directory: docker run: ./build-trt.bash - name: Run unit tests in container env: - # Engines live on the runner machine (they are not in git). - # Must contain image_encoder_fp16.engine, text_encoder_fp16.engine, - # mask_decoder_fp16.engine. - MODELS_DIR: /home/jason/models + MODELS_DIR: /home/jason/Projects/sam3-ws/SAM3-TensorRT/models TEST_SCRIPT: | set -uo pipefail rc=0