Repository: microsoft/hi-ml

Issue:
Bug: Submitting a single node job via AML v2 SDK does not work
Submitting HelloWorld via v2 hangs, jobs are not even starting the main script


Discussion:
Symptoms: view "All properties" on the stuck jobs, it says:
```
    "runDefinition": {
        "command": ...
        "framework": "PyTorch",
        "communicator": "Nccl",
```

Previous working jobs had this:
```
    "runDefinition": {
        "command": ...
        "framework": "Python",
        "communicator": "None",
```
But also: 
```
        "environmentVariables": {
            "JOB_EXECUTION_MODE": "Basic"
        },
```
A simple script that does not do any GPU communication can be submitted fine with the v2 SDK. However, running HelloWorld in hi-ml only gets to the initialization of the first worker:
```
Global seed set to 42
Initializing distributed: GLOBAL_RANK: 0, MEMBER: 1/2
```


Resolving patch:
diff --git a/hi-ml-azure/src/health_azure/himl.py b/hi-ml-azure/src/health_azure/himl.py
--- a/hi-ml-azure/src/health_azure/himl.py
+++ b/hi-ml-azure/src/health_azure/himl.py
@@ -23,6 +23,7 @@
 from azure.ai.ml.constants import AssetTypes, InputOutputModes
 from azure.ai.ml.entities import Data, Job, Command, Sweep
 from azure.ai.ml.entities import Environment as EnvironmentV2
+from azure.ai.ml.entities._job.distribution import MpiDistribution, PyTorchDistribution
 
 from azure.ai.ml.sweep import Choice
 from azureml._base_sdk_common import user_agent
@@ -466,12 +467,11 @@ def submit_run_v2(workspace: Optional[Workspace],
                   tags: Optional[Dict[str, str]] = None,
                   docker_shm_size: str = "",
                   wait_for_completion: bool = False,
-                  wait_for_completion_show_output: bool = False,
                   workspace_config_path: Optional[PathOrString] = None,
                   ml_client: Optional[MLClient] = None,
                   hyperparam_args: Optional[Dict[str, Any]] = None,
                   num_nodes: int = 1,
-                  processes_per_node: int = 1
+                  pytorch_processes_per_node: Optional[int] = None,
                   ) -> Job:
     """
     Starts a v2 AML Job on a given workspace by submitting a command
@@ -493,14 +493,14 @@ def submit_run_v2(workspace: Optional[Workspace],
     :param docker_shm_size: The Docker shared memory size that should be used when creating a new Docker image.
     :param wait_for_completion: If False (the default) return after the run is submitted to AzureML, otherwise wait for
         the completion of this run (if True).
-    :param wait_for_completion_show_output: If wait_for_completion is True this parameter indicates whether to show the
-        run output on sys.stdout.
     :param workspace_config_path: If not provided with an AzureML Workspace, then load one given the information in this
         config
     :param ml_client: An Azure MLClient object for interacting with Azure resources.
     :param hyperparam_args: A dictionary of hyperparameter search args to pass into a sweep job.
-    :param num_nodes: The number of nodes to use for the job in AzureML.
-    :param processes_per_node: The number of processes per node to use for the job in AzureML.
+    :param num_nodes: The number of nodes to use for the job in AzureML. The value must be 1 or greater.
+    :param pytorch_processes_per_node: For plain PyTorch multi-GPU processing: The number of processes per node.
+        If supplied, it will run a command job with the "pytorch" framework (rather than "Python"), and using "nccl"
+        as the communication backend.
     :return: An AzureML Run object.
     """
     if ml_client is None:
