Skip to content

Commit e113b48

Browse files
narendasannarendasan
authored
feat(//py): Use TensorRT to fill in .so libraries automatically if possible (#1085)
* feat(//py): Use TensorRT to fill in .so libraries automatically if possible Signed-off-by: Naren Dasan <naren@narendasan.com> Signed-off-by: Naren Dasan <narens@nvidia.com> * feat: support finding libraries on aarch64 Signed-off-by: Naren Dasan <naren@narendasan.com> * feat: remove unnecessary direct loading cudnn Signed-off-by: Naren Dasan <naren@narendasan.com> * chore: Real testing workflow Signed-off-by: Naren Dasan <naren@narendasan.com> * chore: Fix build paths Signed-off-by: Naren Dasan <naren@narendasan.com> * chore: Dont need to install torchvision twice Signed-off-by: Naren Dasan <naren@narendasan.com> * chore: Reduce concurrent jobs Signed-off-by: Naren Dasan <naren@narendasan.com> * chore: Reduce concurrent jobs Signed-off-by: Naren Dasan <naren@narendasan.com>
1 parent ff2d955 commit e113b48

File tree

4 files changed

+238
-39
lines changed

4 files changed

+238
-39
lines changed

‎.circleci/config.yml‎

Lines changed: 145 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -236,7 +236,7 @@ commands:
236236
mv toolchains/ci_workspaces/WORKSPACE.<< parameters.platform >> WORKSPACE
237237
bazel query 'kind(cc_*, tests(//tests))' --noshow_progress >> /tmp/test_manifest.txt
238238
circleci tests split < /tmp/test_manifest.txt > /tmp/node_test_manifest.txt
239-
bazel test $(cat /tmp/node_test_manifest.txt) --test_arg=--gtest_output=xml:/tmp/artifacts/test_results/ --jobs 4 --config ci_testing --config pre_cxx11_abi --noshow_progress
239+
bazel test $(cat /tmp/node_test_manifest.txt) --test_arg=--gtest_output=xml:/tmp/artifacts/test_results/ --jobs 2 --config ci_testing --config pre_cxx11_abi --noshow_progress
240240
- run:
241241
name: Collect logs
242242
when: on_fail
@@ -265,7 +265,6 @@ commands:
265265
set -e
266266
mkdir -p /tmp/artifacts/test_results
267267
cd tests/py
268-
pip3 install -r requirements.txt
269268
pytest --junitxml=/tmp/artifacts/test_results/api/api_test_results.xml api/
270269
pytest --junitxml=/tmp/artifacts/test_results/integrations/integrations_test_results.xml integrations/
271270
cd ~/project
@@ -482,12 +481,20 @@ jobs:
482481
path: /tmp/dist/nightly
483482
destination: x86_64-pyt-nightly
484483

485-
test-core-cpp-x86_64-pyt-release:
484+
test-core-cpp-x86_64:
486485
parameters:
487486
torch-build:
488487
type: string
489488
torch-build-index:
490489
type: string
490+
trt-version-short:
491+
type: string
492+
trt-version-long:
493+
type: string
494+
cudnn-version:
495+
type: string
496+
channel:
497+
type: string
491498
machine:
492499
image: ubuntu-2004-cuda-11.4:202110-01
493500
resource_class: gpu.nvidia.large
@@ -497,66 +504,75 @@ jobs:
497504
- create-env:
498505
os: "ubuntu2004"
499506
platform: "x86_64"
500-
cudnn-version: << pipeline.parameters.cudnn-release-version >>
501-
trt-version-short: << pipeline.parameters.trt-release-version-short >>
507+
cudnn-version: << parameters.cudnn-version >>
508+
trt-version-short: << parameters.trt-version-short >>
502509
bazel-version: "5.1.1"
503510
bazel-platform: "x86_64"
504511
- create-py-env:
505-
trt-version-long: << pipeline.parameters.trt-release-version-long >>
512+
trt-version-long: << parameters.trt-version-long >>
506513
- install-torch-from-index:
507514
torch-build: << parameters.torch-build >>
508515
torch-build-index: << parameters.torch-build-index >>
509516
- attach_workspace:
510517
at: /tmp/dist
511518
- run:
512519
name: "Install torch-tensorrt"
513-
command: pip3 install /tmp/dist/release/*
520+
command: pip3 install /tmp/dist/<< parameters.channel >>/*
514521
- dump-test-env
515522
- test-ts-core
516523

517-
test-ts-py-x86_64-pyt-release:
524+
test-py-ts-x86_64:
518525
parameters:
526+
channel:
527+
type: string
519528
torch-build:
520529
type: string
521530
torch-build-index:
522531
type: string
532+
trt-version-long:
533+
type: string
523534
machine:
524535
image: ubuntu-2004-cuda-11.4:202110-01
525536
resource_class: gpu.nvidia.large
526537
steps:
527538
- checkout
528539
- create-py-env:
529-
trt-version-long: << pipeline.parameters.trt-release-version-long >>
530-
- install-torch-from-index:
531-
torch-build: << parameters.torch-build >>
532-
torch-build-index: << parameters.torch-build-index >>
540+
trt-version-long: << parameters.trt-version-long >>
533541
- attach_workspace:
534542
at: /tmp/dist
535543
- run:
536544
name: "Install torch-tensorrt"
537-
command: pip3 install /tmp/dist/release/*
545+
command: pip3 install /tmp/dist/<< parameters.channel >>/*
546+
# We install torch after torch-trt because pip automatically enforces the version constraint otherwise
547+
- install-torch-from-index:
548+
torch-build: << parameters.torch-build >>
549+
torch-build-index: << parameters.torch-build-index >>
538550
- dump-test-env
539551
- test-ts-py-api
540552

541-
test-x86_64-pyt-nightly:
553+
test-py-fx-x86_64:
542554
parameters:
555+
channel:
556+
type: string
543557
torch-build:
544558
type: string
545559
torch-build-index:
546560
type: string
561+
trt-version-long:
562+
type: string
547563
machine:
548564
image: ubuntu-2004-cuda-11.4:202110-01
549565
resource_class: gpu.nvidia.large
550566
steps:
551567
- checkout
552568
- create-py-env:
553-
trt-version-long: << pipeline.parameters.trt-nightly-version-long >>
569+
trt-version-long: << parameters.trt-version-long >>
554570
- attach_workspace:
555571
at: /tmp/dist/
556572
- run:
557573
name: "Install torch-tensorrt"
558-
command: pip3 install /tmp/dist/nightly/*
559-
# We install torch after torch-trt because pip automatically enforces the version constraint otherwise, swap back after versions are synced
574+
command: pip3 install /tmp/dist/<< parameters.channel >>/*
575+
# We install torch after torch-trt because pip automatically enforces the version constraint otherwise
560576
- install-torch-from-index:
561577
torch-build: << parameters.torch-build >>
562578
torch-build-index: << parameters.torch-build-index >>
@@ -573,30 +589,30 @@ parameters:
573589
default: "https://download.pytorch.org/whl/nightly/cu113"
574590
cudnn-nightly-version:
575591
type: string
576-
default: "8.2.1"
592+
default: "8.4.1"
577593
trt-nightly-version-short:
578594
type: string
579-
default: "8.2.4"
595+
default: "8.4.1"
580596
trt-nightly-version-long:
581597
type: string
582-
default: "8.2.4.2"
598+
default: "8.4.1.5"
583599

584600
# Release platform config
585601
torch-release-build:
586602
type: string
587-
default: "1.11.0+cu113"
603+
default: "1.12.0+cu113"
588604
torch-release-build-index:
589605
type: string
590606
default: "https://download.pytorch.org/whl/cu113"
591607
cudnn-release-version:
592608
type: string
593-
default: "8.2.1"
609+
default: "8.4.1"
594610
trt-release-version-short:
595611
type: string
596-
default: "8.2.4"
612+
default: "8.4.1"
597613
trt-release-version-long:
598614
type: string
599-
default: "8.2.4.2"
615+
default: "8.4.1.5"
600616

601617
# Jetson platform config
602618
torch-jetson-build:
@@ -632,21 +648,74 @@ workflows:
632648
jetpack-version: << pipeline.parameters.jetpack-version >>
633649
python-version: 3.8.10
634650

651+
652+
635653
- build-x86_64-pyt-release:
636654
torch-build: << pipeline.parameters.torch-release-build >>
637655
torch-build-index: << pipeline.parameters.torch-release-build-index >>
638-
- test-core-cpp-x86_64-pyt-release:
656+
657+
- test-core-cpp-x86_64:
658+
name: test-core-cpp-x86_64-pyt-release
659+
channel: "release"
639660
torch-build: << pipeline.parameters.torch-release-build >>
640661
torch-build-index: << pipeline.parameters.torch-release-build-index >>
662+
trt-version-short: << pipeline.parameters.trt-release-version-short >>
663+
trt-version-long: << pipeline.parameters.trt-release-version-long >>
664+
cudnn-version: << pipeline.parameters.cudnn-release-version >>
641665
requires:
642666
- build-x86_64-pyt-release
643667

668+
- test-py-ts-x86_64:
669+
name: test-py-ts-x86_64-pyt-release
670+
channel: "release"
671+
torch-build: << pipeline.parameters.torch-release-build >>
672+
torch-build-index: << pipeline.parameters.torch-release-build-index >>
673+
trt-version-long: << pipeline.parameters.trt-release-version-long >>
674+
requires:
675+
- build-x86_64-pyt-release
676+
677+
- test-py-ts-x86_64:
678+
name: test-py-fx-x86_64-pyt-release
679+
channel: "release"
680+
torch-build: << pipeline.parameters.torch-release-build >>
681+
torch-build-index: << pipeline.parameters.torch-release-build-index >>
682+
trt-version-long: << pipeline.parameters.trt-release-version-long >>
683+
requires:
684+
- build-x86_64-pyt-release
685+
686+
687+
688+
644689
- build-x86_64-pyt-nightly:
645690
torch-build: << pipeline.parameters.torch-nightly-build >>
646691
torch-build-index: << pipeline.parameters.torch-nightly-build-index >>
647-
- test-x86_64-pyt-nightly:
692+
693+
- test-core-cpp-x86_64:
694+
name: test-core-cpp-x86_64-pyt-nightly
695+
channel: "nightly"
696+
torch-build: << pipeline.parameters.torch-nightly-build >>
697+
torch-build-index: << pipeline.parameters.torch-nightly-build-index >>
698+
trt-version-short: << pipeline.parameters.trt-nightly-version-short >>
699+
trt-version-long: << pipeline.parameters.trt-nightly-version-long >>
700+
cudnn-version: << pipeline.parameters.cudnn-nightly-version >>
701+
requires:
702+
- build-x86_64-pyt-nightly
703+
704+
- test-py-ts-x86_64:
705+
name: test-py-ts-x86_64-pyt-nightly
706+
channel: "nightly"
648707
torch-build: << pipeline.parameters.torch-nightly-build >>
649708
torch-build-index: << pipeline.parameters.torch-nightly-build-index >>
709+
trt-version-long: << pipeline.parameters.trt-nightly-version-long >>
710+
requires:
711+
- build-x86_64-pyt-nightly
712+
713+
- test-py-fx-x86_64:
714+
name: test-py-fx-x86_64-pyt-nightly
715+
channel: "nightly"
716+
torch-build: << pipeline.parameters.torch-nightly-build >>
717+
torch-build-index: << pipeline.parameters.torch-nightly-build-index >>
718+
trt-version-long: << pipeline.parameters.trt-nightly-version-long >>
650719
requires:
651720
- build-x86_64-pyt-nightly
652721

@@ -657,26 +726,73 @@ workflows:
657726
jetpack-version: << pipeline.parameters.jetpack-version >>
658727
python-version: 3.8.10
659728

729+
730+
660731
- build-x86_64-pyt-release:
661732
torch-build: << pipeline.parameters.torch-release-build >>
662733
torch-build-index: << pipeline.parameters.torch-release-build-index >>
663-
- test-core-cpp-x86_64-pyt-release:
734+
735+
- test-core-cpp-x86_64:
736+
name: test-core-cpp-x86_64-pyt-release
737+
channel: "release"
738+
torch-build: << pipeline.parameters.torch-release-build >>
739+
torch-build-index: << pipeline.parameters.torch-release-build-index >>
740+
trt-version-short: << pipeline.parameters.trt-release-version-short >>
741+
trt-version-long: << pipeline.parameters.trt-release-version-long >>
742+
cudnn-version: << pipeline.parameters.cudnn-release-version >>
743+
requires:
744+
- build-x86_64-pyt-release
745+
746+
- test-py-ts-x86_64:
747+
name: test-py-ts-x86_64-pyt-release
748+
channel: "release"
664749
torch-build: << pipeline.parameters.torch-release-build >>
665750
torch-build-index: << pipeline.parameters.torch-release-build-index >>
751+
trt-version-long: << pipeline.parameters.trt-release-version-long >>
666752
requires:
667753
- build-x86_64-pyt-release
668-
- test-ts-py-x86_64-pyt-release:
754+
755+
- test-py-ts-x86_64:
756+
name: test-py-fx-x86_64-pyt-release
757+
channel: "release"
669758
torch-build: << pipeline.parameters.torch-release-build >>
670759
torch-build-index: << pipeline.parameters.torch-release-build-index >>
760+
trt-version-long: << pipeline.parameters.trt-release-version-long >>
671761
requires:
672762
- build-x86_64-pyt-release
673763

764+
765+
674766
- build-x86_64-pyt-nightly:
675767
torch-build: << pipeline.parameters.torch-nightly-build >>
676768
torch-build-index: << pipeline.parameters.torch-nightly-build-index >>
677-
- test-x86_64-pyt-nightly:
769+
770+
- test-core-cpp-x86_64:
771+
name: test-core-cpp-x86_64-pyt-nightly
772+
channel: "nightly"
678773
torch-build: << pipeline.parameters.torch-nightly-build >>
679774
torch-build-index: << pipeline.parameters.torch-nightly-build-index >>
775+
trt-version-short: << pipeline.parameters.trt-nightly-version-short >>
776+
trt-version-long: << pipeline.parameters.trt-nightly-version-long >>
777+
cudnn-version: << pipeline.parameters.cudnn-nightly-version >>
778+
requires:
779+
- build-x86_64-pyt-nightly
780+
781+
- test-py-ts-x86_64:
782+
name: test-py-ts-x86_64-pyt-nightly
783+
channel: "nightly"
784+
torch-build: << pipeline.parameters.torch-nightly-build >>
785+
torch-build-index: << pipeline.parameters.torch-nightly-build-index >>
786+
trt-version-long: << pipeline.parameters.trt-nightly-version-long >>
787+
requires:
788+
- build-x86_64-pyt-nightly
789+
790+
- test-py-fx-x86_64:
791+
name: test-py-fx-x86_64-pyt-nightly
792+
channel: "nightly"
793+
torch-build: << pipeline.parameters.torch-nightly-build >>
794+
torch-build-index: << pipeline.parameters.torch-nightly-build-index >>
795+
trt-version-long: << pipeline.parameters.trt-nightly-version-long >>
680796
requires:
681797
- build-x86_64-pyt-nightly
682798

‎py/setup.py‎

Lines changed: 14 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -22,9 +22,13 @@
2222

2323
JETPACK_VERSION = None
2424

25-
__version__ = '1.2.0a0'
2625
FX_ONLY = False
2726

27+
__version__ = '1.2.0a0'
28+
__cuda_version__ = '11.3'
29+
__cudnn_version__ = '8.2'
30+
__tensorrt_version__ = '8.2'
31+
2832
def get_git_revision_short_hash() -> str:
2933
return subprocess.check_output(['git', 'rev-parse', '--short', 'HEAD']).decode('ascii').strip()
3034

@@ -51,8 +55,10 @@ def get_git_revision_short_hash() -> str:
5155
JETPACK_VERSION = "4.5"
5256
elif version == "4.6":
5357
JETPACK_VERSION = "4.6"
58+
elif version == "5.0":
59+
JETPACK_VERSION = "4.6"
5460
if not JETPACK_VERSION:
55-
warnings.warn("Assuming jetpack version to be 4.6, if not use the --jetpack-version option")
61+
warnings.warn("Assuming jetpack version to be 4.6 or greater, if not use the --jetpack-version option")
5662
JETPACK_VERSION = "4.6"
5763

5864

@@ -103,7 +109,7 @@ def build_libtorchtrt_pre_cxx11_abi(develop=True, use_dist_dir=True, cxx11_abi=F
103109
print("Jetpack version: 4.5")
104110
elif JETPACK_VERSION == "4.6":
105111
cmd.append("--platforms=//toolchains:jetpack_4.6")
106-
print("Jetpack version: 4.6")
112+
print("Jetpack version: >=4.6")
107113

108114
print("building libtorchtrt")
109115
status_code = subprocess.run(cmd).returncode
@@ -118,7 +124,10 @@ def gen_version_file():
118124

119125
with open(dir_path + '/torch_tensorrt/_version.py', 'w') as f:
120126
print("creating version file")
121-
f.write("__version__ = \"" + __version__ + '\"')
127+
f.write("__version__ = \"" + __version__ + '\"\n')
128+
f.write("__cuda_version__ = \"" + __cuda_version__ + '\"\n')
129+
f.write("__cudnn_version__ = \"" + __cudnn_version__ + '\"\n')
130+
f.write("__tensorrt_version__ = \"" + __tensorrt_version__ + '\"\n')
122131

123132

124133
def copy_libtorchtrt(multilinux=False):
@@ -299,7 +308,7 @@ def run(self):
299308
long_description=long_description,
300309
ext_modules=ext_modules,
301310
install_requires=[
302-
'torch>=1.11.0+cu113,<1.12.0',
311+
'torch>=1.12.0+cu113,<1.13.0',
303312
],
304313
setup_requires=[],
305314
cmdclass={

0 commit comments

Comments
 (0)