model_with_add_trace.py
1 import mlflow 2 3 _SAMPLE_TRACE = { 4 "info": { 5 "request_id": "2e72d64369624e6888324462b62dc120", 6 "experiment_id": "0", 7 "timestamp_ms": 1726145090860, 8 "execution_time_ms": 162, 9 "status": "OK", 10 "request_metadata": { 11 "mlflow.trace_schema.version": "2", 12 "mlflow.traceInputs": '{"x": 1}', 13 "mlflow.traceOutputs": '{"prediction": 1}', 14 }, 15 "tags": { 16 "fruit": "apple", 17 "food": "pizza", 18 }, 19 }, 20 "data": { 21 "spans": [ 22 { 23 "name": "remote", 24 "context": { 25 "span_id": "0x337af925d6629c01", 26 "trace_id": "0x05e82d1fc4486f3986fae6dd7b5352b1", 27 }, 28 "parent_id": None, 29 "start_time": 1726145091022155863, 30 "end_time": 1726145091022572053, 31 "status_code": "OK", 32 "status_message": "", 33 "attributes": { 34 "mlflow.traceRequestId": '"2e72d64369624e6888324462b62dc120"', 35 "mlflow.spanType": '"UNKNOWN"', 36 "mlflow.spanInputs": '{"x": 1}', 37 "mlflow.spanOutputs": '{"prediction": 1}', 38 }, 39 "events": [ 40 {"name": "event", "timestamp": 1726145091022287, "attributes": {"foo": "bar"}} 41 ], 42 }, 43 ], 44 "request": '{"x": 1}', 45 "response": '{"prediction": 1}', 46 }, 47 } 48 49 50 class Model(mlflow.pyfunc.PythonModel): 51 def predict(self, context, model_input): 52 mlflow.add_trace(_SAMPLE_TRACE) 53 return 1 54 55 56 mlflow.models.set_model(Model())