Parent job name & multiple docker argument support for command job (Azure#39432)

* testing

* adding 2025_01_01_preview client

* adding multiple docker argument and parent job support for command job

* test case fix

* api review comment for kwargs in job entity

* remove unuse entry in change log

* fixing test findings

* command job adding key word only

* fixing unit test cases

* fixing end to end test

* merging master

* reducing the scope of job configuration 2025 only in caseof multiple docker argument

* Adding 2025 job configuration in from rest object

* fixing failed test

* adding more unit test cases

Author: achauhan-scc

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

@@ -1,6 +1,10 @@
 # Release History
 ## 1.26.0 (unreleased)
 
+### Features Added
+  - Adding parent job support for command job.
+  - Adding multiple docker argument support for command job.
+
 ### Bugs Fixed
   - #37464 - Allowing to update a component in register with anonymousEnvironment environment.
   - #39744 - Fixing a serialization issue in RecurrencePattern schedule.

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

@@ -40,6 +40,7 @@
 from azure.ai.ml._restclient.v2024_04_01_preview import AzureMachineLearningWorkspaces as ServiceClient042024Preview
 from azure.ai.ml._restclient.v2024_07_01_preview import AzureMachineLearningWorkspaces as ServiceClient072024Preview
 from azure.ai.ml._restclient.v2024_10_01_preview import AzureMachineLearningWorkspaces as ServiceClient102024Preview
+from azure.ai.ml._restclient.v2025_01_01_preview import AzureMachineLearningWorkspaces as ServiceClient012025Preview
 from azure.ai.ml._restclient.workspace_dataplane import (
     AzureMachineLearningWorkspaces as ServiceClientWorkspaceDataplane,
 )
@@ -396,6 +397,17 @@ def __init__(
             **kwargs,
         )
 
+        self._service_client_01_2025_preview = ServiceClient012025Preview(
+            credential=self._credential,
+            subscription_id=(
+                self._ws_operation_scope._subscription_id
+                if registry_reference
+                else self._operation_scope._subscription_id
+            ),
+            base_url=base_url,
+            **kwargs,
+        )
+
         # A general purpose, user-configurable pipeline for making
         # http requests
         self._requests_pipeline = HttpPipeline(**kwargs)
@@ -683,6 +695,7 @@ def __init__(
             requests_pipeline=self._requests_pipeline,
             service_client_01_2024_preview=self._service_client_01_2024_preview,
             service_client_10_2024_preview=self._service_client_10_2024_preview,
+            service_client_01_2025_preview=self._service_client_01_2025_preview,
             **ops_kwargs,
         )
         self._operation_container.add(AzureMLResourceType.JOB, self._jobs)

sdk/ml/azure-ai-ml/azure/ai/ml/_restclient/v2025_01_01_preview/models/_models.py

@@ -34032,4 +34032,4 @@ def __init__(
         self.service_managed_resources_settings = kwargs.get('service_managed_resources_settings', None)
         self.soft_delete_retention_in_days = kwargs.get('soft_delete_retention_in_days', None)
         self.system_datastores_auth_mode = kwargs.get('system_datastores_auth_mode', None)
-        self.v1_legacy_mode = kwargs.get('v1_legacy_mode', None)
+        self.v1_legacy_mode = kwargs.get('v1_legacy_mode', None)

sdk/ml/azure-ai-ml/azure/ai/ml/_schema/job/command_job.py

@@ -20,3 +20,4 @@ class CommandJobSchema(ParameterizedCommandSchema, BaseJobSchema):
     parameters = fields.Dict(dump_only=True)
     inputs = InputsField()
     outputs = OutputsField()
+    parent_job_name = fields.Str()

sdk/ml/azure-ai-ml/azure/ai/ml/_schema/job_resource_configuration.py

@@ -4,6 +4,8 @@
 
 from marshmallow import fields, post_load
 
+from azure.ai.ml._schema.core.fields import UnionField
+
 from .resource_configuration import ResourceConfigurationSchema
 
 
@@ -22,7 +24,12 @@ 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."})
+    docker_args = UnionField(
+        [
+            fields.Str(metadata={"description": "arguments to pass to the Docker run command."}),
+            fields.List(fields.Str()),
+        ]
+    )
 
     @post_load
     def make(self, data, **kwargs):

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

@@ -11,8 +11,8 @@
 
 from marshmallow import INCLUDE, Schema
 
-from azure.ai.ml._restclient.v2023_04_01_preview.models import CommandJob as RestCommandJob
-from azure.ai.ml._restclient.v2023_04_01_preview.models import JobBase
+from azure.ai.ml._restclient.v2025_01_01_preview.models import CommandJob as RestCommandJob
+from azure.ai.ml._restclient.v2025_01_01_preview.models import JobBase
 from azure.ai.ml._schema.core.fields import ExperimentalField, NestedField, UnionField
 from azure.ai.ml._schema.job.command_job import CommandJobSchema
 from azure.ai.ml._schema.job.identity import AMLTokenIdentitySchema, ManagedIdentitySchema, UserIdentitySchema
@@ -127,6 +127,8 @@ class Command(BaseNode, NodeWithGroupInputMixin):
         ~azure.ai.ml.entities.VsCodeJobService]]]
     :keyword queue_settings: Queue settings for the job.
     :paramtype queue_settings: Optional[~azure.ai.ml.entities.QueueSettings]
