Repository: mlflow/mlflow
Pull request: Add update_run() Python API
<!-- 🚨 We recommend pull requests be filed from a non-master branch on a repository fork (e.g. <username>:fix-xxx). 🚨 -->

## Related Issues/PRs

<!--
Please reference any related feature requests, issues, or PRs here. For example, `#123`. To automatically close the referenced items when this PR is merged, please use a closing keyword (close, fix, or resolve). For example, `Close #123`. See https://docs.github.com/en/issues/tracking-your-work-with-issues/linking-a-pull-request-to-an-issue for more information.
-->
resolves #6932 

## What changes are proposed in this pull request?

Add a function `update_run(self, run_id, status, name)` in `mlflow.client` (as @dbczumar suggested) that allows users to programmatically update run status and run name.

I also added the same function in the fluent api `mlflow`. Please advise if this is not desirable.

## How is this patch tested?

<!--
If you're unsure about what to test, where to add tests, or how to run tests, please feel free to ask. We'd be happy to help.
-->
- [X] I have written tests (not required for typo or doc fix) and confirmed the proposed feature/bug-fix/change works.
```shell
pytest tests/tracking/test_client.py::test_update_run
```

<!--
Please describe how you confirmed the proposed feature/bug-fix/change works here. For example, if you fixed an MLflow client API, you could attach the code that didn't work prior to the fix but works now, or if you added a new feature on MLflow UI, you could attach a video that demonstrates the feature.
-->

## Does this PR change the documentation?

- [ ] No. You can skip the rest of this section.
- [X] Yes. Make sure the changed pages / sections render correctly by following the steps below.

1. Click the `Details` link on the `Preview docs` check.
2. Find the changed pages / sections and make sure they render correctly.

## Release Notes

### Is this a user-facing change?

- [ ] No. You can skip the rest of this section.
- [X] Yes. Give a description of this change to be included in the release notes for MLflow users.

Allows users to programmatically update a run's status and name.

(Details in 1-2 sentences. You can just refer to another PR with a description if this PR is part of a larger change.)

### What component(s), interfaces, languages, and integrations does this PR 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
- [ ] `area/model-registry`: Model Registry service, APIs, and the fluent client calls for Model Registry
- [ ] `area/models`: MLmodel format, model serialization/deserialization, flavors
- [ ] `area/pipelines`: Pipelines, Pipeline APIs, Pipeline configs, Pipeline Templates
- [ ] `area/projects`: MLproject format, project running backends
- [ ] `area/scoring`: MLflow Model server, model deployment tools, Spark UDFs
- [ ] `area/server-infra`: MLflow Tracking server backend
- [X] `area/tracking`: Tracking Service, tracking client APIs, autologging

Interface
- [ ] `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

Language
- [ ] `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

