Use host network by default

+ make ssh port usage more reliable

Author: amaslenn

src/cloudai/workloads/nccl_test/kubernetes_json_gen_strategy.py

@@ -31,6 +31,10 @@ class NcclTestKubernetesJsonGenStrategy(JsonGenStrategy):
     This strategy generates an MPIJob configuration for running NCCL tests.
     """
 
+    @property
+    def ssh_port(self) -> int:
+        return 2222
+
     def gen_json(self) -> dict[Any, Any]:
         k8s_system = cast(KubernetesSystem, self.system)
         job_name = self.sanitize_k8s_job_name(self.test_run.name)
@@ -68,6 +72,7 @@ def _create_launcher_spec(self) -> dict[str, Any]:
             "replicas": 1,
             "template": {
                 "spec": {
+                    "hostNetwork": True,
                     "containers": [
                         {
                             "image": self.container_url,
@@ -90,10 +95,12 @@ def _create_worker_spec(self) -> dict[str, Any]:
             "replicas": self.test_run.nnodes,
             "template": {
                 "spec": {
+                    "hostNetwork": True,
                     "containers": [
                         {
                             "image": self.container_url,
                             "name": "nccl-test-worker",
+                            "ports": [{"containerPort": self.ssh_port, "name": "ssh"}],
                             "imagePullPolicy": "IfNotPresent",
                             "securityContext": {"privileged": True},
                             "env": self._generate_env_list(env_vars),
@@ -116,8 +123,8 @@ def _generate_worker_command(self) -> str:
 
         If the SSH daemon is not installed, it will be installed and the SSH keys will be generated.
         """
-        return """
-set -e
+        return f"""
+set -ex
 if ! command -v sshd &> /dev/null; then
     apt-get update && apt-get install -y --no-install-recommends openssh-server
 fi
@@ -126,9 +133,11 @@ def _generate_worker_command(self) -> str:
 PermitRootLogin yes
 PubkeyAuthentication yes
 StrictModes no
+Port {self.ssh_port}
 EOF
 ssh-keygen -A
-exec /usr/sbin/sshd -D
+service ssh restart
+sleep infinity
 """.strip()
 
     def _get_merged_env_vars(self) -> dict[str, str | list[str]]:
@@ -154,8 +163,8 @@ def _generate_mpi_args(self) -> List[str]:
         mpi_args = [
             f"-np {total_processes}",
             "-bind-to none",
-            # Disable strict host key checking for SSH
-            "-mca plm_rsh_args '-o StrictHostKeyChecking=no -o UserKnownHostsFile=/dev/null'",
+            # Disable strict host key checking for SSH and ensure correct port is used
+            f"-mca plm_rsh_args '-p {self.ssh_port} -o StrictHostKeyChecking=no -o UserKnownHostsFile=/dev/null'",
         ]
 
         return mpi_args