Filter policy by algo-

Author: sage-maker

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

@@ -18,6 +18,7 @@
 import time
 from typing import List
 
+import paramiko
 from utils import SM_EFA_NCCL_INSTANCES, SM_EFA_RDMA_INSTANCES, get_python_executable, logger
 
 FINISHED_STATUS_FILE = "/tmp/done.algo-1"
@@ -74,6 +75,24 @@ def start_sshd_daemon():
     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:
@@ -82,7 +101,7 @@ def _can_connect(host: str, port: int = DEFAULT_SSH_PORT) -> bool:
         logger.debug("Testing connection to host %s", host)
         client = paramiko.SSHClient()
         client.load_system_host_keys()
-        client.set_missing_host_key_policy(paramiko.RejectPolicy())
+        client.set_missing_host_key_policy(CustomHostKeyPolicy())
         client.connect(host, port=port)
         client.close()
         logger.info("Can connect to host %s", host)