<!--
Insert an empty named anchor here to allow jumping to this section with a fragment URL
(e.g. https://github.com/mlflow/mlflow/pull/123#user-content-release-note-category).
Note that GitHub prefixes anchor names in markdown with "user-content-".
-->
<a name="release-note-category"></a>
### How should the PR be classified in the release notes? Choose one:

- [ ] `rn/breaking-change` - The PR will be mentioned in the "Breaking Changes" section
- [ ] `rn/none` - No description will be included. The PR will be mentioned only by the PR number in the "Small Bugfixes and Documentation Updates" section
- [X] `rn/feature` - A new user-facing feature worth mentioning in the release notes
- [ ] `rn/bug-fix` - A user-facing bug fix worth mentioning in the release notes
- [ ] `rn/documentation` - A user-facing documentation change worth mentioning in the release notes


Problem statement:
[FR] rename_run in API
### Willingness to contribute

Yes. I would be willing to contribute this feature with guidance from the MLflow community.

### Proposal Summary

I need to rename many runs (~100 of them). I can do via the UI, but it would take too long.

Since the mlflow command doesn't provide this function, I planned to create my own basic "rename" python command.

BUT: **the tracking API does not offer a rename_run** .

I would enjoy in the tracking API a method like:

`rename_run(experiment_id, run_id, new_name)
`

Then it could be convenient to extend the mlflow command line to provide access to _rename_experiment_ and _rename _run_.

JL

### Motivation

> #### What is the use case for this feature?
renaming bunch of runs programmatically
> #### Why is this use case valuable to support for MLflow users in general?
An administrative feature useful from time to time.
> #### Why is this use case valuable to support for your project(s) or organization?
Saving time
> #### Why is it currently difficult to achieve this use case?
Need to do it for each run individually via the UI

### Details

Maybe in the API a method like:

`rename_run(experiment_id, run_id, new_name)
`

### What component(s) does this bug affect?

- [ ] `area/artifacts`: Artifact stores and artifact logging
- [ ] `area/build`: Build and test infrastructure for MLflow
- [ ] `area/docs`: MLflow documentation pages
- [ ] `area/examples`: Example code
- [ ] `area/model-registry`: Model Registry service, APIs, and the fluent client calls for Model Registry
- [ ] `area/models`: MLmodel format, model serialization/deserialization, flavors
- [ ] `area/pipelines`: Pipelines, Pipeline APIs, Pipeline configs, Pipeline Templates
- [ ] `area/projects`: MLproject format, project running backends
- [ ] `area/scoring`: MLflow Model server, model deployment tools, Spark UDFs
- [ ] `area/server-infra`: MLflow Tracking server backend
- [X] `area/tracking`: Tracking Service, tracking client APIs, autologging

### What interface(s) does this bug affect?

- [ ] `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

### What language(s) does this bug affect?

- [ ] `language/r`: R APIs and clients
- [ ] `language/java`: Java APIs and clients
- [ ] `language/new`: Proposals for new client languages

### What integration(s) does this bug affect?

- [ ] `integrations/azure`: Azure and Azure ML integrations
- [ ] `integrations/sagemaker`: SageMaker integrations
- [ ] `integrations/databricks`: Databricks integrations

Hints from issue discussion:
The run name is just a tag under run and an attribute in `run.info`. You can write a simple function like this:
```python
import mlflow
from mlflow.tracking import MlflowClient
from mlflow.entities.run_status import RunStatus

def set_run_name(run_id, run_name):
    run = mlflow.get_run(run_id)
    client = MlflowClient()
    client.set_tag(run_id, "mlflow.runName", run_name)
    client._tracking_client.store.update_run_info(
        run_id=run_id,
        run_status=RunStatus.from_string(run.info.status),
        end_time=run.info.end_time,
        run_name=run_name,
    )
```

Anyway, I don't know if this functionality is planned by the dev team. I have drafted an implementation as below, tested with FileStore and Postgres Store. If the dev team deems it desirable, I am happy to finish up the docstring & test cases and create a PR.
https://github.com/mlflow/mlflow/compare/master...gabrielfu:mlflow:feat/rename_run
@gabrielfu  Thanks!

I'm fine with this solution (tested with a server on top of sqlite)

BUT I got an exception with the `run_name` parameter to update_run_info  (I'm using MLFlow 1.29)
`TypeError: RestStore.update_run_info() got an unexpected keyword argument 'run_name'`

So the rename function is as below and works fine.

```
import mlflow
from mlflow.tracking import MlflowClient
from mlflow.entities.run_status import RunStatus

client = MlflowClient(sServerURI)

def set_run_name(client, run_id, run_name):
    run = client.get_run(run_id)
    client.set_tag(run_id, "mlflow.runName", run_name)
    client._tracking_client.store.update_run_info(
        run_id=run_id,
        run_status=RunStatus.from_string(run.info.status),
        end_time=run.info.end_time,
        #run_name=run_name,
    )
```
So here is the command I needed. Happy to share this small piece of Python code.
JL

[do_mlflow_run_rename.py.zip](https://github.com/mlflow/mlflow/files/9684335/do_mlflow_run_rename.py.zip)

```
"""
Utility for mass-renaming of runs in an experiment.

Jean-Luc Meunier
"""

import sys, os
import re
import logging
import argparse

import mlflow
from mlflow.tracking import MlflowClient
from mlflow.entities.run_status import RunStatus
from mlflow.entities import ViewType


logging.basicConfig() # otherwise cannot change logging level... crazy
logger = logging.getLogger('.')
logger.setLevel(logging.INFO)


def rename_run(client, run_id, run_name):
    # thanks to gabrielfu 
    run = client.get_run(run_id)
    client.set_tag(run_id, "mlflow.runName", run_name)
    client._tracking_client.store.update_run_info(
        run_id=run_id,
        run_status=RunStatus.from_string(run.info.status),
        end_time=run.info.end_time,
        #run_name=run_name,
    )
    return "done"


if __name__ == "__main__":
    parser = argparse.ArgumentParser(usage="""
Rename one or several runs from an experiment, given by its experiment id EXP_ID.
The runs to be renamed are selected thanks to a regular expression (the RUN_SELECTOR).
For renaming, the given TERM is replaced by the given REPLACEMENT.

Example:
python do_mlflow_run_rename.py 6 'toto.*foo' foo bar

""")
    parser.add_argument('--uri'
                        , default=None
                        , help='MLFlow server URI (default: %(default)s) By default, read the MLFLOW_TRACKING_URI environment variable')
    # How to select the experiment(s) and the run(s)
    parser.add_argument('exp_id', metavar="EXP_ID"
                        , help='the experience id')
    parser.add_argument('run_name_selector', metavar="RUN_SELECTOR"
                        , help='Regular expression to select the run(s). Must match the full run name. ($ is added automatically at the end)')
    parser.add_argument('term', metavar="TERM_TO_REPLACE"
                        , help='The substring to be replaced')
    parser.add_argument('replacement', metavar="REPLACEMENT"
                        , help='The replacement')
    parser.add_argument('--debug', '-n', action="store_true"
                        , help='Dry run. Only prints what would be done')
    parser.add_argument('--verbose', '-v', action="store_true"
                        , help='Verbose mode')

    args = parser.parse_args()

    if args.debug: logger.info("----- DRY RUN MODE -----")

    sServerURI = args.uri if args.uri else os.getenv("MLFLOW_TRACKING_URI")
    client = MlflowClient(sServerURI)

    lRun = client.search_runs([args.exp_id]
                            , run_view_type=ViewType.ACTIVE_ONLY     # DELETED_ONLY ALL
                            , max_results=1000)
    logger.info(f"{len(lRun)} runs in experiment {args.exp_id}")

    # filter by run name
    run_name_selector = args.run_name_selector
    if run_name_selector[-1] != '$': run_name_selector += '$' 
    cre = re.compile(run_name_selector)

    n = 0
    for run in lRun:
        name = run.data.tags['mlflow.runName']
        newname = name.replace(args.term, args.replacement)
        if cre.match(name):
            if name != newname:
                run_id = run.info.run_id
                if args.debug:
                    logger.info(f"WOULD RENAME: '{name}' --> '{newname}'")
                else:
                    n += 1
                    logger.info(f"RENAMING: '{name}' --> '{newname}'   {rename_run(client, run_id, newname)}")
            elif args.verbose or args.debug:
                logger.info(f"No change, ignored: {name}  {newname}")
        else:
            if args.verbose or args.debug:
                logger.info(f"Ignored: {name}")
    logger.info(f"{n} renaming done")
```


Hi @jlmeunier, we'd be very excited about introducing an `MlflowClient.update_run()` API which can be used to set the status of a run (e.g. running, finished, failed) and / or the run's name, mirroring the UpdateRun API definition which was recently updated on the `master` branch: https://output.circle-artifacts.com/output/job/dcaa1233-23f2-4081-b4fe-22993114ddf8/artifacts/0/docs/build/html/rest-api.html#update-run.

I think the following API signature would be great:

```
def update_run(self, run_id, status=None, name=None):
    """
    :param run_id: The ID of the Run to update.
    :param status: The new status of the run to set, if specified. At least one of ``status`` or ``name`` should be specified.
    :param name: The new name of the run to set, if specified. At least one of ``name`` or ``status`` should be specified.
    """
```

We can implement this by calling through to `self.store.update_run_info()`, as is done here: https://github.com/mlflow/mlflow/blob/df327134d1278f98b6142ddaef77295a146aba18/mlflow/tracking/_tracking_service/client.py#L470-L475

The API should be introduced to https://github.com/mlflow/mlflow/blob/master/mlflow/tracking/client.py and https://github.com/mlflow/mlflow/blob/master/mlflow/tracking/_tracking_service/client.py, and it can be tested in https://github.com/mlflow/mlflow/blob/master/tests/tracking/test_client.py.

Thank you in advance for your contribution, and please let me know if you have any questions!
> @gabrielfu Thanks!
> 
> I'm fine with this solution (tested with a server on top of sqlite)
> 
> BUT I got an exception with the `run_name` parameter to update_run_info (I'm using MLFlow 1.29) `TypeError: RestStore.update_run_info() got an unexpected keyword argument 'run_name'`
> 
> So the rename function is as below and works fine.
> 
> ```
> import mlflow
> from mlflow.tracking import MlflowClient
> from mlflow.entities.run_status import RunStatus
> 
> client = MlflowClient(sServerURI)
> 
> def set_run_name(client, run_id, run_name):
>     run = client.get_run(run_id)
>     client.set_tag(run_id, "mlflow.runName", run_name)
>     client._tracking_client.store.update_run_info(
>         run_id=run_id,
>         run_status=RunStatus.from_string(run.info.status),
>         end_time=run.info.end_time,
>         #run_name=run_name,
>     )
> ```

Interestinig, as below `RestStore.update_run_info()` does have `run_name` argument. Which `mlflow` version are you using?

https://github.com/mlflow/mlflow/blob/42fc25dbc0e32d58669be4e3024f30b46b2c40a1/mlflow/store/tracking/rest_store.py#L166
@dbczumar great suggestion
@jlmeunier are you interested in working on this? if not, do you mind me doing it? :)
@gabrielfu Given that 1 week has passed, feel free to start working on this. Thank you in advance for your contribution!
@dbczumar great, I'll proceed :)
Hi all, thanks! 
sorry I was off.
@gabrielfu  for the  "RestStore.update_run_info() and its run_name argument": I understand now that my server is in 1.29 while the clients are either in 1.24 or 1.26... (different machines)

Commit message: remove fluent api
Signed-off-by: gabrielfu <hfu.gabriel@gmail.com>

Patch to review:
diff --git a/mlflow/tracking/_tracking_service/client.py b/mlflow/tracking/_tracking_service/client.py
index cd853749b5220..3df56c0b6f6e6 100644
--- a/mlflow/tracking/_tracking_service/client.py
+++ b/mlflow/tracking/_tracking_service/client.py
@@ -357,6 +357,29 @@ def delete_tag(self, run_id, key):
         """
         self.store.delete_tag(run_id, key)
 
+    def update_run(self, run_id, status=None, name=None):
+        """
+        Update a run with the specified ID to a new status or name.
+
+        :param run_id: The ID of the Run to update.
+        :param status: The new status of the run to set, if specified.
+                       At least one of ``status`` or ``name`` should be specified.
+        :param name: The new name of the run to set, if specified.
+                     At least one of ``name`` or ``status`` should be specified.
+        """
+        # Exit early
+        if status is None and name is None:
+            return
+
+        run = self.get_run(run_id)
+        status = status or run.info.status
+        self.store.update_run_info(
+            run_id=run_id,
+            run_status=RunStatus.from_string(status),
+            end_time=run.info.end_time,
+            run_name=name,
+        )
+
     def log_batch(self, run_id, metrics=(), params=(), tags=()):
         """
         Log multiple metrics, params, and/or tags.
diff --git a/mlflow/tracking/client.py b/mlflow/tracking/client.py
index 265ec52f94aea..cb8476acbd9f4 100644
--- a/mlflow/tracking/client.py
+++ b/mlflow/tracking/client.py
@@ -992,6 +992,53 @@ def print_run_info(run):
         """
         self._tracking_client.delete_tag(run_id, key)
 
+    def update_run(
+        self, run_id: str, status: Optional[str] = None, name: Optional[str] = None
+    ) -> None:
+        """
+        Update a run with the specified ID to a new status or name.
+
+        :param run_id: The ID of the Run to update.
+        :param status: The new status of the run to set, if specified.
+                       At least one of ``status`` or ``name`` should be specified.
+        :param name: The new name of the run to set, if specified.
+                     At least one of ``name`` or ``status`` should be specified.
+
+        .. code-block:: python
+            :caption: Example
+
+            from mlflow import MlflowClient
+
+            def print_run_info(run):
+                print("run_id: {}".format(run.info.run_id))
+                print("run_name: {}".format(run.info.run_name))
+                print("status: {}".format(run.info.status))
+
+            # Create a run under the default experiment (whose id is '0').
+            client = MlflowClient()
+            experiment_id = "0"
+            run = client.create_run(experiment_id)
+            print_run_info(run)
+            print("--")
+
+            # Update run and fetch info
+            client.update_run(run.info.run_id, "FINISHED", "new_name")
+            run = client.get_run(run.info.run_id)
+            print_run_info(run)
+
+        .. code-block:: text
+            :caption: Output
+
+            run_id: 1cf6bf8bf6484dd8a598bd43be367b20
+            run_name: judicious-hog-915
+            status: RUNNING
+            --
+            run_id: 1cf6bf8bf6484dd8a598bd43be367b20
+            run_name: new_name
+            status: FINISHED
+        """
+        self._tracking_client.update_run(run_id, status, name)
+
     def log_batch(
         self,
         run_id: str,
diff --git a/tests/tracking/test_client.py b/tests/tracking/test_client.py
index 0dfba018a8c45..4d26a22ac04f5 100644
--- a/tests/tracking/test_client.py
+++ b/tests/tracking/test_client.py
@@ -2,7 +2,7 @@
 import pickle
 from unittest import mock
 
-from mlflow.entities import SourceType, ViewType, RunTag, Run, RunInfo, ExperimentTag
+from mlflow.entities import SourceType, ViewType, RunTag, Run, RunInfo, ExperimentTag, RunStatus
 from mlflow.entities.model_registry import ModelVersion, ModelVersionTag
 from mlflow.entities.model_registry.model_version_status import ModelVersionStatus
 from mlflow.exceptions import MlflowException
@@ -731,3 +731,12 @@ def test_delete_model_version_tag(mock_registry_store_with_get_latest_version):
     # 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")
+
+
+def test_update_run(mock_store):
+    MlflowClient().update_run(run_id="run_id", status="FINISHED", name="my name")
+
+    _, kwargs = mock_store.update_run_info.call_args
+    assert kwargs["run_id"] == "run_id"
+    assert kwargs["run_status"] == RunStatus.from_string("FINISHED")
+    assert kwargs["run_name"] == "my name"


Code review comments:
diff_hunk: @@ -731,3 +731,12 @@ def test_delete_model_version_tag(mock_registry_store_with_get_latest_version):
     # 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")
+
+
+def test_update_run(mock_store):
+    MlflowClient().update_run(run_id="run_id", status="FINISHED", name="my name")
+
+    _, kwargs = mock_store.update_run_info.call_args
+    assert kwargs["run_id"] == "run_id"
+    assert kwargs["run_status"] == RunStatus.from_string("FINISHED")
+    assert kwargs["run_name"] == "my name"
path: tests/tracking/test_client.py
text: @user1:
```suggestion
    mock_store.update_run_info.assert_called_once_with(
        run_id="run_id", status=RunStatus.from_string("FINISHED"), name="my name"
    )
```

We can use `assert_called_once_with` here, which is more strict.

@author:
I think `assert_called_once_with` needs to match all arguments so I would need do this:
```python
    mock_store.update_run_info.assert_called_once_with(
        run_id="run_id",
        run_status=RunStatus.from_string("FINISHED"),
        end_time=mock_store.get_run().info.end_time,
        run_name="my name",
    )
```

Does it look good?

@user1:
For `end_time`, we can use `mock.ANY`.

@author:
good point, i updated it