Fix ssh host policy

src/sagemaker/modules/train/container_drivers/mpi_utils.py

@@ -14,12 +14,12 @@
 from __future__ import absolute_import
 
 import os
-import time
 import subprocess
-
+import time
 from typing import List
 
-from utils import logger, SM_EFA_NCCL_INSTANCES, SM_EFA_RDMA_INSTANCES, get_python_executable
+import paramiko
+from utils import SM_EFA_NCCL_INSTANCES, SM_EFA_RDMA_INSTANCES, get_python_executable, logger
 
 FINISHED_STATUS_FILE = "/tmp/done.algo-1"
 READY_FILE = "/tmp/ready.%s"
@@ -75,6 +75,24 @@
     logger.info("Started SSH daemon.")
 
 
+class CustomHostKeyPolicy(paramiko.client.MissingHostKeyPolicy):
+    def missing_host_key(self, client, hostname, key):
+        """Accept host keys for algo-* hostnames, reject others.
+
+        Args:
+            client: The SSHClient instance
+            hostname: The hostname attempting to connect
+            key: The host key
+
+        Raises:
+            paramiko.SSHException: If hostname doesn't match algo-* pattern
+        """
+        if hostname.startswith("algo-"):
+            client.get_host_keys().add(hostname, key.get_name(), key)
+            return
+        raise paramiko.SSHException(f"Unknown host key for {hostname}")
+
+
 def _can_connect(host: str, port: int = DEFAULT_SSH_PORT) -> bool:
     """Check if the connection to the provided host and port is possible."""
     try:
@@ -83,7 +101,7 @@
         logger.debug("Testing connection to host %s", host)
         client = paramiko.SSHClient()
         client.load_system_host_keys()
-        client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
+        client.set_missing_host_key_policy(CustomHostKeyPolicy())
         client.connect(host, port=port)
         client.close()
         logger.info("Can connect to host %s", host)