https://github.com/dataiku/dss-plugin-snowpark-container-services-llm
https://github.com/dataiku/dss-plugin-snowpark-container-services-llm
Last synced: 6 days ago
JSON representation
- Host: GitHub
- URL: https://github.com/dataiku/dss-plugin-snowpark-container-services-llm
- Owner: dataiku
- License: apache-2.0
- Created: 2023-12-19T15:57:57.000Z (over 2 years ago)
- Default Branch: main
- Last Pushed: 2024-04-30T20:11:52.000Z (almost 2 years ago)
- Last Synced: 2024-07-10T11:58:19.963Z (almost 2 years ago)
- Language: Jupyter Notebook
- Size: 123 KB
- Stars: 0
- Watchers: 5
- Forks: 1
- Open Issues: 1
-
Metadata Files:
- Readme: README.md
- Changelog: CHANGELOG.md
- License: LICENSE
- Code of conduct: CODE_OF_CONDUCT.md
Awesome Lists containing this project
README
# Snowflake Snowpark Container Services LLM Plugin
With this plugin, you can leverage LLMs hosted in Snowpark Container Services (SPCS) as part of the LLM Mesh.
# Capabilities
- Custom LLM connection option
- Connect to a chat completion endpoint hosted in SPCS to use in Dataiku Prompt Studios, LLM recipes, and via the LLM Mesh python APIs
- Connect to an embedding endpoint hosted in SPCS to use in Dataiku Embed recipe for Retrieval Augmented Generation (RAG)
# Limitations
- Must use Dataiku >= v12.5.2
- Must use Snowflake Oauth to connect to SPCS
- Must deploy and maintain one’s own LLMs on SPCS (we have sample code below to help)
# Setup
## Setup resources on the Snowflake side
1. Ensure that Snowpark Container Services (including SPCS jobs) is activated in your Snowflake account (ask your Snowflake rep)
2. (In Snowflake as ACCOUNTADMIN) Run the below code to:
a. Create a role with access to SCPS and grant this role to users
b. Create two SPCS compute pools (for the main LLM and embedding LLM)
c. Create a warehouse to run simple admin queries (or you can grant access to an existing warehouse to the DATAIKU_SPCS_ROLE)
d. Create a database and schema to hold the LLM model registry
e. Create a Snowflake OAUTH security integration
``` sql
-- Start of Snowflake SQL instructions
-- We'll need the ACCOUNTADMIN role to create these roles and compute pools
USE ROLE ACCOUNTADMIN;
-- First create a network rule and external access integration. We'll store them in a dedicated DB/Schema
CREATE DATABASE NETWORK_DB;
CREATE SCHEMA NETWORK_SCHEMA;
CREATE OR REPLACE NETWORK RULE SNOWFLAKE_EGRESS_ACCESS
MODE = EGRESS
TYPE = HOST_PORT
VALUE_LIST = ('0.0.0.0:443','0.0.0.0:80');
CREATE EXTERNAL ACCESS INTEGRATION SNOWFLAKE_EGRESS_ACCESS_INTEGRATION
ALLOWED_NETWORK_RULES = (SNOWFLAKE_EGRESS_ACCESS)
ENABLED = true;
-- create a new oauth integration
CREATE SECURITY INTEGRATION OAUTH_DATAIKU_SPCS
TYPE = OAUTH
ENABLED = TRUE
OAUTH_CLIENT = CUSTOM
OAUTH_CLIENT_TYPE = 'CONFIDENTIAL'
OAUTH_REDIRECT_URI = 'https:///dip/api/oauth2-callback'
OAUTH_ISSUE_REFRESH_TOKENS = TRUE
OAUTH_REFRESH_TOKEN_VALIDITY = 7776000
BLOCKED_ROLES_LIST = ('SYSADMIN');
-- get the OAUTH_CLIENT_ID and OAUTH_CLIENT_SECRET
CALL SYSTEM$SHOW_OAUTH_CLIENT_SECRETS('OAUTH_DATAIKU_SPCS');
-- Create a role for Dataiku users to use when calling LLMs on SPCS
CREATE ROLE DATAIKU_SPCS_ROLE;
GRANT ROLE DATAIKU_SPCS_ROLE TO USER "";
-- Grant access to the external access integration to the new role
GRANT USAGE ON INTEGRATION SNOWFLAKE_EGRESS_ACCESS_INTEGRATION TO ROLE DATAIKU_SPCS_ROLE;
-- Create a compute pool with GPU resources to host the LLM for chat completion. GPU_3 is the smallest type with GPUs
-- Grant usage of this compute pool to the DATAIKU_SPCS_ROLE
CREATE COMPUTE POOL DATAIKU_GPU_NV_S_MODEL_COMPUTE_POOL with instance_family=GPU_NV_S min_nodes=1 max_nodes=1;
GRANT USAGE ON COMPUTE POOL DATAIKU_GPU_NV_S_MODEL_COMPUTE_POOL to role DATAIKU_SPCS_ROLE;
GRANT MONITOR ON COMPUTE POOL DATAIKU_GPU_NV_S_MODEL_COMPUTE_POOL to role DATAIKU_SPCS_ROLE;
GRANT BIND SERVICE ENDPOINT on ACCOUNT to DATAIKU_SPCS_ROLE;
-- Create a compute pool with CPU resources to host the LLM for sentence embeddings. STANDARD_1 is the smallest type
-- Grant usage of this compute pool to the DATAIKU_SPCS_ROLE
CREATE COMPUTE POOL DATAIKU_CPU_X64_XS_EMBED_COMPUTE_POOL with instance_family=CPU_X64_XS min_nodes=1 max_nodes=1;
GRANT USAGE ON COMPUTE POOL DATAIKU_CPU_X64_XS_EMBED_COMPUTE_POOL to role DATAIKU_SPCS_ROLE;
GRANT MONITOR ON COMPUTE POOL DATAIKU_CPU_X64_XS_EMBED_COMPUTE_POOL to role DATAIKU_SPCS_ROLE;
CREATE OR REPLACE WAREHOUSE DATAIKU_SPCS_WAREHOUSE WITH
WAREHOUSE_SIZE='X-SMALL'
AUTO_SUSPEND = 60
AUTO_RESUME = true
INITIALLY_SUSPENDED=false;
GRANT ALL ON WAREHOUSE DATAIKU_SPCS_WAREHOUSE TO ROLE DATAIKU_SPCS_ROLE;
CREATE DATABASE DATAIKU_SPCS;
GRANT OWNERSHIP ON DATABASE DATAIKU_SPCS TO ROLE DATAIKU_SPCS_ROLE;
USE ROLE DATAIKU_SPCS_ROLE;
USE DATABASE DATAIKU_SPCS;
USE WAREHOUSE DATAIKU_SPCS_WAREHOUSE;
CREATE SCHEMA MODEL_REGISTRY;
USE SCHEMA MODEL_REGISTRY;
-- End of Snowflake SQL instructions
```
## Setup the plugin in Dataiku
1. Install the plugin - Go to Plugins -> Add Plugin -> Fetch from Git repository, then enter this repo URL:
3. In the plugin settings, add a “Snowflake login with SSO” preset:
4. In the Admin -> Connections page, create a new Snowflake connection called “spcs-access-only”. Enter the Snowflake DB, Schema, Warehouse, and Oauth info from the Snowflake-side setup. The “Auth authorization endpoint” should look like https://.aws.snowflakecomputing.com/oauth/authorize and the “Auth token endpoint” should look like https://.aws.snowflakecomputing.com/oauth/token-request. Uncheck “allow write” on the right in order to prevent users from creating datasets in this connection. We’ll use it to deploy models to SPCS only. Change the credentials mode to “per user”
5. Go to you user profile credentials, and go through the Oauth dance for both the “spcs-access-only” connection and “snowpark-container-services-llm” plugin
6. Create a python 3.8 code environment (I named it “py_38_snowpark_llms”) and add the following packages:
```
scikit-learn==1.2.1
mlflow==1.30.0
statsmodels==0.12.2
protobuf==3.16.0
xgboost==1.7.3
lightgbm==3.3.5
matplotlib==3.7.4
scipy==1.10.1
snowflake-snowpark-python==1.12.1
snowflake-snowpark-python[pandas]
snowflake-connector-python[pandas]
MarkupSafe<2.1.0
cloudpickle==2.0.0
flask>=1.0,<1.1
itsdangerous<2.1.0
Jinja2>=2.11,<2.12
snowflake-ml-python==1.3.0
dash==2.15.0
dash_bootstrap_components==1.5.0
transformers==4.37.2
sentence-transformers==2.3.1
datasets==2.16.1
torch
sentencepiece==0.1.99
presidio-anonymizer==2.2.352
presidio_analyzer==2.2.352
spacy==3.7.3
langchain==0.0.347
bitsandbytes>0.37.0
```
6. Deploy LLM(s) to SPCS, then retrieve the resulting public endpoints for chat completion and embedding models. We have a sample notebook [here](Deploy_LLMs_to_Snowpark_Container_Services.ipynb) that shows how to deploy, for chat completion: Zephyr 7B-beta, Llama2, Phi2, and Falcon; and for text embeddings: MiniLM-L6-v2.
Your URLs should look something like: https://{ENDPOINT_ID}-{SNOWFLAKE_ACCOUNT_NAME/ACCOUNT_ID}.snowflakecomputing.app
7. Create a new “Custom LLM” connection, then choose the “Snowpark Container Services LLM” plugin.
Add a model for the main LLM, give it a name, chose the "Chat completion" capability, "Snowpark Container Services LLM" type, choose your Oauth preset, then enter the generated LLM endpoint URL from earlier. Enter your Snowflake account URL, the compute pool credit cost, and Snowflake credit cost (talk to your SNowflake rep for these). Max parallelism of queries is up to you. Start with 1 or 2.
8. Add another model (in this same connection) for the text embedding model. Give it a name, choose the "Text embedding" capability, "Snowpark Container Services LLM" type, choose your Oauth preset, then enter the other generated LLM endpoint URL from the text embedding model. This model will likely have a different compute pool cost, depending on what you set up on the Snowflake side.
10. Your chat completion LLM is now ready to use in LLM mesh!
12. As is your embedding LLM to generate embeddings for Retrieval Augmented Generation (RAG)!
14. Here is sample SQL code to run from Snowflake to DROP your LLM services and compute pools:
```sql
--must drop a service before dropping a compute pool
SHOW SERVICES;
DROP SERVICE ;
DROP COMPUTE POOL DATAIKU_GPU_NV_S_MODEL_COMPUTE_POOL;
DROP COMPUTE POOL DATAIKU_CPU_X64_XS_EMBED_COMPUTE_POOL;
```