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

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[@]}"

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

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