<aside> 📢 This is as a basic, living, design documentation for the up-coming integration revamp between ZenML and Sagemaker. Please feel free to comment below your thoughts.
</aside>
The SageMakerStepOperator in ZenML currently uses a default Estimator for Training Jobs. However, there are a broader range of Training APIs available in SageMaker. In this document, we will discuss the extension of the SageMakerStepOperator to support additional estimator arguments and training data inputs.
Related issue: https://github.com/zenml-io/zenml/issues/1398
This feature will benefit ZenML users who require more advanced Training APIs for their data processing and training steps. This feature will provide users with a broader range of options for their data processing and training needs.
The design of the extension will involve modifying the current SageMakerStepOperator to support additional estimator arguments and training data inputs. This will be done by adding support for the various training APIs available in SageMaker. The Training Job configuration will be updated to include the new estimator arguments and training data inputs.
The proposed feature is to expand the SageMakerStepOperator to support additional estimator arguments and various training data inputs, such as AutoML, specific Estimators (TensorFlow, PyTorch, HuggingFace, ...), pre-uploaded/third-party algorithms (Algorithm Estimators), hyperparameter tuning (HyperparameterTuner), and processing jobs for data processing steps.
Initially, we start with extending support to the Tensorflow Estimator:
TensorFlow — sagemaker 2.144.0 documentation
The implementation of this feature will involve creating a new branch in the ZenML repository and modifying the SageMakerStepOperator to support the additional estimator arguments and training data inputs.
The proposal will be open to the public, but please suggest a few experience contributors in this technical domain whose comments will help this proposal. Ideally, the list should include ZenML committers.
@Alex Strick van Linschoten @Michael Schuster