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.

Author: Jyothirmaikottu

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

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