Cannot get tracing to work on genai app deployed on databricks

[MisterT]

Hi, I have a gradio app that is deployed on databricks. The app is coming from this example  provided by databricks. The app works fine, but when I want to add tracing I cannot get it to work. I keep getting the error

mlflow.exceptions.MlflowException: The provided MLFLOW_EXPERIMENT_ID environment variable value `--redacted-experiment-id` does not exist in the tracking server. Provide a valid experiment_id.

I'm setting the following environment variables: 

{"name": "MODEL_SERVING_ENDPOINT", "value": endpoint_name},

{"name": "MLFLOW_TRACKING_URI", "value": "databricks"},

{"name": "MLFLOW_REGISTRY_URI", "value": "databricks-uc"},

{"name": "MLFLOW_EXPERIMENT_ID", "value": experiment_id_traces}

 

In the gradio code I've added the following lines 

mlflow.openai.autolog()

@mlflow.trace(name="genai_gradio_interaction", span_type=SpanType.CHAT_MODEL)
def respond(message, history, dropdown😞

  if len(message.strip()) == 0:
  return "ERROR the question should not be empty"

  try: 
    messages = []
    if history:
     for human, assistant in history:
      messages.append(ChatMessage(content=human, role=ChatMessageRole.USER))
      messages.append(ChatMessage(content=assistant, role=ChatMessageRole.ASSISTANT))
      messages.append(ChatMessage(content=message, role=ChatMessageRole.USER))

    response = w.serving_endpoints.query(name=dropdown,messages=messages,temperature=0.2,stream=False)

    generated_text = response.choices[0].message.content

  except Exception as error:

    return f"ERROR requesting endpoint {dropdown}: {error}"

  return generated_text
# Note indents could be wrong 

My imports are

mlflow[databricks]>=3.1.0
databricks-sdk[openai]

The experimentId is coming from a experiment created in the UI. I don't understand why it cannot find the experiment. I also tried using the experiment name, but I get the same error message. What is the recommend way to get traces working in apps running on databricks? 

[NandiniN]

Hi @MisterT , 

In our docs, it is mentioned we use MLflow 3(major upgrade) with GenAI monitoring enabled. Each agent endpoint is assigned an MLflow experiment, and log agent traces from the endpoint to that experiment in real-time.

Internally an  MLFLOW_EXPERIMENT_ID environment variable is set to the endpoint’s experiment. Traces are automatically configured to be written to the Databricks MLflow tracking server. But it is failing because the agent serving endpoint itself is not configured to log metrics/params to Databricks. To resolve this issue, you can try any of the following options,

1) While calling agents.deploy(<existing_params>, environment_vars={"MLFLOW_TRACKING_URI":"databricks"}) to configure the tracking URI to databricks and resolve this, or

2) Remove the log_param call from your agent code.

Also referring https://docs.databricks.com/aws/en/mlflow3/genai/getting-started/connect-environment

Thanks!