Repository: mlflow/mlflow

Issue:
[FR] Allow `stage` as version value in `set_model_version_tag`
Thank you for submitting a feature request. **Before proceeding, please review MLflow's [Issue Policy for feature requests](https://www.github.com/mlflow/mlflow/blob/master/ISSUE_POLICY.md#feature-requests) and the [MLflow Contributing Guide](https://github.com/mlflow/mlflow/blob/master/CONTRIBUTING.rst)**.

**Please fill in this feature request template to ensure a timely and thorough response.**

Currently, the method `client.set_model_version_tag` only accepts an integer value for the parameter `version`. However, downstream jobs that refer to a model in the model registry often refer to it by `stage` instead of `version`. It would simplify things if the `set_model_version_tag` method would accept the stage name in addition to the version as well.

Example:
```python
client.set_model_version_tag(name="my-model", version="staging", key="mykey", value="myvalue")
```

Currently, if someone now wants to set a tag for a `production` model, the person would first have to look up the `model version` for that model before it can set the tag. 

## Willingness to contribute
The MLflow Community encourages new feature contributions. Would you or another member of your organization be willing to contribute an implementation of this feature (either as an MLflow Plugin or an enhancement to the MLflow code base)?

- [ ] Yes. I can contribute this feature independently.
- [x] Yes. I would be willing to contribute this feature with guidance from the MLflow community.
- [ ] No. I cannot contribute this feature at this time.

## Proposal Summary

Allow `client.set_model_version_tag` to accept `staging,production,archived` as values for the `version` parameter as well, or add an additional parameter called `stage` that can be used instead of `version`.

## Motivation
- What is the use case for this feature? _It makes MLflow simpler to use_
- Why is this use case valuable to support for MLflow users in general?
- Why is this use case valuable to support for your project(s) or organization?
- Why is it currently difficult to achieve this use case? (please be as specific as possible about why related MLflow features and components are insufficient)

### What component(s), interfaces, languages, and integrations does this feature affect?
Components 
- [ ] `area/artifacts`: Artifact stores and artifact logging
- [ ] `area/build`: Build and test infrastructure for MLflow
- [ ] `area/docs`: MLflow documentation pages
- [ ] `area/examples`: Example code
- [x] `area/model-registry`: Model Registry service, APIs, and the fluent client calls for Model Registry
- [ ] `area/models`: MLmodel format, model serialization/deserialization, flavors
- [ ] `area/projects`: MLproject format, project running backends
- [ ] `area/scoring`: MLflow Model server, model deployment tools, Spark UDFs
- [ ] `area/server-infra`: MLflow Tracking server backend
- [ ] `area/tracking`: Tracking Service, tracking client APIs, autologging

Interfaces
- [ ] `area/uiux`: Front-end, user experience, plotting, JavaScript, JavaScript dev server
- [ ] `area/docker`: Docker use across MLflow's components, such as MLflow Projects and MLflow Models
- [ ] `area/sqlalchemy`: Use of SQLAlchemy in the Tracking Service or Model Registry
- [ ] `area/windows`: Windows support

Languages 
- [ ] `language/r`: R APIs and clients
- [ ] `language/java`: Java APIs and clients
- [ ] `language/new`: Proposals for new client languages

Integrations
- [ ] `integrations/azure`: Azure and Azure ML integrations
- [ ] `integrations/sagemaker`: SageMaker integrations
- [ ] `integrations/databricks`: Databricks integrations

## Details

