mlflow--mlflow/tests/tracking/test_tracking.py

import filecmp
import io
import json
import os
import pathlib
import posixpath
import random
import re
from datetime import datetime, timezone
from typing import NamedTuple
from unittest import mock

import pytest
import yaml

import mlflow
from mlflow import MlflowClient, tracking
from mlflow.entities import LifecycleStage, Metric, Param, RunStatus, RunTag, ViewType
from mlflow.environment_variables import (
    MLFLOW_ASYNC_LOGGING_THREADPOOL_SIZE,
    MLFLOW_RUN_ID,
)
from mlflow.exceptions import MlflowException
from mlflow.protos.databricks_pb2 import (
    INVALID_PARAMETER_VALUE,
    RESOURCE_DOES_NOT_EXIST,
    ErrorCode,
)
from mlflow.store.tracking.file_store import FileStore
from mlflow.tracking._tracking_service.client import TrackingServiceClient
from mlflow.tracking.fluent import start_run
from mlflow.utils.file_utils import local_file_uri_to_path
from mlflow.utils.mlflow_tags import (
    MLFLOW_PARENT_RUN_ID,
    MLFLOW_RUN_NAME,
    MLFLOW_SOURCE_NAME,
    MLFLOW_SOURCE_TYPE,
    MLFLOW_USER,
)
from mlflow.utils.os import is_windows
from mlflow.utils.time import get_current_time_millis
from mlflow.utils.validation import (
    MAX_METRICS_PER_BATCH,
    MAX_PARAMS_TAGS_PER_BATCH,
)


class MockExperiment(NamedTuple):
    experiment_id: str
    lifecycle_stage: str
    tags: dict[str, str] = {}


def test_create_experiment():
    with pytest.raises(MlflowException, match="Invalid experiment name"):
        mlflow.create_experiment(None)

    with pytest.raises(MlflowException, match="Invalid experiment name"):
        mlflow.create_experiment("")

    exp_id = mlflow.create_experiment(f"Some random experiment name {random.randint(1, int(1e6))}")
    assert exp_id is not None


def test_create_experiment_with_duplicate_name():
    name = "popular_name"
    exp_id = mlflow.create_experiment(name)

    with pytest.raises(MlflowException, match=re.escape(f"Experiment(name={name}) already exists")):
        mlflow.create_experiment(name)

    tracking.MlflowClient().delete_experiment(exp_id)
    with pytest.raises(MlflowException, match=re.escape(f"Experiment(name={name}) already exists")):
        mlflow.create_experiment(name)


def test_create_experiments_with_bad_names():
    # None for name
    with pytest.raises(MlflowException, match="Invalid experiment name: 'None'"):
        mlflow.create_experiment(None)

    # empty string name
    with pytest.raises(MlflowException, match="Invalid experiment name: ''"):
        mlflow.create_experiment("")


@pytest.mark.parametrize("name", [123, 0, -1.2, [], ["A"], {1: 2}])
def test_create_experiments_with_bad_name_types(name):
    with pytest.raises(
        MlflowException,
        match=re.escape(f"Invalid experiment name: {name}. Expects a string."),
    ):
        mlflow.create_experiment(name)


@pytest.mark.usefixtures("reset_active_experiment")
def test_set_experiment_by_name():
    name = "random_exp"
    exp_id = mlflow.create_experiment(name)
    exp1 = mlflow.set_experiment(name)
    assert exp1.experiment_id == exp_id
    with start_run() as run:
        assert run.info.experiment_id == exp_id

    another_name = "another_experiment"
    exp2 = mlflow.set_experiment(another_name)
    with start_run() as another_run:
        assert another_run.info.experiment_id == exp2.experiment_id


@pytest.mark.usefixtures("reset_active_experiment")
def test_set_experiment_by_id():
    name = "random_exp"
    exp_id = mlflow.create_experiment(name)
    active_exp = mlflow.set_experiment(experiment_id=exp_id)
    assert active_exp.experiment_id == exp_id
    with start_run() as run:
        assert run.info.experiment_id == exp_id

    nonexistent_id = "-1337"
    with pytest.raises(MlflowException, match="No Experiment with id=-1337 exists") as exc:
        mlflow.set_experiment(experiment_id=nonexistent_id)
    assert exc.value.error_code == ErrorCode.Name(RESOURCE_DOES_NOT_EXIST)
    with start_run() as run:
        assert run.info.experiment_id == exp_id


def test_set_experiment_parameter_validation():
    with pytest.raises(MlflowException, match="Must specify exactly one") as exc:
        mlflow.set_experiment()
    assert exc.value.error_code == ErrorCode.Name(INVALID_PARAMETER_VALUE)

    with pytest.raises(MlflowException, match="Must specify exactly one") as exc:
        mlflow.set_experiment(None)
    assert exc.value.error_code == ErrorCode.Name(INVALID_PARAMETER_VALUE)

    with pytest.raises(MlflowException, match="Must specify exactly one") as exc:
        mlflow.set_experiment(None, None)
    assert exc.value.error_code == ErrorCode.Name(INVALID_PARAMETER_VALUE)

    with pytest.raises(MlflowException, match="Must specify exactly one") as exc:
        mlflow.set_experiment("name", "id")
    assert exc.value.error_code == ErrorCode.Name(INVALID_PARAMETER_VALUE)


def test_set_experiment_with_deleted_experiment():
    name = "dead_exp"
    mlflow.set_experiment(name)
    with start_run() as run:
        exp_id = run.info.experiment_id

    tracking.MlflowClient().delete_experiment(exp_id)

    with pytest.raises(MlflowException, match="Cannot set a deleted experiment") as exc:
        mlflow.set_experiment(name)
    assert exc.value.error_code == ErrorCode.Name(INVALID_PARAMETER_VALUE)

    with pytest.raises(MlflowException, match="Cannot set a deleted experiment") as exc:
        mlflow.set_experiment(experiment_id=exp_id)
    assert exc.value.error_code == ErrorCode.Name(INVALID_PARAMETER_VALUE)


@pytest.mark.usefixtures("reset_active_experiment")
def test_set_experiment_with_zero_id():
    mock_experiment = MockExperiment(experiment_id=0, lifecycle_stage=LifecycleStage.ACTIVE)
    with (
        mock.patch.object(
            TrackingServiceClient,
            "get_experiment_by_name",
            mock.Mock(return_value=mock_experiment),
        ) as get_experiment_by_name_mock,
        mock.patch.object(TrackingServiceClient, "create_experiment") as create_experiment_mock,
    ):
        mlflow.set_experiment("my_exp")
        get_experiment_by_name_mock.assert_called_once()
        create_experiment_mock.assert_not_called()


