{"id":2054,"date":"2022-04-05T17:38:40","date_gmt":"2022-04-05T17:38:40","guid":{"rendered":"https:\/\/salarydistribution.com\/machine-learning\/2022\/04\/05\/customize-the-amazon-sagemaker-xgboost-algorithm-container\/"},"modified":"2022-04-05T17:38:40","modified_gmt":"2022-04-05T17:38:40","slug":"customize-the-amazon-sagemaker-xgboost-algorithm-container","status":"publish","type":"post","link":"https:\/\/salarydistribution.com\/machine-learning\/2022\/04\/05\/customize-the-amazon-sagemaker-xgboost-algorithm-container\/","title":{"rendered":"Customize the Amazon SageMaker XGBoost algorithm container"},"content":{"rendered":"
\n

The built-in Amazon SageMaker<\/a> XGBoost algorithm provides a managed container to run the popular XGBoost<\/a> machine learning (ML) framework, with added convenience of supporting advanced training or inference features like distributed training, dataset sharding for large-scale datasets, A\/B model testing<\/a>, or multi-model inference<\/a> endpoints. You can also extend this powerful algorithm to accommodate different requirements.<\/p>\n

Packaging the code and dependencies in a single container is a convenient and robust approach for long-term code maintenance, reproducibility, and auditing purposes. Modifying the container directly follows the base container faithfully and avoids duplicating existing functions already supported by the base container. In this post, we review the inner workings of the SageMaker XGBoost algorithm container and provide pragmatic scripts to directly customize the container.<\/p>\n

SageMaker XGBoost container structure<\/h2>\n

The SageMaker built-in XGBoost algorithm is packaged as a stand-alone container, available on GitHub<\/a>, and can be extended under the developer-friendly Apache 2.0 open-source license. The container packages the open-source XGBoost algorithm<\/a> and ancillary tools to run the algorithm in the SageMaker environment integrated with other AWS Cloud services. This allows you to train XGBoost models on a variety of data sources<\/a>, make batch predictions<\/a> on offline data, or host an inference endpoint<\/a> in a real-time pipeline<\/a>.<\/p>\n

The container supports training and inference operations with different entry points. For inference mode, the entry can be found in the main function in the serving.py script<\/a>. For real-time inference serving, the container runs a Flask<\/a>-based web server<\/a> that when invoked<\/a>, receives an HTTP-encoded request containing the data, decodes the data into the XGBoost\u2019s DMatrix<\/a> format, loads the model<\/a>, and returns an HTTP-encoded response back<\/a>. These methods are encapsulated under the ScoringService<\/a> class, which can also be customized through the script mode to a great extent (see the Appendix below).<\/p>\n

The entry point for training mode (algorithm mode) is the main function in the training.py<\/a>. The main function sets up the training environment and calls the training job function. It\u2019s flexible enough to allow for distributed or single-node training, or utilities like cross validation. The heart of the training process can be found in the train_job<\/a> function.<\/p>\n

Docker files packaging the container can be found in the GitHub repo<\/a>. Note that the container is built in two steps: a base<\/a> container is built first, followed by the final<\/a> container on top.<\/p>\n

Solution overview<\/h2>\n

You can modify and rebuild the container through the source code. However, this involves collecting and rebuilding all dependencies and packages from scratch. In this post, we discuss a more straightforward approach that modifies the container on top of the already-built and publicly-available SageMaker XGBoost algorithm container image directly.<\/p>\n

In this approach, we pull<\/a> a copy of the public SageMaker XGBoost image, modify the scripts or add packages, and rebuild the container on top. The modified container can be stored in a private repository. This way, we avoid rebuilding intermediary dependencies and instead build directly on top of the already-built libraries packaged in the official container.<\/p>\n

The following figure shows an overview of the script used to pull the public base image, modify and rebuild the image, and upload it to a private Amazon Elastic Container Registry<\/a> (Amazon ECR) repository. The bash script<\/a> in the accompanying code of this post performs all the workflow steps shown in the diagram. The accompanying notebook<\/a> shows an example where the URI of a specific version of the SageMaker XGBoost algorithm is first retrieved and passed to the bash script<\/a>, which replaces two of the Python scripts in the image, rebuilds it, and pushes the modified image to a private Amazon ECR repository. You can modify the accompanying code to suit your needs.<\/p>\n

\u00ad\"\"<\/a><\/p>\n

Prerequisites<\/h2>\n

The GitHub repository<\/a> contains the code accompanying this post. You can run the sample notebook<\/a> in your AWS account, or use the provided AWS CloudFormation<\/a> stack to deploy the notebook using a SageMaker notebook. You need the following prerequisites:<\/p>\n