Don't break SDK public API for WorkspaceModelReference (Azure#28839)

* don't break public API * added test, removed ml_client docs * added test, removed ml_client docs
msyyc · Feb 21, 2023 · e1946f7 · e1946f7
1 parent 61c3e35
commit e1946f7
Show file tree

Hide file tree

Showing 4 changed files with 12 additions and 4 deletions.
diff --git a/sdk/ml/azure-ai-ml/azure/ai/ml/_ml_client.py b/sdk/ml/azure-ai-ml/azure/ai/ml/_ml_client.py
@@ -769,7 +769,7 @@ def _get_workspace_info(cls, found_path: Optional[str]) -> Tuple[str, str, str]:
 
     # T = valid inputs/outputs for create_or_update
     # Each entry here requires a registered _create_or_update function below
-    T = TypeVar("T", Job, Model, Environment, Component, Datastore, WorkspaceAssetReference)
+    T = TypeVar("T", Job, Model, Environment, Component, Datastore)
 
     def create_or_update(
         self,
@@ -781,7 +781,7 @@ def create_or_update(
         :param entity: The resource to create or update.
         :type entity: typing.Union[~azure.ai.ml.entities.Job
             , ~azure.ai.ml.entities.Model, ~azure.ai.ml.entities.Environment, ~azure.ai.ml.entities.Component
-            , ~azure.ai.ml.entities.Datastore, ~azure.ai.ml.entities.WorkspaceAssetReference]
+            , ~azure.ai.ml.entities.Datastore]
         :return: The created or updated resource.
         :rtype: typing.Union[~azure.ai.ml.entities.Job, ~azure.ai.ml.entities.Model
             , ~azure.ai.ml.entities.Environment, ~azure.ai.ml.entities.Component, ~azure.ai.ml.entities.Datastore]

diff --git a/sdk/ml/azure-ai-ml/azure/ai/ml/entities/__init__.py b/sdk/ml/azure-ai-ml/azure/ai/ml/entities/__init__.py
@@ -6,6 +6,7 @@
 Main areas include managing compute targets, creating/managing workspaces and jobs, and submitting/accessing
 model, runs and run output/logging etc.
 """
+# pylint: disable=naming-mismatch
 __path__ = __import__("pkgutil").extend_path(__path__, __name__)
 
 from azure.ai.ml._restclient.v2022_10_01.models import CreatedByType
@@ -15,7 +16,7 @@
 from ._assets._artifacts.model import Model
 from ._assets.asset import Asset
 from ._assets.environment import BuildContext, Environment
-from ._assets.workspace_asset_reference import WorkspaceAssetReference
+from ._assets.workspace_asset_reference import WorkspaceAssetReference as WorkspaceModelReference
 from ._builders import Command, Parallel, Pipeline, Spark, Sweep
 from ._component.command_component import CommandComponent
 from ._component.component import Component
@@ -254,7 +255,7 @@
     "SynapseSparkCompute",
     "AutoScaleSettings",
     "AutoPauseSettings",
-    "WorkspaceAssetReference",
+    "WorkspaceModelReference",
     # builders
     "Command",
     "Parallel",

diff --git a/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_assets/workspace_asset_reference.py b/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_assets/workspace_asset_reference.py
@@ -19,6 +19,7 @@
 class WorkspaceAssetReference(Asset):
     """Workspace Model Reference.
 
+    This is for SDK internal use only, might be deprecated in the future.
     :param name: Model name
     :type name: str
     :param version: Model version

diff --git a/sdk/ml/azure-ai-ml/tests/model/unittests/test_model_operations.py b/sdk/ml/azure-ai-ml/tests/model/unittests/test_model_operations.py
@@ -307,3 +307,9 @@ def test_promote_model_from_workspace(
             )
             mock_model_operation_reg.create_or_update(model_to_promote)
             mock_model_operation_reg._service_client.resource_management_asset_reference.begin_import_method.assert_called_once()
+
+    def test_model_entity_class_exist(self):
+        try:
+            from azure.ai.ml.entities import WorkspaceModelReference
+        except ImportError:
+            assert False, "WorkspaceModelReference class not found"