def test_start_run_context_manager():
    with start_run() as first_run:
        first_uuid = first_run.info.run_id
        # Check that start_run() causes the run information to be persisted in the store
        persisted_run = tracking.MlflowClient().get_run(first_uuid)
        assert persisted_run is not None
        assert persisted_run.info == first_run.info
    finished_run = tracking.MlflowClient().get_run(first_uuid)
    assert finished_run.info.status == RunStatus.to_string(RunStatus.FINISHED)
    # Launch a separate run that fails, verify the run status is FAILED and the run UUID is
    # different
    with pytest.raises(Exception, match="Failing run!"):
        with start_run() as second_run:
            raise Exception("Failing run!")
    second_run_id = second_run.info.run_id
    assert second_run_id != first_uuid
    finished_run2 = tracking.MlflowClient().get_run(second_run_id)
    assert finished_run2.info.status == RunStatus.to_string(RunStatus.FAILED)


def test_start_and_end_run():
    # Use the start_run() and end_run() APIs without a `with` block, verify they work.

    with start_run() as active_run:
        mlflow.log_metric("name_1", 25)
    finished_run = tracking.MlflowClient().get_run(active_run.info.run_id)
    # Validate metrics
    assert len(finished_run.data.metrics) == 1
    assert finished_run.data.metrics["name_1"] == 25


def test_metric_timestamp():
    with mlflow.start_run() as active_run:
        mlflow.log_metric("name_1", 25)
        mlflow.log_metric("name_1", 30)
        run_id = active_run.info.run_id
    # Check that metric timestamps are between run start and finish
    client = MlflowClient()
    history = client.get_metric_history(run_id, "name_1")
    finished_run = client.get_run(run_id)
    assert len(history) == 2
    assert all(
        m.timestamp >= finished_run.info.start_time and m.timestamp <= finished_run.info.end_time
        for m in history
    )


def test_log_batch():
    expected_metrics = {"metric-key0": 1.0, "metric-key1": 4.0}
    expected_params = {
        "param-key0": "param-val0",
        "param-key1": 123,
        "param-key2": None,
    }
    exact_expected_tags = {"tag-key0": "tag-val0", "tag-key1": "tag-val1"}
    approx_expected_tags = {
        MLFLOW_USER,
        MLFLOW_SOURCE_NAME,
        MLFLOW_SOURCE_TYPE,
        MLFLOW_RUN_NAME,
    }

    t = get_current_time_millis()
    sorted_expected_metrics = sorted(expected_metrics.items(), key=lambda kv: kv[0])
    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
        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 len(finished_run.data.metrics) == 2
    for key, value in finished_run.data.metrics.items():
        assert expected_metrics[key] == value
    metric_history0 = client.get_metric_history(run_id, "metric-key0")
    assert {(m.value, m.timestamp, m.step) for m in metric_history0} == {(1.0, t, 0)}
    metric_history1 = client.get_metric_history(run_id, "metric-key1")
    assert {(m.value, m.timestamp, m.step) for m in metric_history1} == {(4.0, t, 1)}

    # Validate tags (for automatically-set tags)
    assert len(finished_run.data.tags) == len(exact_expected_tags) + len(approx_expected_tags)
    for tag_key, tag_value in finished_run.data.tags.items():
        if tag_key in approx_expected_tags:
            pass
        else:
            assert exact_expected_tags[tag_key] == tag_value
    # Validate params
    assert finished_run.data.params == {k: str(v) for k, v in expected_params.items()}
    # test that log_batch works with fewer params
    new_tags = {"1": "2", "3": "4", "5": "6"}
    tags = [RunTag(key=key, value=value) for key, value in new_tags.items()]
    client.log_batch(run_id=run_id, tags=tags)
    finished_run_2 = client.get_run(run_id)
    # Validate tags (for automatically-set tags)
    assert len(finished_run_2.data.tags) == len(finished_run.data.tags) + 3
    for tag_key, tag_value in finished_run_2.data.tags.items():
        if tag_key in new_tags:
            assert new_tags[tag_key] == tag_value


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 = get_current_time_millis()
    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", return_value=123):
        run_id = active_run.info.run_id
        mlflow.log_metric("name_1", 25)
        mlflow.log_metric("name_2", -3)
        mlflow.log_metric("name_1", 30, 5)
        mlflow.log_metric("name_1", 40, -2)
        mlflow.log_metric("nested/nested/name", 40)
    finished_run = tracking.MlflowClient().get_run(run_id)
    # Validate metrics
    assert len(finished_run.data.metrics) == 3
    expected_pairs = {"name_1": 30, "name_2": -3, "nested/nested/name": 40}
    for key, value in finished_run.data.metrics.items():
        assert expected_pairs[key] == value
    client = tracking.MlflowClient()
    metric_history_name1 = client.get_metric_history(run_id, "name_1")
    assert {(m.value, m.timestamp, m.step) for m in metric_history_name1} == {
        (25, 123 * 1000, 0),
        (30, 123 * 1000, 5),
        (40, 123 * 1000, -2),
    }
    metric_history_name2 = client.get_metric_history(run_id, "name_2")
    assert {(m.value, m.timestamp, m.step) for m in metric_history_name2} == {(-3, 123 * 1000, 0)}


def test_log_metrics_uses_millisecond_timestamp_resolution_fluent():
    with start_run() as active_run, mock.patch("time.time") as time_mock:
        time_mock.side_effect = lambda: 123
        mlflow.log_metrics({"name_1": 25, "name_2": -3})
        mlflow.log_metrics({"name_1": 30})
        mlflow.log_metrics({"name_1": 40})
        run_id = active_run.info.run_id

    client = tracking.MlflowClient()
    metric_history_name1 = client.get_metric_history(run_id, "name_1")
    assert {(m.value, m.timestamp) for m in metric_history_name1} == {
        (25, 123 * 1000),
        (30, 123 * 1000),
        (40, 123 * 1000),
    }
    metric_history_name2 = client.get_metric_history(run_id, "name_2")
    assert {(m.value, m.timestamp) for m in metric_history_name2} == {(-3, 123 * 1000)}


def test_log_metrics_uses_millisecond_timestamp_resolution_client():
    with start_run() as active_run, mock.patch("time.time") as time_mock:
        time_mock.side_effect = lambda: 123
        mlflow_client = tracking.MlflowClient()
        run_id = active_run.info.run_id

        mlflow_client.log_metric(run_id=run_id, key="name_1", value=25)
        mlflow_client.log_metric(run_id=run_id, key="name_2", value=-3)
        mlflow_client.log_metric(run_id=run_id, key="name_1", value=30)
        mlflow_client.log_metric(run_id=run_id, key="name_1", value=40)

    metric_history_name1 = mlflow_client.get_metric_history(run_id, "name_1")
    assert {(m.value, m.timestamp) for m in metric_history_name1} == {
        (25, 123 * 1000),
        (30, 123 * 1000),
        (40, 123 * 1000),
    }

    metric_history_name2 = mlflow_client.get_metric_history(run_id, "name_2")
    assert {(m.value, m.timestamp) for m in metric_history_name2} == {(-3, 123 * 1000)}


