/ tests / tracing / sample_code / model_with_add_trace.py
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())