From 0c30909fe9b1dcbd263bc3132f1c2886aab3642a Mon Sep 17 00:00:00 2001 From: Minwoo Lee <11580164+minwhoo@users.noreply.github.com> Date: Sat, 8 Feb 2025 12:31:01 +0900 Subject: [PATCH 1/4] Reimplement methods required for triton streaming --- .../llms/triton/completion/transformation.py | 25 ++++++++++++++++++- 1 file changed, 24 insertions(+), 1 deletion(-) diff --git a/litellm/llms/triton/completion/transformation.py b/litellm/llms/triton/completion/transformation.py index 0cd6940063..9b100ff1f8 100644 --- a/litellm/llms/triton/completion/transformation.py +++ b/litellm/llms/triton/completion/transformation.py @@ -3,7 +3,7 @@ Translates from OpenAI's `/v1/chat/completions` endpoint to Triton's `/generate` """ import json -from typing import Any, Dict, List, Literal, Optional, Union +from typing import Any, AsyncIterator, Dict, Iterator, List, Literal, Optional, Union from httpx import Headers, Response @@ -52,6 +52,17 @@ class TritonConfig(BaseConfig): ) -> Dict: return {"Content-Type": "application/json"} + def get_complete_url( + self, + api_base: str, + model: str, + optional_params: dict, + stream: Optional[bool] = None, + ) -> str: + if stream: + return api_base + "_stream" + return api_base + def get_supported_openai_params(self, model: str) -> List: return ["max_tokens", "max_completion_tokens"] @@ -149,6 +160,18 @@ class TritonConfig(BaseConfig): else: raise ValueError(f"Invalid Triton API base: {api_base}") + def get_model_response_iterator( + self, + streaming_response: Union[Iterator[str], AsyncIterator[str], ModelResponse], + sync_stream: bool, + json_mode: Optional[bool] = False, + ) -> Any: + return TritonResponseIterator( + streaming_response=streaming_response, + sync_stream=sync_stream, + json_mode=json_mode, + ) + class TritonGenerateConfig(TritonConfig): """ From 268702722504ec2f4bf8f72f7bb15cb6da6843b3 Mon Sep 17 00:00:00 2001 From: Minwoo Lee <11580164+minwhoo@users.noreply.github.com> Date: Thu, 13 Feb 2025 15:40:56 +0900 Subject: [PATCH 2/4] Apply streaming-related transformations only for generate config --- .../llms/triton/completion/transformation.py | 24 +++++++++---------- 1 file changed, 12 insertions(+), 12 deletions(-) diff --git a/litellm/llms/triton/completion/transformation.py b/litellm/llms/triton/completion/transformation.py index 9b100ff1f8..b09f7b0444 100644 --- a/litellm/llms/triton/completion/transformation.py +++ b/litellm/llms/triton/completion/transformation.py @@ -52,17 +52,6 @@ class TritonConfig(BaseConfig): ) -> Dict: return {"Content-Type": "application/json"} - def get_complete_url( - self, - api_base: str, - model: str, - optional_params: dict, - stream: Optional[bool] = None, - ) -> str: - if stream: - return api_base + "_stream" - return api_base - def get_supported_openai_params(self, model: str) -> List: return ["max_tokens", "max_completion_tokens"] @@ -178,6 +167,17 @@ class TritonGenerateConfig(TritonConfig): Transformations for triton /generate endpoint (This is a trtllm model) """ + def get_complete_url( + self, + api_base: str, + model: str, + optional_params: dict, + stream: Optional[bool] = None, + ) -> str: + if stream: + return api_base + "_stream" + return api_base + def transform_request( self, model: str, @@ -227,7 +227,7 @@ class TritonGenerateConfig(TritonConfig): return model_response -class TritonInferConfig(TritonGenerateConfig): +class TritonInferConfig(TritonConfig): """ Transformations for triton /infer endpoint (his is an infer model with a custom model on triton) """ From c1f2ae97c5e3573cbfff173c05cf88ad3b35a249 Mon Sep 17 00:00:00 2001 From: Minwoo Lee <11580164+minwhoo@users.noreply.github.com> Date: Thu, 13 Feb 2025 15:43:42 +0900 Subject: [PATCH 3/4] Add streaming test --- tests/llm_translation/test_triton.py | 40 +++++++++++++++++++++------- 1 file changed, 31 insertions(+), 9 deletions(-) diff --git a/tests/llm_translation/test_triton.py b/tests/llm_translation/test_triton.py index 0835d09fab..7e4ba92f23 100644 --- a/tests/llm_translation/test_triton.py +++ b/tests/llm_translation/test_triton.py @@ -49,16 +49,26 @@ def test_split_embedding_by_shape_fails_with_shape_value_error(): ) -def test_completion_triton_generate_api(): +@pytest.mark.parametrize("stream", [True, False]) +def test_completion_triton_generate_api(stream): try: mock_response = MagicMock() + if stream: + def mock_iter_lines(): + mock_output = ''.join([ + 'data: {"model_name":"ensemble","model_version":"1","sequence_end":false,"sequence_id":0,"sequence_start":false,"text_output":"' + t + '"}\n\n' + for t in ["I", " am", " an", " AI", " assistant"] + ]) + for out in mock_output.split('\n'): + yield out + mock_response.iter_lines = mock_iter_lines + else: + def return_val(): + return { + "text_output": "I am an AI assistant", + } - def return_val(): - return { - "text_output": "I am an AI assistant", - } - - mock_response.json = return_val + mock_response.json = return_val mock_response.status_code = 200 with patch( @@ -71,6 +81,7 @@ def test_completion_triton_generate_api(): max_tokens=10, timeout=5, api_base="http://localhost:8000/generate", + stream=stream, ) # Verify the call was made @@ -81,7 +92,10 @@ def test_completion_triton_generate_api(): call_kwargs = mock_post.call_args.kwargs # Access kwargs directly # Verify URL - assert call_kwargs["url"] == "http://localhost:8000/generate" + if stream: + assert call_kwargs["url"] == "http://localhost:8000/generate_stream" + else: + assert call_kwargs["url"] == "http://localhost:8000/generate" # Parse the request data from the JSON string request_data = json.loads(call_kwargs["data"]) @@ -91,7 +105,15 @@ def test_completion_triton_generate_api(): assert request_data["parameters"]["max_tokens"] == 10 # Verify response - assert response.choices[0].message.content == "I am an AI assistant" + if stream: + tokens = ["I", " am", " an", " AI", " assistant", None] + idx = 0 + for chunk in response: + assert chunk.choices[0].delta.content == tokens[idx] + idx += 1 + assert idx == len(tokens) + else: + assert response.choices[0].message.content == "I am an AI assistant" except Exception as e: print("exception", e) From c62be184c2e9228ad321384f1c385be1ff4f882b Mon Sep 17 00:00:00 2001 From: Minwoo Lee <11580164+minwhoo@users.noreply.github.com> Date: Thu, 13 Feb 2025 16:41:50 +0900 Subject: [PATCH 4/4] Fix get_complete_url --- .../llms/triton/completion/transformation.py | 23 ++++++++++--------- 1 file changed, 12 insertions(+), 11 deletions(-) diff --git a/litellm/llms/triton/completion/transformation.py b/litellm/llms/triton/completion/transformation.py index b09f7b0444..0a65e216df 100644 --- a/litellm/llms/triton/completion/transformation.py +++ b/litellm/llms/triton/completion/transformation.py @@ -67,6 +67,18 @@ class TritonConfig(BaseConfig): optional_params[param] = value return optional_params + def get_complete_url( + self, + api_base: str, + model: str, + optional_params: dict, + stream: Optional[bool] = None, + ) -> str: + llm_type = self._get_triton_llm_type(api_base) + if llm_type == "generate" and stream: + return api_base + "_stream" + return api_base + def transform_response( self, model: str, @@ -167,17 +179,6 @@ class TritonGenerateConfig(TritonConfig): Transformations for triton /generate endpoint (This is a trtllm model) """ - def get_complete_url( - self, - api_base: str, - model: str, - optional_params: dict, - stream: Optional[bool] = None, - ) -> str: - if stream: - return api_base + "_stream" - return api_base - def transform_request( self, model: str,