+    :keyword parent_job_name: parent job id for command job
+    :paramtype parent_job_name: Optional[str]
     :raises ~azure.ai.ml.exceptions.ValidationException: Raised if Command cannot be successfully validated.
         Details will be provided in the error message.
     """
@@ -172,6 +174,7 @@ def __init__(
             Dict[str, Union[JobService, JupyterLabJobService, SshJobService, TensorBoardJobService, VsCodeJobService]]
         ] = None,
         queue_settings: Optional[QueueSettings] = None,
+        parent_job_name: Optional[str] = None,
         **kwargs: Any,
     ) -> None:
         # validate init params are valid type
@@ -203,6 +206,7 @@ def __init__(
         self._resources = resources
         self._services = services
         self.queue_settings = queue_settings
+        self.parent_job_name = parent_job_name
 
         if isinstance(self.component, CommandComponent):
             self.resources = self.resources or self.component.resources  # type: ignore[assignment]
@@ -473,7 +477,7 @@ def set_resources(
         instance_count: Optional[int] = None,
         locations: Optional[List[str]] = None,
         properties: Optional[Dict] = None,
-        docker_args: Optional[str] = None,
+        docker_args: Optional[Union[str, List[str]]] = None,
         shm_size: Optional[str] = None,
         # pylint: disable=unused-argument
         **kwargs: Any,
@@ -492,7 +496,7 @@ def set_resources(
         :keyword properties: The properties of the job.
         :paramtype properties: Optional[dict]
         :keyword docker_args: The Docker arguments for the job.
-        :paramtype docker_args: Optional[str]
+        :paramtype docker_args: Optional[Union[str,List[str]]]
         :keyword shm_size: The size of the docker container's shared memory block. This should be in the
             format of (number)(unit) where the number has to be greater than 0 and the unit can be one of
             b(bytes), k(kilobytes), m(megabytes), or g(gigabytes).
@@ -745,6 +749,7 @@ def _to_job(self) -> CommandJob:
                 creation_context=self.creation_context,
                 parameters=self.parameters,
                 queue_settings=self.queue_settings,
+                parent_job_name=self.parent_job_name,
             )
 
         return CommandJob(
@@ -771,6 +776,7 @@ def _to_job(self) -> CommandJob:
             creation_context=self.creation_context,
             parameters=self.parameters,
             queue_settings=self.queue_settings,
+            parent_job_name=self.parent_job_name,
         )
 
     @classmethod

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

@@ -145,7 +145,7 @@ def command(
     instance_count: Optional[int] = None,
     instance_type: Optional[str] = None,
     locations: Optional[List[str]] = None,
-    docker_args: Optional[str] = None,
+    docker_args: Optional[Union[str, List[str]]] = None,
     shm_size: Optional[str] = None,
     timeout: Optional[int] = None,
     code: Optional[Union[str, os.PathLike]] = None,
@@ -156,6 +156,7 @@ def command(
     ] = None,
     job_tier: Optional[str] = None,
     priority: Optional[str] = None,
+    parent_job_name: Optional[str] = None,
     **kwargs: Any,
 ) -> Command:
     """Creates a Command object which can be used inside a dsl.pipeline function or used as a standalone Command job.
