diff --git a/docker/sglang/Dockerfile b/docker/sglang/Dockerfile index af86cf8921a3..cb6bd336f4cd 100644 --- a/docker/sglang/Dockerfile +++ b/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"] diff --git a/docker/sglang/Dockerfile.amzn2023 b/docker/sglang/Dockerfile.amzn2023 index ef9c421d8abe..3f661ae85556 100644 --- a/docker/sglang/Dockerfile.amzn2023 +++ b/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"] \ No newline at end of file diff --git a/scripts/docker/sglang/sagemaker_diffusion_serve.py b/scripts/docker/sglang/sagemaker_diffusion_serve.py new file mode 100644 index 000000000000..8e0df514d0fc --- /dev/null +++ b/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() diff --git a/scripts/docker/sglang/sagemaker_entrypoint.sh b/scripts/docker/sglang/sagemaker_entrypoint.sh index 23a2fff81e59..041e4d622f05 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,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[@]}" 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!")