Skip to content

feat: TensorRT support#1205

Open
dodamih wants to merge 7 commits intomainfrom
dodam/tensorrt
Open

feat: TensorRT support#1205
dodamih wants to merge 7 commits intomainfrom
dodam/tensorrt

Conversation

@dodamih
Copy link
Collaborator

@dodamih dodamih commented Mar 24, 2026

Occasionally it looks like the workers will fail to find TensorRT, but it's handled gracefully.

trivoldus28 and others added 5 commits March 24, 2026 16:02
Add ldconfig for nvidia/tensorrt .so paths, prune large
site-packages dirs from ldd scan to speed up build.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
Switch from torchscript-based TRT compilation to torch_tensorrt.compile
with ExportedProgram (.ep) format. Add GPU memory cleanup, improve
error handling with fallback to eager mode.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
Make tensorrt semaphore optional with a default count, so existing
specs without it don't break.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
assert input_shape is not None # mypy

trt_fname = (
str(xxhash.xxh128(str((path, tuple(input_shape))).encode("utf-8")).hexdigest())
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think you need to also add the gpu model or gpu capability to the hash. I'm not sure if a T4 can run a model compiled by L4 or reversed

@codecov
Copy link

codecov bot commented Mar 25, 2026

Codecov Report

❌ Patch coverage is 91.66667% with 1 line in your changes missing coverage. Please review.
✅ Project coverage is 99.97%. Comparing base (8c2437b) to head (6991a3e).
⚠️ Report is 79 commits behind head on main.

Files with missing lines Patch % Lines
zetta_utils/convnet/utils.py 88.88% 1 Missing ⚠️
Additional details and impacted files
@@            Coverage Diff             @@
##             main    #1205      +/-   ##
==========================================
- Coverage   99.98%   99.97%   -0.01%     
==========================================
  Files         192      192              
  Lines       10199    10214      +15     
==========================================
+ Hits        10197    10211      +14     
- Misses          2        3       +1     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants