How to get a base image for Python 3.9+ when running Sagemaker training jobs in script mode?

I'm using Sagemaker in AWS Sagemaker Studio to try to run training jobs in script mode. Below is an example code snippet I'm trying to run:

from sagemaker.xgboost.estimator import XGBoost
est = XGBoost(
    entry_point=path_to_script,
    instance_type=instance_type,
    framework_version="1.7-1", #latest
    role=exec_role,
    instance_count=1,
    dependencies=[path_to_requirements],
    sagemaker_session=sm_session,
)
est.fit()

The script at path_to_script is written in Python 3.9+. The training job fails at the time of image creation because it seems XGBoost image is running Python 3.8. How do I pick an image corresponding to 3.9+? Note: The script I'm running does not really depend on XGBoost. I'm using it to run my own training script following the script mode example found here: https://sagemaker-examples.readthedocs.io/en/latest/sagemaker-script-mode/sagemaker-script-mode.html

I tried different versions of XGBoost, the latest one only supports Python 3.8 I also looked up the available images in Sagemaker. I can see Sagemaker has different training images as found here: https://docs.aws.amazon.com/sagemaker/latest/dg/notebooks-available-images.html

but it is not clear to me how to pick one of them to launch a training job.

Solution

Instead of using sagemaker.xgboost.estimator.XGBoost you can also use the generic sagemaker.estimator.Estimator (full documentation available here). When using this estimator, you need to provide an image_uri and can also provide a specific py_version.

You can retrieve the corresponding image URI via sagemaker.image_uris.retrieve.

An example with XGBoost could look like the following:

image_uri = sagemaker.image_uris.retrieve("xgboost", sess.boto_region_name, "1.7-1")

xgb = sagemaker.estimator.Estimator(
    image_uri,
    exec_role,
    py_version="3.10"
    entry_point=path_to_script,
    instance_count=1,
    instance_type=instance_type,
    dependencies=[path_to_requirements],
    sagemaker_session=sm_session,
)

Alternatively, you can use one of the other pre-built training images, and install XGBoost as a dependency. Lastly, you can also build your own XGBoost container.