amazon sagemaker on Digi Hunch

Training and Inference in SageMaker AI

Sun, 20 Apr 2025 07:20:00 -0400

To automate and secure machine learning operations, one needs to know about programming with SageMaker SDK. While algorithms are data scientists’ wheelhouse, infrastructure engineers like myself need to understand the skeleton of the code in order to orchestrate the activities, optimize performances, and secure the communications.

It was overwhelming to begin with. Programming in machine learning heavily depends on specific frameworks, SDKs, and is highly experimental. The experimental nature entails the followings in the notebook workload:

data engineers may want to load data once, and play with different data processing logics
if the program fails halfway through, the data scientist can rerun from the failed step, or even execute the cells out of order
some lines of code (for training or inference), may demand significantly more resource than other lines;

These makes machine learning workload more challenging and require more flexibility at the infrastructure layer. However, most of the tutorials target data scientist aspirant and lack the perspective of infrastructure operations.

In this post I’ll go over what I’ve learned about training and inference workload. I will not focus on specific algorithms. Instead, I’ll use the most easy-to-understand algorithm, solving a typical problem. I’ll also use a simple framework (SK Learn) but try to reveal what is common amongst other supported frameworks.

Frameworks and Data

Amongst many programming frameworks, I need one that is generic, beginner friendly and compatible with SageMaker AI. The candidates are:

Scikit-learn: good for classical machine learning activities;
TensorFlow and PyTorch: industry standards supporting a wide range of tasks; from simple models to advanced deep learning applications;
Hugging Face, PyTorch and TensorFlow can handle NLPs;
XGBoost excels at handling tabular data

The use case statements above might be over-simplifications but at least Scikit-learn is a solid choice for beginners. It’s also good to keep Pytorch and TensorFlow in mind. As to the machine learning problem, I want to go easy with a simple classification problem with Logistic Regression. Despite the word “regression” in the name, logistic regression is an algorithm to solve classification problems (whereas linear regression solves a regression problem).

I use the simplistic iris flower data set to solve the classification problem: identifying the species (setosa, versicolor, or virginica) based on the length and width of sepal and petal. The dataset is very clean and does not reflect real-world messiness, making it a good for testing machine learn concepts, although not for training in production. The input is four numbers. The prediction output is a number (0,1, or 2) representing the species identified.

Training and Inference

Machine Learning concerns with indeterministic results. We first train the models with input data along with known results. Once trained, we can store the model as artifact with version control. Then we serve the model so we can feed it with unseen data and use the output as predictions. This machine learning flow has a lot in common with SDLC. Training resembles building of an application, the output artifact is the trained model, which we store as flat file or tarballs. In a model lifecycle, we track versions of trained models, quality control the models, decomission them and so forth. To perform inference, we serve the model in different ways, such that consuming application can use it for predictions.

To get started, let’s look at the following simple snippet, where we load data, train a simple model and use it for inference.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25


import numpy as np
import pandas as pd
from sklearn.datasets import load_iris
from sklearn.model_selection import train_test_split
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import accuracy_score

# Step 1: Load the Iris dataset
data = load_iris()
X = data.data # Features (sepal length, sepal width, etc.)
y = data.target # Target (species labels)

# Step 2: Split the dataset into training and testing sets
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.3, random_state=42)

# Step 3: Create and train a logistic regression model
model = LogisticRegression(max_iter=200)
model.fit(X_train, y_train)

# Step 4: Predict using the test set
y_pred = model.predict(X_test)

# Step 5: Evaluate the model
accuracy = accuracy_score(y_test, y_pred)
print(f"Model accuracy: {accuracy * 100:.2f}%")

Note that the snippet does not use any SageMaker SDK libraries yet. The example is simple enough to run anywhere. In real life however, the invocations of .fit() method and sometimes .predict() method can be resource-demanding. In such experiments, even a simple change to the hyper-parameter may dramatically increase the resource requirement. This experimental nature calls for a mechanism to allow execution of certain lines of code, in a different runtime environment. The SageMaker SDK provides two mechanisms: a @remote decorator and a RemoteExecutor class. Take the decorator as an example, it employs Python’s decorator to implement a wrapper of the function in the code so that the function can execute remotely on a different machine. The code looks like this:

