[ML] Support QueueSettings in AutoML and Sweep jobs

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

@@ -10,6 +10,7 @@
 - Removed Experimental Tag from Image Metadata on Compute Instances.
 - Added support for data binding on outputs inside dynamic arguments for dsl pipeline
 - Added support for serverless compute in pipeline job
+- Added support for serverless compute in command, automl and sweep job
 
 ### Bugs Fixed
 

sdk/ml/azure-ai-ml/azure/ai/ml/_schema/_sweep/parameterized_sweep.py

@@ -4,9 +4,10 @@
 
 # pylint: disable=unused-argument,no-self-use
 
-from azure.ai.ml._schema.core.fields import NestedField, PathAwareSchema
+from azure.ai.ml._schema.core.fields import ExperimentalField, NestedField, PathAwareSchema
 
 from ..job.job_limits import SweepJobLimitsSchema
+from ..queue_settings import QueueSettingsSchema
 from .sweep_fields_provider import EarlyTerminationField, SamplingAlgorithmField, SearchSpaceField
 from .sweep_objective import SweepObjectiveSchema
 
@@ -26,3 +27,4 @@ class ParameterizedSweepSchema(PathAwareSchema):
         SweepJobLimitsSchema,
         required=True,
     )
+    queue_settings = ExperimentalField(NestedField(QueueSettingsSchema))

sdk/ml/azure-ai-ml/azure/ai/ml/_schema/automl/automl_job.py

@@ -4,10 +4,11 @@
 
 from marshmallow import fields
 
-from azure.ai.ml._schema.core.fields import NestedField, StringTransformedEnum
+from azure.ai.ml._schema.core.fields import ExperimentalField, NestedField, StringTransformedEnum
 from azure.ai.ml._schema.job import BaseJobSchema
 from azure.ai.ml._schema.job.input_output_fields_provider import OutputsField
 from azure.ai.ml._schema.job_resource_configuration import JobResourceConfigurationSchema
+from azure.ai.ml._schema.queue_settings import QueueSettingsSchema
 from azure.ai.ml.constants import JobType
 
 
@@ -17,3 +18,4 @@ class AutoMLJobSchema(BaseJobSchema):
     environment_variables = fields.Dict(keys=fields.Str(), values=fields.Str())
     outputs = OutputsField()
     resources = NestedField(JobResourceConfigurationSchema())
+    queue_settings = ExperimentalField(NestedField(QueueSettingsSchema))

sdk/ml/azure-ai-ml/azure/ai/ml/entities/_builders/command.py

@@ -131,6 +131,8 @@ class Command(BaseNode):
         Please see https://aka.ms/azuremlexperimental for more information.
     :type services:
         Dict[str, Union[JobService, JupyterLabJobService, SshJobService, TensorBoardJobService, VsCodeJobService]]
+    :param queue_settings: Queue settings for the job.
+    :type queue_settings: QueueSettings
     :raises ~azure.ai.ml.exceptions.ValidationException: Raised if Command cannot be successfully validated.
         Details will be provided in the error message.
     """
@@ -416,6 +418,7 @@ def sweep(
         identity: Optional[
             Union[ManagedIdentityConfiguration, AmlTokenConfiguration, UserIdentityConfiguration]
         ] = None,
+        queue_settings: Optional[QueueSettings] = None,
     ) -> Sweep:
         """Turn the command into a sweep node with extra sweep run setting. The command component in current Command
         node will be used as its trial component. A command node can sweep for multiple times, and the generated sweep
@@ -448,6 +451,8 @@ def sweep(
             ManagedIdentityConfiguration,
             AmlTokenConfiguration,
             UserIdentityConfiguration]
+        :param queue_settings: Queue settings for the job.
+        :type queue_settings: QueueSettings
         :return: A sweep node with component from current Command node as its trial component.
         :rtype: Sweep
         """
@@ -477,6 +482,7 @@ def sweep(
             experiment_name=self.experiment_name,
             identity=self.identity if not identity else identity,
             _from_component_func=True,
+            queue_settings=self.queue_settings if queue_settings is None else queue_settings,
         )
         sweep_node.set_limits(
             max_total_trials=max_total_trials,

sdk/ml/azure-ai-ml/azure/ai/ml/entities/_builders/sweep.py

@@ -22,6 +22,7 @@
 from azure.ai.ml.entities._inputs_outputs import Input, Output
 from azure.ai.ml.entities._job.job_limits import SweepJobLimits
 from azure.ai.ml.entities._job.pipeline._io import NodeInput
+from azure.ai.ml.entities._job.queue_settings import QueueSettings
 from azure.ai.ml.entities._job.sweep.early_termination_policy import (
     BanditPolicy,
     EarlyTerminationPolicy,
@@ -97,6 +98,8 @@ class Sweep(ParameterizedSweep, BaseNode):
         ManagedIdentityConfiguration,
         AmlTokenConfiguration,
         UserIdentityConfiguration]
+    :param queue_settings: Queue settings for the job.
+    :type queue_settings: QueueSettings
     """
 
     def __init__(
@@ -121,6 +124,7 @@ def __init__(
         identity: Optional[
             Union[ManagedIdentityConfiguration, AmlTokenConfiguration, UserIdentityConfiguration]
         ] = None,
+        queue_settings: Optional[QueueSettings] = None,
         **kwargs,
     ):
         # TODO: get rid of self._job_inputs, self._job_outputs once we have general Input
@@ -145,6 +149,7 @@ def __init__(
             limits=limits,
             early_termination=early_termination,
             search_space=search_space,
+            queue_settings=queue_settings,
         )
 
         self.identity = identity
@@ -293,6 +298,7 @@ def _to_job(self) -> SweepJob:
             inputs=self._job_inputs,
             outputs=self._job_outputs,
             identity=self.identity,
+            queue_settings=self.queue_settings,
         )
 
     @classmethod

