SDK v2 and CLI changes for distributed training for automl tabular task types. (Azure#29341)

* Initial draft.

* Initial draft.

* Fix imports.

* Add the two flags and enum.

* Add tests for limit settings- max_nodes.

* Add unittests for forecasting job.

* Add unittests for classification jobs.

* Add unittests for regression job.

* Add tests for CLI schema.

* Fix missing new lines in yaml files.

* Add documentation.

* Fix formatting issues.

* Set max_nodes to 1 by default.

* Fix constant import formatting.

* Fix code review comments.

* Fix formatting issue.

* Redefine constants with proper documentation.

* Fix imports.

* Fix documentation.

Author: arun-rajora

sdk/ml/azure-ai-ml/azure/ai/ml/_schema/automl/table_vertical/table_vertical_limit_settings.py

@@ -6,6 +6,7 @@
 
 from marshmallow import fields, post_load
 
+from azure.ai.ml._schema import ExperimentalField
 from azure.ai.ml._schema.core.schema import PatchedSchemaMeta
 from azure.ai.ml.constants._job.automl import AutoMLConstants
 
@@ -15,6 +16,7 @@ class AutoMLTableLimitsSchema(metaclass=PatchedSchemaMeta):
     exit_score = fields.Float()
     max_concurrent_trials = fields.Int()
     max_cores_per_trial = fields.Int()
+    max_nodes = ExperimentalField(fields.Int())
     max_trials = fields.Int(data_key=AutoMLConstants.MAX_TRIALS_YAML)
     timeout_minutes = fields.Int()  # type duration
     trial_timeout_minutes = fields.Int()  # type duration

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

@@ -12,6 +12,8 @@
     RegressionModels,
     StackMetaLearnerType,
 )
+from azure.ai.ml.constants import TabularTrainingMode
+from azure.ai.ml._schema import ExperimentalField
 from azure.ai.ml._schema.core.fields import NestedField, StringTransformedEnum
 from azure.ai.ml._schema.core.schema import PatchedSchemaMeta
 from azure.ai.ml._utils.utils import camel_to_snake
@@ -49,6 +51,12 @@ class TrainingSettingsSchema(metaclass=PatchedSchemaMeta):
     enable_vote_ensemble = fields.Bool()
     ensemble_model_download_timeout = fields.Int(data_key=AutoMLConstants.ENSEMBLE_MODEL_DOWNLOAD_TIMEOUT_YAML)
     stack_ensemble_settings = NestedField(StackEnsembleSettingsSchema())
+    training_mode = ExperimentalField(
+        StringTransformedEnum(
+            allowed_values=[o.value for o in TabularTrainingMode],
+            casing_transform=camel_to_snake,
+        )
+    )
 
 
 class ClassificationTrainingSettingsSchema(TrainingSettingsSchema):

sdk/ml/azure-ai-ml/azure/ai/ml/constants/__init__.py

@@ -17,10 +17,13 @@
     JobType,
     NlpLearningRateScheduler,
     NlpModels,
+    TabularTrainingMode,
 )
 from ._registry import AcrAccountSku, StorageAccountType
 from ._workspace import ManagedServiceIdentityType
 
+TabularTrainingMode.__module__ = __name__
+
 __all__ = [
     "ImportSourceType",
     "JobType",
@@ -40,4 +43,5 @@
     "NlpModels",
     "NlpLearningRateScheduler",
     "Scope",
+    "TabularTrainingMode",
 ]

sdk/ml/azure-ai-ml/azure/ai/ml/constants/_job/__init__.py

@@ -12,6 +12,7 @@
     ImageObjectDetectionModelNames,
     NlpLearningRateScheduler,
     NlpModels,
+    TabularTrainingMode,
 )
 from .job import DistributionType, ImportSourceType, JobType
 from .pipeline import PipelineConstants
@@ -30,4 +31,5 @@
     "SearchSpace",
     "NlpModels",
     "NlpLearningRateScheduler",
+    "TabularTrainingMode",
 ]

