| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """This example shows how to use Ray Data for data parallel batch inference. |
| 4 | |
| 5 | Ray Data is a data processing framework that can process very large datasets |
| 6 | with first-class support for vLLM. |
| 7 | |
| 8 | Ray Data provides functionality for: |
| 9 | * Reading and writing to most popular file formats and cloud object storage. |
| 10 | * Streaming execution, so you can run inference on datasets that far exceed |
| 11 | the aggregate RAM of the cluster. |
| 12 | * Scale up the workload without code changes. |
| 13 | * Automatic sharding, load-balancing, and autoscaling across a Ray cluster, |
| 14 | with built-in fault-tolerance and retry semantics. |
| 15 | * Continuous batching that keeps vLLM replicas saturated and maximizes GPU |
| 16 | utilization. |
| 17 | * Compatible with tensor/pipeline parallel inference. |
| 18 | |
| 19 | Learn more about Ray Data's LLM integration: |
| 20 | https://docs.ray.io/en/latest/data/working-with-llms.html |
| 21 | """ |
| 22 | |
| 23 | import ray |
| 24 | from packaging.version import Version |
| 25 | from ray.data.llm import build_llm_processor, vLLMEngineProcessorConfig |
| 26 | |
| 27 | assert Version(ray.__version__) >= Version("2.44.1"), ( |
| 28 | "Ray version must be at least 2.44.1" |
| 29 | ) |
| 30 | |
| 31 | # Uncomment to reduce clutter in stdout |
| 32 | # ray.init(log_to_driver=False) |
| 33 | # ray.data.DataContext.get_current().enable_progress_bars = False |
| 34 | |
| 35 | # Read one text file from S3. Ray Data supports reading multiple files |
| 36 | # from cloud storage (such as JSONL, Parquet, CSV, binary format). |
| 37 | ds = ray.data.read_text("s3://anonymous@air-example-data/prompts.txt") |
| 38 | print(ds.schema()) |
| 39 | |
| 40 | size = ds.count() |
| 41 | print(f"Size of dataset: {size} prompts") |
| 42 | |
| 43 | # Configure vLLM engine. |
| 44 | config = vLLMEngineProcessorConfig( |
| 45 | model_source="unsloth/Llama-3.1-8B-Instruct", |
| 46 | engine_kwargs={ |
| 47 | "enable_chunked_prefill": True, |
| 48 | "max_num_batched_tokens": 4096, |
| 49 | "max_model_len": 16384, |
| 50 | }, |
| 51 | concurrency=1, # set the number of parallel vLLM replicas |
| 52 | batch_size=64, |
| 53 | ) |
| 54 | |
| 55 | # Create a Processor object, which will be used to |
| 56 | # do batch inference on the dataset |
| 57 | vllm_processor = build_llm_processor( |
| 58 | config, |
| 59 | preprocess=lambda row: dict( |
| 60 | messages=[ |
| 61 | {"role": "system", "content": "You are a bot that responds with haikus."}, |
| 62 | {"role": "user", "content": row["text"]}, |
| 63 | ], |
| 64 | sampling_params=dict( |
| 65 | temperature=0.3, |
| 66 | max_tokens=250, |
| 67 | ), |
| 68 | ), |
| 69 | postprocess=lambda row: dict( |
| 70 | answer=row["generated_text"], |
| 71 | **row, # This will return all the original columns in the dataset. |
| 72 | ), |
| 73 | ) |
| 74 | |
| 75 | ds = vllm_processor(ds) |
| 76 | |
| 77 | # Peek first 10 results. |
| 78 | # NOTE: This is for local testing and debugging. For production use case, |
| 79 | # one should write full result out as shown below. |
| 80 | outputs = ds.take(limit=10) |
| 81 | |
| 82 | for output in outputs: |
| 83 | prompt = output["prompt"] |
| 84 | generated_text = output["generated_text"] |
| 85 | print(f"Prompt: {prompt!r}") |
| 86 | print(f"Generated text: {generated_text!r}") |
| 87 | |
| 88 | # Write inference output data out as Parquet files to S3. |
| 89 | # Multiple files would be written to the output destination, |
| 90 | # and each task would write one or more files separately. |
| 91 | # |
| 92 | # ds.write_parquet("s3://<your-output-bucket>") |