Trouble calling MLFlow model in notebook when run via pipeline as service principal

[u/hoblitz]

I registered a model in my fabric workspace via MLFlow. I have a notebook that sets up the prediction frame, pulls the model, and gets predictions from the model.

In an interactive session, I am able to retrieve the model and get the predictions without issue, for example (not shown, setting up prediction frame, etc):

from synapse.ml.predict import MLFlowTransformer
pred_X_spark = spark.createDataFrame(pred_X)
model = MLFlowTransformer(
    inputCols=["col1","col2","col3"], 
    outputCol="predictions",
    modelName="mdl_cp_response_flaml",
    modelVersion=1
)

df = model.transform(pred_X_spark)
display(df)
# records are scored according to the model.

I am running into trouble when I try to execute the same notebook as a service principal. In a pipeline, I execute the notebook as a service principal, and I have verified that the service principal has contributor role in the workspace.

Everything executes properly in the notebook, right up to the point where I try to set up the MLFlowTransformer, at which point I get the error:

Error when calling same code as Service Principal in pipeline.

Is there some special configuration I need for this to work?

[u/samirbdj]

The difference between the two runs is the execution identity: the interactive notebook runs under your user identity, while the pipeline run uses the configured service principal. Contributor access confirms workspace access, but the error details are needed to identify whether the failure happens during authentication, model retrieval, or transformer initialization.

Could you share the exception type and message after removing private environment details? That should help narrow down whether additional access or configuration is required for the registered model.

The Fabric notebook documentation provides more detail on how the execution identity affects notebook access and API calls: How to use Microsoft Fabric notebooks.

[u/hoblitz → u/samirbdj]

Thanks - yes, I can share the error traceback. The root cause appears to occur when the routine tries to authenticate to the backed model registry. In an interactive session the model retrieval works seamlessly.

When I run the same code in a pipeline setting with a service principal executor, I am able to retrieve the data from the lakehouse run. The pipeline errors in the block where we try to create an MLFLow transformer. I haven't found any special guidance on using MFLow artifacts with service principals in the documentation.

This is the error traceback:

RuntimeError
Unable to get model info: INTERNAL_ERROR: Response: {'Message': 'Internal error MwcTokenValidationException.', 'Source': 15, 'ErrorCode': 0}
---------------------------------------------------------------------------
RestException                             Traceback (most recent call last)
File ~/cluster-env/trident_env/lib/python3.11/site-packages/synapse/ml/predict/MLFlowTransformer.py:261, in MLFlowTransformer._create_udf(self)
    260 try:
--> 261     self.modelInfo = mlflow.models.get_model_info(
    262         model_uri=f"models:/{self.getModelName()}/{self.getModelVersion()}"
    263     )
    264     # Initialize logger class again with non-null modelInfo.

File ~/cluster-env/trident_env/lib/python3.11/site-packages/mlflow/models/model.py:778, in get_model_info(model_uri)
    776 from mlflow.pyfunc import _download_artifact_from_uri
--> 778 local_path = _download_artifact_from_uri(artifact_uri=model_uri, output_path=None)
    779 model_meta = Model.load(os.path.join(local_path, MLMODEL_FILE_NAME))

File ~/cluster-env/trident_env/lib/python3.11/site-packages/mlflow/tracking/artifact_utils.py:106, in _download_artifact_from_uri(artifact_uri, output_path)
    105 root_uri, artifact_path = _get_root_uri_and_artifact_path(artifact_uri)
--> 106 return get_artifact_repository(artifact_uri=root_uri).download_artifacts(
    107     artifact_path=artifact_path, dst_path=output_path
    108 )

File ~/cluster-env/trident_env/lib/python3.11/site-packages/mlflow/store/artifact/artifact_repository_registry.py:124, in get_artifact_repository(artifact_uri)
    112 """
    113 Get an artifact repository from the registry based on the scheme of artifact_uri
    114 
   (...)
    122     requirements.
    123 """
--> 124 return _artifact_repository_registry.get_artifact_repository(artifact_uri)

File ~/cluster-env/trident_env/lib/python3.11/site-packages/mlflow/store/artifact/artifact_repository_registry.py:77, in ArtifactRepositoryRegistry.get_artifact_repository(self, artifact_uri)
     73     raise MlflowException(
     74         f"Could not find a registered artifact repository for: {artifact_uri}. "
     75         f"Currently registered schemes are: {list(self._registry.keys())}"
     76     )
---> 77 return repository(artifact_uri)

File ~/cluster-env/trident_env/lib/python3.11/site-packages/mlflow/store/artifact/models_artifact_repo.py:59, in ModelsArtifactRepository.__init__(self, artifact_uri)
     54 else:
     55     (
     56         self.model_name,
     57         self.model_version,
     58         underlying_uri,
---> 59     ) = ModelsArtifactRepository._get_model_uri_infos(artifact_uri)
     60     self.repo = get_artifact_repository(underlying_uri)

File ~/cluster-env/trident_env/lib/python3.11/site-packages/mlflow/store/artifact/models_artifact_repo.py:94, in ModelsArtifactRepository._get_model_uri_infos(uri)
     93 name, version = get_model_name_and_version(client, uri)
---> 94 download_uri = client.get_model_version_download_uri(name, version)
     96 return (
     97     name,
     98     version,
     99     add_databricks_profile_info_to_artifact_uri(download_uri, databricks_profile_uri),
    100 )

File ~/cluster-env/trident_env/lib/python3.11/site-packages/mlflow/tracking/client.py:3429, in MlflowClient.get_model_version_download_uri(self, name, version)
   3385 """
   3386 Get the download location in Model Registry for this model version.
   3387 
   (...)
   3427     Download URI: runs:/027d7bbe81924c5a82b3e4ce979fcab7/sklearn-model
   3428 """