@pytest.mark.parametrize("step_kwarg", [None, -10, 5])
def test_log_metrics_uses_common_timestamp_and_step_per_invocation(step_kwarg):
    expected_metrics = {"name_1": 30, "name_2": -3, "nested/nested/name": 40}
    with start_run() as active_run:
        run_id = active_run.info.run_id
        mlflow.log_metrics(expected_metrics, step=step_kwarg)
    finished_run = tracking.MlflowClient().get_run(run_id)
    # Validate metric key/values match what we expect, and that all metrics have the same timestamp
    assert len(finished_run.data.metrics) == len(expected_metrics)
    for key, value in finished_run.data.metrics.items():
        assert expected_metrics[key] == value
    common_timestamp = finished_run.data._metric_objs[0].timestamp
    expected_step = step_kwarg if step_kwarg is not None else 0
    for metric_obj in finished_run.data._metric_objs:
        assert metric_obj.timestamp == common_timestamp
        assert metric_obj.step == expected_step


@pytest.fixture
def get_store_mock():
    with mock.patch("mlflow.store.file_store.FileStore.log_batch") as _get_store_mock:
        yield _get_store_mock


def test_set_tags():
    exact_expected_tags = {"name_1": "c", "name_2": "b", "nested/nested/name": 5}
    approx_expected_tags = {
        MLFLOW_USER,
        MLFLOW_SOURCE_NAME,
        MLFLOW_SOURCE_TYPE,
        MLFLOW_RUN_NAME,
    }
    with start_run() as active_run:
        run_id = active_run.info.run_id
        mlflow.set_tags(exact_expected_tags)
    finished_run = tracking.MlflowClient().get_run(run_id)
    # Validate tags
    assert len(finished_run.data.tags) == len(exact_expected_tags) + len(approx_expected_tags)
    for tag_key, tag_val in finished_run.data.tags.items():
        if tag_key in approx_expected_tags:
            pass
        else:
            assert str(exact_expected_tags[tag_key]) == tag_val


def test_log_metric_validation():
    with start_run() as active_run:
        run_id = active_run.info.run_id
        with pytest.raises(
            MlflowException,
            match="Invalid value \"apple\" for parameter 'value' supplied",
        ) as e:
            mlflow.log_metric("name_1", "apple")
    assert e.value.error_code == ErrorCode.Name(INVALID_PARAMETER_VALUE)
    finished_run = tracking.MlflowClient().get_run(run_id)
    assert len(finished_run.data.metrics) == 0


def test_log_param():
    with start_run() as active_run:
        run_id = active_run.info.run_id
        assert mlflow.log_param("name_1", "a") == "a"
        assert mlflow.log_param("name_2", "b") == "b"
        assert mlflow.log_param("nested/nested/name", 5) == 5
    finished_run = tracking.MlflowClient().get_run(run_id)
    # Validate params
    assert finished_run.data.params == {
        "name_1": "a",
        "name_2": "b",
        "nested/nested/name": "5",
    }


def test_log_params():
    expected_params = {"name_1": "c", "name_2": "b", "nested/nested/name": 5}
    with start_run() as active_run:
        run_id = active_run.info.run_id
        mlflow.log_params(expected_params)
    finished_run = tracking.MlflowClient().get_run(run_id)
    # Validate params
    assert finished_run.data.params == {
        "name_1": "c",
        "name_2": "b",
        "nested/nested/name": "5",
    }


def test_log_params_duplicate_keys_raises():
    params = {"a": "1", "b": "2"}
    with start_run() as active_run:
        run_id = active_run.info.run_id
        mlflow.log_params(params)
        with pytest.raises(
            expected_exception=MlflowException,
            match=r"Changing param values is not allowed. Param with key=",
        ) as e:
            mlflow.log_param("a", "3")
        assert e.value.error_code == ErrorCode.Name(INVALID_PARAMETER_VALUE)
    finished_run = tracking.MlflowClient().get_run(run_id)
    assert finished_run.data.params == params


@pytest.mark.skipif(is_windows(), reason="Windows do not support colon in params and metrics")
def test_param_metric_with_colon():
    with start_run() as active_run:
        run_id = active_run.info.run_id
        mlflow.log_param("a:b", 3)
        mlflow.log_metric("c:d", 4)
    finished_run = tracking.MlflowClient().get_run(run_id)

    # Validate param
    assert len(finished_run.data.params) == 1
    assert finished_run.data.params == {"a:b": "3"}

    # Validate metric
    assert len(finished_run.data.metrics) == 1
    assert finished_run.data.metrics["c:d"] == 4


def test_log_batch_duplicate_entries_raises():
    with start_run() as active_run:
        run_id = active_run.info.run_id
        with pytest.raises(
            MlflowException, match=r"Duplicate parameter keys have been submitted."
        ) as e:
            tracking.MlflowClient().log_batch(
                run_id=run_id, params=[Param("a", "1"), Param("a", "2")]
            )
        assert e.value.error_code == ErrorCode.Name(INVALID_PARAMETER_VALUE)


def test_log_batch_validates_entity_names_and_values():
    with start_run() as active_run:
        run_id = active_run.info.run_id

        metrics = [Metric(key="../bad/metric/name", value=0.3, timestamp=3, step=0)]
        with pytest.raises(
            MlflowException,
            match=r"Invalid value \"../bad/metric/name\" for parameter \'metrics\[0\].name\'",
        ) as e:
            tracking.MlflowClient().log_batch(run_id, metrics=metrics)
        assert e.value.error_code == ErrorCode.Name(INVALID_PARAMETER_VALUE)

        metrics = [Metric(key="ok-name", value="non-numerical-value", timestamp=3, step=0)]
        with pytest.raises(
            MlflowException,
            match=r"Invalid value \"non-numerical-value\" "
            + r"for parameter \'metrics\[0\].value\' supplied",
        ) as e:
            tracking.MlflowClient().log_batch(run_id, metrics=metrics)
        assert e.value.error_code == ErrorCode.Name(INVALID_PARAMETER_VALUE)

        metrics = [Metric(key="ok-name", value=0.3, timestamp="non-numerical-timestamp", step=0)]
        with pytest.raises(
            MlflowException,
            match=r"Invalid value \"non-numerical-timestamp\" for "
            + r"parameter \'metrics\[0\].timestamp\' supplied",
        ) as e:
            tracking.MlflowClient().log_batch(run_id, metrics=metrics)
        assert e.value.error_code == ErrorCode.Name(INVALID_PARAMETER_VALUE)

        params = [Param(key="../bad/param/name", value="my-val")]
        with pytest.raises(
            MlflowException,
            match=r"Invalid value \"../bad/param/name\" for parameter \'params\[0\].key\' supplied",
        ) as e:
            tracking.MlflowClient().log_batch(run_id, params=params)
        assert e.value.error_code == ErrorCode.Name(INVALID_PARAMETER_VALUE)

        tags = [Param(key="../bad/tag/name", value="my-val")]
        with pytest.raises(
            MlflowException,
            match=r"Invalid value \"../bad/tag/name\" for parameter \'tags\[0\].key\' supplied",
        ) as e:
            tracking.MlflowClient().log_batch(run_id, tags=tags)
        assert e.value.error_code == ErrorCode.Name(INVALID_PARAMETER_VALUE)

        metrics = [Metric(key=None, value=42.0, timestamp=4, step=1)]
        with pytest.raises(
            MlflowException,
            match="Metric name cannot be None. A key name must be provided.",
        ) as e:
            tracking.MlflowClient().log_batch(run_id, metrics=metrics)
        assert e.value.error_code == ErrorCode.Name(INVALID_PARAMETER_VALUE)


