Notebooks on Container Runtime for ML¶

Overview¶

You can run Snowflake Notebooks on Container Runtime for ML. Container Runtime for ML is powered by Snowpark Container Services, giving you a flexible container infrastructure that supports building and operationalizing a wide variety of workflows entirely within Snowflake. Container Runtime for ML provides software and hardware options to support advanced data science and machine learning workloads. Compared to virtual warehouses, Container Runtime for ML provides a more flexible compute environment where you can install packages from multiple sources and select compute resources, including GPU machine types, while still running SQL queries on warehouses for optimal performance.

This document describes some considerations for using notebooks on Container Runtime for ML. You can also try the Getting Started with Snowflake Notebook Container Runtime quickstart to learn more about using the Container Runtime for ML in your development.

Prerequisites¶

Before you start using Snowflake Notebooks on Container Runtime for ML, the ACCOUNTADMIN role must complete the notebook setup steps for creating the necessary resources and granting privileges to those resources. For detailed steps, see Administrator setup.

Create a notebook on Container Runtime for ML¶

When you create a notebook on Container Runtime for ML, you choose a warehouse, runtime, and compute pool to provide the resources to run your notebook. The runtime you choose gives you access to different Python packages based on your use case. Different warehouse sizes or compute pools have different cost and performance implications. All of these settings can be changed later if needed.

To create a Snowflake Notebook to run on Container Runtime for ML, follow these steps:

1. Sign in to Snowsight.

2. Select Notebooks.

3. Select + Notebook.

4. Enter a name for your notebook.

5. Select a database and schema in which to store your notebook. These cannot be changed after you create the notebook.

   Note

   The database and schema are only required for storing your notebooks. You can query any database and schema your role has access to from within your notebook.

6. Select the Run on container as your Python environment.

7. Select the Runtime type: CPU or GPU.

8. Select a Compute pool.

9. Change the selected warehouse to use to run SQL and Snowpark queries.

10. To create and open your notebook, select Create.

Runtime:

This preview provides two types of runtimes: CPU and GPU. Each runtime image contains a base set of Python packages and versions verified and integrated by Snowflake. All runtime images support data analysis, modeling, and training with Snowpark Python, Snowflake ML, and Streamlit.

To install additional packages from a public repo, you can use pip. An external access integration (EAI) is required for Snowflake Notebooks to install packages from external endpoints. To configure EAIs, see Set up external access for Snowflake Notebooks. However, if a package is already part of the base image, then you can’t change the version on the package by installing a different version with pip install. For a list of the pre-installed packages, see Container Runtime for ML.

Compute pool:

A compute pool provides the compute resources for your notebook kernel and Python code. Use smaller, CPU-based compute pools to get started, and select higher-memory, GPU-based compute pools to optimize for intensive GPU usage scenarios like computer vision or LLMs/VLMs.

Note that each compute node is limited to running one notebook per user at a time. You should set the MAX_NODES parameter to a value greater than one when creating compute pools for notebooks. For an example, see Compute resources. For more details on Snowpark Container Services compute pools, see Snowpark Container Services: Working with compute pools.

When a notebook is not being used, consider shutting it down to free up node resources. You can shut down a notebook by selecting End session in the connection dropdown button.

Run a notebook on Container Runtime for ML¶

After you create your notebook, you can start running code immediately by adding and running cells. For information about adding cells, see Develop and run code in Snowflake Notebooks.

!pip install transformers scipy ftfy accelerate

Running ML workloads¶

Notebooks on Container Runtime for ML are well suited for running ML workloads such as model training and parameter tuning. Runtimes come pre-installed with common ML packages. With EAI set up, you can install any other packages you need using !pip install.

The following examples show how to use some of the available libraries for your ML workload.

from snowflake.snowpark.context import get_active_session
import pandas as pd
import xgboost as xgb
from sklearn.model_selection import train_test_split

session = get_active_session()
df = session.table("my_dataset")
# Pull data into local memory
df_pd = df.to_pandas()
X = df_pd[['feature1', 'feature2']]
y = df_pd['label']
# Split data into test and train in memory
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.15, random_state=34)
# Train in memory
model = xgb.XGBClassifier()
model.fit(X_train, y_train)
# Predict
y_pred = model.predict(X_test)

from snowflake.snowpark.context import get_active_session
from snowflake.ml.modeling.xgboost import XGBClassifier
from snowflake.ml.modeling.metrics import accuracy_score

session = get_active_session()
df = session.table("my_dataset")
feature_cols=['FEATURE1', 'FEATURE2']
label_col = 'LABEL'
predicted_col = [PREDICTED_LABEL]
df = df[features_cols + [label_col]]
# Split is pushed down to associated warehouse
train_df, test_df = df.random_split(weights=[0.85, 0.15], seed=34)

model = XGBClassifier(
    input_cols=feature_cols,
    label_cols=label_col,
    output_cols=predicted_col,
    # This will enable leveraging all GPUs on the node
    tree_method="gpu_hist",
)
# Train
model.fit(train_df)
# Predict
result = model.predict(test_df)

accuracy = accuracy_score(
df=result,
y_true_col_names=label_cols,
y_pred_col_names=predicted_col)

from snowflake.snowpark.context import get_active_session
from snowflake.ml.modeling.lightgbm import LGBMClassifier
from snowflake.ml.modeling.metrics import accuracy_score

session = get_active_session()
df = session.table("my_dataset")
feature_cols=['FEATURE1', 'FEATURE2']
label_col = 'LABEL'
predicted_col = [PREDICTED_LABEL]

df = df[features_cols + [label_col]]
# Split is pushed down to associated warehouse
train_df, test_df = df.random_split(weights=[0.85, 0.15], seed=34)

model = LGBMClassifier(
    input_cols=feature_cols,
    label_cols=label_col,
    ouput_cols=predicted_col,
    # This will enable leveraging all GPUs on the node
    device_type="gpu",
)

# Train
model.fit(train_df)
# Predict
result = model.predict(test_df)

accuracy = accuracy_score(
df=result,
y_true_col_names=label_cols,
y_pred_col_names=predicted_col)

Cost/billing considerations¶

When running notebooks on Container Runtime for ML, you may incur both warehouse compute and SPCS compute costs.

Snowflake Notebooks require a virtual warehouse to run SQL and Snowpark queries for optimized performance. Therefore, you might also incur virtual warehouse compute costs if you use SQL in SQL cells and Snowpark push-down queries executed in Python cells. The following diagram shows where compute happens for each type of cell.

For example, the following Python example uses the xgboost library. The data is pulled into the container and compute occurs on Snowpark Container Services:

from snowflake.snowpark.context import get_active_session
import pandas as pd
import xgboost as xgb
from sklearn.model_selection import train_test_split

session = get_active_session()
df = session.table("my_dataset")
# Pull data into local memory
df_pd = df.to_pandas()
X = df_pd[['feature1', 'feature2']]
y = df_pd['label']
# Split data into test and train in memory
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.15, random_state=34)

To learn more about warehouse costs, see Overview of warehouses.