Compute instance new params (#33605)

Author: nemanjarajic

sdk/ml/azure-ai-ml/azure/ai/ml/_ml_client.py

@@ -452,7 +452,7 @@ def __init__(
         self._compute = ComputeOperations(
             self._operation_scope,
             self._operation_config,
-            self._service_client_10_2022_preview,
+            self._service_client_08_2023_preview,
             **app_insights_handler_kwargs,
         )
         self._operation_container.add(AzureMLResourceType.COMPUTE, self._compute)

sdk/ml/azure-ai-ml/azure/ai/ml/_schema/compute/compute_instance.py

@@ -71,3 +71,7 @@ class ComputeInstanceSchema(ComputeSchema):
     enable_node_public_ip = fields.Bool(
         metadata={"description": "Enable or disable node public IP address provisioning."}
     )
+    enable_sso = fields.Bool(metadata={"description": "Enable or disable single sign-on for the compute instance."})
+    enable_root_access = fields.Bool(
+        metadata={"description": "Enable or disable root access for the compute instance."}
+    )

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

@@ -10,10 +10,10 @@
 from typing import Dict, List, Optional
 
 from azure.ai.ml._restclient.v2022_10_01_preview.models import AssignedUser
-from azure.ai.ml._restclient.v2022_10_01_preview.models import ComputeInstance as CIRest
-from azure.ai.ml._restclient.v2022_10_01_preview.models import ComputeInstanceProperties
-from azure.ai.ml._restclient.v2022_10_01_preview.models import ComputeInstanceSshSettings as CiSShSettings
-from azure.ai.ml._restclient.v2022_10_01_preview.models import (
+from azure.ai.ml._restclient.v2023_08_01_preview.models import ComputeInstance as CIRest
+from azure.ai.ml._restclient.v2023_08_01_preview.models import ComputeInstanceProperties
+from azure.ai.ml._restclient.v2023_08_01_preview.models import ComputeInstanceSshSettings as CiSShSettings
+from azure.ai.ml._restclient.v2023_08_01_preview.models import (
     ComputeResource,
     PersonalComputeInstanceSettings,
     ResourceId,
@@ -160,6 +160,10 @@ class ComputeInstance(Compute):
     :type setup_scripts: Optional[~azure.ai.ml.entities.SetupScripts]
     :param custom_applications: List of custom applications and their endpoints for the compute instance.
     :type custom_applications: Optional[List[~azure.ai.ml.entities.CustomApplications]]
+    :param enable_sso: Enable or disable single sign-on. Defaults to True.
+    :type enable_sso: bool
+    :param enable_root_access: Enable or disable root access. Defaults to True.
+    :type enable_root_access: bool
 
     .. admonition:: Example:
 
@@ -189,6 +193,8 @@ def __init__(
         setup_scripts: Optional[SetupScripts] = None,
         enable_node_public_ip: bool = True,
         custom_applications: Optional[List[CustomApplications]] = None,
+        enable_sso: bool = True,
+        enable_root_access: bool = True,
         **kwargs,
     ) -> None:
         kwargs[TYPE] = ComputeType.COMPUTEINSTANCE
@@ -215,6 +221,8 @@ def __init__(
         self.idle_time_before_shutdown_minutes = idle_time_before_shutdown_minutes
         self.setup_scripts = setup_scripts
         self.enable_node_public_ip = enable_node_public_ip
+        self.enable_sso = enable_sso
+        self.enable_root_access = enable_root_access
         self.custom_applications = custom_applications
         self.subnet = None
 
@@ -295,6 +303,8 @@ def _to_rest_object(self) -> ComputeResource:
             personal_compute_instance_settings=personal_compute_instance_settings,
             idle_time_before_shutdown=idle_time_before_shutdown,
             enable_node_public_ip=self.enable_node_public_ip,
+            enable_sso=self.enable_sso,
+            enable_root_access=self.enable_root_access,
         )
         compute_instance_prop.schedules = self.schedules._to_rest_object() if self.schedules else None
         compute_instance_prop.setup_scripts = self.setup_scripts._to_rest_object() if self.setup_scripts else None
@@ -432,6 +442,12 @@ def _load_from_rest(cls, rest_obj: ComputeResource) -> "ComputeInstance":
             if (prop.properties and prop.properties.enable_node_public_ip is not None)
             else True,
             custom_applications=custom_applications,
+            enable_sso=prop.properties.enable_sso
+            if (prop.properties and prop.properties.enable_sso is not None)
+            else True,
+            enable_root_access=prop.properties.enable_root_access
+            if (prop.properties and prop.properties.enable_root_access is not None)
+            else True,
         )
         return response
 

sdk/ml/azure-ai-ml/azure/ai/ml/operations/_compute_operations.py

@@ -6,7 +6,7 @@
 
 from typing import Any, Dict, Iterable, Optional
 
-from azure.ai.ml._restclient.v2023_02_01_preview import AzureMachineLearningWorkspaces as ServiceClient022023Preview
+from azure.ai.ml._restclient.v2023_08_01_preview import AzureMachineLearningWorkspaces as ServiceClient022023Preview
 from azure.ai.ml._scope_dependent_operations import OperationConfig, OperationScope, _ScopeDependentOperations
 from azure.ai.ml._telemetry import ActivityType, monitor_with_activity
 from azure.ai.ml._utils._logger_utils import OpsLogger

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

@@ -8,7 +8,7 @@
 from test_utilities.utils import verify_entity_load_and_dump
 
 from azure.ai.ml import load_compute
-from azure.ai.ml._restclient.v2022_10_01_preview.models import ComputeResource, ImageMetadata
+from azure.ai.ml._restclient.v2023_08_01_preview.models import ComputeResource, ImageMetadata
 from azure.ai.ml.constants._compute import CustomApplicationDefaults
 from azure.ai.ml.entities import (
     AmlCompute,
@@ -142,6 +142,14 @@ def test_compute_instance_load_from_rest(self):
         assert compute_instance2.tags["test1"] == "test"
         assert compute_instance2.tags["test2"] == "true"
         assert compute_instance2.tags["test3"] == "0"
+        assert compute_instance2.enable_sso is False
+        assert compute_instance2.enable_root_access is False
+
+        compute_instance3: ComputeInstance = load_compute(
+            source="tests/test_configs/compute/compute-ci-defaults-unit.yaml",
+        )._to_rest_object()
+        assert compute_instance3.properties.properties.enable_sso is True
+        assert compute_instance3.properties.properties.enable_root_access is True
 
     def test_compute_instance_with_image_metadata(self):
         os_image_metadata = ImageMetadata(

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

@@ -0,0 +1,5 @@
+name: banchci
+type: computeinstance
+description: some_desc_ci
+
+size: Standard_DS2_v2

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

@@ -3,6 +3,8 @@ type: computeinstance
 description: some_desc_ci
 
 size: Standard_DS2_v2
+enable_sso: false
+enable_root_access: false
 
 network_settings:
    vnet_name: myvnet-wus2