def test_log_artifact_with_dirs(tmp_path):
    # Test log artifact with a directory
    art_dir = tmp_path / "parent"
    art_dir.mkdir()
    file0 = art_dir.joinpath("file0")
    file0.write_text("something")
    file1 = art_dir.joinpath("file1")
    file1.write_text("something")
    sub_dir = art_dir / "child"
    sub_dir.mkdir()
    with start_run():
        artifact_uri = mlflow.get_artifact_uri()
        run_artifact_dir = local_file_uri_to_path(artifact_uri)
        mlflow.log_artifact(str(art_dir))
        base = os.path.basename(str(art_dir))
        assert os.listdir(run_artifact_dir) == [base]
        assert set(os.listdir(os.path.join(run_artifact_dir, base))) == {
            "child",
            "file0",
            "file1",
        }
        with open(os.path.join(run_artifact_dir, base, "file0")) as f:
            assert f.read() == "something"
    # Test log artifact with directory and specified parent folder

    art_dir = tmp_path / "dir"
    art_dir.mkdir()
    with start_run():
        artifact_uri = mlflow.get_artifact_uri()
        run_artifact_dir = local_file_uri_to_path(artifact_uri)
        mlflow.log_artifact(str(art_dir), "some_parent")
        assert os.listdir(run_artifact_dir) == [os.path.basename("some_parent")]
        assert os.listdir(os.path.join(run_artifact_dir, "some_parent")) == [
            os.path.basename(str(art_dir))
        ]
    sub_dir = art_dir.joinpath("another_dir")
    sub_dir.mkdir()
    with start_run():
        artifact_uri = mlflow.get_artifact_uri()
        run_artifact_dir = local_file_uri_to_path(artifact_uri)
        mlflow.log_artifact(str(art_dir), "parent/and_child")
        assert os.listdir(os.path.join(run_artifact_dir, "parent", "and_child")) == [
            os.path.basename(str(art_dir))
        ]
        assert set(
            os.listdir(
                os.path.join(
                    run_artifact_dir,
                    "parent",
                    "and_child",
                    os.path.basename(str(art_dir)),
                )
            )
        ) == {os.path.basename(str(sub_dir))}


def test_log_artifact(tmp_path):
    # Create artifacts
    artifact_dir = tmp_path.joinpath("artifact_dir")
    artifact_dir.mkdir()
    path0 = artifact_dir.joinpath("file0")
    path1 = artifact_dir.joinpath("file1")
    path0.write_text("0")
    path1.write_text("1")
    # Log an artifact, verify it exists in the directory returned by get_artifact_uri
    # after the run finishes
    artifact_parent_dirs = ["some_parent_dir", None]
    for parent_dir in artifact_parent_dirs:
        with start_run():
            artifact_uri = mlflow.get_artifact_uri()
            run_artifact_dir = local_file_uri_to_path(artifact_uri)
            mlflow.log_artifact(path0, parent_dir)
        expected_dir = (
            os.path.join(run_artifact_dir, parent_dir)
            if parent_dir is not None
            else run_artifact_dir
        )
        assert os.listdir(expected_dir) == [os.path.basename(path0)]
        logged_artifact_path = os.path.join(expected_dir, path0)
        assert filecmp.cmp(logged_artifact_path, path0, shallow=False)
    # Log multiple artifacts, verify they exist in the directory returned by get_artifact_uri
    for parent_dir in artifact_parent_dirs:
        with start_run():
            artifact_uri = mlflow.get_artifact_uri()
            run_artifact_dir = local_file_uri_to_path(artifact_uri)

            mlflow.log_artifacts(artifact_dir, parent_dir)
        # Check that the logged artifacts match
        expected_artifact_output_dir = (
            os.path.join(run_artifact_dir, parent_dir)
            if parent_dir is not None
            else run_artifact_dir
        )
        dir_comparison = filecmp.dircmp(artifact_dir, expected_artifact_output_dir)
        assert len(dir_comparison.left_only) == 0
        assert len(dir_comparison.right_only) == 0
        assert len(dir_comparison.diff_files) == 0
        assert len(dir_comparison.funny_files) == 0


@pytest.mark.parametrize("subdir", [None, ".", "dir", "dir1/dir2", "dir/.."])
def test_log_text(subdir):
    filename = "file.txt"
    text = "a"
    artifact_file = filename if subdir is None else posixpath.join(subdir, filename)

    with mlflow.start_run():
        mlflow.log_text(text, artifact_file)

        artifact_path = None if subdir is None else posixpath.normpath(subdir)
        artifact_uri = mlflow.get_artifact_uri(artifact_path)
        run_artifact_dir = local_file_uri_to_path(artifact_uri)
        assert os.listdir(run_artifact_dir) == [filename]

        filepath = os.path.join(run_artifact_dir, filename)
        with open(filepath) as f:
            assert f.read() == text


@pytest.mark.parametrize("subdir", [None, ".", "dir", "dir1/dir2", "dir/.."])
@pytest.mark.parametrize("extension", [".json", ".yml", ".yaml", ".txt", ""])
def test_log_dict(subdir, extension):
    dictionary = {"k": "v"}
    filename = "data" + extension
    artifact_file = filename if subdir is None else posixpath.join(subdir, filename)

    with mlflow.start_run():
        mlflow.log_dict(dictionary, artifact_file)

        artifact_path = None if subdir is None else posixpath.normpath(subdir)
        artifact_uri = mlflow.get_artifact_uri(artifact_path)
        run_artifact_dir = local_file_uri_to_path(artifact_uri)
        assert os.listdir(run_artifact_dir) == [filename]

        filepath = os.path.join(run_artifact_dir, filename)
        extension = os.path.splitext(filename)[1]
        with open(filepath) as f:
            loaded = (
                # Specify `Loader` to suppress the following deprecation warning:
                # https://github.com/yaml/pyyaml/wiki/PyYAML-yaml.load(input)-Deprecation
                yaml.load(f, Loader=yaml.SafeLoader)
                if (extension in [".yml", ".yaml"])
                else json.load(f)
            )
            assert loaded == dictionary


