Merge pull request #759 from NVIDIA/am/nccl-k8s

Automatically install sshd for NCCL k8s workers if no available

Author: amaslenn

src/cloudai/_core/json_gen_strategy.py

@@ -53,7 +53,7 @@ def sanitize_k8s_job_name(self, job_name: str) -> str:
             str: The sanitized job name that complies with Kubernetes naming rules.
         """
         sanitized_name = job_name.lower()
-        sanitized_name = re.sub(r"[^a-z0-9.-]", "-", sanitized_name)
+        sanitized_name = re.sub(r"[^a-z0-9-]", "-", sanitized_name)
         sanitized_name = re.sub(r"^[^a-z0-9]+", "", sanitized_name)
         sanitized_name = re.sub(r"[^a-z0-9]+$", "", sanitized_name)
         return sanitized_name[:253]

src/cloudai/systems/kubernetes/kubernetes_runner.py

@@ -32,7 +32,7 @@ def _submit_test(self, tr: TestRun) -> KubernetesJob:
         job_name = tr.name.replace(".", "-").lower()
         job_spec = self.get_json_gen_strategy(self.system, tr).gen_json()
         job_kind = job_spec.get("kind", "").lower()
-        logging.info(f"Generated JSON string for test {tr.name}: {job_spec}")
+        logging.debug(f"Generated JSON string for test {tr.name}: {job_spec}")
 
         if self.mode == "run":
             k8s_system: KubernetesSystem = cast(KubernetesSystem, self.system)

src/cloudai/systems/kubernetes/kubernetes_system.py

@@ -166,6 +166,7 @@ def _is_mpijob_running(self, job_name: str) -> bool:
             assert isinstance(mpijob, dict)
             status: dict = cast(dict, mpijob.get("status", {}))
             conditions = status.get("conditions", [])
+            logging.debug(f"MPIJob '{job_name}': {conditions=} {status=}")
 
             # Consider an empty conditions list as running
             if not conditions:
@@ -681,7 +682,7 @@ def store_logs_for_job(self, job_name: str, output_dir: Path) -> None:
                     stdout_file.write(logs + "\n")
 
                 except lazy.k8s.client.ApiException as e:
-                    logging.error(f"Error retrieving logs for pod '{pod_name}': {e}")
+                    logging.debug(f"Error retrieving logs for pod '{pod_name}': {e}")
 
         logging.debug(f"All logs concatenated and saved to '{stdout_file_path}'")
 

src/cloudai/workloads/nccl_test/kubernetes_json_gen_strategy.py

@@ -16,25 +16,34 @@
 
 from typing import Any, Dict, List, Union, cast
 
+import yaml
+
 from cloudai.core import JsonGenStrategy
+from cloudai.systems.kubernetes import KubernetesSystem
 
 from .nccl import NCCLTestDefinition
 
 
 class NcclTestKubernetesJsonGenStrategy(JsonGenStrategy):
-    """JSON generation strategy for NCCL tests on Kubernetes systems."""
+    """
+    JSON generation strategy for NCCL tests on Kubernetes systems.
 
