[ml] Add 'max_instance_count' to JobResourceConfiguration for supporting PyTorch Elastic jobs (Azure#29973)

svaruag · Gaurav Singh · web-flow · commit 428474ec2ef9 · 2023-04-24T20:03:55.000-04:00
* updates for using the max_instance_count flag

* update import

* update imports

* undo import job

* fix test

* update deployment imports

* update endpoint svc client

* undo / update svc client

* fix eq

* undo dep ops, update job res cfg

* undo res cfg changes

* undo res cfg changes

* undo formatting changes

---------

Co-authored-by: Gaurav Singh &lt;gasi@microsoft.com&gt;
diff --git a/sdk/ml/azure-ai-ml/azure/ai/ml/_schema/job/identity.py b/sdk/ml/azure-ai-ml/azure/ai/ml/_schema/job/identity.py
@@ -8,8 +8,10 @@
 
 from marshmallow import fields, post_load
 
-from azure.ai.ml._restclient.v2022_01_01_preview.models import ConnectionAuthType
-from azure.ai.ml._restclient.v2022_10_01_preview.models import IdentityConfigurationType
+from azure.ai.ml._restclient.v2023_04_01_preview.models import (
+    ConnectionAuthType,
+    IdentityConfigurationType,
+)
 from azure.ai.ml._schema.core.fields import StringTransformedEnum
 from azure.ai.ml._utils.utils import camel_to_snake
 from azure.ai.ml.entities._credentials import (
diff --git a/sdk/ml/azure-ai-ml/azure/ai/ml/_schema/job_resource_configuration.py b/sdk/ml/azure-ai-ml/azure/ai/ml/_schema/job_resource_configuration.py
@@ -19,6 +19,9 @@ class JobResourceConfigurationSchema(ResourceConfigurationSchema):
             )
         }
     )
+    max_instance_count = fields.Int(
+        metadata={"description": "The maximum number of instances to make available to this job."}
+    )
     docker_args = fields.Str(metadata={"description": "arguments to pass to the Docker run command."})
 
     @post_load
