Repository: optuna/optuna
Pull request:  Temporary fix for problem with length limit of 5000 in MLflow
PR for #1340

Problem statement:
Avoid problem with length limit of 5000 in MLflow
## Motivation
When using the `MLflowCallback` it is possible that the user added `user_attrs` (like lists of float) that are longer then 5000 characters when converted to a `str`. This causes problems with MLflow which limits the length to 5000. See example below:

``` bash
[...]
  File "/home/smay/miniconda3/envs/py38/lib/python3.8/site-packages/mlflow/utils/validation.py", line 136, in _validate_length_limit
    raise MlflowException(
mlflow.exceptions.MlflowException: Tag value '[0.8562690322984875, 0.8544098885636596, 0.8544098885636596, 0.8544098885636596, 0.8544098885636596, 0.859181214773054, 0.86273086038245, 0.86273086038245, 0.86273086038245, 0.86273086038245, 0.86273086038245, 0.8562690322984875, 0.8544098885636596, ' had length 5276, which exceeded length limit of 5000
```

## Description
I suggest to check the strings if they are `len() > 5000` and cut them if needed. If this happens a warning should be printed.

I can provide a PR if wanted. Just give me feedback to my proposal please.

Hints from issue discussion:
I'm not familiar with how "tags" are used in `mlflow`, but would truncating the string representation mitigate your issue? If so, I think that's a reasonable approach. It'd be good to know if this number 5000 can be retrieved via some public API but I guess it's not a must. 

Back to the usage of "tags", would you mind briefly explaining why you need more than 5000 characters since I suppose is beyond what's expected by `mlflow`.
> It'd be good to know if this number 5000 can be retrieved via some public API but I guess it's not a must.

I will have a look into the MLflow code to retrieve the 5000 limit.

> would you mind briefly explaining why you need more than 5000 characters since I suppose is beyond what's expected by mlflow.

Sometimes I am logging a list of feature importances (SHAP values) to the user attr of optuna. Since I have 50000 features (bioinformatics) it is a list of 50000 floats. That is longer then 5000 characters when converted to a string. But there are also other less extreme cases where I want to log the individual results of a nested cross validation.
I asked the MLflow guys if they want to fix this issue on their side: https://github.com/mlflow/mlflow/issues/2931
This issue has not seen any recent activity.
Yes - you are right @github-actions . @hvy what do you think about this?
I saw your ticket on the `mlflow` repository and that you haven't gotten any reply yet. How about truncating in Optuna in the callback, leaving a TODO comment to that ticket?

Commit message: Make sure value in tags is a str

Patch to review:
diff --git a/optuna/integration/mlflow.py b/optuna/integration/mlflow.py
index 85c93486e05..3d7cfa57ce4 100644
--- a/optuna/integration/mlflow.py
+++ b/optuna/integration/mlflow.py
@@ -121,4 +121,16 @@ def __call__(self, study: optuna.study.Study, trial: optuna.trial.FrozenTrial) -
                 (k + "_distribution"): str(v) for (k, v) in trial.distributions.items()
             }
             tags.update(distributions)
+
+            # This is a temporary fix on Optuna side. It avoids an error with user
+            # attributes that are too long. It should be fixed on MLflow side later.
+            # When it is fixed on MLflow side this codeblock can be removed.
+            # see https://github.com/optuna/optuna/issues/1340
+            # see https://github.com/mlflow/mlflow/issues/2931
+            max_mlflow_tag_length = 5000
+            for key, value in tags.items():
+                value = str(value)  # make sure it is a string
+                if len(value) > max_mlflow_tag_length:
+                    tags[key] = value[:max_mlflow_tag_length]
+
             mlflow.set_tags(tags)
diff --git a/tests/integration_tests/test_mlflow.py b/tests/integration_tests/test_mlflow.py
index 705566a93ed..46d392b3a24 100644
--- a/tests/integration_tests/test_mlflow.py
+++ b/tests/integration_tests/test_mlflow.py
@@ -15,6 +15,22 @@ def _objective_func(trial: optuna.trial.Trial) -> float:
     return (x - 2) ** 2 + (y - 25) ** 2 + z
 
 