@@ -538,17 +538,24 @@ def submit_run_v2(workspace: Optional[Workspace],
     display_name = get_display_name_v2(tags)
 
     # number of nodes and processes per node cannot be less than one
+    if num_nodes < 1:
+        raise ValueError("num_nodes must be >= 1")
     num_nodes = num_nodes if num_nodes >= 1 else 1
-    processes_per_node = processes_per_node if processes_per_node >= 1 else 1
-
-    if hyperparam_args:
-        param_sampling = hyperparam_args[PARAM_SAMPLING_ARG]
-
-        for sample_param, choices in param_sampling.items():
-            input_datasets_v2[sample_param] = choices.values[0]
-            cmd += f" --{sample_param}=" + "${{inputs." + sample_param + "}}"
-
-        command_job = command(
+    if pytorch_processes_per_node is not None:
+        if pytorch_processes_per_node < 1:
+            raise ValueError("pytorch_processes_per_node must be >= 1")
+
+    def create_command_job(cmd: str) -> Command:
+        if pytorch_processes_per_node is None:
+            if num_nodes > 1:
+                distribution: Any = MpiDistribution(process_count_per_instance=1)
+            else:
+                # An empty dictionary for single node jobs would be in line with the type annotations on the
+                # 'command' function, but this is not recognized by the SDK. So we need to pass None instead.
+                distribution = None
+        else:
+            distribution = PyTorchDistribution(process_count_per_instance=pytorch_processes_per_node)
+        return command(
             code=str(snapshot_root_directory),
             command=cmd,
             inputs=input_datasets_v2,
@@ -560,12 +567,18 @@ def submit_run_v2(workspace: Optional[Workspace],
             shm_size=docker_shm_size,
             display_name=display_name,
             instance_count=num_nodes,
-            distribution={
-                "type": "PyTorch",
-                "process_count_per_instance": processes_per_node,
-            },
+            distribution=distribution,  # type: ignore
         )
 
+    if hyperparam_args:
+        param_sampling = hyperparam_args[PARAM_SAMPLING_ARG]
+
+        for sample_param, choices in param_sampling.items():
+            input_datasets_v2[sample_param] = choices.values[0]
+            cmd += f" --{sample_param}=" + "${{inputs." + sample_param + "}}"
+
+        command_job = create_command_job(cmd)
+
         del hyperparam_args[PARAM_SAMPLING_ARG]
         # override command with parameter expressions
         command_job = command_job(
@@ -583,23 +596,7 @@ def submit_run_v2(workspace: Optional[Workspace],
         job_to_submit.set_limits(max_total_trials=hyperparam_args.get(MAX_TOTAL_TRIALS_ARG, None))
 
     else:
-        job_to_submit = command(
-            code=str(snapshot_root_directory),
-            command=cmd,
-            inputs=input_datasets_v2,
-            outputs=output_datasets_v2,
-            environment=environment.name + "@latest",
-            compute=compute_target,
-            experiment_name=experiment_name,
-            tags=tags or {},
-            shm_size=docker_shm_size,
-            display_name=display_name,
-            instance_count=num_nodes,
-            distribution={
-                "type": "PyTorch",
-                "process_count_per_instance": processes_per_node,
-            },
-        )
+        job_to_submit = create_command_job(cmd)
 
     returned_job = ml_client.jobs.create_or_update(job_to_submit)
     logging.info(f"URL to job: {returned_job.services['Studio'].endpoint}")  # type: ignore
@@ -775,7 +772,7 @@ def submit_to_azure_if_needed(  # type: ignore
         hyperdrive_config: Optional[HyperDriveConfig] = None,
         hyperparam_args: Optional[Dict[str, Any]] = None,
         strictly_aml_v1: bool = False,
-        processes_per_node_v2: int = 1,
+        pytorch_processes_per_node_v2: Optional[int] = None,
 ) -> AzureRunInfo:  # pragma: no cover
     """
     Submit a folder to Azure, if needed and run it.
@@ -824,7 +821,13 @@ def submit_to_azure_if_needed(  # type: ignore
         will also register the data in this folder as an AzureML dataset.
     :param output_datasets: The script will create a temporary folder when running in AzureML, and while the job writes
         data to that folder, upload it to blob storage, in the data store.
-    :param num_nodes: The number of nodes to use in distributed training on AzureML.
+    :param num_nodes: The number of nodes to use in distributed training on AzureML. When using a value > 1, multiple
+        nodes in AzureML will be started. If `pytorch_processes_per_node_v2=None`, the job will be submitted
+        as a multi-node MPI job, with 1 process per node. This is suitable for PyTorch Lightning jobs.
+        If `pytorch_processes_per_node_v2` is not None,
+        a job with framework "PyTorch" and communication backend "nccl" will be started.
+        `pytorch_processes_per_node_v2` will guide the number of processes per node. This is suitable for plain PyTorch
+        training jobs without the use of frameworks like PyTorch Lightning.
     :param wait_for_completion: If False (the default) return after the run is submitted to AzureML, otherwise wait for
         the completion of this run (if True).
     :param wait_for_completion_show_output: If wait_for_completion is True this parameter indicates whether to show the
@@ -834,7 +837,9 @@ def submit_to_azure_if_needed(  # type: ignore
         will be triggered if the commandline flag '--azureml' is present in sys.argv
     :param hyperdrive_config: A configuration object for Hyperdrive (hyperparameter search).
     :param strictly_aml_v1: If True, use Azure ML SDK v1. Otherwise, attempt to use Azure ML SDK v2.
-    :param processes_per_node_v2: The number of processes per node to use in distributed training on AzureML SDK v2.
+    :param pytorch_processes_per_node_v2: For plain PyTorch multi-GPU processing: The number of processes per node. This
+        is only supported with AML SDK v2, and ignored in v1. If supplied, the job will be submitted as using the
+        "pytorch" framework (rather than "Python"), and using "nccl" as the communication backend.
     :return: If the script is submitted to AzureML then we terminate python as the script should be executed in AzureML,
         otherwise we return a AzureRunInfo object.
     """
@@ -977,10 +982,9 @@ def submit_to_azure_if_needed(  # type: ignore
                                 tags=tags,
                                 docker_shm_size=docker_shm_size,
                                 wait_for_completion=wait_for_completion,
-                                wait_for_completion_show_output=wait_for_completion_show_output,
                                 hyperparam_args=hyperparam_args,
                                 num_nodes=num_nodes,
-                                processes_per_node=processes_per_node_v2,
+                                pytorch_processes_per_node=pytorch_processes_per_node_v2,
                                 )
             if after_submission is not None:
                 after_submission(job, ml_client)  # type: ignore
diff --git a/hi-ml/src/health_ml/runner.py b/hi-ml/src/health_ml/runner.py
--- a/hi-ml/src/health_ml/runner.py
+++ b/hi-ml/src/health_ml/runner.py
@@ -268,7 +268,6 @@ def after_submission_hook(azure_run: Run, *args: Any) -> None:
                 after_submission=after_submission_hook,
                 tags=self.additional_run_tags(script_params),
                 strictly_aml_v1=self.experiment_config.strictly_aml_v1,
-                processes_per_node_v2=self.lightning_container.num_gpus_per_node(),
             )
         else:
             azure_run_info = submit_to_azure_if_needed(


Test patch:
diff --git a/hi-ml-azure/testazure/testazure/test_himl.py b/hi-ml-azure/testazure/testazure/test_himl.py
--- a/hi-ml-azure/testazure/testazure/test_himl.py
+++ b/hi-ml-azure/testazure/testazure/test_himl.py
@@ -26,6 +26,7 @@
 from azure.ai.ml import Input, Output, MLClient
 from azure.ai.ml.constants import AssetTypes, InputOutputModes
 from azure.ai.ml.entities import Data, Job
+from azure.ai.ml.entities._job.distribution import MpiDistribution, PyTorchDistribution
 from azure.ai.ml.sweep import Choice
 from azureml._restclient.constants import RunStatus
 from azureml.core import ComputeTarget, Environment, RunConfiguration, ScriptRunConfig, Workspace
@@ -797,7 +798,7 @@ def _mock_sweep(*args: Any, **kwargs: Any) -> MagicMock:
                 tags=dummy_tags,
                 shm_size=dummy_docker_shm_size,
                 display_name=dummy_display_name,
-                distribution={"type": "PyTorch", "process_count_per_instance": 1},
+                distribution=None,
                 instance_count=1
             )
 
@@ -857,7 +858,7 @@ def _mock_sweep(*args: Any, **kwargs: Any) -> MagicMock:
                 tags=dummy_tags,
                 shm_size=dummy_docker_shm_size,
                 display_name=dummy_display_name,
-                distribution={"type": "PyTorch", "process_count_per_instance": 1},
+                distribution=None,
                 instance_count=1
             )
 
@@ -1793,7 +1794,6 @@ def test_submit_to_azure_v2_distributed() -> None:
                 # If num_nodes and processes_per_node are not passed to submit_to_azure_if_needed, then
                 # submit_run_v2 should receive the default values of 1 and 1 respectively
                 expected_num_nodes = 1
-                expected_processes_per_node = 1
                 _ = himl.submit_to_azure_if_needed(
                     workspace_config_file="mockconfig.json",
                     snapshot_root_directory="dummy",
@@ -1804,7 +1804,7 @@ def test_submit_to_azure_v2_distributed() -> None:
                 _, call_kwargs = mock_submit_run_v2.call_args
                 print(call_kwargs)
                 assert call_kwargs.get("num_nodes") == expected_num_nodes
-                assert call_kwargs.get("processes_per_node") == expected_processes_per_node
+                assert call_kwargs.get("pytorch_processes_per_node") is None
 
                 # If num_nodes and processes_per_node are both passed in to submit_to_azure_if_needed,
                 # they should be passed through to submit_run_v2
@@ -1816,14 +1816,13 @@ def test_submit_to_azure_v2_distributed() -> None:
                     submit_to_azureml=True,
                     strictly_aml_v1=False,
                     num_nodes=num_nodes,
-                    processes_per_node_v2=processes_per_node
+                    pytorch_processes_per_node_v2=processes_per_node
                 )
                 _, call_kwargs = mock_submit_run_v2.call_args
                 assert call_kwargs.get("num_nodes") == num_nodes
-                assert call_kwargs.get("processes_per_node") == processes_per_node
+                assert call_kwargs.get("pytorch_processes_per_node") == processes_per_node
 
-            # If num_nodes and/or processes_per_node are passed in to submit_to_azure_if_needed with
-            # values less than one, they should be updated to values of 1 before the command is created
+            # Single node job: The "distribution" argument of "command" should be set to None
             with patch("health_azure.himl.command") as mock_command:
                 _ = himl.submit_to_azure_if_needed(
                     workspace_config_file="mockconfig.json",
@@ -1831,10 +1830,67 @@ def test_submit_to_azure_v2_distributed() -> None:
                     snapshot_root_directory=Path.cwd(),
                     submit_to_azureml=True,
                     strictly_aml_v1=False,
-                    num_nodes=-1,
-                    processes_per_node_v2=0
                 )
                 mock_command.assert_called_once()
                 _, call_kwargs = mock_command.call_args
                 assert call_kwargs.get("instance_count") == 1
-                assert call_kwargs.get("distribution").get("process_count_per_instance") == 1
+                assert call_kwargs.get("distribution") is None
+
+            with pytest.raises(ValueError, match="num_nodes must be >= 1"):
+                _ = himl.submit_to_azure_if_needed(
+                    workspace_config_file="mockconfig.json",
+                    entry_script=Path(__file__),
+                    snapshot_root_directory=Path.cwd(),
+                    submit_to_azureml=True,
+                    strictly_aml_v1=False,
+                    num_nodes=0
+                )
+
+            with pytest.raises(ValueError, match="pytorch_processes_per_node must be >= 1"):
+                _ = himl.submit_to_azure_if_needed(
+                    workspace_config_file="mockconfig.json",
+                    entry_script=Path(__file__),
+                    snapshot_root_directory=Path.cwd(),
+                    submit_to_azureml=True,
+                    strictly_aml_v1=False,
+                    num_nodes=1,
+                    pytorch_processes_per_node_v2=0,
+                )
+
+            # When specifying the number of pytorch processes, a PyTorchDistribution job should be created
+            num_nodes = 2
+            num_processes = 3
+            with patch("health_azure.himl.command") as mock_command:
+                _ = himl.submit_to_azure_if_needed(
+                    workspace_config_file="mockconfig.json",
+                    entry_script=Path(__file__),
+                    snapshot_root_directory=Path.cwd(),
+                    submit_to_azureml=True,
+                    strictly_aml_v1=False,
+                    num_nodes=num_nodes,
+                    pytorch_processes_per_node_v2=num_processes
+                )
+                mock_command.assert_called_once()
+                _, call_kwargs = mock_command.call_args
+                assert call_kwargs.get("instance_count") == num_nodes
+                distribution = call_kwargs.get("distribution")
+                assert isinstance(distribution, PyTorchDistribution)
+                # pytorch_processes_per_node_v2=0 should be updated to 1
+                assert distribution.process_count_per_instance == num_processes
+
+            # If num_nodes is specified without processes_per_node, an MPI job should be created.
+            with patch("health_azure.himl.command") as mock_command:
+                _ = himl.submit_to_azure_if_needed(
+                    workspace_config_file="mockconfig.json",
+                    entry_script=Path(__file__),
+                    snapshot_root_directory=Path.cwd(),
+                    submit_to_azureml=True,
+                    strictly_aml_v1=False,
+                    num_nodes=num_nodes,
+                )
+                mock_command.assert_called_once()
+                _, call_kwargs = mock_command.call_args
+                assert call_kwargs.get("instance_count") == num_nodes
+                distribution = call_kwargs.get("distribution")
+                assert isinstance(distribution, MpiDistribution)
+                assert distribution.process_count_per_instance == 1