diff --git a/.github/trigger_files/beam_PostCommit_Python.json b/.github/trigger_files/beam_PostCommit_Python.json index 2bb052d5f715..00c0cdc3f9c8 100644 --- a/.github/trigger_files/beam_PostCommit_Python.json +++ b/.github/trigger_files/beam_PostCommit_Python.json @@ -1,5 +1,5 @@ { "comment": "Modify this file in a trivial way to cause this test suite to run.", - "pr": "37345", - "modification": 53 + "pr": "38701", + "modification": 55 } diff --git a/runners/google-cloud-dataflow-java/build.gradle b/runners/google-cloud-dataflow-java/build.gradle index 3d1f8d78777c..d5b0d9e82363 100644 --- a/runners/google-cloud-dataflow-java/build.gradle +++ b/runners/google-cloud-dataflow-java/build.gradle @@ -113,6 +113,9 @@ dependencies { implementation library.java.google_cloud_logging permitUnusedDeclared library.java.google_cloud_logging // BEAM-11761 implementation library.java.opentelemetry_context + runtimeOnly library.java.opentelemetry_exporter_otlp + runtimeOnly library.java.opentelemetry_extension_autoconfigure + runtimeOnly project(":sdks:java:extensions:opentelemetry-gcp-auth-extension") implementation library.java.hamcrest implementation library.java.jackson_annotations implementation library.java.jackson_core diff --git a/sdks/go.mod b/sdks/go.mod index 7924367a3ef2..b7094dc368cf 100644 --- a/sdks/go.mod +++ b/sdks/go.mod @@ -33,10 +33,10 @@ require ( cloud.google.com/go/spanner v1.92.0 cloud.google.com/go/storage v1.63.0 github.com/aws/aws-sdk-go-v2 v1.42.0 - github.com/aws/aws-sdk-go-v2/config v1.32.25 - github.com/aws/aws-sdk-go-v2/credentials v1.19.24 - github.com/aws/aws-sdk-go-v2/feature/s3/manager v1.22.28 - github.com/aws/aws-sdk-go-v2/service/s3 v1.104.0 + github.com/aws/aws-sdk-go-v2/config v1.32.26 + github.com/aws/aws-sdk-go-v2/credentials v1.19.25 + github.com/aws/aws-sdk-go-v2/feature/s3/manager v1.22.29 + github.com/aws/aws-sdk-go-v2/service/s3 v1.104.1 github.com/aws/smithy-go v1.27.3 github.com/docker/go-connections v0.7.0 // indirect github.com/dustin/go-humanize v1.0.1 @@ -46,7 +46,7 @@ require ( github.com/johannesboyne/gofakes3 v0.0.0-20250106100439-5c39aecd6999 github.com/lib/pq v1.12.3 github.com/linkedin/goavro/v2 v2.15.0 - github.com/nats-io/nats-server/v2 v2.14.2 + github.com/nats-io/nats-server/v2 v2.14.3 github.com/nats-io/nats.go v1.52.0 github.com/proullon/ramsql v0.1.4 github.com/spf13/cobra v1.10.2 @@ -91,7 +91,7 @@ require ( github.com/GoogleCloudPlatform/opentelemetry-operations-go/internal/resourcemapping v0.57.0 // indirect github.com/antithesishq/antithesis-sdk-go v0.7.0-default-no-op // indirect github.com/apache/arrow/go/v15 v15.0.2 // indirect - github.com/aws/aws-sdk-go-v2/service/signin v1.2.0 // indirect + github.com/aws/aws-sdk-go-v2/service/signin v1.2.1 // indirect github.com/containerd/errdefs v1.0.0 // indirect github.com/containerd/errdefs/pkg v0.3.0 // indirect github.com/containerd/log v0.1.0 // indirect @@ -156,10 +156,10 @@ require ( github.com/aws/aws-sdk-go-v2/service/internal/accept-encoding v1.13.12 // indirect github.com/aws/aws-sdk-go-v2/service/internal/checksum v1.9.22 // indirect github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.13.29 // indirect - github.com/aws/aws-sdk-go-v2/service/internal/s3shared v1.19.29 // indirect - github.com/aws/aws-sdk-go-v2/service/sso v1.31.3 // indirect - github.com/aws/aws-sdk-go-v2/service/ssooidc v1.36.6 // indirect - github.com/aws/aws-sdk-go-v2/service/sts v1.43.3 // indirect + github.com/aws/aws-sdk-go-v2/service/internal/s3shared v1.19.30 // indirect + github.com/aws/aws-sdk-go-v2/service/sso v1.31.4 // indirect + github.com/aws/aws-sdk-go-v2/service/ssooidc v1.36.7 // indirect + github.com/aws/aws-sdk-go-v2/service/sts v1.43.4 // indirect github.com/cenkalti/backoff/v4 v4.3.0 // indirect github.com/cespare/xxhash/v2 v2.3.0 // indirect github.com/cncf/xds/go v0.0.0-20260202195803-dba9d589def2 // indirect diff --git a/sdks/go.sum b/sdks/go.sum index 75ca43710f28..03df389b0e56 100644 --- a/sdks/go.sum +++ b/sdks/go.sum @@ -207,20 +207,20 @@ github.com/aws/aws-sdk-go-v2/aws/protocol/eventstream v1.7.13 h1:p1BBrg/Hhp6uK7z github.com/aws/aws-sdk-go-v2/aws/protocol/eventstream v1.7.13/go.mod h1:8cIfkE9MDhkRZGpQ22aV6/lkYeYSozpz16Smrs5x4Ls= github.com/aws/aws-sdk-go-v2/config v1.15.3/go.mod h1:9YL3v07Xc/ohTsxFXzan9ZpFpdTOFl4X65BAKYaz8jg= github.com/aws/aws-sdk-go-v2/config v1.25.3/go.mod h1:tAByZy03nH5jcq0vZmkcVoo6tRzRHEwSFx3QW4NmDw8= -github.com/aws/aws-sdk-go-v2/config v1.32.25 h1:ACCejvStYoilgwrfegSt5ZntCbPrk52qfwyNcnl3omM= -github.com/aws/aws-sdk-go-v2/config v1.32.25/go.mod h1:LJyU8sDRbXUxFn8xMJIGP+v9QYYwveNLI8a/giAOiAs= +github.com/aws/aws-sdk-go-v2/config v1.32.26 h1:JI+W5B3jUA8UBz2ggbICGd9UCR6/+SB21G8EFl0SFTQ= +github.com/aws/aws-sdk-go-v2/config v1.32.26/go.mod h1:RLE2Ls/wRstvdSz1GPrIWNnXcKZ/znDdWyMuiQxdBoY= github.com/aws/aws-sdk-go-v2/credentials v1.11.2/go.mod h1:j8YsY9TXTm31k4eFhspiQicfXPLZ0gYXA50i4gxPE8g= github.com/aws/aws-sdk-go-v2/credentials v1.16.2/go.mod h1:sDdvGhXrSVT5yzBDR7qXz+rhbpiMpUYfF3vJ01QSdrc= -github.com/aws/aws-sdk-go-v2/credentials v1.19.24 h1:2hQqYCV9yqyePQ9o6dCrZc/zO8U3TwPr9mIKlZnPu/I= -github.com/aws/aws-sdk-go-v2/credentials v1.19.24/go.mod h1:IDwpACtwqHLISdzfwUUNq4P9DsB/h5BLg4FwJPNfqFY= +github.com/aws/aws-sdk-go-v2/credentials v1.19.25 h1:TzPVjfUZ1hsKafvYE+DIzKXIik2KufQxsPHanlkttbo= +github.com/aws/aws-sdk-go-v2/credentials v1.19.25/go.mod h1:K4hw0buguVvtC74HnVfTRr0LzQQHAWPqJbBU9QGk2Pg= github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.12.3/go.mod h1:uk1vhHHERfSVCUnqSqz8O48LBYDSC+k6brng09jcMOk= github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.14.4/go.mod h1:t4i+yGHMCcUNIX1x7YVYa6bH/Do7civ5I6cG/6PMfyA= github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.18.29 h1:r6qZHbT+wxgWO/e9vYNUEtg7lv5+UN3pRqKhLXvnArg= github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.18.29/go.mod h1:QRnaRcTVGKPGRy8w78HMQtKUGRYcnMZAANATkeVA6Mo= github.com/aws/aws-sdk-go-v2/feature/s3/manager v1.11.3/go.mod h1:0dHuD2HZZSiwfJSy1FO5bX1hQ1TxVV1QXXjpn3XUE44= github.com/aws/aws-sdk-go-v2/feature/s3/manager v1.14.0/go.mod h1:UcgIwJ9KHquYxs6Q5skC9qXjhYMK+JASDYcXQ4X7JZE= -github.com/aws/aws-sdk-go-v2/feature/s3/manager v1.22.28 h1:ez4y5o7sa0uaRI8BquYOXtZpioUPhbQEh7Igm88oV9U= -github.com/aws/aws-sdk-go-v2/feature/s3/manager v1.22.28/go.mod h1:TpmZOrQA12XKEpVypgBGZSQBsm1WUTndCiSnbDsbvug= +github.com/aws/aws-sdk-go-v2/feature/s3/manager v1.22.29 h1:YteQL/8ZD9nS/eiLq3Ab9ldHBUvfWorEOufALZOXoXY= +github.com/aws/aws-sdk-go-v2/feature/s3/manager v1.22.29/go.mod h1:3hj0jtS3hQmQAAZ0yz/jTA+uTAHfDpwxbISkgZ2E0d0= github.com/aws/aws-sdk-go-v2/internal/configsources v1.1.9/go.mod h1:AnVH5pvai0pAF4lXRq0bmhbes1u9R8wTE+g+183bZNM= github.com/aws/aws-sdk-go-v2/internal/configsources v1.2.3/go.mod h1:7sGSz1JCKHWWBHq98m6sMtWQikmYPpxjqOydDemiVoM= github.com/aws/aws-sdk-go-v2/internal/configsources v1.4.29 h1:f3vKqSo13fhTYb+JEcXwXefZQE26I1FB5eTSniU67ko= @@ -248,30 +248,30 @@ github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.13.29 h1:DRebniUG github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.13.29/go.mod h1:LfRkPCD8YHDM2E5eTkos2UpwYeZnBcVarTa8L59bJHA= github.com/aws/aws-sdk-go-v2/service/internal/s3shared v1.13.3/go.mod h1:Bm/v2IaN6rZ+Op7zX+bOUMdL4fsrYZiD0dsjLhNKwZc= github.com/aws/aws-sdk-go-v2/service/internal/s3shared v1.16.3/go.mod h1:KZgs2ny8HsxRIRbDwgvJcHHBZPOzQr/+NtGwnP+w2ec= -github.com/aws/aws-sdk-go-v2/service/internal/s3shared v1.19.29 h1:hiME6pBzC7OTl9LMtlyTWBuEl1f4QBcUmFDKC7MLXtc= -github.com/aws/aws-sdk-go-v2/service/internal/s3shared v1.19.29/go.mod h1:G7RP+uhagpKtKhd1BM9N6JQqjCcGEU47K5lBVZQyRQw= +github.com/aws/aws-sdk-go-v2/service/internal/s3shared v1.19.30 h1:4HbXxyipSYxexU0juMIpdS05dilL6dbB2VQHxxN2vGU= +github.com/aws/aws-sdk-go-v2/service/internal/s3shared v1.19.30/go.mod h1:G7RP+uhagpKtKhd1BM9N6JQqjCcGEU47K5lBVZQyRQw= github.com/aws/aws-sdk-go-v2/service/kms v1.16.3/go.mod h1:QuiHPBqlOFCi4LqdSskYYAWpQlx3PKmohy+rE2F+o5g= github.com/aws/aws-sdk-go-v2/service/s3 v1.26.3/go.mod h1:g1qvDuRsJY+XghsV6zg00Z4KJ7DtFFCx8fJD2a491Ak= github.com/aws/aws-sdk-go-v2/service/s3 v1.43.0/go.mod h1:NXRKkiRF+erX2hnybnVU660cYT5/KChRD4iUgJ97cI8= -github.com/aws/aws-sdk-go-v2/service/s3 v1.104.0 h1:ta8csKy5vN91F3i5gGR85lFV0srBqySEji7Jroes6rE= -github.com/aws/aws-sdk-go-v2/service/s3 v1.104.0/go.mod h1:77ZAgynvx1txMvDG8gGWoWkO1augYDxkp9JElWFgjQU= +github.com/aws/aws-sdk-go-v2/service/s3 v1.104.1 h1:yb03KevaOAG5e8suo79Af74vjIQvoeKmjl79WQchLrs= +github.com/aws/aws-sdk-go-v2/service/s3 v1.104.1/go.mod h1:mreYODw0Y4yv7xeczvqC6vciwFao8lPE9k1l1ulfY6E= github.com/aws/aws-sdk-go-v2/service/secretsmanager v1.15.4/go.mod h1:PJc8s+lxyU8rrre0/4a0pn2wgwiDvOEzoOjcJUBr67o= -github.com/aws/aws-sdk-go-v2/service/signin v1.2.0 h1:3nXpRcFwRCW8n7HgO2QGy0Dc20eQNfBuUemGQhpF8m8= -github.com/aws/aws-sdk-go-v2/service/signin v1.2.0/go.mod h1:LxYujSTLPRlp2vTtcUO/+1ilrew8ytt6SvQyOgejzFQ= +github.com/aws/aws-sdk-go-v2/service/signin v1.2.1 h1:BeJmkm5YOZs6lGRGcNoIuLSoTTtGLLCEqlSiRKYodfM= +github.com/aws/aws-sdk-go-v2/service/signin v1.2.1/go.mod h1:LxYujSTLPRlp2vTtcUO/+1ilrew8ytt6SvQyOgejzFQ= github.com/aws/aws-sdk-go-v2/service/sns v1.17.4/go.mod h1:kElt+uCcXxcqFyc+bQqZPFD9DME/eC6oHBXvFzQ9Bcw= github.com/aws/aws-sdk-go-v2/service/sqs v1.18.3/go.mod h1:skmQo0UPvsjsuYYSYMVmrPc1HWCbHUJyrCEp+ZaLzqM= github.com/aws/aws-sdk-go-v2/service/ssm v1.24.1/go.mod h1:NR/xoKjdbRJ+qx0pMR4mI+N/H1I1ynHwXnO6FowXJc0= github.com/aws/aws-sdk-go-v2/service/sso v1.11.3/go.mod h1:7UQ/e69kU7LDPtY40OyoHYgRmgfGM4mgsLYtcObdveU= github.com/aws/aws-sdk-go-v2/service/sso v1.17.2/go.mod h1:/pE21vno3q1h4bbhUOEi+6Zu/aT26UK2WKkDXd+TssQ= -github.com/aws/aws-sdk-go-v2/service/sso v1.31.3 h1:ey1XLTYXb9PcLt4535632o5kCGXNXEhNb620Dqwuylo= -github.com/aws/aws-sdk-go-v2/service/sso v1.31.3/go.mod h1:Lk7PlmoTYryQmyBG0EXqj5BcUbj3whXdU2s3yGI3EAc= +github.com/aws/aws-sdk-go-v2/service/sso v1.31.4 h1:i465b/3c7xJd++pobNIDOggouekCuiWOnB0goQJy+94= +github.com/aws/aws-sdk-go-v2/service/sso v1.31.4/go.mod h1:Lk7PlmoTYryQmyBG0EXqj5BcUbj3whXdU2s3yGI3EAc= github.com/aws/aws-sdk-go-v2/service/ssooidc v1.20.0/go.mod h1:dWqm5G767qwKPuayKfzm4rjzFmVjiBFbOJrpSPnAMDs= -github.com/aws/aws-sdk-go-v2/service/ssooidc v1.36.6 h1:yLr03zQE/5Eu5l3QU0Si+xMbLMbSDF2YXsigqXngs6g= -github.com/aws/aws-sdk-go-v2/service/ssooidc v1.36.6/go.mod h1:Q5N6icH+KJZDLh+ESNwzdv6cZ6vLFF/egy3IOxWhmz4= +github.com/aws/aws-sdk-go-v2/service/ssooidc v1.36.7 h1:xbmJAnBbyYPkTzoCNCF/bpJ6ymQHRdXX1vquYfDIGYk= +github.com/aws/aws-sdk-go-v2/service/ssooidc v1.36.7/go.mod h1:Q5N6icH+KJZDLh+ESNwzdv6cZ6vLFF/egy3IOxWhmz4= github.com/aws/aws-sdk-go-v2/service/sts v1.16.3/go.mod h1:bfBj0iVmsUyUg4weDB4NxktD9rDGeKSVWnjTnwbx9b8= github.com/aws/aws-sdk-go-v2/service/sts v1.25.3/go.mod h1:4EqRHDCKP78hq3zOnmFXu5k0j4bXbRFfCh/zQ6KnEfQ= -github.com/aws/aws-sdk-go-v2/service/sts v1.43.3 h1:VrIhKRCSK1umelSgB9RghvA9RTUYeQffyAS5ApXehNI= -github.com/aws/aws-sdk-go-v2/service/sts v1.43.3/go.mod h1:r8wkDOuLaaMFqFiYAb8dGY2A3gJCOujMc6CFOVC4Zhc= +github.com/aws/aws-sdk-go-v2/service/sts v1.43.4 h1:Np0vmL7op0Zs5xGacYMMX3v5O5pvZ46xhb5LwDgPj8M= +github.com/aws/aws-sdk-go-v2/service/sts v1.43.4/go.mod h1:r8wkDOuLaaMFqFiYAb8dGY2A3gJCOujMc6CFOVC4Zhc= github.com/aws/smithy-go v1.11.2/go.mod h1:3xHYmszWVx2c0kIwQeEVf9uSm4fYZt67FBJnwub1bgM= github.com/aws/smithy-go v1.17.0/go.mod h1:NukqUGpCZIILqqiV0NIjeFh24kd/FAa4beRb6nbIUPE= github.com/aws/smithy-go v1.27.3 h1:F3Zb497UhhskkfpJmfkXswyo+t0sh9OTBnIHjogWbVY= @@ -712,8 +712,8 @@ github.com/montanaflynn/stats v0.9.0 h1:tsBJ0RXwph9BmAuFoCmqGv6e8xa0MENQ8m0ptKq2 github.com/montanaflynn/stats v0.9.0/go.mod h1:etXPPgVO6n31NxCd9KQUMvCM+ve0ruNzt6R8Bnaayow= github.com/nats-io/jwt/v2 v2.8.2 h1:XXRgB60MSTnqsRwejQurVDs/hcv2dkt+86GjI+I/bMc= github.com/nats-io/jwt/v2 v2.8.2/go.mod h1:Ag/56sq9OblL4JgdYufDd16Egb17Kr/8WwwuO/forVc= -github.com/nats-io/nats-server/v2 v2.14.2 h1:Q7dRhCY03Y00rETFW3KV+KGaCIajlDfWgWUVgbMxyuk= -github.com/nats-io/nats-server/v2 v2.14.2/go.mod h1:lWpb1bSpRELZfRdlMkdz8E7lbXKKyNe8RIn0vvepIHs= +github.com/nats-io/nats-server/v2 v2.14.3 h1:+xjydPt7rkit67G+04TN0mcO2n+8nveZE7tK/PPV53A= +github.com/nats-io/nats-server/v2 v2.14.3/go.mod h1:5IlCtBzfwyzQzPMjmoJ9W2/LKmnJRtNyuOs/OT+NHDY= github.com/nats-io/nats.go v1.52.0 h1:n3avV4VBsCgsdwh71TppsTwtv+QdPs7ntSKM8qJLGsc= github.com/nats-io/nats.go v1.52.0/go.mod h1:26HypzazeOkyO3/mqd1zZd53STJN0EjCYF9Uy2ZOBno= github.com/nats-io/nkeys v0.4.16 h1:rd5oAuLOb8mnAycB0xleuEBNS1pVVnN0fv/FF34Eypg= diff --git a/sdks/python/apache_beam/examples/inference/vllm_text_completion.py b/sdks/python/apache_beam/examples/inference/vllm_text_completion.py index a7468f521ebb..21f0d41d3cb2 100644 --- a/sdks/python/apache_beam/examples/inference/vllm_text_completion.py +++ b/sdks/python/apache_beam/examples/inference/vllm_text_completion.py @@ -138,6 +138,20 @@ def parse_known_args(argv): 'Passed to the vLLM OpenAI server as --gpu-memory-utilization ' '(fraction of total GPU memory for KV cache). Lower this if the ' 'engine fails to start with CUDA out of memory.')) + parser.add_argument( + '--use_dynamo', + dest='use_dynamo', + action='store_true', + help=( + 'Use embedded NVIDIA Dynamo as the vLLM engine. Requires ' + 'ai-dynamo[vllm] and the etcd binary in the runtime environment. ' + 'See VLLMCompletionsModelHandler for limitations of embedded mode.')) + parser.add_argument( + '--max_tokens', + dest='max_tokens', + type=int, + default=16, + help='Maximum number of tokens to generate for each example.') return parser.parse_known_args(argv) @@ -178,14 +192,17 @@ def run( build_vllm_server_kwargs(known_args)) model_handler = VLLMCompletionsModelHandler( - model_name=known_args.model, vllm_server_kwargs=effective_vllm_kwargs) + model_name=known_args.model, + vllm_server_kwargs=effective_vllm_kwargs, + use_dynamo=known_args.use_dynamo) input_examples = COMPLETION_EXAMPLES if known_args.chat: model_handler = VLLMChatModelHandler( model_name=known_args.model, chat_template_path=known_args.chat_template, - vllm_server_kwargs=dict(effective_vllm_kwargs)) + vllm_server_kwargs=dict(effective_vllm_kwargs), + use_dynamo=known_args.use_dynamo) input_examples = CHAT_EXAMPLES pipeline = test_pipeline @@ -193,7 +210,8 @@ def run( pipeline = beam.Pipeline(options=pipeline_options) examples = pipeline | "Create examples" >> beam.Create(input_examples) - predictions = examples | "RunInference" >> RunInference(model_handler) + predictions = examples | "RunInference" >> RunInference( + model_handler, inference_args={'max_tokens': known_args.max_tokens}) process_output = predictions | "Process Predictions" >> beam.ParDo( PostProcessor()) _ = process_output | "WriteOutput" >> beam.io.WriteToText( diff --git a/sdks/python/apache_beam/io/gcp/pubsub_test.py b/sdks/python/apache_beam/io/gcp/pubsub_test.py index c35de62fca1f..050a69aff6cf 100644 --- a/sdks/python/apache_beam/io/gcp/pubsub_test.py +++ b/sdks/python/apache_beam/io/gcp/pubsub_test.py @@ -502,6 +502,70 @@ def finish_bundle(self): @unittest.skipIf(pubsub is None, 'GCP dependencies are not installed') @mock.patch('google.cloud.pubsub.SubscriberClient') class TestReadFromPubSub(unittest.TestCase): + def setUp(self): + _PubSubReadEvaluator._subscription_cache.clear() + _PubSubReadEvaluator._subscriber_client_cache.clear() + + def test_subscriber_client_is_reused_for_transform(self, mock_pubsub): + class Transform(object): + pass + + transform = Transform() + first_client = _PubSubReadEvaluator._get_subscriber_client(transform) + second_client = _PubSubReadEvaluator._get_subscriber_client(transform) + + self.assertIs(first_client, second_client) + mock_pubsub.assert_called_once_with() + first_client.close.assert_not_called() + + def test_subscription_creation_does_not_hold_global_cache_lock( + self, mock_pubsub): + class Transform(object): + pass + + def subscription_path(project, subscription): + return 'projects/%s/subscriptions/%s' % (project, subscription) + + transform = Transform() + client = mock_pubsub.return_value + client.subscription_path.side_effect = subscription_path + client.topic_path.return_value = 'projects/topic_project/topics/topic' + global_lock_available = [] + + def create_subscription(name, topic): + self.assertTrue(name.startswith('projects/sub_project/subscriptions/')) + self.assertEqual('projects/topic_project/topics/topic', topic) + acquired = _PubSubReadEvaluator._subscriber_client_cache_lock.acquire( + blocking=False) + global_lock_available.append(acquired) + if acquired: + _PubSubReadEvaluator._subscriber_client_cache_lock.release() + + client.create_subscription.side_effect = create_subscription + + sub_name = _PubSubReadEvaluator.get_subscription( + transform, 'topic_project', 'topic', 'sub_project', None) + + self.assertTrue(global_lock_available) + self.assertTrue(global_lock_available[0]) + self.assertTrue(sub_name.startswith('projects/sub_project/subscriptions/')) + + def test_subscriber_client_cleanup_is_idempotent(self, unused_mock_pubsub): + client = mock.Mock() + subscriber_client = transform_evaluator._PubSubSubscriberClient(client) + subscriber_client.set_temporary_subscription('subscription') + + subscriber_client.close() + subscriber_client.close() + + client.assert_has_calls([ + mock.call.delete_subscription(subscription='subscription'), + mock.call.close() + ]) + client.delete_subscription.assert_called_once_with( + subscription='subscription') + client.close.assert_called_once_with() + def test_read_messages_success(self, mock_pubsub): data = b'data' publish_time_secs = 1520861821 @@ -533,7 +597,8 @@ def test_read_messages_success(self, mock_pubsub): mock_pubsub.return_value.acknowledge.assert_has_calls( [mock.call(subscription=mock.ANY, ack_ids=[ack_id])]) - mock_pubsub.return_value.close.assert_has_calls([mock.call()]) + mock_pubsub.assert_called_once_with() + mock_pubsub.return_value.close.assert_not_called() def test_read_strings_success(self, mock_pubsub): data = '🤷 ¯\\_(ツ)_/¯' @@ -555,7 +620,7 @@ def test_read_strings_success(self, mock_pubsub): mock_pubsub.return_value.acknowledge.assert_has_calls( [mock.call(subscription=mock.ANY, ack_ids=[ack_id])]) - mock_pubsub.return_value.close.assert_has_calls([mock.call()]) + mock_pubsub.return_value.close.assert_not_called() def test_read_data_success(self, mock_pubsub): data_encoded = '🤷 ¯\\_(ツ)_/¯'.encode('utf-8') @@ -575,7 +640,7 @@ def test_read_data_success(self, mock_pubsub): mock_pubsub.return_value.acknowledge.assert_has_calls( [mock.call(subscription=mock.ANY, ack_ids=[ack_id])]) - mock_pubsub.return_value.close.assert_has_calls([mock.call()]) + mock_pubsub.return_value.close.assert_not_called() def test_read_messages_timestamp_attribute_milli_success(self, mock_pubsub): data = b'data' @@ -610,7 +675,7 @@ def test_read_messages_timestamp_attribute_milli_success(self, mock_pubsub): mock_pubsub.return_value.acknowledge.assert_has_calls( [mock.call(subscription=mock.ANY, ack_ids=[ack_id])]) - mock_pubsub.return_value.close.assert_has_calls([mock.call()]) + mock_pubsub.return_value.close.assert_not_called() def test_read_messages_timestamp_attribute_rfc3339_success(self, mock_pubsub): data = b'data' @@ -645,7 +710,7 @@ def test_read_messages_timestamp_attribute_rfc3339_success(self, mock_pubsub): mock_pubsub.return_value.acknowledge.assert_has_calls( [mock.call(subscription=mock.ANY, ack_ids=[ack_id])]) - mock_pubsub.return_value.close.assert_has_calls([mock.call()]) + mock_pubsub.return_value.close.assert_not_called() def test_read_messages_timestamp_attribute_missing(self, mock_pubsub): data = b'data' @@ -681,7 +746,7 @@ def test_read_messages_timestamp_attribute_missing(self, mock_pubsub): mock_pubsub.return_value.acknowledge.assert_has_calls( [mock.call(subscription=mock.ANY, ack_ids=[ack_id])]) - mock_pubsub.return_value.close.assert_has_calls([mock.call()]) + mock_pubsub.return_value.close.assert_not_called() def test_read_messages_timestamp_attribute_fail_parse(self, mock_pubsub): data = b'data' @@ -710,7 +775,7 @@ def test_read_messages_timestamp_attribute_fail_parse(self, mock_pubsub): p.run() mock_pubsub.return_value.acknowledge.assert_not_called() - mock_pubsub.return_value.close.assert_has_calls([mock.call()]) + mock_pubsub.return_value.close.assert_not_called() def test_read_message_id_label_unsupported(self, unused_mock_pubsub): # id_label is unsupported in DirectRunner. diff --git a/sdks/python/apache_beam/ml/inference/test_resources/vllm.dockerfile.old b/sdks/python/apache_beam/ml/inference/test_resources/vllm.dockerfile.old index b9c99e49e02f..d0080debc092 100644 --- a/sdks/python/apache_beam/ml/inference/test_resources/vllm.dockerfile.old +++ b/sdks/python/apache_beam/ml/inference/test_resources/vllm.dockerfile.old @@ -34,14 +34,22 @@ RUN python3 --version RUN apt-get install -y curl RUN curl -sS https://bootstrap.pypa.io/get-pip.py | python3.12 && pip install --upgrade pip -RUN pip install --no-cache-dir -vvv apache-beam[gcp]==2.58.1 -RUN pip install openai vllm +RUN pip install --no-cache-dir -vvv apache-beam[gcp]==2.71.0 +RUN pip install --no-cache-dir openai vllm ai-dynamo[vllm] RUN apt install libcairo2-dev pkg-config python3-dev -y RUN pip install pycairo +# etcd binary required by embedded NVIDIA Dynamo for runtime discovery. +ENV ETCD_VERSION=v3.5.13 +RUN curl -L https://github.com/etcd-io/etcd/releases/download/${ETCD_VERSION}/etcd-${ETCD_VERSION}-linux-amd64.tar.gz -o /tmp/etcd.tar.gz && \ + tar xzf /tmp/etcd.tar.gz -C /tmp && \ + mv /tmp/etcd-${ETCD_VERSION}-linux-amd64/etcd /usr/local/bin/etcd && \ + chmod +x /usr/local/bin/etcd && \ + rm -rf /tmp/etcd* + # Copy the Apache Beam worker dependencies from the Beam Python 3.12 SDK image. -COPY --from=apache/beam_python3.12_sdk:2.58.1 /opt/apache/beam /opt/apache/beam +COPY --from=apache/beam_python3.12_sdk:2.71.0 /opt/apache/beam /opt/apache/beam # Set the entrypoint to Apache Beam SDK worker launcher. ENTRYPOINT [ "/opt/apache/beam/boot" ] \ No newline at end of file diff --git a/sdks/python/apache_beam/ml/inference/vllm_inference.py b/sdks/python/apache_beam/ml/inference/vllm_inference.py index 38283f1efd42..e5d8918d5ec5 100644 --- a/sdks/python/apache_beam/ml/inference/vllm_inference.py +++ b/sdks/python/apache_beam/ml/inference/vllm_inference.py @@ -20,10 +20,12 @@ import asyncio import logging import os +import shutil import subprocess import sys import threading import time +import urllib.request import uuid from collections.abc import Callable from collections.abc import Iterable @@ -109,36 +111,216 @@ def getAsyncVLLMClient(port) -> AsyncOpenAI: ) +# Embedded Dynamo runtime defaults proven on the smoke test: etcd discovery, +# TCP request plane, ZMQ event plane, KV events disabled. KV-aware routing, +# disaggregated prefill/decode, and the Planner are not active in this mode. +_DYNAMO_FRONTEND_DEFAULT_KWARGS: dict[str, Optional[str]] = { + 'discovery-backend': 'etcd', + 'request-plane': 'tcp', + 'event-plane': 'zmq', + 'router-mode': 'round-robin', + 'no-router-kv-events': None, +} + +_DYNAMO_ENGINE_DEFAULT_KWARGS: dict[str, Optional[str]] = { + 'discovery-backend': 'etcd', + 'request-plane': 'tcp', + 'event-plane': 'zmq', + 'kv-events-config': '{"enable_kv_cache_events": false}', +} + + +def _append_kwargs(cmd: list[str], kwargs: dict[str, Optional[str]]) -> None: + for k, v in kwargs.items(): + cmd.append(f'--{k}') + # Only add values for commands with value part. + if v is not None: + cmd.append(v) + + +def _uses_etcd_discovery(kwargs: dict[str, Optional[str]]) -> bool: + return kwargs.get('discovery-backend') == 'etcd' + + class _VLLMModelServer(): - def __init__(self, model_name: str, vllm_server_kwargs: dict[str, str]): + def __init__( + self, + model_name: str, + vllm_server_kwargs: dict[str, Optional[str]], + dynamo_frontend_kwargs: Optional[dict[str, Optional[str]]] = None, + use_dynamo: bool = False): self._model_name = model_name self._vllm_server_kwargs = vllm_server_kwargs + self._dynamo_frontend_kwargs = dynamo_frontend_kwargs or {} self._server_started = False self._server_process = None + self._dynamo_process = None + self._etcd_process = None + self._etcd_data_dir: Optional[str] = None + self._managed_etcd_endpoint = None self._server_port: int = -1 self._server_process_lock = threading.RLock() + self._use_dynamo = use_dynamo self.start_server() + @staticmethod + def _stop_process(process: Optional[subprocess.Popen]) -> None: + if process is None or process.poll() is not None: + return + # A process may exit between poll() and terminate() / kill(), in which + # case the OS raises ProcessLookupError (or another OSError). Treat that + # as already-stopped so we don't bail out of the broader cleanup. + try: + process.terminate() + try: + process.wait(timeout=10) + except subprocess.TimeoutExpired: + process.kill() + process.wait() + except OSError: + pass + + def _stop_processes(self) -> None: + self._stop_process(self._dynamo_process) + self._stop_process(self._server_process) + self._stop_process(self._etcd_process) + if (self._managed_etcd_endpoint is not None and + os.environ.get('ETCD_ENDPOINTS') == self._managed_etcd_endpoint): + os.environ.pop('ETCD_ENDPOINTS', None) + if self._etcd_data_dir is not None: + shutil.rmtree(self._etcd_data_dir, ignore_errors=True) + self._etcd_data_dir = None + self._dynamo_process = None + self._server_process = None + self._etcd_process = None + self._managed_etcd_endpoint = None + self._server_started = False + self._server_port = -1 + + def _process_status(self) -> str: + process_status = [] + if self._server_process is not None: + process_status.append( + 'frontend/server exit code: %s' % self._server_process.poll()) + if self._dynamo_process is not None: + process_status.append( + 'dynamo worker exit code: %s' % self._dynamo_process.poll()) + if self._etcd_process is not None: + process_status.append('etcd exit code: %s' % self._etcd_process.poll()) + return ', '.join(process_status) or 'no process status available' + + def __del__(self): + # __del__ may run during interpreter shutdown when module globals can + # already be torn down; swallow any cleanup failures so we don't print + # a noisy traceback. + try: + self._stop_processes() + except Exception: # pylint: disable=broad-except + pass + + def _uses_embedded_etcd(self) -> bool: + return ( + self._use_dynamo and + _uses_etcd_discovery(self._dynamo_frontend_kwargs) and + _uses_etcd_discovery(self._vllm_server_kwargs) and + 'ETCD_ENDPOINTS' not in os.environ) + + def _wait_for_etcd(self, endpoint: str, timeout_secs=30) -> None: + deadline = time.time() + timeout_secs + health_url = endpoint.rstrip('/') + '/health' + while time.time() < deadline and self._etcd_process.poll() is None: + try: + with urllib.request.urlopen(health_url, timeout=2) as response: + if response.status < 500: + return + except Exception: # pylint: disable=broad-except + time.sleep(1) + + process_status = self._process_status() + self._stop_processes() + raise RuntimeError( + "Failed to start embedded etcd for Dynamo. Process status: " + f"{process_status}. Install etcd in the worker container or set " + "ETCD_ENDPOINTS to an external etcd service.") + + def _ensure_etcd(self) -> None: + if not self._uses_embedded_etcd(): + return + if shutil.which('etcd') is None: + raise RuntimeError( + "Embedded Dynamo mode requires etcd when ETCD_ENDPOINTS is not " + "set. Install etcd in the worker container or set ETCD_ENDPOINTS " + "to an external etcd service.") + + etcd_name = f'beam-dynamo-etcd-{uuid.uuid4().hex}' + self._etcd_data_dir = f'/tmp/{etcd_name}' + peer_port, = subprocess_server.pick_port(None) + etcd_cmd = [ + 'etcd', + '--name', + etcd_name, + '--listen-client-urls', + 'http://127.0.0.1:{{PORT}}', + '--advertise-client-urls', + 'http://127.0.0.1:{{PORT}}', + '--listen-peer-urls', + f'http://127.0.0.1:{peer_port}', + '--initial-advertise-peer-urls', + f'http://127.0.0.1:{peer_port}', + '--initial-cluster', + f'{etcd_name}=http://127.0.0.1:{peer_port}', + '--data-dir', + self._etcd_data_dir, + '--log-level', + 'warn', + ] + self._etcd_process, etcd_port = start_process(etcd_cmd) + endpoint = f'http://127.0.0.1:{etcd_port}' + os.environ['ETCD_ENDPOINTS'] = endpoint + self._managed_etcd_endpoint = endpoint + self._wait_for_etcd(endpoint) + def start_server(self, retries=3): with self._server_process_lock: if not self._server_started: - server_cmd = [ - sys.executable, - '-m', - 'vllm.entrypoints.openai.api_server', - '--model', - self._model_name, - '--port', - '{{PORT}}', - ] - for k, v in self._vllm_server_kwargs.items(): - server_cmd.append(f'--{k}') - # Only add values for commands with value part. - if v is not None: - server_cmd.append(v) + self._stop_processes() + self._ensure_etcd() + if self._use_dynamo: + # Dynamo embedded mode uses the frontend as its OpenAI-compatible + # local endpoint and a separate vLLM worker process. + server_cmd = [ + sys.executable, + '-m', + 'dynamo.frontend', + '--http-port', + '{{PORT}}', + ] + _append_kwargs(server_cmd, self._dynamo_frontend_kwargs) + else: + server_cmd = [ + sys.executable, + '-m', + 'vllm.entrypoints.openai.api_server', + '--model', + self._model_name, + '--port', + '{{PORT}}', + ] + _append_kwargs(server_cmd, self._vllm_server_kwargs) self._server_process, self._server_port = start_process(server_cmd) + if self._use_dynamo: + server_cmd = [ + sys.executable, + '-m', + 'dynamo.vllm', + '--model', + self._model_name, + ] + _append_kwargs(server_cmd, self._vllm_server_kwargs) + self._dynamo_process, _ = start_process(server_cmd) + self.check_connectivity(retries) def get_server_port(self) -> int: @@ -146,9 +328,14 @@ def get_server_port(self) -> int: self.start_server() return self._server_port - def check_connectivity(self, retries=3): + def check_connectivity(self, retries=3, timeout_secs=600): + start_time = time.time() with getVLLMClient(self._server_port) as client: - while self._server_process.poll() is None: + while (time.time() - start_time < timeout_secs and + self._server_process.poll() is None and + (self._dynamo_process is None or + self._dynamo_process.poll() is None) and + (self._etcd_process is None or self._etcd_process.poll() is None)): try: models = client.models.list().data logging.info('models: %s' % models) @@ -160,12 +347,13 @@ def check_connectivity(self, retries=3): # Sleep while bringing up the process time.sleep(5) + process_status = self._process_status() + self._stop_processes() if retries == 0: - self._server_started = False raise Exception( - "Failed to start vLLM server, polling process exited with code " + - "%s. Next time a request is tried, the server will be restarted" % - self._server_process.poll()) + "Failed to start vLLM server. Process status: " + f"{process_status}. Next time a request is tried, the server " + "will be restarted") else: self.start_server(retries - 1) @@ -176,8 +364,10 @@ class VLLMCompletionsModelHandler(ModelHandler[str, def __init__( self, model_name: str, - vllm_server_kwargs: Optional[dict[str, str]] = None, + vllm_server_kwargs: Optional[dict[str, Optional[str]]] = None, *, + use_dynamo: bool = False, + dynamo_frontend_kwargs: Optional[dict[str, Optional[str]]] = None, min_batch_size: Optional[int] = None, max_batch_size: Optional[int] = None, max_batch_duration_secs: Optional[int] = None, @@ -197,15 +387,28 @@ def __init__( https://docs.vllm.ai/en/latest/models/supported_models.html for supported models. vllm_server_kwargs: Any additional kwargs to be passed into your vllm - server when it is being created. Will be invoked using - `python -m vllm.entrypoints.openai.api_serverv - `. For example, you could pass - `{'echo': 'true'}` to prepend new messages with the previous message. - On ~16GB GPUs, pass lower ``max-num-seqs`` and - ``gpu-memory-utilization`` values (see - ``apache_beam.examples.inference.vllm_text_completion``). For a list of - possible kwargs, see + server when it is being created. When ``use_dynamo`` is disabled, + this is invoked using ``python -m vllm.entrypoints.openai.api_server + ``. When ``use_dynamo`` is + enabled, these kwargs are passed to the ``dynamo.vllm`` worker + process. For example, you could pass ``{'echo': 'true'}`` to prepend + new messages with the previous message. On ~16GB GPUs, pass lower + ``max-num-seqs`` and ``gpu-memory-utilization`` values (see + ``apache_beam.examples.inference.vllm_text_completion``). For a list + of possible kwargs, see https://docs.vllm.ai/en/latest/serving/openai_compatible_server.html#extra-parameters-for-completions-api + use_dynamo: Whether to use NVIDIA Dynamo as the underlying vLLM engine. + Requires installing Dynamo in your runtime environment + (``pip install ai-dynamo[vllm]``). This is an opt-in single-worker + embedded mode; KV-aware routing, disaggregated prefill/decode, KVBM + offload across nodes, the Planner, and Grove are not active in + embedded mode. Dynamo also requires an etcd-style discovery service: + when ``ETCD_ENDPOINTS`` is unset, Beam starts a local etcd, which + requires the ``etcd`` binary in the worker environment. + dynamo_frontend_kwargs: Additional kwargs to be passed to the + ``dynamo.frontend`` process when ``use_dynamo`` is enabled. By + default, embedded Dynamo uses etcd discovery, TCP request plane, ZMQ + event plane, round-robin routing, and disables router KV events. min_batch_size: optional. the minimum batch size to use when batching inputs. max_batch_size: optional. the maximum batch size to use when batching @@ -229,10 +432,20 @@ def __init__( batch_length_fn=batch_length_fn, batch_bucket_boundaries=batch_bucket_boundaries) self._model_name = model_name - self._vllm_server_kwargs: dict[str, str] = vllm_server_kwargs or {} + self._vllm_server_kwargs: dict[str, Optional[str]] = ({ + **_DYNAMO_ENGINE_DEFAULT_KWARGS, **(vllm_server_kwargs or {}) + } if use_dynamo else vllm_server_kwargs or {}) + self._dynamo_frontend_kwargs: dict[str, Optional[str]] = { + **_DYNAMO_FRONTEND_DEFAULT_KWARGS, **(dynamo_frontend_kwargs or {}) + } + self._use_dynamo = use_dynamo def load_model(self) -> _VLLMModelServer: - return _VLLMModelServer(self._model_name, self._vllm_server_kwargs) + return _VLLMModelServer( + self._model_name, + self._vllm_server_kwargs, + self._dynamo_frontend_kwargs, + self._use_dynamo) async def _async_run_inference( self, @@ -274,6 +487,12 @@ def run_inference( """ return asyncio.run(self._async_run_inference(batch, model, inference_args)) + def validate_inference_args(self, inference_args: Optional[dict[str, Any]]): + # Override the base validator so OpenAI-compatible request kwargs such as + # ``max_tokens`` can be passed through ``RunInference`` to the vLLM / + # Dynamo server. + pass + def share_model_across_processes(self) -> bool: return True @@ -285,8 +504,10 @@ def __init__( self, model_name: str, chat_template_path: Optional[str] = None, - vllm_server_kwargs: Optional[dict[str, str]] = None, + vllm_server_kwargs: Optional[dict[str, Optional[str]]] = None, *, + use_dynamo: bool = False, + dynamo_frontend_kwargs: Optional[dict[str, Optional[str]]] = None, min_batch_size: Optional[int] = None, max_batch_size: Optional[int] = None, max_batch_duration_secs: Optional[int] = None, @@ -311,12 +532,26 @@ def __init__( For info on chat templates, see: https://docs.vllm.ai/en/latest/serving/openai_compatible_server.html#chat-template vllm_server_kwargs: Any additional kwargs to be passed into your vllm - server when it is being created. Will be invoked using - `python -m vllm.entrypoints.openai.api_serverv - `. For example, you could pass - `{'echo': 'true'}` to prepend new messages with the previous message. - For a list of possible kwargs, see + server when it is being created. When ``use_dynamo`` is disabled, + this is invoked using ``python -m vllm.entrypoints.openai.api_server + ``. When ``use_dynamo`` is + enabled, these kwargs are passed to the ``dynamo.vllm`` worker + process. For example, you could pass ``{'echo': 'true'}`` to prepend + new messages with the previous message. For a list of possible + kwargs, see https://docs.vllm.ai/en/latest/serving/openai_compatible_server.html#extra-parameters-for-chat-api + use_dynamo: Whether to use NVIDIA Dynamo as the underlying vLLM engine. + Requires installing Dynamo in your runtime environment + (``pip install ai-dynamo[vllm]``). This is an opt-in single-worker + embedded mode; KV-aware routing, disaggregated prefill/decode, KVBM + offload across nodes, the Planner, and Grove are not active in + embedded mode. Dynamo also requires an etcd-style discovery service: + when ``ETCD_ENDPOINTS`` is unset, Beam starts a local etcd, which + requires the ``etcd`` binary in the worker environment. + dynamo_frontend_kwargs: Additional kwargs to be passed to the + ``dynamo.frontend`` process when ``use_dynamo`` is enabled. By + default, embedded Dynamo uses etcd discovery, TCP request plane, ZMQ + event plane, round-robin routing, and disables router KV events. min_batch_size: optional. the minimum batch size to use when batching inputs. max_batch_size: optional. the maximum batch size to use when batching @@ -340,9 +575,15 @@ def __init__( batch_length_fn=batch_length_fn, batch_bucket_boundaries=batch_bucket_boundaries) self._model_name = model_name - self._vllm_server_kwargs: dict[str, str] = vllm_server_kwargs or {} + self._vllm_server_kwargs: dict[str, Optional[str]] = ({ + **_DYNAMO_ENGINE_DEFAULT_KWARGS, **(vllm_server_kwargs or {}) + } if use_dynamo else vllm_server_kwargs or {}) + self._dynamo_frontend_kwargs: dict[str, Optional[str]] = { + **_DYNAMO_FRONTEND_DEFAULT_KWARGS, **(dynamo_frontend_kwargs or {}) + } self._chat_template_path = chat_template_path self._chat_file = f'template-{uuid.uuid4().hex}.jinja' + self._use_dynamo = use_dynamo def load_model(self) -> _VLLMModelServer: chat_template_contents = '' @@ -355,7 +596,11 @@ def load_model(self) -> _VLLMModelServer: f.write(chat_template_contents) self._vllm_server_kwargs['chat_template'] = local_chat_template_path - return _VLLMModelServer(self._model_name, self._vllm_server_kwargs) + return _VLLMModelServer( + self._model_name, + self._vllm_server_kwargs, + self._dynamo_frontend_kwargs, + self._use_dynamo) async def _async_run_inference( self, @@ -400,5 +645,11 @@ def run_inference( """ return asyncio.run(self._async_run_inference(batch, model, inference_args)) + def validate_inference_args(self, inference_args: Optional[dict[str, Any]]): + # Override the base validator so OpenAI-compatible request kwargs such as + # ``max_tokens`` can be passed through ``RunInference`` to the vLLM / + # Dynamo server. + pass + def share_model_across_processes(self) -> bool: return True diff --git a/sdks/python/apache_beam/ml/inference/vllm_inference_test.py b/sdks/python/apache_beam/ml/inference/vllm_inference_test.py new file mode 100644 index 000000000000..4ff5186178e3 --- /dev/null +++ b/sdks/python/apache_beam/ml/inference/vllm_inference_test.py @@ -0,0 +1,157 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +import os +import sys +import types +import unittest +from unittest import mock + +# Protect against environments where the OpenAI python library is not +# available. The command-construction tests below do not actually need a +# real OpenAI client; stubbing the module is enough for vllm_inference to +# import cleanly. +# pylint: disable=wrong-import-order, wrong-import-position +try: + import openai # pylint: disable=unused-import +except ImportError: + openai = types.ModuleType('openai') + + class _FakeOpenAI: + pass + + openai.AsyncOpenAI = _FakeOpenAI + openai.OpenAI = _FakeOpenAI + sys.modules['openai'] = openai + +from apache_beam.ml.inference import vllm_inference + + +class _FakeProcess: + def __init__(self): + self.returncode = None + + def poll(self): + return self.returncode + + def terminate(self): + self.returncode = 0 + + def wait(self, timeout=None): + return self.returncode + + def kill(self): + self.returncode = -9 + + +class _FakeModels: + def list(self): + return types.SimpleNamespace(data=[object()]) + + +class _FakeClient: + def __init__(self): + self.models = _FakeModels() + + def __enter__(self): + return self + + def __exit__(self, exc_type, exc_value, traceback): + return False + + +def _record_start_process(commands): + def start_process(cmd): + commands.append(list(cmd)) + return _FakeProcess(), 10000 + len(commands) + + return start_process + + +class VLLMInferenceTest(unittest.TestCase): + def test_native_vllm_starts_single_server_process(self): + commands = [] + with mock.patch.object(vllm_inference, + 'start_process', + _record_start_process(commands)): + with mock.patch.object(vllm_inference, 'getVLLMClient'): + vllm_inference.getVLLMClient.return_value = _FakeClient() + vllm_inference.VLLMCompletionsModelHandler( + model_name='test-model', + vllm_server_kwargs={ + 'gpu-memory-utilization': '0.9' + }).load_model() + self.assertEqual(1, len(commands)) + self.assertIn('vllm.entrypoints.openai.api_server', commands[0]) + self.assertIn('--model', commands[0]) + self.assertIn('test-model', commands[0]) + self.assertIn('--gpu-memory-utilization', commands[0]) + self.assertIn('0.9', commands[0]) + self.assertNotIn('dynamo.frontend', commands[0]) + self.assertNotIn('dynamo.vllm', commands[0]) + + def test_dynamo_starts_frontend_and_engine_with_separate_kwargs(self): + commands = [] + with mock.patch.dict(os.environ, + {'ETCD_ENDPOINTS': 'http://127.0.0.1:2379'}): + with mock.patch.object(vllm_inference, + 'start_process', + _record_start_process(commands)): + with mock.patch.object(vllm_inference, 'getVLLMClient'): + vllm_inference.getVLLMClient.return_value = _FakeClient() + vllm_inference.VLLMCompletionsModelHandler( + model_name='test-model', + vllm_server_kwargs={ + 'tensor-parallel-size': '1' + }, + use_dynamo=True, + dynamo_frontend_kwargs={ + 'router-mode': 'round-robin' + }).load_model() + self.assertEqual(2, len(commands)) + frontend_cmd = commands[0] + engine_cmd = commands[1] + self.assertIn('dynamo.frontend', frontend_cmd) + self.assertIn('--http-port', frontend_cmd) + self.assertIn('--discovery-backend', frontend_cmd) + self.assertIn('--request-plane', frontend_cmd) + self.assertIn('--event-plane', frontend_cmd) + self.assertIn('--router-mode', frontend_cmd) + self.assertIn('--no-router-kv-events', frontend_cmd) + self.assertNotIn('--model', frontend_cmd) + self.assertNotIn('--tensor-parallel-size', frontend_cmd) + self.assertNotIn('--kv-events-config', frontend_cmd) + self.assertIn('dynamo.vllm', engine_cmd) + self.assertIn('--model', engine_cmd) + self.assertIn('test-model', engine_cmd) + self.assertIn('--discovery-backend', engine_cmd) + self.assertIn('--request-plane', engine_cmd) + self.assertIn('--event-plane', engine_cmd) + self.assertIn('--kv-events-config', engine_cmd) + self.assertIn('--tensor-parallel-size', engine_cmd) + self.assertNotIn('--http-port', engine_cmd) + self.assertNotIn('--router-mode', engine_cmd) + self.assertNotIn('--no-router-kv-events', engine_cmd) + + def test_validate_inference_args_accepts_openai_request_kwargs(self): + vllm_inference.VLLMCompletionsModelHandler( + 'test-model').validate_inference_args({'max_tokens': 8}) + vllm_inference.VLLMChatModelHandler('test-model').validate_inference_args( + {'max_tokens': 8}) + + +if __name__ == '__main__': + unittest.main() diff --git a/sdks/python/apache_beam/runners/direct/transform_evaluator.py b/sdks/python/apache_beam/runners/direct/transform_evaluator.py index 2349a0881e40..6702ec3362bc 100644 --- a/sdks/python/apache_beam/runners/direct/transform_evaluator.py +++ b/sdks/python/apache_beam/runners/direct/transform_evaluator.py @@ -19,11 +19,12 @@ # pytype: skip-file -import atexit import collections import logging import random +import threading import time +import weakref from collections import abc from typing import TYPE_CHECKING from typing import Any @@ -577,13 +578,46 @@ def finish_bundle(self): self, self.bundles, unprocessed_bundles, None, {None: self.watermark}) +class _PubSubSubscriberClient(object): + """SubscriberClient state cached for one DirectRunner Pub/Sub read.""" + def __init__(self, client): + self.client = client + self._temporary_subscription = None + self._closed = False + self._lock = threading.Lock() + + def set_temporary_subscription(self, subscription): + self._temporary_subscription = subscription + + def close(self): + if self._closed: + return + self._closed = True + + try: + if self._temporary_subscription: + self.client.delete_subscription( + subscription=self._temporary_subscription) + except Exception: + _LOGGER.warning( + 'Failed to delete temporary Pub/Sub subscription %s', + self._temporary_subscription, + exc_info=True) + + try: + self.client.close() + except Exception: + _LOGGER.warning( + 'Failed to close Pub/Sub subscriber client', exc_info=True) + + class _PubSubReadEvaluator(_TransformEvaluator): """TransformEvaluator for PubSub read.""" - # A mapping of transform to _PubSubSubscriptionWrapper. - # TODO(https://github.com/apache/beam/issues/19751): Prevents garbage - # collection of pipeline instances. - _subscription_cache: dict[AppliedPTransform, str] = {} + # Weak-keyed per-transform caches avoid keeping completed pipelines alive. + _subscription_cache = weakref.WeakKeyDictionary() + _subscriber_client_cache = weakref.WeakKeyDictionary() + _subscriber_client_cache_lock = threading.Lock() def __init__( self, @@ -627,18 +661,46 @@ def get_subscription( if short_sub_name: return pubsub.SubscriberClient.subscription_path(project, short_sub_name) - if transform in cls._subscription_cache: - return cls._subscription_cache[transform] + with cls._subscriber_client_cache_lock: + sub_name = cls._subscription_cache.get(transform) + if sub_name: + return sub_name - sub_client = pubsub.SubscriberClient() - sub_name = sub_client.subscription_path( - sub_project, - 'beam_%d_%x' % (int(time.time()), random.randrange(1 << 32))) - topic_name = sub_client.topic_path(project, short_topic_name) - sub_client.create_subscription(name=sub_name, topic=topic_name) - atexit.register(sub_client.delete_subscription, subscription=sub_name) - cls._subscription_cache[transform] = sub_name - return cls._subscription_cache[transform] + subscriber_client = cls._get_subscriber_client_state_unlocked(transform) + + with subscriber_client._lock: + with cls._subscriber_client_cache_lock: + sub_name = cls._subscription_cache.get(transform) + if sub_name: + return sub_name + + sub_client = subscriber_client.client + sub_name = sub_client.subscription_path( + sub_project, + 'beam_%d_%x' % (int(time.time()), random.randrange(1 << 32))) + topic_name = sub_client.topic_path(project, short_topic_name) + sub_client.create_subscription(name=sub_name, topic=topic_name) + subscriber_client.set_temporary_subscription(sub_name) + + with cls._subscriber_client_cache_lock: + cls._subscription_cache[transform] = sub_name + + return sub_name + + @classmethod + def _get_subscriber_client(cls, transform): + with cls._subscriber_client_cache_lock: + return cls._get_subscriber_client_state_unlocked(transform).client + + @classmethod + def _get_subscriber_client_state_unlocked(cls, transform): + subscriber_client = cls._subscriber_client_cache.get(transform) + if subscriber_client is None: + from google.cloud import pubsub + subscriber_client = _PubSubSubscriberClient(pubsub.SubscriberClient()) + cls._subscriber_client_cache[transform] = subscriber_client + weakref.finalize(transform, subscriber_client.close) + return subscriber_client def start_bundle(self): pass @@ -648,8 +710,6 @@ def process_element(self, element): def _read_from_pubsub( self, timestamp_attribute) -> list[tuple[Timestamp, 'PubsubMessage']]: - from google.cloud import pubsub - from apache_beam.io.gcp.pubsub import PubsubMessage def _get_element(message): @@ -678,16 +738,13 @@ def _get_element(message): # evaluator fails with an exception before emitting a bundle. However, # the DirectRunner currently doesn't retry work items anyway, so the # pipeline would enter an inconsistent state on any error. - sub_client = pubsub.SubscriberClient() - try: - response = sub_client.pull( - subscription=self._sub_name, max_messages=10, timeout=30) - results = [_get_element(rm.message) for rm in response.received_messages] - ack_ids = [rm.ack_id for rm in response.received_messages] - if ack_ids: - sub_client.acknowledge(subscription=self._sub_name, ack_ids=ack_ids) - finally: - sub_client.close() + sub_client = self._get_subscriber_client(self._applied_ptransform) + response = sub_client.pull( + subscription=self._sub_name, max_messages=10, timeout=30) + results = [_get_element(rm.message) for rm in response.received_messages] + ack_ids = [rm.ack_id for rm in response.received_messages] + if ack_ids: + sub_client.acknowledge(subscription=self._sub_name, ack_ids=ack_ids) return results diff --git a/sdks/python/apache_beam/utils/subprocess_server.py b/sdks/python/apache_beam/utils/subprocess_server.py index 6e78ce88e0d7..38afe305b0f3 100644 --- a/sdks/python/apache_beam/utils/subprocess_server.py +++ b/sdks/python/apache_beam/utils/subprocess_server.py @@ -201,45 +201,56 @@ def __exit__(self, *unused_args): self.stop() def start(self): - try: - process, endpoint = self.start_process() - wait_secs = .1 - channel_options = [ - ("grpc.max_receive_message_length", -1), - ("grpc.max_send_message_length", -1), - # Default: 20000ms (20s), increased to 10 minutes for stability - ("grpc.keepalive_timeout_ms", 600_000), - # Default: 2, set to 0 to allow unlimited pings without data - ("grpc.http2.max_pings_without_data", 0), - # Default: False, set to True to allow keepalive pings when no calls - ("grpc.keepalive_permit_without_calls", True), - # Default: 2, set to 0 to allow unlimited ping strikes - ("grpc.http2.max_ping_strikes", 0), - # Default: 0 (disabled), enable socket reuse for better handling - ("grpc.so_reuseport", 1), - ] - self._grpc_channel = grpc.insecure_channel( - endpoint, options=channel_options) - channel_ready = grpc.channel_ready_future(self._grpc_channel) - while True: - if process is not None and process.poll() is not None: - _LOGGER.error("Failed to start job service with %s", process.args) - raise RuntimeError( - 'Service failed to start up with error %s' % process.poll()) - try: - channel_ready.result(timeout=wait_secs) - break - except (grpc.FutureTimeoutError, grpc.RpcError): - wait_secs *= 1.2 - logging.log( - logging.WARNING if wait_secs > 1 else logging.DEBUG, - 'Waiting for grpc channel to be ready at %s.', - endpoint) - return self._stub_class(self._grpc_channel) - except: # pylint: disable=bare-except - _LOGGER.exception("Error bringing up service") - self.stop_force() - raise + max_retries = 3 + for attempt in range(max_retries): + try: + process, endpoint = self.start_process() + wait_secs = .1 + channel_options = [ + ("grpc.max_receive_message_length", -1), + ("grpc.max_send_message_length", -1), + # Default: 20000ms (20s), increased to 10 minutes for stability + ("grpc.keepalive_timeout_ms", 600_000), + # Default: 2, set to 0 to allow unlimited pings without data + ("grpc.http2.max_pings_without_data", 0), + # Default: False, set to True to allow keepalive pings when no calls + ("grpc.keepalive_permit_without_calls", True), + # Default: 2, set to 0 to allow unlimited ping strikes + ("grpc.http2.max_ping_strikes", 0), + # Default: 0 (disabled), enable socket reuse for better handling + ("grpc.so_reuseport", 1), + ] + self._grpc_channel = grpc.insecure_channel( + endpoint, options=channel_options) + channel_ready = grpc.channel_ready_future(self._grpc_channel) + while True: + if process is not None and process.poll() is not None: + _LOGGER.error("Failed to start job service with %s", process.args) + raise RuntimeError( + 'Service failed to start up with error %s' % process.poll()) + try: + channel_ready.result(timeout=wait_secs) + break + except (grpc.FutureTimeoutError, grpc.RpcError): + wait_secs *= 1.2 + logging.log( + logging.WARNING if wait_secs > 1 else logging.DEBUG, + 'Waiting for grpc channel to be ready at %s.', + endpoint) + return self._stub_class(self._grpc_channel) + except Exception as e: + _LOGGER.warning( + "Error bringing up service (attempt %d of %d): %s", + attempt + 1, + max_retries, + e) + self.stop_force() + if attempt == max_retries - 1: + raise + time.sleep(1) + except: # pylint: disable=bare-except + self.stop_force() + raise def start_process(self): if self._owner_id is not None: diff --git a/sdks/python/apache_beam/utils/subprocess_server_test.py b/sdks/python/apache_beam/utils/subprocess_server_test.py index a008ae05c52d..6f623e336070 100644 --- a/sdks/python/apache_beam/utils/subprocess_server_test.py +++ b/sdks/python/apache_beam/utils/subprocess_server_test.py @@ -514,9 +514,16 @@ def __init__(self): def poll(self): return 1 # Simulate that process exited/failed + constructor_calls = 0 + + def custom_constructor(*args): + nonlocal constructor_calls + constructor_calls += 1 + return (dummy_process, "localhost:12345") + dummy_process = DummyProcess() cache = subprocess_server._SharedCache( - lambda *args: (dummy_process, "localhost:12345"), custom_destructor) + custom_constructor, custom_destructor) # 1. Register an independent, unrelated owner in the cache first. other_owner = cache.register() @@ -536,11 +543,14 @@ def __init__(self): self.assertEqual(cache._cache[cache_key].owners, {other_owner}) # 2. Verify starting the server (which registers its own owner and retrieves from cache) raises RuntimeError - with self.assertRaises(RuntimeError): - server.start() - - # 3. Verify that the destructor was called on the process, meaning no leak (even though other_owner was still registered!) - self.assertEqual(destructor_calls, [(dummy_process, "localhost:12345")]) + with patch('time.sleep'): + with self.assertRaises(RuntimeError): + server.start() + self.assertEqual(constructor_calls, 3) + + # 3. Verify that the destructor was called on the process for each retry attempt (3 total), + # meaning there is no leak (even though other_owner was still registered). + self.assertEqual(destructor_calls, [(dummy_process, "localhost:12345")] * 3) # 4. Verify that the server has cleaned up its owner_id self.assertIsNone(server._owner_id) diff --git a/sdks/python/test-suites/dataflow/common.gradle b/sdks/python/test-suites/dataflow/common.gradle index 480e2a62a2ef..c450eb3612fa 100644 --- a/sdks/python/test-suites/dataflow/common.gradle +++ b/sdks/python/test-suites/dataflow/common.gradle @@ -450,25 +450,28 @@ def tensorRTTests = tasks.create("tensorRTtests") { } } +def vllmBaseArgMap = [ + "runner": "DataflowRunner", + "machine_type":"n1-standard-4", + // TODO(https://github.com/apache/beam/issues/22651): Build docker image for VLLM tests during Run time. + // This would also enable to use wheel "--sdk_location" as other tasks, and eliminate distTarBall dependency + // declaration for this project. + // Right now, this is built from https://github.com/apache/beam/blob/master/sdks/python/apache_beam/ml/inference/test_resources/vllm.dockerfile.old + "sdk_container_image": "us.gcr.io/apache-beam-testing/python-postcommit-it/vllm:latest", + "sdk_location": files(configurations.distTarBall.files).singleFile, + "project": "apache-beam-testing", + "region": "us-central1", + "disk_size_gb": 75 +] + def vllmTests = tasks.create("vllmTests") { dependsOn 'installGcpTest' dependsOn ':sdks:python:sdist' doLast { def testOpts = basicPytestOpts - def argMap = [ - "runner": "DataflowRunner", - "machine_type":"n1-standard-4", - // TODO(https://github.com/apache/beam/issues/22651): Build docker image for VLLM tests during Run time. - // This would also enable to use wheel "--sdk_location" as other tasks, and eliminate distTarBall dependency - // declaration for this project. - // Right now, this is built from https://github.com/apache/beam/blob/master/sdks/python/apache_beam/ml/inference/test_resources/vllm.dockerfile.old - "sdk_container_image": "us.gcr.io/apache-beam-testing/python-postcommit-it/vllm:latest", - "sdk_location": files(configurations.distTarBall.files).singleFile, - "project": "apache-beam-testing", - "region": "us-central1", + def argMap = vllmBaseArgMap + [ "model": "facebook/opt-125m", "output": "gs://apache-beam-ml/outputs/vllm_predictions.txt", - "disk_size_gb": 75 ] def cmdArgs = mapToArgString(argMap) // Exec one version with and one version without the chat option @@ -480,6 +483,19 @@ def vllmTests = tasks.create("vllmTests") { executable 'sh' args '-c', ". ${envdir}/bin/activate && pip install openai && python -m apache_beam.examples.inference.vllm_text_completion $cmdArgs --chat true --chat_template 'gs://apache-beam-ml/additional_files/sample_chat_template.jinja' --experiment='worker_accelerator=type:nvidia-tesla-t4;count:1;install-nvidia-driver:5xx'" } + // Embedded NVIDIA Dynamo path. Reuses the same sdk_container_image + // (vllm.dockerfile.old now installs etcd and ai-dynamo[vllm]) and the + // same nvidia-tesla-t4 accelerator as the native vLLM ITs above. + // Validated end-to-end on Dataflow with Qwen/Qwen3-0.6B on T4. + def dynamoArgMap = vllmBaseArgMap + [ + "model": "Qwen/Qwen3-0.6B", + "output": "gs://apache-beam-ml/outputs/vllm_dynamo_predictions.txt", + ] + def dynamoCmdArgs = mapToArgString(dynamoArgMap) + exec { + executable 'sh' + args '-c', ". ${envdir}/bin/activate && pip install openai && python -m apache_beam.examples.inference.vllm_text_completion $dynamoCmdArgs --use_dynamo --max_tokens 8 --experiment='worker_accelerator=type:nvidia-tesla-t4;count:1;install-nvidia-driver:5xx'" + } } } diff --git a/settings.gradle.kts b/settings.gradle.kts index 3d4346661a47..d9dbbf9021ef 100644 --- a/settings.gradle.kts +++ b/settings.gradle.kts @@ -39,7 +39,7 @@ pluginManagement { plugins { id("com.gradle.develocity") version "3.19" - id("com.gradle.common-custom-user-data-gradle-plugin") version "2.6.0" + id("com.gradle.common-custom-user-data-gradle-plugin") version "2.7.0" } // JENKINS_HOME and BUILD_ID set automatically during Jenkins execution diff --git a/website/www/yarn.lock b/website/www/yarn.lock index 04f6526dfd56..54100d4a20b2 100644 --- a/website/www/yarn.lock +++ b/website/www/yarn.lock @@ -442,9 +442,9 @@ isexe@^2.0.0: integrity sha1-6PvzdNxVb/iUehDcsFctYz8s+hA= js-yaml@^3.13.1: - version "3.14.2" - resolved "https://registry.yarnpkg.com/js-yaml/-/js-yaml-3.14.2.tgz#77485ce1dd7f33c061fd1b16ecea23b55fcb04b0" - integrity sha512-PMSmkqxr106Xa156c2M265Z+FTrPl+oxd/rgOQy2tijQeK5TxQ43psO1ZCwhVOSdnn+RzkzlRz/eY4BgJBYVpg== + version "3.15.0" + resolved "https://registry.yarnpkg.com/js-yaml/-/js-yaml-3.15.0.tgz#586e5214eafe3e893756a41e979b50d89d3e4a67" + integrity sha512-ttBQIIQPDeLjpPOohtUdXuXUVoA2uIB6fEH9HyJ7234s5mBJ5wTx20njxplLZQgLaOfpmPQA7X2t5AX6tIPbog== dependencies: argparse "^1.0.7" esprima "^4.0.0"