sdk/ml/azure-ai-ml/azure/ai/ml/constants/_job/automl.py

@@ -5,6 +5,8 @@
 
 # pylint: disable=unused-import
 from azure.ai.ml._restclient.v2022_10_01_preview.models import NlpLearningRateScheduler
+from azure.ai.ml._restclient.v2023_02_01_preview.models import TrainingMode
+from azure.ai.ml._utils._experimental import experimental
 
 
 class AutoMLConstants:
@@ -107,3 +109,8 @@ class NlpModels(Enum):
     XLM_ROBERTA_LARGE = "xlm-roberta-large"
     XLNET_BASE_CASED = "xlnet-base-cased"
     XLNET_LARGE_CASED = "xlnet-large-cased"
+
+
+TrainingMode.__doc__ = "Mode to enable/disable distributed training."
+TabularTrainingMode = experimental(TrainingMode)
+TabularTrainingMode.__name__ = "TabularTrainingMode"

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

@@ -14,6 +14,7 @@
     LogVerbosity,
 )
 from azure.ai.ml._utils.utils import camel_to_snake
+from azure.ai.ml.constants import TabularTrainingMode
 from azure.ai.ml.constants._job.automl import AutoMLConstants
 from azure.ai.ml.entities._inputs_outputs import Input
 from azure.ai.ml.entities._job.automl.automl_vertical import AutoMLVertical
