| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | |
| 5 | import openai # use the official client for correctness check |
| 6 | import pytest |
| 7 | import requests |
| 8 | |
| 9 | from tests.utils import RemoteOpenAIServer |
| 10 | |
| 11 | from .conftest import ( |
| 12 | HEADER_SAGEMAKER_CLOSED_SESSION_ID, |
| 13 | HEADER_SAGEMAKER_NEW_SESSION_ID, |
| 14 | HEADER_SAGEMAKER_SESSION_ID, |
| 15 | MODEL_NAME_SMOLLM, |
| 16 | ) |
| 17 | |
| 18 | CLOSE_BADREQUEST_CASES = [ |
| 19 | ( |
| 20 | "nonexistent_session_id", |
| 21 | {"session_id": "nonexistent-session-id"}, |
| 22 | {}, |
| 23 | "session not found", |
| 24 | ), |
| 25 | ("malformed_close_request", {}, {"extra-field": "extra-field-data"}, None), |
| 26 | ] |
| 27 | |
| 28 | |
| 29 | @pytest.mark.asyncio |
| 30 | async def test_create_session_badrequest(basic_server_with_lora: RemoteOpenAIServer): |
| 31 | bad_response = requests.post( |
| 32 | basic_server_with_lora.url_for("invocations"), |
| 33 | json={"requestType": "NEW_SESSION", "extra-field": "extra-field-data"}, |
| 34 | ) |
| 35 | |
| 36 | assert bad_response.status_code == 400 |
| 37 | |
| 38 | |
| 39 | @pytest.mark.asyncio |
| 40 | @pytest.mark.parametrize( |
| 41 | "test_name,session_id_change,request_body_change,expected_error", |
| 42 | CLOSE_BADREQUEST_CASES, |
| 43 | ) |
| 44 | async def test_close_session_badrequest( |
| 45 | basic_server_with_lora: RemoteOpenAIServer, |
| 46 | test_name: str, |
| 47 | session_id_change: dict[str, str], |
| 48 | request_body_change: dict[str, str], |
| 49 | expected_error: str | None, |
| 50 | ): |
| 51 | # first attempt to create a session |
| 52 | url = basic_server_with_lora.url_for("invocations") |
| 53 | create_response = requests.post(url, json={"requestType": "NEW_SESSION"}) |
| 54 | create_response.raise_for_status() |
| 55 | valid_session_id, expiration = create_response.headers.get( |
| 56 | HEADER_SAGEMAKER_NEW_SESSION_ID, "" |
| 57 | ).split(";") |
| 58 | assert valid_session_id |
| 59 | |
| 60 | close_request_json = {"requestType": "CLOSE"} |
| 61 | if request_body_change: |
| 62 | close_request_json.update(request_body_change) |
| 63 | bad_session_id = session_id_change.get("session_id") |
| 64 | bad_close_response = requests.post( |
| 65 | url, |
| 66 | headers={HEADER_SAGEMAKER_SESSION_ID: bad_session_id or valid_session_id}, |
| 67 | json=close_request_json, |
| 68 | ) |
| 69 | |
| 70 | # clean up created session, should succeed |
| 71 | clean_up_response = requests.post( |
| 72 | url, |
| 73 | headers={HEADER_SAGEMAKER_SESSION_ID: valid_session_id}, |
| 74 | json={"requestType": "CLOSE"}, |
| 75 | ) |
| 76 | clean_up_response.raise_for_status() |
| 77 | |
| 78 | assert bad_close_response.status_code == 400 |
| 79 | if expected_error: |
| 80 | assert expected_error in bad_close_response.json()["error"]["message"] |
| 81 | |
| 82 | |
| 83 | @pytest.mark.asyncio |
| 84 | async def test_close_session_invalidrequest( |
| 85 | basic_server_with_lora: RemoteOpenAIServer, async_client: openai.AsyncOpenAI |
| 86 | ): |
| 87 | # first attempt to create a session |
| 88 | url = basic_server_with_lora.url_for("invocations") |
| 89 | create_response = requests.post(url, json={"requestType": "NEW_SESSION"}) |
| 90 | create_response.raise_for_status() |
| 91 | valid_session_id, expiration = create_response.headers.get( |
| 92 | HEADER_SAGEMAKER_NEW_SESSION_ID, "" |
| 93 | ).split(";") |
| 94 | assert valid_session_id |
| 95 | |
| 96 | close_request_json = {"requestType": "CLOSE"} |
| 97 | invalid_close_response = requests.post( |
| 98 | url, |
| 99 | # no headers to specify session_id |
| 100 | json=close_request_json, |
| 101 | ) |
| 102 | |
| 103 | # clean up created session, should succeed |
| 104 | clean_up_response = requests.post( |
| 105 | url, |
| 106 | headers={HEADER_SAGEMAKER_SESSION_ID: valid_session_id}, |
| 107 | json={"requestType": "CLOSE"}, |
| 108 | ) |
| 109 | clean_up_response.raise_for_status() |
| 110 | |
| 111 | assert invalid_close_response.status_code == 424 |
| 112 | assert "invalid session_id" in invalid_close_response.json()["error"]["message"] |
| 113 | |
| 114 | |
| 115 | @pytest.mark.asyncio |
| 116 | async def test_session(basic_server_with_lora: RemoteOpenAIServer): |
| 117 | # first attempt to create a session |
| 118 | url = basic_server_with_lora.url_for("invocations") |
| 119 | create_response = requests.post(url, json={"requestType": "NEW_SESSION"}) |
| 120 | create_response.raise_for_status() |
| 121 | valid_session_id, expiration = create_response.headers.get( |
| 122 | HEADER_SAGEMAKER_NEW_SESSION_ID, "" |
| 123 | ).split(";") |
| 124 | assert valid_session_id |
| 125 | |
| 126 | # test invocation with session id |
| 127 | |
| 128 | request_args = { |
| 129 | "model": MODEL_NAME_SMOLLM, |
| 130 | "prompt": "what is 1+1?", |
| 131 | "max_completion_tokens": 5, |
| 132 | "temperature": 0.0, |
| 133 | "logprobs": False, |
| 134 | } |
| 135 | |
| 136 | invocation_response = requests.post( |
| 137 | basic_server_with_lora.url_for("invocations"), |
| 138 | headers={HEADER_SAGEMAKER_SESSION_ID: valid_session_id}, |
| 139 | json=request_args, |
| 140 | ) |
| 141 | invocation_response.raise_for_status() |
| 142 | |
| 143 | # close created session, should succeed |
| 144 | close_response = requests.post( |
| 145 | url, |
| 146 | headers={HEADER_SAGEMAKER_SESSION_ID: valid_session_id}, |
| 147 | json={"requestType": "CLOSE"}, |
| 148 | ) |
| 149 | close_response.raise_for_status() |
| 150 | |
| 151 | assert ( |
| 152 | close_response.headers.get(HEADER_SAGEMAKER_CLOSED_SESSION_ID) |
| 153 | == valid_session_id |
| 154 | ) |