@pytest.mark.parametrize("subdir", [None, ".", "dir", "dir1/dir2", "dir/.."])
def test_log_stream_bytes(subdir):
    filename = "file.bin"
    content = b"binary content"
    artifact_file = filename if subdir is None else posixpath.join(subdir, filename)

    with mlflow.start_run():
        stream = io.BytesIO(content)
        mlflow.log_stream(stream, artifact_file)

        artifact_path = None if subdir is None else posixpath.normpath(subdir)
        artifact_uri = mlflow.get_artifact_uri(artifact_path)
        run_artifact_dir = pathlib.Path(local_file_uri_to_path(artifact_uri))
        assert list(run_artifact_dir.iterdir()) == [run_artifact_dir / filename]
        assert (run_artifact_dir / filename).read_bytes() == content


def test_log_stream_empty():
    with mlflow.start_run():
        artifact_uri = mlflow.get_artifact_uri()
        run_artifact_dir = pathlib.Path(local_file_uri_to_path(artifact_uri))

        stream = io.BytesIO(b"")
        mlflow.log_stream(stream, "empty.bin")
        assert (run_artifact_dir / "empty.bin").read_bytes() == b""


def test_log_stream_large_content():
    with mlflow.start_run():
        # Large binary content (larger than chunk size of 8192)
        large_content = b"x" * 100000
        stream = io.BytesIO(large_content)
        mlflow.log_stream(stream, "large.bin")

        artifact_uri = mlflow.get_artifact_uri()
        run_artifact_dir = pathlib.Path(local_file_uri_to_path(artifact_uri))
        assert (run_artifact_dir / "large.bin").read_bytes() == large_content


def test_with_startrun():
    run_id = None
    t0 = get_current_time_millis()
    with mlflow.start_run() as active_run:
        assert mlflow.active_run() == active_run
        run_id = active_run.info.run_id
    t1 = get_current_time_millis()
    run_info = mlflow.tracking._get_store().get_run(run_id).info
    assert run_info.status == "FINISHED"
    assert t0 <= run_info.end_time
    assert run_info.end_time <= t1
    assert mlflow.active_run() is None


def test_parent_create_run(monkeypatch):
    with mlflow.start_run() as parent_run:
        parent_run_id = parent_run.info.run_id
    monkeypatch.setenv(MLFLOW_RUN_ID.name, parent_run_id)
    with mlflow.start_run() as parent_run:
        assert parent_run.info.run_id == parent_run_id
        with pytest.raises(Exception, match="To start a nested run"):
            mlflow.start_run()
        with mlflow.start_run(nested=True) as child_run:
            assert child_run.info.run_id != parent_run_id
            with mlflow.start_run(nested=True) as grand_child_run:
                pass

    def verify_has_parent_id_tag(child_id, expected_parent_id):
        tags = tracking.MlflowClient().get_run(child_id).data.tags
        assert tags[MLFLOW_PARENT_RUN_ID] == expected_parent_id

    verify_has_parent_id_tag(child_run.info.run_id, parent_run.info.run_id)
    verify_has_parent_id_tag(grand_child_run.info.run_id, child_run.info.run_id)
    assert mlflow.active_run() is None


def test_start_deleted_run():
    run_id = None
    with mlflow.start_run() as active_run:
        run_id = active_run.info.run_id
    tracking.MlflowClient().delete_run(run_id)
    with pytest.raises(MlflowException, match="because it is in the deleted state."):
        with mlflow.start_run(run_id=run_id):
            pass
    assert mlflow.active_run() is None


@pytest.mark.usefixtures("reset_active_experiment")
def test_start_run_exp_id_0():
    mlflow.set_experiment("some-experiment")
    # Create a run and verify that the current active experiment is the one we just set
    with mlflow.start_run() as active_run:
        exp_id = active_run.info.experiment_id
        assert exp_id != FileStore.DEFAULT_EXPERIMENT_ID
        assert MlflowClient().get_experiment(exp_id).name == "some-experiment"
    # Set experiment ID to 0 when creating a run, verify that the specified experiment ID is honored
    with mlflow.start_run(experiment_id=0) as active_run:
        assert active_run.info.experiment_id == FileStore.DEFAULT_EXPERIMENT_ID


def test_get_artifact_uri_with_artifact_path_unspecified_returns_artifact_root_dir():
    with mlflow.start_run() as active_run:
        assert mlflow.get_artifact_uri(artifact_path=None) == active_run.info.artifact_uri


def test_get_artifact_uri_uses_currently_active_run_id():
    artifact_path = "artifact"
    with mlflow.start_run() as active_run:
        assert mlflow.get_artifact_uri(
            artifact_path=artifact_path
        ) == tracking.artifact_utils.get_artifact_uri(
            run_id=active_run.info.run_id, artifact_path=artifact_path
        )


def _assert_get_artifact_uri_appends_to_uri_path_component_correctly(
    artifact_location, expected_uri_format
):
    client = MlflowClient()
    client.create_experiment("get-artifact-uri-test", artifact_location=artifact_location)
    mlflow.set_experiment("get-artifact-uri-test")
    with mlflow.start_run():
        run_id = mlflow.active_run().info.run_id
        for artifact_path in ["path/to/artifact", "/artifact/path", "arty.txt"]:
            artifact_uri = mlflow.get_artifact_uri(artifact_path)
            assert artifact_uri == tracking.artifact_utils.get_artifact_uri(run_id, artifact_path)
            assert artifact_uri == expected_uri_format.format(
                run_id=run_id,
                path=artifact_path.lstrip("/"),
                drive=pathlib.Path.cwd().drive,
            )


@pytest.mark.parametrize(
    ("artifact_location", "expected_uri_format"),
    [
        (
            "mysql://user:password@host:port/dbname?driver=mydriver",
            "mysql://user:password@host:port/dbname/{run_id}/artifacts/{path}?driver=mydriver",
        ),
        (
            "mysql+driver://user:pass@host:port/dbname/subpath/#fragment",
            "mysql+driver://user:pass@host:port/dbname/subpath/{run_id}/artifacts/{path}#fragment",
        ),
        (
            "s3://bucketname/rootpath",
            "s3://bucketname/rootpath/{run_id}/artifacts/{path}",
        ),
    ],
)
def test_get_artifact_uri_appends_to_uri_path_component_correctly(
    artifact_location, expected_uri_format
):
    _assert_get_artifact_uri_appends_to_uri_path_component_correctly(
        artifact_location, expected_uri_format
    )


@pytest.mark.skipif(not is_windows(), reason="This test only passes on Windows")
def test_get_artifact_uri_appends_to_local_path_component_correctly_on_windows():
    _assert_get_artifact_uri_appends_to_uri_path_component_correctly(
        "/dirname/rootpa#th?",
        "file:///{drive}/dirname/rootpa/{run_id}/artifacts/{path}#th?",
    )