(Use this section to include any additional information about the feature. If you have a proposal for how to implement this feature, please include it here. For implementation guidelines, please refer to the [Contributing Guide](https://github.com/mlflow/mlflow/blob/master/CONTRIBUTING.rst#contribution-guidelines).)



Discussion:
For any questions, concerns, or clarification on implementing this issue, please ping @BenWilson2.

Resolving patch:
diff --git a/mlflow/entities/model_registry/model_version_stages.py b/mlflow/entities/model_registry/model_version_stages.py
--- a/mlflow/entities/model_registry/model_version_stages.py
+++ b/mlflow/entities/model_registry/model_version_stages.py
@@ -17,6 +17,9 @@ def get_canonical_stage(stage):
     key = stage.lower()
     if key not in _CANONICAL_MAPPING:
         raise MlflowException(
-            "Invalid Model Version stage {}.".format(stage), INVALID_PARAMETER_VALUE
+            "Invalid Model Version stage: {}. Value must be one of {}.".format(
+                stage, ", ".join(ALL_STAGES)
+            ),
+            INVALID_PARAMETER_VALUE,
         )
     return _CANONICAL_MAPPING[key]
diff --git a/mlflow/tracking/client.py b/mlflow/tracking/client.py
--- a/mlflow/tracking/client.py
+++ b/mlflow/tracking/client.py
@@ -37,6 +37,7 @@
 )
 from mlflow.utils.logging_utils import eprint
 from mlflow.utils.uri import is_databricks_uri, construct_run_url
+from mlflow.utils.validation import _validate_model_version_or_stage_exists
 
 if TYPE_CHECKING:
     import matplotlib  # pylint: disable=unused-import
@@ -2681,14 +2682,19 @@ def get_model_version_stages(
         """
         return ALL_STAGES
 
-    def set_model_version_tag(self, name: str, version: str, key: str, value: Any) -> None:
+    def set_model_version_tag(
+        self, name: str, version: str = None, key: str = None, value: Any = None, stage: str = None
+    ) -> None:
         """
         Set a tag for the model version.
+        When stage is set, tag will be set for latest model version of the stage.
+        Setting both version and stage parameter will result in error.
 
         :param name: Registered model name.
         :param version: Registered model version.
-        :param key: Tag key to log.
-        :param value: Tag value to log.
+        :param key: Tag key to log. key is required.
+        :param value: Tag value to log. value is required.
+        :param stage: Registered model stage.
         :return: None
 
         .. code-block:: python
@@ -2723,7 +2729,13 @@ def print_model_version_info(mv):
             mv = client.create_model_version(name, model_uri, run.info.run_id)
             print_model_version_info(mv)
             print("--")
+
+            # Tag using model version
             client.set_model_version_tag(name, mv.version, "t", "1")
+
+            # Tag using model stage
+            client.set_model_version_tag(name, key="t1", value="1", stage=mv.current_stage)
+
             mv = client.get_model_version(name, mv.version)
             print_model_version_info(mv)
 
@@ -2736,17 +2748,33 @@ def print_model_version_info(mv):
             --
             Name: RandomForestRegression
             Version: 1
-            Tags: {'t': '1'}
+            Tags: {'t': '1', 't1': '1'}
         """
+        _validate_model_version_or_stage_exists(version, stage)
+        if stage:
+            latest_versions = self.get_latest_versions(name, stages=[stage])
+            if not latest_versions:
+                raise MlflowException(
+                    "Could not find any model version for {} stage".format(
+                        stage,
+                    )
+                )
+            version = latest_versions[0].version
+
         self._get_registry_client().set_model_version_tag(name, version, key, value)
 
-    def delete_model_version_tag(self, name: str, version: str, key: str) -> None:
+    def delete_model_version_tag(
+        self, name: str, version: str = None, key: str = None, stage: str = None
+    ) -> None:
         """
         Delete a tag associated with the model version.
+        When stage is set, tag will be deleted for latest model version of the stage.
+        Setting both version and stage parameter will result in error.
 
         :param name: Registered model name.
         :param version: Registered model version.
-        :param key: Tag key.
+        :param key: Tag key. key is required.
+        :param stage: Registered model stage.
         :return: None
 
         .. code-block:: python
@@ -2778,11 +2806,15 @@ def print_model_version_info(mv):
             # Create a new version of the rfr model under the registered model name
             # and delete a tag
             model_uri = "runs:/{}/sklearn-model".format(run.info.run_id)
-            tags = {'t': "t1"}
+            tags = {'t': "1", "t1" : "2"}
             mv = client.create_model_version(name, model_uri, run.info.run_id, tags=tags)
             print_model_version_info(mv)
             print("--")
+            #using version to delete tag
             client.delete_model_version_tag(name, mv.version, "t")
+
+            #using stage to delete tag
+            client.delete_model_version_tag(name, key="t1", stage=mv.current_stage)
             mv = client.get_model_version(name, mv.version)
             print_model_version_info(mv)
 
@@ -2791,10 +2823,20 @@ def print_model_version_info(mv):
 
             Name: RandomForestRegression
             Version: 1
-            Tags: {'t': 't1'}
+            Tags: {'t': '1', 't1': '2'}
             --
             Name: RandomForestRegression
             Version: 1
             Tags: {}
         """
+        _validate_model_version_or_stage_exists(version, stage)
+        if stage:
+            latest_versions = self.get_latest_versions(name, stages=[stage])
+            if not latest_versions:
+                raise MlflowException(
+                    "Could not find any model version for {} stage".format(
+                        stage,
+                    )
+                )
+            version = latest_versions[0].version
         self._get_registry_client().delete_model_version_tag(name, version, key)
diff --git a/mlflow/utils/validation.py b/mlflow/utils/validation.py
--- a/mlflow/utils/validation.py
+++ b/mlflow/utils/validation.py
@@ -180,6 +180,7 @@ def _validate_model_version_tag(key, value):
     Check that a tag with the specified key & value is valid and raise an exception if it isn't.
     """
     _validate_tag_name(key)
+    _validate_tag_value(value)
     _validate_length_limit("Model version key", MAX_MODEL_REGISTRY_TAG_KEY_LENGTH, key)
     _validate_length_limit("Model version value", MAX_MODEL_REGISTRY_TAG_VALUE_LENGTH, value)
 
@@ -386,3 +387,16 @@ def _validate_db_type_string(db_type):
     if db_type not in DATABASE_ENGINES:
         error_msg = "Invalid database engine: '%s'. '%s'" % (db_type, _UNSUPPORTED_DB_TYPE_MSG)
         raise MlflowException(error_msg, INVALID_PARAMETER_VALUE)
+
+
+def _validate_model_version_or_stage_exists(version, stage):
+    if version and stage:
+        raise MlflowException("version and stage cannot be set together", INVALID_PARAMETER_VALUE)
+
+    if not (version or stage):
+        raise MlflowException("version or stage must be set", INVALID_PARAMETER_VALUE)
+
+
+def _validate_tag_value(value):
+    if value == None:
+        raise MlflowException("Tag value cannot be None", INVALID_PARAMETER_VALUE)


Test patch:
diff --git a/tests/store/model_registry/test_sqlalchemy_store.py b/tests/store/model_registry/test_sqlalchemy_store.py
--- a/tests/store/model_registry/test_sqlalchemy_store.py
+++ b/tests/store/model_registry/test_sqlalchemy_store.py
@@ -605,7 +605,9 @@ def test_update_model_version(self):
 
         # only valid stages can be set
         with self.assertRaisesRegex(
-            MlflowException, r"Invalid Model Version stage unknown"
+            MlflowException,
+            "Invalid Model Version stage: unknown. "
+            "Value must be one of None, Staging, Production, Archived.",
         ) as exception_context:
             self.store.transition_model_version_stage(
                 mv1.name, mv1.version, stage="unknown", archive_existing_versions=False
diff --git a/tests/tracking/test_client.py b/tests/tracking/test_client.py
--- a/tests/tracking/test_client.py
+++ b/tests/tracking/test_client.py
@@ -592,3 +592,70 @@ class MockUnpickleableModelRegistryStore(SqlAlchemyModelRegistryStore):
     client.create_experiment("test_experiment")
     client.create_registered_model("test_model")
     pickle.dumps(client)
+
+
+@pytest.fixture
+def mock_registry_store_with_get_latest_version(mock_registry_store):
+    mock_get_latest_versions = mock.Mock()
+    mock_get_latest_versions.return_value = [
+        ModelVersion(
+            "model_name",
+            1,
+            0,
+        )
+    ]
+
+    mock_registry_store.get_latest_versions = mock_get_latest_versions
+    yield mock_registry_store
+
+
+def test_set_model_version_tag(mock_registry_store_with_get_latest_version):
+
+    # set_model_version_tag using version
+    MlflowClient().set_model_version_tag("model_name", 1, "tag1", "foobar")
+    mock_registry_store_with_get_latest_version.set_model_version_tag.assert_called_once_with(
+        "model_name", 1, ModelVersionTag(key="tag1", value="foobar")
+    )
+
+    mock_registry_store_with_get_latest_version.reset_mock()
+
+    # set_model_version_tag using stage
+    MlflowClient().set_model_version_tag("model_name", key="tag1", value="foobar", stage="Staging")
+    mock_registry_store_with_get_latest_version.set_model_version_tag.assert_called_once_with(
+        "model_name", 1, ModelVersionTag(key="tag1", value="foobar")
+    )
+
+    # set_model_version_tag with version and stage set
+    with pytest.raises(MlflowException, match="version and stage cannot be set together"):
+        MlflowClient().set_model_version_tag("model_name", 1, "tag1", "foobar", stage="Staging")
+
+    # set_model_version_tag with version and stage not set
+    with pytest.raises(MlflowException, match="version or stage must be set"):
+        MlflowClient().set_model_version_tag("model_name", key="tag1", value="foobar")
+
+
+def test_delete_model_version_tag(mock_registry_store_with_get_latest_version):
+
+    # delete_model_version_tag using version
+    MlflowClient().delete_model_version_tag("model_name", 1, "tag1")
+    mock_registry_store_with_get_latest_version.delete_model_version_tag.assert_called_once_with(
+        "model_name", 1, "tag1"
+    )
+
+    mock_registry_store_with_get_latest_version.reset_mock()
+
+    # delete_model_version_tag using stage
+    MlflowClient().delete_model_version_tag("model_name", key="tag1", stage="Staging")
+    mock_registry_store_with_get_latest_version.delete_model_version_tag.assert_called_once_with(
+        "model_name", 1, "tag1"
+    )
+
+    # delete_model_version_tag with version and stage set
+    with pytest.raises(MlflowException, match="version and stage cannot be set together"):
+        MlflowClient().delete_model_version_tag(
+            "model_name", version=1, key="tag1", stage="staging"
+        )
+
+    # delete_model_version_tag with version and stage not set
+    with pytest.raises(MlflowException, match="version or stage must be set"):
+        MlflowClient().delete_model_version_tag("model_name", key="tag1")