test(sglang): add FLUX.2 diffusion serving test on SageMaker#6367
Open
Jyothirmaikottu wants to merge 9 commits into
Open
test(sglang): add FLUX.2 diffusion serving test on SageMaker#6367Jyothirmaikottu wants to merge 9 commits into
Jyothirmaikottu wants to merge 9 commits into
Conversation
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.
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.
Jyothirmaikottu
changed the title
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.
Jyothirmaikottu
changed the title
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.
…erving 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.
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.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Adds a real SGLang serving test for FLUX.2 diffusion on SageMaker. The SGLang DLC image ships the
sglang.multimodal_gendiffusion stack but the SageMaker entrypoint hardcoded the LLM engine, so diffusion pipelines could not be served.Changes
scripts/docker/sglang/sagemaker_entrypoint.sh: newSM_SGLANG_ENGINEgate — defaultllm(existing behavior unchanged);diffusionlaunchessglang.multimodal_gen.runtime.launch_server. The control var is excluded from theSM_SGLANG_*→--flagmapping.test/sglang/sagemaker/test_sm_endpoint.py: deployFLUX.2-klein-4B(the ungated, servable FLUX.2 pipeline) with the diffusion engine and assert/v1/images/generationsreturns image data.Testing: FLUX.2-klein-4B verified serving end-to-end on an L40S via
sglang.multimodal_gen— pipeline recognized (Qwen3 text encoder →Flux2Transformer2DModel→AutoencoderKLFlux2),/healthok,/v1/images/generationsreturned a generated image. Backward-compatible — existing LLM SageMaker images are unaffected (engine defaults tollm).