Exception: mlflow.exceptions.MlflowException: An API request to https://canada.cloud.databricks.com/api/2.0/mlflow/model-versions/list-artifacts failed due to a timeout. The error message was: HTTPSConnectionPool(host='canada.cloud.databricks.com', port=443): Max retries exceeded with url: /api/2.0/mlflow/model-versions/list-artifacts.

Expected: The request should be directed to the current workspace URI instead of cloud.databricks.com.

I am encountering a strange issue while executing my code within a workflow. I am attempting to load an MLflow model registered in the Unity Catalog using mlflow.pyfunc.spark_udf(). I have ensured that the model URI and other parameters are correct. This code is part of a Python class.

Additionally, I am setting the registry URI using mlflow.set_registry_uri("databricks-uc"). Below are the environment variables I am configuring at the compute level:

DATABRICKS_HOST="<workspace_url>" DATABRICKS_TOKEN="<PAT>"

Spark configuration:

spark.executorEnv.DATABRICKS_HOST="<workspace_url>" spark.executorEnv.DATABRICKS_TOKEN="<PAT>"

Databricks runtime version: 14.3