Repository: microsoft/hi-ml

Issue:
Bug: Jobs submitted to a Kubernetes compute target don't use the GPUs
Jobs that work have this section in the run definition:
```
        "cmk8sCompute": {
            "configuration": {
                "gpu_count": 16,
                "preemption_allowed": false,
                "job_priority": 200,
                "enable_ssh": true,
                "ssh_public_key": "
            }
        },
```



Discussion:

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
@@ -547,12 +547,10 @@ def submit_run_v2(workspace: Optional[Workspace],
 
     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
+            # On AML managed compute, we can set distribution to None for single node jobs.
+            # However, on Kubernetes compute, single node jobs don't see any GPUs. GPUs are visible for MpiDistribution
+            # jobs, so we set MpiDistribution even for single node jobs.
+            distribution: Union[MpiDistribution, PyTorchDistribution] = MpiDistribution(process_count_per_instance=1)
         else:
             distribution = PyTorchDistribution(process_count_per_instance=pytorch_processes_per_node)
         return command(
@@ -567,7 +565,7 @@ def create_command_job(cmd: str) -> Command:
             shm_size=docker_shm_size,
             display_name=display_name,
             instance_count=num_nodes,
-            distribution=distribution,  # type: ignore
+            distribution=distribution,
         )
 
     if hyperparam_args:


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
@@ -798,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=None,
+                distribution=MpiDistribution(process_count_per_instance=1),
                 instance_count=1
             )
 
@@ -858,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=None,
+                distribution=MpiDistribution(process_count_per_instance=1),
                 instance_count=1
             )
 
@@ -1867,7 +1867,7 @@ def test_submit_to_azure_v2_distributed() -> None:
                 mock_command.assert_called_once()
                 _, call_kwargs = mock_command.call_args
                 assert call_kwargs.get("instance_count") == 1
-                assert call_kwargs.get("distribution") is None
+                assert call_kwargs.get("distribution") == MpiDistribution(process_count_per_instance=1)
 
             with pytest.raises(ValueError, match="num_nodes must be >= 1"):
                 _ = himl.submit_to_azure_if_needed(
@@ -1925,5 +1925,4 @@ def test_submit_to_azure_v2_distributed() -> None:
                 _, 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
+                assert distribution == MpiDistribution(process_count_per_instance=1)