Configure access to resources from model serving endpoints
This notebook example demonstrates how you can securely store credentials in a Databricks secrets scope and reference those secrets in model serving. This allows credentials to be fetched at serving time from model serving endpoints.
Requirements
- This functionality currently is only supported via the Databricks REST API.
- To use this feature, you must store credentials like your API key or other tokens as a Databricks secret.
- The endpoint creator must have
Read
access to the Databricks secrets being referenced in the endpoint configuration. - This notebook requireds Databricks SDK version 0.6.0 and above.
Step 0: Install and upgrade any dependencies if necessary
Step 1: Add secrets to Databricks Secret Store
You can modify the following variables to assign your secret and its corresponding key and values.
Step 2: Upload sample document data into DBFS
- Download sample document data locally from https://github.com/mlflow/mlflow/blob/master/tests/langchain/state_of_the_union.txt.
- Click Catalog > Browse DBFS
- Upload the file into a FileStore location in DBFS
Step 3: Log and register the LangChain model
This example model is adapted from https://github.com/mlflow/mlflow/blob/master/examples/langchain/retrieval_qa_chain.py.
In this section, you can create a vector database and persist that database to a local file store folder. You also create a RetrievalQA chain and log it start your model run.
Next, you can load that logged model using MLflow pyfunc
.
Step 4: Create and query the serving endpoint
In this section you create a serving endpoint to serve your model and query it.
You can use the following code to check the endpoint status to verify it is ready.
Finally, you can query the endpoint with sample data.