Repository: mlflow/mlflow

Issue:
[FR] Make log_param() return the passed parameter value (expression for functional use)
## MLflow Roadmap Item

This is an MLflow Roadmap item that has been prioritized by the MLflow maintainers. We’ve identified this feature as a highly requested addition to the MLflow package based on community feedback. 
We're seeking a community contribution for the implementation of this feature and will enthusiastically support the development and review of a submitted PR for this.

## Feature scope

This roadmap feature’s complexity is classified as:

- [X] `good-first-issue`: This feature is limited in complexity and effort required to implement.

## Feature Request

This proposal is to make `mlflow.log_param()` return the passed parameter value so it can be used in a functional way, as an expression.  This will reduce the likelihood of usage errors and increase the speed of development for machine learning practitioners.

## Proposal Summary

This feature will allow `mlflow.log_param()` to return the logged parameter value.  This would allow the in-line use of `log_param`, ensuring that the logged value is actually identical to the value used in the model or process.  

## Motivation
**- What is the use case for this feature?**

This feature can be used in machine learning experiments to ensure that the parameter value used is the same as the value that is logged.

**- Why is this use case valuable to support for MLflow users in general?**

This use case is valuable because it enables rapid, agile development while reducing the likelihood of errors.  It's also backwards compatible.

**- Why is this use case valuable to support for your project(s) or organization?**

This use case supports all data science efforts across our companies' projects.

**- 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)**

Currently this use case cannot be achieved as `mlflow.log_param()` is imperative and does not currently return the value.

### 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
- [ ] `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`: Local serving, model deployment tools, spark UDFs
- [ ] `area/server-infra`: MLflow server, JavaScript dev server
- [x] `area/tracking`: Tracking Service, tracking client APIs, autologging

Interfaces
- [ ] `area/uiux`: Front-end, user experience, JavaScript, plotting
- [ ] `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

Currently `mlflow.log_param()` is an imperative statement. Because of this, multiple lines of code are needed to use `log_param` and parametrize a model. In particular:

1. A variable must be created and set to the desired value (ideally a constant, which isn't really supported in Python)
2. The machine learning model is parametrized with that variable's value
3. `mlflow.log_param()` is called using that variable's value

For example:
```python
NUM_HIDDEN = 30
model.add(Dense(NUM_HIDDEN))
mlflow.log_param("NUM_HIDDEN", NUM_HIDDEN)  # Current MLflow usage pattern
```

Several things can go wrong when an implementer uses the current interface:
- the implementer might inadvertently use a different variable **name** from the name used in the `log_param` call.  This could cause confusion in interpreting the logged parameters w.r.t. the source code (e.g., `mlflow.log_param("NUM_HIDDEN_UNITS", NUM_HIDDEN)`).
- the implementer might inadvertently change the **value** of `NUM_HIDDEN` between steps 2 and 3.  This would result in the wrong value being logged.
- the implementer might inadvertently **omit** or **skip** step 3, resulting in the model running but no value being logged for that parameter.

If the proposal describe in this issue is adopted, the 3 lines of Python above could be rewritten as:

```python
model.add(Dense(mlflow.log_param("NUM_HIDDEN", 30)))   # Proposed MLflow usage pattern
```

This would greatly reduce the possibility of errors in the use of MLflow.



Discussion:
I also think it's pretty helpful to have these params deal with the same pointers, so when params are logged earlier during the code and then changed by frameworks later, the changes get reflected in MLFlow.
@serena-ruan volunteers to do this task, assign it to her. :)

Resolving patch:
diff --git a/mlflow/tracking/client.py b/mlflow/tracking/client.py
--- a/mlflow/tracking/client.py
+++ b/mlflow/tracking/client.py
@@ -798,7 +798,7 @@ def print_run_info(r):
         """
         self._tracking_client.log_metric(run_id, key, value, timestamp, step)
 
-    def log_param(self, run_id: str, key: str, value: Any) -> None:
+    def log_param(self, run_id: str, key: str, value: Any) -> Any:
         """
         Log a parameter (e.g. model hyperparameter) against the run ID.
 
@@ -810,6 +810,7 @@ def log_param(self, run_id: str, key: str, value: Any) -> None:
         :param value: Parameter value (string, but will be string-ified if not).
                       All backend stores support values up to length 500, but some
                       may support larger values.
+        :return: the parameter value that is logged.
 
         .. code-block:: python
             :caption: Example
@@ -833,7 +834,8 @@ def print_run_info(r):
             # Log the parameter. Unlike mlflow.log_param this method
             # does not start a run if one does not exist. It will log
             # the parameter in the backend store
-            client.log_param(run.info.run_id, "p", 1)
+            p_value = client.log_param(run.info.run_id, "p", 1)
+            assert p_value == 1
             client.set_terminated(run.info.run_id)
             run = client.get_run(run.info.run_id)
             print_run_info(run)
@@ -850,6 +852,7 @@ def print_run_info(r):
             status: FINISHED
         """
         self._tracking_client.log_param(run_id, key, value)
+        return value
 
     def set_experiment_tag(self, experiment_id: str, key: str, value: Any) -> None:
         """
diff --git a/mlflow/tracking/fluent.py b/mlflow/tracking/fluent.py
--- a/mlflow/tracking/fluent.py
+++ b/mlflow/tracking/fluent.py
@@ -514,7 +514,7 @@ def get_run(run_id: str) -> Run:
     return MlflowClient().get_run(run_id)
 
 
-def log_param(key: str, value: Any) -> None:
+def log_param(key: str, value: Any) -> Any:
     """
     Log a parameter (e.g. model hyperparameter) under the current run. If no run is active,
     this method will create a new active run.
@@ -527,16 +527,19 @@ def log_param(key: str, value: Any) -> None:
                   All backend stores support values up to length 500, but some
                   may support larger values.
 
+    :return: the parameter value that is logged.
+
     .. code-block:: python
         :caption: Example
 
         import mlflow
 
         with mlflow.start_run():
-            mlflow.log_param("learning_rate", 0.01)
+            value = mlflow.log_param("learning_rate", 0.01)
+            assert value == 0.01
     """
     run_id = _get_or_start_run().info.run_id
-    MlflowClient().log_param(run_id, key, value)
+    return MlflowClient().log_param(run_id, key, value)
 
 
 def set_experiment_tag(key: str, value: Any) -> None:


Test patch:
diff --git a/tests/tracking/test_tracking.py b/tests/tracking/test_tracking.py
--- a/tests/tracking/test_tracking.py
+++ b/tests/tracking/test_tracking.py
@@ -501,9 +501,9 @@ def test_log_metric_validation():
 def test_log_param():
     with start_run() as active_run:
         run_id = active_run.info.run_id
-        mlflow.log_param("name_1", "a")
-        mlflow.log_param("name_2", "b")
-        mlflow.log_param("nested/nested/name", 5)
+        assert mlflow.log_param("name_1", "a") == "a"
+        assert mlflow.log_param("name_2", "b") == "b"
+        assert mlflow.log_param("nested/nested/name", 5) == 5
     finished_run = tracking.MlflowClient().get_run(run_id)
     # Validate params
     assert finished_run.data.params == {"name_1": "a", "name_2": "b", "nested/nested/name": "5"}