@pytest.mark.skipif(is_windows(), reason="This test fails on Windows")
def test_get_artifact_uri_appends_to_local_path_component_correctly():
    _assert_get_artifact_uri_appends_to_uri_path_component_correctly(
        "/dirname/rootpa#th?", "{drive}/dirname/rootpa#th?/{run_id}/artifacts/{path}"
    )


@pytest.mark.usefixtures("reset_active_experiment")
def test_search_runs():
    mlflow.set_experiment("exp-for-search")
    # Create a run and verify that the current active experiment is the one we just set
    logged_runs = {}
    with mlflow.start_run() as active_run:
        logged_runs["first"] = active_run.info.run_id
        mlflow.log_metric("m1", 0.001)
        mlflow.log_metric("m2", 0.002)
        mlflow.log_metric("m1", 0.002)
        mlflow.log_param("p1", "a")
        mlflow.set_tag("t1", "first-tag-val")
    with mlflow.start_run() as active_run:
        logged_runs["second"] = active_run.info.run_id
        mlflow.log_metric("m1", 0.008)
        mlflow.log_param("p2", "aa")
        mlflow.set_tag("t2", "second-tag-val")

    def verify_runs(runs, expected_set):
        assert {r.info.run_id for r in runs} == {logged_runs[r] for r in expected_set}

    experiment_id = MlflowClient().get_experiment_by_name("exp-for-search").experiment_id

    # 2 runs in this experiment
    assert len(MlflowClient().search_runs([experiment_id], run_view_type=ViewType.ACTIVE_ONLY)) == 2

    # 2 runs that have metric "m1" > 0.001
    runs = MlflowClient().search_runs([experiment_id], "metrics.m1 > 0.0001")
    verify_runs(runs, ["first", "second"])

    # 1 run with has metric "m1" > 0.002
    runs = MlflowClient().search_runs([experiment_id], "metrics.m1 > 0.002")
    verify_runs(runs, ["second"])

    # no runs with metric "m1" > 0.1
    runs = MlflowClient().search_runs([experiment_id], "metrics.m1 > 0.1")
    verify_runs(runs, [])

    # 1 run with metric "m2" > 0
    runs = MlflowClient().search_runs([experiment_id], "metrics.m2 > 0")
    verify_runs(runs, ["first"])

    # 1 run each with param "p1" and "p2"
    runs = MlflowClient().search_runs([experiment_id], "params.p1 = 'a'", ViewType.ALL)
    verify_runs(runs, ["first"])
    runs = MlflowClient().search_runs([experiment_id], "params.p2 != 'a'", ViewType.ALL)
    verify_runs(runs, ["second"])
    runs = MlflowClient().search_runs([experiment_id], "params.p2 = 'aa'", ViewType.ALL)
    verify_runs(runs, ["second"])

    # 1 run each with tag "t1" and "t2"
    runs = MlflowClient().search_runs([experiment_id], "tags.t1 = 'first-tag-val'", ViewType.ALL)
    verify_runs(runs, ["first"])
    runs = MlflowClient().search_runs([experiment_id], "tags.t2 != 'qwerty'", ViewType.ALL)
    verify_runs(runs, ["second"])
    runs = MlflowClient().search_runs([experiment_id], "tags.t2 = 'second-tag-val'", ViewType.ALL)
    verify_runs(runs, ["second"])

    # delete "first" run
    MlflowClient().delete_run(logged_runs["first"])
    runs = MlflowClient().search_runs([experiment_id], "params.p1 = 'a'", ViewType.ALL)
    verify_runs(runs, ["first"])

    runs = MlflowClient().search_runs([experiment_id], "params.p1 = 'a'", ViewType.DELETED_ONLY)
    verify_runs(runs, ["first"])

    runs = MlflowClient().search_runs([experiment_id], "params.p1 = 'a'", ViewType.ACTIVE_ONLY)
    verify_runs(runs, [])


@pytest.mark.usefixtures("reset_active_experiment")
def test_search_runs_multiple_experiments():
    experiment_ids = [mlflow.create_experiment(f"exp__{exp_id}") for exp_id in range(1, 4)]
    for eid in experiment_ids:
        with mlflow.start_run(experiment_id=eid):
            mlflow.log_metric("m0", 1)
            mlflow.log_metric(f"m_{eid}", 2)

    assert len(MlflowClient().search_runs(experiment_ids, "metrics.m0 > 0", ViewType.ALL)) == 3

    assert len(MlflowClient().search_runs(experiment_ids, "metrics.m_1 > 0", ViewType.ALL)) == 1
    assert len(MlflowClient().search_runs(experiment_ids, "metrics.m_2 = 2", ViewType.ALL)) == 1
    assert len(MlflowClient().search_runs(experiment_ids, "metrics.m_3 < 4", ViewType.ALL)) == 1


def read_data(artifact_path):
    import pandas as pd

    if artifact_path.endswith(".json"):
        return pd.read_json(artifact_path, orient="split")
    if artifact_path.endswith(".parquet"):
        return pd.read_parquet(artifact_path)
    raise ValueError(f"Unsupported file type in {artifact_path}. Expected .json or .parquet")


@pytest.mark.skipif(
    "MLFLOW_SKINNY" in os.environ,
    reason="Skinny client does not support the np or pandas dependencies",
)
@pytest.mark.parametrize("file_type", ["json", "parquet"])
def test_log_table(file_type):
    import pandas as pd

    table_dict = {
        "inputs": ["What is MLflow?", "What is Databricks?"],
        "outputs": ["MLflow is ...", "Databricks is ..."],
        "toxicity": [0.0, 0.0],
    }
    artifact_file = f"qabot_eval_results.{file_type}"
    TAG_NAME = "mlflow.loggedArtifacts"
    run_id = None

    with pytest.raises(
        MlflowException, match="data must be a pandas.DataFrame or a dictionary"
    ) as e:
        with mlflow.start_run() as run:
            # Log the incorrect data format as a table
            mlflow.log_table(data="incorrect-data-format", artifact_file=artifact_file)
    assert e.value.error_code == ErrorCode.Name(INVALID_PARAMETER_VALUE)

    with mlflow.start_run() as run:
        # Log the dictionary as a table
        mlflow.log_table(data=table_dict, artifact_file=artifact_file)
        run_id = run.info.run_id

    run = mlflow.get_run(run_id)
    artifact_path = mlflow.artifacts.download_artifacts(run_id=run_id, artifact_path=artifact_file)
    table_data = read_data(artifact_path)
    assert table_data.shape[0] == 2
    assert table_data.shape[1] == 3

    # Get the current value of the tag
    current_tag_value = json.loads(run.data.tags.get(TAG_NAME, "[]"))
    assert {"path": artifact_file, "type": "table"} in current_tag_value
    assert len(current_tag_value) == 1

    table_df = pd.DataFrame.from_dict(table_dict)
    with mlflow.start_run(run_id=run_id):
        # Log the dataframe as a table
        mlflow.log_table(data=table_df, artifact_file=artifact_file)

    run = mlflow.get_run(run_id)
    artifact_path = mlflow.artifacts.download_artifacts(run_id=run_id, artifact_path=artifact_file)
    table_data = read_data(artifact_path)
    assert table_data.shape[0] == 4
    assert table_data.shape[1] == 3
    # Get the current value of the tag
    current_tag_value = json.loads(run.data.tags.get(TAG_NAME, "[]"))
    assert {"path": artifact_file, "type": "table"} in current_tag_value
    assert len(current_tag_value) == 1

    artifact_file_new = f"qabot_eval_results_new.{file_type}"
    with mlflow.start_run(run_id=run_id):
        # Log the dataframe as a table to new artifact file
        mlflow.log_table(data=table_df, artifact_file=artifact_file_new)

    run = mlflow.get_run(run_id)
    artifact_path = mlflow.artifacts.download_artifacts(
        run_id=run_id, artifact_path=artifact_file_new
    )
    table_data = read_data(artifact_path)
    assert table_data.shape[0] == 2
    assert table_data.shape[1] == 3
    # Get the current value of the tag
    current_tag_value = json.loads(run.data.tags.get(TAG_NAME, "[]"))
    assert {"path": artifact_file_new, "type": "table"} in current_tag_value
    assert len(current_tag_value) == 2