@@ -199,7 +200,7 @@ def command(
     :keyword docker_args: Extra arguments to pass to the Docker run command. This would override any
         parameters that have already been set by the system, or in this section. This parameter is only
         supported for Azure ML compute types. Defaults to None.
-    :paramtype docker_args: Optional[str]
+    :paramtype docker_args: Optional[Union[str,List[str]]]
     :keyword shm_size: The size of the Docker container's shared memory block. This should be in the
         format of (number)(unit) where the number has to be greater than 0 and the unit can be one of
         b(bytes), k(kilobytes), m(megabytes), or g(gigabytes).
@@ -229,6 +230,8 @@ def command(
     :keyword priority: The priority of the job on the compute. Accepted values are "low", "medium", and "high".
         Defaults to "medium".
     :paramtype priority: Optional[str]
+    :keyword parent_job_name: parent job id for command job
+    :paramtype parent_job_name: Optional[str]
     :return: A Command object.
     :rtype: ~azure.ai.ml.entities.Command
 
@@ -283,6 +286,7 @@ def command(
         environment=environment,
         environment_variables=environment_variables,
         services=services,
+        parent_job_name=parent_job_name,
         **kwargs,
     )
 

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

@@ -9,8 +9,8 @@
 from pathlib import Path
 from typing import TYPE_CHECKING, Any, Dict, Optional, Union
 
-from azure.ai.ml._restclient.v2023_04_01_preview.models import CommandJob as RestCommandJob
-from azure.ai.ml._restclient.v2023_04_01_preview.models import JobBase
+from azure.ai.ml._restclient.v2025_01_01_preview.models import CommandJob as RestCommandJob
+from azure.ai.ml._restclient.v2025_01_01_preview.models import JobBase
 from azure.ai.ml._schema.job.command_job import CommandJobSchema
 from azure.ai.ml._utils.utils import map_single_brackets_and_warn
 from azure.ai.ml.constants import JobType
@@ -77,6 +77,8 @@ class CommandJob(Job, ParameterizedCommand, JobIOMixin):
         ~azure.ai.ml.UserIdentityConfiguration]]
     :keyword limits: The limits for the job.
     :paramtype limits: Optional[~azure.ai.ml.entities.CommandJobLimits]
+    :keyword parent_job_name: parent job id for command job
+    :paramtype parent_job_name: Optional[str]
     :keyword kwargs: A dictionary of additional configuration parameters.
     :paramtype kwargs: dict
 
@@ -103,10 +105,12 @@ def __init__(
         services: Optional[
             Dict[str, Union[JobService, JupyterLabJobService, SshJobService, TensorBoardJobService, VsCodeJobService]]
         ] = None,
+        parent_job_name: Optional[str] = None,
         **kwargs: Any,
     ) -> None:
         kwargs[TYPE] = JobType.COMMAND
         self._parameters: dict = kwargs.pop("parameters", {})
+        self.parent_job_name = parent_job_name
 
         super().__init__(**kwargs)
 
@@ -173,6 +177,7 @@ def _to_rest_object(self) -> JobBase:
             limits=self.limits._to_rest_object() if self.limits else None,
             services=JobServiceBase._to_rest_job_services(self.services),
             queue_settings=self.queue_settings._to_rest_object() if self.queue_settings else None,
+            parent_job_name=self.parent_job_name,
         )
         result = JobBase(properties=properties)
         result.name = self.name
@@ -215,6 +220,7 @@ def _load_from_rest(cls, obj: JobBase) -> "CommandJob":
             inputs=from_rest_inputs_to_dataset_literal(rest_command_job.inputs),
             outputs=from_rest_data_outputs(rest_command_job.outputs),
             queue_settings=QueueSettings._from_rest_object(rest_command_job.queue_settings),
