Bump restclient version for model package and model ops (Azure#32531)

* fixes

* fixes

* fixes

* fixes

* test recording

* test recording

* target env fix

* pylint & black

* e2e test

* e2e test

* pylint

Author: nemanjarajic

sdk/ml/azure-ai-ml/assets.json

@@ -2,5 +2,5 @@
   "AssetsRepo": "Azure/azure-sdk-assets",
   "AssetsRepoPrefixPath": "python",
   "TagPrefix": "python/ml/azure-ai-ml",
-  "Tag": "python/ml/azure-ai-ml_748fafdb82"
+  "Tag": "python/ml/azure-ai-ml_26cf777aef"
 }

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

@@ -456,11 +456,11 @@ def __init__(
             self._operation_config,
             self._service_client_10_2021_dataplanepreview
             if registry_name or registry_reference
-            else self._service_client_04_2023_preview,
+            else self._service_client_08_2023_preview,
             self._datastores,
             self._operation_container,
             requests_pipeline=self._requests_pipeline,
-            control_plane_client=self._service_client_04_2023_preview,
+            control_plane_client=self._service_client_08_2023_preview,
             workspace_rg=self._ws_rg,
             workspace_sub=self._ws_sub,
             registry_reference=registry_reference,

sdk/ml/azure-ai-ml/azure/ai/ml/_schema/assets/package/model_package.py

@@ -9,7 +9,7 @@
 from marshmallow import fields, post_load
 
 from azure.ai.ml._schema.core.schema import PathAwareSchema
-from azure.ai.ml._schema.core.fields import VersionField, NestedField
+from azure.ai.ml._schema.core.fields import UnionField, NestedField, StringTransformedEnum
 from .inference_server import InferenceServerSchema
 from .model_configuration import ModelConfigurationSchema
 from .model_package_input import ModelPackageInputSchema
@@ -19,8 +19,12 @@
 
 
 class ModelPackageSchema(PathAwareSchema):
-    target_environment_name = fields.Str(required=True, dump_default="packaged-env")
-    target_environment_version = VersionField()
+    target_environment = UnionField(
+        union_fields=[
+            fields.Dict(keys=StringTransformedEnum(allowed_values=["name"]), values=fields.Str()),
+            fields.Str(required=True),
+        ]
+    )  # pylint: disable=line-too-long
     base_environment_source = NestedField(BaseEnvironmentSourceSchema)
     inferencing_server = NestedField(InferenceServerSchema)
     model_configuration = NestedField(ModelConfigurationSchema)

sdk/ml/azure-ai-ml/azure/ai/ml/entities/_assets/_artifacts/_package/base_environment_source.py

@@ -6,7 +6,7 @@
 
 from typing import Dict
 
-from azure.ai.ml._restclient.v2023_02_01_preview.models import BaseEnvironmentId as RestBaseEnvironmentId
+from azure.ai.ml._restclient.v2023_08_01_preview.models import BaseEnvironmentId as RestBaseEnvironmentId
 from azure.ai.ml._schema.assets.package.base_environment_source import BaseEnvironmentSourceSchema
 from azure.ai.ml._utils._experimental import experimental
 from azure.ai.ml.constants._common import BASE_PATH_CONTEXT_KEY

sdk/ml/azure-ai-ml/azure/ai/ml/entities/_assets/_artifacts/_package/inferencing_server.py

@@ -4,7 +4,7 @@
 
 # pylint: disable=protected-access,redefined-builtin,unused-argument
 
-from azure.ai.ml._restclient.v2023_02_01_preview.models import (
+from azure.ai.ml._restclient.v2023_08_01_preview.models import (
     AzureMLOnlineInferencingServer as RestAzureMLOnlineInferencingServer,
     AzureMLBatchInferencingServer as RestAzureMLBatchInferencingServer,
     CustomInferencingServer as RestCustomInferencingServer,

sdk/ml/azure-ai-ml/azure/ai/ml/entities/_assets/_artifacts/_package/model_package.py

@@ -7,14 +7,14 @@
 from os import PathLike
 from pathlib import Path
 from typing import IO, AnyStr, Dict, List, Optional, Union
-
-from azure.ai.ml._restclient.v2021_10_01_dataplanepreview.models import PackageRequest as DataPlanePackageRequest
-from azure.ai.ml._restclient.v2023_02_01_preview.models import CodeConfiguration
-from azure.ai.ml._restclient.v2023_02_01_preview.models import ModelPackageInput as RestModelPackageInput
-from azure.ai.ml._restclient.v2023_02_01_preview.models import PackageInputPathId as RestPackageInputPathId
-from azure.ai.ml._restclient.v2023_02_01_preview.models import PackageInputPathUrl as RestPackageInputPathUrl
-from azure.ai.ml._restclient.v2023_02_01_preview.models import PackageInputPathVersion as RestPackageInputPathVersion
-from azure.ai.ml._restclient.v2023_02_01_preview.models import PackageRequest, PackageResponse
+import re
+
+from azure.ai.ml._restclient.v2023_08_01_preview.models import CodeConfiguration
+from azure.ai.ml._restclient.v2023_08_01_preview.models import ModelPackageInput as RestModelPackageInput
+from azure.ai.ml._restclient.v2023_08_01_preview.models import PackageInputPathId as RestPackageInputPathId
+from azure.ai.ml._restclient.v2023_08_01_preview.models import PackageInputPathUrl as RestPackageInputPathUrl
+from azure.ai.ml._restclient.v2023_08_01_preview.models import PackageInputPathVersion as RestPackageInputPathVersion
+from azure.ai.ml._restclient.v2023_08_01_preview.models import PackageRequest, PackageResponse
 from azure.ai.ml._schema.assets.package.model_package import ModelPackageSchema
 from azure.ai.ml._utils._experimental import experimental
 from azure.ai.ml._utils.utils import dump_yaml_to_file, snake_to_pascal
@@ -221,30 +221,38 @@ class ModelPackage(Resource, PackageRequest):
     def __init__(
         self,
         *,
-        target_environment_name: str = None,
-        target_environment_id: str = None,
+        target_environment: Union[str, Dict[str, str]],
         inferencing_server: Union[AzureMLOnlineInferencingServer, AzureMLBatchInferencingServer],
         base_environment_source: BaseEnvironment = None,
-        target_environment_version: Optional[str] = None,
         environment_variables: Optional[Dict[str, str]] = None,
         inputs: Optional[List[ModelPackageInput]] = None,
         model_configuration: Optional[ModelConfiguration] = None,
         tags: Optional[Dict[str, str]] = None,
+        **kwargs,
     ):
-        if target_environment_name is None:
-            target_environment_name = "dummy_name"
+        if isinstance(target_environment, dict):
+            target_environment = target_environment["name"]
+            env_version = None
+        else:
+            parse_id = re.match(r"azureml:(\w+):(\d+)$", target_environment)
+
+            if parse_id:
+                target_environment = parse_id.group(1)
+                env_version = parse_id.group(2)
+            else:
+                env_version = None
+
         super().__init__(
-            name=target_environment_name,
-            target_environment_name=target_environment_name,
-            target_environment_version=target_environment_version,
+            name=target_environment,
+            target_environment_id=target_environment,
             base_environment_source=base_environment_source,
             inferencing_server=inferencing_server,
             model_configuration=model_configuration,
             inputs=inputs,
             tags=tags,
             environment_variables=environment_variables,
         )
-        self.target_environment_id = target_environment_id
+        self.environment_version = env_version
 
     @classmethod
     def _load(
@@ -307,29 +315,16 @@ def _to_rest_object(self) -> PackageRequest:
             )
             self.inferencing_server.code_configuration = code
 
-        if self.target_environment_id:
-            package_request = DataPlanePackageRequest(
-                target_environment_id=self.target_environment_id,
-                base_environment_source=self.base_environment_source._to_rest_object()
-                if self.base_environment_source
-                else None,
-                inferencing_server=self.inferencing_server._to_rest_object() if self.inferencing_server else None,
-                model_configuration=self.model_configuration._to_rest_object() if self.model_configuration else None,
-                inputs=[input._to_rest_object() for input in self.inputs] if self.inputs else None,
-                tags=self.tags,
-                environment_variables=self.environment_variables,
-            )
-        else:
-            package_request = PackageRequest(
-                target_environment_name=self.name,
-                base_environment_source=self.base_environment_source._to_rest_object()
-                if self.base_environment_source
-                else None,
-                inferencing_server=self.inferencing_server._to_rest_object() if self.inferencing_server else None,
-                model_configuration=self.model_configuration._to_rest_object() if self.model_configuration else None,
-                inputs=[input._to_rest_object() for input in self.inputs] if self.inputs else None,
-                tags=self.tags,
-                environment_variables=self.environment_variables,
-            )
+        package_request = PackageRequest(
+            target_environment_id=self.target_environment_id,
+            base_environment_source=self.base_environment_source._to_rest_object()
+            if self.base_environment_source
+            else None,
+            inferencing_server=self.inferencing_server._to_rest_object() if self.inferencing_server else None,
+            model_configuration=self.model_configuration._to_rest_object() if self.model_configuration else None,
+            inputs=[input._to_rest_object() for input in self.inputs] if self.inputs else None,
+            tags=self.tags,
+            environment_variables=self.environment_variables,
+        )
 
         return package_request

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

@@ -24,16 +24,19 @@
 from azure.ai.ml._restclient.v2021_10_01_dataplanepreview import (
     AzureMachineLearningWorkspaces as ServiceClient102021Dataplane,
 )
-from azure.ai.ml._restclient.v2023_04_01_preview import AzureMachineLearningWorkspaces as ServiceClient042023Preview
-from azure.ai.ml._restclient.v2023_04_01_preview.models import ListViewType, ModelVersion
+from azure.ai.ml._restclient.v2023_08_01_preview import AzureMachineLearningServices as ServiceClient082023Preview
+from azure.ai.ml._restclient.v2023_08_01_preview.models import ListViewType, ModelVersion
 from azure.ai.ml._scope_dependent_operations import (
     OperationConfig,
     OperationsContainer,
     OperationScope,
     _ScopeDependentOperations,
 )
 from azure.ai.ml._telemetry import ActivityType, monitor_with_activity
-from azure.ai.ml._utils._arm_id_utils import is_ARM_id_for_resource
+from azure.ai.ml._utils._arm_id_utils import (
+    is_ARM_id_for_resource,
+    AMLVersionedArmId,
+)
 from azure.ai.ml._utils._asset_utils import (
     _archive_or_restore,
     _get_latest,
@@ -82,7 +85,7 @@ class ModelOperations(_ScopeDependentOperations):
     :param operation_config: Common configuration for operations classes of an MLClient object.
     :type operation_config: ~azure.ai.ml._scope_dependent_operations.OperationConfig
     :param service_client: Service client to allow end users to operate on Azure Machine Learning Workspace
-        resources (ServiceClient042023Preview or ServiceClient102021Dataplane).
+        resources (ServiceClient082023Preview or ServiceClient102021Dataplane).
     :type service_client: typing.Union[
         ~azure.ai.ml._restclient.v2023_04_01_preview._azure_machine_learning_workspaces.AzureMachineLearningWorkspaces,
         ~azure.ai.ml._restclient.v2021_10_01_dataplanepreview._azure_machine_learning_workspaces.
@@ -98,7 +101,7 @@ def __init__(
         self,
         operation_scope: OperationScope,
         operation_config: OperationConfig,
-        service_client: Union[ServiceClient042023Preview, ServiceClient102021Dataplane],
+        service_client: Union[ServiceClient082023Preview, ServiceClient102021Dataplane],
         datastore_operations: DatastoreOperations,
         all_operations: OperationsContainer = None,
         **kwargs: Dict,
@@ -721,11 +724,22 @@ def package(self, name: str, version: str, package_request: ModelPackage, **kwar
                     else package_request.base_environment_source.resource_id
                 )
 
-            if self._registry_name:
-                # create ARM id for the target environment
-                if package_request.target_environment_name:
-                    package_request.target_environment_id = f"azureml://locations/{self._operation_scope._workspace_location}/workspaces/{self._operation_scope._workspace_id}/environments/{package_request.target_environment_name}"
+            # create ARM id for the target environment
+            if self._operation_scope._workspace_location and self._operation_scope._workspace_id:
+                package_request.target_environment_id = f"azureml://locations/{self._operation_scope._workspace_location}/workspaces/{self._operation_scope._workspace_id}/environments/{package_request.target_environment_id}"
+            else:
+                ws = self._all_operations.all_operations.get("workspaces")
+                ws_details = ws.get(self._workspace_name)
+                workspace_location, workspace_id = (
+                    ws_details.location,
+                    ws_details._workspace_id,
+                )
+                package_request.target_environment_id = f"azureml://locations/{workspace_location}/workspaces/{workspace_id}/environments/{package_request.target_environment_id}"
 
+            if package_request.environment_version is not None:
+                package_request.target_environment_id = (
+                    package_request.target_environment_id + f"/versions/{package_request.environment_version}"
+                )
             package_request = package_request._to_rest_object()
 
         if self._registry_reference:
@@ -749,15 +763,14 @@ def package(self, name: str, version: str, package_request: ModelPackage, **kwar
         )
         if is_deployment_flow:  # No need to go through the schema, as this is for deployment notification only
             return package_out
-        if hasattr(package_out, "target_environment_name"):
-            environment_name = package_out.target_environment_name
+        if hasattr(package_out, "target_environment_id"):
+            environment_id = package_out.target_environment_id
         else:
-            environment_name = package_out.additional_properties["targetEnvironmentName"]
+            environment_id = package_out.additional_properties["targetEnvironmentId"]
 
-        if hasattr(package_out, "target_environment_version"):
-            environment_version = package_out.target_environment_version
-        else:
-            environment_version = package_out.additional_properties["targetEnvironmentVersion"]
+        parsed_id = AMLVersionedArmId(environment_id)
+        environment_name = parsed_id.asset_name
+        environment_version = parsed_id.asset_version
 
         module_logger.info("\nPackage Created")
         if package_out is not None and package_out.__class__.__name__ == "PackageResponse":

sdk/ml/azure-ai-ml/tests/model_package/e2etests/test_model_package_operation.py

@@ -0,0 +1,25 @@
+from typing import Callable, Union
+from devtools_testutils import AzureRecordedTestCase, is_live
+import pytest
+
+from azure.ai.ml import MLClient
+
+from azure.ai.ml.entities import (
+    ModelPackage,
+    ModelConfiguration,
+    AzureMLOnlineInferencingServer,
+)
+
+
+@pytest.mark.timeout(600)
+@pytest.mark.usefixtures("recorded_test")
+@pytest.mark.production_experiences_test
+class TestModelPackage(AzureRecordedTestCase):
+    def test_model_package_workspace(self, client: MLClient):
+        package_config = ModelPackage(
+            target_environment="my-package-name",
+            inferencing_server=AzureMLOnlineInferencingServer(),
+            model_configuration=ModelConfiguration(mode="copy"),
+        )
+
+        client.models.package("test-model2", "1", package_config)

sdk/ml/azure-ai-ml/tests/test_configs/model_package/model_package_config_copy.yml

@@ -1,4 +1,5 @@
-target_environment_name: diabetes-online-mlflow
+target_environment: 
+    name: diabetes-online-mlflow
 inferencing_server: 
     type: azureml_online
 model_configuration:

sdk/ml/azure-ai-ml/tests/test_configs/model_package/model_package_simple.yml

@@ -1,4 +1,4 @@
-target_environment_name: diabetes-online-mlflow
+target_environment: diabetes-online-mlflow
 inferencing_server: 
     type: azureml_online
 model_configuration: