Repository: pytorch/TensorRT (language: python)

Issue:
❓ [Question] How to solve this warning: Detected this engine is being instantitated in a multi-GPU system with multi-device safe mode disabled.
## ❓ Question

I used Torch-TensorRT to compile the torchscript model in C++. When compiling or loading torchtrt model, it displays many warnings.

```
WARNING: [Torch-TensorRT] - Detected this engine is being instantitated in a multi-GPU system with multi-device safe mode disabled. For more on the implications of this as well as workarounds, see the linked documentation (https://pytorch.org/TensorRT/user_guide/runtime.html#multi-device-safe-mode)
WARNING: [Torch-TensorRT] - Detected this engine is being instantitated in a multi-GPU system with multi-device safe mode disabled. For more on the implications of this as well as workarounds, see the linked documentation (https://pytorch.org/TensorRT/user_guide/runtime.html#multi-device-safe-mode)
WARNING: [Torch-TensorRT] - Detected this engine is being instantitated in a multi-GPU system with multi-device safe mode disabled. For more on the implications of this as well as workarounds, see the linked documentation (https://pytorch.org/TensorRT/user_guide/runtime.html#multi-device-safe-mode)
```

## What you have already tried

I found this [link](https://pytorch.org/TensorRT/user_guide/runtime.html#multi-device-safe-mode) is useful, but it only provides Python API.

I checked the source code, but I still haven't figured out how to set up MULTI_DEVICE_SAFE_MODE in C++.

What can I do to address this warning?

## Environment

> Build information about Torch-TensorRT can be found by turning on debug messages

 - PyTorch Version (e.g., 1.0):
 - CPU Architecture: x86
 - OS (e.g., Linux): ubuntu18
 - How you installed PyTorch (`conda`, `pip`, `libtorch`, source): libtorch
 - Build command you used (if compiling from source):
 - Are you using local sources or building from archives:
 - Python version:
 - CUDA version: 12.2
 - GPU models and configuration: 1080Ti
 - Any other relevant information:



Discussion / hints:
Moreover, when the model infers using forward function, there are these warnings:

```
WARNING: [Torch-TensorRT] - Using default stream in enqueue()/enqueueV2()/enqueueV3() may lead to performance issues due to additional cudaDeviceSynchronize() calls by TensorRT to ensure correct synchronizations. Please use non-default stream instead.
WARNING: [Torch-TensorRT] - Using default stream in enqueue()/enqueueV2()/enqueueV3() may lead to performance issues due to additional cudaDeviceSynchronize() calls by TensorRT to ensure correct synchronizations. Please use non-default stream instead.
WARNING: [Torch-TensorRT] - Using default stream in enqueue()/enqueueV2()/enqueueV3() may lead to performance issues due to additional cudaDeviceSynchronize() calls by TensorRT to ensure correct synchronizations. Please use non-default stream instead.
```
Hi @demuxin - thanks for the report - we likely need to add `getter`/`setter` methods to toggle this value in C++ as well:
https://github.com/pytorch/TensorRT/blob/07c5b07b5ca81165682be90049a6bc67cff9928e/core/runtime/runtime.cpp#L10
Similar to the following functions:
https://github.com/pytorch/TensorRT/blob/07c5b07b5ca81165682be90049a6bc67cff9928e/core/runtime/register_jit_hooks.cpp#L121-L122
We are aware of the default stream warning and are working on this. It should not have a substantial effect on inference from what I've seen
These may actually already be accessible in C++, prior to inference, could you try adding the line:
`torch_tensorrt::core::runtime::set_multi_device_safe_mode(true)`
Hi @gs-olive, this is not right, there is complie error:
```
error: ‘set_multi_device_safe_mode’ is not a member of ‘torch_tensorrt::core::runtime’
   27 |         torch_tensorrt::core::runtime::set_multi_device_safe_mode(true);
```
Does `torch::ops::tensorrt::get_multi_device_safe_mode()` exist, or does this also cause a compilation error?

Fix patch:
diff --git a/core/runtime/runtime.cpp b/core/runtime/runtime.cpp
index afe559ec4d..24eba16cc3 100644
--- a/core/runtime/runtime.cpp
+++ b/core/runtime/runtime.cpp
@@ -121,6 +121,14 @@ void multi_gpu_device_check() {
   }
 }
 
