Repository: mlflow/mlflow
Pull request: #4299 log_batch actually writes in batches
## What changes are proposed in this pull request?

- Relates to https://github.com/mlflow/mlflow/pull/3024
- Fixes https://github.com/mlflow/mlflow/issues/4299 

In this PR, SqlAlchemyStore.log_batch is implemented in a way that it actually writes in batches. Three new public methods are added:

* log_params: log the given params at once. In case of any IntegrityError, rollback the session and call log_param for
each param and let it handle IntegrityError
* log_metrics: log the given metrics at once, and call _update_latest_metric_if_necessary if the metric instance is recently created.
* set_tags: set the given tags at once.

In all of them, logging logic of single param, metric, or tag is preserved.

The main intuition is to share the same session and not to open new session for each atomic operation and commit the session. To maintain the current behavior, log_params actually commits its own session to catch the integrity error and let it handled by `log_param` method, by logging all params one by one by calling `log_param`.

SqlAlchemyStore lacks docstrings and typing, so I did not add them. But I can add if it is desired. I only added it to protected `_get_or_create_many` method to clarify intention of it.

## How is this patch tested?

It passes current unit tests cases of `tests/store/tracking/test_sqlalchemy_store.py` and new test cases are added for `log_params`, `log_metrics`, and `set_tags`.

I also did a small load test:

```python
import os
from timeit import default_timer as timer

import mlflow
from mlflow.tracking import MlflowClient
from mlflow.entities import Param, RunTag, Metric


os.environ["MLFLOW_TRACKING_URI"] = "sqlite:///mock.db"

params = []
metrics = []
tags = []
N = 100 * 100
MAX_BATCH_SIZE = 100

for i in range(N):
    params.append(Param(str(i), "val"))
    metrics.append(Metric(str(i), i, 0, 0))
    tags.append(RunTag(str(i), "val"))


with mlflow.start_run() as run:
    start = timer()
    for i in range(0, N, MAX_BATCH_SIZE):
        MlflowClient().log_batch(
            run.info.run_id,
            metrics=metrics[i:i+MAX_BATCH_SIZE],
            params=params[i:i+MAX_BATCH_SIZE],
            tags=tags[i:i+MAX_BATCH_SIZE]
        )
    runtime = round(timer() - start, 3)
    print(f"Runtime: {runtime} seconds...")
```

Results of master branch: `Runtime: 246.888 seconds...`
Results of this branch: `Runtime: 4.769 seconds...` 

## Does this PR change the documentation?

- [x] No. You can skip the rest of this section.
- [ ] Yes. Make sure the changed pages / sections render correctly by following the steps below.

1. Check the status of the `ci/circleci: build_doc` check. If it's successful, proceed to the
   next step, otherwise fix it.
2. Click `Details` on the right to open the job page of CircleCI.
3. Click the `Artifacts` tab.
4. Click `docs/build/html/index.html`.
5. 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.

In this change, SqlAlchemyStore.log_batch method is modified and it actually writes in batches now. Also, new public methods are added to SqlAlchemyStore, `log_params`, `log_metrics`, and `set_tags`.

### 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/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="sqlalchemy-store-log-batch"></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
- [ ] `rn/feature` - A new user-facing feature worth mentioning in the release notes
- [x] `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] SqlAlchemyStore.log_batch should really write in batches
## MLflow Roadmap Item

This is an MLflow Roadmap item that has been prioritized by the MLflow maintainers. We're seeking help with the implementation of roadmap items tagged with the `help wanted` label.

For requirements clarifications and implementation questions, or to request a PR review, please tag @WeichenXu123 in your communications related to this issue.

## Proposal Summary
Since MLFlow can easily get really slow, it would be great if the log_batch method really would log in batches. This method currently only iterates over all metrics, params and tags and then calls the log_metric/param/tag methods on each. In my understanding each of these calls opens a new database connection which is very slow. It seriously slows down my training loop when I log the losses in each batch of each epoch.

## Motivation
See the following snippet from `sqlalchemy_store.py`:

```python
def log_batch(self, run_id, metrics, params, tags):
    _validate_run_id(run_id)
    _validate_batch_log_data(metrics, params, tags)
    _validate_batch_log_limits(metrics, params, tags)
    with self.ManagedSessionMaker() as session:
        run = self._get_run(run_uuid=run_id, session=session)
        self._check_run_is_active(run)
    try:
        for param in params:
            self.log_param(run_id, param)
        for metric in metrics:
            self.log_metric(run_id, metric)
        for tag in tags:
            self.set_tag(run_id, tag)
    except MlflowException as e:
        raise e
    except Exception as e:
        raise MlflowException(e, INTERNAL_ERROR)