@pytest.mark.skipif(
    "MLFLOW_SKINNY" in os.environ,
    reason="Skinny client does not support the np or pandas dependencies",
)
@pytest.mark.parametrize("file_type", ["json", "parquet"])
def test_log_table_with_subdirectory(file_type):
    import pandas as pd

    table_dict = {
        "inputs": ["What is MLflow?", "What is Databricks?"],
        "outputs": ["MLflow is ...", "Databricks is ..."],
        "toxicity": [0.0, 0.0],
    }
    artifact_file = f"dir/foo.{file_type}"
    TAG_NAME = "mlflow.loggedArtifacts"
    run_id = None

    with mlflow.start_run() as run:
        # Log the dictionary as a table
        mlflow.log_table(data=table_dict, artifact_file=artifact_file)
        run_id = run.info.run_id

    run = mlflow.get_run(run_id)
    artifact_path = mlflow.artifacts.download_artifacts(run_id=run_id, artifact_path=artifact_file)
    table_data = read_data(artifact_path)
    assert table_data.shape[0] == 2
    assert table_data.shape[1] == 3

    # Get the current value of the tag
    current_tag_value = json.loads(run.data.tags.get(TAG_NAME, "[]"))
    assert {"path": artifact_file, "type": "table"} in current_tag_value
    assert len(current_tag_value) == 1

    table_df = pd.DataFrame.from_dict(table_dict)
    with mlflow.start_run(run_id=run_id):
        # Log the dataframe as a table
        mlflow.log_table(data=table_df, artifact_file=artifact_file)

    run = mlflow.get_run(run_id)
    artifact_path = mlflow.artifacts.download_artifacts(run_id=run_id, artifact_path=artifact_file)
    table_data = read_data(artifact_path)
    assert table_data.shape[0] == 4
    assert table_data.shape[1] == 3
    # Get the current value of the tag
    current_tag_value = json.loads(run.data.tags.get(TAG_NAME, "[]"))
    assert {"path": artifact_file, "type": "table"} in current_tag_value
    assert len(current_tag_value) == 1


@pytest.mark.skipif(
    "MLFLOW_SKINNY" in os.environ,
    reason="Skinny client does not support the np or pandas dependencies",
)
@pytest.mark.parametrize("file_type", ["json", "parquet"])
def test_load_table(file_type):
    table_dict = {
        "inputs": ["What is MLflow?", "What is Databricks?"],
        "outputs": ["MLflow is ...", "Databricks is ..."],
        "toxicity": [0.0, 0.0],
    }
    artifact_file = f"qabot_eval_results.{file_type}"
    artifact_file_2 = f"qabot_eval_results_2.{file_type}"
    run_id_2 = None

    with mlflow.start_run() as run:
        # Log the dictionary as a table
        mlflow.log_table(data=table_dict, artifact_file=artifact_file)
        mlflow.log_table(data=table_dict, artifact_file=artifact_file_2)

    with mlflow.start_run() as run:
        # Log the dictionary as a table
        mlflow.log_table(data=table_dict, artifact_file=artifact_file)
        run_id_2 = run.info.run_id

    with mlflow.start_run() as run:
        # Log the dictionary as a table
        mlflow.log_table(data=table_dict, artifact_file=artifact_file)
        run_id_3 = run.info.run_id

    extra_columns = ["run_id", "tags.mlflow.loggedArtifacts"]

    # test 1: load table with extra columns
    output_df = mlflow.load_table(artifact_file=artifact_file, extra_columns=extra_columns)

    assert output_df.shape[0] == 6
    assert output_df.shape[1] == 5
    assert output_df["run_id"].nunique() == 3
    assert output_df["tags.mlflow.loggedArtifacts"].nunique() == 2

    # test 2: load table with extra columns and single run_id
    output_df = mlflow.load_table(
        artifact_file=artifact_file, run_ids=[run_id_2], extra_columns=extra_columns
    )

    assert output_df.shape[0] == 2
    assert output_df.shape[1] == 5
    assert output_df["run_id"].nunique() == 1
    assert output_df["tags.mlflow.loggedArtifacts"].nunique() == 1

    # test 3: load table with extra columns and multiple run_ids
    output_df = mlflow.load_table(
        artifact_file=artifact_file,
        run_ids=[run_id_2, run_id_3],
        extra_columns=extra_columns,
    )

    assert output_df.shape[0] == 4
    assert output_df.shape[1] == 5
    assert output_df["run_id"].nunique() == 2
    assert output_df["tags.mlflow.loggedArtifacts"].nunique() == 1

    # test 4: load table with no extra columns and run_ids specified but different artifact file
    output_df = mlflow.load_table(artifact_file=artifact_file_2)
    import pandas as pd

    pd.testing.assert_frame_equal(output_df, pd.DataFrame(table_dict), check_dtype=False)

    # test 5: load table with no extra columns and run_ids specified
    output_df = mlflow.load_table(artifact_file=artifact_file)

    assert output_df.shape[0] == 6
    assert output_df.shape[1] == 3

    # test 6: load table with no matching results found. Error case
    with pytest.raises(
        MlflowException, match="No runs found with the corresponding table artifact"
    ):
        mlflow.load_table(artifact_file=f"error_case.{file_type}")

    # test 7: load table with no matching extra_column found. Error case
    with pytest.raises(KeyError, match="error_column"):
        mlflow.load_table(artifact_file=artifact_file, extra_columns=["error_column"])


