Repository: ray-project/ray
File: python/ray/llm/tests/batch/cpu/stages/test_http_request_stage.py (Python)

Code before:
        data = [
            {
                "payload": {
                    "model": "test-model",
                    "messages": [
                        {
                            "role": "user",
                            "content": [
                                {
                                    "type": "image_url",
                                    "image_url": {
                                        "url": "http://example.com/image.png"
                                    },
                                },
                                {
                                    "type": "text",
                                    "text": "hello",
                                },
                            ],
                        }
                    ],
                }
            }
        ]

        # Simulate ray data map batches which converts data from pyarrow table then to numpy
        table = pa.Table.from_pydict({"__data": data})
        batch = {k: v.to_numpy() for k, v in table.to_pandas().items()}

        # Execute the UDF
        results = []
        async for result in udf(batch):
            results.extend(result["__data"])

        # Verify results
        assert len(results) == 1
        assert results[0]["http_response"]["response"] == "success"
    finally:
        await server.close()


if __name__ == "__main__":
    sys.exit(pytest.main(["-v", __file__]))


Review comment (performance):
Could we directly hit the `map_batches` path without simulation? I'm not a big fan of simulating the pyarrow / numpy conversion. We can create a `HTTPRequestProcessor` and pass in `ray.data.Dataset`.

Code after:
    async def handler(request):
        data = await request.json()
        assert data["model"] == "test-model"
        assert data["embedding"] == [1.0, 2.0, 3.0]
        assert data["flags"] == [True, False]
        assert len(data["messages"]) == 1
        assert data["messages"][0]["role"] == "user"
        assert len(data["messages"][0]["content"]) == 2
        return aiohttp.web.json_response({"response": "success"})

    # Create test app and server
    app = aiohttp.web.Application()
    app.router.add_post("/", handler)
    server = TestServer(app)
    await server.start_server()
    yield server
    await server.close()


@pytest.mark.asyncio
async def test_http_request_udf_with_numpy_payload_server(numpy_payload_server):
    """Test HttpRequestUDF with numpy arrays using a real aiohttp server."""
    data = [
        {
            "payload": {
                "model": "test-model",
                "embedding": np.array([1.0, 2.0, 3.0]),
                "flags": np.array([True, False]),
                "messages": [
                    {
                        "role": "user",
                        "content": [
                            {
                                "type": "image_url",
                                "image_url": {"url": "http://example.com/image.png"},
                            },
                            {
                                "type": "text",
                                "text": "hello",
                            },
                        ],
                    }
                ],
            }
        }
    ]
    config = HttpRequestProcessorConfig(
        url=str(numpy_payload_server.make_url("/")),
        headers={"Content-Type": "application/json"},
        qps=None,
    )