@Dharma25 

The error clearly shows that MLflow requests are being incorrectly routed to "canada.cloud.databricks.com" instead of your actual workspace URL, causing the timeout.
This is a known issue that can occur with MLflow in Databricks, particularly when working with Unity Catalog models. Here's a more targeted approach to fix it:

1. First, ensure your environment variables are correctly set before any MLflow operations:
import os
import mlflow

# Explicitly define the workspace URL (don't rely on environment variables that might not propagate)
workspace_url = "https://your-actual-workspace-url.cloud.databricks.com" # Replace with your actual URL

# Set both tracking and registry URIs explicitly
os.environ["DATABRICKS_HOST"] = workspace_url
mlflow.set_tracking_uri(workspace_url)
mlflow.set_registry_uri("databricks-uc")

2. When loading the model, you might need to be more explicit:
# Try with the full model URI format
model_uri = f"models://{model_name}/{model_version}"
model_udf = mlflow.pyfunc.spark_udf(spark, model_uri, env_manager="conda")

3. For Spark environments, also ensure the configuration is properly set:
spark.conf.set("spark.databricks.mlflow.trackingUri", workspace_url)

This issue is often related to how the MLflow client resolves endpoints when working with Unity Catalog,
where it sometimes falls back to a default domain instead of using the workspace-specific URL you've provided.

Thanks for the answer. I will try this solution 🙂