import numpy as np
import pandas as pd
from sklearn.datasets import load_iris
from sklearn.model_selection import train_test_split
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import accuracy_score

# Step 1: Load the Iris dataset
data = load_iris()
X = data.data # Features (sepal length, sepal width, etc.)
y = data.target # Target (species labels)

# Step 2: Split the dataset into training and testing sets
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.3, random_state=42)

# Step 3: Create and train a logistic regression model: put the action in a function and mark it to execute remotely
from sagemaker.remote_function import remote

@remote(instance_type="ml.m5.large")
def train(X_train,y_train):
 model = LogisticRegression(max_iter=200)
 model.fit(X_train, y_train)
 return model
model=train(X_train, y_train)

# Step 4: Predict using the test set: put the action in a function and mark it to execute remotely
@remote(instance_type="ml.m5.large")
def predict(model):
 return model.predict(X_test)
y_pred = predict(model)

# Step 6: Evaluate the model
accuracy = accuracy_score(y_test, y_pred)
print(f"Model accuracy: {accuracy * 100:.2f}%")

Note that this snippet requires a SageMaker SDK library (sagemaker.remote_function) and must run on a Jupyter Lab notebook. That comes with the flexibility to run any function on any type of instance allowed in the SageMaker domain. When a function runs remotely from Jupyter workspace notebook, SageMaker has to spend time provisioning and bootstrapping an EC2 instance, which usually takes a noticeable delay and the stdout gives a summary of billable seconds. SageMaker also capture these remote executions as training job uniformly.

SageMaker SDK Abstractions

We can summarize the activity pattern from the two example scripts as below:

This pattern however, does not fully integrate with other SageMaker AI capabilities. In a full-blown integration, we use the abstraction classes in SageMaker SDK, so that the platform has insights into our activities in the script, i.e. when we’re training models, building pipelines, running inference, etc. This allows SageMaker AI platform to manage these activities. The verb manage represents many specific operations to keep track of outside of the Python scripting runtime. For example, we can keep the training runs, even after the Jupyter workspace runtime is stopped or deleted.

We have to know the abstraction classes for different frameworks. They all follow the same pattern regardless of the framework name. Here is a diagram of the relations between the base classes in SageMaker SDK:

In this programming model, the users offload activities like training and loading a model to a separate Python script called entry point script. Then the users reference the script in the constructor of the estimator. The platform requires the script to follow a particular pattern. For example, to load a model, the function named model_fn() must exist. The training code goes under main() function. A simple entry point script looks like this:

# entry_point.py script for training and loading models on SageMaker AI
import argparse
import joblib
import pandas as pd
from sklearn.linear_model import LogisticRegression
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score
import os

def main(args):
 # Load data
 df = pd.read_csv(args.train)
 X = df.drop(columns='species')
 y = df['species']
 
 # Train-test split
 X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.3, random_state=42)
 
 # Train the logistic regression model
 model = LogisticRegression(max_iter=200)
 model.fit(X_train, y_train)
 
 # Evaluate the model
 y_pred = model.predict(X_test)
 accuracy = accuracy_score(y_test, y_pred)
 
 print(f'Model accuracy: {accuracy * 100:.2f}%')
 
 # Save the model to a file
 joblib.dump(model, '/opt/ml/model/model.joblib')

if __name__ == '__main__':
 parser = argparse.ArgumentParser()
 parser.add_argument('--train', type=str, default='/opt/ml/input/data/train/iris_train.csv')
 args = parser.parse_args()
 main(args)

# function to load a model
# https://sagemaker.readthedocs.io/en/stable/frameworks/sklearn/using_sklearn.html#load-a-model
def model_fn(model_dir):
 """Load the trained model from the model_dir"""
 model_path = os.path.join(model_dir, "model.joblib")
 return joblib.load(model_path)

This script takes an argument for the location of the training data. The main function performs the model training. Let’s name this script entry_point.py and in the next couple sections we can discuss two patterns to perform inference in a scalable fashion: Batch Transform and Inference Endpoint.

