Support larger memory instance types for RepackModelStep

[michael-lum]

Describe the feature you'd like
Please either provide a larger default memory instance type or allow for configuration for the instance type of the RepackModelStep. This step is currently added automatically for the CreateModel step if custom code is provided but does not allow configuration of instance type.

How would this feature be used? Please describe.
For Large Language Models, model sizes can be larger than 30MB. This would allow providing custom code for model artifacts of large sizes with an automated repackaging step.

Describe alternatives you've considered
The only other alternative is to manually package the inference and loading code directly into the model zip or implementing the packaging logic ourself.

Additional context
Code Link

[aaronpavez-woebot]

We ran into this problem as well. I believe the system memory for the instance is fine, but there needs to be more room on the attached EBS volume for the instance to unpack the model, copy it around, and repack it.

We were able to successfully work around it in the short-term by editing the Repack step in the SDK:

sagemaker-python-sdk/src/sagemaker/workflow/model_step.py

Line 271 in 8d5d7fd

output_kms_key=model.model_kms_key,

adding the parameter:
volume_size=50

to increase the EBS volume size to 50 GB.

[ShwetaSingh801]

Hi @michael-lum,

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]

Posted update in another thread but also updating this issue for broader audience:

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,
    )