+            parent_job_name=rest_command_job.parent_job_name,
         )
         # Handle special case of local job
         if (

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

@@ -7,6 +7,9 @@
 from typing import Any, Dict, List, Optional, Union, cast
 
 from azure.ai.ml._restclient.v2023_04_01_preview.models import JobResourceConfiguration as RestJobResourceConfiguration
+from azure.ai.ml._restclient.v2025_01_01_preview.models import (
+    JobResourceConfiguration as RestJobResourceConfiguration202501,
+)
 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
@@ -98,7 +101,7 @@ class JobResourceConfiguration(RestTranslatableMixin, DictMixin):
     :keyword docker_args: Extra arguments to pass to the Docker run command. This would override any
         parameters that have already been set by the system, or in this section. This parameter is only
         supported for Azure ML compute types.
-    :paramtype docker_args: Optional[str]
+    :paramtype docker_args: Optional[Union[str, List[str]]]
     :keyword shm_size: The size of the docker container's shared memory block. This should be in the
         format of (number)(unit) where the number has to be greater than 0 and the unit can be one of
         b(bytes), k(kilobytes), m(megabytes), or g(gigabytes).
@@ -125,7 +128,7 @@ def __init__(
         instance_count: Optional[int] = None,
         instance_type: Optional[Union[str, List]] = None,
         properties: Optional[Union[Properties, Dict]] = None,
-        docker_args: Optional[str] = None,
+        docker_args: Optional[Union[str, List[str]]] = None,
         shm_size: Optional[str] = None,
         max_instance_count: Optional[int] = None,
         **kwargs: Any
@@ -162,7 +165,16 @@ def properties(self, properties: Dict[str, Any]) -> None:
         else:
             raise TypeError("properties must be a dict.")
 
-    def _to_rest_object(self) -> RestJobResourceConfiguration:
+    def _to_rest_object(self) -> Union[RestJobResourceConfiguration, RestJobResourceConfiguration202501]:
+        if self.docker_args and isinstance(self.docker_args, list):
+            return RestJobResourceConfiguration202501(
+                instance_count=self.instance_count,
+                instance_type=self.instance_type,
+                max_instance_count=self.max_instance_count,
+                properties=self.properties.as_dict() if isinstance(self.properties, Properties) else None,
+                docker_args_list=self.docker_args,
+                shm_size=self.shm_size,
+            )
         return RestJobResourceConfiguration(
             locations=self.locations,
             instance_count=self.instance_count,
@@ -174,18 +186,20 @@ def _to_rest_object(self) -> RestJobResourceConfiguration:
         )
 
     @classmethod
-    def _from_rest_object(cls, obj: Optional[RestJobResourceConfiguration]) -> Optional["JobResourceConfiguration"]:
+    def _from_rest_object(
+        cls, obj: Optional[Union[RestJobResourceConfiguration, RestJobResourceConfiguration202501]]
+    ) -> Optional["JobResourceConfiguration"]:
         if obj is None:
             return None
         if isinstance(obj, dict):
             return cls(**obj)
         return JobResourceConfiguration(
-            locations=obj.locations,
+            locations=obj.locations if hasattr(obj, "locations") else None,
             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,
+            docker_args=obj.docker_args_list if hasattr(obj, "docker_args_list") else obj.docker_args,
             shm_size=obj.shm_size,
             deserialize_properties=True,
         )

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

@@ -9,6 +9,7 @@
 from typing import Any
 
 from azure.ai.ml._restclient.v2023_08_01_preview.models import JobBase as JobBaseData
+from azure.ai.ml._restclient.v2025_01_01_preview.models import JobBase as JobBaseData202501
 from azure.ai.ml.constants._common import DEFAULT_EXPERIMENT_NAME
 from azure.ai.ml.entities._builders.command import Command
 from azure.ai.ml.entities._builders.pipeline import Pipeline
@@ -47,7 +48,7 @@ def _(job: Job) -> JobBaseData:
 
 
 @to_rest_job_object.register(Command)
-def _(command: Command) -> JobBaseData:
+def _(command: Command) -> JobBaseData202501:
     rest_job = command._to_job()._to_rest_object()
     generate_defaults(command, rest_job)
     return rest_job