Batch Transform

Training is an expensive activity that we cannot afford to repeat upon receiving every inference requests. That’d be too slow. To serve the model for inference request in a scalable way, we hold the inference requests in batches, and sometimes in the form of a multi-line input file from S3 bucket. Below is an example:

# Import necessary libraries for SageMaker
import sagemaker
from sagemaker import get_execution_role
from sklearn.datasets import load_iris
import pandas as pd
import os
import joblib
# Get the execution role for the SageMaker notebook
role = get_execution_role()
########################################################
import boto3
from sagemaker.inputs import TrainingInput
# Load the Iris dataset
iris = load_iris()
X = iris.data
y = iris.target
# Create a DataFrame to easily export
df = pd.DataFrame(X, columns=iris.feature_names)
df['species'] = y
# Save dataset to a CSV file and upload to S3 bucket
train_file = 'iris_train.csv'
df.to_csv(train_file, index=False)
s3_bucket = 'sagemaker-ca-central-1-train' # Make sure the bucket exists
s3_data_path = f's3://{s3_bucket}/iris/'
s3 = boto3.client('s3')
s3.upload_file(train_file, s3_bucket, 'iris/iris_train.csv')
# Define the location of training data in S3
s3_input_train = TrainingInput(s3_data=s3_data_path, content_type='csv')
########################################################
# Define the SKlearn Estimator and start training the model
from sagemaker.sklearn import SKLearn
sklearn_estimator = SKLearn(
 entry_point='train.py', # Your script
 role=role,
 instance_type='ml.m5.large', # Choose an appropriate instance
 instance_count=1,
 framework_version='0.23-1', # Use a compatible version of sklearn
 base_job_name='sklearn-iris-job',
 hyperparameters={},
 source_dir='.', # If using any other source code files, include them here
)
# Start training the model
sklearn_estimator.fit({'train': s3_input_train})
########################################################
# Preparing inference input file and output directory
input_file='batch_input.csv'
pd.DataFrame(X).to_csv(input_file, header=False, index=False)

# Upload the CSV file to S3
s3_bucket = 'sagemaker-ca-central-1-train' # Make sure you replace with your bucket
s3_data_path = f's3://{s3_bucket}/iris_input/'
s3 = boto3.client('s3')
s3.upload_file(input_file, s3_bucket, 'iris_input/batch_input.csv')
s3_output_path = f's3://{s3_bucket}/iris_ouput/'
########################################################
transformer = sklearn_estimator.transformer(
 instance_count=1,
 instance_type='ml.m5.large',
 output_path=s3_output_path,
 strategy='MultiRecord',
 assemble_with='Line',
 accept='text/csv',
)
transformer.transform(
 data=f'{s3_data_path}{input_file}',
 content_type='text/csv',
 split_type='Line',
 input_filter="$[0,1,2,3]"
)
transformer.wait()
########################################################

We specify the instance size while creating the estimator with constructor. The later transform step takes long and requires provisioning of a computing resource. After the transformer run, the result is in the S3 bucket in the output directory.

Inference Endpoint

If inference requests come in at arbitrary times and expect timely responses, it makes sense to introduce client-server architecture. We use web service framework to host trained models behind a passively-open inference endpoint, allowing client applications to make REST API calls with inference input in the payload, and get inference output in the HTTP response. Essentially, we serve the model behind web servers.

In SageMaker AI, we can set up inference endpoint in three ways. The most common is real-time inference endpoint. We can configure a fleet of web servers in an autoscaling group to serve the inference requests. Below is an example:

