Compute Instance OS patching and Release quota (#34563)

* new CI args

* run black

* feedback

Author: nemanjarajic

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

@@ -75,3 +75,9 @@ class ComputeInstanceSchema(ComputeSchema):
     enable_root_access = fields.Bool(
         metadata={"description": "Enable or disable root access for the compute instance."}
     )
+    release_quota_on_stop = fields.Bool(
+        metadata={"description": "Release quota on stop for the compute instance. Defaults to False."}
+    )
+    enable_os_patching = fields.Bool(
+        metadata={"description": "Enable or disable OS patching for the compute instance. Defaults to False."}
+    )

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

@@ -167,6 +167,10 @@ class ComputeInstance(Compute):
     :type enable_sso: bool
     :param enable_root_access: Enable or disable root access. Defaults to True.
     :type enable_root_access: bool
+    :param release_quota_on_stop: Release quota on stop for the compute instance. Defaults to False.
+    :type release_quota_on_stop: bool
+    :param enable_os_patching: Enable or disable OS patching for the compute instance. Defaults to False.
+    :type enable_os_patching: bool
 
     .. admonition:: Example:
 
@@ -198,6 +202,8 @@ def __init__(
         custom_applications: Optional[List[CustomApplications]] = None,
         enable_sso: bool = True,
         enable_root_access: bool = True,
+        release_quota_on_stop: bool = False,
+        enable_os_patching: bool = False,
         **kwargs: Any,
     ) -> None:
         kwargs[TYPE] = ComputeType.COMPUTEINSTANCE
@@ -226,6 +232,8 @@ def __init__(
         self.enable_node_public_ip = enable_node_public_ip
         self.enable_sso = enable_sso
         self.enable_root_access = enable_root_access
+        self.release_quota_on_stop = release_quota_on_stop
+        self.enable_os_patching = enable_os_patching
         self.custom_applications = custom_applications
         self.subnet = None
 
@@ -308,6 +316,8 @@ def _to_rest_object(self) -> ComputeResource:
             enable_node_public_ip=self.enable_node_public_ip,
             enable_sso=self.enable_sso,
             enable_root_access=self.enable_root_access,
+            release_quota_on_stop=self.release_quota_on_stop,
+            enable_os_patching=self.enable_os_patching,
         )
         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
@@ -453,6 +463,12 @@ def _load_from_rest(cls, rest_obj: ComputeResource) -> "ComputeInstance":
             enable_root_access=prop.properties.enable_root_access
             if (prop.properties and prop.properties.enable_root_access is not None)
             else True,
+            release_quota_on_stop=prop.properties.release_quota_on_stop
+            if (prop.properties and prop.properties.release_quota_on_stop is not None)
+            else False,
+            enable_os_patching=prop.properties.enable_os_patching
+            if (prop.properties and prop.properties.enable_os_patching is not None)
+            else False,
         )
         return response
 

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

@@ -144,12 +144,16 @@ def test_compute_instance_load_from_rest(self):
         assert compute_instance2.tags["test3"] == "0"
         assert compute_instance2.enable_sso is False
         assert compute_instance2.enable_root_access is False
+        assert compute_instance2.enable_os_patching is True
+        assert compute_instance2.release_quota_on_stop is True
 
         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
+        assert compute_instance3.properties.properties.enable_os_patching is False
+        assert compute_instance3.properties.properties.release_quota_on_stop is False
 
     def test_compute_instance_with_image_metadata(self):
         os_image_metadata = ImageMetadata(

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

@@ -5,6 +5,8 @@ description: some_desc_ci
 size: Standard_DS2_v2
 enable_sso: false
 enable_root_access: false
+release_quota_on_stop: True
+enable_os_patching: True
 
 network_settings:
    vnet_name: myvnet-wus2