-    SSH_PORT: int = 2222
+    This strategy generates an MPIJob configuration for running NCCL tests.
+    """
 
     def gen_json(self) -> dict[Any, Any]:
-        return {
+        k8s_system = cast(KubernetesSystem, self.system)
+        job_name = self.sanitize_k8s_job_name(self.test_run.name)
+
+        deployment = {
             "apiVersion": "kubeflow.org/v2beta1",
             "kind": "MPIJob",
             "metadata": {
-                "name": self.sanitize_k8s_job_name("nccl-test"),
+                "name": job_name,
+                "namespace": k8s_system.default_namespace,
             },
             "spec": {
-                "slotsPerWorker": 1,
+                "slotsPerWorker": k8s_system.gpus_per_node,
                 "runPolicy": {"cleanPodPolicy": "Running"},
                 "mpiReplicaSpecs": {
                     "Launcher": self._create_launcher_spec(),
@@ -43,47 +52,53 @@ def gen_json(self) -> dict[Any, Any]:
             },
         }
 
-    def _create_launcher_spec(self) -> dict[str, Any]:
+        with open(self.test_run.output_path / "deployment.yaml", "w") as f:
+            yaml.dump(deployment, f)
+
+        return deployment
+
+    @property
+    def container_url(self) -> str:
         tdef: NCCLTestDefinition = cast(NCCLTestDefinition, self.test_run.test)
+        return tdef.cmd_args.docker_image_url.replace("#", "/")
+
+    def _create_launcher_spec(self) -> dict[str, Any]:
         env_vars = self._get_merged_env_vars()
         return {
             "replicas": 1,
             "template": {
                 "spec": {
-                    "hostNetwork": True,
                     "containers": [
                         {
-                            "image": tdef.cmd_args.docker_image_url,
+                            "image": self.container_url,
                             "name": "nccl-test-launcher",
                             "imagePullPolicy": "IfNotPresent",
                             "securityContext": {"privileged": True},
                             "env": self._generate_env_list(env_vars),
                             "command": ["/bin/bash", "-c"],
-                            "args": [self._generate_launcher_command(env_vars)],
+                            "args": [self._generate_launcher_command()],
+                            "resources": self._prepare_launcher_resources(),
                         }
                     ],
                 },
             },
         }
 
     def _create_worker_spec(self) -> dict[str, Any]:
-        tdef: NCCLTestDefinition = cast(NCCLTestDefinition, self.test_run.test)
         env_vars = self._get_merged_env_vars()
         return {
-            "replicas": self.test_run.num_nodes,
+            "replicas": self.test_run.nnodes,
             "template": {
                 "spec": {
-                    "hostNetwork": True,
                     "containers": [
                         {
-                            "image": tdef.cmd_args.docker_image_url,
+                            "image": self.container_url,
                             "name": "nccl-test-worker",
                             "imagePullPolicy": "IfNotPresent",
                             "securityContext": {"privileged": True},
-                            "ports": [{"containerPort": self.SSH_PORT, "name": "ssh"}],
                             "env": self._generate_env_list(env_vars),
-                            "command": ["/bin/bash"],
-                            "args": ["-c", f"/usr/sbin/sshd -p {self.SSH_PORT}; sleep infinity"],
+                            "command": ["/bin/bash", "-c"],
+                            "args": [self._generate_worker_command()],
                             "resources": self._prepare_worker_resources(),
                             "volumeMounts": [
                                 {"mountPath": "/dev/shm", "name": "dev-shm"},
@@ -95,31 +110,54 @@ def _create_worker_spec(self) -> dict[str, Any]:
             },
         }
 
+    def _generate_worker_command(self) -> str:
+        """
+        Generate command for worker pods that starts the SSH daemon.
+
+        If the SSH daemon is not installed, it will be installed and the SSH keys will be generated.
+        """
+        return """
+set -e
+if ! command -v sshd &> /dev/null; then
+    apt-get update && apt-get install -y --no-install-recommends openssh-server
+fi
+mkdir -p /var/run/sshd
+cat >> /etc/ssh/sshd_config << EOF
+PermitRootLogin yes
+PubkeyAuthentication yes
+StrictModes no
+EOF
+ssh-keygen -A
+exec /usr/sbin/sshd -D
+""".strip()
+
     def _get_merged_env_vars(self) -> dict[str, str | list[str]]:
         final_env_vars = self.system.global_env_vars.copy()
         final_env_vars.update(self.test_run.test.extra_env_vars)
         return final_env_vars
 
     def _generate_env_list(self, env_vars: Dict[str, Union[str, List[str]]]) -> List[Dict[str, str]]:
-        env_list = [{"name": "OMPI_ALLOW_RUN_AS_ROOT", "value": "1"}]
+        env_list = [
+            {"name": "OMPI_ALLOW_RUN_AS_ROOT", "value": "1"},
+            {"name": "OMPI_ALLOW_RUN_AS_ROOT_CONFIRM", "value": "1"},
+        ]
         for key, value in env_vars.items():
             if isinstance(value, list):
                 value = ",".join(value)
             env_list.append({"name": key, "value": value})
         return env_list
 
-    def _generate_mpi_args(self, env_vars: Dict[str, Union[str, List[str]]]) -> List[str]:
+    def _generate_mpi_args(self) -> List[str]:
+        k8s_system = cast(KubernetesSystem, self.system)
+        total_processes = self.test_run.nnodes * k8s_system.gpus_per_node
+
         mpi_args = [
-            "--allow-run-as-root",
-            f"--mca plm_rsh_args '-p {self.SSH_PORT}'",
-            "-c 2",
-            "-bind-to none -map-by slot",
-            "-mca btl tcp,self",
+            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'",
         ]
 
-        if "NCCL_SOCKET_IFNAME" in env_vars:
-            mpi_args.append(f"-mca btl_tcp_if_include {env_vars['NCCL_SOCKET_IFNAME']}")
-
         return mpi_args
 
     def _generate_nccl_args(self, cmd_args_dict: Dict[str, Any]) -> List[str]:
@@ -136,7 +174,7 @@ def _generate_extra_args(self, extra_cmd_args: Dict[str, str]) -> List[str]:
             extra_args.append(f"{key} {value}" if value else key)
         return extra_args
 
-    def _generate_launcher_command(self, env_vars: dict[str, str | list[str]]) -> str:
+    def _generate_launcher_command(self) -> str:
         tdef: NCCLTestDefinition = cast(NCCLTestDefinition, self.test_run.test)
         tdef_cmd_args = tdef.cmd_args
 
@@ -146,7 +184,7 @@ def _generate_launcher_command(self, env_vars: dict[str, str | list[str]]) -> st
 
         command_parts = [
             "mpirun",
-            " ".join(self._generate_mpi_args(env_vars)),
+            " ".join(self._generate_mpi_args()),
             tdef_cmd_args.subtest_name,
             " ".join(self._generate_nccl_args(cmd_args_dict)),
         ]
@@ -163,4 +201,6 @@ def _prepare_launcher_resources(self) -> Dict[str, Dict[str, str]]:
         }
 
     def _prepare_worker_resources(self) -> Dict[str, Dict[str, str]]:
-        return {"requests": {"nvidia.com/gpu": "8"}, "limits": {"nvidia.com/gpu": "8"}}
+        k8s_system = cast(KubernetesSystem, self.system)
+        gpu_count = str(k8s_system.gpus_per_node)
+        return {"requests": {"nvidia.com/gpu": gpu_count}, "limits": {"nvidia.com/gpu": gpu_count}}

tests/json_gen_strategy/test_nccl_kubernetes_json_gen_strategy.py

@@ -15,6 +15,8 @@
 # limitations under the License.
 
 
+from typing import cast
+
 import pytest
 
 from cloudai.core import TestRun
@@ -68,12 +70,13 @@ def json_gen_strategy(
         return NcclTestKubernetesJsonGenStrategy(kubernetes_system, test_run)
 
     def test_gen_json_basic_structure(self, basic_test_run: TestRun, k8s_system: KubernetesSystem) -> None:
-        json_payload = self.json_gen_strategy(k8s_system, basic_test_run).gen_json()
+        json_gen_strategy = self.json_gen_strategy(k8s_system, basic_test_run)
+        json_payload = json_gen_strategy.gen_json()
 
         assert json_payload["apiVersion"] == "kubeflow.org/v2beta1"
         assert json_payload["kind"] == "MPIJob"
-        assert json_payload["metadata"]["name"] == "nccl-test"
-        assert json_payload["spec"]["slotsPerWorker"] == 1
+        assert json_payload["metadata"]["name"] == json_gen_strategy.sanitize_k8s_job_name(basic_test_run.name)
+        assert json_payload["spec"]["slotsPerWorker"] == k8s_system.gpus_per_node
         assert json_payload["spec"]["runPolicy"]["cleanPodPolicy"] == "Running"
         assert "Launcher" in json_payload["spec"]["mpiReplicaSpecs"]
         assert "Worker" in json_payload["spec"]["mpiReplicaSpecs"]
@@ -83,7 +86,6 @@ def test_launcher_spec(self, basic_test_run: TestRun, k8s_system: KubernetesSyst
         launcher_spec = json_payload["spec"]["mpiReplicaSpecs"]["Launcher"]
 
         assert launcher_spec["replicas"] == 1
-        assert launcher_spec["template"]["spec"]["hostNetwork"] is True
 
         container = launcher_spec["template"]["spec"]["containers"][0]
         assert container["image"] == "fake_image_url"
@@ -93,22 +95,24 @@ def test_launcher_spec(self, basic_test_run: TestRun, k8s_system: KubernetesSyst
         assert container["command"] == ["/bin/bash", "-c"]
 
     def test_worker_spec(self, basic_test_run: TestRun, k8s_system: KubernetesSystem) -> None:
-        json_payload = self.json_gen_strategy(k8s_system, basic_test_run).gen_json()
+        json_gen_strategy = self.json_gen_strategy(k8s_system, basic_test_run)
+        json_payload = json_gen_strategy.gen_json()
         worker_spec = json_payload["spec"]["mpiReplicaSpecs"]["Worker"]
 
-        assert worker_spec["replicas"] == 2
-        assert worker_spec["template"]["spec"]["hostNetwork"] is True
+        assert worker_spec["replicas"] == basic_test_run.nnodes
 
         container = worker_spec["template"]["spec"]["containers"][0]
         assert container["image"] == "fake_image_url"
         assert container["name"] == "nccl-test-worker"
         assert container["imagePullPolicy"] == "IfNotPresent"
         assert container["securityContext"]["privileged"] is True
-        assert container["ports"][0] == {"containerPort": 2222, "name": "ssh"}
-        assert container["command"] == ["/bin/bash"]
-        assert container["args"] == ["-c", "/usr/sbin/sshd -p 2222; sleep infinity"]
+        assert container["command"] == ["/bin/bash", "-c"]
+        assert container["args"] == [json_gen_strategy._generate_worker_command()]
 
-        assert container["resources"] == {"requests": {"nvidia.com/gpu": "8"}, "limits": {"nvidia.com/gpu": "8"}}
+        assert container["resources"] == {
+            "requests": {"nvidia.com/gpu": str(k8s_system.gpus_per_node)},
+            "limits": {"nvidia.com/gpu": str(k8s_system.gpus_per_node)},
+        }
 
         assert container["volumeMounts"] == [{"mountPath": "/dev/shm", "name": "dev-shm"}]
         assert worker_spec["template"]["spec"]["volumes"] == [
@@ -127,18 +131,17 @@ def test_env_variables(self, test_run_with_env_vars: TestRun, k8s_system: Kubern
             assert env_dict["LIST_VAR"] == "item1,item2"
 
     def test_launcher_command_generation(self, test_run_with_extra_args: TestRun, k8s_system: KubernetesSystem) -> None:
-        json_payload = self.json_gen_strategy(k8s_system, test_run_with_extra_args).gen_json()
+        json_gen_strategy = self.json_gen_strategy(k8s_system, test_run_with_extra_args)
+        json_payload = json_gen_strategy.gen_json()
         launcher_args = json_payload["spec"]["mpiReplicaSpecs"]["Launcher"]["template"]["spec"]["containers"][0][
             "args"
         ][0]
+        nccl = cast(NCCLTestDefinition, json_gen_strategy.test_run.test)
 
         assert "mpirun" in launcher_args
-        assert "--allow-run-as-root" in launcher_args
-        assert "--mca plm_rsh_args '-p 2222'" in launcher_args
-        assert "-bind-to none -map-by slot" in launcher_args
-        assert "all_reduce_perf" in launcher_args
-        assert "--nthreads 4" in launcher_args
-        assert "--ngpus 2" in launcher_args
-        assert "--minbytes 32M" in launcher_args
-        assert "--maxbytes 64M" in launcher_args
-        assert "--extra-flag value" in launcher_args
+        assert f"-np {test_run_with_extra_args.nnodes * k8s_system.gpus_per_node}" in launcher_args
+        assert "-bind-to none" in launcher_args
+        assert "-mca plm_rsh_args '-o StrictHostKeyChecking=no -o UserKnownHostsFile=/dev/null'" in launcher_args
+        assert nccl.cmd_args.subtest_name in launcher_args
+        assert f"--nthreads {nccl.cmd_args.nthreads}" in launcher_args
+        assert f"--ngpus {nccl.cmd_args.ngpus}" in launcher_args