Support Model Stage in the SDK (#30254)

* Support Model Stage in SDK/CLI

* update

* fix lint

* should fix test_ipp_model

* fix list, test, and add recording

* fix test

* fix lint

* fix it

* run black

* Revert "fix it"

This reverts commit d3c337c.

* Revert "run black"

This reverts commit 8cb4113.

* update

* cleanup

* remove extra spaces

* fixes

* update changelog

Author: namikhai

sdk/ml/azure-ai-ml/CHANGELOG.md

@@ -4,6 +4,7 @@
 
 ### Features Added
 - Added support to enable set workspace connection secret expiry time.
+- Added support for `stage` on model version
 
 ### Bugs Fixed
 

sdk/ml/azure-ai-ml/azure/ai/ml/_schema/assets/model.py

@@ -35,6 +35,7 @@ class ModelSchema(PathAwareSchema):
     description = fields.Str()
     properties = fields.Dict()
     tags = fields.Dict()
+    stage = fields.Str()
     utc_time_created = fields.DateTime(format="iso", dump_only=True)
     flavors = fields.Dict()
     creation_context = NestedField(CreationContextSchema, dump_only=True)

sdk/ml/azure-ai-ml/azure/ai/ml/entities/_assets/_artifacts/model.py

@@ -29,7 +29,7 @@
 from .artifact import ArtifactStorageInfo
 
 
-class Model(Artifact):
+class Model(Artifact):  # pylint: disable=too-many-instance-attributes
     """Model for training and scoring.
 
     :param name: Name of the resource.
@@ -54,6 +54,8 @@ class Model(Artifact):
     :type tags: dict[str, str]
     :param properties: The asset property dictionary.
     :type properties: dict[str, str]
+    :param stage: The stage of the resource.
+    :type stage: str
     :param kwargs: A dictionary of additional configuration parameters.
     :type kwargs: dict
     """
@@ -70,6 +72,7 @@ def __init__(
         description: Optional[str] = None,
         tags: Optional[Dict] = None,
         properties: Optional[Dict] = None,
+        stage: Optional[str] = None,
         **kwargs,
     ):
         self.job_name = kwargs.pop("job_name", None)
@@ -87,6 +90,7 @@ def __init__(
         self.flavors = dict(flavors) if flavors else None
         self._arm_type = ArmConstants.MODEL_VERSION_TYPE
         self.type = type or AssetTypes.CUSTOM_MODEL
+        self.stage = stage
         if self._is_anonymous and self.path:
             _ignore_file = get_ignore_file(self.path)
             _upload_hash = get_object_hash(self.path, _ignore_file)
@@ -115,6 +119,7 @@ def _to_dict(self) -> Dict:
     def _from_rest_object(cls, model_rest_object: ModelVersion) -> "Model":
         rest_model_version: ModelVersionProperties = model_rest_object.properties
         arm_id = AMLVersionedArmId(arm_id=model_rest_object.id)
+        model_stage = rest_model_version.stage if hasattr(rest_model_version, "stage") else None
         if hasattr(rest_model_version, "flavors"):
             flavors = {key: flavor.data for key, flavor in rest_model_version.flavors.items()}
         model = Model(
@@ -126,6 +131,7 @@ def _from_rest_object(cls, model_rest_object: ModelVersion) -> "Model":
             tags=rest_model_version.tags,
             flavors=flavors,
             properties=rest_model_version.properties,
+            stage=model_stage,
             # pylint: disable=protected-access
             creation_context=SystemData._from_rest_object(model_rest_object.system_data),
             type=rest_model_version.model_type,
@@ -162,6 +168,7 @@ def _to_rest_object(self) -> ModelVersion:
             else None,  # flatten OrderedDict to dict
             model_type=self.type,
             model_uri=self.path,
+            stage=self.stage,
             is_anonymous=self._is_anonymous,
         )
         model_version_resource = ModelVersion(properties=model_version)

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

@@ -248,7 +248,6 @@ def _get(self, name: str, version: Optional[str] = None) -> ModelVersion:  # nam
                     name=name,
                     version=version,
                     workspace_name=self._workspace_name,
-                    api_version="2023-02-01-preview",
                     **self._scope_kwargs,
                 )
             )
@@ -418,6 +417,7 @@ def restore(
     def list(
         self,
         name: Optional[str] = None,
+        stage: Optional[str] = None,
         *,
         list_view_type: ListViewType = ListViewType.ACTIVE_ONLY,
     ) -> Iterable[Model]:
@@ -444,6 +444,7 @@ def list(
                     workspace_name=self._workspace_name,
                     cls=lambda objs: [Model._from_rest_object(obj) for obj in objs],
                     list_view_type=list_view_type,
+                    stage=stage,
                     **self._scope_kwargs,
                 )
             )

sdk/ml/azure-ai-ml/tests/model/e2etests/test_model.py

@@ -69,6 +69,35 @@ def test_crud_file(self, client: MLClient, randstr: Callable[[], str], tmp_path:
         # with pytest.raises(Exception):
         #     client.models.get(name=model.name, version="3")
 
+    def test_crud_model_with_stage(self, client: MLClient, randstr: Callable[[], str], tmp_path: Path) -> None:
+        path = Path("./tests/test_configs/model/model_with_stage.yml")
+        model_name = randstr("model_prod_name")
+
+        model = load_model(path)
+        model.name = model_name
+        model = client.models.create_or_update(model)
+        assert model.name == model_name
+        assert model.version == "3"
+        assert model.description == "this is my test model with stage"
+        assert model.type == "mlflow_model"
+        assert model.stage == "Production"
+        assert re.match(LONG_URI_REGEX_FORMAT, model.path)
+
+        with pytest.raises(Exception):
+            with patch("azure.ai.ml._artifacts._artifact_utilities.get_object_hash", return_value="DIFFERENT_HASH"):
+                model = load_model(source=artifact_path)
+                model = client.models.create_or_update(model)
+
+        model = client.models.get(model.name, "3")
+        assert model.name == model_name
+        assert model.version == "3"
+        assert model.description == "this is my test model with stage"
+        assert model.stage == "Production"
+
+        model_list = client.models.list(name=model.name, stage="Production")
+        model_stage_list = [m.stage for m in model_list if m is not None]
+        assert model.stage in model_stage_list
+
     def test_list_no_name(self, client: MLClient) -> None:
         models = client.models.list()
         assert isinstance(models, Iterator)

sdk/ml/azure-ai-ml/tests/model/unittests/test_model_schema.py

@@ -32,6 +32,11 @@ def test_deserialize_no_version(self) -> None:
         assert model._auto_increment_version
         assert model.type == AssetTypes.CUSTOM_MODEL  # assert the default model type
 
+    def test_deserialize_with_stage(self) -> None:
+        path = Path("./tests/test_configs/model/model_with_stage.yml")
+        model = load_model(path)
+        assert model.stage == "Production"
+
     def test_ipp_model(self) -> None:
         rest_ipp_model = {
             "id": "azureml://registries/fake_registry/models/fake_ipp_model/versions/611575",