Repository: mlflow/mlflow

Issue:
[FR] Automatically split metrics, tags, etc into smaller chunks to avoid request limit error
### Willingness to contribute

Yes. I would be willing to contribute this feature with guidance from the MLflow community.

### Proposal Summary

As in the title, MLflow might be able to avoid request limits by automatically splitting data that contains too many elements into smaller ones when we call `mlflow.log_metrics`, `mlflow.log_params`, and `mlflow.set_tags` (possible other API, too).   


### Motivation

> #### What is the use case for this feature?

Suppose `mlflow.log_metrics` case. We would like to save many parameters at once as follows:

```python
import mlflow


metrics = {f"param_{i}": i for i in range(mlflow.utils.validation.MAX_METRICS_PER_BATCH+1)}

# Log a batch of metrics
with mlflow.start_run():
    mlflow.log_metrics(metrics)
```

However, mlflow yields the following error:

```
File /opt/homebrew/Caskroom/miniconda/base/envs/optuna/lib/python3.9/site-packages/mlflow/utils/validation.py:296, in _validate_batch_limit(entity_name, limit, length)
    290 if length > limit:
    291     error_msg = (
    292         "A batch logging request can contain at most {limit} {name}. "
    293         "Got {count} {name}. Please split up {name} across multiple requests and try "
    294         "again."
    295     ).format(name=entity_name, count=length, limit=limit)
--> 296     raise MlflowException(error_msg, error_code=INVALID_PARAMETER_VALUE)

MlflowException: A batch logging request can contain at most 1000 metrics. Got 1001 metrics. Please split up metrics across multiple requests and try again.
```

> #### Why is this use case valuable to support for MLflow users in general?

To avoid this, users need to split `metrics` into smaller subsets like 

```python
from itertools import islice
import mlflow


metrics = {f"param_{i}": i for i in range(mlflow.utils.validation.MAX_METRICS_PER_BATCH+1)}


# Log a batch of metrics
with mlflow.start_run():
    if len(metrics) > mlflow.utils.validation.MAX_METRICS_PER_BATCH:
        it = iter(metrics)
        for _ in range(0, len(metrics), mlflow.utils.validation.MAX_METRICS_PER_BATCH):
            sub_metrics = {k: metrics[k] for k in islice(it, mlflow.utils.validation.MAX_METRICS_PER_BATCH)}
            mlflow.log_metrics(sub_metrics)
    else:
        mlflow.log_metrics(metrics)
```

I suppose this can handle in `mlflow.log_metrics`. By doing so, users do not need to care about the number of elements of `metrics`. 

> #### Why is this use case valuable to support for your project(s) or organization?

I'm from [optuna](https://github.com/optuna/optuna), a black-box optimisation framework library, community. Optuna provides an MLFlow callback that enables us to save optimisation results by using MLFlow API. If this feature request can be done by MLFlow side, we do not need the added changes by https://github.com/optuna/optuna/pull/3651.

> #### Why is it currently difficult to achieve this use case?

N/A

### 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/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`: 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


Discussion:
@nzw0301 This is an excellent idea, and we would be excited to review a PR that implements this capability. Thank you in advance for your contribution!

Resolving patch:
diff --git a/mlflow/tracking/_tracking_service/client.py b/mlflow/tracking/_tracking_service/client.py
--- a/mlflow/tracking/_tracking_service/client.py
+++ b/mlflow/tracking/_tracking_service/client.py
@@ -6,6 +6,7 @@
 
 import time
 import os
+from itertools import zip_longest
 
 from mlflow.store.tracking import SEARCH_MAX_RESULTS_DEFAULT
 from mlflow.tracking._tracking_service import utils
@@ -19,9 +20,15 @@
 from mlflow.exceptions import MlflowException
 from mlflow.protos.databricks_pb2 import INVALID_PARAMETER_VALUE, ErrorCode
 from mlflow.store.artifact.artifact_repository_registry import get_artifact_repository
+from mlflow.utils import chunk_list
 from mlflow.utils.mlflow_tags import MLFLOW_USER
 from mlflow.utils.string_utils import is_string_type
 from mlflow.utils.uri import add_databricks_profile_info_to_artifact_uri
