wehub-resource-sync
123 lines
3.4 KiB
Python
123 lines
3.4 KiB
Python
import random
|
|
import subprocess
|
|
import tempfile
|
|
import time
|
|
from unittest import mock
|
|
|
|
import pytest
|
|
|
|
import mlflow
|
|
from mlflow.environment_variables import (
|
|
MLFLOW_ENABLE_ASYNC_LOGGING,
|
|
MLFLOW_ENABLE_ASYNC_TRACE_LOGGING,
|
|
)
|
|
from mlflow.tracing.fluent import _flush_pending_async_trace_writes
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def enable_async_trace_logging(monkeypatch):
|
|
"""Enable async trace logging for all tests in tests/tracing/ to exercise the async path.
|
|
|
|
Overrides the global disable_async_trace_logging fixture from tests/conftest.py.
|
|
Tests that need both sync and async coverage use the async_logging_enabled fixture.
|
|
Terminates async queues on teardown to prevent thread leaks between tests.
|
|
"""
|
|
monkeypatch.setenv(MLFLOW_ENABLE_ASYNC_TRACE_LOGGING.name, "true")
|
|
monkeypatch.setenv(MLFLOW_ENABLE_ASYNC_LOGGING.name, "true")
|
|
|
|
yield
|
|
|
|
_flush_pending_async_trace_writes(terminate=True)
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def reset_active_experiment():
|
|
yield
|
|
mlflow.tracking.fluent._active_experiment_id = None
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def reset_tracking_uri():
|
|
# Some API like set_destination("databricks") updates the tracking URI,
|
|
# we should reset it between tests
|
|
original_tracking_uri = mlflow.get_tracking_uri()
|
|
|
|
yield
|
|
|
|
mlflow.set_tracking_uri(original_tracking_uri)
|
|
|
|
|
|
@pytest.fixture
|
|
def databricks_tracking_uri():
|
|
with mock.patch("mlflow.get_tracking_uri", return_value="databricks"):
|
|
yield
|
|
|
|
|
|
# Fixture to run the test case with and without async logging enabled.
|
|
# When async logging is enabled, the batch span processor is also active (the default),
|
|
# so tests exercise the full production pipeline.
|
|
@pytest.fixture(params=[True, False])
|
|
def async_logging_enabled(request, monkeypatch):
|
|
monkeypatch.setenv(MLFLOW_ENABLE_ASYNC_TRACE_LOGGING.name, str(request.param))
|
|
# TODO: V2 Trace depends on this env var rather than MLFLOW_ENABLE_ASYNC_TRACE_LOGGING
|
|
# We should remove this once the backend is fully migrated to V3
|
|
monkeypatch.setenv(MLFLOW_ENABLE_ASYNC_LOGGING.name, str(request.param))
|
|
return request.param
|
|
|
|
|
|
@pytest.fixture
|
|
def otel_collector():
|
|
"""Start an OpenTelemetry collector in a Docker container."""
|
|
subprocess.check_call(["docker", "pull", "otel/opentelemetry-collector"])
|
|
|
|
# Use a random port to avoid conflicts
|
|
port = random.randint(20000, 30000)
|
|
|
|
docker_collector_config = """receivers:
|
|
otlp:
|
|
protocols:
|
|
grpc:
|
|
endpoint: 0.0.0.0:4317
|
|
|
|
exporters:
|
|
debug:
|
|
verbosity: detailed
|
|
sampling_initial: 5
|
|
sampling_thereafter: 1
|
|
|
|
service:
|
|
pipelines:
|
|
traces:
|
|
receivers: [otlp]
|
|
exporters: [debug]"""
|
|
|
|
with tempfile.NamedTemporaryFile() as output_file:
|
|
# Use echo to pipe config to Docker stdin
|
|
docker_cmd = [
|
|
"bash",
|
|
"-c",
|
|
f'echo "{docker_collector_config}" | '
|
|
f"docker run --rm -p 127.0.0.1:{port}:4317 -i "
|
|
f"otel/opentelemetry-collector --config=/dev/stdin",
|
|
]
|
|
|
|
process = subprocess.Popen(
|
|
docker_cmd,
|
|
stdout=output_file,
|
|
stderr=subprocess.STDOUT,
|
|
text=True,
|
|
)
|
|
|
|
# Wait for the collector to start
|
|
time.sleep(5)
|
|
|
|
yield process, output_file.name, port
|
|
|
|
# Stop the collector
|
|
process.terminate()
|
|
try:
|
|
process.wait(timeout=5)
|
|
except subprocess.TimeoutExpired:
|
|
process.kill()
|
|
process.wait()
|