```

- What is the use case for this feature?
Imagine you are doing a hyperparameter search (e.g. using hyperopt or optuna) and you want to log each of the trials in detail. In my example I am traning 20 models in parallel on 4 GPUs and I want to log the losses. When using the default file storage (mlruns folder), this works very well, but then the UI gets extremely slow. Therefore I am switching to a database backend which makes the UI really fast, but unfortunately the logging very slow.

- Why is this use case valuable to support for MLflow users in general?
Performance is always important and I have already seen many questions about performance issues in MLFlow.

- Why is this use case valuable to support for your project(s) or organization?
Same as above.

- 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)
Logging slows down my training loop a lot. Logging should be highly optimized and the number of DB transactions should be as low as possible.

### What component(s), interfaces, languages, and integrations does this feature affect?
Components 
- [ ] `area/docs`: MLflow documentation pages
- [ ] `area/artifacts`: Artifact stores and artifact logging
- [ ] `area/examples`: Example code
- [ ] `area/model-registry`: Model Registry service, APIs, and the fluent client calls for Model Registry
- [ ] `area/build`: Build and test infrastructure for MLflow
- [ ] `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
- [x] `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

Hints from issue discussion:
@simonhessner Currently, in the log_batch we are opening a single session doing multiple writes. A more efficient strategy would be
to write all in a single batch as opposed to iterative over each MLflow entity in a single DB session.

cc: @harupy @WeichenXu123 for commentary.
https://github.com/mlflow/mlflow/issues/3010 seems to be a related issue
Is there any solution?
What I am currently doing to not slow down my training is to run MLFlow in a separate Thread and do the ML training on the main thread. Both communicate via a thread-safe queue. This works and is fast, but a very ugly solution.
is that mean using something like  'bulk_insert_mapping' or 'add_all' in  one db transaction? cause by using this way,i can speed up maybe  10 times.
#3024 also seems to be related
I opened this PR: https://github.com/mlflow/mlflow/pull/5460

Commit message: Remove deduplicating params since log_batch controls that param keys are
unique (#4299)

Control is done by calling _validate_param_keys_unique

Signed-off-by: Eren Sahin <sahineren.09@gmail.com>

Patch to review:
diff --git a/mlflow/entities/metric.py b/mlflow/entities/metric.py
index e06f9878b9798..355cb3a7cb02e 100644
--- a/mlflow/entities/metric.py
+++ b/mlflow/entities/metric.py
@@ -44,3 +44,12 @@ def to_proto(self):
     @classmethod
     def from_proto(cls, proto):
         return cls(proto.key, proto.value, proto.timestamp, proto.step)
+
+    def __eq__(self, __o):
+        if isinstance(__o, self.__class__):
+            return self.__dict__ == __o.__dict__
+
+        return False
+
+    def __hash__(self):
+        return hash((self._key, self._value, self._timestamp, self._step))
diff --git a/mlflow/entities/param.py b/mlflow/entities/param.py
index e74496cb04ec6..157bc588200bc 100644
--- a/mlflow/entities/param.py
+++ b/mlflow/entities/param.py
@@ -38,3 +38,12 @@ def to_proto(self):
     @classmethod
     def from_proto(cls, proto):
         return cls(proto.key, proto.value)
+
+    def __eq__(self, __o):
+        if isinstance(__o, self.__class__):
+            return self._key == __o._key
+
+        return False
+
+    def __hash__(self):
+        return hash(self._key)
diff --git a/mlflow/store/tracking/sqlalchemy_store.py b/mlflow/store/tracking/sqlalchemy_store.py
index 31bc9cd45a2fa..ced50904d8f51 100644
--- a/mlflow/store/tracking/sqlalchemy_store.py
+++ b/mlflow/store/tracking/sqlalchemy_store.py
@@ -1,5 +1,7 @@
 import json
 import logging
+import random
+import time
 import uuid
 import threading
 
@@ -136,9 +138,9 @@ def __init__(self, db_uri, default_artifact_root):
         if any(table not in inspected_tables for table in expected_tables):
             mlflow.store.db.utils._initialize_tables(self.engine)
         Base.metadata.bind = self.engine
-        SessionMaker = sqlalchemy.orm.sessionmaker(bind=self.engine)
+        self.SessionMaker = sqlalchemy.orm.sessionmaker(bind=self.engine)
         self.ManagedSessionMaker = mlflow.store.db.utils._get_managed_session_maker(
-            SessionMaker, self.db_type
+            self.SessionMaker, self.db_type
         )
         mlflow.store.db.utils._verify_schema(self.engine)
 
@@ -579,7 +581,7 @@ def _get_deleted_runs(self):
             )
             return [run_id[0] for run_id in run_ids]
 
-    def log_metric(self, run_id, metric):
+    def _get_metric_value_details(self, metric):
         _validate_metric(metric.key, metric.value, metric.timestamp, metric.step)
         is_nan = math.isnan(metric.value)
         if is_nan:
@@ -589,28 +591,79 @@ def log_metric(self, run_id, metric):
             value = 1.7976931348623157e308 if metric.value > 0 else -1.7976931348623157e308
         else:
             value = metric.value
-        with self.ManagedSessionMaker() as session:
-            run = self._get_run(run_uuid=run_id, session=session)
-            self._check_run_is_active(run)
-            # ToDo: Consider prior checks for null, type, metric name validations, ... etc.
-            logged_metric, just_created = self._get_or_create(
-                model=SqlMetric,
-                run_uuid=run_id,
-                key=metric.key,
-                value=value,
-                timestamp=metric.timestamp,
-                step=metric.step,
-                session=session,
-                is_nan=is_nan,
-            )
-            # Conditionally update the ``latest_metrics`` table if the logged metric  was not
-            # already present in the ``metrics`` table. If the logged metric was already present,
-            # we assume that the ``latest_metrics`` table already accounts for its presence
-            if just_created:
-                self._update_latest_metric_if_necessary(logged_metric, session)
+        return metric, value, is_nan
 
-    @staticmethod
-    def _update_latest_metric_if_necessary(logged_metric, session):
+    def log_metric(self, run_id, metric):
+        # simply call _log_metrics and let it handle the rest
+        self._log_metrics(run_id, [metric])
+
+    def _log_metrics(self, run_id, metrics):
+        if not metrics:
+            return
+
+        # Duplicate metric values are eliminated here to maintain
+        # the same behavior in log_metric
+        metric_instances = []
+        seen = set()
+        for metric in metrics:
+            metric, value, is_nan = self._get_metric_value_details(metric)
+            if metric not in seen:
+                metric_instances.append(
+                    SqlMetric(
+                        run_uuid=run_id,
+                        key=metric.key,
+                        value=value,
+                        timestamp=metric.timestamp,
+                        step=metric.step,
+                        is_nan=is_nan,
+                    )
+                )
+            seen.add(metric)
+
+        # use SessionManager instead of ManagedSessionManager to set
+        # isolation level for individual transaction.
+        # Here, we open a connection with SERIALIZABLE isolation level.
+        # it will be released when session.commit is called.
+        session = self.SessionMaker()
+        session.connection(execution_options={"isolation_level": "SERIALIZABLE"})
+
+        run = self._get_run(run_uuid=run_id, session=session)
+        self._check_run_is_active(run)
+
+        # commit the session to make sure that we catch any IntegrityError
+        # and try to handle them.
+        try:
+            self._save_to_db(session=session, objs=metric_instances)
+            self._update_latest_metrics_if_necessary(metric_instances, session)
+            session.commit()
+        except sqlalchemy.exc.IntegrityError:
+            # Primary key can be violated if it is tried to log a metric with same value,
+            # timestamp, step, and key within the same run.
+            # Roll back the current session to make it usable for further transactions. In the
+            # event of an error during "commit", a rollback is required in order to continue
+            # using the session. In this case, we re-use the session to query SqlMetric
+            session.rollback()
+            # obtain the metric history corresponding to the given metrics
+            metric_history = (
+                session.query(SqlMetric)
+                .filter(
+                    SqlMetric.run_uuid == run_id,
+                    SqlMetric.key.in_([m.key for m in metric_instances]),
+                )
+                .all()
+            )
+            # convert to a set of Metric instance to take advantage of its hashable
+            # and then obtain the metrics that were not logged earlier within this run_id
+            metric_history = {m.to_mlflow_entity() for m in metric_history}
+            non_existing_metrics = [
+                m for m in metric_instances if m.to_mlflow_entity() not in metric_history
+            ]
+            # if there exist metrics that were tried to be logged & rolled back even though
+            # they were not violating the PK, log them.
+            if non_existing_metrics:
+                self._log_metrics(run_id, non_existing_metrics)
+
+    def _update_latest_metrics_if_necessary(self, logged_metrics, session):
         def _compare_metrics(metric_a, metric_b):
             """
             :return: True if ``metric_a`` is strictly more recent than ``metric_b``, as determined
@@ -622,21 +675,55 @@ def _compare_metrics(metric_a, metric_b):
                 metric_b.value,
             )
 
-        # Fetch the latest metric value corresponding to the specified run_id and metric key and
-        # lock its associated row for the remainder of the transaction in order to ensure
+        def _overwrite_metric(new_metric, old_metric):
+            """
+            writes content of new_metric over old_metric. The content are
+            `value`, `step`, `timestamp`, and `is_nan`.
+
+            :return: old_metric with its content updated.
+            """
+            old_metric.value = new_metric.value
+            old_metric.step = new_metric.step
+            old_metric.timestamp = new_metric.timestamp
+            old_metric.is_nan = new_metric.is_nan
+            return old_metric
+
+        if not logged_metrics:
+            return
+
+        # Fetch the latest metric value corresponding to the specified run_id and metric keys and
+        # lock their associated rows for the remainder of the transaction in order to ensure
         # isolation
-        latest_metric = (
+        latest_metrics = (
             session.query(SqlLatestMetric)
             .filter(
-                SqlLatestMetric.run_uuid == logged_metric.run_uuid,
-                SqlLatestMetric.key == logged_metric.key,
+                SqlLatestMetric.run_uuid == logged_metrics[0].run_uuid,
+                SqlLatestMetric.key.in_([m.key for m in logged_metrics]),
             )
             .with_for_update()
-            .one_or_none()
+            .all()
         )
-        if latest_metric is None or _compare_metrics(logged_metric, latest_metric):
-            session.merge(
-                SqlLatestMetric(
+        latest_metrics = {m.key: m for m in latest_metrics}
+
+        # iterate over all logged metrics and compare them with corresponding
+        # SqlLatestMetric entries
+        # if there's no SqlLatestMetric entry for the current metric key,
+        # create a new SqlLatestMetric instance and put it in
+        # new_latest_metric_dict so that they can be saved later.
+        new_latest_metric_dict = {}
+        for logged_metric in logged_metrics:
+            latest_metric = latest_metrics.get(logged_metric.key)
+            # a metric key can be passed more then once within logged metrics
+            # with different step/timestamp/value. However SqlLatestMetric
+            # entries are inserted after this loop is completed.
+            # so, retrieve the instances they were just created and use them
+            # for comparison.
+            new_latest_metric = new_latest_metric_dict.get(logged_metric.key)
+
+            # just create a new SqlLatestMetric instance since both
+            # latest_metric row or recently created instance does not exist
+            if not latest_metric and not new_latest_metric:
+                new_latest_metric = SqlLatestMetric(
                     run_uuid=logged_metric.run_uuid,
                     key=logged_metric.key,
                     value=logged_metric.value,
@@ -644,7 +731,24 @@ def _compare_metrics(metric_a, metric_b):
                     step=logged_metric.step,
                     is_nan=logged_metric.is_nan,
                 )
-            )
+                new_latest_metric_dict[logged_metric.key] = new_latest_metric
+
+            # there's no row but a new instance is recently created.
+            # so, update the recent instance in new_latest_metric_dict if
+            # metric comparison is successful.
+            elif not latest_metric and new_latest_metric:
+                if _compare_metrics(logged_metric, new_latest_metric):
+                    new_latest_metric = _overwrite_metric(logged_metric, new_latest_metric)
+                    new_latest_metric_dict[logged_metric.key] = new_latest_metric
+
+            # compare with the row
+            elif _compare_metrics(logged_metric, latest_metric):
+                # editing the attributes of latest_metric, which is a
+                # SqlLatestMetric instance will result in UPDATE in DB side.
+                latest_metric = _overwrite_metric(logged_metric, latest_metric)
+
+        if new_latest_metric_dict:
+            self._save_to_db(session=session, objs=list(new_latest_metric_dict.values()))
 
     def get_metric_history(self, run_id, metric_key):
         with self.ManagedSessionMaker() as session:
@@ -697,6 +801,57 @@ def log_param(self, run_id, param):
                 else:
                     raise
 
+    def _log_params(self, run_id, params):
+        if not params:
+            return
+
+        param_instances = [
+            SqlParam(run_uuid=run_id, key=param.key, value=param.value) for param in params
+        ]
+
+        with self.ManagedSessionMaker() as session:
+            run = self._get_run(run_uuid=run_id, session=session)
+            self._check_run_is_active(run)
+            # commit the session to make sure that we catch any IntegrityError
+            # and try to handle them.
+            try:
+                self._save_to_db(session=session, objs=param_instances)
+                session.commit()
+            except sqlalchemy.exc.IntegrityError:
+                # Roll back the current session to make it usable for further transactions. In the
+                # event of an error during "commit", a rollback is required in order to continue
+                # using the session. In this case, we re-use the session because the SqlRun, `run`,
+                # is lazily evaluated during the invocation of `run.params`.
+                session.rollback()
+
+                # in case of an integrity error, compare the parameters of the
+                # run. If the parameters match the ones whom being saved,
+                # ignore the exception since idempotency is reached.
+                # Also, multiple params for the same key can still be passed within
+                # the same batch. So, handle them by selecting the first param
+                # for the given key
+                run_params = {param.key: param.value for param in run.params}
+                non_matching_params = []
+                for param in param_instances:
+                    existing_value = run_params.get(param.key)
+                    if param.value != existing_value:
+                        non_matching_params.append(
+                            {
+                                "key": param.key,
+                                "old_value": existing_value,
+                                "new_value": param.value,
+                            }
+                        )
+
+                if non_matching_params:
+                    raise MlflowException(
+                        "Changing param values is not allowed. Params were already logged='{}'"
+                        " for run ID='{}'.".format(non_matching_params, run_id),
+                        INVALID_PARAMETER_VALUE,
+                    )
+                # if there's no mismatch, do not raise an Exception since
+                # we are sure that idempotency is reached.
+
     def set_experiment_tag(self, experiment_id, tag):
         """
         Set a tag for the specified experiment
@@ -727,6 +882,76 @@ def set_tag(self, run_id, tag):
             self._check_run_is_active(run)
             session.merge(SqlTag(run_uuid=run_id, key=tag.key, value=tag.value))
 
+    def _set_tags(self, run_id, tags, attempt=None):
+        """
+        Set multiple tags on a run
+
+        :param run_id: String ID of the run
+        :param tags: List of RunTag instances to log
+        """
+        if not tags:
+            return
+
+        valid_tags = []
+        for tag in tags:
+            _validate_tag(tag.key, tag.value)
+            valid_tags.append(tag)
+        tags = valid_tags
+
+        with self.ManagedSessionMaker() as session:
+            run = self._get_run(run_uuid=run_id, session=session)
+            self._check_run_is_active(run)
+
+            try:
+                current_tags = (
+                    session.query(SqlTag)
+                    .filter(SqlTag.run_uuid == run_id, SqlTag.key.in_([t.key for t in tags]))
+                    .all()
+                )
+                current_tags = {t.key: t for t in current_tags}
+
+                new_tag_dict = {}
+                for tag in tags:
+                    current_tag = current_tags.get(tag.key)
+                    new_tag = new_tag_dict.get(tag.key)
+
+                    # update the SqlTag if it is already present in DB
+                    if current_tag:
+                        current_tag.value = tag.value
+                        continue
+
+                    # if a SqlTag instance is already present in `new_tag_dict`,
+                    # this means that multiple tags with the same key were passed to `set_tags`.
+                    # In this case, we resolve potential conflicts by updating the value of the
+                    # existing instance to the value of `tag`
+                    if new_tag:
+                        new_tag.value = tag.value
+                    # otherwise, put it into the dict
+                    else:
+                        new_tag = SqlTag(run_uuid=run_id, key=tag.key, value=tag.value)
+
+                    new_tag_dict[tag.key] = new_tag
+
+                # finally, save new entries to DB.
+                self._save_to_db(session=session, objs=list(new_tag_dict.values()))
+                session.commit()
+            except sqlalchemy.exc.IntegrityError:
+                session.rollback()
+                # two concurrent operations may try to attempt to insert tags.
+                # apply retry here.
+                max_retries = 3
+                attempt = (attempt or 0) + 1
+                if attempt > max_retries:
+                    raise MlflowException(
+                        "Failed to set tags with given within {} retries. Keys: {}".format(
+                            max_retries, [t.key for t in tags]
+                        )
+                    )
+                sleep_duration = (2 ** attempt) - 1
+                sleep_duration += random.uniform(0, 1)
+                time.sleep(sleep_duration)
+                self._set_tags(run_id, tags, attempt=attempt)
+
     def delete_tag(self, run_id, key):
         """
         Delete a tag from a run. This is irreversible.
@@ -814,20 +1039,18 @@ def log_batch(self, run_id, metrics, params, tags):
         _validate_batch_log_data(metrics, params, tags)
         _validate_batch_log_limits(metrics, params, tags)
         _validate_param_keys_unique(params)
+
         with self.ManagedSessionMaker() as session:
             run = self._get_run(run_uuid=run_id, session=session)
             self._check_run_is_active(run)
-        try:
-            for param in params:
-                self.log_param(run_id, param)
-            for metric in metrics:
-                self.log_metric(run_id, metric)
-            for tag in tags:
-                self.set_tag(run_id, tag)
-        except MlflowException as e:
-            raise e
-        except Exception as e:
-            raise MlflowException(e, INTERNAL_ERROR)
+            try:
+                self._log_params(run_id, params)
+                self._log_metrics(run_id, metrics)
+                self._set_tags(run_id, tags)
+            except MlflowException as e:
+                raise e
+            except Exception as e:
+                raise MlflowException(e, INTERNAL_ERROR)
 
     def record_logged_model(self, run_id, mlflow_model):
         from mlflow.models import Model
diff --git a/tests/store/tracking/test_sqlalchemy_store.py b/tests/store/tracking/test_sqlalchemy_store.py
index f6ae45d43deb0..78d0efc823eae 100644
--- a/tests/store/tracking/test_sqlalchemy_store.py
+++ b/tests/store/tracking/test_sqlalchemy_store.py
@@ -1665,7 +1665,9 @@ def test_log_batch_param_overwrite_disallowed(self):
             self.store.log_batch(
                 run.info.run_id, metrics=[metric], params=[overwrite_param], tags=[tag]
             )
-        self.assertIn("Changing param values is not allowed. Param with key=", e.exception.message)
+        self.assertIn(
+            "Changing param values is not allowed. Params were already logged=", e.exception.message
+        )
         assert e.exception.error_code == ErrorCode.Name(INVALID_PARAMETER_VALUE)
         self._verify_logged(self.store, run.info.run_id, metrics=[], params=[param], tags=[])
 
@@ -1699,9 +1701,9 @@ def _raise_exception_fn(*args, **kwargs):  # pylint: disable=unused-argument
             raise Exception("Some internal error")
 
         package = "mlflow.store.tracking.sqlalchemy_store.SqlAlchemyStore"
-        with mock.patch(package + ".log_metric") as metric_mock, mock.patch(
-            package + ".log_param"
-        ) as param_mock, mock.patch(package + ".set_tag") as tags_mock:
+        with mock.patch(package + "._log_metrics") as metric_mock, mock.patch(
+            package + "._log_params"
+        ) as param_mock, mock.patch(package + "._set_tags") as tags_mock:
             metric_mock.side_effect = _raise_exception_fn
             param_mock.side_effect = _raise_exception_fn
             tags_mock.side_effect = _raise_exception_fn
@@ -1758,6 +1760,35 @@ def test_log_batch_allows_tag_overwrite_single_req(self):
         self.store.log_batch(run.info.run_id, metrics=[], params=[], tags=tags)
         self._verify_logged(self.store, run.info.run_id, metrics=[], params=[], tags=[tags[-1]])
 
+    def test_log_batch_metrics(self):
+        run = self._run_factory()
+
+        tkey = "blahmetric"
+        tval = 100.0
+        metric = entities.Metric(tkey, tval, int(1000 * time.time()), 0)
+        metric2 = entities.Metric(tkey, tval, int(1000 * time.time()) + 2, 0)
+        nan_metric = entities.Metric("NaN", float("nan"), 0, 0)
+        pos_inf_metric = entities.Metric("PosInf", float("inf"), 0, 0)
+        neg_inf_metric = entities.Metric("NegInf", -float("inf"), 0, 0)
+
+        # duplicate metric and metric2 values should be eliminated
+        metrics = [metric, metric2, nan_metric, pos_inf_metric, neg_inf_metric, metric, metric2]
+        self.store._log_metrics(run.info.run_id, metrics)
+
+        run = self.store.get_run(run.info.run_id)
+        self.assertTrue(tkey in run.data.metrics and run.data.metrics[tkey] == tval)
+
+        # SQL store _get_run method returns full history of recorded metrics.
+        # Should return duplicates as well
+        # MLflow RunData contains only the last reported values for metrics.
+        with self.store.ManagedSessionMaker() as session:
+            sql_run_metrics = self.store._get_run(session, run.info.run_id).metrics
+            self.assertEqual(5, len(sql_run_metrics))
+            self.assertEqual(4, len(run.data.metrics))
+            self.assertTrue(math.isnan(run.data.metrics["NaN"]))
+            self.assertTrue(run.data.metrics["PosInf"] == 1.7976931348623157e308)
+            self.assertTrue(run.data.metrics["NegInf"] == -1.7976931348623157e308)
+
     def test_log_batch_same_metric_repeated_single_req(self):
         run = self._run_factory()
         metric0 = Metric(key="metric-key", value=1, timestamp=2, step=0)
@@ -1778,6 +1809,38 @@ def test_log_batch_same_metric_repeated_multiple_reqs(self):
             self.store, run.info.run_id, params=[], metrics=[metric0, metric1], tags=[]
         )
 
+    def test_log_batch_same_metrics_repeated_multiple_reqs(self):
+        run = self._run_factory()
+        metric0 = Metric(key="metric-key", value=1, timestamp=2, step=0)
+        metric1 = Metric(key="metric-key", value=2, timestamp=3, step=0)
+        self.store.log_batch(run.info.run_id, params=[], metrics=[metric0, metric1], tags=[])
+        self._verify_logged(
+            self.store, run.info.run_id, params=[], metrics=[metric0, metric1], tags=[]
+        )
+        self.store.log_batch(run.info.run_id, params=[], metrics=[metric0, metric1], tags=[])
+        self._verify_logged(
+            self.store, run.info.run_id, params=[], metrics=[metric0, metric1], tags=[]
+        )
+
+    def test_log_batch_null_metrics(self):
+        run = self._run_factory()
+
+        tkey = "blahmetric"
+        tval = None
+        metric_1 = entities.Metric(tkey, tval, int(1000 * time.time()), 0)
+
+        tkey = "blahmetric2"
+        tval = None
+        metric_2 = entities.Metric(tkey, tval, int(1000 * time.time()), 0)
+
+        metrics = [metric_1, metric_2]
+
+        warnings.simplefilter("ignore")
+        with self.assertRaises(MlflowException) as exception_context, warnings.catch_warnings():
+            self.store.log_batch(run.info.run_id, metrics=metrics, params=[], tags=[])
+            warnings.resetwarnings()
+        assert exception_context.exception.error_code == ErrorCode.Name(INVALID_PARAMETER_VALUE)
+
     def test_upgrade_cli_idempotence(self):
         # Repeatedly run `mlflow db upgrade` against our database, verifying that the command
         # succeeds and that the DB has the latest schema


Code review comments:
diff_hunk: @@ -727,6 +882,76 @@ def set_tag(self, run_id, tag):
             self._check_run_is_active(run)
             session.merge(SqlTag(run_uuid=run_id, key=tag.key, value=tag.value))
 
+    def _set_tags(self, run_id, tags, attempt=None):
+        """
+        Set multiple tags on a run
+
+        :param run_id: String ID of the run
+        :param tags: List of RunTag instances to log
+        """
+        if not tags:
+            return
+
+        valid_tags = []
+        for tag in tags:
+            _validate_tag(tag.key, tag.value)
+            valid_tags.append(tag)
+        tags = valid_tags
path: mlflow/store/tracking/sqlalchemy_store.py
text: @user1:
Nit: I don't think we need to create a new list

```suggestion
        for tag in tags:
            _validate_tag(tag.key, tag.value)
