[Dist] Add env var for non-default SSH configs in tests (#4098)

Signed-off-by: Serge Panev <spanev@nvidia.com>
Co-authored-by: Rhett Ying <85214957+Rhett-Ying@users.noreply.github.com>

tests/dist/utils.py

 import subprocess
 import multiprocessing as mp
 from typing import Optional
-
+import os
 
 def run(ssh_cmd):
     subprocess.check_call(ssh_cmd, shell=True)
@@ -29,8 +29,22 @@ def execute_remote(
     ip_prefix = ""
     if username:
         ip_prefix += "{username}@".format(username=username)
+
+    custom_port = os.getenv('DIST_DGL_TEST_SSH_PORT', '')
+    if custom_port:
+        port = custom_port
+
+    custom_ssh_key = os.getenv('DIST_DGL_TEST_SSH_KEY', '')
+    if custom_ssh_key:
+        custom_ssh_key = os.path.expanduser(custom_ssh_key)
+        custom_ssh_key = "-i " + custom_ssh_key
+
+    ssh_setup = os.getenv('DIST_DGL_TEST_SSH_SETUP', '')
+    if ssh_setup:
+        cmd = ssh_setup + ';' + cmd
     # Construct ssh command that executes `cmd` on the remote host
-    ssh_cmd = "ssh -o StrictHostKeyChecking=no -p {port} {ip_prefix}{ip} '{cmd}'".format(
+    ssh_cmd = "ssh -o StrictHostKeyChecking=no {ssh_key} -p {port} {ip_prefix}{ip} '{cmd}'".format(
+        ssh_key=custom_ssh_key,
         port=str(port),
         ip_prefix=ip_prefix,
         ip=ip,

tests/scripts/task_dist_test.sh

@@ -13,11 +13,24 @@ export DIST_DGL_TEST_CPP_BIN_DIR=${PWD}/build
 export DIST_DGL_TEST_IP_CONFIG=/home/ubuntu/workspace/ip_config.txt
 export DIST_DGL_TEST_PY_BIN_DIR=${PWD}/tests/dist/python
 
+if [[ -v DIST_DGL_TEST_SSH_PORT ]]; then
+    SSH_PORT_LINE="-p $DIST_DGL_TEST_SSH_PORT";
+fi
+
+if [[ -v DIST_DGL_TEST_SSH_KEY ]]; then
+    SSH_KEY_LINE="-i $DIST_DGL_TEST_SSH_KEY";
+fi
+
+if [[ -v DIST_DGL_TEST_SSH_SETUP ]]; then
+    SSH_SETUP_LINE="$DIST_DGL_TEST_SSH_SETUP;";
+fi
+
+
 while IFS= read line
 do
     for pkg in 'pytest' 'psutil' 'torch'
     do
-        ret_pkg=$(ssh ${line} "python3 -m pip list | grep -i ${pkg} ") || fail "${pkg} not installed in ${line}"
+        ret_pkg=$(ssh -o StrictHostKeyChecking=no ${line} ${SSH_PORT_LINE} ${SSH_KEY_LINE} "${SSH_SETUP_LINE}python3 -m pip list | grep -i ${pkg} ") || fail "${pkg} not installed in ${line}"
     done
 done < ${DIST_DGL_TEST_IP_CONFIG}