@@ -245,6 +246,7 @@ def set_limits(
         exit_score: Optional[float] = None,
         max_concurrent_trials: Optional[int] = None,
         max_cores_per_trial: Optional[int] = None,
+        max_nodes: Optional[int] = None,
         max_trials: Optional[int] = None,
         timeout_minutes: Optional[int] = None,
         trial_timeout_minutes: Optional[int] = None,
@@ -301,6 +303,14 @@ def set_limits(
 
             * Equal to 1, the default.
         :paramtype max_cores_per_trial: typing.Optional[int]
+        :keyword max_nodes: [Experimental] The maximum number of nodes to use for distributed training.
+
+            * For forecasting, each model is trained using max(2, int(max_nodes / max_concurrent_trials)) nodes.
+
+            * For classification/regression, each model is trained using max_nodes nodes.
+
+            Note- This parameter is in public preview and might change in future.
+        :paramtype max_nodes: typing.Optional[int]
         :keyword max_trials: The total number of different algorithm and parameter combinations to test during an
             automated ML experiment. If not specified, the default is 1000 iterations.
         :paramtype max_trials: typing.Optional[int]
@@ -324,6 +334,7 @@ def set_limits(
         self._limits.max_cores_per_trial = (
             max_cores_per_trial if max_cores_per_trial is not None else self._limits.max_cores_per_trial
         )
+        self._limits.max_nodes = max_nodes if max_nodes is not None else self._limits.max_nodes
         self._limits.max_trials = max_trials if max_trials is not None else self._limits.max_trials
         self._limits.timeout_minutes = timeout_minutes if timeout_minutes is not None else self._limits.timeout_minutes
         self._limits.trial_timeout_minutes = (
@@ -342,6 +353,7 @@ def set_training(
         ensemble_model_download_timeout: Optional[int] = None,
         allowed_training_algorithms: Optional[List[str]] = None,
         blocked_training_algorithms: Optional[List[str]] = None,
+        training_mode: Optional[Union[str, TabularTrainingMode]] = None,
     ) -> None:
         """The method to configure training related settings.
 
@@ -383,6 +395,17 @@ def set_training(
         :paramtype allowed_training_algorithms: typing.Optional[List[str]]
         :keyword blocked_training_algorithms: A list of algorithms to ignore for an experiment, defaults to None
         :paramtype blocked_training_algorithms: typing.Optional[List[str]]
+        :keyword training_mode: [Experimental] The training mode to use.
+            The possible values are-
+
+            * distributed- enables distributed training for supported algorithms.
+
+            * non_distributed- disables distributed training.
+
+            * auto- Currently, it is same as non_distributed. In future, this might change.
+
+            Note: This parameter is in public preview and may change in future.
+        :paramtype training_mode: typing.Optional[typing.Union[str, azure.ai.ml.constants.TabularTrainingMode]]
         """
         # get training object by calling training getter of respective tabular task
         self._training = self.training
@@ -417,6 +440,7 @@ def set_training(
 
         self._training.allowed_training_algorithms = allowed_training_algorithms
         self._training.blocked_training_algorithms = blocked_training_algorithms
+        self._training.training_mode = training_mode if training_mode is not None else self._training.training_mode
 
     def set_featurization(
         self,

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

@@ -10,6 +10,7 @@
 from azure.ai.ml._restclient.v2023_02_01_preview.models import Forecasting as RestForecasting
 from azure.ai.ml._restclient.v2023_02_01_preview.models import ForecastingPrimaryMetrics, JobBase, TaskType
 from azure.ai.ml._utils.utils import camel_to_snake, is_data_binding_expression
+from azure.ai.ml.constants import TabularTrainingMode
 from azure.ai.ml.constants._common import BASE_PATH_CONTEXT_KEY
 from azure.ai.ml.constants._job.automl import AutoMLConstants
 from azure.ai.ml.entities._credentials import _BaseJobIdentityConfiguration
@@ -351,6 +352,7 @@ def set_training(
         ensemble_model_download_timeout: Optional[int] = None,
         allowed_training_algorithms: Optional[List[str]] = None,
         blocked_training_algorithms: Optional[List[str]] = None,
+        training_mode: Optional[Union[str, TabularTrainingMode]] = None,
     ) -> None:
         super().set_training(
             enable_onnx_compatible_models=enable_onnx_compatible_models,
@@ -362,6 +364,7 @@ def set_training(
             ensemble_model_download_timeout=ensemble_model_download_timeout,
             allowed_training_algorithms=allowed_training_algorithms,
             blocked_training_algorithms=blocked_training_algorithms,
+            training_mode=training_mode,
         )
 
         # Disable stack ensemble by default, since it is currently not supported for forecasting tasks

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

@@ -21,6 +21,14 @@ class TabularLimitSettings(RestTranslatableMixin):
     :type max_concurrent_trials: int, optional
     :param max_cores_per_trial: The maximum number of threads to use for a given training iteration.
     :type max_cores_per_trial: int, optional
+    :param max_nodes: [Experimental] The maximum number of nodes to use for distributed training.
+
+        * For forecasting, each model is trained using max(2, int(max_nodes / max_concurrent_trials)) nodes.
+
+        * For classification/regression, each model is trained using max_nodes nodes.
+
+        Note- This parameter is in public preview and might change in future.
+    :type max_nodes: int, optional
     :param max_trials: Maximum number of AutoML iterations.
     :type max_trials: int, optional
     :param timeout_minutes: AutoML job timeout.
@@ -36,6 +44,7 @@ def __init__(
         exit_score: Optional[float] = None,
         max_concurrent_trials: Optional[int] = None,
         max_cores_per_trial: Optional[int] = None,
+        max_nodes: Optional[int] = None,
         max_trials: Optional[int] = None,
         timeout_minutes: Optional[int] = None,
         trial_timeout_minutes: Optional[int] = None,
@@ -44,6 +53,7 @@ def __init__(
         self.exit_score = exit_score
         self.max_concurrent_trials = max_concurrent_trials
         self.max_cores_per_trial = max_cores_per_trial
+        self.max_nodes = max_nodes
         self.max_trials = max_trials
         self.timeout_minutes = timeout_minutes
         self.trial_timeout_minutes = trial_timeout_minutes
@@ -54,6 +64,7 @@ def _to_rest_object(self) -> RestTabularLimitSettings:
             exit_score=self.exit_score,
             max_concurrent_trials=self.max_concurrent_trials,
             max_cores_per_trial=self.max_cores_per_trial,
+            max_nodes=self.max_nodes,
             max_trials=self.max_trials,
             timeout=to_iso_duration_format_mins(self.timeout_minutes),
             trial_timeout=to_iso_duration_format_mins(self.trial_timeout_minutes),
@@ -66,6 +77,7 @@ def _from_rest_object(cls, obj: RestTabularLimitSettings) -> "TabularLimitSettin
             exit_score=obj.exit_score,
             max_concurrent_trials=obj.max_concurrent_trials,
             max_cores_per_trial=obj.max_cores_per_trial,
+            max_nodes=obj.max_nodes,
             max_trials=obj.max_trials,
             timeout_minutes=from_iso_duration_format_mins(obj.timeout),
             trial_timeout_minutes=from_iso_duration_format_mins(obj.trial_timeout),
@@ -79,6 +91,7 @@ def __eq__(self, other: object) -> bool:
             and self.exit_score == other.exit_score
             and self.max_concurrent_trials == other.max_concurrent_trials
             and self.max_cores_per_trial == other.max_cores_per_trial
+            and self.max_nodes == other.max_nodes
             and self.max_trials == other.max_trials
             and self.timeout_minutes == other.timeout_minutes
             and self.trial_timeout_minutes == other.trial_timeout_minutes

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

@@ -19,9 +19,12 @@
     RegressionTrainingSettings as RestRegressionTrainingSettings,
 )
 from azure.ai.ml._restclient.v2023_02_01_preview.models import TrainingSettings as RestTrainingSettings
+from azure.ai.ml._utils._experimental import experimental
 from azure.ai.ml._utils.utils import camel_to_snake, from_iso_duration_format_mins, to_iso_duration_format_mins
+from azure.ai.ml.constants import TabularTrainingMode
 from azure.ai.ml.entities._job.automl.stack_ensemble_settings import StackEnsembleSettings
 from azure.ai.ml.entities._mixins import RestTranslatableMixin
+from azure.ai.ml.exceptions import ErrorCategory, ErrorTarget, ValidationException
 
 
 class TrainingSettings(RestTranslatableMixin):
@@ -39,6 +42,7 @@ def __init__(
         ensemble_model_download_timeout: Optional[int] = None,
         allowed_training_algorithms: Optional[List[str]] = None,
         blocked_training_algorithms: Optional[List[str]] = None,
+        training_mode: Optional[Union[str, TabularTrainingMode]] = None,
     ):
         """TrainingSettings class for Azure Machine Learning.
 
@@ -51,6 +55,16 @@ def __init__(
         :param ensemble_model_download_timeout: Timeout for downloading ensemble models
         :param allowed_training_algorithms: Models to train
         :param blocked_training_algorithms: Models that will not be considered for training
+        :param training_mode: [Experimental] The training mode to use.
+            The possible values are-
+
+            * distributed- enables distributed training for supported algorithms.
+
+            * non_distributed- disables distributed training.
+
+            * auto- Currently, it is same as non_distributed. In future, this might change.
+
+            Note: This parameter is in public preview and may change in future.
         """
         self.enable_onnx_compatible_models = enable_onnx_compatible_models
         self.enable_dnn_training = enable_dnn_training
@@ -61,6 +75,31 @@ def __init__(
         self.ensemble_model_download_timeout = ensemble_model_download_timeout
         self.allowed_training_algorithms = allowed_training_algorithms
         self.blocked_training_algorithms = blocked_training_algorithms
+        self.training_mode = training_mode
+
+    @experimental
+    @property
+    def training_mode(self):
+        return self._training_mode
+
+    @training_mode.setter
+    def training_mode(self, value: Optional[Union[str, TabularTrainingMode]]):
+        if value is None or value is TabularTrainingMode:
+            self._training_mode = value
+        elif hasattr(TabularTrainingMode, camel_to_snake(value).upper()):
+            self._training_mode = TabularTrainingMode[camel_to_snake(value).upper()]
+        else:
+            supported_values = ", ".join([f'"{camel_to_snake(mode.value)}"' for mode in TabularTrainingMode])
+            msg = (
+                f"Unsupported training mode: {value}. Supported values are- {supported_values}. "
+                "Or you can use azure.ai.ml.constants.TabularTrainingMode enum."
+            )
+            raise ValidationException(
+                message=msg,
+                no_personal_data_message=msg,
+                target=ErrorTarget.AUTOML,
+                error_category=ErrorCategory.USER_ERROR,
+            )
 
     @property
     def allowed_training_algorithms(self):
@@ -81,6 +120,7 @@ def _to_rest_object(self) -> RestTrainingSettings:
             if self.stack_ensemble_settings
             else None,
             ensemble_model_download_timeout=to_iso_duration_format_mins(self.ensemble_model_download_timeout),
+            training_mode=self.training_mode,
         )
 
     @classmethod
@@ -97,6 +137,7 @@ def _from_rest_object(cls, obj: RestTrainingSettings) -> "TrainingSettings":
                 if obj.stack_ensemble_settings
                 else None
             ),
+            training_mode=obj.training_mode,
         )
 
     def __eq__(self, other: object) -> bool:
@@ -112,6 +153,7 @@ def __eq__(self, other: object) -> bool:
             and self.stack_ensemble_settings == other.stack_ensemble_settings
             and self.allowed_training_algorithms == other.allowed_training_algorithms
             and self.blocked_training_algorithms == other.blocked_training_algorithms
+            and self.training_mode == other.training_mode
         )
 
     def __ne__(self, other: object) -> bool:
@@ -162,6 +204,7 @@ def _to_rest_object(self) -> RestClassificationTrainingSettings:
             ensemble_model_download_timeout=to_iso_duration_format_mins(self.ensemble_model_download_timeout),
             allowed_training_algorithms=self.allowed_training_algorithms,
             blocked_training_algorithms=self.blocked_training_algorithms,
+            training_mode=self.training_mode,
         )
 
     @classmethod
@@ -176,6 +219,7 @@ def _from_rest_object(cls, obj: RestClassificationTrainingSettings) -> "Classifi
             stack_ensemble_settings=obj.stack_ensemble_settings,
             allowed_training_algorithms=obj.allowed_training_algorithms,
             blocked_training_algorithms=obj.blocked_training_algorithms,
+            training_mode=obj.training_mode,
         )
 
 
@@ -211,6 +255,7 @@ def _to_rest_object(self) -> RestForecastingTrainingSettings:
             ensemble_model_download_timeout=to_iso_duration_format_mins(self.ensemble_model_download_timeout),
             allowed_training_algorithms=self.allowed_training_algorithms,
             blocked_training_algorithms=self.blocked_training_algorithms,
+            training_mode=self.training_mode,
         )
 
     @classmethod
@@ -225,6 +270,7 @@ def _from_rest_object(cls, obj: RestForecastingTrainingSettings) -> "Forecasting
             stack_ensemble_settings=obj.stack_ensemble_settings,
             allowed_training_algorithms=obj.allowed_training_algorithms,
             blocked_training_algorithms=obj.blocked_training_algorithms,
+            training_mode=obj.training_mode,
         )
 
 
@@ -260,6 +306,7 @@ def _to_rest_object(self) -> RestRegressionTrainingSettings:
             ensemble_model_download_timeout=to_iso_duration_format_mins(self.ensemble_model_download_timeout),
             allowed_training_algorithms=self.allowed_training_algorithms,
             blocked_training_algorithms=self.blocked_training_algorithms,
+            training_mode=self.training_mode,
         )
 
     @classmethod
@@ -274,4 +321,5 @@ def _from_rest_object(cls, obj: RestRegressionTrainingSettings) -> "RegressionTr
             stack_ensemble_settings=obj.stack_ensemble_settings,
             allowed_training_algorithms=obj.allowed_training_algorithms,
             blocked_training_algorithms=obj.blocked_training_algorithms,
+            training_mode=obj.training_mode,
         )