feat: add ability to specify mlflow tag for tesseract (#426)

<!--
Please use a PR title that conforms to *conventional commits*:
"<commit_type>: Describe your change"; for example: "fix: prevent race
condition". Some other commit types are: fix, feat, ci, doc, refactor...
For a full list of commit types visit
https://www.conventionalcommits.org/en/v1.0.0/
-->

#### Relevant issue or PR
<!-- If the changes resolve an issue or follow some other PR, link to
them here. Only link something if it is directly relevant. -->

#### Description of changes
<!-- Add a high-level description of changes, focusing on the *what* and
*why*. -->

This adds the ability to specify mlflow tags for Tesseract runs via the
`TESSERACT_MLFLOW_TAGS` env variable. This is useful to embed additional
metadata into the MLFlow run.

n.b. As we want to ensure that all runs have the tags if set, we need to
make a small change here to check if there is an active run. Previously,
if there was no an active run and you logged something, it would
implicitly create one. That is a problem as it will not include the
tags. Instead we now explicitly check if there an active run; and, if
not, we create one with the tags.

#### Testing done
<!-- Describe how the changes were tested; e.g., "CI passes", "Tested
manually in stagingrepo#123", screenshots of a terminal session that
verify the changes, or any other evidence of testing the changes. -->

Local and unit testing

---------

Co-authored-by: Dion Häfner <dion.haefner@simulation.science>

Author: nbowen3

docs/content/using-tesseracts/advanced.md

@@ -75,6 +75,30 @@ $ tesseract serve --env=TESSERACT_MLFLOW_TRACKING_URI="..." \
     metrics
 ````
 
+If you wish to pass additional parameters to the MLflow run (such as tags, run name, or description), you can do so via the `TESSERACT_MLFLOW_RUN_EXTRA_ARGS` environment variable. This accepts a Python dictionary string that is passed directly to `mlflow.start_run()`. See supported
+parameters in the [mlflow documentation](https://mlflow.org/docs/latest/api_reference/python_api/mlflow.html#mlflow.start_run).
+
+**Example: Setting tags only**
+```bash
+$ tesseract serve --env=TESSERACT_MLFLOW_TRACKING_URI="..." \
+    --env=TESSERACT_MLFLOW_RUN_EXTRA_ARGS='{"tags": {"key1": "value1", "key2": "value2"}}' \
+    metrics
+```
+
+**Example: Setting run name and tags**
+```bash
+$ tesseract serve --env=TESSERACT_MLFLOW_TRACKING_URI="..." \
+    --env=TESSERACT_MLFLOW_RUN_EXTRA_ARGS='{"run_name": "my_experiment", "tags": {"env": "production"}}' \
+    metrics
+```
+
+**Example: Multiple parameters**
+```bash
+$ tesseract serve --env=TESSERACT_MLFLOW_TRACKING_URI="..." \
+    --env=TESSERACT_MLFLOW_RUN_EXTRA_ARGS='{"run_name": "test_run", "description": "Testing new feature", "tags": {"version": "1.0"}}' \
+    metrics
+```
+
 ## Volume mounts and user permissions
 
 When mounting a volume into a Tesseract container, default behavior depends on the Docker engine being used. Specifically, Docker Desktop, Docker Engine, and Podman have different ways of handling user permissions for mounted volumes.

tesseract_core/runtime/cli.py

@@ -141,6 +141,10 @@ def main_callback(**kwargs: Any) -> None:
         if get_origin(field_type) is Literal:
             field_type = make_choice_enum(f"{field_name}Choices", get_args(field_type))
 
+        if get_origin(field_type) is dict:
+            # Dicts are parsed as strings
+            field_type = str
+
         params.append(
             inspect.Parameter(
                 field_name,

tesseract_core/runtime/config.py

@@ -1,15 +1,26 @@
 # Copyright 2025 Pasteur Labs. All Rights Reserved.
 # SPDX-License-Identifier: Apache-2.0
 
+import ast
 import os
 from pathlib import Path
-from typing import Any
+from typing import Annotated, Any
 
-from pydantic import BaseModel, ConfigDict, FilePath
+from pydantic import BaseModel, BeforeValidator, ConfigDict, Field, FilePath
 
 from tesseract_core.runtime.file_interactions import supported_format_type
 
 
+def _eval_str(obj: Any) -> Any:
+    """Evaluate a string into the corresponding Python object."""
+    if isinstance(obj, str):
+        try:
+            return ast.literal_eval(obj)
+        except SyntaxError as exc:
+            raise ValueError("Could not parse string as Python object") from exc
+    return obj
+
+
 class RuntimeConfig(BaseModel):
     """Available runtime configuration."""
 
@@ -25,6 +36,9 @@ class RuntimeConfig(BaseModel):
     mlflow_tracking_uri: str = ""
     mlflow_tracking_username: str = ""
     mlflow_tracking_password: str = ""
+    mlflow_run_extra_args: Annotated[dict[str, Any], BeforeValidator(_eval_str)] = (
+        Field(default_factory=dict)
+    )
 
     model_config = ConfigDict(frozen=True, extra="forbid")
 

tesseract_core/runtime/mpa.py

@@ -215,8 +215,10 @@ def log_artifact(self, local_path: str) -> None:
         self.mlflow.log_artifact(local_path)
 
     def start_run(self) -> None:
-        """Start a new MLflow run."""
-        self.mlflow.start_run()
+        """Start a new MLflow run with optional extra arguments from config."""
+        config = get_config()
+        run_extra_args = config.mlflow_run_extra_args
+        self.mlflow.start_run(**run_extra_args)
 
     def end_run(self) -> None:
         """End the current MLflow run."""

tests/runtime_tests/test_mpa.py

@@ -9,6 +9,7 @@
 import sqlite3
 
 import pytest
+from pydantic import ValidationError
 
 from tesseract_core.runtime import mpa
 from tesseract_core.runtime.config import update_config
@@ -293,3 +294,43 @@ def test_build_tracking_uri_sqlite_ignores_credentials():
     assert "testuser" not in tracking_uri
     assert "testpass" not in tracking_uri
     assert tracking_uri.startswith("sqlite:///")
+
+
+def test_mlflow_run_extra_args(mocker):
+    """Test passing a dict with basic tags."""
+    pytest.importorskip("mlflow")
+
+    kwargs = {"tags": {"env": "prod", "team": "ml"}}
+    kwargs_str = repr(kwargs)
+
+    update_config(mlflow_run_extra_args=kwargs_str)
+    mocked_mlflow = mocker.Mock()
+
+    backend = mpa.MLflowBackend()
+    backend.mlflow = mocked_mlflow
+
+    # Make sure kwargs are forwarded correctly to mlflow.start_run
+    backend.start_run()
+    mocked_mlflow.start_run.assert_called_with(**kwargs)
+
+
+def test_mlflow_run_extra_args_parsing():
+    # This is actually a test for config.py but we add it here for now
+
+    with pytest.raises(ValidationError):
+        # Not a valid Python object
+        update_config(mlflow_run_extra_args="{'unbalanced dict': True")
+
+    with pytest.raises(ValidationError):
+        # Not a dict
+        update_config(mlflow_run_extra_args="['this is a list']")
+
+    with pytest.raises(ValidationError):
+        # Not str keys
+        update_config(mlflow_run_extra_args="{0: 'hey there'}")
+
+    # All good
+    update_config(mlflow_run_extra_args="{'hey there': 'general kenobi'}")
+
+    # Passing dicts directly is fine, too
+    update_config(mlflow_run_extra_args={"hey there": "general kenobi"})