Implement MlflowMetricWriter.write_tags

src/jax_loop_utils/metric_writers/mlflow/metric_writer.py

@@ -16,6 +16,7 @@
 import mlflow.tracking.fluent
 import numpy as np
 from absl import logging
+from mlflow.entities import RunTag
 
 from jax_loop_utils import asynclib
 from jax_loop_utils.metric_writers.interface import (
@@ -82,6 +83,12 @@ def __init__(
             experiment_id=experiment_id, run_name=run_name
         ).info.run_id
 
+    def write_tags(self, tags: dict[str, Any]):
+        """Set tags on the MLFlow run"""
+        self._client.log_batch(
+            self._run_id, [], [], [RunTag(k, str(v)) for k, v in tags.items()], synchronous=False
+        )
+
     def write_scalars(self, step: int, scalars: Mapping[str, Scalar]):
         """Write scalar metrics to MLflow."""
         timestamp = int(time.time() * 1000)

src/jax_loop_utils/metric_writers/mlflow/metric_writer_test.py

@@ -34,6 +34,13 @@ def create_experiment(self, *args, **kwargs):
 
 
 class MlflowMetricWriterTest(absltest.TestCase):
+    def test_set_tags(self):
+        with tempfile.TemporaryDirectory() as temp_dir:
+            tracking_uri = f"file://{temp_dir}"
+            experiment_name = "experiment_name"
+            writer = MlflowMetricWriter(experiment_name, tracking_uri=tracking_uri)
+            writer.set_tags({"ooh": "aah"})
+
     def test_write_scalars(self):
         with tempfile.TemporaryDirectory() as temp_dir:
             tracking_uri = f"file://{temp_dir}"