Allow customizing repack model step model output S3 URI #4218
Description
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:
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.