```

@author:
Yeah, absolutely better one.

diff_hunk: @@ -1778,6 +1809,38 @@ def test_log_batch_same_metric_repeated_multiple_reqs(self):
             self.store, run.info.run_id, params=[], metrics=[metric0, metric1], tags=[]
         )
 
+    def test_log_batch_same_metrics_repeated_multiple_reqs(self):
+        run = self._run_factory()
+        metric0 = Metric(key="metric-key", value=1, timestamp=2, step=0)
+        metric1 = Metric(key="metric-key", value=2, timestamp=3, step=0)
+        self.store.log_batch(run.info.run_id, params=[], metrics=[metric0, metric1], tags=[])
+        self._verify_logged(
+            self.store, run.info.run_id, params=[], metrics=[metric0, metric1], tags=[]
+        )
+        self.store.log_batch(run.info.run_id, params=[], metrics=[metric0, metric1], tags=[])
+        self._verify_logged(
+            self.store, run.info.run_id, params=[], metrics=[metric0, metric1], tags=[]
+        )
+
+    def test_log_batch_null_metrics(self):
+        run = self._run_factory()
+
+        tkey = "blahmetric"
+        tval = None
+        metric_1 = entities.Metric(tkey, tval, int(1000 * time.time()), 0)
+
+        tkey = "blahmetric2"
+        tval = None
+        metric_2 = entities.Metric(tkey, tval, int(1000 * time.time()), 0)
+
+        metrics = [metric_1, metric_2]
+
+        warnings.simplefilter("ignore")
+        with self.assertRaises(MlflowException) as exception_context, warnings.catch_warnings():
+            self.store.log_batch(run.info.run_id, metrics=metrics, params=[], tags=[])
+            warnings.resetwarnings()
path: tests/store/tracking/test_sqlalchemy_store.py
text: @user1:
```suggestion
        with self.assertRaises(MlflowException) as exception_context:
            self.store.log_batch(run.info.run_id, metrics=metrics, params=[], tags=[])
```

nit: can we remove unnecessary warning capturing?

@user1:
We removed them in https://github.com/mlflow/mlflow/pull/5685