Create a custom AWS Sagemaker Estimator with a configurable entrypoint - amazon-sagemaker

Question

I'm writing a custom Estimator in AWS Sagemaker , for a framework that isn't supported out-of-the-box. I have my own docker image for training, with the training code bundled into the image, which forces me to rebuild the image every time the code changes.
What I would like to do, is create an Estimator that uses this image, and accepts a file as an entry-point, like the built-in framework estimators do (e.g Tensorflow).
From reading the source code for the Sagemaker python SDK, I found the sagemaker.estimator.Framework class, which accepts the entry_point argument, and from which the built-in frameworks estimators inherit. However, the documentation doesn't really show how to inherit from that class in my own code.
Is it possible to write a custom Estimator class that inherits from Framework, or is there another way to create a custom estimator that receives an entry-point argument?

Answer 1

Instead of using an Estimator, use a TensorFlow estimator and then pass your custom image as an input to image_uri. This way you have the best of the both worlds. And if you have a requirements.txt file in your source_dir, that gets installed before training.
Here is a custom image docker file
FROM python:3.9-buster
RUN apt-get -y update && apt-get install -y --no-install-recommends \
wget \
python3-pip \
python3-setuptools \
nginx \
ca-certificates \
vim \
&& rm -rf /var/lib/apt/lists/*
#RUN ln -s /usr/bin/python3 /usr/bin/python
RUN ln -s /usr/bin/pip3 /usr/bin/pip
ENV PYTHONUNBUFFERED=TRUE
ENV PYTHONDONTWRITEBYTECODE=TRUE
ENV PATH="/opt/program:${PATH}"
RUN pip install sagemaker-training
and here is how you call your custom image
tf_estimator = TensorFlow(entry_point='mnist_keras_tf.py',
source_dir='.',
role=role,
instance_count=1,
instance_type='ml.p3.2xlarge',
image_uri='<arn of custom image>',
script_mode=True,
hyperparameters={
'epochs': 2,
'batch-size': 256,
'learning-rate': 0.01}

Answer 2

The existing framework estimators in the SageMaker Python SDK can be a good place to start, e.g. https://github.com/aws/sagemaker-python-sdk/blob/master/src/sagemaker/pytorch/estimator.py
However, it would probably be easier, since you've already bundled your training code into the image, to set an entrypoint in your Dockerfile that runs the training code. For more about how Amazon SageMaker runs your training image, see https://docs.aws.amazon.com/sagemaker/latest/dg/your-algorithms-training-algo-dockerfile.html