Allow registered component to be passed into the sdk. (Azure#32756)

* allow for a registered component to be passed into the sdk for deployment creation

* allow for a completed job to be passed into job definition

* fix pylint error

* account for cli httperror

* black reformatting

* resolve pylint errors

Author: TajiHarrisMicrosoft

sdk/ml/azure-ai-ml/azure/ai/ml/_schema/_deployment/batch/pipeline_component_batch_deployment_schema.py

@@ -6,17 +6,20 @@
 import logging
 from typing import Any
 
-from marshmallow import fields, post_load
+from marshmallow import INCLUDE, fields, post_load
 
 from azure.ai.ml._schema import (
     ArmVersionedStr,
     ArmStr,
     UnionField,
     RegistryStr,
+    NestedField,
 )
+from azure.ai.ml._schema.core.fields import PipelineNodeNameStr, TypeSensitiveUnionField
 from azure.ai.ml._schema._deployment.deployment import DeploymentSchema
 from azure.ai.ml._schema.pipeline.pipeline_component import PipelineComponentFileRefField
 from azure.ai.ml.constants._common import AzureMLResourceType
+from azure.ai.ml.constants._component import NodeType
 
 module_logger = logging.getLogger(__name__)
 
@@ -37,7 +40,7 @@ class PipelineComponentBatchDeploymentSchema(DeploymentSchema):
     job_definition = UnionField(
         [
             ArmStr(azureml_type=AzureMLResourceType.JOB),
-            PipelineComponentFileRefField(),
+            NestedField("PipelineSchema", unknown=INCLUDE),
         ]
     )
 
@@ -48,3 +51,19 @@ def make(self, data: Any, **kwargs: Any) -> Any:  # pylint: disable=unused-argum
         )
 
         return PipelineComponentBatchDeployment(**data)
+
+
+class NodeNameStr(PipelineNodeNameStr):
+    def _get_field_name(self) -> str:
+        return "Pipeline node"
+
+
+def PipelineJobsField():
+    pipeline_enable_job_type = {NodeType.PIPELINE: [NestedField("PipelineSchema", unknown=INCLUDE)]}
+
+    pipeline_job_field = fields.Dict(
+        keys=NodeNameStr(),
+        values=TypeSensitiveUnionField(pipeline_enable_job_type),
+    )
+
+    return pipeline_job_field

sdk/ml/azure-ai-ml/azure/ai/ml/entities/_deployment/pipeline_component_batch_deployment.py

@@ -7,6 +7,7 @@
 from typing import Dict, Optional, Union
 
 from azure.ai.ml.entities._component.component import Component
+from azure.ai.ml.entities._builders import BaseNode
 from azure.ai.ml._schema._deployment.batch.pipeline_component_batch_deployment_schema import (
     PipelineComponentBatchDeploymentSchema,
 )  # pylint: disable=line-too-long
@@ -39,6 +40,8 @@ class PipelineComponentBatchDeployment(Deployment):
     :type description: Optional[str]
     :param tags: A set of tags. The tags which will be applied to the job.
     :type tags: Optional[Dict[str, Any]]
+    :param job_definition: Arm ID or PipelineJob entity of an existing pipeline job.
+    :param job_definition: Optional[Dict[str, ~azure.ai.ml.entities._builders.BaseNode]]
     """
 
     def __init__(
@@ -48,12 +51,13 @@ def __init__(
         endpoint_name: Optional[str] = None,
         component: Optional[Union[Component, str]] = None,
         settings: Optional[Dict[str, str]] = None,
+        job_definition: Optional[Dict[str, BaseNode]] = None,
         **kwargs,  # pylint: disable=unused-argument
     ):
-        self.job_definition = kwargs.pop("job_definition", None)
         super().__init__(endpoint_name=endpoint_name, name=name, **kwargs)
         self.component = component
         self.settings = settings
+        self.job_definition = job_definition
 
     def _to_rest_object(self, location: str) -> "RestBatchDeployment":  # pylint: disable=arguments-differ
         if isinstance(self.component, PipelineComponent):

sdk/ml/azure-ai-ml/azure/ai/ml/operations/_batch_deployment_operations.py

@@ -23,13 +23,14 @@
 from azure.ai.ml._utils._package_utils import package_deployment
 from azure.ai.ml._utils.utils import _get_mfe_base_url_from_discovery_service, modified_operation_client
 from azure.ai.ml.constants._common import ARM_ID_PREFIX, AzureMLResourceType, LROConfigurations
-from azure.ai.ml.entities import BatchDeployment, BatchJob, ModelBatchDeployment, PipelineComponent
+from azure.ai.ml.entities import BatchDeployment, BatchJob, ModelBatchDeployment, PipelineComponent, PipelineJob
 from azure.ai.ml.entities._deployment.deployment import Deployment
 from azure.ai.ml.entities._deployment.pipeline_component_batch_deployment import PipelineComponentBatchDeployment
 from azure.core.credentials import TokenCredential
 from azure.core.paging import ItemPaged
 from azure.core.polling import LROPoller
 from azure.core.tracing.decorator import distributed_trace
+from azure.core.exceptions import ResourceNotFoundError, HttpResponseError
 
 from ._operation_orchestrator import OperationOrchestrator
 
@@ -332,14 +333,25 @@ def _validate_component(self, deployment: Deployment, orchestrators: OperationOr
         :type orchestrators: _operation_orchestrator.OperationOrchestrator
         """
         if isinstance(deployment.component, PipelineComponent):
-            deployment.component = self._all_operations.all_operations[AzureMLResourceType.COMPONENT].create_or_update(
-                name=deployment.component.name,
-                resource_group_name=self._resource_group_name,
-                workspace_name=self._workspace_name,
-                component=deployment.component,
-                version=deployment.component.version,
-                **self._init_kwargs,
-            )
+            try:
+                registered_component = self._all_operations.all_operations[AzureMLResourceType.COMPONENT].get(
+                    name=deployment.component.name, version=deployment.component.version
+                )
+                deployment.component = registered_component.id
+            except Exception as err:  # pylint: disable=broad-except
+                if isinstance(err, (ResourceNotFoundError, HttpResponseError)):
+                    deployment.component = self._all_operations.all_operations[
+                        AzureMLResourceType.COMPONENT
+                    ].create_or_update(
+                        name=deployment.component.name,
+                        resource_group_name=self._resource_group_name,
+                        workspace_name=self._workspace_name,
+                        component=deployment.component,
+                        version=deployment.component.version,
+                        **self._init_kwargs,
+                    )
+                else:
+                    raise err
         elif isinstance(deployment.component, str):
             component_id = orchestrators.get_asset_arm_id(
                 deployment.component, azureml_type=AzureMLResourceType.COMPONENT
@@ -356,3 +368,22 @@ def _validate_component(self, deployment: Deployment, orchestrators: OperationOr
                 **self._init_kwargs,
             )
             deployment.component = job_component.id
+
+        elif isinstance(deployment.job_definition, PipelineJob):
+            try:
+                registered_job = self._all_operations.all_operations[AzureMLResourceType.JOB].get(
+                    name=deployment.job_definition.name
+                )
+                if registered_job:
+                    job_component = PipelineComponent(source_job_id=registered_job.name)
+                    job_component = self._component_operations.create_or_update(
+                        name=job_component.name,
+                        resource_group_name=self._resource_group_name,
+                        workspace_name=self._workspace_name,
+                        body=job_component._to_rest_object(),
+                        version=job_component.version,
+                        **self._init_kwargs,
+                    )
+                    deployment.component = job_component.id
+            except ResourceNotFoundError as err:
+                raise err