RAPIDS can be used in a few ways with AWS SageMaker.

SageMaker Notebooks#

SageMaker Notebook Instances can be augmented with a RAPIDS conda environment.

We can add a RAPIDS conda environment to the set of Jupyter ipython kernels available in our SageMaker notebook instance by installing in a lifecycle configuration script.

To get started head to SageMaker and create a new SageMaker Notebook Instance.

Select your instance#

Select a RAPIDS compatible GPU (NVIDIA Pascal or greater with compute capability 6.0+) as the SageMaker Notebook instance type (e.g., ml.p3.2xlarge).

Screenshot of the create new notebook screen with a ml.p3.2xlarge selected

Create a RAPIDS lifecycle configuration#

Create a new lifecycle configuration (via the ‘Additional Options’ dropdown).

Screenshot of the create lifecycle configuration screen

Give your configuration a name like rapids and paste the following script into the “start notebook” script.


set -e

sudo -u ec2-user -i <<'EOF'

mamba create -y -n rapids -c rapidsai -c conda-forge -c nvidia rapids=24.04 python=3.10 cuda-version=11.8 ipykernel
conda activate rapids

# optionally install AutoGluon for AutoML GPU demo
# python -m pip install --pre autogluon

python -m ipykernel install --user --name rapids
echo "kernel install completed"

Set the volume size to at least 15GB, to accommodate the conda environment.

Then launch the instance.

Select the RAPIDS environment#

Once your Notebook Instance is InService select “Open JupyterLab”

Then in Jupyter select the rapids kernel when working with a new notebook.

Screenshot of Jupyter with the rapids kernel highlighted

SageMaker Estimators#

RAPIDS can also be used in SageMaker Estimators. Estimators allow you to launch training jobs on ephemeral VMs which SageMaker manages for you. The benefit of this is that your Notebook Instance doesn’t need to have a GPU, so you are only charged for GPU instances for the time that your training job is running.

All you’ll need to do is bring in your RAPIDS training script and libraries as a Docker container image and ask Amazon SageMaker to run copies of it in-parallel on a specified number of GPU instances in the cloud. Let’s take a closer look at how this works through a step-by-step approach:

  • Training script should accept hyperparameters as command line arguments. Starting with the base RAPIDS container (pulled from Docker Hub), use a Dockerfile to augment it by copying your training code and set WORKDIR path to the code.

  • Install sagemaker-training toolkit to make the container compatible with Sagemaker. Add other packages as needed for your workflow needs e.g. python, flask (model serving), dask-ml etc.

  • Push the image to a container registry (ECR).

Screenshot of summarized step to build Estimator

  • Having built our container and custom logic, we can now assemble all components into an Estimator. We can now test the Estimator and run parallel hyperparameter optimization tuning jobs.

estimator = sagemaker.estimator.Estimator(
), job_name=job_name)

hpo = sagemaker.tuner.HyperparameterTuner(
), job_name=tuning_job_name, wait=True, logs="All")

Upload data to S3#

We offer the dataset for this demo in a public bucket hosted in either the us-east-1 or us-west-2 regions

Create Notebook Instance#

Sign in to the Amazon SageMaker console. Choose Notebook > Notebook Instances > Create notebook instance. If a field is not mentioned below, leave the default values:

  • NOTEBOOK_INSTANCE_NAME = Name of the notebook instance

  • NOTEBOOK_INSTANCE_TYPE = Type of notebook instance. We recommend a lightweight instance (e.g., ml.t2.medium) since the instance will only be used to build the container and launch work

  • PLATFORM_IDENTIFIER = ‘Amazon Linux 2, Jupyter Lab 3’

  • IAM_ROLE = Create a new role > Create role

Launch Jupyter Notebook#

In a few minutes, Amazon SageMaker launches an ML compute instance — when its ready you should see several links appear in the Actions tab of the ‘Notebook Instances’ section, click on Open JupyerLab to launch into the notebook.


If you see Pending to the right of the notebook instance in the Status column, your notebook is still being created. The status will change to InService when the notebook is ready for use.

Run the Example Notebook#

Once inside JupyterLab you should be able to upload the Running RAPIDS hyperparameter experiments at scale example notebook and continue following those instructions.

Further reading#

We’ve also written a detailed blog post on how to use SageMaker with RAPIDS.

Related Examples#

Running RAPIDS hyperparameter experiments at scale on Amazon SageMaker

cloud/aws/sagemaker workflow/hpo library/cudf library/cuml library/scikit-learn data-format/csv data-storage/s3

Deep Dive into running Hyper Parameter Optimization on AWS SageMaker

cloud/aws/sagemaker workflow/hpo library/xgboost library/cuml library/cupy library/cudf library/dask data-storage/s3 data-format/parquet