From 11376ae98660d2ec9f6c9521ad848740fd7e3c20 Mon Sep 17 00:00:00 2001
From: gs-olive <113141689+gs-olive@users.noreply.github.com>
Date: Mon, 18 Dec 2023 14:55:24 -0800
Subject: [PATCH] Move Torch-TRT install to file (#2092)

Summary:
- Add custom installer support for userbenchmark testing
- Add support for installing Torch-TRT outside of main container installs for nightly runs
- Add necessary hooks and subprocess commands in code

Pull Request resolved: https://github.com/pytorch/benchmark/pull/2092

Reviewed By: aaronenyeshi

Differential Revision: D52266771

Pulled By: xuzhao9

fbshipit-source-id: 369009a5bd1d5681a8aa7f72f736592da952ffda
---
 .../userbenchmark/schedule-benchmarks.py      | 10 +++++
 userbenchmark/torch_trt/install.py            | 44 +++++++++++++++++++
 utils/cuda_utils.py                           | 28 ------------
 3 files changed, 54 insertions(+), 28 deletions(-)
 create mode 100644 userbenchmark/torch_trt/install.py

diff --git a/.github/scripts/userbenchmark/schedule-benchmarks.py b/.github/scripts/userbenchmark/schedule-benchmarks.py
index b29fa81241..d81b06ac7c 100644
--- a/.github/scripts/userbenchmark/schedule-benchmarks.py
+++ b/.github/scripts/userbenchmark/schedule-benchmarks.py
@@ -1,5 +1,6 @@
 import argparse
 import sys
+import os
 import subprocess
 from pathlib import Path
 
@@ -25,6 +26,15 @@ def __exit__(self, exc_type, exc_value, traceback):
 
 def run_userbenchmark(ub_name, dryrun=True):
     workdir = REPO_ROOT
+
+    # Check if userbenchmark has an installer
+    candidate_installer_path = os.path.join(workdir, "userbenchmark", ub_name, "install.py")
+    if os.path.exists(candidate_installer_path):
+        install_command = [sys.executable, "install.py"]
+        print(f"Running user benchmark installer: {command}")
+        if not dryrun:
+            subprocess.check_call(install_command, cwd=Path(candidate_installer_path).parent.resolve())
+
     command = [sys.executable, "run_benchmark.py", ub_name]
     print(f"Running user benchmark command: {command}")
     if not dryrun:
diff --git a/userbenchmark/torch_trt/install.py b/userbenchmark/torch_trt/install.py
new file mode 100644
index 0000000000..aebfa70af7
--- /dev/null
+++ b/userbenchmark/torch_trt/install.py
@@ -0,0 +1,44 @@
+import subprocess
+
+import torch
+
+
+def install_torch_tensorrt():
+    # Install Torch-TensorRT with validation
+    uninstall_torchtrt_cmd = ["pip", "uninstall", "-y", "torch_tensorrt"]
+    subprocess.check_call(uninstall_torchtrt_cmd)
+
+    if torch.version.cuda.startswith("12"):
+        cuda_index_modifier = "cu121"
+    elif torch.version.cuda.startswith("11"):
+        cuda_index_modifier = "cu118"
+    else:
+        raise AssertionError(
+            f"Detected Torch-TRT unsupported CUDA version {torch.version.cuda}"
+        )
+
+    pytorch_nightly_url = (
+        f"https://download.pytorch.org/whl/nightly/{cuda_index_modifier}"
+    )
+    install_torchtrt_cmd = [
+        "pip",
+        "install",
+        "--pre",
+        "--no-cache-dir",
+        "torch_tensorrt",
+        "--extra-index-url",
+        pytorch_nightly_url,
+    ]
+    validate_torchtrt_cmd = ["python", "-c", "'import torch_tensorrt'"]
+
+    # Install and validate Torch-TensorRT
+    try:
+        subprocess.check_call(install_torchtrt_cmd)
+        subprocess.check_call(validate_torchtrt_cmd)
+    except subprocess.CalledProcessError:
+        subprocess.check_call(uninstall_torchtrt_cmd)
+        print("Failed to install torch-tensorrt, skipping install")
+
+
+if __name__ == "__main__":
+    install_torch_tensorrt()
diff --git a/utils/cuda_utils.py b/utils/cuda_utils.py
index 1a0ee23278..043725c9f4 100644
--- a/utils/cuda_utils.py
+++ b/utils/cuda_utils.py
@@ -88,34 +88,6 @@ def install_pytorch_nightly(cuda_version: str, env, dryrun=False):
     else:
         subprocess.check_call(install_torch_cmd, env=env)
 
-    # Install Torch-TensorRT with validation
-    uninstall_torchtrt_cmd = ["pip", "uninstall", "-y", "torch_tensorrt"]
-    if dryrun:
-        print(f"Uninstall torch-tensorrt: {uninstall_torchtrt_cmd}")
-    else:
-        subprocess.check_call(uninstall_torchtrt_cmd)
-
-    install_torchtrt_cmd = [
-        "pip",
-        "install",
-        "--pre",
-        "--no-cache-dir",
-        "torch_tensorrt",
-        "--extra-index-url",
-        pytorch_nightly_url,
-    ]
-    validate_torchtrt_cmd = ["python", "-c", "'import torch_tensorrt'"]
-    if dryrun:
-        print(f"Install torch-tensorrt nightly: {install_torchtrt_cmd}")
-        print(f"Validate torch-tensorrt nightly install: {validate_torchtrt_cmd}")
-    else:
-        try:
-            subprocess.check_call(install_torchtrt_cmd, env=env)
-            subprocess.check_call(validate_torchtrt_cmd, env=env)
-        except subprocess.CalledProcessError:
-            subprocess.check_call(uninstall_torchtrt_cmd, env=env)
-            print(f"Failed to install torch-tensorrt, skipping install")
-
 def install_torch_deps(cuda_version: str):
     # install magma
     magma_pkg = CUDA_VERSION_MAP[cuda_version]["magma_version"]