wehub-resource-sync
52 lines
1.8 KiB
Python
52 lines
1.8 KiB
Python
import mlflow
|
|
from mlflow.tracing.attachments import Attachment
|
|
|
|
|
|
def test_attachment_roundtrip_with_local_tracking():
|
|
image_bytes = b"\x89PNG\r\n\x1a\n fake png content"
|
|
audio_bytes = b"RIFF fake wav content"
|
|
|
|
with mlflow.start_span(name="integration-span") as span:
|
|
span.set_inputs({
|
|
"prompt": "describe this",
|
|
"image": Attachment(content_type="image/png", content_bytes=image_bytes),
|
|
})
|
|
span.set_outputs({
|
|
"audio": Attachment(content_type="audio/wav", content_bytes=audio_bytes),
|
|
"text": "a cat",
|
|
})
|
|
trace_id = span.trace_id
|
|
|
|
# Retrieve the trace and verify reference URIs are stored
|
|
mlflow.flush_trace_async_logging()
|
|
trace = mlflow.get_trace(trace_id)
|
|
root_span = trace.data.spans[0]
|
|
|
|
assert root_span.inputs["prompt"] == "describe this"
|
|
assert root_span.outputs["text"] == "a cat"
|
|
|
|
image_ref = root_span.inputs["image"]
|
|
audio_ref = root_span.outputs["audio"]
|
|
|
|
image_parsed = Attachment.parse_ref(image_ref)
|
|
audio_parsed = Attachment.parse_ref(audio_ref)
|
|
|
|
assert image_parsed["content_type"] == "image/png"
|
|
assert image_parsed["trace_id"] == trace_id
|
|
assert audio_parsed["content_type"] == "audio/wav"
|
|
assert audio_parsed["trace_id"] == trace_id
|
|
|
|
# Verify the attachment files were written to the artifact repo
|
|
from mlflow.tracing.client import TracingClient
|
|
|
|
tracking_uri = mlflow.get_tracking_uri()
|
|
client = TracingClient(tracking_uri)
|
|
trace_info = client.get_trace_info(trace_id)
|
|
artifact_repo = client._get_artifact_repo_for_trace(trace_info)
|
|
|
|
stored_image = artifact_repo.download_trace_attachment(image_parsed["attachment_id"])
|
|
stored_audio = artifact_repo.download_trace_attachment(audio_parsed["attachment_id"])
|
|
|
|
assert stored_image == image_bytes
|
|
assert stored_audio == audio_bytes
|