For ComputeInstance and AmlCompute update disableLocalAuth property based on ssh_public_access (Azure#39934)

* add disableLocalAuth for computeInstance

* fix disableLocalAuthAuth issue for amlCompute

* update compute instance

* update recordings

* temp changes

* Revert "temp changes"

This reverts commit 64e3c38.

* update recordings

* fix tests

Author: pdhotems

sdk/ml/azure-ai-ml/CHANGELOG.md

@@ -3,6 +3,7 @@
 ### Features Added
 
 ### Bugs Fixed
+ - Fix for compute Instance, disableLocalAuth property should be depend on ssh public access enabled.
 
 ## 1.26.0 (2025-03-11)
 

sdk/ml/azure-ai-ml/assets.json

@@ -2,5 +2,5 @@
   "AssetsRepo": "Azure/azure-sdk-assets",
   "AssetsRepoPrefixPath": "python",
   "TagPrefix": "python/ml/azure-ai-ml",
-  "Tag": "python/ml/azure-ai-ml_a2c955e6e2"
+  "Tag": "python/ml/azure-ai-ml_305b890d5b"
 }

sdk/ml/azure-ai-ml/azure/ai/ml/entities/_compute/aml_compute.py

@@ -251,7 +251,7 @@ def _to_rest_object(self) -> ComputeResource:
             ),
         )
         remote_login_public_access = "Enabled"
-        disableLocalAuth = not (self.ssh_public_access_enabled and self.ssh_settings is not None)
+        disableLocalAuth = not (self.ssh_settings)
         if self.ssh_public_access_enabled is not None:
             remote_login_public_access = "Enabled" if self.ssh_public_access_enabled else "Disabled"
 

sdk/ml/azure-ai-ml/azure/ai/ml/entities/_compute/compute_instance.py

@@ -280,12 +280,14 @@ def _to_rest_object(self) -> ComputeResource:
             subnet_resource = None
 
         ssh_settings = None
+        disable_local_auth = True
         if self.ssh_public_access_enabled is not None or self.ssh_settings is not None:
             ssh_settings = CiSShSettings()
             ssh_settings.ssh_public_access = "Enabled" if self.ssh_public_access_enabled else "Disabled"
             ssh_settings.admin_public_key = (
                 self.ssh_settings.ssh_key_value if self.ssh_settings and self.ssh_settings.ssh_key_value else None
             )
+            disable_local_auth = not self.ssh_public_access_enabled
 
         personal_compute_instance_settings = None
         if self.create_on_behalf_of:
@@ -330,6 +332,7 @@ def _to_rest_object(self) -> ComputeResource:
             description=self.description,
             compute_type=self.type,
             properties=compute_instance_prop,