+bool get_multi_device_safe_mode() {
+  return MULTI_DEVICE_SAFE_MODE;
+}
+
+void set_multi_device_safe_mode(bool multi_device_safe_mode) {
+  MULTI_DEVICE_SAFE_MODE = multi_device_safe_mode;
+}
+
 namespace {
 static DeviceList cuda_device_list;
 }
diff --git a/core/runtime/runtime.h b/core/runtime/runtime.h
index 5e5676b11e..8c9b33328a 100644
--- a/core/runtime/runtime.h
+++ b/core/runtime/runtime.h
@@ -38,6 +38,10 @@ std::vector<at::Tensor> execute_engine(std::vector<at::Tensor> inputs, c10::intr
 
 void multi_gpu_device_check();
 
+bool get_multi_device_safe_mode();
+
+void set_multi_device_safe_mode(bool multi_device_safe_mode);
+
 class DeviceList {
   using DeviceMap = std::unordered_map<int, RTDevice>;
   DeviceMap device_list;


Test patch:
diff --git a/tests/core/runtime/BUILD b/tests/core/runtime/BUILD
index cd9f123b59..deddca4cfb 100644
--- a/tests/core/runtime/BUILD
+++ b/tests/core/runtime/BUILD
@@ -1,3 +1,5 @@
+load("//tests/core/runtime:runtime_test.bzl", "runtime_test")
+
 package(default_visibility = ["//visibility:public"])
 
 config_setting(
@@ -6,3 +8,14 @@ config_setting(
         "define": "abi=pre_cxx11_abi",
     },
 )
+
+runtime_test(
+    name = "test_multi_device_safe_mode",
+)
+
+test_suite(
+    name = "runtime_tests",
+    tests = [
+        ":test_multi_device_safe_mode",
+    ],
+)
diff --git a/tests/core/runtime/runtime_test.bzl b/tests/core/runtime/runtime_test.bzl
new file mode 100644
index 0000000000..2c4753cc9b
--- /dev/null
+++ b/tests/core/runtime/runtime_test.bzl
@@ -0,0 +1,25 @@
+"""
+runtime test macros
+"""
+
+load("@rules_cc//cc:defs.bzl", "cc_test")
+
+def runtime_test(name, visibility = None):
+    """Macro to define a runtime test
+
+    Args:
+        name: Name of test file
+        visibility: Visibility of the test target
+    """
+    cc_test(
+        name = name,
+        srcs = [name + ".cpp"],
+        visibility = visibility,
+        deps = [
+            "//tests/util",
+            "@googletest//:gtest_main",
+        ] + select({
+            ":use_pre_cxx11_abi": ["@libtorch_pre_cxx11_abi//:libtorch"],
+            "//conditions:default": ["@libtorch//:libtorch"],
+        }),
+    )
diff --git a/tests/core/runtime/test_multi_device_safe_mode.cpp b/tests/core/runtime/test_multi_device_safe_mode.cpp
new file mode 100644
index 0000000000..a8c72cac7b
--- /dev/null
+++ b/tests/core/runtime/test_multi_device_safe_mode.cpp
@@ -0,0 +1,8 @@
+#include "core/runtime/runtime.h"
+#include "gtest/gtest.h"
+
+TEST(Runtime, MultiDeviceSafeMode) {
+  ASSERT_TRUE(!torch_tensorrt::core::runtime::get_multi_device_safe_mode());
+  torch_tensorrt::core::runtime::set_multi_device_safe_mode(true);
+  ASSERT_TRUE(torch_tensorrt::core::runtime::get_multi_device_safe_mode());
+}