test-serverless-endpoint-example(Python)

Import Notebook

Test Serverless endpoint by querying your model

The notebook loads an input example that was logged with the registered model, ElasticNetDiabetes, and queries the model to test the servereless endpoint.

Prerequisites

  • Logged and registered Python-based model in the Model Registry
  • Model Serving endpoint (AWS|Azure)

Provide shard URL

shard_url = '<put-shard-url-here (e.g. https://company.cloud.databricks.com)>'
model_name = 'ElasticNetDiabetes'

The following assumes an input example was logged with your model. When possible, make sure to log an input example so that it is easy to construct a sample query to your model version.

You also need a Databricks token to issue requests to your model endpoint. You can replace <YOUR_TOKEN> with your token in the following example.

You can generate a token on the User Settings page of your Databricks workspace (AWS|Azure|GCP).

import os
import requests
import numpy as np
import pandas as pd

def create_tf_serving_json(data):
  return {'inputs': {name: data[name].tolist() for name in data.keys()} if isinstance(data, dict) else data.tolist()}

def process_input(dataset):
  if isinstance(dataset, pd.DataFrame):
    return {'dataframe_split': dataset.to_dict(orient='split')}
  elif isinstance(dataset, str):
    return dataset
  else:
    return create_tf_serving_json(dataset)

def score_model(dataset):
  url = f'{shard_url}/model-endpoint/{model_name}/1/invocations'
  databricks_token = <YOUR-TOKEN>
  headers = {'Authorization': f'Bearer {databricks_token}'}
  data_json = process_input(dataset)
  response = requests.request(method='POST', headers=headers, url=url, json=data_json)
  if response.status_code != 200:
    raise Exception(f'Request failed with status {response.status_code}, {response.text}')
  return response.json()

Call the score_model() function and pass your input example as the parameter

import mlflow
path = mlflow.artifacts.download_artifacts(f'models:/{model_name}/1')
model = mlflow.pyfunc.load_model(f'models:/{model_name}/1')
input_example = model.metadata.load_input_example(path)
score_model(input_example)