nccl ops address review

Author: apbose

tests/py/dynamo/distributed/distributed_utils.py

@@ -16,18 +16,6 @@ def set_environment_variables_pytest():
     os.environ["USE_TRTLLM_PLUGINS"] = "1"
 
 
-def find_repo_root(max_depth=10):
-    dir_path = os.path.dirname(os.path.realpath(__file__))
-    for i in range(max_depth):
-        files = os.listdir(dir_path)
-        if "MODULE.bazel" in files:
-            return dir_path
-        else:
-            dir_path = os.path.dirname(dir_path)
-
-    raise RuntimeError("Could not find repo root")
-
-
 def initialize_logger(rank, logger_file_name):
     logger = logging.getLogger()
     logger.setLevel(logging.INFO)
@@ -49,9 +37,7 @@ def initialize_distributed_env(logger_file_name, rank=0, world_size=1, port=2950
     os.environ["WORLD_SIZE"] = str(world_size)
     os.environ["MASTER_ADDR"] = "127.0.0.1"
     os.environ["MASTER_PORT"] = str(port)
-    os.environ["TRTLLM_PLUGINS_PATH"] = (
-        find_repo_root() + "/lib/libnvinfer_plugin_tensorrt_llm.so"
-    )
+    os.environ["TRTLLM_PLUGINS_PATH"] = "./tmp/lib/libnvinfer_plugin_tensorrt_llm.so"
 
     # Necessary to assign a device to each rank.
     torch.cuda.set_device(local_rank)

tests/py/dynamo/distributed/test_nccl_ops.sh

@@ -72,20 +72,15 @@ echo "Setting up the environment"
 
 OS="$(uname -s)"
 ARCH="$(uname -m)"
-PYTHON_VERSION="$(python3 -c 'import sys; print(f"cp{sys.version_info.major}{sys.version_info.minor}")')"
 
 
 #getting the file name for TensorRT-LLM download
-if [[ "$OS" == "Linux" && "$ARCH" == "x86_64" && "$PYTHON_VERSION" == "cp312" ]]; then
+if [[ "$OS" == "Linux" && "$ARCH" == "x86_64"]]; then
     FILE="tensorrt_llm-0.17.0.post1-cp312-cp312-linux_x86_64.whl"
-elif [[ "$OS" == "Linux" && "$ARCH" == "aarch64" && "$PYTHON_VERSION" == "cp312" ]]; then
+elif [[ "$OS" == "Linux" && "$ARCH" == "aarch64"]]; then
     FILE="tensorrt_llm-0.17.0.post1-cp312-cp312-linux_aarch64.whl"
-elif [[ "$OS" == "Linux" && "$ARCH" == "x86_64" && "$PYTHON_VERSION" == "cp310" ]]; then
-    FILE="tensorrt_llm-0.17.0.post1-cp310-cp310-linux_x86_64.whl"
-elif [[ "$OS" == "Linux" && "$ARCH" == "aarch64" && "$PYTHON_VERSION" == "cp310" ]]; then
-    FILE="tensorrt_llm-0.17.0.post1-cp310-cp310-linux_aarch64.whl"
 else:
-    echo "Unsupported platform: OS=$OS ARCH=$ARCH PYTHON=$PYTHON_VERSION"
+    echo "Unsupported platform: OS=$OS ARCH=$ARCH
     exit 1
 fi