sdk/ml/azure-ai-ml/azure/ai/ml/entities/_job/automl/automl_job.py

@@ -8,7 +8,13 @@
 from abc import ABC, abstractmethod
 from typing import Any, Dict, Optional, Union
 
-from azure.ai.ml._restclient.v2023_02_01_preview.models import JobBase, MLTableJobInput, ResourceConfiguration, TaskType
+from azure.ai.ml._restclient.v2023_02_01_preview.models import (
+    JobBase,
+    MLTableJobInput,
+    QueueSettings,
+    ResourceConfiguration,
+    TaskType,
+)
 from azure.ai.ml._utils.utils import camel_to_snake
 from azure.ai.ml.constants import JobType
 from azure.ai.ml.constants._common import TYPE, AssetTypes
@@ -37,6 +43,7 @@ def __init__(
         identity: Optional[
             Union[ManagedIdentityConfiguration, AmlTokenConfiguration, UserIdentityConfiguration]
         ] = None,
+        queue_settings: Optional[QueueSettings] = None,
         **kwargs: Any,
     ) -> None:
         """Initialize an AutoML job entity.
@@ -45,6 +52,8 @@ def __init__(
         :param resources: Resource configuration for the job.
         :param identity: Identity that training job will use while running on compute.
         :type identity: Union[ManagedIdentity, AmlToken, UserIdentity]
+        :param queue_settings: Queue settings for the job.
+        :type queue_settings: QueueSettings
         :param kwargs:
         """
         kwargs[TYPE] = JobType.AUTOML
@@ -56,6 +65,7 @@ def __init__(
 
         self.resources = resources
         self.identity = identity
+        self.queue_settings = queue_settings
 
     @property
     @abstractmethod

sdk/ml/azure-ai-ml/azure/ai/ml/entities/_job/automl/image/image_classification_job.py

@@ -104,6 +104,7 @@ def _to_rest_object(self) -> JobBase:
             resources=self.resources,
             task_details=image_classification_task,
             identity=self.identity._to_job_rest_object() if self.identity else None,
+            queue_settings=self.queue_settings,
         )
 
         result = JobBase(properties=properties)
@@ -132,6 +133,7 @@ def _from_rest_object(cls, obj: JobBase) -> "ImageClassificationJob":
             "identity": _BaseJobIdentityConfiguration._from_rest_object(properties.identity)
             if properties.identity
             else None,
+            "queue_settings": properties.queue_settings,
         }
 
         image_classification_job = cls(

sdk/ml/azure-ai-ml/azure/ai/ml/entities/_job/automl/image/image_classification_multilabel_job.py

@@ -106,6 +106,7 @@ def _to_rest_object(self) -> JobBase:
             resources=self.resources,
             task_details=image_classification_multilabel_task,
             identity=self.identity._to_job_rest_object() if self.identity else None,
+            queue_settings=self.queue_settings,
         )
 
         result = JobBase(properties=properties)
@@ -134,6 +135,7 @@ def _from_rest_object(cls, obj: JobBase) -> "ImageClassificationMultilabelJob":
             "identity": _BaseJobIdentityConfiguration._from_rest_object(properties.identity)
             if properties.identity
             else None,
+            "queue_settings": properties.queue_settings,
         }
 
         image_classification_multilabel_job = cls(

sdk/ml/azure-ai-ml/azure/ai/ml/entities/_job/automl/image/image_instance_segmentation_job.py

@@ -104,6 +104,7 @@ def _to_rest_object(self) -> JobBase:
             resources=self.resources,
             task_details=image_instance_segmentation_task,
             identity=self.identity._to_job_rest_object() if self.identity else None,
+            queue_settings=self.queue_settings,
         )
 
         result = JobBase(properties=properties)
