[AI] Convert Deployment Create/Update to an LRO (#33419)

needuv · web-flow · commit 36a0df021492 · 2023-12-07T17:58:16.000Z
* add ARM template creation

* verify local upload for deployment works
diff --git a/sdk/ai/azure-ai-resources/azure/ai/resources/_utils/_deployment_utils.py b/sdk/ai/azure-ai-resources/azure/ai/resources/_utils/_deployment_utils.py
@@ -10,6 +10,93 @@
 from ._registry_utils import get_registry_model
 
 
+def get_empty_deployment_arm_template():
+    return {
+        "$schema": "https://schema.management.azure.com/schemas/2015-01-01/deploymentTemplate.json#",
+        "contentVersion": "1.0.0.0",
+        "parameters": {
+            "onlineEndpointProperties": {
+                "defaultValue": {
+                    "authMode": "Key",
+                    "publicNetworkAccess": "Enabled",
+                    "properties": {
+                        "enforce_access_to_default_secret_stores": "enabled"
+                    },
+                },
+                "type": "Object",
+            },
+            "onlineEndpointPropertiesTrafficUpdate": {
+                "defaultValue": {
+                    "traffic": {"[parameters('onlineDeploymentName')]": 100},
+                    "authMode": "Key",
+                    "publicNetworkAccess": "Enabled",
+                    "properties": {
+                        "enforce_access_to_default_secret_stores": "enabled"
+                    },
+                },
+                "type": "Object",
+            },
+        },
+        "resources": [
+            {
+                "type": "Microsoft.MachineLearningServices/workspaces/onlineEndpoints",
+                "apiVersion": "2023-04-01-Preview",
+                "name": "[concat(parameters('workspaceName'), '/', parameters('onlineEndpointName'))]",
+                "location": "[parameters('location')]",
+                "identity": {"type": "SystemAssigned"},
+                "properties": "[parameters('onlineEndpointProperties')]",
+                "copy": {"name": "onlineEndpointCopy", "count": 1, "mode": "serial"},
+            },
+            {
+                "type": "Microsoft.MachineLearningServices/workspaces/onlineEndpoints/deployments",
+                "apiVersion": "2023-04-01-Preview",
+                "name": "[concat(parameters('workspaceName'), '/', parameters('onlineEndpointName'), '/', parameters('onlineDeploymentName'))]",
+                "location": "[parameters('location')]",
+                "dependsOn": [
+                    "onlineEndpointCopy",
+                ],
+                "sku": {"capacity": "[parameters('deploymentInstanceCount')]", "name": "default"},
+                "identity": {"type": "None"},
+                "properties": "[parameters('onlineDeploymentProperties')]",
+                "copy": {"name": "onlineDeploymentCopy", "count": 1, "mode": "serial"},
+            },
+            {
+                "type": "Microsoft.Resources/deployments",
+                "apiVersion": "2015-01-01",
+                "name": "[concat('updateEndpointWithTraffic', '-', parameters('onlineEndpointName'))]",
+                "dependsOn": ["onlineDeploymentCopy"],
+                "properties": {
+                    "mode": "Incremental",
+                    "template": {
+                        "$schema": "https://schema.management.azure.com/schemas/2015-01-01/deploymentTemplate.json#",
+                        "contentVersion": "1.0.0.0",
+                        "resources": [
+                            {
+                                "type": "Microsoft.MachineLearningServices/workspaces/onlineEndpoints",
+                                "apiVersion": "2023-04-01-Preview",
+                                "location": "[parameters('location')]",
+                                "name": "[concat(parameters('workspaceName'), '/', parameters('onlineEndpointName'))]",
+                                "properties": "[parameters('onlineEndpointPropertiesTrafficUpdate')]",
+                                "identity": {"type": "SystemAssigned"},
+                            }
+                        ],
+                    },
+                },
+            },
+        ],
+        "outputs": {
+            "online_endpoint_name": {
+                "type": "string",
+                "value": "[parameters('onlineEndpointName')]",
+            },
+            "online_deployment_name": {
+                "type": "string",
+                "value": "[parameters('onlineDeploymentName')]",
+            },
+        },
+    }
+
+
 def get_default_allowed_instance_type_for_hugging_face(
     model_details: Model, credential: Any
 ) -> Tuple[str, str]:
@@ -23,31 +110,46 @@ def get_default_allowed_instance_type_for_hugging_face(
         hf_engines = hf_engines.split(",")
         if len(hf_engines) > 1:
             for engine_id in hf_engines:
-                instance_type, instance_type_list = get_default_allowed_instance_type_from_model_engine(engine_id, credential)
+                (
+                    instance_type,
+                    instance_type_list,
+                ) = get_default_allowed_instance_type_from_model_engine(
+                    engine_id, credential
+                )
                 if "cpu" in engine_id:
                     default_instance_type = instance_type
                 allowed_instance_types.append(instance_type_list)
         else:
             # if the model has only one engine, we can proceed with that as the default engine for SKU
             # selection
-            default_instance_type, allowed_instance_types = get_default_allowed_instance_type_from_model_engine(hf_engines[0], credential)
+            (
+                default_instance_type,
+                allowed_instance_types,
+            ) = get_default_allowed_instance_type_from_model_engine(
+                hf_engines[0], credential
+            )
     else:
-        default_instance_type, allowed_instance_types = parse_deployment_config(deployment_config)
+        default_instance_type, allowed_instance_types = parse_deployment_config(
+            deployment_config
+        )
     return (default_instance_type, allowed_instance_types)
 
 
 def parse_deployment_config(deployment_config: str):
     deployment_config = json.loads(deployment_config)
-    allowed_instance_types = deployment_config["PipelineMetadata"]["PipelineDefinition"]["ec"]["AllowedInstanceTypes"]
-    default_instance_type = deployment_config["PipelineMetadata"]["PipelineDefinition"]["ec"]["DefaultInstanceType"]
+    allowed_instance_types = deployment_config["PipelineMetadata"][
+        "PipelineDefinition"
+    ]["ec"]["AllowedInstanceTypes"]
+    default_instance_type = deployment_config["PipelineMetadata"]["PipelineDefinition"][
+        "ec"
+    ]["DefaultInstanceType"]
 
     return (default_instance_type, allowed_instance_types)
 
 
-def get_default_allowed_instance_type_from_model_engine(engine_id: str, credential: Any):
-    model_details = get_registry_model(
-        credential,
-        id=engine_id
-    )
+def get_default_allowed_instance_type_from_model_engine(
+    engine_id: str, credential: Any
+):
+    model_details = get_registry_model(credential, id=engine_id)
     deployment_config = model_details.properties.get("modelDeploymentConfig", None)
-    return parse_deployment_config(deployment_config)
+    return parse_deployment_config(deployment_config)
diff --git a/sdk/ai/azure-ai-resources/azure/ai/resources/entities/models.py b/sdk/ai/azure-ai-resources/azure/ai/resources/entities/models.py
@@ -7,12 +7,6 @@
 from typing import Union
 
 
-@dataclass
-class LangchainModel:
-    chain: "langchain.chains.Chain"
-    conda_file: Union[str, os.PathLike]
-
-
 @dataclass
 class PromptflowModel:
     path: Union[str, os.PathLike]
diff --git a/sdk/ai/azure-ai-resources/azure/ai/resources/operations/_deployment_operations.py b/sdk/ai/azure-ai-resources/azure/ai/resources/operations/_deployment_operations.py
@@ -8,15 +8,21 @@
 import shutil
 import tempfile
 from typing import Any, List, Union, Iterable
+import uuid
 
 
 from azure.ai.ml import MLClient
-from azure.ai.ml.entities import ManagedOnlineDeployment, ManagedOnlineEndpoint, Model as AzureMLModel, DataCollector, DeploymentCollection, Environment, BuildContext
+from azure.ai.ml.entities import ManagedOnlineDeployment, Model as AzureMLModel, DataCollector, DeploymentCollection, Environment, BuildContext
+from azure.ai.ml.operations._operation_orchestrator import OperationOrchestrator
+from azure.ai.ml._utils._endpoint_utils import upload_dependencies
 from azure.core.tracing.decorator import distributed_trace
+from azure.mgmt.resource import ResourceManagementClient
+from azure.mgmt.resource.resources.models import DeploymentMode
+from azure.core.polling import LROPoller
 
 from .._utils._scoring_script_utils import create_chat_scoring_script, create_mlmodel_file
 from .._utils._registry_utils import get_registry_model
-from .._utils._deployment_utils import get_default_allowed_instance_type_for_hugging_face
+from .._utils._deployment_utils import get_default_allowed_instance_type_for_hugging_face, get_empty_deployment_arm_template
 from ..entities.deployment import Deployment
 from ..entities.deployment_keys import DeploymentKeys
 from ..entities.models import Model, PromptflowModel
@@ -32,10 +38,11 @@ def __init__(self, ml_client: MLClient, connections, **kwargs) -> None:
         self._ml_client = ml_client
         self._connections = connections
         ops_logger.update_info(kwargs)
+        self._resource_management_client = ResourceManagementClient(self._ml_client._credential, self._ml_client.subscription_id)
 
     @distributed_trace
-    @monitor_with_activity(logger, "Deployment.CreateOrUpdate", ActivityType.PUBLICAPI)
-    def create_or_update(self, deployment: Deployment) -> Any:
+    @monitor_with_activity(logger, "Deployment.BeginCreateOrUpdate", ActivityType.PUBLICAPI)
+    def begin_create_or_update(self, deployment: Deployment) -> LROPoller[Deployment]:
         model = deployment.model
         endpoint_name = deployment.endpoint_name if deployment.endpoint_name else deployment.name
 
@@ -53,13 +60,6 @@ def create_or_update(self, deployment: Deployment) -> Any:
                 sampling_rate=1,
             )
 
-        v2_endpoint = ManagedOnlineEndpoint(
-            name=endpoint_name,
-            properties={
-                "enforce_access_to_default_secret_stores": "enabled"
-            }
-        )
-        created_endpoint = self._ml_client.begin_create_or_update(v2_endpoint).result()
         model = deployment.model
         v2_deployment = None
         temp_dir = tempfile.TemporaryDirectory()
@@ -227,18 +227,46 @@ def create_or_update(self, deployment: Deployment) -> Any:
                 app_insights_enabled=deployment.app_insights_enabled,
                 data_collector=data_collector,
             )