# Import necessary libraries for SageMaker
import sagemaker
from sagemaker import get_execution_role
from sklearn.datasets import load_iris
import pandas as pd
import os
import joblib
# Get the execution role for the SageMaker notebook
role = get_execution_role()
########################################################
import boto3
from sagemaker.inputs import TrainingInput
# Load the Iris dataset
iris = load_iris()
X = iris.data
y = iris.target
# Create a DataFrame to easily export
df = pd.DataFrame(X, columns=iris.feature_names)
df['species'] = y
# Save dataset to a CSV file and upload to S3 bucket
train_file = 'iris_train.csv'
df.to_csv(train_file, index=False)
s3_bucket = 'sagemaker-ca-central-1-train' # Make sure you replace with your bucket
s3_data_path = f's3://{s3_bucket}/iris/'
s3 = boto3.client('s3')
s3.upload_file(train_file, s3_bucket, 'iris/iris_train.csv')
# Define the location of your training data in S3
s3_input_train = TrainingInput(s3_data=s3_data_path, content_type='csv')
########################################################
# Define the SKlearn Estimator and start training the model
from sagemaker.sklearn import SKLearn
sklearn_estimator = SKLearn(
 entry_point='train.py', # Your script
 role=role,
 instance_type='ml.m5.large', # Choose an appropriate instance
 instance_count=1,
 framework_version='0.23-1', # Use a compatible version of sklearn
 base_job_name='sklearn-iris-job',
 hyperparameters={},
 source_dir='.', # If using any other source code files, include them here
)
# Start training the model
sklearn_estimator.fit({'train': s3_input_train})
########################################################
# Create a model based on training result
model = sklearn_estimator.create_model()
predictor = model.deploy(
 role=role,
 initial_instance_count=1,
 instance_type='ml.m5.large',
 endpoint_name='TESTENDPOINT1',
 wait=True,
)
########################################################
# Make predictions with the deployed model
test_data = [[5.1, 3.5, 1.4, 0.2]] # Example test data (sepal length, width, petal length, petal width)
prediction = predictor.predict(test_data)
print('Prediction: ', prediction)
# Make predictions with the deployed model again
test_data_2 = [[6.7, 3.1, 5.6, 2.4]] # Example test data (sepal length, width, petal length, petal width)
prediction_2 = predictor.predict(test_data_2)
print('Prediction_2: ', prediction_2)
# Clean up by deleting the endpoint when done
predictor.delete_endpoint()

The real-time inference endpoint is backed by specialized hardware and remains up all the time. That is expensive, especially if inference requests are only sporadic and can tolerate latency. To lower the cost, SageMaker AI platform supports two variants of inference: serverless and asynchronous inferences. The former stems from the (psudo-)concept of serverless computing, with the ability to scale to zero when idle. The latter uses a queue to hold up requests until a meaningful cumulation of jobs. Then it notifies the requestor of the processing results.

To use asynchronous inference, you specify the serverless_inference_config parameter when creating a model. To use asynchronous inference, you need to use the AsyncPredictor class. Whether every model (self-trained, JumpStart, etc) can work with all three patterns of inference, has to be reviewed case by base.

Summary

Most of the machine learning training materials were created by data scientist focusing on the models and algorithms. For infrastructure engineering like myself to learn the topic, I need to understand different ways that one can perform training and inference without getting swapped into the algorithmic details. In this post I took some simplistic examples to look at what training and inference looks like, as local script and as integrated with SageMaker AI platform using its SDKs, to familiarize myself with MLOps operations.
There is also a Udemy course I used to familiarize myself with the topic. The material in the course does not rely on SageMaker AI at all and built model (with scikit-learn) and app (host on Docker and Kubernetes) ground up. In the next few posts, I’ll continue explore training activities on Kubernetes based infrastructure.

SageMaker AI and Networking

Tue, 01 Apr 2025 00:53:00 -0400

Machine learning workflows are highly experimental. To smooth out the processes, Amazon SageMaker AI packages many features as managed services. As an infrastructure specialist, I want to remain compliant. At a basic level, compliant architecture means multi-account structure and hub-and-spoke VPC topology in a landing zone. While the multi-account operating model for SageMaker is well documented, these ML managed services obscure the network configuration. I’m always looking for low-level insights on where the computing activity is happening and how the application traffic flows in and out of our VPCs. I don’t always get straight answers and I often have to experiment them out. This post is a review of the networking aspects of SageMaker AI I recently learned.

Naming Shenanigans

Unfortunately, I have to start with how AWS has renamed these services, so the terms remain clear throughout the post.

In Dec 2024, AWS renamed SageMaker to SageMaker AI. The name of SageMaker going forward represents the overarching AWS service for machine learning, data, analytics and generative AI. Here is a video for clarification. I’m not a fan of how they repurpose the names. As of date there are still a lot of content referencing SageMaker AI as SageMaker, whose meaning has changed.

If that’s not enough, here’s another one. The SageMaker Studio launched originally in 2019 for model development. In 2023 that became SageMaker Studio classic, in favour of the newly launched studio, taking the name of SageMaker studio. In Feb 2025, AWS deprecated SageMaker Studio classic. You can only create SageMaker Studio in SageMaker AI.

At the SageMaker level, AWS launched SageMaker Unified Studio, the all-encompassing development environment for data analytics, generative AI, and so on. In this post though, we talk about many features under SageMaker AI and SageMaker Studio. While the service UIs are picturesque, we remain focused on two questions: how these services interact with resources on our VPCs, and how they connect to the Internet.

Workload Categories

We divide machine learning workload into three categories, based on network connectivity pattern: notebooks, model hosting, and pipeline jobs.

Notebooks are where data scientists carry out experiments by running experimental scripts on performing hardwares (depending on the tasks), usually within IDE application as Jupyter Labs, Code Editors. Data scientist users may perform any machine learning related activities such as model evaluation, etc. It is possible that one part of a notebook only requires consumer grade CPU and another part of the notebook program requires a performant GPU. It all depends on the nature of the program code.

Model is the key artifact in Machine Learning workflows. Models themselves are files stored in S3 buckets. The machine learning engineers performs two most common activities. They train the model, and feed the model with unseen data for new output (inference). In simple workflows, data scientists may build, train a model and run inference all from within the same notebook. As the experiment concludes and the team wants to operationalize the inference, it makes senses the run inference in a client-server architecture. This calls for a inference endpoint acting as the server, backed by the trained model, operating on a single or an autoscaling group of instances. The client application feeds the endpoint with unseen data, often using REST API calls, and expects inference results.

Pipeline steps like training do not operate on a server. They are similar to Notebook workloads. The difference is that pipeline steps are headless executions. The steps are non-interactive without engaging the Studio GUI. Many other types of activities in machine learning are in similar pattern, such as model evaluation, model optimization or any general processing such as a Python script. I consider them similar to training activities. Since we orchestrate these headless activities with pipelines (e.g. SageMaker pipeline), and each step may execute on some specialized instance depending on the computing requirements. Collectively, I call these activities the pipeline jobs.

Let’s look at these workloads through the networking lens.

Studio Notebooks

The most common Studio app is some kind of notebooks, such as Jupyper Lab. However, this category can generally include all kinds of SageMaker Studio apps, e.g. Canvas, Code Editor. I use the term Studio app and Studio notebook interchangeably but the APIs mostly refer to these as apps, such as AppNetworkAccessType.

In the app, a user may create one or more spaces each specifying the backing instance type. The configuration that influences the instances’ networking setup is in SageMaker AI domain’s Network Setting. There are two parts of this AppNetworkAccessType setting:

Network Mode (also called AppNetworkAccessType in AWS SDK):
- PublicInternetOnly (default): only EFS traffic goes through the specified VPC and subnets. Other studio traffic (e.g. API calls) goes through the Internet Gateway of the VPC that the studio manages internally
- VpcOnly: all studio traffic goes through the specified VPC and subnets. This delegates the responsibility of connectivity to endpoints to the VPC’s owner.
VPC and Subnet: to place EFS mount points on. Also route other studio traffic in VpcOnly mode.

There are one diagrams on the documentation for each network mode (PublicInternetOnly on the left; VPCOnly on the right):

The diagrams (as of March 2025) are not accurate because the PublicInternetOnly mode also has a domain managed ENI per app space. The VpcOnly mode is when the ML do not like the idea that a Notebook instance can bypass centrally managed Internet path. The team must ensure the endpoints are reachable, either via Internet, or via routable VPC endpoint (e.g. Gateway Endpoint for S3 and Interface endpoint for the rest). In a hub-and-spoke setup it might be another dedicated VPC that provides the interface endpoints centrally.

The domain settings include a few configurations on the underlying instance. For example, SecurityGroupIds specifies the security groups associated with the ENIs. DockerSettings enables Docker daemon on the instance, allowing users to test container workload in local mode. Note that in the more constraint VPC-only mode, Docker pull and push operations outside of Amazon Elastic Container Registry aren’t supported. To pull or push from ECRs users also need to white-list account IDs of the private ECRs in the VpcOnlyTrustedAccounts setting.

In comparison with the other two types of workload, the network traffic for studio notebooks are the easiest to control because they are all configured at the SageMaker AI domain level. Once users with user profiles under a domain creates a notebook, the domain or user profile determines the network mode, and subnet values and security groups. Users themselves cannot change these settings.

Inference Endpoint

The machine learning realm has a few established programming frameworks to host a model file behind an inference endpoint. For example, TensorFlow, PyTorch, Scikit-learn, and even Fast API. Amazon SageMaker AI supports many such frameworks and makes it straightforward. Managing the frameworks requires complex dependency management, a typical use case of containers. Apart from choosing a proper container image, user also selects instance types. These machine learning special purpose instances are pricier than their commodity counterpart. In low-traffic workflows, users may provision inference endpoint on-demand and use it in a controlled time-window, or just use serverless inference endpoint if the model supports it.

When creating inference endpoint, the CreateModel API is used. Under the VpcConfig attribute, two parameters are at play: Network Isolation and VPC-Subnet configuration. I summarize them as below based on the documentation:

Network Isolation	VPC-Subnet Configuration	Description
Disabled	Not specified	SageMaker AI containers are able to access external service and resources on the public Internet; but not able to access resources inside your VPC
Disabled	Specified	SageMaker AI containers communicate with resources inside your VPC through an ENI (Elastic Network Interface). Users are responsible for managing network access to your VPC and Internet.
Enabled	Not specified	SageMaker AI container cannot communicate with resources inside your VPC or on the public Internet
Enabled	Specified	The download and upload operations are routed through your VPC, but the inference (and training) containers themselves continue to be isolated from the network, and do not have access to any resource within your VPC or on the internet.

The Network Isolation option governs the container connectivity option. If we do not expect the inference activity to make outgoing network calls (except for downloading artifact and packages), then we should enable network isolation. On the other hand, if the inference container needs resources on VPC or on the Internet, disable network isolation. Either way, we specify the VPC so that we manage the routing through VPC.

In a compliant networking setup where Internet access must be centralized, the VPC-Subnet configuration must always be configured. The network isolation value depends on the nature of inference workload. However, what seems to be missing in the Studio UI is the activity to enforce that VPC-Subnet configuration is always specified.

As a workaround, we can use the SageMaker Domain’s IAM role to contain such attempt at API level. Below is an example of deny policy:

{
 "Version": "2012-10-17",
 "Statement": [
 {
 "Action": [
 "sagemaker:CreateModel"
 ],
 "Condition": {
 "BoolIfExists": {
 "sagemaker:VpcSubnets": "false"
 }
 },
 "Effect": "Deny",
 "Resource": "*",
 "Sid": "DenyModelcreationIfNotOnVPC"
 },
 {
 "Action": [
 "sagemaker:CreateModel"
 ],
 "Condition": {
 "ForAnyValue:StringNotEquals": {
 "sagemaker:VpcSubnets": [
 "subnet-999999999999999",
 "subnet-111111111111111"
 ]
 }
 },
 "Effect": "Deny",
 "Resource": "*",
 "Sid": "DenyModelcreationIfAnySpecifiedSubnetIsNotIntended"
 }
 ]
}

We can use this policy in conjunction with the AmazonSageMakerFullAccess managed policy. The request to create a mode gets denied, either if VpcSubnets are not specified, or they are but not from the preset list of subnet IDs. Once the user selects subnets, corresponding ENIs will get created in the subnets too and user needs to specify security groups for the ENIs.

Enforcing in IAM policy requires that the user who creates endpoint either on SageMaker Studio UI or SageMaker SDK must know the exact subnet IDs as well as appropriate security groups. This requires access to the VPC and can turn into an operation pain point if the users are not well versed with networking. Ideally subnet configuration should also be enforceable at the domain level.

Pipeline Jobs

Most machine learning jobs do not need to function behind an endpoint (i.e. server-side), for example, training, labeling job, model optimization, hyper parameter tuning, etc. In operation, we often use a pipeline to orchestrate these short-lived, non-interactive, headless jobs. Therefore, I simply refer to them as pipeline jobs. They sometimes rely on special purpose instance types. In most cases, they need access to either the Internet or other resources available via customer VPC, a connectivity pattern similar to that of interface endpoints.

To create such resources, SageMaker AI domain user either operate on SageMaker Studio, or program with SageMaker SDK. To make it easy to specify network isolation and subnet configurations, the SDK even has a class for NetworkConfig that can pass to many types of processors (steps).

from sagemaker.network import NetworkConfig

security_group_ids = ['sg-#']
subnets = ['subnet-#']
enable_network_isolation = True

network_config = NetworkConfig(
 security_group_ids=security_group_ids, 
 subnets=subnets, 
 enable_network_isolation=enable_network_isolation
)

script_processor = ScriptProcessor(
 image_uri='my-script-processor-image',
 command=['python3', 'script.py'],
 instance_type='ml.m5.large',
 role=role,
 network_config=network_config
)

The network_config parameter exists as an argument in the creation method of many other resources via SDK. However, the SageMaker UI domain does not have a mechanism to enforce it one way or another. Leaving this option open to users is not what every organization wants either. We could exercise control as much as we can with condition keys such as sagemaker:VpcSubnets in the deny policy for SageMaker IAM role as the example above shows. The Service Authorization Reference document lists out in which SageMaker SDK calls the sagemaker:VpcSubnets condition key (or equivalent) exists. A proactive IAM policy to safeguard all the applicable SDK calls would be helpful as a workaround to the missing enforceability at domain level for SageMaker AI.

Note that on the SageMaker Studio’s Pipeline tool there is a Network configuration seemingly for the pipeline. However, the CreatePipeline SDK call does not have an argument about network configuration. The PipelineDefinition argument requires a JSON format input to define the pipeline configuration and the network configuration is defined per step in the definition.

Another perspective to look at this issue is how we can give a pipeline step (or in general any job runtime) the flexibility to connect to Internet, and in the mean time remain in control of its network connectivity. It depends on the intended security posture but we mainly look at these two questions on the requirement:

can the job runtime access the VPC?
can the job runtime access the Internet on its own path?

Depending on the answer, we can configure a Job in three ways:

	Configuration	Description
a	Neither #1 or #2 are allowed	Most secure but might be overly restrictive because the job runtime may need to download artifacts. This requires enabling Network Isolation and specify VPC-Subnet configuration.
b	Either #1 or #2 is allowed, but not both at the same time.	Exclusively allowing #1 is more secure because the VPC can manage access to Internet for the job runtime. Exclusively allowing #2 breaks the central Internet access pattern and should not be allowed if central Internet access is a compliance requirement.
c	Both #1 and #2 are allowed at the same time	This configuration should not be made possible due to exfiltration risk.

This classification of network configuration for ML pipeline job, is quite similar to DevOps pipeline job agent (think of Azure DevOps agent or Terraform agents). The user may use service provider’s agent which come with its own Internet access, or choose to self-host the agent to allow access to VPC but the VPC’s owner is then responsible for managing Internet routing through the custom VPC.

SageMaker AI makes c impossible, which is good. SageMaker administrator needs to evaluate the requirement between a and b and determine how to enforce it with IAM policy.

Summary

As an infrastructure security specialist, I investigated networking options in SageMaker AI. When any user is performing any task in SageMaker AI, I am concerned with two questions:

How does the instance behind the endpoint connects to the internet and to our VPCs;
How do we enforce the connectivity pattern (enforceability);

For the studio notebooks workload, the control is through SageMaker domain. For Inference endpoint and pipeline jobs, the control is reactive. We use IAM policies and role to deny attempts to create resources with bad network configuration. The limitation is that the SageMaker users have to know what network configuration options are available.

As a result, the infrastructure security specialist must clearly define the required network configuration, and communicate it out to SageMaker AI users.