@@ -132,6 +133,7 @@ def _from_rest_object(cls, obj: JobBase) -> "ImageInstanceSegmentationJob":
             "identity": _BaseJobIdentityConfiguration._from_rest_object(properties.identity)
             if properties.identity
             else None,
+            "queue_settings": properties.queue_settings,
         }
 
         image_instance_segmentation_job = cls(

sdk/ml/azure-ai-ml/azure/ai/ml/entities/_job/automl/image/image_object_detection_job.py

@@ -103,6 +103,7 @@ def _to_rest_object(self) -> JobBase:
             resources=self.resources,
             task_details=image_object_detection_task,
             identity=self.identity._to_job_rest_object() if self.identity else None,
+            queue_settings=self.queue_settings,
         )
 
         result = JobBase(properties=properties)
@@ -131,6 +132,7 @@ def _from_rest_object(cls, obj: JobBase) -> "ImageObjectDetectionJob":
             "identity": _BaseJobIdentityConfiguration._from_rest_object(properties.identity)
             if properties.identity
             else None,
+            "queue_settings": properties.queue_settings,
         }
 
         image_object_detection_job = cls(

sdk/ml/azure-ai-ml/azure/ai/ml/entities/_job/automl/nlp/text_classification_job.py

@@ -107,6 +107,7 @@ def _to_rest_object(self) -> JobBase:
             resources=self.resources,
             task_details=text_classification,
             identity=self.identity._to_ob_rest_object() if self.identity else None,
+            queue_settings=self.queue_settings,
         )
 
         result = JobBase(properties=properties)
@@ -162,6 +163,7 @@ def _from_rest_object(cls, obj: JobBase) -> "TextClassificationJob":
             identity=_BaseJobIdentityConfiguration._from_rest_object(properties.identity)
             if properties.identity
             else None,
+            queue_settings=properties.queue_settings,
         )
 
         text_classification_job._restore_data_inputs()

sdk/ml/azure-ai-ml/azure/ai/ml/entities/_job/automl/nlp/text_classification_multilabel_job.py

@@ -106,6 +106,7 @@ def _to_rest_object(self) -> JobBase:
             resources=self.resources,
             task_details=text_classification_multilabel,
             identity=self.identity._to_job_rest_object() if self.identity else None,
+            queue_settings=self.queue_settings,
         )
 
         result = JobBase(properties=properties)
@@ -161,6 +162,7 @@ def _from_rest_object(cls, obj: JobBase) -> "TextClassificationMultilabelJob":
             identity=_BaseJobIdentityConfiguration._from_rest_object(properties.identity)
             if properties.identity
             else None,
+            queue_settings=properties.queue_settings,
         )
 
         text_classification_multilabel_job._restore_data_inputs()

sdk/ml/azure-ai-ml/azure/ai/ml/entities/_job/automl/nlp/text_ner_job.py

@@ -103,6 +103,7 @@ def _to_rest_object(self) -> JobBase:
             resources=self.resources,
             task_details=text_ner,
             identity=self.identity._to_job_rest_object() if self.identity else None,
+            queue_settings=self.queue_settings,
         )
 
         result = JobBase(properties=properties)
@@ -158,6 +159,7 @@ def _from_rest_object(cls, obj: JobBase) -> "TextNerJob":
             identity=_BaseJobIdentityConfiguration._from_rest_object(properties.identity)
             if properties.identity
             else None,
+            queue_settings=properties.queue_settings,
         )
 
         text_ner_job._restore_data_inputs()

sdk/ml/azure-ai-ml/azure/ai/ml/entities/_job/automl/tabular/classification_job.py

@@ -113,6 +113,7 @@ def _to_rest_object(self) -> JobBase:
             resources=self.resources,
             task_details=classification_task,
             identity=self.identity._to_job_rest_object() if self.identity else None,
+            queue_settings=self.queue_settings,
         )
 
         result = JobBase(properties=properties)
@@ -141,6 +142,7 @@ def _from_rest_object(cls, obj: JobBase) -> "ClassificationJob":
             "identity": _BaseJobIdentityConfiguration._from_rest_object(properties.identity)
             if properties.identity
             else None,
+            "queue_settings": properties.queue_settings,
         }
 
         classification_job = cls(

sdk/ml/azure-ai-ml/azure/ai/ml/entities/_job/automl/tabular/forecasting_job.py

@@ -404,6 +404,7 @@ def _to_rest_object(self) -> JobBase:
             resources=self.resources,
             task_details=forecasting_task,
             identity=self.identity._to_job_rest_object() if self.identity else None,
+            queue_settings=self.queue_settings,
         )
 
         result = JobBase(properties=properties)
@@ -432,6 +433,7 @@ def _from_rest_object(cls, obj: JobBase) -> "ForecastingJob":
             "identity": _BaseJobIdentityConfiguration._from_rest_object(properties.identity)
             if properties.identity
             else None,
