Repository: mlflow/mlflow
Pull request: Make parent id a parameter
<details><summary>&#x1F6E0 DevTools &#x1F6E0</summary>
<p>

[![Open in GitHub Codespaces](https://github.com/codespaces/badge.svg)](https://codespaces.new/Flametaa/mlflow/pull/12721?quickstart=1)

#### Install mlflow from this PR

```
pip install git+https://github.com/mlflow/mlflow.git@refs/pull/12721/merge
```

#### Checkout with GitHub CLI

```
gh pr checkout 12721
```

</p>
</details>

### Related Issues/PRs

<!-- Uncomment 'Resolve' if this PR can close the linked items. -->
Resolve  #12639 

### What changes are proposed in this pull request?

<!-- Please fill in changes proposed in this PR. -->

### How is this PR tested?

- [x] Existing unit/integration tests
- [x] New unit/integration tests
- [ ] Manual tests

<!-- Attach code, screenshot, video used for manual testing here. -->

### Does this PR require documentation update?

- [ ] No. You can skip the rest of this section.
- [ ] Yes. I've updated:
  - [ ] Examples
  - [ ] API references
  - [ ] Instructions

### Release Notes

#### Is this a user-facing change?

- [ ] No. You can skip the rest of this section.
- [x] Yes. This impacts Python API users, they can now use the argument `parent_run_id` to start a nested run under a specific parent run (without having to start the parent run first)

<!-- 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/deployments`: MLflow Deployments client APIs, server, and third-party Deployments integrations
- [ ] `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/recipes`: Recipes, Recipe APIs, Recipe configs, Recipe 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/none` - No description will be included. The PR will be mentioned only by the PR number in the "Small Bugfixes and Documentation Updates" section
- [ ] `rn/breaking-change` - The PR will be mentioned in the "Breaking Changes" 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

#### Should this PR be included in the next patch release?

`Yes` should be selected for bug fixes, documentation updates, and other small changes. `No` should be selected for new features and larger changes. If you're unsure about the release classification of this PR, leave this unchecked to let the maintainers decide.

<details>
<summary>What is a minor/patch release?</summary>

- Minor release: a release that increments the second part of the version number (e.g., 1.2.0 -> 1.3.0).
  Bug fixes, doc updates and new features usually go into minor releases.
- Patch release: a release that increments the third part of the version number (e.g., 1.2.0 -> 1.2.1).
  Bug fixes and doc updates usually go into patch releases.

</details>

<!-- patch -->

- [x] Yes (this PR will be cherry-picked and included in the next patch release)
- [ ] No (this PR will be included in the next minor release)


Problem statement:
[FR] Make parent_run_id a parameter in start_run for hyperparameter tuning using multiple jobs
### Willingness to contribute

Yes. I can contribute this feature independently.

### Proposal Summary

Add `parent_run_id ` as a parameter in `mlflow.start_run` to avoid having to start a parent run in order to be able to use nested runs.

### Motivation

> #### What is the use case for this feature?
Performing hyperparameter tuning using multiple jobs in parallel on Kubernetes. Each trial needs to create nested runs under the same parent run. It would be great if we can simply specify the parent_run_id for nested runs where the run already exists.

> #### Why is this use case valuable to support for MLflow users in general?
I believe that this change can make the user experience better.

> #### Why is this use case valuable to support for your project(s) or organization?
Our training pipelines are time consuming and it is a requirement for us to parallelize hyperparameter tuning experiments.

> #### Why is it currently difficult to achieve this use case?
It is possible to use the current approach to implement this, however the code looks more complex as we simply need to start the nested runs at each trial iteration. 

### Details

_No response_

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

- [ ] `area/artifacts`: Artifact stores and artifact logging
- [ ] `area/build`: Build and test infrastructure for MLflow
- [ ] `area/deployments`: MLflow Deployments client APIs, server, and third-party Deployments integrations
- [ ] `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/recipes`: Recipes, Recipe APIs, Recipe configs, Recipe 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:
@Flametaa The proposal makes a lot sense! Would you mind contributing for this feature?

A few pointers to start:
* The parent_id is recorded to the run as a system tag here. https://github.com/mlflow/mlflow/blob/master/mlflow/tracking/fluent.py#L362C9-L362C22
* We may want to deprecate the existing `nested` parameter as `parent_run_id` appears to be more precise and convenient. https://github.com/mlflow/mlflow/blob/master/mlflow/tracking/fluent.py#L238

Please let me know if there is any help needed to implement the feature. Thank you so much!
I don't think it is a good idea to deprecate `nested` parameter, as we might still want to be able to infer the `parent_id` from the active runs. In that sense, if we already have an active run, we might still want to be able to create other runs that are not nested. What I propose is:

- keep both `nested` and `parent_id` parameters
- If `parent_id` is specified then the run will be nested under that parent run
- else if `nested` is True, then we infer the `parent_id` from the active run
- else if `nested` is False then the run is not nested (so the last two steps will remain the same as now)
What do you think @B-Step62 ?
@Flametaa Thank you for giving some thoughts on the design!

> If parent_id is specified then the run will be nested under that parent run
> else if nested is True, then we infer the parent_id from the active run
> else if nested is False then the run is not nested (so the last two steps will remain the same as now)

The third option is not allowed in current implementation. If the active run exists but a new run is started with `nested=False`, MLflow raise an exception [here](https://github.com/mlflow/mlflow/blob/master/mlflow/tracking/fluent.py#L305). This is because some of the Run operation relies on global context and may cause pollution between two runs that are intended to be separate. Therefore, practically we don't really support `nested=False` case.

That said, I'm fine with keeping the parameter for now. Then this change can be framed as a pure enhancement and no breaking change. Let's proceed with that approach🙂
<!-- assign-maintainer -->
@mlflow/mlflow-team Please assign a maintainer and start triaging this issue.

Commit message: add example

Signed-off-by: SIALA KARIM <37486371+Flametaa@users.noreply.github.com>

Patch to review:
diff --git a/mlflow/tracking/fluent.py b/mlflow/tracking/fluent.py
index ebe55a044f0b9..2fd0eed806203 100644
--- a/mlflow/tracking/fluent.py
+++ b/mlflow/tracking/fluent.py
@@ -204,6 +204,7 @@ def start_run(
     experiment_id: Optional[str] = None,
     run_name: Optional[str] = None,
     nested: bool = False,
+    parent_run_id: Optional[str] = None,
     tags: Optional[Dict[str, Any]] = None,
     description: Optional[str] = None,
     log_system_metrics: Optional[bool] = None,
@@ -236,6 +237,8 @@ def start_run(
         run_name: Name of new run. Used only when ``run_id`` is unspecified. If a new run is
             created and ``run_name`` is not specified, a random name will be generated for the run.
         nested: Controls whether run is nested in parent run. ``True`` creates a nested run.
+        parent_run_id: If specified, the current run will be nested under the the run with
+            the specified UUID.
         tags: An optional dictionary of string keys and values to set as tags on the run.
             If a run is being resumed, these tags are set on the resumed run. If a new run is
             being created, these tags are set on the new run.
@@ -279,12 +282,22 @@ def start_run(
         print("version tag value: {}".format(parent_run.data.tags.get("version")))
         print("priority tag value: {}".format(parent_run.data.tags.get("priority")))
         print("--")
+
         # Search all child runs with a parent id
         query = f"tags.mlflow.parentRunId = '{parent_run.info.run_id}'"
         results = mlflow.search_runs(experiment_ids=[experiment_id], filter_string=query)
         print("child runs:")
         print(results[["run_id", "params.child", "tags.mlflow.runName"]])
 
+        # Create a nested run under the existing parent run
+        with mlflow.start_run(
+            run_name="NEW_CHILD_RUN",
+            experiment_id=experiment_id,
+            description="new child",
+            parent_run_id=parent_run.info.run_id,
+        ) as child_run:
+            mlflow.log_param("new-child", "yes")
+
     .. code-block:: text
         :caption: Output
 
@@ -359,7 +372,28 @@ def start_run(
             )
         active_run_obj = client.get_run(existing_run_id)
     else:
-        parent_run_id = _active_run_stack[-1].info.run_id if len(_active_run_stack) > 0 else None
+        if parent_run_id:
+            _validate_run_id(parent_run_id)
+            # Make sure parent_run_id matches the current run id, if there is an active run
+            if len(_active_run_stack) > 0 and parent_run_id != _active_run_stack[0].info.run_id:
+                raise Exception(
+                    (
+                        "Current run with UUID {} does not match the specified parent_run_id {}"
+                        + " To start a new nested run under the parent run with UUID {}, "
+                        + "first end the current run with mlflow.end_run()."
+                    ).format(_active_run_stack[0].info.run_id, parent_run_id)
+                )
+            parent_run_obj = client.get_run(parent_run_id)
+            # Check if the specified parent_run has been deleted.
+            if parent_run_obj.info.lifecycle_stage == LifecycleStage.DELETED:
+                raise MlflowException(
+                    f"Cannot start run under parent run with ID {parent_run_id} "
+                    f"because it is in the deleted state."
+                )
+        else:
+            parent_run_id = (
+                _active_run_stack[-1].info.run_id if len(_active_run_stack) > 0 else None
+            )
 
         exp_id_for_run = experiment_id if experiment_id is not None else _get_experiment_id()
 
diff --git a/tests/tracking/fluent/test_fluent.py b/tests/tracking/fluent/test_fluent.py
index 0c407a0101cf5..f0f90a0de81ac 100644
--- a/tests/tracking/fluent/test_fluent.py
+++ b/tests/tracking/fluent/test_fluent.py
@@ -749,6 +749,16 @@ def test_start_run_with_parent():
         assert is_from_run(active_run, MlflowClient.create_run.return_value)
 
 
+@pytest.mark.usefixtures(empty_active_run_stack.__name__)
+def test_start_run_with_parent_id():
+    parent_run_id = mlflow.start_run().info.run_id
+    mlflow.end_run()
+    nested_run_id = mlflow.start_run(parent_run_id=parent_run_id).info.run_id
+    mlflow.end_run()
+
+    assert mlflow.get_parent_run(nested_run_id).info.run_id == parent_run_id
+
+
 def test_start_run_with_parent_non_nested():
     with mock.patch("mlflow.tracking.fluent._active_run_stack", [mock.Mock()]):
         with pytest.raises(Exception, match=r"Run with UUID .+ is already active"):


Code review comments:
diff_hunk: @@ -236,6 +237,8 @@ def start_run(
         run_name: Name of new run. Used only when ``run_id`` is unspecified. If a new run is
             created and ``run_name`` is not specified, a random name will be generated for the run.
         nested: Controls whether run is nested in parent run. ``True`` creates a nested run.
+        parent_run_id: If specified, the current run will be nested under the the run with
+            the specified UUID.
path: mlflow/tracking/fluent.py
text: @user1:
```suggestion
        parent_run_id: If specified, the current run will be nested under the the run with
            the specified UUID. The parent run must be in the ACTIVE state.
```

diff_hunk: @@ -359,7 +372,28 @@ def start_run(
             )
         active_run_obj = client.get_run(existing_run_id)
     else:
-        parent_run_id = _active_run_stack[-1].info.run_id if len(_active_run_stack) > 0 else None
+        if parent_run_id:
+            _validate_run_id(parent_run_id)
+            # Make sure parent_run_id matches the current run id, if there is an active run
+            if len(_active_run_stack) > 0 and parent_run_id != _active_run_stack[0].info.run_id:
+                raise Exception(
+                    (
+                        "Current run with UUID {} does not match the specified parent_run_id {}"
+                        + " To start a new nested run under the parent run with UUID {}, "
+                        + "first end the current run with mlflow.end_run()."
+                    ).format(_active_run_stack[0].info.run_id, parent_run_id)
path: mlflow/tracking/fluent.py
text: @user1:
```suggestion
            if len(_active_run_stack) > 0 and parent_run_id != _active_run_stack[-1].info.run_id:
                raise Exception(
                    (
                        "Current run with UUID {} does not match the specified parent_run_id {}"
                        + " To start a new nested run under the parent run with UUID {}, "
                        + "first end the current run with mlflow.end_run()."
                    ).format(_active_run_stack[-1].info.run_id, parent_run_id)
```

Should this be consistent with the behavior of `nested=True`?