wehub-resource-sync
124 lines
3.6 KiB
Python
124 lines
3.6 KiB
Python
import time
|
|
|
|
import pytest
|
|
|
|
import mlflow
|
|
from mlflow.tracing.utils.copy import copy_trace_to_experiment
|
|
|
|
from tests.tracing.helper import purge_traces
|
|
|
|
|
|
def _create_test_span_dict(request_id="test-trace", parent_id=None):
|
|
"""Helper to create a minimal valid span dict for testing"""
|
|
return {
|
|
"name": "root_span" if parent_id is None else "child_span",
|
|
"context": {
|
|
"span_id": "0d48a6670588966b" if parent_id is None else "6fc32f36ef591f60",
|
|
"trace_id": "63076d0c1b90f1df0970f897dc428bd6",
|
|
},
|
|
"parent_id": parent_id,
|
|
"start_time": 100,
|
|
"end_time": 200,
|
|
"status_code": "OK",
|
|
"status_message": "",
|
|
"attributes": {
|
|
"mlflow.traceRequestId": f'"{request_id}"',
|
|
"mlflow.spanType": '"UNKNOWN"',
|
|
},
|
|
"events": [],
|
|
}
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def setup_experiment():
|
|
"""Set up a test experiment before each test"""
|
|
exp = mlflow.set_experiment(f"test_copy_trace_{time.time()}")
|
|
yield exp
|
|
purge_traces(exp.experiment_id)
|
|
|
|
|
|
def test_copy_trace_with_metadata():
|
|
trace_dict = {
|
|
"info": {
|
|
"request_id": "test-trace-789",
|
|
"experiment_id": "0",
|
|
"timestamp_ms": 100,
|
|
"execution_time_ms": 200,
|
|
"status": "OK",
|
|
"trace_metadata": {
|
|
"mlflow.trace.session": "session123",
|
|
"custom.metadata": "metadata_value",
|
|
"user.key": "user_value",
|
|
},
|
|
},
|
|
"data": {"spans": [_create_test_span_dict("test-trace-789")]},
|
|
}
|
|
|
|
new_trace_id = copy_trace_to_experiment(trace_dict)
|
|
|
|
# Verify metadata was copied correctly
|
|
trace = mlflow.get_trace(new_trace_id, flush=True)
|
|
metadata = trace.info.trace_metadata
|
|
|
|
assert metadata["mlflow.trace.session"] == "session123"
|
|
assert metadata["custom.metadata"] == "metadata_value"
|
|
assert metadata["user.key"] == "user_value"
|
|
|
|
|
|
def test_copy_trace_missing_info():
|
|
trace_dict = {"data": {"spans": [_create_test_span_dict("test-trace-no-info")]}}
|
|
|
|
# Should not raise an error, just skip tag/metadata copying
|
|
new_trace_id = copy_trace_to_experiment(trace_dict)
|
|
|
|
assert new_trace_id is not None
|
|
trace = mlflow.get_trace(new_trace_id, flush=True)
|
|
assert trace is not None
|
|
|
|
|
|
def test_copy_trace_missing_metadata():
|
|
trace_dict = {
|
|
"info": {
|
|
"request_id": "test-trace-no-metadata",
|
|
"experiment_id": "0",
|
|
"tags": {
|
|
"user.tag": "tag_value",
|
|
},
|
|
},
|
|
"data": {"spans": [_create_test_span_dict("test-trace-no-metadata")]},
|
|
}
|
|
|
|
# Should not raise an error, just skip metadata copying
|
|
new_trace_id = copy_trace_to_experiment(trace_dict)
|
|
|
|
assert new_trace_id is not None
|
|
trace = mlflow.get_trace(new_trace_id, flush=True)
|
|
|
|
# Tags should still be copied
|
|
tags = trace.info.tags
|
|
assert tags["user.tag"] == "tag_value"
|
|
|
|
|
|
def test_copy_trace_empty_metadata_dict():
|
|
trace_dict = {
|
|
"info": {
|
|
"request_id": "test-trace-empty-metadata",
|
|
"experiment_id": "0",
|
|
"tags": {
|
|
"user.tag": "value",
|
|
},
|
|
"trace_metadata": {},
|
|
},
|
|
"data": {"spans": [_create_test_span_dict("test-trace-empty-metadata")]},
|
|
}
|
|
|
|
# Should not raise an error
|
|
new_trace_id = copy_trace_to_experiment(trace_dict)
|
|
|
|
assert new_trace_id is not None
|
|
trace = mlflow.get_trace(new_trace_id, flush=True)
|
|
|
|
# Tags should still be copied
|
|
tags = trace.info.tags
|
|
assert tags["user.tag"] == "value"
|