+            disable_local_auth=disable_local_auth,
         )
         return ComputeResource(
             location=self.location,

sdk/ml/azure-ai-ml/tests/compute/e2etests/test_compute.py

@@ -30,6 +30,8 @@ def test_aml_compute_create_and_delete(self, client: MLClient, rand_compute_name
         assert compute_resource_get.name == compute_name
         assert compute_resource_get.tier == "dedicated"
         assert compute_resource_get.location == compute.location
+        assert compute_resource_get.ssh_public_access_enabled == True
+        assert compute_resource_get.ssh_settings.admin_username == "azureuser"
 
         compute_resource_get.idle_time_before_scale_down = 200
         compute_update_poller = client.compute.begin_update(compute_resource_get)
@@ -46,7 +48,6 @@ def test_aml_compute_create_and_delete(self, client: MLClient, rand_compute_name
         # so this is a preferred approach to assert
         assert isinstance(outcome, LROPoller)
 
-    @pytest.mark.skip(reason="not enough capacity")
     def test_compute_instance_create_and_delete(
         self, client: MLClient, rand_compute_name: Callable[[str], str]
     ) -> None:
@@ -65,20 +66,11 @@ def test_compute_instance_create_and_delete(
         assert isinstance(compute_resource_list, ItemPaged)
         compute_resource_get = client.compute.get(name=compute_name)
         assert compute_resource_get.name == compute_name
-        assert compute_resource_get.identity.type == "system_assigned"
         outcome = client.compute.begin_delete(name=compute_name)
         # the compute is getting deleted , but not waiting on the poller! so immediately returning
         # so this is a preferred approach to assert
         assert isinstance(outcome, LROPoller)
 
-    @pytest.mark.skipif(
-        condition=not is_live(),
-        reason=(
-            "Test takes 5 minutes in automation. "
-            "Already have unit tests verifying correct _restclient method is called. "
-            "Can be validated in live build only."
-        ),
-    )
     def test_compute_instance_stop_start_restart(
         self, client: MLClient, rand_compute_name: Callable[[str], str]
     ) -> None:

sdk/ml/azure-ai-ml/tests/compute/unittests/test_compute_entity.py

@@ -75,6 +75,8 @@ def test_compute_from_yaml(self):
         assert rest_intermediate.location == compute.location
         assert rest_intermediate.tags is not None
         assert rest_intermediate.tags["test"] == "true"
+        assert rest_intermediate.properties.disable_local_auth is False
+        assert rest_intermediate.properties.properties.remote_login_port_public_access == "Enabled"
 
         serializer = Serializer({"ComputeResource": ComputeResource})
         body = serializer.body(rest_intermediate, "ComputeResource")
@@ -98,6 +100,19 @@ def test_aml_compute_from_yaml_with_disable_public_access(self):
         assert rest_intermediate.properties.properties.enable_node_public_ip
         assert rest_intermediate.properties.disable_local_auth is True
         assert rest_intermediate.location == compute.location
+        assert rest_intermediate.properties.properties.remote_login_port_public_access == "NotSpecified"
+
+    def test_aml_compute_from_yaml_with_creds_and_disable_public_access(self):
+        compute: AmlCompute = load_compute("tests/test_configs/compute/compute-aml-no-identity.yaml")
+        compute.ssh_public_access_enabled = False
+
+        rest_intermediate = compute._to_rest_object()
+
+        assert rest_intermediate.properties.compute_type == "AmlCompute"
+        assert rest_intermediate.properties.properties.enable_node_public_ip
+        assert rest_intermediate.properties.disable_local_auth is False
+        assert rest_intermediate.location == compute.location
+        assert rest_intermediate.properties.properties.remote_login_port_public_access == "Disabled"
 
     def test_aml_compute_from_yaml_with_disable_public_access_when_no_sshSettings(self):
 
@@ -112,6 +127,7 @@ def test_aml_compute_from_yaml_with_disable_public_access_when_no_sshSettings(se
         assert rest_intermediate.properties.compute_type == "AmlCompute"
         assert rest_intermediate.properties.properties.enable_node_public_ip
         assert rest_intermediate.properties.disable_local_auth is True
+        assert rest_intermediate.properties.properties.remote_login_port_public_access == "Enabled"
         assert rest_intermediate.location == compute.location
 
     def test_compute_vm_from_yaml(self):
@@ -126,6 +142,7 @@ def test_compute_vm_from_yaml(self):
         assert compute.ssh_settings.ssh_private_key_file == "tests/test_configs/compute/ssh_fake_key.txt"
 
         rest_intermediate = compute._to_rest_object()
+        assert rest_intermediate.properties.compute_type == "VirtualMachine"
         assert rest_intermediate.properties.resource_id == resource_id
         assert rest_intermediate.properties.properties.ssh_port == 8888
         assert rest_intermediate.properties.properties.administrator_account.password == "azureuserpassword"
@@ -183,6 +200,7 @@ def test_compute_instance_load_from_rest(self):
         compute_instance3: ComputeInstance = load_compute(
             source="tests/test_configs/compute/compute-ci-defaults-unit.yaml",
         )._to_rest_object()
+        assert compute_instance3.properties.compute_type == "ComputeInstance"
         assert compute_instance3.properties.properties.enable_sso is True
         assert compute_instance3.properties.properties.enable_root_access is True
         assert compute_instance3.properties.properties.enable_os_patching is False
@@ -291,14 +309,28 @@ def test_compute_instance_uai_from_yaml(self):
         )
         assert compute_from_rest.ssh_public_access_enabled == False
 
-    def test_compute_instance_sai_from_yaml(self):
+    def test_compute_instace_to_rest(self):
         compute: ComputeInstance = load_compute("tests/test_configs/compute/compute-ci.yaml")
+        compute.ssh_public_access_enabled = True
+
+        compute_rest = compute._to_rest_object()
+
+        assert compute_rest.location == compute.location
+        assert compute_rest.properties.properties.ssh_settings.ssh_public_access == "Enabled"
+        assert compute_rest.properties.properties.ssh_settings.admin_public_key == compute.ssh_settings.ssh_key_value
+        assert compute_rest.properties.description == compute.description
+        assert compute_rest.properties.compute_type == "ComputeInstance"
+        assert compute_rest.properties.disable_local_auth == False
+
+    def test_compute_instance_sai_from_yaml(self):
+        compute: ComputeInstance = load_compute("tests/test_configs/compute/compute-ci-id-sys-assigned.yaml")
         assert compute.name == "banchci"
         assert compute.type == "computeinstance"
         assert compute.identity.type == "system_assigned"
 
         compute_resource = compute._to_rest_object()
         assert compute_resource.identity.type == "SystemAssigned"
+        assert compute_resource.properties.disable_local_auth == True
 
         compute_from_rest = Compute._from_rest_object(compute_resource)
         assert compute_from_rest.type == "computeinstance"

sdk/ml/azure-ai-ml/tests/test_configs/compute/compute-aml-no-identity.yaml

@@ -2,7 +2,7 @@ name: banchaml
 type: amlcompute
 tier: dedicated
 description: some_desc_aml
-size: Standard_DS2_v2
+size: Standard_DS3_v2
 location: eastus
 ssh_public_access_enabled: true
 ssh_settings:

sdk/ml/azure-ai-ml/tests/test_configs/compute/compute-ci-id-sys-assigned.yaml

@@ -0,0 +1,10 @@
+name: banchci
+type: computeinstance
+description: some_desc_ci
+size: Standard_DS3_v2
+ssh_public_access_enabled: false
+ssh_settings:
+   admin_username: azureuser
+   ssh_key_value: ssh-rsa AAAAB3NzaC1yc2EAAAADAQABAAABAQDWm/4TTHMZdzZVJcob2aFRWDmIyJLxk028AKh7K03RDgR8bz/Knd9DgB2V3sipjY9wYJ1U9YbsUoAt24+CWZFpnoB79J6vaoqwUb7c6nGkaNrWoGZKI+v26GOX8O2MUjjOyBuPEinJtQ432J4affHfeRI+1iDQXuwFUKhNbRVpxh2h9otXF+J1UvSUaPYggS7Iivyha/x8HJzFcNnIPrAZkPiT/Nb/Qk7FyoFTEw64cIl1ByvmF3ewSOeVXKDpb2d4vrSDTVmXFKRrWduhM3sHO5dckREKc+tSQ0M+SpfBqaBgRTAie4jFVSIiSHCvL2BRQXBMRiyPgNlpiJWTPLp/ administrator@MININT-7IP9G6S
+identity:
+   type: system_assigned

sdk/ml/azure-ai-ml/tests/test_configs/compute/compute-ci.yaml

@@ -1,12 +1,8 @@
 name: banchci
 type: computeinstance
 description: some_desc_ci
-location: eastus
-
-size: Standard_DS2_v2
-
-ssh_public_access_enabled: False
+size: Standard_DS3_v2
+ssh_public_access_enabled: false
 ssh_settings:
+   admin_username: azureuser
    ssh_key_value: ssh-rsa AAAAB3NzaC1yc2EAAAADAQABAAABAQDWm/4TTHMZdzZVJcob2aFRWDmIyJLxk028AKh7K03RDgR8bz/Knd9DgB2V3sipjY9wYJ1U9YbsUoAt24+CWZFpnoB79J6vaoqwUb7c6nGkaNrWoGZKI+v26GOX8O2MUjjOyBuPEinJtQ432J4affHfeRI+1iDQXuwFUKhNbRVpxh2h9otXF+J1UvSUaPYggS7Iivyha/x8HJzFcNnIPrAZkPiT/Nb/Qk7FyoFTEw64cIl1ByvmF3ewSOeVXKDpb2d4vrSDTVmXFKRrWduhM3sHO5dckREKc+tSQ0M+SpfBqaBgRTAie4jFVSIiSHCvL2BRQXBMRiyPgNlpiJWTPLp/ administrator@MININT-7IP9G6S
-identity:
-   type: system_assigned