allow setting dict in command.identity , sweep.search_space in pipleine (Azure#27696)

* add identity setter for command, add search_space setter for sweep

* resolve comments

* add early_termination setter for sweep, resolve comments

* fix pylint error

* fix pylint error

* resolve comments

Co-authored-by: chenyangliao <[email protected]>

Author: lllilithyang

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

@@ -19,6 +19,7 @@
 from azure.ai.ml._schema.core.fields import NestedField, UnionField
 from azure.ai.ml._schema.job.command_job import CommandJobSchema
 from azure.ai.ml._schema.job.services import JobServiceSchema
+from azure.ai.ml._schema.job.identity import AMLTokenIdentitySchema, ManagedIdentitySchema, UserIdentitySchema
 from azure.ai.ml.constants._common import BASE_PATH_CONTEXT_KEY, LOCAL_COMPUTE_PROPERTY, LOCAL_COMPUTE_TARGET
 from azure.ai.ml.constants._component import ComponentSource, NodeType
 from azure.ai.ml.entities._assets import Environment
@@ -240,6 +241,32 @@ def resources(self, value: Union[Dict, JobResourceConfiguration]):
             value = JobResourceConfiguration(**value)
         self._resources = value
 
+    @property
+    def identity(
+        self,
+    ) -> Optional[Union[ManagedIdentityConfiguration, AmlTokenConfiguration, UserIdentityConfiguration]]:
+        """
+        Configuration of the hyperparameter identity.
+        """
+        return self._identity
+
+    @identity.setter
+    def identity(self, value: Union[
+                                Dict[str, str],
+                                ManagedIdentityConfiguration,
+                                AmlTokenConfiguration,
+                                UserIdentityConfiguration, None]):
+        if isinstance(value, dict):
+            identity_schema = UnionField(
+                [
+                    NestedField(ManagedIdentitySchema, unknown=INCLUDE),
+                    NestedField(AMLTokenIdentitySchema, unknown=INCLUDE),
+                    NestedField(UserIdentitySchema, unknown=INCLUDE),
+                ]
+            )
+            value = identity_schema._deserialize(value=value, attr=None, data=None)
+        self._identity = value
+
     @property
     def services(self) -> Dict:
         return self._services

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

@@ -11,6 +11,7 @@
 
 from azure.ai.ml.constants._common import BASE_PATH_CONTEXT_KEY
 from azure.ai.ml.constants._component import NodeType
+from azure.ai.ml.constants._job.sweep import SearchSpace
 from azure.ai.ml.entities._component.command_component import CommandComponent
 from azure.ai.ml.entities._inputs_outputs import Input, Output
 from azure.ai.ml.entities._credentials import (
@@ -24,6 +25,7 @@
     BanditPolicy,
     MedianStoppingPolicy,
     TruncationSelectionPolicy,
+    EarlyTerminationPolicy,
 )
 from azure.ai.ml.entities._job.sweep.objective import Objective
 from azure.ai.ml.entities._job.sweep.parameterized_sweep import ParameterizedSweep
@@ -41,8 +43,14 @@
     SweepDistribution,
     Uniform,
 )
-from azure.ai.ml.exceptions import ErrorTarget, UserErrorException, ValidationErrorType, ValidationException
+from azure.ai.ml.exceptions import (
+    ErrorTarget,
+    UserErrorException,
+    ValidationErrorType,
+    ValidationException,
+)
 from azure.ai.ml.sweep import SweepJob
+from azure.ai.ml._schema._sweep.sweep_fields_provider import EarlyTerminationField
 
 from ..._schema import PathAwareSchema
 from ..._schema._utils.data_binding_expression import support_data_binding_expression_for_fields
@@ -149,6 +157,32 @@ def trial(self):
         """Id or instance of the command component/job to be run for the step."""
         return self._component
 