+# This is tool function for a temporary fix on Optuna side. It avoids an error with user
+# attributes that are too long. It should be fixed on MLflow side later.
+# When it is fixed on MLflow side this test can be removed.
+# see https://github.com/optuna/optuna/issues/1340
+# see https://github.com/mlflow/mlflow/issues/2931
+def _objective_func_long_user_attr(trial: optuna.trial.Trial) -> float:
+
+    x = trial.suggest_uniform("x", -1.0, 1.0)
+    y = trial.suggest_loguniform("y", 20, 30)
+    z = trial.suggest_categorical("z", (-1.0, 1.0))
+    assert isinstance(z, float)
+    long_str = str(list(range(5000)))
+    trial.set_user_attr("my_user_attr", long_str)
+    return (x - 2) ** 2 + (y - 25) ** 2 + z
+
+
 def test_study_name(tmpdir: py.path.local) -> None:
 
     tracking_file_name = "file:{}".format(tmpdir)
@@ -82,3 +98,37 @@ def test_metric_name(tmpdir: py.path.local) -> None:
     first_run_dict = first_run.to_dictionary()
 
     assert metric_name in first_run_dict["data"]["metrics"]
+
+
+# This is a test for a temporary fix on Optuna side. It avoids an error with user
+# attributes that are too long. It should be fixed on MLflow side later.
+# When it is fixed on MLflow side this test can be removed.
+# see https://github.com/optuna/optuna/issues/1340
+# see https://github.com/mlflow/mlflow/issues/2931
+def test_tag_truncation(tmpdir: py.path.local) -> None:
+
+    tracking_file_name = "file:{}".format(tmpdir)
+    study_name = "my_study"
+    n_trials = 3
+
+    mlflc = MLflowCallback(tracking_uri=tracking_file_name)
+    study = optuna.create_study(study_name=study_name)
+    study.optimize(_objective_func_long_user_attr, n_trials=n_trials, callbacks=[mlflc])
+
+    mlfl_client = MlflowClient(tracking_file_name)
+    experiments = mlfl_client.list_experiments()
+    assert len(experiments) == 1
+
+    experiment = experiments[0]
+    assert experiment.name == study_name
+    experiment_id = experiment.experiment_id
+
+    run_infos = mlfl_client.list_run_infos(experiment_id)
+    assert len(run_infos) == n_trials
+
+    first_run_id = run_infos[0].run_id
+    first_run = mlfl_client.get_run(first_run_id)
+    first_run_dict = first_run.to_dictionary()
+
+    my_user_attr = first_run_dict["data"]["tags"]["my_user_attr"]
+    assert len(my_user_attr) <= 5000


Code review comments:
diff_hunk: @@ -121,4 +121,16 @@ def __call__(self, study: optuna.study.Study, trial: optuna.trial.FrozenTrial) -
                 (k + "_distribution"): str(v) for (k, v) in trial.distributions.items()
             }
             tags.update(distributions)
+
+            # This is a temporary fix on Optuna side. It avoids an error with user
+            # attributes that are too long. It should be fixed on MLflow side later.
+            # When it is fixed on MLflow side this codeblock can be removed.
+            # see https://github.com/optuna/optuna/issues/1340
+            # see https://github.com/mlflow/mlflow/issues/2931
+            max_mlflow_tag_length = 5000
+            for key, value in tags.items():
+                value = str(value)  # make sure it is a string
+                if len(value) > max_mlflow_tag_length:
+                    tags[key] = value[:max_mlflow_tag_length]
path: optuna/integration/mlflow.py
text: @user1:
`textwrap.shorten` with `placeholder` could be used but it's a detail, https://docs.python.org/3/library/textwrap.html#textwrap.shorten.

@author:
I will add that tonight.

@author:
I made the change.