Skip to content

Added support for Sagemaker ENV variables + custom formatter scripts#2921

Merged
smouaa merged 3 commits intomasterfrom
cx-sagemaker
Oct 30, 2025
Merged

Added support for Sagemaker ENV variables + custom formatter scripts#2921
smouaa merged 3 commits intomasterfrom
cx-sagemaker

Conversation

@smouaa
Copy link
Contributor

@smouaa smouaa commented Oct 27, 2025

Description

  1. Added support for following ENV variables for CX compatibility:

SAGEMAKER_MAX_REQUEST_SIZE
SAGEMAKER_NUM_MODEL_WORKERS
SAGEMAKER_DEFAULT_INVOCATIONS_ACCEPT
SAGEMAKER_MODEL_SERVER_TIMEOUT_SECONDS
SAGEMAKER_MODEL_SERVER_VMARGS
SAGEMAKER_STARTUP_TIMEOUT
SAGEMAKER_MAX_PAYLOAD_IN_MB

  1. Extended custom formatters to work with existing Sagemaker scripts
  2. Integration tests for CX env variables, custom formatters

SKLearn handler notebook: https://github.com/deepjavalibrary/djl-demo/blob/master/djl-serving/python-mode/sklearn_handler_guide.ipynb
XGBoost handler notebook: https://github.com/deepjavalibrary/djl-demo/blob/master/djl-serving/python-mode/xgboost_handler_guide.ipynb

Type of change

Please delete options that are not relevant.

  • Bug fix (non-breaking change which fixes an issue)
  • Breaking change (fix or feature that would cause existing functionality to not work as expected)
  • New feature (non-breaking change which adds functionality)
  • This change requires a documentation update

Checklist:

  • Please add the link of Integration Tests Executor run with related tests.
  • Have you manually built the docker image and verify the change?
  • Have you run related tests? Check how to set up the test environment here; One example would be pytest tests.py -k "TestCorrectnessLmiDist" -m "lmi_dist"
  • Have you added tests that prove your fix is effective or that this feature works?
  • Has code been commented, particularly in hard-to-understand areas?
  • Have you made corresponding changes to the documentation?

Feature/Issue validation/testing

Please describe the Unit or Integration tests that you ran to verify your changes and relevant result summary. Provide instructions so it can be reproduced.
Please also list any relevant details for your test configuration.

  • Test A
    Logs for Test A

  • Test B
    Logs for Test B

@smouaa smouaa requested review from a team and zachgk as code owners October 27, 2025 08:48
@smouaa smouaa force-pushed the cx-sagemaker branch 4 times, most recently from a204f8d to 7dde753 Compare October 27, 2025 17:27
for (Map.Entry<String, String> env : Utils.getenv().entrySet()) {
String key = env.getKey();
// Apply SageMaker compatibility for server-level configurations
SageMakerCompatibility.applyServerCompatibility(prop);
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Do you want this to be higher priority than SERVING_XXX env? If yes, you need to move it after line 141. Otherwise SERVING_MAX_REQUEST_SIZE will override SAGEMAKER_MAX_REQUEST_SIZE.

}

// Apply SageMaker compatibility for model-level configurations
SageMakerCompatibility.applyModelCompatibility(prop);
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Same as above. Move after line 915 if needed.


String defaultAccept = System.getenv("SAGEMAKER_DEFAULT_INVOCATIONS_ACCEPT");
if (defaultAccept != null) {
String entryPoint = properties.getProperty("option.entryPoint");
Copy link
Contributor

@xyang16 xyang16 Oct 30, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Need to move SageMakerCompatibility.applyModelCompatibility(prop) in ModelInfo.java after line 915. Because if user set entrypoint via OPTION_ENTRYPOINT env, it is not populated to properties yet.

@smouaa smouaa dismissed Lokiiiiii’s stale review October 30, 2025 22:00

Requested changes were implemented and approved by Xin. Merge is necessary in order to publish the container by the deadline (10/31)

@smouaa smouaa merged commit 63554fc into master Oct 30, 2025
42 of 49 checks passed
@smouaa smouaa deleted the cx-sagemaker branch October 30, 2025 22:00
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants