SageMakerCore Overview of Resource Level Abstractions - XGBoost Training Example

Introductions

SageMakerCore is a Python SDK designed as a lightweight layer over boto3, the AWS SDK for Python. It is built on the concept of resource level abstractions, where SageMaker Resources are represented as Python classes. This approach enables SageMakerCore to simplify the management of SageMaker Resources and provide a more object-oriented programming interface.

Resource Level Abstraction

Resource Level Abstractions can be best understood by examining how the AWS TrainingJob APIs are transfromed into a TrainingJob Python class abstraction in SageMakerCore.

For instance, an AWS TrainingJob has the following APIs:

1. CreateTrainingJob
2. DescribeTrainingJob
3. UpdateTrainingJob
4. StopTrainingJob
5. ListTrainingJobs

In SageMakerCore, these APIs are encapsulated within a TrainingJob class that exposes these operations as methods and attributes. The details of the TrainingJob class are below:

class TrainingJob(Base):
    # Class attributes are mapped to describe_training_job response
    training_job_name: str
    training_job_arn: Optional[str] = Unassigned()
    tuning_job_arn: Optional[str] = Unassigned()
    labeling_job_arn: Optional[str] = Unassigned()
    auto_ml_job_arn: Optional[str] = Unassigned()
    model_artifacts: Optional[ModelArtifacts] = Unassigned()
    training_job_status: Optional[str] = Unassigned()
    ...

    @classmethod
    def create():       # Calls `create_training_job`

    @classmethod
    def get():          # Calls `describe_training_job`

    @classmethod
    def get_all():      # Calls `list_training_job`

    
    def update():       # Calls `update_training_job`


    def stop():         # Calls `stop_training_job`


    def refresh():      # Calls `describe_training_job` and refreshes instance attributes


    def wait():         # Calls `describe_training_job` and waits for TrainingJob to enter terminal state

Comparing Boto3 and SageMakerCore SDKs

In this notebook, we create an AWS TrainingJob to train an XGBoost Container. We will be using both Boto3 and the SageMakerCore SDKs with the goal of highlighting and comparing the differences in user experience for performing operations such as creating, updating, waiting, and listing AWS TrainingJobs.

Install Latest SageMakerCore

All SageMakerCore beta distributions will be released to a private s3 bucket. After being allowlisted, run the cells below to install the latest version of SageMakerCore from s3://sagemaker-core-beta-artifacts/sagemaker_core-latest.tar.gz

Ensure you are using a kernel with python version >=3.8

Install Additional Packages

Setup

Let's start by specifying:

• AWS region.
• The IAM role arn used to give learning and hosting access to your data. Ensure your enviornment has AWS Credentials configured.
• The S3 bucket that you want to use for storing training and model data.

Load and Prepare Dataset

For this example, we will be using the IRIS data set from sklearn.datasets to train our XGBoost container.

Upload Data to S3

In this step, we will upload the train and test data to the S3 bucket configured earlier using sagemaker_session.default_bucket()

Fetch the XGBoost Image URI

In this step, we will fetch the XGBoost Image URI we will use as an input parameter when creating an AWS TrainingJob

Create TrainingJob with Boto3

With the necessary setup completed, we can now create an AWS TrainingJob. First we will begin by creating a TrainingJob with Boto3 to understand what the experience is like when interecting directly with low-level APIs through Boto3.

When executing the following cells there are a few things to note about the experience with Boto3:

1. Boto3 dynamically generates the API operation methods like create_training_job. When a client is instantiated, the methods are generated from the JSON service model description and are not statically coded into the boto3 library.
2. Boto3 returns a JSON response. As a result, users must either be familiar with the structure of these responses or refer to the documentation to parse them correctly.
3. Boto3 client methods expect keyword arguments. Similar to the experience with JSON response, users must be familiar with what keyword argumnets are expected or refer to the documentation to pass them correctly.

Wait for TrainingJob with Boto3

When a user creates a TrainingJob it is often the case that they would wish to wait on the TrainingJob to complete. Below is an example of how a user wait on a TrainingJob using Boto3. Notebly, this requires creating some logic to poll the TrainingJob using describe_training_job until the TrainingJobStatus is 'Failed', 'Completed', or 'Stopped'.

Create TrainingJob with SageMakerCore

In this step we will use SageMakerCore to create a TrainingJob to understand what experience the object-oriented resource level abstractions provide for users.

When executing the following cells, there are a few things to note about the experience with SageMakerCore:

1. SageMakerCore generates Python classes and methods from the service model JSON, similar to Boto3. However, this generation is done prior to a release, resulting in a statically coded interface in the library.
2. SageMakerCore adopts an object-oriented approach, providing users with clear visibility of available methods and attributes through type hinting and IDE IntelliSense
3. Instead of returning JSON responses like Boto3, SageMakerCore returns objects. This allows users to access response attributes directly from the returned object, eliminating the need to parse JSON or refer to the documentation for structure details.

Wait for TrainingJob with SageMakerCore

In SageMakerCore, the logic required to wait on a resource is abstracted away using a wait() method. As a result, a user can directly call the wait() method on a TrainingJob object instance like below.

List TrainingJobs with Boto3

When a user lists TrainingJobs, there are 2 main approaches provided by Boto3.

1. The first is calling list_training_jobs directly and implementing some logic to handle the NextToken provided in the response to enable pagination.
2. The second is by utilizing the Boto3 get_paginator method to get a paginator that encapsulates the NextToken and simplifies the logic required.

Both approaches are shown below. Although the boto3 provided paginator simplifies the logic over using a NextToken, in both cases the user must understand the structure of the list responses or refer to the docs (ie, understand to access TrainingJobSummaries by doing response["TrainingJobSummaries"])

List TrainingJobs with SageMakerCore

In SageMakerCore, listing is done similar to the boto3 paginator approach but instead with a ResourceIterator which implements the python iterator protocol to instantiate and return resource objects only as they are accessed.

Below, is an example of how the get_all() method would be used to list TrainingJobs.

Delete All SageMaker Resources

The following code block will call the delete() method for any SageMaker Core Resources created during the execution of this notebook which were assigned to local or global variables. If you created any additional deleteable resources without assigning the returning object to a unique variable, you will need to delete the resource manually by doing something like:

resource = Resource.get("resource-name")
resource.delete()