Allow customizing repack model step model output S3 URI

[l3ku]

Describe the feature you'd like
I am training a HuggingFace estimator in a SageMaker pipeline and registering the model in model registry with a custom inference entrypoint. It seems that the HuggingFace model sets repack=True in the prepare_container_def when it uploads the code using self._upload_code:

https://github.com/aws/sagemaker-python-sdk/blob/master/src/sagemaker/huggingface/model.py#L498

It seems that there is no way to customize the S3 URI where the model produced during the repack step will be stored, other than by setting the default SageMaker bucket and key prefix to some value initially. Ideally, I would like to overwrite the model artifact from the previous training step, which is currently impossible since the S3 URI depends on the training job name.

So currently what happens, is when I provide the output path to the estimator during training:

s3://MODEL_ARTIFACT_BUCKET_NAME/MODEL_NAME/models/

SageMaker pipelines will store the training job model artifact at:

s3://MODEL_ARTIFACT_BUCKET_NAME/MODEL_NAME/models/TRAINING_JOB_NAME/output/model.tar.gz

The repack step (training job) will store the repacked model artifacts at:

s3://SAGEMAKER_DEFAULT_BUCKET/huggingface-pytorch-inference-2023-10-18-14-03-06-845

I am using a centralized S3 bucket intended only for model artifacts (MODEL_ARTIFACT_BUCKET_NAME), hence I would only want to store the model artifacts there and not any uploaded code or other scripts that would eventually end up there if I would use the bucket as the default SageMaker session bucket. Hence the request for the feature to customize only the S3 URI of the repacked model artifacts.

How would this feature be used? Please describe.
The repack step would either use the same model artifacts S3 URI as the previous training step, or it would be possible to provide the output S3 URI of the repacked model artifact either in the constructor of HuggingFaceModel or in the register() method.

Describe alternatives you've considered
No alternatives at the moment.

Additional context

[l3ku]

I noticed that it is possible to change the S3OutputPath of the repack step using the code_location parameter of the HuggingFaceModel. When I set it to

HuggingFaceModel(
    ...
    code_location="s3://MODEL_ARTIFACT_BUCKET_NAME/MODEL_NAME/code",
    ...
)

The model artifact produced by the repack step will be uploaded to:

s3://MODEL_ARTIFACT_BUCKET_NAME/MODEL_NAME/code/huggingface-pytorch-inference-2023-10-19-05-46-58-983

However, then it will also create some additional bloat in that same bucket since it will also upload the _repack_script_launcher.sh and _repack_model.py into the same bucket into a location I can't seem to control:

s3://MODEL_ARTIFACT_BUCKET_NAME/ModelStep-RepackModel-0-652893bd0590941fe8bbd169af113197/source/sourcedir.tar.gz

Anyways, why should the repack step output be controlled by the code_location parameter of the HuggingFaceModel? To me, it doesn't make sense that the training job produces model artifacts somewhere in S3 where the user wants the pipeline to store them, and then the repack model step moves those model artifacts to another location?

[vladyslav-melnyk-ait]

facing the same challenge

[ShwetaSingh801]

Hi @l3ku,

Thanks for using SageMaker and taking the time to suggest ways to improve SageMaker Python SDK. We have added your feature request it to our backlog of feature requests and may consider putting it into future SDK versions. I will go ahead and close the issue now, please let me know if you have any more feedback. Let me know if you have any other questions.

Best,
Shweta

[qidewenwhen]

Hi @l3ku, this feature request was released in v2.204.0.

After upgrading the SDK version, to customize the repack model step, you can use the newly added repack_model_step_settings (please find it in doc)

Below is an example on how to use it:

repack_model_step_settings = {
            "output_path": f"s3://{BUCKET}/repack-model-output",
            "code_location": f"s3://{BUCKET}/repack-model-code-location",
            "instance_type": "ml.m4.xlarge",
            "volume_size": 50,
            "output_kms_key": "repack-model-output-kms-key",
},
ModelStep(
        name="MyModelStep",
        step_args=step_args,
        repack_model_step_settings=repack_model_step_settings,
    )