test(sglang): add FLUX.2 diffusion serving test on SageMaker

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"]

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"]

scripts/docker/sglang/sagemaker_diffusion_serve.py

@@ -0,0 +1,67 @@
+"""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()

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,12 @@ 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[@]}"
+# diffusion routes through the wrapper that adds SageMaker's /ping + /invocations routes.
+if [ "$ENGINE" = "diffusion" ]; then
+    LAUNCH_TARGET=(/usr/local/bin/sagemaker_diffusion_serve.py)
+else
+    LAUNCH_TARGET=(-m sglang.launch_server)
+fi
+
+echo "Running command: exec python3 ${LAUNCH_TARGET[@]} ${ARGS[@]}"
+exec python3 "${LAUNCH_TARGET[@]}" "${ARGS[@]}"

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.",

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!")