@pytest.mark.skipif(
    "MLFLOW_SKINNY" in os.environ,
    reason="Skinny client does not support the np or pandas dependencies",
)
@pytest.mark.parametrize("file_type", ["json", "parquet"])
def test_log_table_with_datetime_columns(file_type):
    import pandas as pd

    start_time = str(datetime.now(timezone.utc))
    table_dict = {
        "inputs": ["What is MLflow?", "What is Databricks?"],
        "outputs": ["MLflow is ...", "Databricks is ..."],
        "start_time": [start_time, start_time],
    }
    artifact_file = f"test_time.{file_type}"

    with mlflow.start_run() as run:
        # Log the dictionary as a table
        mlflow.log_table(data=table_dict, artifact_file=artifact_file)
        run_id = run.info.run_id

    artifact_path = mlflow.artifacts.download_artifacts(run_id=run_id, artifact_path=artifact_file)
    if file_type == "parquet":
        table_data = pd.read_parquet(artifact_path)
    else:
        table_data = pd.read_json(artifact_path, orient="split", convert_dates=False)
    assert table_data["start_time"][0] == start_time

    # append the same table to the same artifact file
    mlflow.log_table(data=table_dict, artifact_file=artifact_file, run_id=run_id)
    artifact_path = mlflow.artifacts.download_artifacts(run_id=run_id, artifact_path=artifact_file)
    if file_type == "parquet":
        df = pd.read_parquet(artifact_path)
    else:
        df = pd.read_json(artifact_path, orient="split", convert_dates=False)
    assert df["start_time"][2] == start_time


@pytest.mark.skipif(
    "MLFLOW_SKINNY" in os.environ,
    reason="Skinny client does not support the np or pandas dependencies",
)
@pytest.mark.parametrize("file_type", ["json", "parquet"])
def test_log_table_with_image_columns(file_type):
    import numpy as np
    from PIL import Image

    image = mlflow.Image([[1, 2, 3]])
    table_dict = {
        "inputs": ["What is MLflow?", "What is Databricks?"],
        "outputs": ["MLflow is ...", "Databricks is ..."],
        "image": [image, image],
    }
    artifact_file = f"test_time.{file_type}"

    with mlflow.start_run() as run:
        # Log the dictionary as a table
        mlflow.log_table(data=table_dict, artifact_file=artifact_file)
        run_id = run.info.run_id

    artifact_path = mlflow.artifacts.download_artifacts(run_id=run_id, artifact_path=artifact_file)
    table_data = read_data(artifact_path)
    assert table_data["image"][0]["type"] == "image"
    image_path = mlflow.artifacts.download_artifacts(
        run_id=run_id, artifact_path=table_data["image"][0]["filepath"]
    )
    image2 = Image.open(image_path)
    assert np.abs(image.to_array() - np.array(image2)).sum() == 0

    # append the same table to the same artifact file
    mlflow.log_table(data=table_dict, artifact_file=artifact_file, run_id=run_id)
    artifact_path = mlflow.artifacts.download_artifacts(run_id=run_id, artifact_path=artifact_file)
    df = read_data(artifact_path)
    assert df["image"][2]["type"] == "image"


@pytest.mark.skipif(
    "MLFLOW_SKINNY" in os.environ,
    reason="Skinny client does not support the np or pandas dependencies",
)
@pytest.mark.parametrize("file_type", ["json", "parquet"])
def test_log_table_with_pil_image_columns(file_type):
    import numpy as np
    from PIL import Image

    image = Image.fromarray(np.array([[1.0, 2.0, 3.0]]))
    image = image.convert("RGB")

    table_dict = {
        "inputs": ["What is MLflow?", "What is Databricks?"],
        "outputs": ["MLflow is ...", "Databricks is ..."],
        "image": [image, image],
    }
    artifact_file = f"test_time.{file_type}"

    with mlflow.start_run() as run:
        # Log the dictionary as a table
        mlflow.log_table(data=table_dict, artifact_file=artifact_file)
        run_id = run.info.run_id

    artifact_path = mlflow.artifacts.download_artifacts(run_id=run_id, artifact_path=artifact_file)
    table_data = read_data(artifact_path)
    assert table_data["image"][0]["type"] == "image"
    image_path = mlflow.artifacts.download_artifacts(
        run_id=run_id, artifact_path=table_data["image"][0]["filepath"]
    )
    image2 = Image.open(image_path)
    assert np.abs(np.array(image) - np.array(image2)).sum() == 0

    # append the same table to the same artifact file
    mlflow.log_table(data=table_dict, artifact_file=artifact_file, run_id=run_id)
    artifact_path = mlflow.artifacts.download_artifacts(run_id=run_id, artifact_path=artifact_file)
    df = read_data(artifact_path)
    assert df["image"][2]["type"] == "image"


@pytest.mark.skipif(
    "MLFLOW_SKINNY" in os.environ,
    reason="Skinny client does not support the np or pandas dependencies",
)
@pytest.mark.parametrize("file_type", ["json", "parquet"])
def test_log_table_with_invalid_image_columns(file_type):
    image = mlflow.Image([[1, 2, 3]])
    table_dict = {
        "inputs": ["What is MLflow?", "What is Databricks?"],
        "outputs": ["MLflow is ...", "Databricks is ..."],
        "image": [image, "text"],
    }
    artifact_file = f"test_time.{file_type}"
    with pytest.raises(ValueError, match="Column `image` contains a mix of images and non-images"):
        with mlflow.start_run():
            # Log the dictionary as a table
            mlflow.log_table(data=table_dict, artifact_file=artifact_file)


@pytest.mark.skipif(
    "MLFLOW_SKINNY" in os.environ,
    reason="Skinny client does not support the np or pandas dependencies",
)
@pytest.mark.parametrize("file_type", ["json", "parquet"])
def test_log_table_with_valid_image_columns(file_type):
    class ImageObj:
        def __init__(self):
            self.size = (1, 1)

        def resize(self, size):
            return self

        def save(self, path):
            with open(path, "w+") as f:
                f.write("dummy data")

    image_obj = ImageObj()
    image = mlflow.Image([[1, 2, 3]])

    table_dict = {
        "inputs": ["What is MLflow?", "What is Databricks?"],
        "outputs": ["MLflow is ...", "Databricks is ..."],
        "image": [image, image_obj],
    }
    # No error should be raised
    artifact_file = f"test_time.{file_type}"
    with mlflow.start_run():
        # Log the dictionary as a table
        mlflow.log_table(data=table_dict, artifact_file=artifact_file)


def test_set_async_logging_threadpool_size():
    MLFLOW_ASYNC_LOGGING_THREADPOOL_SIZE.set(6)
    assert MLFLOW_ASYNC_LOGGING_THREADPOOL_SIZE.get() == 6

    with mlflow.start_run():
        mlflow.log_param("key", "val", synchronous=False)

    store = mlflow.tracking._get_store()
    async_queue = store._async_logging_queue
    assert async_queue._batch_logging_worker_threadpool._max_workers == 6
    mlflow.flush_async_logging()
    MLFLOW_ASYNC_LOGGING_THREADPOOL_SIZE.unset()