From 3a66fab5ede6e661946285c2b302a9e212633fd4 Mon Sep 17 00:00:00 2001 From: Jyothirmai Kottu Date: Thu, 9 Jul 2026 17:29:41 +0000 Subject: [PATCH 1/7] test(sglang): add FLUX.2 diffusion inference tests Add FLUX.2 model inference testing for the SGLang DLC image, which already ships the sglang.multimodal_gen diffusion stack (diffusers, etc.) but previously had no diffusion test coverage or SageMaker serving path. - model-tests: validate FLUX.2-small-decoder (AutoencoderKLFlux2) via a VAE decode test, since the distilled decoder is a pipeline component rather than a servable text-to-image model. - sagemaker_entrypoint.sh: add an SM_SGLANG_ENGINE gate (default "llm", unchanged) so SM_SGLANG_ENGINE=diffusion launches sglang.multimodal_gen.runtime.launch_server instead of the LLM engine. - sagemaker test: deploy FLUX.2-klein-4B via the diffusion engine and assert /v1/images/generations returns image data. --- .../config/model-tests/sglang-model-tests.yml | 10 +++ scripts/docker/sglang/sagemaker_entrypoint.sh | 20 ++++- test/sglang/sagemaker/test_sm_endpoint.py | 82 +++++++++++++++++++ .../scripts/sglang_flux_vae_decode_test.sh | 44 ++++++++++ 4 files changed, 154 insertions(+), 2 deletions(-) create mode 100755 test/sglang/scripts/sglang_flux_vae_decode_test.sh diff --git a/.github/config/model-tests/sglang-model-tests.yml b/.github/config/model-tests/sglang-model-tests.yml index 5197e46046ab..555717d905f8 100644 --- a/.github/config/model-tests/sglang-model-tests.yml +++ b/.github/config/model-tests/sglang-model-tests.yml @@ -38,6 +38,16 @@ smoke-test: fleet: "x86-g6xl-runner" extra_args: "--tp 1 --context-length 4096 --dtype bfloat16" + # FLUX.2-small-decoder is a distilled VAE decoder (AutoencoderKLFlux2), a + # FLUX.2 pipeline component rather than a servable model — validated via a + # decode test (see test_script) instead of the default LLM serving smoke + # test. Full-pipeline serving is covered by the SageMaker endpoint test. + - name: "flux2-small-decoder" + s3_model: "flux2-small-decoder.tar.gz" + fleet: "x86-g6xl-runner" + test_script: "sglang_flux_vae_decode_test.sh" + extra_args: "" + runner-scale-sets: [] benchmark: diff --git a/scripts/docker/sglang/sagemaker_entrypoint.sh b/scripts/docker/sglang/sagemaker_entrypoint.sh index 23a2fff81e59..359affa05ab1 100755 --- a/scripts/docker/sglang/sagemaker_entrypoint.sh +++ b/scripts/docker/sglang/sagemaker_entrypoint.sh @@ -12,9 +12,19 @@ echo "Starting server" PREFIX="SM_SGLANG_" ARG_PREFIX="--" +# Engine selector (default: llm). Set SM_SGLANG_ENGINE=diffusion to serve a +# FLUX.2 / diffusion pipeline via sglang.multimodal_gen instead of the LLM +# engine. This var controls the launch module and is NOT forwarded as a flag. +ENGINE=$(echo "${SM_SGLANG_ENGINE:-llm}" | tr '[:upper:]' '[:lower:]') + ARGS=() while IFS='=' read -r key value; do + # SM_SGLANG_ENGINE selects the launch module; it is not a server flag. + if [ "$key" = "${PREFIX}ENGINE" ]; then + continue + fi + arg_name=$(echo "${key#"${PREFIX}"}" | tr '[:upper:]' '[:lower:]' | tr '_' '-') # Handle boolean flags: true -> flag only, false -> skip entirely @@ -46,5 +56,11 @@ if ! [[ " ${ARGS[@]} " =~ " --model-path " ]]; then ARGS+=(--model-path "${SM_SGLANG_MODEL_PATH:-/opt/ml/model}") fi -echo "Running command: exec python3 -m sglang.launch_server ${ARGS[@]}" -exec python3 -m sglang.launch_server "${ARGS[@]}" +if [ "$ENGINE" = "diffusion" ]; then + LAUNCH_MODULE="sglang.multimodal_gen.runtime.launch_server" +else + LAUNCH_MODULE="sglang.launch_server" +fi + +echo "Running command: exec python3 -m ${LAUNCH_MODULE} ${ARGS[@]}" +exec python3 -m "${LAUNCH_MODULE}" "${ARGS[@]}" diff --git a/test/sglang/sagemaker/test_sm_endpoint.py b/test/sglang/sagemaker/test_sm_endpoint.py index eac323d47d97..38f0f3abc34c 100644 --- a/test/sglang/sagemaker/test_sm_endpoint.py +++ b/test/sglang/sagemaker/test_sm_endpoint.py @@ -118,3 +118,85 @@ def test_sglang_sagemaker_endpoint(model_endpoint, model_id): LOGGER.info(f"Model response: {pformat(body)}") LOGGER.info("Inference test successful!") + + +@pytest.fixture(scope="function") +def flux_endpoint(aws_session, image_uri): + """Deploy a FLUX.2 diffusion endpoint via the sglang.multimodal_gen engine.""" + model_id = "black-forest-labs/FLUX.2-klein-4B" + instance_type = "ml.g6e.xlarge" + + cleaned_id = clean_string(model_id.split("/")[1], "_./") + endpoint_name = random_suffix_name(f"sglang-{cleaned_id}", 50) + model_name = endpoint_name + + hf_token = get_hf_token(aws_session) + role_arn = aws_session.resolve_role_arn(SAGEMAKER_ROLE) + + model = endpoint_config = endpoint = None + try: + LOGGER.info(f"Creating FLUX.2 model: {model_name}") + model = Model.create( + model_name=model_name, + primary_container=ContainerDefinition( + image=image_uri, + environment={ + "SM_SGLANG_MODEL_PATH": model_id, + "SM_SGLANG_ENGINE": "diffusion", + "HF_TOKEN": hf_token, + }, + ), + execution_role_arn=role_arn, + ) + + LOGGER.info(f"Creating endpoint config: {endpoint_name}") + endpoint_config = EndpointConfig.create( + endpoint_config_name=endpoint_name, + production_variants=[ + ProductionVariant( + variant_name="AllTraffic", + model_name=model_name, + initial_instance_count=1, + instance_type=instance_type, + inference_ami_version=INFERENCE_AMI_VERSION, + container_startup_health_check_timeout_in_seconds=900, + ), + ], + ) + + LOGGER.info(f"Deploying endpoint: {endpoint_name} (this may take 10-15 minutes)...") + endpoint = Endpoint.create( + endpoint_name=endpoint_name, + endpoint_config_name=endpoint_name, + ) + endpoint.wait_for_status("InService") + LOGGER.info("FLUX.2 endpoint deployment completed successfully") + + yield endpoint + finally: + _cleanup([endpoint, endpoint_config, model]) + + +def test_sglang_sagemaker_flux_diffusion_endpoint(flux_endpoint): + """FLUX.2 text-to-image generation through the SageMaker diffusion endpoint.""" + endpoint = flux_endpoint + + payload = json.dumps( + { + "prompt": "a red cube on a white table", + "num_inference_steps": 4, + "width": 512, + "height": 512, + } + ) + LOGGER.debug(f"Sending image-generation request with payload: {payload}") + + result = endpoint.invoke(body=payload, content_type="application/json") + body = json.loads(result.body.read()) + LOGGER.info("Image-generation request invoked successfully") + + assert body, "Model response is empty, failing FLUX.2 endpoint test!" + # OpenAI-style image response returns a non-empty `data` list of images. + assert body.get("data"), f"No image data in FLUX.2 response: {pformat(body)}" + + LOGGER.info("FLUX.2 diffusion inference test successful!") diff --git a/test/sglang/scripts/sglang_flux_vae_decode_test.sh b/test/sglang/scripts/sglang_flux_vae_decode_test.sh new file mode 100755 index 000000000000..e3ef66b818e2 --- /dev/null +++ b/test/sglang/scripts/sglang_flux_vae_decode_test.sh @@ -0,0 +1,44 @@ +#!/bin/bash +set -euo pipefail + +# SGLang FLUX.2 VAE decode test +# +# FLUX.2-small-decoder is a distilled VAE decoder (class AutoencoderKLFlux2), a +# drop-in decoder for the FLUX.2 diffusion pipeline. It is a pipeline COMPONENT, +# not a servable text-to-image model, so this validates it the way it is used: +# load the decoder and decode a latent into an image tensor. End-to-end serving +# of a full FLUX.2 pipeline is covered by the SageMaker endpoint test. +# +# Usage: sglang_flux_vae_decode_test.sh [extra_args...] + +MODEL_DIR="${1:?Usage: $0 [extra_args...]}" +MODEL_NAME="${2:?Usage: $0 [extra_args...]}" + +echo "=== Model directory: ${MODEL_DIR} ===" +ls -la "${MODEL_DIR}" + +echo "=== Loading AutoencoderKLFlux2 and running a decode ===" +python3 - "${MODEL_DIR}" "${MODEL_NAME}" <<'PY' +import sys + +import torch +from diffusers import AutoencoderKLFlux2 + +model_dir, model_name = sys.argv[1], sys.argv[2] +device = "cuda" if torch.cuda.is_available() else "cpu" + +vae = AutoencoderKLFlux2.from_pretrained(model_dir).to(device).eval() +n_params = sum(p.numel() for p in vae.parameters()) +print(f"Loaded {type(vae).__name__} params={n_params} device={device}") + +# Build a latent that matches the decoder's expected channel count, then decode. +latent_channels = vae.config.latent_channels +latent = torch.randn(1, latent_channels, 16, 16, device=device, dtype=next(vae.parameters()).dtype) +with torch.no_grad(): + image = vae.decode(latent).sample + +assert image.ndim == 4 and image.shape[0] == 1, f"unexpected decode output shape {tuple(image.shape)}" +assert torch.isfinite(image).all(), "decode produced non-finite values" +print(f"Decoded latent {tuple(latent.shape)} -> image {tuple(image.shape)}") +print(f"=== PASSED: {model_name} ===") +PY From 4175faa7fd369d16748d8844f5cc40d7234f53a5 Mon Sep 17 00:00:00 2001 From: Jyothirmai Kottu Date: Thu, 9 Jul 2026 17:38:33 +0000 Subject: [PATCH 2/7] test(sglang): scope FLUX.2 tests to the small-decoder VAE decode only Revert the SageMaker diffusion endpoint test and the SM_SGLANG_ENGINE entrypoint gate. Those served FLUX.2-klein-4B (the full pipeline), not FLUX.2-small-decoder, which is what this change is about. The distilled VAE decoder is a pipeline component and is validated by the model-tests VAE decode test; no entrypoint change is needed for it. --- .../config/model-tests/sglang-model-tests.yml | 2 +- scripts/docker/sglang/sagemaker_entrypoint.sh | 20 +---- test/sglang/sagemaker/test_sm_endpoint.py | 82 ------------------- .../scripts/sglang_flux_vae_decode_test.sh | 3 +- 4 files changed, 4 insertions(+), 103 deletions(-) diff --git a/.github/config/model-tests/sglang-model-tests.yml b/.github/config/model-tests/sglang-model-tests.yml index 0e6c53760bb1..ba64be0b9fe5 100644 --- a/.github/config/model-tests/sglang-model-tests.yml +++ b/.github/config/model-tests/sglang-model-tests.yml @@ -41,7 +41,7 @@ smoke-test: # FLUX.2-small-decoder is a distilled VAE decoder (AutoencoderKLFlux2), a # FLUX.2 pipeline component rather than a servable model — validated via a # decode test (see test_script) instead of the default LLM serving smoke - # test. Full-pipeline serving is covered by the SageMaker endpoint test. + # test. - name: "flux2-small-decoder" s3_model: "flux2-small-decoder.tar.gz" fleet: "x86-g6xl-runner" diff --git a/scripts/docker/sglang/sagemaker_entrypoint.sh b/scripts/docker/sglang/sagemaker_entrypoint.sh index 359affa05ab1..23a2fff81e59 100755 --- a/scripts/docker/sglang/sagemaker_entrypoint.sh +++ b/scripts/docker/sglang/sagemaker_entrypoint.sh @@ -12,19 +12,9 @@ echo "Starting server" PREFIX="SM_SGLANG_" ARG_PREFIX="--" -# Engine selector (default: llm). Set SM_SGLANG_ENGINE=diffusion to serve a -# FLUX.2 / diffusion pipeline via sglang.multimodal_gen instead of the LLM -# engine. This var controls the launch module and is NOT forwarded as a flag. -ENGINE=$(echo "${SM_SGLANG_ENGINE:-llm}" | tr '[:upper:]' '[:lower:]') - ARGS=() while IFS='=' read -r key value; do - # SM_SGLANG_ENGINE selects the launch module; it is not a server flag. - if [ "$key" = "${PREFIX}ENGINE" ]; then - continue - fi - arg_name=$(echo "${key#"${PREFIX}"}" | tr '[:upper:]' '[:lower:]' | tr '_' '-') # Handle boolean flags: true -> flag only, false -> skip entirely @@ -56,11 +46,5 @@ if ! [[ " ${ARGS[@]} " =~ " --model-path " ]]; then ARGS+=(--model-path "${SM_SGLANG_MODEL_PATH:-/opt/ml/model}") fi -if [ "$ENGINE" = "diffusion" ]; then - LAUNCH_MODULE="sglang.multimodal_gen.runtime.launch_server" -else - LAUNCH_MODULE="sglang.launch_server" -fi - -echo "Running command: exec python3 -m ${LAUNCH_MODULE} ${ARGS[@]}" -exec python3 -m "${LAUNCH_MODULE}" "${ARGS[@]}" +echo "Running command: exec python3 -m sglang.launch_server ${ARGS[@]}" +exec python3 -m sglang.launch_server "${ARGS[@]}" diff --git a/test/sglang/sagemaker/test_sm_endpoint.py b/test/sglang/sagemaker/test_sm_endpoint.py index 38f0f3abc34c..eac323d47d97 100644 --- a/test/sglang/sagemaker/test_sm_endpoint.py +++ b/test/sglang/sagemaker/test_sm_endpoint.py @@ -118,85 +118,3 @@ def test_sglang_sagemaker_endpoint(model_endpoint, model_id): LOGGER.info(f"Model response: {pformat(body)}") LOGGER.info("Inference test successful!") - - -@pytest.fixture(scope="function") -def flux_endpoint(aws_session, image_uri): - """Deploy a FLUX.2 diffusion endpoint via the sglang.multimodal_gen engine.""" - model_id = "black-forest-labs/FLUX.2-klein-4B" - instance_type = "ml.g6e.xlarge" - - cleaned_id = clean_string(model_id.split("/")[1], "_./") - endpoint_name = random_suffix_name(f"sglang-{cleaned_id}", 50) - model_name = endpoint_name - - hf_token = get_hf_token(aws_session) - role_arn = aws_session.resolve_role_arn(SAGEMAKER_ROLE) - - model = endpoint_config = endpoint = None - try: - LOGGER.info(f"Creating FLUX.2 model: {model_name}") - model = Model.create( - model_name=model_name, - primary_container=ContainerDefinition( - image=image_uri, - environment={ - "SM_SGLANG_MODEL_PATH": model_id, - "SM_SGLANG_ENGINE": "diffusion", - "HF_TOKEN": hf_token, - }, - ), - execution_role_arn=role_arn, - ) - - LOGGER.info(f"Creating endpoint config: {endpoint_name}") - endpoint_config = EndpointConfig.create( - endpoint_config_name=endpoint_name, - production_variants=[ - ProductionVariant( - variant_name="AllTraffic", - model_name=model_name, - initial_instance_count=1, - instance_type=instance_type, - inference_ami_version=INFERENCE_AMI_VERSION, - container_startup_health_check_timeout_in_seconds=900, - ), - ], - ) - - LOGGER.info(f"Deploying endpoint: {endpoint_name} (this may take 10-15 minutes)...") - endpoint = Endpoint.create( - endpoint_name=endpoint_name, - endpoint_config_name=endpoint_name, - ) - endpoint.wait_for_status("InService") - LOGGER.info("FLUX.2 endpoint deployment completed successfully") - - yield endpoint - finally: - _cleanup([endpoint, endpoint_config, model]) - - -def test_sglang_sagemaker_flux_diffusion_endpoint(flux_endpoint): - """FLUX.2 text-to-image generation through the SageMaker diffusion endpoint.""" - endpoint = flux_endpoint - - payload = json.dumps( - { - "prompt": "a red cube on a white table", - "num_inference_steps": 4, - "width": 512, - "height": 512, - } - ) - LOGGER.debug(f"Sending image-generation request with payload: {payload}") - - result = endpoint.invoke(body=payload, content_type="application/json") - body = json.loads(result.body.read()) - LOGGER.info("Image-generation request invoked successfully") - - assert body, "Model response is empty, failing FLUX.2 endpoint test!" - # OpenAI-style image response returns a non-empty `data` list of images. - assert body.get("data"), f"No image data in FLUX.2 response: {pformat(body)}" - - LOGGER.info("FLUX.2 diffusion inference test successful!") diff --git a/test/sglang/scripts/sglang_flux_vae_decode_test.sh b/test/sglang/scripts/sglang_flux_vae_decode_test.sh index e3ef66b818e2..79c503da8a04 100755 --- a/test/sglang/scripts/sglang_flux_vae_decode_test.sh +++ b/test/sglang/scripts/sglang_flux_vae_decode_test.sh @@ -6,8 +6,7 @@ set -euo pipefail # FLUX.2-small-decoder is a distilled VAE decoder (class AutoencoderKLFlux2), a # drop-in decoder for the FLUX.2 diffusion pipeline. It is a pipeline COMPONENT, # not a servable text-to-image model, so this validates it the way it is used: -# load the decoder and decode a latent into an image tensor. End-to-end serving -# of a full FLUX.2 pipeline is covered by the SageMaker endpoint test. +# load the decoder and decode a latent into an image tensor. # # Usage: sglang_flux_vae_decode_test.sh [extra_args...] From a020ac7c5fe7f842648a704f619c87c2a2e2554a Mon Sep 17 00:00:00 2001 From: Jyothirmai Kottu Date: Thu, 9 Jul 2026 17:48:04 +0000 Subject: [PATCH 3/7] test(sglang): serve FLUX.2 diffusion on SageMaker instead of VAE decode Replace the diffusers-only VAE decode test with a real SGLang serving test. The decode test loaded AutoencoderKLFlux2 via diffusers and never exercised SGLang, so it did not validate SGLang's diffusion serving path. - sagemaker_entrypoint.sh: add an SM_SGLANG_ENGINE gate (default "llm", unchanged) so SM_SGLANG_ENGINE=diffusion launches sglang.multimodal_gen.runtime.launch_server instead of the LLM engine. - sagemaker test: deploy FLUX.2-klein-4B (the ungated, servable FLUX.2 pipeline) via the diffusion engine and assert /v1/images/generations returns image data. - drop the sglang_flux_vae_decode_test.sh script and its model-tests entry. --- .../config/model-tests/sglang-model-tests.yml | 9 -- scripts/docker/sglang/sagemaker_entrypoint.sh | 20 ++++- test/sglang/sagemaker/test_sm_endpoint.py | 82 +++++++++++++++++++ .../scripts/sglang_flux_vae_decode_test.sh | 43 ---------- 4 files changed, 100 insertions(+), 54 deletions(-) delete mode 100755 test/sglang/scripts/sglang_flux_vae_decode_test.sh diff --git a/.github/config/model-tests/sglang-model-tests.yml b/.github/config/model-tests/sglang-model-tests.yml index ba64be0b9fe5..a3b3cc034129 100644 --- a/.github/config/model-tests/sglang-model-tests.yml +++ b/.github/config/model-tests/sglang-model-tests.yml @@ -38,15 +38,6 @@ smoke-test: fleet: "x86-g6xl-runner" extra_args: "--tp 1 --context-length 4096 --dtype bfloat16" - # FLUX.2-small-decoder is a distilled VAE decoder (AutoencoderKLFlux2), a - # FLUX.2 pipeline component rather than a servable model — validated via a - # decode test (see test_script) instead of the default LLM serving smoke - # test. - - name: "flux2-small-decoder" - s3_model: "flux2-small-decoder.tar.gz" - fleet: "x86-g6xl-runner" - test_script: "sglang_flux_vae_decode_test.sh" - extra_args: "" - name: "locate-anything-3b" s3_prefix: "s3://dlc-cicd-models/vlm-models" s3_model: "locate-anything-3b.tar.gz" diff --git a/scripts/docker/sglang/sagemaker_entrypoint.sh b/scripts/docker/sglang/sagemaker_entrypoint.sh index 23a2fff81e59..359affa05ab1 100755 --- a/scripts/docker/sglang/sagemaker_entrypoint.sh +++ b/scripts/docker/sglang/sagemaker_entrypoint.sh @@ -12,9 +12,19 @@ echo "Starting server" PREFIX="SM_SGLANG_" ARG_PREFIX="--" +# Engine selector (default: llm). Set SM_SGLANG_ENGINE=diffusion to serve a +# FLUX.2 / diffusion pipeline via sglang.multimodal_gen instead of the LLM +# engine. This var controls the launch module and is NOT forwarded as a flag. +ENGINE=$(echo "${SM_SGLANG_ENGINE:-llm}" | tr '[:upper:]' '[:lower:]') + ARGS=() while IFS='=' read -r key value; do + # SM_SGLANG_ENGINE selects the launch module; it is not a server flag. + if [ "$key" = "${PREFIX}ENGINE" ]; then + continue + fi + arg_name=$(echo "${key#"${PREFIX}"}" | tr '[:upper:]' '[:lower:]' | tr '_' '-') # Handle boolean flags: true -> flag only, false -> skip entirely @@ -46,5 +56,11 @@ if ! [[ " ${ARGS[@]} " =~ " --model-path " ]]; then ARGS+=(--model-path "${SM_SGLANG_MODEL_PATH:-/opt/ml/model}") fi -echo "Running command: exec python3 -m sglang.launch_server ${ARGS[@]}" -exec python3 -m sglang.launch_server "${ARGS[@]}" +if [ "$ENGINE" = "diffusion" ]; then + LAUNCH_MODULE="sglang.multimodal_gen.runtime.launch_server" +else + LAUNCH_MODULE="sglang.launch_server" +fi + +echo "Running command: exec python3 -m ${LAUNCH_MODULE} ${ARGS[@]}" +exec python3 -m "${LAUNCH_MODULE}" "${ARGS[@]}" diff --git a/test/sglang/sagemaker/test_sm_endpoint.py b/test/sglang/sagemaker/test_sm_endpoint.py index eac323d47d97..38f0f3abc34c 100644 --- a/test/sglang/sagemaker/test_sm_endpoint.py +++ b/test/sglang/sagemaker/test_sm_endpoint.py @@ -118,3 +118,85 @@ def test_sglang_sagemaker_endpoint(model_endpoint, model_id): LOGGER.info(f"Model response: {pformat(body)}") LOGGER.info("Inference test successful!") + + +@pytest.fixture(scope="function") +def flux_endpoint(aws_session, image_uri): + """Deploy a FLUX.2 diffusion endpoint via the sglang.multimodal_gen engine.""" + model_id = "black-forest-labs/FLUX.2-klein-4B" + instance_type = "ml.g6e.xlarge" + + cleaned_id = clean_string(model_id.split("/")[1], "_./") + endpoint_name = random_suffix_name(f"sglang-{cleaned_id}", 50) + model_name = endpoint_name + + hf_token = get_hf_token(aws_session) + role_arn = aws_session.resolve_role_arn(SAGEMAKER_ROLE) + + model = endpoint_config = endpoint = None + try: + LOGGER.info(f"Creating FLUX.2 model: {model_name}") + model = Model.create( + model_name=model_name, + primary_container=ContainerDefinition( + image=image_uri, + environment={ + "SM_SGLANG_MODEL_PATH": model_id, + "SM_SGLANG_ENGINE": "diffusion", + "HF_TOKEN": hf_token, + }, + ), + execution_role_arn=role_arn, + ) + + LOGGER.info(f"Creating endpoint config: {endpoint_name}") + endpoint_config = EndpointConfig.create( + endpoint_config_name=endpoint_name, + production_variants=[ + ProductionVariant( + variant_name="AllTraffic", + model_name=model_name, + initial_instance_count=1, + instance_type=instance_type, + inference_ami_version=INFERENCE_AMI_VERSION, + container_startup_health_check_timeout_in_seconds=900, + ), + ], + ) + + LOGGER.info(f"Deploying endpoint: {endpoint_name} (this may take 10-15 minutes)...") + endpoint = Endpoint.create( + endpoint_name=endpoint_name, + endpoint_config_name=endpoint_name, + ) + endpoint.wait_for_status("InService") + LOGGER.info("FLUX.2 endpoint deployment completed successfully") + + yield endpoint + finally: + _cleanup([endpoint, endpoint_config, model]) + + +def test_sglang_sagemaker_flux_diffusion_endpoint(flux_endpoint): + """FLUX.2 text-to-image generation through the SageMaker diffusion endpoint.""" + endpoint = flux_endpoint + + payload = json.dumps( + { + "prompt": "a red cube on a white table", + "num_inference_steps": 4, + "width": 512, + "height": 512, + } + ) + LOGGER.debug(f"Sending image-generation request with payload: {payload}") + + result = endpoint.invoke(body=payload, content_type="application/json") + body = json.loads(result.body.read()) + LOGGER.info("Image-generation request invoked successfully") + + assert body, "Model response is empty, failing FLUX.2 endpoint test!" + # OpenAI-style image response returns a non-empty `data` list of images. + assert body.get("data"), f"No image data in FLUX.2 response: {pformat(body)}" + + LOGGER.info("FLUX.2 diffusion inference test successful!") diff --git a/test/sglang/scripts/sglang_flux_vae_decode_test.sh b/test/sglang/scripts/sglang_flux_vae_decode_test.sh deleted file mode 100755 index 79c503da8a04..000000000000 --- a/test/sglang/scripts/sglang_flux_vae_decode_test.sh +++ /dev/null @@ -1,43 +0,0 @@ -#!/bin/bash -set -euo pipefail - -# SGLang FLUX.2 VAE decode test -# -# FLUX.2-small-decoder is a distilled VAE decoder (class AutoencoderKLFlux2), a -# drop-in decoder for the FLUX.2 diffusion pipeline. It is a pipeline COMPONENT, -# not a servable text-to-image model, so this validates it the way it is used: -# load the decoder and decode a latent into an image tensor. -# -# Usage: sglang_flux_vae_decode_test.sh [extra_args...] - -MODEL_DIR="${1:?Usage: $0 [extra_args...]}" -MODEL_NAME="${2:?Usage: $0 [extra_args...]}" - -echo "=== Model directory: ${MODEL_DIR} ===" -ls -la "${MODEL_DIR}" - -echo "=== Loading AutoencoderKLFlux2 and running a decode ===" -python3 - "${MODEL_DIR}" "${MODEL_NAME}" <<'PY' -import sys - -import torch -from diffusers import AutoencoderKLFlux2 - -model_dir, model_name = sys.argv[1], sys.argv[2] -device = "cuda" if torch.cuda.is_available() else "cpu" - -vae = AutoencoderKLFlux2.from_pretrained(model_dir).to(device).eval() -n_params = sum(p.numel() for p in vae.parameters()) -print(f"Loaded {type(vae).__name__} params={n_params} device={device}") - -# Build a latent that matches the decoder's expected channel count, then decode. -latent_channels = vae.config.latent_channels -latent = torch.randn(1, latent_channels, 16, 16, device=device, dtype=next(vae.parameters()).dtype) -with torch.no_grad(): - image = vae.decode(latent).sample - -assert image.ndim == 4 and image.shape[0] == 1, f"unexpected decode output shape {tuple(image.shape)}" -assert torch.isfinite(image).all(), "decode produced non-finite values" -print(f"Decoded latent {tuple(latent.shape)} -> image {tuple(image.shape)}") -print(f"=== PASSED: {model_name} ===") -PY From 3d1304e9a5a9edd29627dfbd4a8b2f1b21da8b15 Mon Sep 17 00:00:00 2001 From: Jyothirmai Kottu Date: Thu, 9 Jul 2026 18:28:41 +0000 Subject: [PATCH 4/7] test(sglang): allowlist CVE-2026-39822 (mooncake go/stdlib) Pre-existing go/stdlib CVE compiled into mooncake's libetcd_wrapper.so, same class as the other mooncake go/stdlib entries already allowlisted. Not patchable without an upstream mooncake rebuild on Go 1.26.5+. Added to both the sglang (ec2) and sglang_server (sagemaker) allowlists. --- .../data/ecr_scan_allowlist/sglang/framework_allowlist.json | 5 +++++ .../sglang_server/framework_allowlist.json | 5 +++++ 2 files changed, 10 insertions(+) diff --git a/test/security/data/ecr_scan_allowlist/sglang/framework_allowlist.json b/test/security/data/ecr_scan_allowlist/sglang/framework_allowlist.json index 297c727dbbb3..2e40063a4b23 100644 --- a/test/security/data/ecr_scan_allowlist/sglang/framework_allowlist.json +++ b/test/security/data/ecr_scan_allowlist/sglang/framework_allowlist.json @@ -34,6 +34,11 @@ "reason": "go/stdlib 1.25.9 embedded in mooncake libetcd_wrapper.so, MIME header parsing CPU exhaustion, cannot patch without upstream mooncake rebuild with Go 1.26.4+", "review_by": "2026-09-10" }, + { + "vulnerability_id": "CVE-2026-39822", + "reason": "go/stdlib 1.25.9 embedded in mooncake libetcd_wrapper.so, os.Root symlink traversal, cannot patch without upstream mooncake rebuild with Go 1.26.5+", + "review_by": "2026-09-10" + }, { "vulnerability_id": "RUSTSEC-2026-0195", "reason": "quick-xml 0.39.2 is bundled inside the uv binary at /usr/local/bin/uv (uv is statically compiled Rust). Advisory is a denial-of-service in NsReader's namespace resolution during Start/Empty XML event handling; the vulnerable code path is only reachable when uv parses attacker-controlled XML (e.g. a hostile PyPI mirror index). uv is retained in the image for customer-side venv management; the standard PyPI flow uses JSON/HTML endpoints and never exercises the XML path. No uv release containing the patched quick-xml>=0.41.0 has been published upstream yet; will swap to a pinned fixed version and drop this entry when upstream ships.", diff --git a/test/security/data/ecr_scan_allowlist/sglang_server/framework_allowlist.json b/test/security/data/ecr_scan_allowlist/sglang_server/framework_allowlist.json index c6ce30216d2e..0333b92a8ab0 100644 --- a/test/security/data/ecr_scan_allowlist/sglang_server/framework_allowlist.json +++ b/test/security/data/ecr_scan_allowlist/sglang_server/framework_allowlist.json @@ -64,6 +64,11 @@ "reason": "go/stdlib 1.24.12 embedded in mooncake libetcd_wrapper.so, MIME header parsing CPU exhaustion, cannot patch without upstream mooncake rebuild with Go 1.26.4+", "review_by": "2026-09-10" }, + { + "vulnerability_id": "CVE-2026-39822", + "reason": "go/stdlib embedded in mooncake libetcd_wrapper.so, os.Root symlink traversal, cannot patch without upstream mooncake rebuild with Go 1.26.5+", + "review_by": "2026-09-10" + }, { "vulnerability_id": "RUSTSEC-2026-0185", "reason": "quinn-proto in uv binary, upstream uv has not released a fix yet. QUIC reassembly DoS requires malicious server, low risk for pip installer.", From 72437ef1d9c4480ebd268611afd02152334ece8f Mon Sep 17 00:00:00 2001 From: Jyothirmai Kottu Date: Fri, 10 Jul 2026 19:00:23 +0000 Subject: [PATCH 5/7] fix(sglang): add SageMaker /ping + /invocations to FLUX.2 diffusion serving The multimodal_gen diffusion server does not expose SageMaker's /ping health probe or /invocations route (only sglang.launch_server does), so a FLUX.2 SageMaker endpoint never reaches InService and the endpoint test fails. Add a thin launch wrapper (sagemaker_diffusion_serve.py) that monkeypatches multimodal_gen's create_app() to register GET /ping (200) and POST /invocations (delegates to the images handler), then hands off to the real launch_server(). The diffusion entrypoint path now execs this wrapper; both Dockerfiles COPY it into the image. LLM path is unchanged. Verified on g6e.xlarge with FLUX.2-klein-4B: /ping -> 200, /invocations -> 200 returning a 512x512 JPEG. --- docker/sglang/Dockerfile | 1 + docker/sglang/Dockerfile.amzn2023 | 1 + .../sglang/sagemaker_diffusion_serve.py | 68 +++++++++++++++++++ scripts/docker/sglang/sagemaker_entrypoint.sh | 11 +-- 4 files changed, 76 insertions(+), 5 deletions(-) create mode 100644 scripts/docker/sglang/sagemaker_diffusion_serve.py diff --git a/docker/sglang/Dockerfile b/docker/sglang/Dockerfile index c2dc8661513f..fa2e01a78ef2 100644 --- a/docker/sglang/Dockerfile +++ b/docker/sglang/Dockerfile @@ -133,6 +133,7 @@ RUN dpkg -l | grep -E "cuda|nvidia|libnv" | awk '{print $2}' | xargs apt-mark ho RUN rm -rf /tmp/* COPY ./scripts/docker/sglang/sagemaker_entrypoint.sh /usr/bin/serve +COPY ./scripts/docker/sglang/sagemaker_diffusion_serve.py /usr/local/bin/sagemaker_diffusion_serve.py RUN chmod +x /usr/bin/serve ENTRYPOINT ["/usr/bin/serve"] diff --git a/docker/sglang/Dockerfile.amzn2023 b/docker/sglang/Dockerfile.amzn2023 index ef9c421d8abe..3f661ae85556 100644 --- a/docker/sglang/Dockerfile.amzn2023 +++ b/docker/sglang/Dockerfile.amzn2023 @@ -417,6 +417,7 @@ RUN dnf upgrade -y --security --releasever latest \ && ln -sf /usr/bin/python${PYTHON_VERSION} /usr/bin/python COPY ./scripts/docker/sglang/sagemaker_entrypoint.sh /usr/bin/serve +COPY ./scripts/docker/sglang/sagemaker_diffusion_serve.py /usr/local/bin/sagemaker_diffusion_serve.py RUN chmod +x /usr/bin/serve ENTRYPOINT ["/usr/bin/serve"] \ No newline at end of file diff --git a/scripts/docker/sglang/sagemaker_diffusion_serve.py b/scripts/docker/sglang/sagemaker_diffusion_serve.py new file mode 100644 index 000000000000..fbaed96361c7 --- /dev/null +++ b/scripts/docker/sglang/sagemaker_diffusion_serve.py @@ -0,0 +1,68 @@ +"""SageMaker launch wrapper for the SGLang diffusion (multimodal_gen) server. + +SGLang's LLM server (sglang.launch_server) natively exposes the SageMaker +serving contract — GET /ping and POST /invocations — but the separate +diffusion server (sglang.multimodal_gen) does not: its FastAPI app only serves +a Vertex AI route and the OpenAI-images routes. As a result a FLUX.2 (or other +diffusion) SageMaker endpoint fails the /ping health check. + +This wrapper adds the two SageMaker routes to the diffusion server without +forking it: it monkeypatches multimodal_gen's create_app() to register +GET /ping (200) and POST /invocations (delegates to the same handler as +POST /v1/images/generations), then hands off to the real launch_server() so +the GPU workers/scheduler are bootstrapped exactly as usual. + +TODO(remove-when-upstream): delete this wrapper and point the entrypoint back +at `python3 -m sglang.multimodal_gen.runtime.launch_server` once upstream +SGLang adds /ping + /invocations to multimodal_gen's create_app() (mirroring +srt/entrypoints/http_server.py) and the DLC image is bumped to that version. +""" + +import sys + +from fastapi import Request, Response + +from sglang.multimodal_gen.runtime import launch_server as _launch +from sglang.multimodal_gen.runtime.entrypoints.http_server import create_app as _create_app +from sglang.multimodal_gen.runtime.entrypoints.openai.image_api import generations +from sglang.multimodal_gen.runtime.entrypoints.openai.protocol import ( + ImageGenerationsRequest, + ImageResponse, +) + + +async def _sagemaker_ping() -> Response: + """SageMaker startup health probe.""" + return Response(status_code=200) + + +async def _sagemaker_invocations( + request: ImageGenerationsRequest, raw_request: Request +) -> ImageResponse: + """SageMaker inference route — same code path as POST /v1/images/generations.""" + return await generations(request, raw_request) + + +def _create_app_with_sagemaker_routes(server_args): + app = _create_app(server_args) + app.add_api_route("/ping", _sagemaker_ping, methods=["GET"]) + app.add_api_route( + "/invocations", + _sagemaker_invocations, + methods=["POST"], + response_model=ImageResponse, + ) + return app + + +def main(): + # Patch the name launch_server resolves at call time so the real launch + # flow (worker/scheduler bootstrap + uvicorn.run) builds our augmented app. + _launch.create_app = _create_app_with_sagemaker_routes + + server_args = _launch.prepare_server_args(sys.argv[1:]) + _launch.launch_server(server_args) + + +if __name__ == "__main__": + main() diff --git a/scripts/docker/sglang/sagemaker_entrypoint.sh b/scripts/docker/sglang/sagemaker_entrypoint.sh index 359affa05ab1..3c525013b4eb 100755 --- a/scripts/docker/sglang/sagemaker_entrypoint.sh +++ b/scripts/docker/sglang/sagemaker_entrypoint.sh @@ -57,10 +57,11 @@ if ! [[ " ${ARGS[@]} " =~ " --model-path " ]]; then fi if [ "$ENGINE" = "diffusion" ]; then - LAUNCH_MODULE="sglang.multimodal_gen.runtime.launch_server" + # multimodal_gen's server lacks the SageMaker /ping + /invocations routes, + # so launch it through the wrapper that adds them. + echo "Running command: exec python3 /usr/local/bin/sagemaker_diffusion_serve.py ${ARGS[@]}" + exec python3 /usr/local/bin/sagemaker_diffusion_serve.py "${ARGS[@]}" else - LAUNCH_MODULE="sglang.launch_server" + echo "Running command: exec python3 -m sglang.launch_server ${ARGS[@]}" + exec python3 -m sglang.launch_server "${ARGS[@]}" fi - -echo "Running command: exec python3 -m ${LAUNCH_MODULE} ${ARGS[@]}" -exec python3 -m "${LAUNCH_MODULE}" "${ARGS[@]}" From 7d3811100f818ae93612630a2bf9b8a62b15d187 Mon Sep 17 00:00:00 2001 From: Jyothirmai Kottu Date: Fri, 10 Jul 2026 21:25:38 +0000 Subject: [PATCH 6/7] fix(sglang): keep single column-0 exec in entrypoint for sanity dry-run Wrapping the terminal exec in an if/else indented both exec lines, so the sanity test's entrypoint dry-run regex (^exec python3 ...) no longer matched and TestEntrypointArgHandling execed the real launch_server, timing out (7 errors). Select the launch target into a variable and keep one unindented exec python3 line, matching the vllm entrypoint convention. --- scripts/docker/sglang/sagemaker_entrypoint.sh | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/scripts/docker/sglang/sagemaker_entrypoint.sh b/scripts/docker/sglang/sagemaker_entrypoint.sh index 3c525013b4eb..041e4d622f05 100755 --- a/scripts/docker/sglang/sagemaker_entrypoint.sh +++ b/scripts/docker/sglang/sagemaker_entrypoint.sh @@ -56,12 +56,12 @@ if ! [[ " ${ARGS[@]} " =~ " --model-path " ]]; then ARGS+=(--model-path "${SM_SGLANG_MODEL_PATH:-/opt/ml/model}") fi +# diffusion routes through the wrapper that adds SageMaker's /ping + /invocations routes. if [ "$ENGINE" = "diffusion" ]; then - # multimodal_gen's server lacks the SageMaker /ping + /invocations routes, - # so launch it through the wrapper that adds them. - echo "Running command: exec python3 /usr/local/bin/sagemaker_diffusion_serve.py ${ARGS[@]}" - exec python3 /usr/local/bin/sagemaker_diffusion_serve.py "${ARGS[@]}" + LAUNCH_TARGET=(/usr/local/bin/sagemaker_diffusion_serve.py) else - echo "Running command: exec python3 -m sglang.launch_server ${ARGS[@]}" - exec python3 -m sglang.launch_server "${ARGS[@]}" + LAUNCH_TARGET=(-m sglang.launch_server) fi + +echo "Running command: exec python3 ${LAUNCH_TARGET[@]} ${ARGS[@]}" +exec python3 "${LAUNCH_TARGET[@]}" "${ARGS[@]}" From 017cc65de11ff4c1da3d995e337e5b06cca3b6be Mon Sep 17 00:00:00 2001 From: Jyothirmai Kottu Date: Fri, 10 Jul 2026 21:35:08 +0000 Subject: [PATCH 7/7] style(sglang): fix import grouping in diffusion wrapper (ruff isort) --- scripts/docker/sglang/sagemaker_diffusion_serve.py | 1 - 1 file changed, 1 deletion(-) diff --git a/scripts/docker/sglang/sagemaker_diffusion_serve.py b/scripts/docker/sglang/sagemaker_diffusion_serve.py index fbaed96361c7..8e0df514d0fc 100644 --- a/scripts/docker/sglang/sagemaker_diffusion_serve.py +++ b/scripts/docker/sglang/sagemaker_diffusion_serve.py @@ -21,7 +21,6 @@ import sys from fastapi import Request, Response - from sglang.multimodal_gen.runtime import launch_server as _launch from sglang.multimodal_gen.runtime.entrypoints.http_server import create_app as _create_app from sglang.multimodal_gen.runtime.entrypoints.openai.image_api import generations