Distributed Deep Learning with PyTorch on Snowflake GPUs: A Hands-On Tutorial

Distributed deep learning is increasingly important as models grow in complexity and datasets scale to millions of records. Traditionally, training large models on GPUs required setting up dedicated infrastructure and moving data out of data warehouses. Snowflake’s new Machine Learning (ML) Container Runtime changes the game by letting you train PyTorch models directly where your data lives – inside Snowflake – using multiple GPUs in parallel. This means you can shard and distribute training across GPUs without the usual heavy infrastructure management, all while leveraging Snowflake’s robust data platform. In this tutorial, we focus on this high-demand capability: distributed PyTorch training on Snowflake’s GPU-accelerated Container Runtime. We’ll walk through a practical end-to-end example, from environment setup to model deployment, illustrating how to harness Snowflake’s GPUs for deep learning. The tutorial balances accessibility with technical depth, so both data engineers and data scientists can follow along. 

What We’ll Cover: 

Let’s get started!

Step 1: Setting Up Snowflake Notebook with GPU Container Runtime

Before we begin coding, we need a Snowflake environment that supports running notebooks with GPU acceleration. Snowflake’s Container Runtime for ML (currently in public preview) provides a flexible containerized IDE within Snowflake that supports scalable ML workflows on both CPUs and GPUs.

Prerequisites

Make sure you have:

• Snowflake account with ML runtime access – You’ll need a non-trial Snowflake account in a supported AWS region for Snowflake ML (GPU support is available in specific regions). Ensure your Snowflake administrator has enabled the Container Runtime by creating the required roles and an external access integration (this grants the container permission to fetch packages from the internet).
• Appropriate privileges – You (or your role) should have permission to create notebooks and services. Typically, an ACCOUNTADMIN (or a role with CREATE NOTEBOOK & CREATE SERVICE on a schema) is needed for initial setup. Snowflake’s documentation provides SQL commands to set up these roles and integrations (as shown in the quickstart).
• Snowsight access – Snowsight is Snowflake’s web interface where Notebooks reside. We will use it to create and run our notebook.

Creating a GPU-Backed Notebook

1. In Snowsight, navigate to the Notebooks section and create a new notebook. When prompted for the compute environment, select the Container Runtime and choose a GPU image type. Snowflake allows choosing between CPU or GPU images for the container; the GPU image comes pre-loaded with popular ML libraries like PyTorch.
2. After the notebook is created, locate the Notebook Settings (usually via a “⋮” menu or a settings icon in the notebook interface). In settings, enable the PyPI access integration (often labeled as PYPI_ACCESS_INTEGRATION). Enabling this integration allows the container to install Python packages from PyPI (internet access), which we’ll need for adding any extra libraries.
3. Optionally, verify that your notebook is indeed running on a GPU container. You can run a small snippet to check for CUDA availability:

import torch
print("CUDA available:", torch.cuda.is_available())
print("GPU devices:", torch.cuda.device_count())

This should output CUDA available: True and list the number of GPUs allocated to your notebook (e.g., 2 GPUs if you chose a 2x GPU environment). If this is False, re-check that you selected the GPU runtime image in step 1.

At this point, our Snowflake Notebook environment is ready with GPU support. Snowflake’s Container Runtime comes with many ML packages pre-installed and allows us to install any others as needed, all within Snowflake’s managed infrastructure. We’ll leverage this in the next step.

Snowflake’s GPU container already includes common frameworks (PyTorch, TensorFlow, scikit-learn, etc.) by default. However, in any practical ML project, you might need additional libraries (for example, PyTorch Lightning, Hugging Face Transformers, or specific utility packages). Thanks to the PyPI integration enabled earlier, installing extra dependencies is straightforward. In your Snowflake Notebook, you can use shell/pip commands to add packages. For example, if we need the Torchvision library for image transforms (often used with PyTorch), run a cell with:

!pip install torchvision

The !pip install syntax tells the notebook to execute a pip command in the container. After running, the package will be available in the environment. Similarly, install any other packages you require (e.g., !pip install pandas matplotlib scikit-learn if needed for data handling or visualization).