-
+        if deployment.data_collector_enabled:
+            self._ml_client.online_deployments._register_collection_data_assets(v2_deployment)
+        orchestrators = OperationOrchestrator(
+            operation_container=self._ml_client.online_deployments._all_operations,
+            operation_scope=self._ml_client.online_deployments._operation_scope,
+            operation_config=self._ml_client.online_deployments._operation_config,
+        )
+        upload_dependencies(v2_deployment, orchestrators)
+        location = self._ml_client.online_deployments._get_workspace_location()
         v2_deployment.tags = deployment.tags
         v2_deployment.properties = deployment.properties
-        create_deployment_poller = self._ml_client.begin_create_or_update(v2_deployment)
         shutil.rmtree(temp_dir.name)
-        created_deployment = create_deployment_poller.result()
+    
+        template = get_empty_deployment_arm_template()
+        template["parameters"]["workspaceName"] = {"defaultValue": self._ml_client.workspace_name, "type": "String"}
+        template["parameters"]["onlineEndpointName"] = {"defaultValue": endpoint_name, "type": "String"}
+        template["parameters"]["onlineDeploymentName"] = {"defaultValue": deployment.name, "type": "String"}
+        template["parameters"]["onlineDeploymentProperties"] = {"defaultValue": v2_deployment._to_rest_object(location=location).properties.serialize(), "type": "Object"}
+        template["parameters"]["location"] = {"defaultValue": location, "type": "String"}
+        template["parameters"]["deploymentInstanceCount"] = {"defaultValue": deployment.instance_count, "type": "int"}
 