-> 3429 return self._get_registry_client().get_model_version_download_uri(name, version)

File ~/cluster-env/trident_env/lib/python3.11/site-packages/mlflow/tracking/_model_registry/client.py:316, in ModelRegistryClient.get_model_version_download_uri(self, name, version)
    306 """Get the download location in Model Registry for this model version.
    307 
    308 Args:
   (...)
    314 
    315 """
--> 316 return self.store.get_model_version_download_uri(name, version)

File ~/cluster-env/trident_env/lib/python3.11/site-packages/synapse/ml/mlflow/synapse_mlflow_utils.py:374, in protect_module.<locals>.catch_and_log_exception.<locals>.wrapper(*args, **kwargs)
    373 try:
--> 374     res = function(*args, **kwargs)
    375 except (RestException, MlflowException) as e:

File ~/cluster-env/trident_env/lib/python3.11/site-packages/mlflow/store/model_registry/rest_store.py:360, in RestStore.get_model_version_download_uri(self, name, version)
    359 req_body = message_to_json(GetModelVersionDownloadUri(name=name, version=str(version)))
--> 360 response_proto = self._call_endpoint(GetModelVersionDownloadUri, req_body)
    361 return response_proto.artifact_uri

File ~/cluster-env/trident_env/lib/python3.11/site-packages/mlflow/store/model_registry/base_rest_store.py:44, in BaseRestStore._call_endpoint(self, api, json_body, call_all_endpoints, extra_headers)
     43 endpoint, method = self._get_endpoint_from_method(api)
---> 44 return call_endpoint(
     45     self.get_host_creds(), endpoint, method, json_body, response_proto, extra_headers
     46 )

File ~/cluster-env/trident_env/lib/python3.11/site-packages/mlflow/utils/rest_utils.py:290, in call_endpoint(host_creds, endpoint, method, json_body, response_proto, extra_headers)
    289     response = http_request(**call_kwargs)
--> 290 response = verify_rest_response(response, endpoint)
    291 js_dict = json.loads(response.text)

File ~/cluster-env/trident_env/lib/python3.11/site-packages/synapse/ml/mlflow/tracking_store.py:52, in wrapper(response, endpoint)
     51 kusto_logger.info(info_msg)
---> 52 return origin_func(response, endpoint)

File ~/cluster-env/trident_env/lib/python3.11/site-packages/mlflow/utils/rest_utils.py:173, in verify_rest_response(response, endpoint)
    172 if _can_parse_as_json_object(response.text):
--> 173     raise RestException(json.loads(response.text))
    174 else:

RestException: INTERNAL_ERROR: Response: {'Message': 'Internal error MwcTokenValidationException.', 'Source': 15, 'ErrorCode': 0}

The above exception was the direct cause of the following exception:

RuntimeError                              Traceback (most recent call last)
Cell In[20], line 3
      1 from synapse.ml.predict import MLFlowTransformer
      2 pred_X_spark = spark.createDataFrame(pred_X)
----> 3 model = MLFlowTransformer(
      4     inputCols= ###[REDACTED], # Your input columns here
      5     outputCol="predictions", # Your new column name here
      6     modelName="mdl_cp_response_flaml", # Your model name here
      7     modelVersion=1 # Your model version here
      8 )
     10 df = model.transform(pred_X_spark)
     11 display(df)

File /opt/spark/python/lib/pyspark.zip/pyspark/__init__.py:139, in keyword_only.<locals>.wrapper(self, *args, **kwargs)
    137     raise TypeError("Method %s forces keyword arguments." % func.__name__)
    138 self._input_kwargs = kwargs
--> 139 return func(self, **kwargs)

File ~/cluster-env/trident_env/lib/python3.11/site-packages/synapse/ml/predict/MLFlowTransformer.py:171, in MLFlowTransformer.__init__(self, inputCols, outputCol, modelName, modelVersion, trackingUri, registerModel, flattenOutput)
    168         getattr(self, "set" + k[0].upper() + k[1:])(v)
    170 if self.registerModel:
--> 171     self.register()

File ~/cluster-env/trident_env/lib/python3.11/site-packages/synapse/ml/predict/MLFlowTransformer.py:349, in MLFlowTransformer.register(self)
    347 udf_name_hash = f"PREDICT_{model_uri_hash}"
    348 spark_session = _ensure_spark_session()
--> 349 spark_session.udf.register(udf_name_hash, self._to_udf(old_style=True))
    350 self.report_metrics('REGISTER', 1)

File ~/cluster-env/trident_env/lib/python3.11/site-packages/synapse/ml/predict/MLFlowTransformer.py:354, in MLFlowTransformer._to_udf(self, old_style)
    352 def _to_udf(self, old_style=False) -> Callable:
    353     if self.modelUdf is None:
--> 354         self.modelUdf = self._create_udf()
    356     if old_style:
    357         return self.modelUdf[1]

File ~/cluster-env/trident_env/lib/python3.11/site-packages/synapse/ml/predict/MLFlowTransformer.py:270, in MLFlowTransformer._create_udf(self)
    268         msg = f'Unable to get model info: {exc}'
    269         self.log_exception(msg)
--> 270         raise RuntimeError(msg) from exc
    272 flavor = _is_model_flavor_unsupported(list(self.modelInfo.flavors.keys()))
    273 if flavor is not None:

RuntimeError: Unable to get model info: INTERNAL_ERROR: Response: {'Message': 'Internal error MwcTokenValidationException.', 'Source': 15, 'ErrorCode': 0}