Snowflake’s container runtime is flexible – it lets you bring in any open-source Python package of your choice. This means you can work with the same libraries you’re used to, but directly against data in Snowflake. Just ensure you enable the PyPI integration (as we did) so the container can fetch packages from the internet.

Note: It’s good practice to list all your needed libraries at the top of your notebook and install them in one go. The container will cache these (for the duration of the session), and subsequent notebook runs will not need to reinstall unless the container is restarted. Also, remember that the base environment already includes many ML frameworks – for instance, PyTorch is pre-installed – so you might not need to install it manually. You can verify by importing the library (e.g., import torch) and checking its version.

With our environment set up and dependencies in place, we’re ready to access data and begin building our distributed training pipeline.

Step 3: Preparing Datasets with Snowflake’s Optimized Data Loading

One key advantage of doing ML inside Snowflake is easy access to data. We can pull data from Snowflake tables directly into our training code without any CSV exports or external ETL. Snowflake provides an optimized data loading API that bridges Snowflake data and popular ML frameworks like PyTorch. This is done through the DataConnector interface.

a. Query or Create Your Training Data in Snowflake:

Typically, you’ll have your dataset stored in Snowflake (perhaps in a table or as the result of some SQL query/feature engineering pipeline). Using Snowflake’s Python API (Snowpark), you can retrieve a Snowpark DataFrame for the data. For example:

from snowflake.snowpark.session import Session from snowflake.snowpark.functions import col
#Assume session is already available in the notebook (Snowflake provides an active session)
#Replace with your table or query
df = session.table("MY_DATABASE.MY_SCHEMA.TRAINING_DATA")
#(Optional) filter or select relevant columns
df = df.select("feature1", "feature2", "label").filter(col("label").is_not_null())

Here, df is a Snowpark DataFrame containing the features and label we need for training. You can use Snowpark to join tables, aggregate, or preprocess data using SQL and Python as needed – all executed inside Snowflake. It’s often efficient to do heavy preprocessing in Snowflake (using SQL or Snowflake ML functions) so that the data fed to PyTorch is already clean and prepared. This leverages Snowflake’s parallel processing to handle big data prep tasks, ensuring your ML pipeline is consistent and efficient (for example, using Snowflake’s Feature Store to reuse features and maintain consistency between training and inference).

b. Use DataConnector to create a PyTorch Dataset

Once you have a Snowpark DataFrame for your training data, the next step is to convert it into a format that PyTorch can work with. Snowflake’s ML library provides a DataConnector that can transform a Snowpark DataFrame directly into a PyTorch Dataset object. This avoids the need to materialize the entire dataset as a pandas DataFrame or CSV first – it streams data efficiently.

For example:

from snowflake.ml.data import DataConnector
#Create a PyTorch Dataset from the Snowpark DataFrame
train_dataset = DataConnector.from_dataframe(df).to_torch_dataset()

With these two lines, train_dataset becomes a torch.utils.data.Dataset that you can use just like any other PyTorch dataset. Under the hood, Snowflake’s connector handles batching data from Snowflake into memory. It’s optimized to utilize multiple cores or threads for loading, which helps keep the GPUs fed with data during training. If you have separate datasets for training and validation, you could similarly create val_dataset by filtering your Snowpark DataFrame for the validation split or using a different DataFrame. For example, if your table has a column split indicating train/val, you can do:

df_val = session.table("MYDB.MYSCHEMA.TRAINING_DATA").filter(col("split") == "val") val_dataset = DataConnector.from_dataframe(df_val).to_torch_dataset()

At this point, we have our data ready to go as PyTorch Dataset objects. We can wrap them in PyTorch DataLoaders later inside our training function. The key takeaway is that we did not export any data out of Snowflake – we’re pulling it directly into our training process efficiently, thanks to Snowflake’s DataConnector.

Tip: For initial exploration or if your dataset is small, you could also convert to pandas using .to_pandas() or to a NumPy array using .to_numpy(), as the DataConnector supports those too. But for large-scale training, staying in Snowflake’s optimized path (to a Torch dataset or TensorFlow dataset) is recommended to avoid memory issues.

Now that data is ready, let’s move on to configuring distributed training with PyTorch on Snowflake’s GPUs.

Step 4: Distributed PyTorch Training on Snowflake GPUs