-        v2_endpoint.traffic = {deployment.name: 100}
-        update_endpoint_poller = self._ml_client.begin_create_or_update(v2_endpoint)
-        updated_endpoint = update_endpoint_poller.result()
 
-        return Deployment._from_v2_endpoint_deployment(updated_endpoint, deployment)
+        def lro_callback(raw_response, deserialized, headers):
+            outputs = deserialized.properties.outputs
+            return Deployment._from_v2_endpoint_deployment(
+                self._ml_client.online_endpoints.get(outputs["online_endpoint_name"]["value"]),
+                self._ml_client.online_deployments.get(outputs["online_deployment_name"]["value"], outputs["online_endpoint_name"]["value"])
+            )
+
+        return self._resource_management_client.deployments.begin_create_or_update(
+            self._ml_client.resource_group_name,
+            str(uuid.uuid4()),
+            {
+                "properties": {
+                    "template": template,
+                    "mode": DeploymentMode.incremental,
+                }
+            },
+            cls=lro_callback,
+        )
 
     @distributed_trace
     @monitor_with_activity(logger, "Deployment.Get", ActivityType.PUBLICAPI)
diff --git a/sdk/ai/azure-ai-resources/setup.py b/sdk/ai/azure-ai-resources/setup.py
@@ -69,6 +69,7 @@
         "azure-ai-ml>=1.12.0",
         "mlflow-skinny<3",
         "azureml-telemetry~=1.0,>=1.51.0",
+        "azure-mgmt-resource<23.0.0,>=22.0.0",
     ],
     project_urls={
         "Bug Reports": "https://github.com/Azure/azure-sdk-for-python/issues",