diff --git a/sdk/ml/azure-ai-ml/azure/ai/ml/automl/__init__.py b/sdk/ml/azure-ai-ml/azure/ai/ml/automl/__init__.py
@@ -6,9 +6,7 @@
 
 Main areas include managing AutoML tasks.
 """
-from azure.ai.ml.entities._job.automl import (
-    TrainingSettings,
-)
+from azure.ai.ml.entities._job.automl import TrainingSettings
 from azure.ai.ml.entities._job.automl.image import (
     ImageClassificationJob,
     ImageClassificationMultilabelJob,
@@ -32,7 +30,9 @@
     TextNerJob,
 )
 from azure.ai.ml.entities._job.automl.search_space import SearchSpace
-from azure.ai.ml.entities._job.automl.stack_ensemble_settings import StackEnsembleSettings
+from azure.ai.ml.entities._job.automl.stack_ensemble_settings import (
+    StackEnsembleSettings,
+)
 from azure.ai.ml.entities._job.automl.tabular import (
     ClassificationJob,
     ColumnTransformer,
@@ -43,7 +43,7 @@
     TabularLimitSettings,
 )
 
-from .._restclient.v2022_10_01_preview.models import (
+from .._restclient.v2023_04_01_preview.models import (
     BlockedTransformers,
     ClassificationModels,
     ClassificationMultilabelPrimaryMetrics,
diff --git a/sdk/ml/azure-ai-ml/azure/ai/ml/constants/_job/automl.py b/sdk/ml/azure-ai-ml/azure/ai/ml/constants/_job/automl.py
@@ -4,8 +4,10 @@
 from enum import Enum
 
 # 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._restclient.v2023_04_01_preview.models import (
+    NlpLearningRateScheduler,
+    TrainingMode,
+)
 from azure.ai.ml._utils._experimental import experimental
 
 
diff --git a/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_job/data_transfer/data_transfer_job.py b/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_job/data_transfer/data_transfer_job.py
@@ -8,28 +8,28 @@
 from pathlib import Path
 from typing import Dict, Optional, Union
 
-from azure.ai.ml._restclient.v2022_10_01_preview.models import JobBase
+from azure.ai.ml._restclient.v2023_04_01_preview.models import JobBase
 from azure.ai.ml._schema.job.data_transfer_job import (
     DataTransferCopyJobSchema,
-    DataTransferImportJobSchema,
     DataTransferExportJobSchema,
+    DataTransferImportJobSchema,
 )
 from azure.ai.ml.constants import JobType
 from azure.ai.ml.constants._common import BASE_PATH_CONTEXT_KEY, TYPE
 from azure.ai.ml.constants._component import (
-    ExternalDataType,
     DataTransferBuiltinComponentUri,
     DataTransferTaskType,
+    ExternalDataType,
 )
 from azure.ai.ml.entities._inputs_outputs import Input, Output
+from azure.ai.ml.entities._inputs_outputs.external_data import Database, FileSystem
 from azure.ai.ml.entities._util import load_from_dict
 from azure.ai.ml.exceptions import (
     ErrorCategory,
     ErrorTarget,
     ValidationErrorType,
     ValidationException,
 )
-from azure.ai.ml.entities._inputs_outputs.external_data import Database, FileSystem
 
 from ..job import Job
 from ..job_io_mixin import JobIOMixin
diff --git a/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_job/job_limits.py b/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_job/job_limits.py
@@ -6,9 +6,17 @@
 from abc import ABC
 from typing import Optional, Union
 
-from azure.ai.ml._restclient.v2022_12_01_preview.models import CommandJobLimits as RestCommandJobLimits
-from azure.ai.ml._restclient.v2022_12_01_preview.models import SweepJobLimits as RestSweepJobLimits
-from azure.ai.ml._utils.utils import from_iso_duration_format, is_data_binding_expression, to_iso_duration_format
+from azure.ai.ml._restclient.v2023_04_01_preview.models import (
+    CommandJobLimits as RestCommandJobLimits,
+)
+from azure.ai.ml._restclient.v2023_04_01_preview.models import (
+    SweepJobLimits as RestSweepJobLimits,
+)
+from azure.ai.ml._utils.utils import (
+    from_iso_duration_format,
+    is_data_binding_expression,
+    to_iso_duration_format,
+)
 from azure.ai.ml.constants import JobType
 from azure.ai.ml.entities._mixins import RestTranslatableMixin
 
diff --git a/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_job/job_resource_configuration.py b/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_job/job_resource_configuration.py
@@ -6,7 +6,9 @@
 import logging
 from typing import Any, Dict, List, Optional
 
-from azure.ai.ml._restclient.v2023_04_01_preview.models import JobResourceConfiguration as RestJobResourceConfiguration
+from azure.ai.ml._restclient.v2023_04_01_preview.models import (
+    JobResourceConfiguration as RestJobResourceConfiguration,
+)
 from azure.ai.ml.constants._job.job import JobComputePropertyFields
 from azure.ai.ml.entities._mixins import DictMixin, RestTranslatableMixin
 from azure.ai.ml.entities._util import convert_ordered_dict_to_dict
@@ -92,6 +94,8 @@ class JobResourceConfiguration(RestTranslatableMixin, DictMixin):
     :vartype locations: List[str]
     :param instance_type: Optional type of VM used as supported by the compute target.
     :type instance_type: str
+    :param max_instance_count: Optional maximum number of instances or nodes used by the compute target.
+    :type max_instance_count: int
     :param properties: Additional properties bag.
     :type properties: Dict[str, Any]
     :param docker_args: Extra arguments to pass to the Docker run command. This would override any
@@ -113,12 +117,14 @@ def __init__(
         properties: Optional[Dict[str, Any]] = None,
         docker_args: Optional[str] = None,
         shm_size: Optional[str] = None,
+        max_instance_count: Optional[int] = None,
         **kwargs
     ):  # pylint: disable=unused-argument
         self.locations = locations
         self.instance_count = instance_count
         self.instance_type = instance_type
         self.shm_size = shm_size
+        self.max_instance_count = max_instance_count
         self.docker_args = docker_args
         self._properties = None
         self.properties = properties
@@ -141,6 +147,7 @@ def _to_rest_object(self) -> RestJobResourceConfiguration:
             locations=self.locations,
             instance_count=self.instance_count,
             instance_type=self.instance_type,
+            max_instance_count=self.max_instance_count,
             properties=self.properties.as_dict(),
             docker_args=self.docker_args,
             shm_size=self.shm_size,
@@ -156,6 +163,7 @@ def _from_rest_object(cls, obj: Optional[RestJobResourceConfiguration]) -> Optio
             locations=obj.locations,
             instance_count=obj.instance_count,
             instance_type=obj.instance_type,
+            max_instance_count=obj.max_instance_count if hasattr(obj, "max_instance_count") else None,
             properties=obj.properties,
             docker_args=obj.docker_args,
             shm_size=obj.shm_size,
@@ -169,6 +177,7 @@ def __eq__(self, other: object) -> bool:
             self.locations == other.locations
             and self.instance_count == other.instance_count
             and self.instance_type == other.instance_type
+            and self.max_instance_count == other.max_instance_count
             and self.docker_args == other.docker_args
             and self.shm_size == other.shm_size
         )
@@ -186,6 +195,8 @@ def _merge_with(self, other: "JobResourceConfiguration") -> None:
                 self.instance_count = other.instance_count
             if other.instance_type:
                 self.instance_type = other.instance_type
+            if other.max_instance_count:
+                self.max_instance_count = other.max_instance_count
             if other.properties:
                 self.properties = other.properties
             if other.docker_args:
diff --git a/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_job/job_service.py b/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_job/job_service.py
@@ -6,12 +6,21 @@
 
 import logging
 from typing import Dict, Optional, Union
+
 from typing_extensions import Literal
-from azure.ai.ml._restclient.v2022_12_01_preview.models import AllNodes
-from azure.ai.ml._restclient.v2022_12_01_preview.models import JobService as RestJobService
+
+from azure.ai.ml._restclient.v2023_04_01_preview.models import AllNodes
+from azure.ai.ml._restclient.v2023_04_01_preview.models import (
+    JobService as RestJobService,
+)
 from azure.ai.ml.constants._job.job import JobServiceTypeNames
 from azure.ai.ml.entities._mixins import DictMixin, RestTranslatableMixin
-from azure.ai.ml.exceptions import ErrorCategory, ErrorTarget, ValidationErrorType, ValidationException
+from azure.ai.ml.exceptions import (
+    ErrorCategory,
+    ErrorTarget,
+    ValidationErrorType,
+    ValidationException,
+)
 
 module_logger = logging.getLogger(__name__)
 
diff --git a/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_job/resource_configuration.py b/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_job/resource_configuration.py
@@ -6,7 +6,9 @@
 import logging
 from typing import Any, Dict, Optional
 
-from azure.ai.ml._restclient.v2022_10_01.models import ResourceConfiguration as RestResourceConfiguration
+from azure.ai.ml._restclient.v2023_04_01_preview.models import (
+    ResourceConfiguration as RestResourceConfiguration,
+)
 from azure.ai.ml.constants._job.job import JobComputePropertyFields
 from azure.ai.ml.entities._mixins import DictMixin, RestTranslatableMixin
 
diff --git a/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_job/sweep/objective.py b/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_job/sweep/objective.py
@@ -3,7 +3,9 @@
 # ---------------------------------------------------------
 from typing import Optional
 
-from azure.ai.ml._restclient.v2022_02_01_preview.models import Objective as RestObjective
+from azure.ai.ml._restclient.v2023_04_01_preview.models import (
+    Objective as RestObjective,
+)
 from azure.ai.ml.entities._mixins import RestTranslatableMixin
 
 

Original file line number	Diff line number	Diff line change
`@@ -19,6 +19,9 @@ class JobResourceConfigurationSchema(ResourceConfigurationSchema):`
`19`	`19`	`)`
`20`	`20`	`}`
`21`	`21`	`)`
	`22`	`+ max_instance_count = fields.Int(`
	`23`	`+ metadata={"description": "The maximum number of instances to make available to this job."}`
	`24`	`+ )`
`22`	`25`	`docker_args = fields.Str(metadata={"description": "arguments to pass to the Docker run command."})`
`23`	`26`
`24`	`27`	`@post_load`