This is the heart of the tutorial – we’ll train a PyTorch model using multiple GPUs within Snowflake. Snowflake ML provides a high-level API called PyTorchDistributor (along with supporting classes) that simplifies distributed training. Instead of manually launching multiple processes and handling inter-process communication, we can define our training logic once, and Snowflake will orchestrate it across the available GPUs for us.

a. Define the Training Function

Snowflake’s distributed training expects you to provide a training function (train_func) that contains the code to train your model on one shard of data (i.e., what one GPU/worker should do). This function will be executed in parallel on each worker (GPU). Within this function, you’ll typically: load the data shard for that worker, define your PyTorch model, set up the optimizer and loss, and run the training loop.

One special tool available inside train_func is Snowflake’s distributed context. We can get a context object that provides information like the worker’s rank (ID) and helps retrieve the data shard meant for this worker. For instance, if we used a Sharded Data Connector for data (more on this soon), the context lets each worker access only its portion of the dataset.

Let’s write a simplified training function for demonstration. Suppose we are training a simple neural network on our dataset:

import torch import torch.nn as nn import torch.optim as optim from torch.utils.data import DataLoader from snowflake.ml.modeling.distributors.pytorch import get_context
def train_func(): 
# 1. Get distributed context and data shard for this worker context = get_context() # Snowflake provides this function dataset_map = context.get_dataset_map() # Get dataset(s) assigned to this worker train_dataset = dataset_map["train"].get_shard().to_torch_dataset() # retrieve this worker's shard :contentReference[oaicite:29]{index=29}
# 2. Create DataLoader for batch loading
train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)
# 3. Define model (simple 2-layer MLP for example)
model = nn.Sequential(
    nn.Linear(100, 64),   # assuming 100 features
    nn.ReLU(),
    nn.Linear(64, 1)
)
model = model.to(torch.device("cuda" if torch.cuda.is_available() else "cpu"))
# 4. Set up optimizer and loss function
optimizer = optim.Adam(model.parameters(), lr=0.001)
criterion = nn.BCEWithLogitsLoss()  # example loss for binary classification
# 5. Training loop (one epoch for simplicity)
model.train()
for data, label in train_loader:
    data = data.to(torch.device("cuda"))   # move batch to GPU
    label = label.to(torch.device("cuda"))
    optimizer.zero_grad()
    outputs = model(data.float())
    loss = criterion(outputs.flatten(), label.float())
    loss.backward()
    optimizer.step()
# 6. (Optional) Return the trained model state or metrics
return model.state_dict()

In this train_func, a few things to note:

• We used context.get_dataset_map() to fetch the datasets passed to this worker. We assume we’ll pass a dataset named “train” when launching the job. Each worker gets a different shard of train_dataset via get_shard(). This is how we ensure data is split among GPUs (so each GPU processes a unique subset of data).
• We moved the model and data to GPU using torch.device(“cuda”). In Snowflake’s container, cuda should be available. Each worker will execute on a separate GPU if configured correctly.
• We kept it simple with one epoch and returned the model weights. In a real scenario, you might run multiple epochs and gather metrics (like training loss) for monitoring. Snowflake’s PyTorchDistributor can aggregate returns from each worker if needed.

b. Configure the Distributed Training Job

Now we need to tell Snowflake how many GPUs/workers to use and assign resources. This is done via PyTorchScalingConfig and WorkerResourceConfig:

from snowflake.ml.modeling.distributors.pytorch import PyTorchDistributor, PyTorchScalingConfig, WorkerResourceConfig
#Specify resources for each worker (each worker = one process, ideally one GPU)
worker_resources = WorkerResourceConfig(num_cpus=4, num_gpus=1) # e.g., 4 CPU cores and 1 GPU per worker :contentReference[oaicite:31]{index=31}
#Specify the scaling (number of workers and nodes)
scaling = PyTorchScalingConfig( num_nodes=1, # run on one node (our notebook's container) num_workers_per_node=2, # e.g., 2 workers on this node if 2 GPUs are available resource_requirements_per_worker=worker_resources )

In this example, we plan to use 2 GPUs on a single node, meaning two parallel workers. Each worker will get 1 GPU and 4 CPU cores for data loading. You can adjust these numbers based on your Snowflake compute size and needs:

• num_nodes: Snowflake can scale out to multiple nodes (machines) for even larger jobs, but often one node with multiple GPUs is sufficient for experiments.
• num_workers_per_node: typically set to the number of GPUs per node (so each GPU runs one worker).
• WorkerResourceConfig(num_cpus, num_gpus): how many CPU cores and GPUs each worker process should use. Important: If you don’t set num_gpus here, it defaults to 0 (meaning the worker would run on CPU). So to utilize GPUs, make sure to set num_gpus=1 (or more if somehow a worker uses multiple GPUs, but usually it’s one).

c. Launch the Distributed Training Job

We have our train_func and scaling configuration. Now we tie it together with PyTorchDistributor:

#Create the PyTorch distributor with our training function and scaling config
pytorch_job = PyTorchDistributor(train_func=train_func, scaling_config=scaling)
#Launch the training job, providing the dataset(s) to distribute
result = pytorch_job.run(dataset_map={"train": train_dataset})

A few things happen when we call run():

• Snowflake will spin up the specified number of worker processes (in this case, 2 workers on 1 node, i.e., 2 parallel processes) within the container.
• It will automatically shard the train_dataset we passed across those workers. Internally, because we passed a DataConnector dataset (from Step 3), Snowflake can split it such that each worker gets a portion of the data (approximately half the data each, since we have 2 workers). We accessed this via context.get_dataset_map()[“train”].get_shard() in our train_func.
• The train_func is executed on each worker (each with its own GPU). Snowflake’s PyTorchDistributor sets up the environment so that the workers can communicate if needed (for example, for collective operations or to avoid overlapping work), although in our simple loop we didn’t explicitly use distributed communication beyond each worker handling its shard.
• After training, the results (in our case, the returned model state dict from each worker) are collected. You could aggregate model parameters (if doing something like model averaging) or just take one. Often in distributed data-parallel training, each worker starts from the same initial weights and processes different data; typically you’d use a distributed algorithm (like AllReduce) to synchronize gradients. Snowflake’s API likely uses PyTorch’s DistributedDataParallel under the hood to ensure the model parameters are kept in sync across workers (so effectively it trains as if on one dataset). The Snowflake ML docs indicate that PyTorchDistributor handles communication and result collection for you, so you don’t have to manually sync models.

d. Monitoring and Logging

While the job runs, you can monitor the output in the notebook. Each worker’s stdout/stderr will be streamed to your notebook cell output. If you included print statements or logging in train_func (e.g., printing loss per batch or epoch), you should see them. If any worker encounters an error, it will also surface in the output.

In a more complex scenario, you might incorporate callbacks or custom logging (writing metrics to Snowflake tables or using MLflow). But for our purposes, we’ll assume the training completes successfully and we have a trained model.

We now have a trained PyTorch model (or models on each worker). Next, let’s discuss best practices to ensure this kind of distributed training runs as efficiently as possible in Snowflake’s environment.

Step 5: Best Practices for Optimizing Training Performance

Running distributed deep learning inside a data platform is new territory, so it’s worth highlighting some best practices to get the most out of Snowflake’s ML environment:

• Use Sharded Data Loading for Distribution: When using multiple workers/GPUs, ensure each worker gets a unique slice of data. Snowflake’s ShardedDataConnector is a great tool for this. We used DataConnector.from_dataframe(df) directly for simplicity, but Snowflake also allows creating a ShardedDataConnector which automatically partitions the DataFrame. For example, train_data = ShardedDataConnector.from_dataframe(df) and then passing dataset_map={“train”: train_data} to .run() (as we did) ensures each worker calls get_shard() to retrieve only its portion. This prevents redundant data processing and scales linearly with more GPUs. In short, always shard your dataset in distributed jobs – either manually or via Snowflake’s connectors – to maximize throughput.
• Leverage Snowflake for Data Prep and Feature Engineering: A common bottleneck in GPU training is reading and preprocessing data. With Snowflake, you can push a lot of this work to the database layer before training. Use Snowflake’s SQL, Snowpark DataFrame transformations, or even the Snowflake Feature Store to compute features and filter data on the warehouse, rather than in Python loops. By the time data reaches your PyTorch training, it’s already cleaned and batched efficiently. This not only speeds up training (less Python overhead per batch) but also ensures consistency – the same feature transformations can be applied during inference via Snowflake (no training-serving skew).
• Right-size your Cluster and Resources: Snowflake’s Container Runtime lets you configure how many nodes and GPUs to use. More GPUs can accelerate training, but ensure your dataset and model can scale to them. If you only have a small dataset, using too many GPUs might lead to overhead from coordination without much speed gain. Start with a modest setup (say 1 node with 2 GPUs) and measure performance. Also, allocate sufficient CPUs (num_cpus) for each worker – these CPUs handle data loading and any CPU-bound tasks. If you see GPUs are underutilized (waiting for data), consider increasing num_cpus or using a smaller batch size to reduce per-batch load time. Conversely, if GPUs are maxed out, you could increase batch size or add more GPUs.
• Take Advantage of Familiar APIs with Minimal Changes: One of the strengths of Snowflake’s ML offering is that it integrates with open-source tools with very little modification. For example, the PyTorchDistributor API is designed to feel similar to PyTorch’s usual training flow. Best practice is to reuse as much of your existing PyTorch code as possible – wrap it in the train_func – and only add Snowflake-specific calls for things like obtaining the data shard (as shown). This way, you’re not locked into proprietary interfaces, and it’s easier to maintain the code. Snowflake’s distributed classes (for XGBoost, LightGBM, PyTorch, etc.) simply add a layer to configure scaling, without requiring you to learn a brand new training framework.
• Monitor and Tune Iteratively: Just like any ML training on a new platform, monitor your jobs. Snowflake’s notebook will show output from workers; you can print timing info to see if data loading is a bottleneck. If training is not scaling well with more GPUs, investigate if perhaps the data shards are imbalanced or if there’s a lot of overhead (e.g., too frequent synchronization). You might experiment with more epochs on fewer GPUs vs. fewer epochs on more GPUs to see what yields better runtime. Also ensure that the warehouse underlying the container has enough resources – if you configured the container on a very small warehouse (with limited CPU/memory), that could throttle performance.
• Cleanup and Checkpointing: Long training jobs might benefit from checkpointing model weights to avoid losing progress if something interrupts. You can periodically save the model to Snowflake stages or the model registry during training (though this requires some coordination – one way is to have only rank 0 worker perform the save). Also be mindful of container time limits; if your notebook session has an auto-suspend or time-to-live, very long training could be cut off. Plan accordingly or split training into multiple runs (e.g., train for N epochs, save model, then possibly resume training with another run if needed).

By following these best practices, you’ll ensure that your distributed training on Snowflake runs smoothly and efficiently. Now that we have a trained model and confidence in our training pipeline, the next step is to register the model in Snowflake’s Model Registry and deploy it for inference.

Step 6: Registering and Deploying the Model in Snowflake’s ML Registry

Snowflake provides a Model Registry to version and manage machine learning models within the platform. This is a central repository (backed by Snowflake’s databases) where you can store your trained model, along with metadata, and later retrieve or deploy it. By registering the model in Snowflake, you make it a first-class Snowflake object – which means you can integrate it into Snowflake’s production workflows easily (including invoking the model via SQL or deploying it as a service).

a. Save (Log) the Model to the Registry

After training, assume we have our PyTorch model object (say, model if we consolidated the trained model weights). To log this model to Snowflake’s registry:

from snowflake.ml.registry import Registry
#Open a registry in a chosen Snowflake database/schema for models
reg = Registry(session=session, database_name="ML_MODELS_DB", schema_name="MODEL_REGISTRY")
#Register the model
model_version = reg.log_model( model, model_name="pytorch_recommendation_model", version_name="v1", description="DLRM model trained with Snowflake GPU runtime" )

A few points:

• The Registry is pointed at a specific database and schema. You might create a dedicated schema (e.g., ML_MODELS.ML_REGISTRY) to store models. If that schema doesn’t exist, you should create it first (with proper privileges). Any schema can be used as a registry, but organizing models in one place is convenient.
• log_model takes the model object (model) and a model name and version. You can also pass dependencies if needed (for example, if the model requires specific Python packages at inference time, you could list them in conda_dependencies or similar arguments). In our case, since we used PyTorch and that’s available in the Snowflake runtime, we might not need extra dependencies. This call will serialize the PyTorch model (essentially pickling it along with any needed files) and store it in Snowflake.
• After logging, model_version will be a reference to the model we just saved. In Snowsight’s UI, you can navigate to the Models page and see this model listed (with name and version). Snowflake tracks versions, so if you retrain and log again with version_name=”v2″, it will keep both versions under the same model name for you.

Under the hood, the model registry stores the model binary and metadata in Snowflake. It becomes accessible by other Snowflake users or services (with appropriate permissions) and serves as a single source of truth for the model’s state. This is great for governance and reproducibility.

b. Deploy the Model for Inference

Registering the model is like saving it; deployment is about making it usable for predictions. Snowflake offers a couple of ways to use a registered model:

• Batch or on-demand inference via SQL: You can use the model in a Snowflake SQL query, thanks to Snowflake’s support for UDFs (user-defined functions) with ML models. For example, Snowflake might allow a syntax like SELECT PREDICT(MODEL my_model VERSION ‘v1’, input_features) FROM …. This would execute the model within Snowflake, returning predictions. (The exact SQL syntax depends on Snowflake’s ML extensions, but conceptually Snowflake can treat the model as a function.)
• Managed model serving: Snowflake is introducing the ability to deploy models as services on its Snowpark Container Services. This would effectively host the model in a container (with GPU if needed) so it can serve real-time predictions via an API endpoint or via Snowflake calls.

In our scenario, since we used GPUs for training, we might also want GPU for inference (especially if it’s a heavy deep learning model). Snowflake’s managed deployment can ensure the same environment (same packages, versions) is used in production as was in development. This consistency is a big plus – no more “works on my machine” issues, because Snowflake will deploy with the same container image that our notebook used for training.

For example, to deploy the model, you could do:

#(Pseudo-code for deployment – actual API may vary as Snowflake’s model serving evolves)
from snowflake.ml.registry import Model
#Retrieve the model we just logged
my_model = reg.get_model("pytorch_recommendation_model") # get the model by name model_version = my_model.get_version("v1")
#Deploy the model (this could create a Snowflake Managed Service or a UDF)
deployment = model_version.deploy(name="pytorch_reco_service", compute_pool="GPU_SMALL")

This pseudo-code assumes a method like deploy exists. In practice, Snowflake might allow deployment via SQL commands or via Snowsight UI (“Deploy to warehouse” or “Deploy to service” buttons). The idea is to allocate compute (e.g., a Snowflake warehouse or a compute pool with GPU) where the model will be loaded and serve predictions. Given Snowflake’s preview features, you might deploy it as a Snowpark Container Service (which is indicated by granting CREATE SERVICE permission earlier) that listens for prediction requests.

Once deployed, you can test the model. If it’s deployed as a service, you might get an endpoint URL or Snowflake function to call. If deployed as a SQL function, you can run a query to get predictions.

c. Using the Model for Inference

With the model deployed in Snowflake, using it could be as simple as running a SQL query such as:

SELECT pytorch_reco_service(features_column) 
FROM MY_DATABASE.MY_SCHEMA.NEW_DATA_TO_SCORE;

This would send the features_column through the model and return the prediction (e.g., a score or class). If using the Python API, you might do:

#Using Snowpark to call the model on new data
new_df = session.table("NEW_DATA_TO_SCORE") predictions_df = new_df.select(model_version.predict(new_df["features_column"]).alias("prediction")) predictions = predictions_df.collect()

(The above is illustrative; the actual API might differ, but Snowflake is moving towards seamless predict functions on Model objects).

The key benefit here: the same Snowflake environment is running the inference. We don’t have to spin up a separate server or export the model to another serving system. Snowflake ensures that the packages (like PyTorch) and the model code are consistent with what we had during training. This managed deployment saves a ton of engineering effort in productionalizing ML models.

We’ve now trained and deployed a PyTorch model all within Snowflake. The final piece (optional but often useful) is to create a simple application to interact with our model. We’ll demonstrate that with a Streamlit app.

Step 7 (Optional): Building a Streamlit App for Interactive Inference

To make our solution interactive, especially for demonstration or internal tooling, we can use Streamlit – a popular Python framework for data apps – to build a small web application that calls our Snowflake-hosted model and displays results. Snowflake’s container environment can support running Streamlit as well (since we can pip install streamlit and even execute it in the container), or we could run Streamlit externally and have it query Snowflake.

For simplicity, let’s outline how one might do it within Snowflake’s environment:

1. Install Streamlit in the Snowflake notebook container if not already present: !pip install streamlit.
2. Write a Streamlit script that connects to Snowflake and uses the model. For example, a script app.py might:
   1. Use the Snowflake Python connector or Snowpark to query predictions. Or, if the model is exposed as an API endpoint, use requests to call that.
   2. Take some user input via Streamlit widgets (text inputs, sliders, file uploader for an image if it’s CV, etc.).
   3. When the user submits, retrieve predictions from Snowflake (perhaps via a session.sql(“SELECT model_service(…”) call or using the Python Model API).
   4. Display the results nicely (tables, charts, etc., depending on the use-case).
3. Run Streamlit from the notebook. You might use !streamlit run app.py –server.headless=true –server.enableCORS=false and then use the Snowflake Worksheets UI “Open in Streamlit” option if available. (Snowflake Notebooks have an option to render Streamlit apps output).

If running externally, the Streamlit app would simply query Snowflake using your model UDF or service.

The purpose of the Streamlit app is to provide a non-technical user interface to our model. For instance, if our model is a recommendation engine, the app could allow a user to select a customer and then display the recommended products with their scores, by querying the model’s predictions. According to Snowflake’s example, they built Streamlit dashboards to let business users explore the model’s recommendations in real time. Streamlit makes it easy to create such interactive dashboards with minimal code.

Here’s a very basic pseudo-code for a Streamlit app that queries the model via Snowpark (assuming we have a predict method):

import streamlit as st 
from snowflake.snowpark import Session
#Snowflake connection (you might use Streamlit secrets or config for credentials)
session = Session.builder.configs({...}).create()
st.title("Snowflake ML Model Inference") customer_id = st.text_input("Enter Customer ID:") if st.button("Get Recommendations"): if customer_id: query = f"SELECT recommend_model('{customer_id}') AS RECS" # using a hypothetical UDF/model result = session.sql(query).collect() recs = result[0]["RECS"] st.write("Recommendations:", recs) else: st.error("Please enter a customer ID.")

The above is illustrative. The main point is that the Streamlit app can call the model deployed in Snowflake and display the output in a user-friendly way (perhaps with charts or images if it’s a CV model).

Deploying the Streamlit app might involve running it on a separate server or within the Snowflake container. Snowflake’s documentation suggests that you can run Streamlit inside Snowflake’s UI for demonstration purposes. If you’re using the Snowflake quickstart, after running the notebook, they show how to launch a Streamlit app within Snowsight to visualize the recommendations.

Conclusion of technical steps: We’ve now gone through an entire workflow: environment setup, data prep, multi-GPU training, model registration, deployment, and even a simple front-end. All of this was accomplished within the Snowflake ecosystem, highlighting the power of Snowflake’s unified platform for data and ML.

In practice, such an approach can drastically simplify ML operations: no need to maintain separate ETL pipelines, external training clusters, and separate model serving infrastructure – Snowflake centralizes it. This is especially appealing to teams that want to govern data and models in one place and minimize the friction of moving between tools.

B EYE’s Expertise in Snowflake ML Solutions

Implementing advanced AI/ML workflows in Snowflake – like the distributed PyTorch training we just walked through – can be a game-changer for an organization. However, it helps to have experts who have done it before to ensure you’re following best practices and getting the most value. This is where B EYE comes in.

B EYE is a recognized Snowflake partner with deep expertise in data analytics and AI solutions. Our team has hands-on experience with Snowflake’s latest ML features, including the Snowpark Python API, Container Runtime for ML, and Model Registry. We specialize in helping organizations design, implement, and optimize end-to-end ML workflows within the Snowflake ecosystem.

Have questions about Snowflake?

Let’s talk!

Ask an expert at +1 888 564 1235 (for US) or +359 2 493 0393 (for Europe) or fill in our form below to tell us more about your project.