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: