Configuring Local Mode Execution in Sagemaker Studio

Towards the end of 2023, AWS announced a new feature in Sageamker called local mode which allows. developers and data scientists to run SageMaker jobs directly on their local development environment or SageMaker Studio IDE instance. This feature is particularly handy for testing and debugging machine learning workflows before deploying them to the full SageMaker managed environment. [1]

This enables the execution of preprocessing, training and batch transform jobs without provisioning remote compute instances, allowing quick iterations and feedback on code changes in development/testing phase.

However, in the past I have encountered errors in configuring local mode to work correctly in Sagemaker studio with no single document source detailing the series of steps to be carried out in order for successful execution. The purpose of this blog is to detail the steps to successfully install docker in Sageamker Studio and run Sagemaker pipeline in local mode. An example dataset and workflow from the Sagemaker docs will be used to illustrate this.

Create Sagemaker Domain and Default User for Studio

Before using Sagemaker Studio, you will need to create a domain (and associated user profile) NB if you already have a user setup and have used studio before, you can skip to the next section.

In Sagemaker console click domain option in the left navigation bar, select create domain and setup for single user (quick access).
Wait for the domain to setup, it should take a few minutes for the status to change from Pending status to InService

Once the domain status updates to InService (you may need to refresh the browser), click on the domain name

Setup Docker for use in Local Mode

By default, local mode and Docker is not enabled in Studio. To enable this, you need to update the domain created. From the cli, run the following command [2]. You can get the domain id for the domain created in the previous step from the Sagemaker console in the domain tab.

Now navigate back to the Sagemaker main console and select Studio from the left pane. It should now show up with a default user profile (associated with the domain which was setup).Select this from the dropdown and click Open Studio

LocalMode is currently only supported in following applications in Studio: JupyterLab, Studio Classic and Code Editor [2]. We will use Studio Classic for this demonstration. Follow the instructions in [3] to start a studio classic notebook. The default Data Science Image, Python 3 kernel and ml.t3.medium instance type will be sufficient.
Once the notebook is open and Data Science kernel setup and activated, run the following command cat /etc/os-release in the first cell and you should see output similar as below.

Based on the VERSION_CODENAME from the output above, you will need to copy the code from the respective shell scripts here. In this case, I have jammy so I will need the sagemaker-ubuntu-jammy-docker-cli-install.sh. If your VERSION_CODENAME is bullseye (on debian), select the sagemaker-debian-bullseye-cli-install.sh
In your local workspace in Jupyter lab in studio, create a docker_cli_install.sh file and copy the code into that file. Note you may need to prepend sudo to a few commands to avoid permissions denied error when running from the terminal

The top menu of the notebook will have the shell script icon (red arrow). Click in this. Note do not launch the terminal from the file menu, as this will not have root access and you will get an error when trying to start the docker daemon once docker is installed (I still have not figured out how to do this without root access).

Execute the script . docker_cli_install.sh . If this executes successfully, you should see the output as below with the docker cli installed.

You can confirm docker is installed and the docker daemon is up and running, by running a few docker command e.g. docker images , docker ps -a
Navigate back to the Studio Notebook, and run !pip install -U sagemaker in the next cell. This will install the latest version of sagemaker sdk library. If you see conflicts with some other libraries like auto-gluon etc you can ignore them as we will not require those libraries in this example.

Now restart the jupyter kernel by clicking the kernel tab at the top and selecting Restart kernel

Executing Demo Pipeline in Local Mode

We will use the Census-Income-KDD-dataset and the preprocessing and training steps in the aws-sample-notebook with a few modifications to the code to run these as a Sagemaker Pipeline in local mode.

First we will create the preprocess.py and train.py scripts which will be executed in the preprocessing and training jobs. Copy the cells containing %writefile preprocessing.py and %writefile train.py from the sample-notebook [4].

This should create a train.py and preprocess.py script with the respective code in the working directory of your studio notebook.

In the next cell, paste the following code [4]. This imports the modules and initialises LocalPipelineSession object to be passed to the pipeline’s and pipeline steps’ initializer [5]. This replaces the PipelineSession which you would normally use when executing the pipeline on the managed SageMaker Pipelines service.

We also define a Sagemaker execution role, which is created by calling get_execution_role . This creates a new Sagemaker IAM role, with default permissions (which includes Sagemaker Full Access). If you have an existing Sagemaker role with fine grained permissions (S3 access), then you can pass the arn for that to the role variable instead.

Paste the following code in the next cell, which is adapted from [4]. This will define the preprocess step, with input to the raw data in the public S3 bucket (and destination path in the preprocessing container which the script will read from) and the output paths for the test and train data (which are saved in the /opt/ml/preprocessing in the container). Sagemaker will output the split datasets into S3 bucket when this job completes.

Define the training step, which executes the train.py script which trains a logistic regression model on the training data. Once training is complete, the model is saved to /opt/ml/model directory. Sagemaker compresses and uploaded to S3 at the end of the training job.

Finally, the code block below will create a pipeline with the steps defined above and setting the sagemaker_session argument to local_pipeline_session object defined previously to execute in local mode [5].

The first time the pipeline runs, it authenticates docker to ECR and pulls the base sagemaker-sklearn image, Subsequent pipeline re-runs will be faster. This pipeline should finish running in less than a minute and faster compared to if we were to use the fully managed Sagemaker service as we are not waiting for the compute to be provisioned for each job. If the pipeline has executed successfully, you will see the artifacts (code, data, model) from the jobs in a /tmp folder in the same working directory as the notebook. You can check the logs in the notebook, to find which subfolders the model, test/train data are copied into.

Limitations

Not all pipeline steps are supported in local mode. The list of the steps supported are listed in [5].
Distributed training is not supported
Some built in algorithms like Xgboost are not supported in local mode. You will need to use Script Mode instead.
Only use local mode for debugging and testing purposes on smaller datasets. Once tested and ready to move to production, shift to Sagemaker managed service which is scalable, offers parallelisation capabilities (local mode executes sequentually) and flexibility of compute (e.g. GPU, compute optimised) .

References

Sagemaker Local mode announcement https://aws.amazon.com/about-aws/whats-new/2023/12/sagemaker-studio-local-mode-docker/
Local Mode Support in Sagemaker Studiohttps://docs.aws.amazon.com/sagemaker/latest/dg/studio-updated-local.html
Open Studio Classic Notebook https://docs.aws.amazon.com/sagemaker/latest/dg/notebooks-create-open.html
Sageamaker sample https://github.com/aws-samples/amazon-sagemaker-develop-your-ml-project/blob/main/sagemaker-pipelines-preprocess-train-evaluate-register.ipynb
Sagemaker pipelines in local mode https://docs.aws.amazon.com/sagemaker/latest/dg/pipelines-local-mode.html
Local Mode blog https://aws.amazon.com/blogs/machine-learning/use-the-amazon-sagemaker-local-mode-to-train-on-your-notebook-instance/

Site Terms, Privacy, and more.