+    @property
+    def search_space(self):
+        """Dictionary of the hyperparameter search space.
+        The key is the name of the hyperparameter and the value is the parameter expression.
+        """
+        return self._search_space
+
+    @search_space.setter
+    def search_space(self, values: Dict[str, Dict[str, Union[str, int, float, dict]]]):
+        search_space = {}
+        for name, value in values.items():
+            # If value is a SearchSpace object, directly pass it to job.search_space[name]
+            search_space[name] = self._value_type_to_class(value) if isinstance(value, dict) else value
+        self._search_space = search_space
+
+    @classmethod
+    def _value_type_to_class(cls, value):
+        value_type = value['type']
+        search_space_dict = {
+            SearchSpace.CHOICE: Choice, SearchSpace.RANDINT: Randint, SearchSpace.LOGNORMAL: LogNormal,
+            SearchSpace.NORMAL: Normal, SearchSpace.LOGUNIFORM: LogUniform, SearchSpace.UNIFORM: Uniform,
+            SearchSpace.QLOGNORMAL: QLogNormal, SearchSpace.QNORMAL: QNormal, SearchSpace.QLOGUNIFORM: QLogUniform,
+            SearchSpace.QUNIFORM: QUniform
+        }
+        return search_space_dict[value_type](**value)
+
     @classmethod
     def _get_supported_inputs_types(cls):
         supported_types = super()._get_supported_inputs_types() or ()
@@ -320,3 +354,14 @@ def __setattr__(self, key, value):
                 self.early_termination.slack_amount = None
             if self.early_termination.slack_factor == 0.0:
                 self.early_termination.slack_factor = None
+
+    @property
+    def early_termination(self) -> Union[str, EarlyTerminationPolicy]:
+        return self._early_termination
+
+    @early_termination.setter
+    def early_termination(self, value: Union[EarlyTerminationPolicy, Dict[str, Union[str, float, int, bool]]]):
+        if isinstance(value, dict):
+            early_termination_schema = EarlyTerminationField()
+            value = early_termination_schema._deserialize(value=value, attr=None, data=None)
+        self._early_termination = value

sdk/ml/azure-ai-ml/tests/component/unittests/test_command_component_entity.py

@@ -367,6 +367,28 @@ def test_sweep_help_function(self):
                 in std_out.getvalue()
             )
 
+    def test_sweep_early_termination_setter(self):
+        yaml_file = "./tests/test_configs/components/helloworld_component.yml"
+
+        component_to_sweep: CommandComponent = load_component(source=yaml_file)
+        cmd_node1: Command = component_to_sweep(
+            component_in_number=Choice([2, 3, 4, 5]), component_in_path=Input(path="/a/path/on/ds")
+        )
+
+        sweep_job1: Sweep = cmd_node1.sweep(
+            primary_metric="AUC",  # primary_metric,
+            goal="maximize",
+            sampling_algorithm="random",
+        )
+        sweep_job1.early_termination = {
+            'type': "bandit", 'evaluation_interval': 100, 'delay_evaluation': 200, 'slack_factor': 40.0
+            }
+        from azure.ai.ml.entities._job.sweep.early_termination_policy import BanditPolicy
+        assert isinstance(sweep_job1.early_termination, BanditPolicy)
+        assert [sweep_job1.early_termination.evaluation_interval,
+                sweep_job1.early_termination.delay_evaluation,
+                sweep_job1.early_termination.slack_factor] == [100, 200, 40.0]
+
     def test_invalid_component_inputs(self) -> None:
         yaml_path = "./tests/test_configs/components/invalid/helloworld_component_conflict_input_names.yml"
         component = load_component(yaml_path)

sdk/ml/azure-ai-ml/tests/dsl/unittests/test_command_builder.py

@@ -970,3 +970,32 @@ def my_pipeline():
                 "type": "command",
             }
         }
+
+    def test_set_identity(self, test_command):
+        from azure.ai.ml.entities._credentials import AmlTokenConfiguration
+        node1 = test_command()
+        node2 = node1()
+        node2.identity = AmlTokenConfiguration()
+        node3 = node1()
+        node3.identity = {'type': 'AMLToken'}
+        assert node2.identity == node3.identity
+
+    def test_sweep_set_search_space(self, test_command):
+        from azure.ai.ml.entities._job.sweep.search_space import Choice
+        node1 = test_command()
+        command_node_to_sweep_1 = node1()
+        sweep_node_1 = command_node_to_sweep_1.sweep(
+            primary_metric="AUC",
+            goal="maximize",
+            sampling_algorithm="random",
+        )
+        sweep_node_1.search_space = {'batch_size': {'type': 'choice', 'values': [25, 35]}}
+
+        command_node_to_sweep_2 = node1()
+        sweep_node_2 = command_node_to_sweep_2.sweep(
+            primary_metric="AUC",
+            goal="maximize",
+            sampling_algorithm="random",
+        )
+        sweep_node_2.search_space = {'batch_size': Choice(values=[25, 35])}
+        assert sweep_node_1.search_space == sweep_node_2.search_space