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