+            "queue_settings": properties.queue_settings,
         }
 
         forecasting_job = cls(

sdk/ml/azure-ai-ml/azure/ai/ml/entities/_job/automl/tabular/regression_job.py

@@ -107,6 +107,7 @@ def _to_rest_object(self) -> JobBase:
             resources=self.resources,
             task_details=regression_task,
             identity=self.identity._to_job_rest_object() if self.identity else None,
+            queue_settings=self.queue_settings,
         )
 
         result = JobBase(properties=properties)
@@ -135,6 +136,7 @@ def _from_rest_object(cls, obj: JobBase) -> "RegressionJob":
             "identity": _BaseJobIdentityConfiguration._from_rest_object(properties.identity)
             if properties.identity
             else None,
+            "queue_settings": properties.queue_settings,
         }
 
         regression_job = cls(

sdk/ml/azure-ai-ml/azure/ai/ml/entities/_job/parameterized_command.py

@@ -100,5 +100,6 @@ def _load_from_sweep_job(cls, sweep_job: SweepJob) -> "ParameterizedCommand":
             environment=sweep_job.trial.environment_id,
             distribution=DistributionConfiguration._from_rest_object(sweep_job.trial.distribution),
             resources=JobResourceConfiguration._from_rest_object(sweep_job.trial.resources),
+            queue_settings=QueueSettings._from_rest_object(sweep_job.queue_settings),
         )
         return parameterized_command

sdk/ml/azure-ai-ml/azure/ai/ml/entities/_job/sweep/parameterized_sweep.py

@@ -6,6 +6,7 @@
 from azure.ai.ml.exceptions import ErrorCategory, ErrorTarget, ValidationErrorType, ValidationException
 
 from ..job_limits import SweepJobLimits
+from ..queue_settings import QueueSettings
 from .early_termination_policy import (
     BanditPolicy,
     EarlyTerminationPolicy,
@@ -70,11 +71,13 @@ def __init__(
                 ],
             ]
         ] = None,
+        queue_settings: Optional[QueueSettings] = None,
     ):
         self.sampling_algorithm = sampling_algorithm
         self.early_termination = early_termination
         self._limits = limits
         self.search_space = search_space
+        self.queue_settings = queue_settings
 
         if isinstance(objective, Dict):
             self.objective = Objective(**objective)

sdk/ml/azure-ai-ml/azure/ai/ml/entities/_job/sweep/sweep_job.py

@@ -41,6 +41,7 @@
 # from ..identity import AmlToken, Identity, ManagedIdentity, UserIdentity
 from ..job_limits import SweepJobLimits
 from ..parameterized_command import ParameterizedCommand
+from ..queue_settings import QueueSettings
 from .early_termination_policy import (
     BanditPolicy,
     EarlyTerminationPolicy,
@@ -112,6 +113,8 @@ class SweepJob(Job, ParameterizedSweep, JobIOMixin):
     ~azure.mgmt.machinelearningservices.models.TruncationSelectionPolicy]
     :param limits: Limits for the sweep job.
     :type limits: ~azure.ai.ml.entities.SweepJobLimits
+    :param queue_settings: Queue settings for the job.
+    :type queue_settings: QueueSettings
     :param kwargs: A dictionary of additional configuration parameters.
     :type kwargs: dict
     """
@@ -143,6 +146,7 @@ def __init__(
         objective: Optional[Objective] = None,
         trial: Optional[Union[CommandJob, CommandComponent]] = None,
         early_termination: Optional[Union[BanditPolicy, MedianStoppingPolicy, TruncationSelectionPolicy]] = None,
+        queue_settings: Optional[QueueSettings] = None,
         **kwargs: Any,
     ):
         kwargs[TYPE] = JobType.SWEEP
@@ -169,6 +173,7 @@ def __init__(
             objective=objective,
             early_termination=early_termination,
             search_space=search_space,
+            queue_settings=queue_settings,
         )
 
     def _to_dict(self) -> Dict:
@@ -208,6 +213,7 @@ def _to_rest_object(self) -> JobBase:
             inputs=to_rest_dataset_literal_inputs(self.inputs, job_type=self.type),
             outputs=to_rest_data_outputs(self.outputs),
             identity=self.identity._to_job_rest_object() if self.identity else None,
+            queue_settings=self.queue_settings._to_rest_object() if self.queue_settings else None,
         )
         sweep_job_resource = JobBase(properties=sweep_job)
         sweep_job_resource.name = self.name
@@ -268,6 +274,7 @@ def _load_from_rest(cls, obj: JobBase) -> "SweepJob":
             identity=_BaseJobIdentityConfiguration._from_rest_object(properties.identity)
             if properties.identity
             else None,
+            queue_settings=properties.queue_settings,
         )
 
     def _override_missing_properties_from_trial(self):