+from mlflow.utils.validation import (
+    MAX_METRICS_PER_BATCH,
+    MAX_PARAMS_TAGS_PER_BATCH,
+    MAX_ENTITIES_PER_BATCH,
+)
 from collections import OrderedDict
 
 
@@ -292,7 +299,25 @@ def log_batch(self, run_id, metrics=(), params=(), tags=()):
         """
         if len(metrics) == 0 and len(params) == 0 and len(tags) == 0:
             return
-        self.store.log_batch(run_id=run_id, metrics=metrics, params=params, tags=tags)
+
+        param_batches = chunk_list(params, MAX_PARAMS_TAGS_PER_BATCH)
+        tag_batches = chunk_list(tags, MAX_PARAMS_TAGS_PER_BATCH)
+
+        for params_batch, tags_batch in zip_longest(param_batches, tag_batches, fillvalue=[]):
+            metrics_batch_size = min(
+                MAX_ENTITIES_PER_BATCH - len(params_batch) - len(tags_batch),
+                MAX_METRICS_PER_BATCH,
+            )
+            metrics_batch_size = max(metrics_batch_size, 0)
+            metrics_batch = metrics[:metrics_batch_size]
+            metrics = metrics[metrics_batch_size:]
+
+            self.store.log_batch(
+                run_id=run_id, metrics=metrics_batch, params=params_batch, tags=tags_batch
+            )
+
+        for metrics_batch in chunk_list(metrics, chunk_size=MAX_METRICS_PER_BATCH):
+            self.store.log_batch(run_id=run_id, metrics=metrics_batch, params=[], tags=[])
 
     def _record_logged_model(self, run_id, mlflow_model):
         from mlflow.models import Model


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
@@ -27,6 +27,10 @@
     MLFLOW_SOURCE_NAME,
     MLFLOW_SOURCE_TYPE,
 )
+from mlflow.utils.validation import (
+    MAX_METRICS_PER_BATCH,
+    MAX_PARAMS_TAGS_PER_BATCH,
+)
 from mlflow.tracking.fluent import _RUN_ID_ENV_VAR
 
 MockExperiment = namedtuple("MockExperiment", ["experiment_id", "lifecycle_stage"])
@@ -340,6 +344,45 @@ def test_log_batch():
             assert new_tags[tag_key] == tag_value
 
 
+@pytest.mark.usefixtures("tmpdir")
+def test_log_batch_with_many_elements():
+    num_metrics = MAX_METRICS_PER_BATCH * 2
+    num_params = num_tags = MAX_PARAMS_TAGS_PER_BATCH * 2
+    expected_metrics = {f"metric-key{i}": float(i) for i in range(num_metrics)}
+    expected_params = {f"param-key{i}": f"param-val{i}" for i in range(num_params)}
+    exact_expected_tags = {f"tag-key{i}": f"tag-val{i}" for i in range(num_tags)}
+
+    t = int(time.time())
+    sorted_expected_metrics = sorted(expected_metrics.items(), key=lambda kv: kv[1])
+    metrics = [
+        Metric(key=key, value=value, timestamp=t, step=i)
+        for i, (key, value) in enumerate(sorted_expected_metrics)
+    ]
+    params = [Param(key=key, value=value) for key, value in expected_params.items()]
+    tags = [RunTag(key=key, value=value) for key, value in exact_expected_tags.items()]
+
+    with start_run() as active_run:
+        run_id = active_run.info.run_id
+        mlflow.tracking.MlflowClient().log_batch(
+            run_id=run_id, metrics=metrics, params=params, tags=tags
+        )
+    client = tracking.MlflowClient()
+    finished_run = client.get_run(run_id)
+    # Validate metrics
+    assert expected_metrics == finished_run.data.metrics
+    for i in range(num_metrics):
+        metric_history = client.get_metric_history(run_id, f"metric-key{i}")
+        assert {(m.value, m.timestamp, m.step) for m in metric_history} == {(float(i), t, i)}
+
+    # Validate tags
+    logged_tags = finished_run.data.tags
+    for tag_key, tag_value in exact_expected_tags.items():
+        assert logged_tags[tag_key] == tag_value
+
+    # Validate params
+    assert finished_run.data.params == expected_params
+
+
 def test_log_metric():
     with start_run() as active_run, mock.patch("time.time") as time_mock:
         time_mock.side_effect = [123 for _ in range(100)]