diff --git a/.codespellrc b/.codespellrc index 4f60f6084..0edc8e73d 100644 --- a/.codespellrc +++ b/.codespellrc @@ -1,3 +1,3 @@ [codespell] -ignore-words-list = ans, als, hel, boostrap, childs, te, vas, hsa, ment, cann, thi, makro, wil, rouge, PRIS, ather, MIS, medias, allready, inout, nd, fo, visibles, nothink, renderD, ond, tbe, CopyIn, notin, subtile, subtiles, dout, IST +ignore-words-list = ans, als, hel, boostrap, childs, te, vas, hsa, ment, cann, thi, makro, wil, rouge, PRIS, ather, MIS, medias, allready, inout, nd, fo, visibles, nothink, renderD, ond, tbe, CopyIn, notin, subtile, subtiles, dout, IST, kInf skip = *.json, *.jsonl, *.patch, *.txt, *.lock diff --git a/python/sglang/srt/parser/conversation.py b/python/sglang/srt/parser/conversation.py index e947737ab..920ac33e3 100644 --- a/python/sglang/srt/parser/conversation.py +++ b/python/sglang/srt/parser/conversation.py @@ -358,8 +358,6 @@ class Conversation: ret = system_prompt + self.sep for role, message in self.messages: if message: - if type(message) is tuple: - message, _, _ = message ret += role + message + self.sep else: ret += role diff --git a/test/registered/unit/entrypoints/anthropic/test_serving.py b/test/registered/unit/entrypoints/anthropic/test_serving.py index a6ad8e358..c1fa902de 100644 --- a/test/registered/unit/entrypoints/anthropic/test_serving.py +++ b/test/registered/unit/entrypoints/anthropic/test_serving.py @@ -134,16 +134,6 @@ async def _collect_anthropic_events(serving, anthropic_request): class TestAnthropicServing(unittest.TestCase): - # System-first guard (Qwen-style): rejects non-first system → must merge. - QWEN_SYSTEM_FIRST_TEMPLATE = ( - "{%- for message in messages %}" - "{%- if message.role == 'system' and not loop.first %}" - "{{- raise_exception('system must be first') }}" - "{%- endif %}" - "{{- message.role }}: {{ message.content }}\n" - "{%- endfor %}" - ) - # Renders system at any position (GLM/Kimi/Qwen3) → can pass through. INLINE_SYSTEM_TEMPLATE = ( "{%- for message in messages %}" @@ -675,19 +665,6 @@ class TestAnthropicServing(unittest.TestCase): self.assertEqual(anthropic_response.content[1].type, "text") self.assertEqual(anthropic_response.content[1].text, "the answer is 4") - def test_request_thinking_enabled_invokes_apply_reasoning_enabled(self): - """``thinking={"type":"enabled", "budget_tokens":N}`` flips reasoning on. - - ``budget_tokens`` is required by the SDK shape on ``enabled``; the - local backend does not enforce it but accepts the value. - """ - serving = self._serving() - request = self._anthropic_request( - thinking={"type": "enabled", "budget_tokens": 1024}, stream=False - ) - serving._convert_to_chat_completion_request(request) - self.assertEqual(serving.openai_serving_chat.apply_reasoning_calls, [True]) - def test_request_thinking_disabled_invokes_apply_reasoning_enabled(self): """``thinking={"type": "disabled"}`` must flip the reasoning toggle off.""" serving = self._serving() @@ -828,34 +805,17 @@ class TestAnthropicServing(unittest.TestCase): self.assertEqual(chat_request.max_tokens, 16) self.assertTrue(any("task_budget" in r and "32768" in r for r in log.output)) - def test_request_task_budget_with_remaining_is_accepted(self): - """SDK's ``BetaTokenTaskBudgetParam`` has a ``remaining`` field - used for client-side compaction. Must round-trip cleanly.""" - serving = self._serving() - request = self._anthropic_request( - output_config={ - "task_budget": {"type": "tokens", "total": 32768, "remaining": 12000} - }, - stream=False, - ) - # Must not raise; pre-existing logging still works. - serving._convert_to_chat_completion_request(request) - self.assertEqual(request.output_config.task_budget.remaining, 12000) - def test_request_betas_is_accepted_and_logged(self): - """The Anthropic SDK attaches ``betas`` to many requests; must not 400.""" + """``betas`` is accepted and logged; the local backend has no beta system.""" import logging serving = self._serving() - request = self._anthropic_request( - betas=["thinking-2025-08-04", "computer-use-2025-01-24"], - stream=False, - ) + request = self._anthropic_request(betas=["thinking-2025-08-04"], stream=False) with self.assertLogs( "sglang.srt.entrypoints.anthropic.serving", level=logging.INFO ) as log: serving._convert_to_chat_completion_request(request) - self.assertTrue(any("betas" in r for r in log.output)) + self.assertTrue(any("thinking-2025-08-04" in r for r in log.output)) def test_assistant_thinking_history_is_rewrapped_for_chat_template(self): """Past-turn thinking blocks get re-emitted via wrap_reasoning_history.""" @@ -1152,18 +1112,6 @@ class TestAnthropicServing(unittest.TestCase): serving._convert_to_chat_completion_request(request) self.assertIn("tool_choice", str(ctx.exception)) - def test_server_tool_only_with_tool_choice_auto_is_allowed(self): - """tool_choice=auto over server-only tools is a no-op (model decides).""" - serving = self._serving() - request = self._anthropic_request( - stream=False, - tools=[{"type": "web_search_20250305", "name": "web_search"}], - tool_choice={"type": "auto"}, - ) - # Must not raise; the request just runs with no client-side tools. - chat_request = serving._convert_to_chat_completion_request(request) - self.assertIsNone(chat_request.tools) - def test_tool_choice_named_custom_tool_is_resolved(self): """tool_choice={type:'tool', name:'X'} where X is a custom tool wires through.""" serving = self._serving() diff --git a/test/registered/unit/entrypoints/openai/test_serving_chat.py b/test/registered/unit/entrypoints/openai/test_serving_chat.py index 8dd0cc22c..16b38857d 100644 --- a/test/registered/unit/entrypoints/openai/test_serving_chat.py +++ b/test/registered/unit/entrypoints/openai/test_serving_chat.py @@ -240,49 +240,6 @@ class ServingChatTestCase(unittest.TestCase): self.assertTrue(adapted.require_reasoning) - def test_kimi_tool_call_keeps_explicit_reasoning(self): - self.template_manager.reasoning_config = ReasoningToggleConfig( - toggle_param="thinking", default_enabled=True - ) - self.tm.server_args.reasoning_parser = "kimi_k2" - self.tm.server_args.tool_call_parser = "kimi_k2" - self.chat.reasoning_parser = "kimi_k2" - self.chat.tool_call_parser = "kimi_k2" - - req = ChatCompletionRequest( - model="x", - messages=[{"role": "user", "content": "What is 2+2?"}], - tools=[ - { - "type": "function", - "function": { - "name": "add", - "parameters": { - "type": "object", - "properties": {"a": {"type": "integer"}}, - }, - }, - } - ], - tool_choice="required", - chat_template_kwargs={"thinking": True}, - ) - - with patch.object(self.chat, "_process_messages") as proc_mock: - proc_mock.return_value = MessageProcessingResult( - "", - [1, 2, 3], - None, - None, - [], - [], - None, - ) - - adapted, _ = self.chat._convert_to_internal_request(req) - - self.assertTrue(adapted.require_reasoning) - def test_kimi_tool_call_respects_explicit_reasoning_disable(self): self.template_manager.reasoning_config = ReasoningToggleConfig( toggle_param="thinking", default_enabled=True @@ -895,98 +852,6 @@ class ServingChatTestCase(unittest.TestCase): ) # ------------- kimi_k2 tool_call_id formatting ------------- - def test_kimi_k2_non_streaming_tool_call_id_format(self): - """Ensure non-streaming tool_call.id matches functions.{name}:{index} for kimi_k2 parser.""" - - # Force kimi_k2 parser - self.chat.tool_call_parser = "kimi_k2" - - # Mock FunctionCallParser.parse_non_stream to return one tool call - with patch( - "sglang.srt.entrypoints.openai.serving_chat.FunctionCallParser" - ) as ParserMock: - parser_instance = ParserMock.return_value - - # Build a mock ToolCallItem-like object - call_info = Mock() - call_info.name = "get_weather" - call_info.parameters = '{"city":"Paris"}' - call_info.tool_index = 0 - - parser_instance.has_tool_call.return_value = True - parser_instance.parse_non_stream.return_value = ("", [call_info]) - - finish_reason = {"type": "stop", "matched": None} - tools = [ - {"type": "function", "function": {"name": "get_weather"}}, - ] - - tool_calls, remaining_text, finish_reason = self.chat._process_tool_calls( - text="<|tool_calls_section_begin|>...", - tools=tools, - finish_reason=finish_reason, - ) - - self.assertIsNotNone(tool_calls) - self.assertEqual(len(tool_calls), 1) - self.assertEqual(tool_calls[0].id, "functions.get_weather:0") - self.assertEqual(tool_calls[0].function.name, "get_weather") - - def test_kimi_k2_streaming_tool_call_id_format(self): - """Ensure streaming first chunk tool_call.id matches functions.{name}:{index} for kimi_k2 parser.""" - - # Force kimi_k2 parser - self.chat.tool_call_parser = "kimi_k2" - - # Prepare request with tools - req = ChatCompletionRequest( - model="x", - messages=[{"role": "user", "content": "Hi?"}], - tools=[{"type": "function", "function": {"name": "get_weather"}}], - stream=True, - ) - - # Patch FunctionCallParser used inside _process_tool_call_stream - with patch( - "sglang.srt.entrypoints.openai.serving_chat.FunctionCallParser" - ) as ParserMock: - parser_instance = ParserMock.return_value - - # First call returns one ToolCallItem-like chunk (with name) - first_chunk_call = Mock() - first_chunk_call.tool_index = 0 - first_chunk_call.name = "get_weather" - first_chunk_call.parameters = "" - parser_instance.parse_stream_chunk.side_effect = [ - ("", [first_chunk_call]), - ("", []), - ] - - async def collect_first_tool_chunk(): - gen = self.chat._process_tool_call_stream( - index=0, - delta="irrelevant", - parser_dict={}, - content={"meta_info": {"id": "chatcmpl-test"}}, - request=req, - has_tool_calls={}, - ) - # Get first yielded SSE line - line = None - async for emitted in gen: - line = emitted - break - return line - - loop = get_or_create_event_loop() - line = loop.run_until_complete(collect_first_tool_chunk()) - self.assertIsNotNone(line) - self.assertTrue(line.startswith("data: ")) - - payload = json.loads(line[len("data: ") :]) - tool_calls = payload["choices"][0]["delta"]["tool_calls"] - self.assertEqual(tool_calls[0]["id"], "functions.get_weather:0") - def test_kimi_k2_non_streaming_tool_call_id_with_history(self): """Ensure non-streaming tool_call.id increase with tool calls history for kimi_k2 parser.""" @@ -1225,28 +1090,6 @@ class ServingChatTestCase(unittest.TestCase): task="bogus", ) - def test_latest_reminder_role_accepted(self): - """`latest_reminder` is a first-class message role on generic param.""" - from sglang.srt.entrypoints.openai.protocol import ( - ChatCompletionMessageGenericParam, - ) - - msg = ChatCompletionMessageGenericParam( - role="latest_reminder", content="Be terse." - ) - self.assertEqual(msg.role, "latest_reminder") - - # Full request with reminder before user parses cleanly. - req = ChatCompletionRequest( - model="x", - messages=[ - {"role": "latest_reminder", "content": "Be terse."}, - {"role": "user", "content": "Hi"}, - ], - ) - self.assertEqual(req.messages[0].role, "latest_reminder") - self.assertEqual(req.messages[1].role, "user") - def test_attach_task_to_last_user_message(self): """Helper attaches task to the nearest user/developer message.""" from sglang.srt.entrypoints.openai import encoding_dsv4 @@ -1954,14 +1797,6 @@ class ServingChatTestCase(unittest.TestCase): ) self.assertIsNone(result) - def test_extract_routed_dp_rank_from_header_with_header(self): - """Test that header value is extracted correctly.""" - self.fastapi_request.headers = {"x-data-parallel-rank": "2"} - result = self.chat.extract_routed_dp_rank_from_header( - self.fastapi_request, body_routed_dp_rank=None - ) - self.assertEqual(result, 2) - def test_extract_routed_dp_rank_header_overrides_body(self): """Test that header value has higher priority than body.""" self.fastapi_request.headers = {"x-data-parallel-rank": "3"} @@ -2404,10 +2239,6 @@ class TestProcessToolCallsWithRequiredToolChoice(unittest.TestCase): class TestNormalizeToolContent(unittest.TestCase): """Unit tests for normalize_tool_content().""" - def test_openai_text_parts_flattened(self): - result = normalize_tool_content("tool", [{"type": "text", "text": "10525"}]) - self.assertEqual(result, "10525") - def test_multiple_text_parts_joined(self): result = normalize_tool_content( "tool", diff --git a/test/registered/unit/function_call/test_function_call_parser.py b/test/registered/unit/function_call/test_function_call_parser.py index d9edb5a2b..96e027481 100644 --- a/test/registered/unit/function_call/test_function_call_parser.py +++ b/test/registered/unit/function_call/test_function_call_parser.py @@ -79,6 +79,12 @@ class TestPythonicDetector(unittest.TestCase): ] self.detector = PythonicDetector() + def test_has_tool_call_detects_marker(self): + # Guards the pythonic ``[func(kw=v)]`` regex predicate, which + # detect_and_parse does not exercise in isolation. + self.assertTrue(self.detector.has_tool_call('[get_weather(location="Tokyo")]')) + self.assertFalse(self.detector.has_tool_call("plain text only")) + def test_parse_streaming_no_brackets(self): """Test parsing text with no brackets (no tool calls).""" text = "This is just normal text without any tool calls." @@ -105,19 +111,6 @@ class TestPythonicDetector(unittest.TestCase): self.assertEqual(params["location"], "New York") self.assertEqual(params["unit"], "celsius") - def test_parse_streaming_text_before_tool_call(self): - """Test parsing text that appears before a tool call.""" - text = "This is some text before [get_weather(location='London')]" - result = self.detector.parse_streaming_increment(text, self.tools) - - self.assertEqual(result.normal_text, "This is some text before ") - self.assertEqual(len(result.calls), 1) - self.assertEqual(result.calls[0].name, "get_weather") - - # Check the parameters - params = json.loads(result.calls[0].parameters) - self.assertEqual(params["location"], "London") - def test_parse_streaming_partial_tool_call(self): """Test parsing a partial tool call that spans multiple chunks.""" # First chunk with opening bracket but no closing bracket @@ -430,20 +423,6 @@ class TestMistralDetector(unittest.TestCase): result.normal_text, "", "Normal text should be empty for pure tool call" ) - def test_detect_and_parse_simple_case(self): - """Test parsing a simple Mistral format tool call without nested brackets.""" - test_text = '[TOOL_CALLS] [{"name":"make_next_step_decision", "arguments":{"decision":"TOOL", "content":"Use weather API"}}]' - - result = self.detector.detect_and_parse(test_text, self.tools) - - self.assertEqual(len(result.calls), 1) - call = result.calls[0] - self.assertEqual(call.name, "make_next_step_decision") - - params = json.loads(call.parameters) - self.assertEqual(params["decision"], "TOOL") - self.assertEqual(params["content"], "Use weather API") - def test_detect_and_parse_no_tool_calls(self): """Test parsing text without any tool calls.""" test_text = "This is just normal text without any tool calls." @@ -680,43 +659,6 @@ class TestBaseFormatDetector(unittest.TestCase): tourist_calls[0].tool_index, 1, "Second tool should have tool_index=1" ) - def test_tool_name_streaming_with_correct_index(self): - """Test that tool names are streamed with correct tool_index values.""" - # Process first tool - self.detector.parse_streaming_increment("", self.tools) - result1 = self.detector.parse_streaming_increment( - '{"name": "get_weather", ', self.tools - ) - - # First tool name should have tool_index=0 - weather_calls = [call for call in result1.calls if call.name == "get_weather"] - self.assertEqual(len(weather_calls), 1, "Should have one weather call") - self.assertEqual( - weather_calls[0].tool_index, 0, "First tool should have tool_index=0" - ) - - # Complete first tool - self.detector.parse_streaming_increment( - '"arguments": {"city": "Paris"}}', self.tools - ) - - # Start second tool - self.detector.parse_streaming_increment(", ", self.tools) - result2 = self.detector.parse_streaming_increment( - '{"name": "get_tourist_attractions", ', self.tools - ) - - # Second tool name should have tool_index=1 - tourist_calls = [ - call for call in result2.calls if call.name == "get_tourist_attractions" - ] - self.assertEqual( - len(tourist_calls), 1, "Should have one tourist attractions call" - ) - self.assertEqual( - tourist_calls[0].tool_index, 1, "Second tool should have tool_index=1" - ) - def test_buffer_reset_on_invalid_tool(self): """Test that buffer and state are reset when an invalid tool name is encountered.""" # Start fresh with an invalid tool name from the beginning @@ -819,56 +761,6 @@ class TestBaseFormatDetector(unittest.TestCase): params["city"], "杭州", "Should correctly parse Chinese city name" ) - def test_multiple_chinese_parameters(self): - """Test multiple tool calls with Chinese parameters.""" - # Test with multiple tool calls containing Chinese characters - chunks = [ - "", - '{"name": "get_weather", "arguments": {"city": "北京"}}, ', - '{"name": "get_tourist_attractions", "arguments": {"city": "上海"}}', - "", - ] - - accumulated_parameters = {} - for chunk in chunks: - result = self.detector.parse_streaming_increment(chunk, self.tools) - if result.calls: - for call in result.calls: - if call.parameters: - tool_idx = call.tool_index if call.tool_index is not None else 0 - if tool_idx not in accumulated_parameters: - accumulated_parameters[tool_idx] = "" - accumulated_parameters[tool_idx] += call.parameters - - # Verify both tool calls have correct Chinese characters - self.assertGreaterEqual( - len(accumulated_parameters), 1, "Should have parsed parameters" - ) - - # Check first tool call (北京 - Beijing) - if 0 in accumulated_parameters: - params0 = json.loads(accumulated_parameters[0]) - self.assertIn( - "北京", - accumulated_parameters[0], - "Should contain actual Chinese characters", - ) - self.assertEqual( - params0["city"], "北京", "Should correctly parse first Chinese city" - ) - - # Check second tool call (上海 - Shanghai) if present - if 1 in accumulated_parameters: - params1 = json.loads(accumulated_parameters[1]) - self.assertIn( - "上海", - accumulated_parameters[1], - "Should contain actual Chinese characters", - ) - self.assertEqual( - params1["city"], "上海", "Should correctly parse second Chinese city" - ) - class TestLlama32Detector(unittest.TestCase): def setUp(self): @@ -1008,15 +900,6 @@ class TestKimiK2Detector(unittest.TestCase): ] self.detector = KimiK2Detector() - def test_single_tool_call(self): - """Test parsing a single tool call in a complete text.""" - text = '<|tool_calls_section_begin|><|tool_call_begin|>functions.get_weather:0<|tool_call_argument_begin|>{"city": "Paris"}<|tool_call_end|><|tool_calls_section_end|>' - result = self.detector.detect_and_parse(text, self.tools) - self.assertEqual(len(result.calls), 1) - self.assertEqual(result.calls[0].name, "get_weather") - self.assertEqual(result.calls[0].parameters, '{"city": "Paris"}') - self.assertEqual(result.normal_text, "") - def test_multiple_tool_calls(self): """Test parsing multiple tool calls in a complete text.""" text = '<|tool_calls_section_begin|><|tool_call_begin|>functions.get_weather:0<|tool_call_argument_begin|>{"city": "Paris"}<|tool_call_end|><|tool_call_begin|>functions.get_tourist_attractions:1<|tool_call_argument_begin|>{"city": "London"}<|tool_call_end|><|tool_calls_section_end|>' @@ -1028,35 +911,6 @@ class TestKimiK2Detector(unittest.TestCase): self.assertEqual(result.calls[1].parameters, '{"city": "London"}') self.assertEqual(result.normal_text, "") - def test_streaming_tool_call(self): - """Test streaming incremental parsing of a tool call.""" - chunks = [ - "<|tool_calls_section_begin|><|tool_call_begin|>functions.get_weather:0<|tool_call_argument_begin|>{", - '"city": "Paris"', - "}", - "<|tool_call_end|><|tool_calls_section_end|>", - ] - - tool_calls = [] - for chunk in chunks: - result = self.detector.parse_streaming_increment(chunk, self.tools) - for tool_call_chunk in result.calls: - if tool_call_chunk.tool_index is not None: - - while len(tool_calls) <= tool_call_chunk.tool_index: - tool_calls.append({"name": "", "parameters": ""}) - - tc = tool_calls[tool_call_chunk.tool_index] - - if tool_call_chunk.name: - tc["name"] += tool_call_chunk.name - if tool_call_chunk.parameters: - tc["parameters"] += tool_call_chunk.parameters - - self.assertEqual(len(tool_calls), 1) - self.assertEqual(tool_calls[0]["name"], "get_weather") - self.assertEqual(tool_calls[0]["parameters"], '{"city": "Paris"}') - def test_streaming_multiple_tool_calls(self): """Test streaming incremental parsing of multiple tool calls.""" chunks = [ @@ -1797,44 +1651,6 @@ class TestDeepSeekV4Detector(unittest.TestCase): self.assertEqual(params2["topn"], 10) self.assertEqual(params2["source"], "web") - def test_detect_and_parse_json_format(self): - """Test parsing JSON format inside invoke tags""" - text = """I'll help you with information about San Francisco and get its favorite tourist spot for you. - - <｜DSML｜tool_calls> - <｜DSML｜invoke name="get_favorite_tourist_spot"> - { - "city": "San Francisco" - } - - <｜DSML｜invoke name="search"> - { - "query": "WebNav benchmark", - "topn": 10, - "source": "web" - } - - - """ - result = self.detector.detect_and_parse(text, self.tools) - - self.assertIn("I'll help you with information", result.normal_text) - self.assertEqual(len(result.calls), 2) - - # Check first call - call1 = result.calls[0] - self.assertEqual(call1.name, "get_favorite_tourist_spot") - params1 = json.loads(call1.parameters) - self.assertEqual(params1["city"], "San Francisco") - - # Check second call - call2 = result.calls[1] - self.assertEqual(call2.name, "search") - params2 = json.loads(call2.parameters) - self.assertEqual(params2["query"], "WebNav benchmark") - self.assertEqual(params2["topn"], 10) - self.assertEqual(params2["source"], "web") - def test_streaming_xml_format(self): """Test streaming parsing of XML format""" text = """<｜DSML｜tool_calls> @@ -1885,211 +1701,6 @@ class TestDeepSeekV4Detector(unittest.TestCase): self.assertEqual(params["obj"]["name"], "John") self.assertEqual(params["obj"]["age"], 30) - def test_streaming_json_format(self): - """Test streaming parsing of JSON format""" - text = """<｜DSML｜tool_calls> - <｜DSML｜invoke name="get_favorite_tourist_spot"> - { - "city": "San Francisco", - "another_city": "London", - "topn": 10, - "obj": { - "name": "John", - "age": 30 - } - } - - """ - - input_ids = self.tokenizer.encode(text, add_special_tokens=False) - chunk_ids = [ - input_ids[i : i + self.interval] - for i in range(0, len(input_ids), self.interval) - ] - chunks = [self.tokenizer.decode(chunk_id) for chunk_id in chunk_ids] - - tool_calls_by_index = {} - - num_tool_call_chunks = 0 - for chunk in chunks: - result = self.detector.parse_streaming_increment(chunk, self.tools) - for call in result.calls: - num_tool_call_chunks += 1 - if call.tool_index is not None: - if call.tool_index not in tool_calls_by_index: - tool_calls_by_index[call.tool_index] = { - "name": "", - "parameters": "", - } - - if call.name: - tool_calls_by_index[call.tool_index]["name"] = call.name - if call.parameters: - tool_calls_by_index[call.tool_index][ - "parameters" - ] += call.parameters - - self.assertGreater(num_tool_call_chunks, 8) - self.assertEqual(len(tool_calls_by_index), 1) - self.assertEqual(tool_calls_by_index[0]["name"], "get_favorite_tourist_spot") - - # Clean up parameters string if needed (trim whitespace) - params_str = tool_calls_by_index[0]["parameters"].strip() - params = json.loads(params_str) - self.assertEqual(params["city"], "San Francisco") - - def test_detect_and_parse_no_parameters(self): - """Test parsing function calls with no parameters (non-streaming)""" - # Add a no-parameter tool - tools_with_no_param = self.tools + [ - Tool( - type="function", - function=Function( - name="get_date", - description="Get the current date.", - parameters={"type": "object", "properties": {}}, - ), - ), - ] - - text = """Let me get the current date for you. - -<｜DSML｜tool_calls> -<｜DSML｜invoke name="get_date"> - -""" - - result = self.detector.detect_and_parse(text, tools_with_no_param) - - self.assertIn("Let me get the current date", result.normal_text) - self.assertEqual(len(result.calls), 1) - - call = result.calls[0] - self.assertEqual(call.name, "get_date") - params = json.loads(call.parameters) - self.assertEqual(params, {}) - - def test_streaming_no_parameters(self): - """Test streaming parsing of function calls with no parameters. - - This test verifies the fix for the bug where functions with no parameters - were being silently skipped in streaming mode. - """ - # Add a no-parameter tool - tools_with_no_param = self.tools + [ - Tool( - type="function", - function=Function( - name="get_date", - description="Get the current date.", - parameters={"type": "object", "properties": {}}, - ), - ), - ] - - text = """<｜DSML｜tool_calls> -<｜DSML｜invoke name="get_date"> - -""" - - # Reset detector state - self.detector = DeepSeekV4Detector() - - # Simulate streaming by splitting into small chunks - input_ids = self.tokenizer.encode(text, add_special_tokens=False) - chunk_ids = [ - input_ids[i : i + self.interval] - for i in range(0, len(input_ids), self.interval) - ] - chunks = [self.tokenizer.decode(chunk_id) for chunk_id in chunk_ids] - - tool_calls_by_index = {} - - for chunk in chunks: - result = self.detector.parse_streaming_increment(chunk, tools_with_no_param) - for call in result.calls: - if call.tool_index is not None: - if call.tool_index not in tool_calls_by_index: - tool_calls_by_index[call.tool_index] = { - "name": "", - "parameters": "", - } - - if call.name: - tool_calls_by_index[call.tool_index]["name"] = call.name - if call.parameters: - tool_calls_by_index[call.tool_index][ - "parameters" - ] += call.parameters - - # Verify that the no-parameter function was correctly parsed - self.assertEqual( - len(tool_calls_by_index), 1, "Should have exactly one tool call" - ) - self.assertEqual(tool_calls_by_index[0]["name"], "get_date") - - # Parameters should be empty JSON object - params_str = tool_calls_by_index[0]["parameters"].strip() - params = json.loads(params_str) - self.assertEqual(params, {}) - - def test_streaming_no_parameters_with_whitespace(self): - """Test streaming parsing when invoke content has only whitespace (newlines).""" - tools_with_no_param = self.tools + [ - Tool( - type="function", - function=Function( - name="get_date", - description="Get the current date.", - parameters={"type": "object", "properties": {}}, - ), - ), - ] - - # This format has newlines inside the invoke tag (common model output) - text = """<｜DSML｜tool_calls> -<｜DSML｜invoke name="get_date"> - - -""" - - # Reset detector state - self.detector = DeepSeekV4Detector() - - input_ids = self.tokenizer.encode(text, add_special_tokens=False) - chunk_ids = [ - input_ids[i : i + self.interval] - for i in range(0, len(input_ids), self.interval) - ] - chunks = [self.tokenizer.decode(chunk_id) for chunk_id in chunk_ids] - - tool_calls_by_index = {} - - for chunk in chunks: - result = self.detector.parse_streaming_increment(chunk, tools_with_no_param) - for call in result.calls: - if call.tool_index is not None: - if call.tool_index not in tool_calls_by_index: - tool_calls_by_index[call.tool_index] = { - "name": "", - "parameters": "", - } - - if call.name: - tool_calls_by_index[call.tool_index]["name"] = call.name - if call.parameters: - tool_calls_by_index[call.tool_index][ - "parameters" - ] += call.parameters - - # Should still parse correctly even with whitespace-only content - self.assertEqual( - len(tool_calls_by_index), 1, "Should have exactly one tool call" - ) - self.assertEqual(tool_calls_by_index[0]["name"], "get_date") - params = json.loads(tool_calls_by_index[0]["parameters"]) - self.assertEqual(params, {}) - def test_get_model_structural_tag(self): import xgrammar as xgr @@ -2310,30 +1921,6 @@ class TestQwen3CoderDetector(unittest.TestCase): self.assertEqual(result.normal_text, text) self.assertEqual(len(result.calls), 0) - def test_single_tool_call(self): - """ - Test parsing of a single tool call. - - Scenario: Input contains one complete tool call with parameters. - Purpose: Verify correct extraction of tool name and parameters. - """ - text = """ - -Boston -celsius -3 - -""" - result = self.detector.detect_and_parse(text, self.tools) - - self.assertEqual(len(result.calls), 1) - self.assertEqual(result.calls[0].name, "get_current_weather") - - params = json.loads(result.calls[0].parameters) - self.assertEqual(params["location"], "Boston") - self.assertEqual(params["unit"], "celsius") - self.assertEqual(params["days"], 3) - def test_single_tool_call_with_text_prefix(self): """ Test parsing of tool call with preceding text. @@ -2573,18 +2160,6 @@ class TestQwen3CoderDetector(unittest.TestCase): result = self.detector.detect_and_parse(text, self.tools) self.assertIsInstance(result, StreamingParseResult) - def test_has_tool_call_detection(self): - """ - Test the has_tool_call method for detecting tool call markers. - - Scenario: Various inputs with and without tool call markers. - Purpose: Verify correct detection of tool call presence. - """ - self.assertTrue(self.detector.has_tool_call("")) - self.assertTrue(self.detector.has_tool_call("text more")) - self.assertFalse(self.detector.has_tool_call("plain text only")) - self.assertFalse(self.detector.has_tool_call("")) - # ==================== Structural tag (xgrammar builtin) ==================== # Qwen3 Coder uses the new builtin structural tag path. supports_structural_tag() # is True so required/named tool_choice routes through FunctionCallParser @@ -2680,6 +2255,16 @@ class TestGptOssDetector(unittest.TestCase): ] self.detector = GptOssDetector() + def test_has_tool_call_detects_marker(self): + # Guards the bot_token substring predicate, which detect_and_parse + # does not exercise in isolation. + self.assertTrue( + self.detector.has_tool_call( + "<|start|>assistant<|channel|>commentary to=get_weather<|return|>" + ) + ) + self.assertFalse(self.detector.has_tool_call("no tool call here")) + def test_get_model_structural_tag(self): import xgrammar as xgr @@ -2749,21 +2334,6 @@ class TestGlm4MoeDetector(unittest.TestCase): ] self.detector = Glm4MoeDetector() - def test_single_tool_call(self): - text = ( - "get_weather\n" - "city\nBeijing\n" - "date\n2024-06-27\n" - "" - ) - result = self.detector.detect_and_parse(text, self.tools) - self.assertEqual(len(result.calls), 1) - self.assertEqual(result.calls[0].name, "get_weather") - self.assertEqual( - result.calls[0].parameters, '{"city": "Beijing", "date": "2024-06-27"}' - ) - self.assertEqual(result.normal_text, "") - def test_multiple_tool_calls(self): text = ( "get_weather\n" @@ -2787,35 +2357,6 @@ class TestGlm4MoeDetector(unittest.TestCase): ) self.assertEqual(result.normal_text, "") - def test_streaming_tool_call(self): - """Test streaming incremental parsing of a tool call.""" - chunks = [ - "get_weather\n", - "city\nBeijing\n", - "date\n2024-06-27\n", - "", - ] - tool_calls = [] - for chunk in chunks: - result = self.detector.parse_streaming_increment(chunk, self.tools) - for tool_call_chunk in result.calls: - if ( - hasattr(tool_call_chunk, "tool_index") - and tool_call_chunk.tool_index is not None - ): - while len(tool_calls) <= tool_call_chunk.tool_index: - tool_calls.append({"name": "", "parameters": ""}) - tc = tool_calls[tool_call_chunk.tool_index] - if tool_call_chunk.name: - tc["name"] = tool_call_chunk.name - if tool_call_chunk.parameters: - tc["parameters"] += tool_call_chunk.parameters - self.assertEqual(len(tool_calls), 1) - self.assertEqual(tool_calls[0]["name"], "get_weather") - self.assertEqual( - tool_calls[0]["parameters"], '{"city": "Beijing", "date": "2024-06-27"}' - ) - def test_streaming_multiple_tool_calls(self): """Test streaming incremental parsing of multiple tool calls.""" chunks = [ @@ -3073,21 +2614,6 @@ class TestGlm47MoeDetector(unittest.TestCase): ] self.detector = Glm47MoeDetector() - def test_single_tool_call(self): - text = ( - "get_weather" - "cityBeijing" - "date2024-06-27" - "" - ) - result = self.detector.detect_and_parse(text, self.tools) - self.assertEqual(len(result.calls), 1) - self.assertEqual(result.calls[0].name, "get_weather") - self.assertEqual( - result.calls[0].parameters, '{"city": "Beijing", "date": "2024-06-27"}' - ) - self.assertEqual(result.normal_text, "") - def test_multiple_tool_calls(self): text = ( "get_weather" @@ -3111,35 +2637,6 @@ class TestGlm47MoeDetector(unittest.TestCase): ) self.assertEqual(result.normal_text, "") - def test_streaming_tool_call(self): - """Test streaming incremental parsing of a tool call.""" - chunks = [ - "get_weather", - "cityBeijing", - "date2024-06-27", - "", - ] - tool_calls = [] - for chunk in chunks: - result = self.detector.parse_streaming_increment(chunk, self.tools) - for tool_call_chunk in result.calls: - if ( - hasattr(tool_call_chunk, "tool_index") - and tool_call_chunk.tool_index is not None - ): - while len(tool_calls) <= tool_call_chunk.tool_index: - tool_calls.append({"name": "", "parameters": ""}) - tc = tool_calls[tool_call_chunk.tool_index] - if tool_call_chunk.name: - tc["name"] = tool_call_chunk.name - if tool_call_chunk.parameters: - tc["parameters"] += tool_call_chunk.parameters - self.assertEqual(len(tool_calls), 1) - self.assertEqual(tool_calls[0]["name"], "get_weather") - self.assertEqual( - tool_calls[0]["parameters"], '{"city": "Beijing", "date": "2024-06-27"}' - ) - def test_streaming_multiple_tool_calls(self): """Test streaming incremental parsing of multiple tool calls.""" chunks = [ @@ -3468,13 +2965,6 @@ class TestJsonArrayParser(unittest.TestCase): ] self.detector = JsonArrayParser() - def test_json_detector_has_no_ebnf(self): - """JsonArrayParser no longer exposes EBNF generation helpers.""" - self.assertFalse( - hasattr(self.detector, "build_ebnf"), - "JsonArrayParser should not expose EBNF helpers after cleanup", - ) - def test_parse_streaming_increment_malformed_json(self): """Test parsing with malformed JSON""" # Test with malformed JSON @@ -3495,49 +2985,6 @@ class TestJsonArrayParser(unittest.TestCase): self.assertEqual(len(result.calls), 0) self.assertEqual(result.normal_text, "") - def test_parse_streaming_increment_whitespace_handling(self): - """Test parsing with various whitespace scenarios""" - # Test with leading/trailing whitespace split across chunks - chunk1 = ' [{"name": "get_weather", "parameters": ' - result1 = self.detector.parse_streaming_increment(chunk1, self.tools) - self.assertIsInstance(result1, StreamingParseResult) - chunk2 = '{"location": "Tokyo"}}] ' - result2 = self.detector.parse_streaming_increment(chunk2, self.tools) - - # The base class should handle this - self.assertIsInstance(result2, StreamingParseResult) - - def test_parse_streaming_increment_nested_objects(self): - """Test parsing with nested JSON objects""" - chunk1 = '[{"name": "get_weather", "parameters": {"location": "Tokyo", ' - result1 = self.detector.parse_streaming_increment(chunk1, self.tools) - self.assertIsInstance(result1, StreamingParseResult) - chunk2 = '"nested": {"key": "value"}}}]' - result2 = self.detector.parse_streaming_increment(chunk2, self.tools) - - # The base class should handle this - self.assertIsInstance(result2, StreamingParseResult) - - def test_json_parsing_with_commas(self): - """Test that JSON parsing works correctly with comma separators""" - # Stream two complete objects, at least 2 chunks per tool call - chunk1 = '[{"name": "get_weather", "parameters": {"location": "Tok' - result1 = self.detector.parse_streaming_increment(chunk1, self.tools) - self.assertIsInstance(result1, StreamingParseResult) - chunk2 = 'yo"}},' - result2 = self.detector.parse_streaming_increment(chunk2, self.tools) - self.assertIsInstance(result2, StreamingParseResult) - - chunk3 = '{"name": "get_weather", "parameters": {"location": "Par' - result3 = self.detector.parse_streaming_increment(chunk3, self.tools) - self.assertIsInstance(result3, StreamingParseResult) - chunk4 = 'is"}}]' - result4 = self.detector.parse_streaming_increment(chunk4, self.tools) - self.assertIsInstance(result4, StreamingParseResult) - self.assertGreater( - len(result4.calls), 0, "Should parse tool calls from text with separators" - ) - def test_braces_in_strings(self): """Test that JSON with } characters inside strings works correctly""" # Test case: JSON array with } inside string values - streamed across chunks @@ -3582,57 +3029,6 @@ class TestJsonArrayParser(unittest.TestCase): "Should parse tool calls with separator in same chunk", ) - def test_separator_in_separate_chunk(self): - """Test that separator in separate chunk works correctly""" - # Test case: separator in separate chunk - this tests streaming behavior - chunk1 = '[{"name": "get_weather", "parameters": {"location": "Tokyo"}}' - chunk2 = "," - chunk3 = '{"name": "get_weather", "parameters": {"location": "Paris"}}' - - # Process first chunk - result1 = self.detector.parse_streaming_increment(chunk1, self.tools) - self.assertIsInstance(result1, StreamingParseResult) - - # Process separator chunk - result2 = self.detector.parse_streaming_increment(chunk2, self.tools) - self.assertIsInstance(result2, StreamingParseResult) - - # Process second chunk (streaming in progress) - result3 = self.detector.parse_streaming_increment(chunk3, self.tools) - self.assertIsInstance(result3, StreamingParseResult) - - def test_incomplete_json_across_chunks(self): - """Test that incomplete JSON across chunks works correctly""" - # Test case: incomplete JSON across chunks - this tests streaming behavior - chunk1 = '[{"name": "get_weather", "parameters": {"location": "Tokyo"' - chunk2 = '}},{"name": "get_weather"' - - # Process first chunk (incomplete) - result1 = self.detector.parse_streaming_increment(chunk1, self.tools) - self.assertIsInstance(result1, StreamingParseResult) - - # Process second chunk (completes first object and starts second, streaming in progress) - result2 = self.detector.parse_streaming_increment(chunk2, self.tools) - self.assertIsInstance(result2, StreamingParseResult) - - def test_malformed_json_recovery(self): - """Test that malformed JSON recovers gracefully""" - # Test with malformed JSON - should handle gracefully - malformed_text = ( - '[{"name": "get_weather", "parameters": {"location": "unclosed string' - ) - - result1 = self.detector.parse_streaming_increment(malformed_text, self.tools) - self.assertIsInstance(result1, StreamingParseResult) - - # Test valid JSON after malformed - streamed across 2 chunks (streaming in progress) - valid_chunk1 = '[{"name": "get_weather", "parameters": {"location": "Tok' - result2 = self.detector.parse_streaming_increment(valid_chunk1, self.tools) - self.assertIsInstance(result2, StreamingParseResult) - valid_chunk2 = 'yo"}}' - result3 = self.detector.parse_streaming_increment(valid_chunk2, self.tools) - self.assertIsInstance(result3, StreamingParseResult) - def test_nested_objects_with_commas(self): """Test that nested objects with commas inside work correctly""" # Test with nested objects that have commas - should work with json.loads() @@ -3646,69 +3042,6 @@ class TestJsonArrayParser(unittest.TestCase): len(result2.calls), 0, "Should parse tool call with nested objects" ) - def test_empty_objects(self): - """Test that empty objects work correctly""" - # Test with empty objects - should work with json.loads() - chunk1 = '[{"name": "get_weather", "parameters": ' - result1 = self.detector.parse_streaming_increment(chunk1, self.tools) - self.assertIsInstance(result1, StreamingParseResult) - chunk2 = "{}}" - result2 = self.detector.parse_streaming_increment(chunk2, self.tools) - self.assertIsInstance(result2, StreamingParseResult) - - def test_whitespace_handling(self): - """Test that various whitespace scenarios work correctly""" - # Test with various whitespace patterns - should work with json.loads() - chunk1 = ' \n\n [{"name": "get_weather", "parameters": ' - result1 = self.detector.parse_streaming_increment(chunk1, self.tools) - self.assertIsInstance(result1, StreamingParseResult) - chunk2 = '{"location": "Tokyo"}}' - result2 = self.detector.parse_streaming_increment(chunk2, self.tools) - self.assertIsInstance(result2, StreamingParseResult) - - def test_multiple_commas_in_chunk(self): - """Test that multiple commas in a single chunk work correctly""" - # Stream multiple tool calls ensuring at least 2 chunks per complete tool call - chunk1 = '[{"name": "get_weather", "parameters": {"location": "To' - result1 = self.detector.parse_streaming_increment(chunk1, self.tools) - self.assertIsInstance(result1, StreamingParseResult) - chunk2 = 'kyo"}},' - result2 = self.detector.parse_streaming_increment(chunk2, self.tools) - self.assertIsInstance(result2, StreamingParseResult) - - chunk3 = '{"name": "get_weather", "parameters": {"location": "Pa' - result3 = self.detector.parse_streaming_increment(chunk3, self.tools) - self.assertIsInstance(result3, StreamingParseResult) - chunk4 = 'ris"}},' - result4 = self.detector.parse_streaming_increment(chunk4, self.tools) - self.assertIsInstance(result4, StreamingParseResult) - - chunk5 = '{"name": "get_weather"' - result5 = self.detector.parse_streaming_increment(chunk5, self.tools) - self.assertIsInstance(result5, StreamingParseResult) - self.assertGreater( - len(result5.calls), 0, "Should parse tool calls with multiple commas" - ) - - def test_complete_tool_call_with_trailing_comma(self): - """Test that complete tool call with trailing comma parses correctly""" - # Test case: complete tool call followed by comma at end of chunk (split across 2 chunks) - chunk1 = '[{"name": "get_weather", "parameters": {"location": "Tokyo"}' - result1 = self.detector.parse_streaming_increment(chunk1, self.tools) - self.assertIsInstance(result1, StreamingParseResult) - chunk2 = "}, " - result2 = self.detector.parse_streaming_increment(chunk2, self.tools) - self.assertIsInstance(result2, StreamingParseResult) - self.assertGreater(len(result2.calls), 0, "Should parse complete tool call") - - # Test that next chunk with opening brace gets the separator prepended - next_chunk = '{"name": "get_weather", "parameters": {"location": "Paris"}}' - result_next = self.detector.parse_streaming_increment(next_chunk, self.tools) - self.assertIsInstance(result_next, StreamingParseResult) - self.assertGreater( - len(result_next.calls), 0, "Should parse subsequent tool call" - ) - def test_three_tool_calls_separate_chunks_with_commas(self): """Test parsing 3 tool calls in separate chunks with commas at the end""" # First tool call: 2 chunks @@ -3806,11 +3139,6 @@ class TestLfm2Detector(unittest.TestCase): # ==================== has_tool_call tests ==================== - def test_has_tool_call_true(self): - """Test detection of tool call markers.""" - text = '<|tool_call_start|>[get_weather(city="Paris")]<|tool_call_end|>' - self.assertTrue(self.detector.has_tool_call(text)) - def test_has_tool_call_false(self): """Test no false positives for regular text.""" text = "The weather in Paris is nice today." @@ -3900,14 +3228,6 @@ class TestLfm2Detector(unittest.TestCase): params = json.loads(result.calls[0].parameters) self.assertIn("weather", params["query"]) - def test_detect_and_parse_numeric_values(self): - """Test parsing with numeric argument values.""" - text = '<|tool_call_start|>[calculator(expression="5 * 7")]<|tool_call_end|>' - result = self.detector.detect_and_parse(text, self.tools) - - self.assertEqual(len(result.calls), 1) - self.assertEqual(result.calls[0].name, "calculator") - # ==================== detect_and_parse tests (JSON format) ==================== def test_detect_and_parse_json_simple(self): @@ -3978,14 +3298,6 @@ class TestLfm2Detector(unittest.TestCase): # blocks, then parses the complete block. This allows proper handling of both # JSON and Pythonic formats. - def test_streaming_json_complete_in_one_chunk(self): - """Test streaming with complete JSON tool call in one chunk.""" - text = '<|tool_call_start|>{"name": "get_weather", "arguments": {"city": "Rome"}}<|tool_call_end|>' - result = self.detector.parse_streaming_increment(text, self.tools) - - self.assertEqual(len(result.calls), 1) - self.assertEqual(result.calls[0].name, "get_weather") - def test_streaming_json_split_across_chunks(self): """Test streaming with JSON tool call split across multiple chunks - waits for complete block.""" # Reset detector state @@ -4037,16 +3349,6 @@ class TestLfm2Detector(unittest.TestCase): # ==================== Pythonic streaming tests ==================== - def test_streaming_pythonic_complete_in_one_chunk(self): - """Test streaming with complete Pythonic tool call in one chunk.""" - self.detector = Lfm2Detector() - text = '<|tool_call_start|>[get_weather(city="Berlin")]<|tool_call_end|>' - result = self.detector.parse_streaming_increment(text, self.tools) - - self.assertEqual(len(result.calls), 1) - self.assertEqual(result.calls[0].name, "get_weather") - self.assertEqual(json.loads(result.calls[0].parameters), {"city": "Berlin"}) - def test_streaming_pythonic_split_across_chunks(self): """Test streaming with Pythonic tool call split across multiple chunks.""" self.detector = Lfm2Detector() @@ -4148,8 +3450,9 @@ class TestGigaChat3Detector(unittest.TestCase): ] self.detector = GigaChat3Detector() - def test_has_tool_call(self): - """Test detection of tool call markers.""" + def test_has_tool_call_detects_both_markers(self): + # Guards both marker forms the regex ORs together; the parser exercises + # each via detect_and_parse but never the bare predicate. self.assertTrue(self.detector.has_tool_call("function call<|role_sep|>\n{}")) self.assertTrue(self.detector.has_tool_call("<|function_call|>{}")) self.assertFalse(self.detector.has_tool_call("No tool call here")) @@ -4208,28 +3511,6 @@ function call<|role_sep|> self.assertEqual(params["content"]["short_answers"], True) self.assertEqual(params["content"]["hate_emojis"], True) - def test_detect_and_parse_with_content_before(self): - """Test parsing tool call with text content before it.""" - text = 'I\'ll check that for you.<|message_sep|>\n\nfunction call<|role_sep|>\n{"name": "manage_user_memory", "arguments": {"action": "create", "id": "preferences"}}' - result = self.detector.detect_and_parse(text, self.tools) - - self.assertEqual(result.normal_text, "I'll check that for you.") - self.assertEqual(len(result.calls), 1) - self.assertEqual(result.calls[0].name, "manage_user_memory") - - def test_detect_and_parse_with_eos_token(self): - """Test parsing tool call with EOS token at the end.""" - text = '<|message_sep|>\n\nfunction call<|role_sep|>\n{"name": "manage_user_memory", "arguments": {"action": "create", "id": "preferences"}}' - result = self.detector.detect_and_parse(text, self.tools) - - self.assertEqual(result.normal_text, "") - self.assertEqual(len(result.calls), 1) - self.assertEqual(result.calls[0].name, "manage_user_memory") - - params = json.loads(result.calls[0].parameters) - self.assertEqual(params["action"], "create") - self.assertEqual(params["id"], "preferences") - def test_detect_and_parse_with_content_and_eos(self): """Test parsing tool call with content and EOS token.""" text = 'I\'ll remember that.<|message_sep|>\n\nfunction call<|role_sep|>\n{"name": "manage_user_memory", "arguments": {"action": "create", "id": "test"}}' @@ -4561,46 +3842,6 @@ function call<|role_sep|> params = json.loads(tool_calls_by_index[0]["parameters"]) self.assertEqual(params["city"], "NYC") - def test_streaming_json_split_at_quotes(self): - """Test streaming when JSON is split at quote boundaries.""" - chunks = [ - "<|message_sep|>\n\nfunction call<|role_sep|>\n", - '{"name', - '": "', - "get_weather", - '", "arguments', - '": {"city', - '": "', - "Rome", - '"}}', - ] - - tool_calls_by_index = {} - - for chunk in chunks: - result = self.detector.parse_streaming_increment(chunk, self.tools) - - for call in result.calls: - if call.tool_index is not None: - if call.tool_index not in tool_calls_by_index: - tool_calls_by_index[call.tool_index] = { - "name": "", - "parameters": "", - } - - if call.name: - tool_calls_by_index[call.tool_index]["name"] = call.name - if call.parameters: - tool_calls_by_index[call.tool_index][ - "parameters" - ] += call.parameters - - self.assertEqual(len(tool_calls_by_index), 1) - self.assertEqual(tool_calls_by_index[0]["name"], "get_weather") - - params = json.loads(tool_calls_by_index[0]["parameters"]) - self.assertEqual(params["city"], "Rome") - def test_detect_and_parse_function_call_marker_simple_tool_call(self): """Test parsing a simple <|function_call|> tool call (GigaChat3.1-style).""" text = '<|function_call|>{"name": "manage_user_memory", "arguments": {"action": "create", "id": "preferences"}}' @@ -4614,42 +3855,6 @@ function call<|role_sep|> self.assertEqual(params["action"], "create") self.assertEqual(params["id"], "preferences") - def test_detect_and_parse_function_call_marker_with_content_before(self): - """Test parsing <|function_call|> tool call with prefix content.""" - text = ( - 'I\'ll check that for you.<|function_call|>{"name": "get_weather", ' - '"arguments": {"city": "Tokyo"}}' - ) - result = self.detector.detect_and_parse(text, self.tools) - - self.assertEqual(result.normal_text, "I'll check that for you.") - self.assertEqual(len(result.calls), 1) - self.assertEqual(result.calls[0].name, "get_weather") - - params = json.loads(result.calls[0].parameters) - self.assertEqual(params["city"], "Tokyo") - - def test_detect_and_parse_function_call_marker_with_eos_token(self): - """Test parsing <|function_call|> tool call with EOS token at the end.""" - text = '<|function_call|>{"name": "manage_user_memory", "arguments": {"action": "create", "id": "preferences"}}' - result = self.detector.detect_and_parse(text, self.tools) - - self.assertEqual(result.normal_text, "") - self.assertEqual(len(result.calls), 1) - self.assertEqual(result.calls[0].name, "manage_user_memory") - - params = json.loads(result.calls[0].parameters) - self.assertEqual(params["action"], "create") - self.assertEqual(params["id"], "preferences") - - def test_detect_and_parse_function_call_marker_invalid_json(self): - """Test parsing invalid JSON after <|function_call|> marker.""" - text = '<|function_call|>{"name": "manage_user_memory", "arguments": {invalid json}}' - result = self.detector.detect_and_parse(text, self.tools) - - self.assertIn("<|function_call|>", result.normal_text) - self.assertEqual(len(result.calls), 0) - def test_streaming_function_call_marker_simple_tool_call(self): """Test streaming parsing of the <|function_call|> marker form.""" chunks = [ @@ -4689,46 +3894,6 @@ function call<|role_sep|> self.assertEqual(params["action"], "create") self.assertEqual(params["id"], "prefs") - def test_streaming_function_call_marker_json_split_at_quotes(self): - """Test streaming when JSON is split at quote boundaries (<|function_call|>).""" - chunks = [ - "<|function_call|>", - '{"name', - '": "', - "get_weather", - '", "arguments', - '": {"city', - '": "', - "Rome", - '"}}', - ] - - tool_calls_by_index = {} - - for chunk in chunks: - result = self.detector.parse_streaming_increment(chunk, self.tools) - - for call in result.calls: - if call.tool_index is not None: - if call.tool_index not in tool_calls_by_index: - tool_calls_by_index[call.tool_index] = { - "name": "", - "parameters": "", - } - - if call.name: - tool_calls_by_index[call.tool_index]["name"] = call.name - if call.parameters: - tool_calls_by_index[call.tool_index][ - "parameters" - ] += call.parameters - - self.assertEqual(len(tool_calls_by_index), 1) - self.assertEqual(tool_calls_by_index[0]["name"], "get_weather") - - params = json.loads(tool_calls_by_index[0]["parameters"]) - self.assertEqual(params["city"], "Rome") - class TestGetStructureConstraint(unittest.TestCase): """Tests for FunctionCallParser.get_structure_constraint() logic. @@ -4797,17 +3962,6 @@ class TestGetStructureConstraint(unittest.TestCase): self.assertIn('"type":"triggered_tags"', serialized) self.assertIn("<|tool_calls_section_begin|>", serialized) - def test_kimi_routes_through_native_with_section_markers(self): - """xgrammar 0.2.1's Kimi builtin keeps auto tool calls section-wrapped.""" - import xgrammar as xgr - - parser = self._make_parser("kimi_k2", strict=True) - result = parser.get_structure_constraint("auto") - self.assertIsInstance(result[1], xgr.StructuralTag) - serialized = self._constraint_json(result) - self.assertIn("<|tool_calls_section_begin|>", serialized) - self.assertIn("<|tool_calls_section_end|>", serialized) - def test_kimi_auto_no_strict_returns_none(self): """auto without strict should not constrain.""" parser = self._make_parser("kimi_k2", strict=False) @@ -4919,14 +4073,6 @@ class TestQwen25Detector(unittest.TestCase): # -- Non-streaming tests -- - def test_detect_and_parse_single_tool_call(self): - text = '\n{"name": "get_current_weather", "arguments": {"city": "NYC", "state": "NY", "unit": "fahrenheit"}}\n' - result = self.detector.detect_and_parse(text, self.tools) - self.assertEqual(len(result.calls), 1) - self.assertEqual(result.calls[0].name, "get_current_weather") - params = json.loads(result.calls[0].parameters) - self.assertEqual(params["city"], "NYC") - def test_detect_and_parse_multiple_tool_calls(self): text = ( '\n{"name": "get_current_weather", "arguments": {"city": "NYC", "state": "NY", "unit": "fahrenheit"}}\n\n' @@ -4970,21 +4116,6 @@ class TestQwen25Detector(unittest.TestCase): ] += call.parameters return tool_calls_by_index - def test_streaming_single_tool_call(self): - chunks = [ - "\n", - '{"name": "get_current_weather",', - ' "arguments": {"city": "NYC",', - ' "state": "NY",', - ' "unit": "fahrenheit"}}', - "\n", - ] - result = self._collect_streaming_tool_calls(chunks) - self.assertEqual(len(result), 1) - self.assertEqual(result[0]["name"], "get_current_weather") - params = json.loads(result[0]["parameters"]) - self.assertEqual(params["city"], "NYC") - def test_streaming_multiple_tool_calls(self): """Core regression test: multiple tool calls must all be parsed in streaming mode.""" chunks = [ @@ -5132,27 +4263,12 @@ class TestGemma4Detector(unittest.TestCase): self.assertTrue(found_params) - def test_has_tool_call(self): - self.assertTrue( - self.detector.has_tool_call( - '<|tool_call>call:get_weather{location:<|"|>Tokyo<|"|>}' - ) - ) - self.assertFalse(self.detector.has_tool_call("no tool call here")) - def test_detect_and_parse_no_tool_call(self): text = "This is plain text without any tool calls." result = self.detector.detect_and_parse(text, self.tools) self.assertEqual(result.normal_text, text) self.assertEqual(len(result.calls), 0) - def test_detect_and_parse_tool_index(self): - text = '<|tool_call>call:get_weather{location:<|"|>Tokyo<|"|>}' - result = self.detector.detect_and_parse(text, self.tools) - self.assertEqual(len(result.calls), 1) - self.assertEqual(result.calls[0].tool_index, 0) - self.assertEqual(result.calls[0].name, "get_weather") - def test_detect_and_parse_unknown_tool_index(self): text = '<|tool_call>call:unknown_func{arg:<|"|>val<|"|>}' result = self.detector.detect_and_parse(text, self.tools) @@ -5202,11 +4318,6 @@ class TestGemma4Detector(unittest.TestCase): self.assertIs(result["flag"], True) self.assertIs(result["other"], False) - def test_parse_gemma4_args_numbers(self): - result = _parse_gemma4_args("count:42,ratio:3.14") - self.assertEqual(result["count"], 42) - self.assertAlmostEqual(result["ratio"], 3.14) - def test_parse_gemma4_args_string_with_colon(self): result = _parse_gemma4_args('url:<|"|>http://example.com<|"|>') self.assertEqual(result["url"], "http://example.com") diff --git a/test/registered/unit/hardware_backend/mlx/test_attention_patching.py b/test/registered/unit/hardware_backend/mlx/test_attention_patching.py index 0920f7a5c..f31df9571 100644 --- a/test/registered/unit/hardware_backend/mlx/test_attention_patching.py +++ b/test/registered/unit/hardware_backend/mlx/test_attention_patching.py @@ -35,7 +35,6 @@ if _HAS_MLX: MlxAuxiliaryStateReqToTokenPool, MlxModelCacheLayout, find_attention_layers, - is_attention_module, patch_model_attention, ) from sglang.srt.hardware_backend.mlx.model_runner import ( @@ -156,9 +155,6 @@ class TestMlxAttentionPatching(unittest.TestCase): self.assertFalse(isinstance(model.layers[0].linear_attn, MLXAttentionWrapper)) self.assertIsInstance(model.layers[1].self_attn, MLXAttentionWrapper) - def test_projection_only_mixer_is_not_attention(self): - self.assertFalse(is_attention_module(ProjectionOnlyMixer())) - def test_cache_layout_separates_attention_and_auxiliary_layers(self): layout = MlxModelCacheLayout.from_attention_discovery( [object(), object(), object(), object()], @@ -1012,41 +1008,6 @@ class TestMlxAuxiliaryStateRunnerCache(unittest.TestCase): self.assertIsNone(req.mamba_last_track_seqlen) self.assertEqual(pool.auxiliary_state_pool.available_size(), 2) - def test_auxiliary_state_component_keeps_new_live_slot_owned_by_radix(self): - pool = MlxAuxiliaryStateReqToTokenPool( - size=2, - max_context_len=8, - device="cpu", - enable_memory_saver=False, - auxiliary_state_size=4, - ) - req = FakeRequest() - pool.alloc([req]) - component = MlxAuxiliaryStateComponent( - SimpleNamespace(req_to_token_pool=pool), - SimpleNamespace(enable_mamba_extra_buffer=False), - ) - insert_params = InsertParams() - - cache_len = component.prepare_for_caching_req( - req=req, - insert_params=insert_params, - token_ids_len=7, - is_finished=True, - ) - component.cleanup_after_caching_req( - req=req, - is_finished=True, - insert_result=InsertResult(prefix_len=0, mamba_exist=False), - insert_params=insert_params, - ) - - self.assertEqual(cache_len, 7) - self.assertFalse(getattr(insert_params, "mlx_auxiliary_state_uses_track_slot")) - self.assertEqual(insert_params.mamba_value.tolist(), [1]) - self.assertIsNone(req.mamba_pool_idx) - self.assertEqual(pool.auxiliary_state_pool.available_size(), 3) - def test_auxiliary_state_component_frees_stale_track_slot_when_live_slot_inserted( self, ): diff --git a/test/registered/unit/mem_cache/test_hicache_nixl_storage.py b/test/registered/unit/mem_cache/test_hicache_nixl_storage.py index 841133420..63e530fc6 100644 --- a/test/registered/unit/mem_cache/test_hicache_nixl_storage.py +++ b/test/registered/unit/mem_cache/test_hicache_nixl_storage.py @@ -9,10 +9,8 @@ import shutil import socket import subprocess import tempfile -import threading import time import unittest -from typing import List import torch @@ -513,177 +511,6 @@ class TestNixlUnified(CustomTestCase): self.assertEqual(self.hicache.batch_exists(["key1", "key2"]), 1) - def _run_concurrent_stress( - self, is_zero_copy_mode: bool, hicache: HiCacheNixl = None - ): - """One getter thread + one setter thread share the same HiCacheNixl - for ``is_zero_copy_mode``. Defaults to ``self.hicache`` (FILE backend); - pass ``hicache`` to exercise a different backend (e.g. OBJ). - - Phase 1 pre-seeds N preset pages and stores them under fixed keys. - Phase 2 runs the getter (reads the presets back and verifies content) - concurrently with the setter (writes a stream of fresh distinct keys - from a disjoint source region). The kv_buffer regions touched by the - two threads are disjoint so any data corruption observed is from the - backend's shared state (bounce buffers, devId maps, fd pool). - """ - if hicache is None: - hicache = self.hicache - - # 8 preset pages, 8 getter dst pages, 8 setter src pages -> 24 in use. - mock_host = MockMemPoolHost(is_zero_copy_mode=is_zero_copy_mode, num_pages=32) - hicache.register_mem_pool_host(mock_host) - hicache.is_zero_copy = is_zero_copy_mode - - page_size = mock_host.page_size - dtype = mock_host.dtype - num_pages = 8 - - # Disjoint per-thread regions in kv_buffer (indexed by token index). - preset_src = (0, num_pages) - getter_dst = (num_pages, 2 * num_pages) - setter_src = (2 * num_pages, 3 * num_pages) - - # zero_copy=page_first uses dim 1 for the token axis; non-zero-copy= - # layer_first uses dim 2. All buffer accesses below go through this so - # the rest of the harness stays layout-agnostic. - def token_index(start_token: int, n_tokens: int): - s = slice(start_token, start_token + n_tokens) - if is_zero_copy_mode: - return (slice(None), s, slice(None), slice(None), slice(None)) - return (slice(None), slice(None), s, slice(None), slice(None)) - - def page_index(start_page: int, n_pages: int): - return token_index(start_page * page_size, n_pages * page_size) - - def fill_pages(start_page: int, n_pages: int, value_fn): - """value_fn(i) -> scalar value for page i.""" - for i in range(n_pages): - idx = page_index(start_page + i, 1) - shape = mock_host.kv_buffer[idx].shape - mock_host.kv_buffer[idx] = torch.full( - shape, float(value_fn(i)), dtype=dtype - ) - - # Phase 1: distinct value per preset page so a wrong-page result is - # detectable; setter source is constant (value irrelevant to the - # test, just needs to be valid). - fill_pages(preset_src[0], num_pages, lambda i: i + 1) - fill_pages(setter_src[0], num_pages, lambda i: -1.0) - - preset_keys = [f"preset_{int(is_zero_copy_mode)}_{i}" for i in range(num_pages)] - preset_indices = torch.arange( - preset_src[0] * page_size, - preset_src[1] * page_size, - dtype=torch.int64, - ) - self.assertTrue( - all(hicache.batch_set_v1(preset_keys, preset_indices)), - "phase 1: presetting keys failed", - ) - - # Expected per-page-i payload after a successful get into getter_dst. - expected_pages = [ - mock_host.kv_buffer[page_index(preset_src[0] + i, 1)].clone() - for i in range(num_pages) - ] - - # Phase 2. - stop = threading.Event() - errors: List[str] = [] - errors_lock = threading.Lock() - - def record_error(msg: str): - with errors_lock: - errors.append(msg) - - def getter_loop(): - dst_indices = torch.arange( - getter_dst[0] * page_size, - getter_dst[1] * page_size, - dtype=torch.int64, - ) - loops = 0 - while not stop.is_set(): - # Zero the dst pages so a no-op get is observable. - mock_host.kv_buffer[page_index(getter_dst[0], num_pages)] = 0.0 - ok = hicache.batch_get_v1(preset_keys, dst_indices) - if not all(ok): - record_error(f"getter loop {loops}: batch_get_v1 returned {ok}") - return - for i in range(num_pages): - got = mock_host.kv_buffer[page_index(getter_dst[0] + i, 1)] - if not torch.equal(got, expected_pages[i]): - record_error(f"getter loop {loops}: preset page {i} corrupted") - return - loops += 1 - - def setter_loop(): - src_indices = torch.arange( - setter_src[0] * page_size, - setter_src[1] * page_size, - dtype=torch.int64, - ) - loops = 0 - while not stop.is_set(): - keys = [ - f"setter_{int(is_zero_copy_mode)}_{loops}_{i}" - for i in range(num_pages) - ] - ok = hicache.batch_set_v1(keys, src_indices) - if not all(ok): - record_error(f"setter loop {loops}: batch_set_v1 returned {ok}") - return - loops += 1 - - t_get = threading.Thread(target=getter_loop, daemon=True) - t_set = threading.Thread(target=setter_loop, daemon=True) - t_get.start() - t_set.start() - - # Bounded run: long enough to interleave many ops under NIXL I/O - # GIL release, short enough for a unit test. - time.sleep(3.0) - stop.set() - t_get.join(timeout=10) - t_set.join(timeout=10) - - self.assertFalse( - t_get.is_alive() or t_set.is_alive(), - "stress threads failed to stop", - ) - self.assertEqual(errors, [], f"concurrency errors: {errors}") - - @unittest.skipUnless(STRESS_ENABLED, "set SGLANG_RUN_NIXL_STRESS=1 to run") - def test_concurrent_getter_setter_file_zero_copy(self): - """Stress: concurrent getter+setter, FILE backend, zero-copy.""" - self._run_concurrent_stress(is_zero_copy_mode=True) - - @unittest.skipUnless(STRESS_ENABLED, "set SGLANG_RUN_NIXL_STRESS=1 to run") - def test_concurrent_getter_setter_file_non_zero_copy(self): - """Stress: concurrent getter+setter, FILE backend, non-zero-copy.""" - self._run_concurrent_stress(is_zero_copy_mode=False) - - @unittest.skipUnless(STRESS_ENABLED, "set SGLANG_RUN_NIXL_STRESS=1 to run") - @unittest.skipUnless( - MinioFixture.is_available(), "minio binary or boto3 not available" - ) - def test_concurrent_getter_setter_obj_zero_copy(self): - """Stress: concurrent getter+setter, OBJ backend (MinIO), zero-copy.""" - self._run_concurrent_stress( - is_zero_copy_mode=True, hicache=self._make_obj_hicache() - ) - - @unittest.skipUnless(STRESS_ENABLED, "set SGLANG_RUN_NIXL_STRESS=1 to run") - @unittest.skipUnless( - MinioFixture.is_available(), "minio binary or boto3 not available" - ) - def test_concurrent_getter_setter_obj_non_zero_copy(self): - """Stress: concurrent getter+setter, OBJ backend (MinIO), non-zero-copy.""" - self._run_concurrent_stress( - is_zero_copy_mode=False, hicache=self._make_obj_hicache() - ) - @unittest.skipUnless(hasattr(os, "O_DIRECT"), "O_DIRECT not available on this platform") class TestNixlDirectIO(CustomTestCase): diff --git a/test/registered/unit/mem_cache/test_hicache_staged_write_back_dispatch.py b/test/registered/unit/mem_cache/test_hicache_staged_write_back_dispatch.py index 3a4bc7a55..ac3a1b650 100644 --- a/test/registered/unit/mem_cache/test_hicache_staged_write_back_dispatch.py +++ b/test/registered/unit/mem_cache/test_hicache_staged_write_back_dispatch.py @@ -776,58 +776,6 @@ class TestHiCacheStagedWriteBackDispatch(unittest.TestCase): self.assertEqual(captured["host_indices"].device.type, "cpu") self.assertEqual(captured["pool_transfers"][0].host_indices.device.type, "cpu") - def test_hybrid_write_moves_indices_without_page_first_layout(self): - captured = {} - - class FakeHostGroup: - layout = "layer_first" - can_use_write_back_jit = True - - def backup_from_device_all_layer( - self, - device_pool, - host_indices, - device_indices, - io_backend, - pool_transfers=None, - ): - captured["host_indices"] = host_indices - captured["pool_transfers"] = pool_transfers - - op = CacheOperation( - host_indices=_indices(0, 4), - device_indices=_indices(4, 8), - node_id=1, - pool_transfers=[ - PoolTransfer( - name=PoolName.DEEPSEEK_V4_C4, - host_indices=_indices(0, 4), - device_indices=_indices(4, 8), - ) - ], - ) - controller = HybridCacheController.__new__(HybridCacheController) - controller.write_queue = [op] - controller.io_backend = "kernel" - controller.mem_pool_host = FakeHostGroup() - controller.mem_pool_device = None - controller.has_draft = False - controller.write_stream = object() - controller.ack_write_queue = [] - controller._record_transfer_indices_on_stream = lambda *args: None - controller.move_hybrid_indices = mock.Mock( - return_value=(op.host_indices, op.device_indices, op.pool_transfers) - ) - - with mock.patch.object( - hybrid_cache_controller, "device_module", _FakeDeviceModule - ): - controller.start_writing() - - controller.move_hybrid_indices.assert_called_once() - self.assertEqual(captured["host_indices"].device.type, "cpu") - self.assertEqual(captured["pool_transfers"][0].host_indices.device.type, "cpu") - def test_write_back_jit_cache_controller_keeps_host_indices_on_cpu(self): captured = {} @@ -905,43 +853,6 @@ class TestHiCacheStagedWriteBackDispatch(unittest.TestCase): controller.move_indices.assert_called_once() self.assertEqual(captured["host_indices"].device.type, "cpu") - def test_cache_controller_moves_indices_without_page_first_layout(self): - captured = {} - - class FakeHostPool: - layout = "layer_first" - can_use_write_back_jit = True - - def backup_from_device_all_layer( - self, device_pool, host_indices, device_indices, io_backend - ): - captured["host_indices"] = host_indices - - op = ManagerCacheOperation( - host_indices=_indices(0, 4), - device_indices=_indices(4, 8), - node_id=1, - ) - controller = HiCacheController.__new__(HiCacheController) - controller.write_queue = [op] - controller.io_backend = "kernel" - controller.mem_pool_host = FakeHostPool() - controller.mem_pool_device = None - controller.has_draft = False - controller.write_stream = object() - controller.ack_write_queue = [] - controller.move_indices = mock.Mock( - return_value=(op.host_indices, op.device_indices) - ) - - with mock.patch.object( - manager_cache_controller, "device_module", _FakeDeviceModule - ): - controller.start_writing() - - controller.move_indices.assert_called_once() - self.assertEqual(captured["host_indices"].device.type, "cpu") - if __name__ == "__main__": unittest.main() diff --git a/test/registered/unit/mem_cache/test_multi_ended_allocator.py b/test/registered/unit/mem_cache/test_multi_ended_allocator.py index db64ceee3..115de9e21 100644 --- a/test/registered/unit/mem_cache/test_multi_ended_allocator.py +++ b/test/registered/unit/mem_cache/test_multi_ended_allocator.py @@ -408,16 +408,6 @@ class TestMultiEndedAllocator(unittest.TestCase): expected = full_alloc.virtual_to_physical[v] self.assertTrue(bool((buf == expected).all().item())) - def test_translate_kv_loc_without_out_returns_fresh_tensor(self): - """REGRESSION: without `out=`, behavior returns a fresh tensor.""" - _, full_alloc, _, full_kv, _ = self._build_pair() - v = self._alloc(full_alloc, full_kv, 5) - ret = full_alloc.translate_kv_loc(v) - # Fresh tensor: different storage from v2p table - self.assertNotEqual(ret.data_ptr(), full_alloc.virtual_to_physical.data_ptr()) - expected = full_alloc.virtual_to_physical[v] - self.assertTrue(bool((ret == expected).all().item())) - def test_translate_kv_loc_out_matches_no_out(self): """REGRESSION: result of translate_kv_loc(v, out=buf) byte-equals translate_kv_loc(v).""" @@ -833,62 +823,6 @@ class TestUnifiedSWATokenToKVPoolAllocator(unittest.TestCase): self.assertEqual(allocator.full_attn_allocator.allocated_count(), 0) self.assertEqual(allocator.swa_attn_allocator.allocated_count(), 0) - # 7. Joint byte-budget pre-check. - def test_swa_joint_byte_budget_pre_check(self): - # Pick sizes where the byte gap, not slot-index headroom, is the bind. - full_spec = MHASubPoolSpec( - name="full", - layer_num=2, - head_num=2, - head_dim=4, - store_dtype=torch.float16, - grow_direction="up", - ) - swa_spec = MHASubPoolSpec( - name="swa", - layer_num=2, - head_num=2, - head_dim=4, - store_dtype=torch.float16, - grow_direction="down", - ) - n_full, n_swa = 10, 10 - total = n_full * full_spec.entry_bytes() + n_swa * swa_spec.entry_bytes() - pool = UnifiedKVPool( - total_bytes=total, - sub_pool_specs=[full_spec, swa_spec], - device=_DEV, - enable_memory_saver=False, - ) - kvcache = _FakeUnifiedSWAKVPool(pool) - allocator = UnifiedSWATokenToKVPoolAllocator( - unified_buffer=pool, - kvcache=kvcache, - device=_DEV, - full_max_total_num_tokens=n_full, - swa_max_total_num_tokens=n_swa, - need_sort=False, - forward_stream=None, - ) - fa = allocator.full_attn_allocator - sa = allocator.swa_attn_allocator - # Compute the "naive min" against the joint budget — at idle, the - # joint budget is strictly less than min(full.available, swa.available) - # because the joint uses (entry_full + entry_swa) per slot. - naive = min(fa.available_size(), sa.available_size()) - joint = allocator.available_size() - # The joint must be no greater than naive (typically strictly less). - self.assertLessEqual(joint, naive) - # And it must equal `gap_bytes // (entry_full + entry_swa)` clamped - # by slot-room. - gap = sa._byte_low_frontier() - fa._byte_high_frontier() - expected = min( - gap // (fa.entry_bytes + sa.entry_bytes), - fa.max_slots - fa.min_slot_index - fa.allocated_count(), - sa.max_slots - sa.min_slot_index - sa.allocated_count(), - ) - self.assertEqual(joint, expected) - # 8. Watermark rollback on partial alloc failure. def test_swa_alloc_swa_failure_is_fail_loud(self): """The SWA composite runs a tight JOINT pre-check before allocating, so @@ -1194,6 +1128,11 @@ class TestPagedMultiEndedAllocator(unittest.TestCase): self.assertNotEqual(int(full_alloc.virtual_to_physical[v_page].item()), -1) # 5. free() recovers pages via unique(// page_size) — matches upstream. + # REGRESSION: `allocated_count()` MUST return + # TOKENS, not pages -- matching upstream's convention that all external + # capacity methods report tokens. At page_size > 1, returning pages + # here breaks the leak invariant + # (`available + evictable + ... == total`, with all terms in tokens). def test_paged_free_unique_by_page(self): _, full_alloc, _, full_kv, _ = self._build() a = full_alloc.alloc(self.PAGE_SIZE * 2) # 2 pages = 2*PS tokens @@ -1501,31 +1440,6 @@ class TestPagedMultiEndedAllocator(unittest.TestCase): int(free_before.numel()), ) - # 12. translate_kv_loc preserves token-level identity end-to-end. - def test_paged_translate_kv_loc_token_round_trip(self): - _, full_alloc, _, _, _ = self._build() - v = full_alloc.alloc(self.PAGE_SIZE * 2) - # Build the composite-style translation manually: virt_page * ps + offset. - ps = self.PAGE_SIZE - virt_pages = v // ps - offsets = v % ps - phys_pages = full_alloc.virtual_to_physical[virt_pages] - phys_tokens = phys_pages * ps + offsets - # `phys_tokens` should be a coherent set of two contiguous PAGES. - phys_pages_unique = sorted(set(phys_pages.tolist())) - self.assertEqual(len(phys_pages_unique), 2) - # Within each page the tokens go through offsets 0..7 in order. - for p in phys_pages_unique: - page_phys = sorted( - int(t) - for i, t in enumerate(phys_tokens.tolist()) - if int(phys_pages[i].item()) == p - ) - self.assertEqual( - page_phys, - [p * ps + i for i in range(ps)], - ) - # REGRESSION: `translate_kv_loc(virt, out=buf)` must work # under page_size > 1 — the page-math branch writes via # `index_select(out=out)` + in-place `mul_` / `add_` and must match the @@ -1679,28 +1593,6 @@ class TestPagedMultiEndedAllocator(unittest.TestCase): ) self.assertTrue(bool((buf[:ps] == 0).all().item())) - # 13. REGRESSION: `allocated_count()` MUST return - # TOKENS, not pages — matching upstream's convention that all external - # capacity methods report tokens. At page_size > 1, returning pages - # here breaks the leak invariant - # (`available + evictable + ... == total`, with all terms in tokens). - def test_paged_allocated_count_returns_tokens(self): - _, full_alloc, _, _, _ = self._build() - PS = self.PAGE_SIZE - # Idle → allocated_count == 0. - self.assertEqual(full_alloc.allocated_count(), 0) - # Alloc 2 pages = 2 * PS tokens. - v = full_alloc.alloc(2 * PS) - self.assertIsNotNone(v) - # allocated_count() must report TOKENS (= 2 * PS), not pages (= 2). - self.assertEqual( - full_alloc.allocated_count(), - 2 * PS, - "REGRESSION: allocated_count() must return TOKENS at page_size > 1", - ) - # _allocated_pages() is the page-granular internal helper. - self.assertEqual(full_alloc._allocated_pages(), 2) - # 14. REGRESSION: the leak-invariant terms used by the # scheduler runtime checker must all be in TOKENS. Specifically # `full_available_size() + allocated_tokens == static_cap` must hold for @@ -2032,24 +1924,6 @@ class TestLazyCompaction(unittest.TestCase): p = int(alloc.virtual_to_physical[v].item()) kv.buf[p] = int(v) - def test_lazy_state_initialized(self): - """Lazy allocator initializes the new state cleanly.""" - _pool, fa, _kv = self._make_full(lazy=True) - self.assertTrue(fa.lazy_compaction) - self.assertEqual(len(fa._free_phys_pages), 0) - self.assertEqual(fa._pending_reuse, {}) - self.assertEqual(fa.live_page_count, 0) - # Watermark + free virtual list start equivalent to eager. - self.assertEqual(fa.watermark_physical, fa.min_page_index) - - def test_lazy_alloc_increments_live_page_count(self): - _pool, fa, _kv = self._make_full(lazy=True) - tokens = fa.alloc(8) - self.assertIsNotNone(tokens) - self.assertEqual(int(tokens.numel()), 8) - self.assertEqual(fa.live_page_count, 8) - self.assertEqual(len(fa._free_phys_pages), 0) - def test_lazy_free_boundary_shortcut(self): """Boundary absorption is DEFERRED to `_flush` (the hot path `_free_lazy` does only a `torch.cat`, no watermark mutation). @@ -2075,21 +1949,6 @@ class TestLazyCompaction(unittest.TestCase): self.assertEqual(len(fa._free_phys_pages), 0) self.assertEqual(fa.live_page_count, 2) - def test_lazy_free_non_boundary_pushes_hole(self): - """Freeing a non-boundary page enters _free_phys_pages, watermark - stays put. - """ - _pool, fa, _kv = self._make_full(lazy=True) - a = fa.alloc(5) - wm_before = fa.watermark_physical - # Free a middle id (NOT the topmost), boundary-shortcut should - # NOT fire. - mid = a[2:3].clone() - fa.free(mid) - self.assertEqual(fa.watermark_physical, wm_before) - self.assertEqual(len(fa._free_phys_pages), 1) - self.assertEqual(fa.live_page_count, 4) - def test_lazy_free_inward_walk(self): """The inward walk (multiple contiguous holes absorbed into the watermark in one pass) is DEFERRED to `_flush`. After @@ -2248,29 +2107,6 @@ class TestLazyCompaction(unittest.TestCase): ) self.assertEqual(lazy_data[v], lazy_stamps[v], f"lazy: KV[v={v}] != stamp") - def test_lazy_hole_set_directional_pop(self): - """The _HoleSet pops smallest-first for grow-up; alloc must drain - the deepest hole first (the greedy clustering rule keeps near- - boundary holes available for cheap absorption by compaction). - """ - _pool, fa, _kv = self._make_full(lazy=True) - a = fa.alloc(6) - # Free middle and lower middles so the holes are NOT at boundary. - fa.free(a[1:2].clone()) # frees physical at index v2p[a[1]] - fa.free(a[3:4].clone()) - # Capture which physical pages are now in the hole set. - # `_free_phys_pages` is a torch.Tensor; `.tolist()` returns - # Python ints so `sorted` produces ints (not 0-dim tensors). - holes_before = sorted(fa._free_phys_pages.tolist()) - self.assertEqual(len(holes_before), 2) - # Alloc 1 — should drain a hole (grow-up). - # With sort-after-merge OFF (default), the drain order - # is FIFO over the free-list tensor — NOT "smallest first". - # We only assert that the bound physical is ONE OF the holes. - a2 = fa.alloc(1) - bound_phys = int(fa.virtual_to_physical[int(a2.item())].item()) - self.assertIn(bound_phys, holes_before) - def test_lazy_non_urgent_stops_at_write_set_blocker(self): """Write-race case: when the topmost survivor IS in an in-flight batch's write-set, non-urgent `_flush` STOPS the @@ -2358,57 +2194,6 @@ class TestLazyCompaction(unittest.TestCase): self.assertEqual(len(fa._pending_reuse), 0) self.assertEqual(len(fa._pending_reuse_pages_cpu), 0) - def test_lazy_pending_reuse_urgent_wait(self): - """Under urgent drain, an unfired event triggers wait_event; we - simulate this by checking that the drain ALSO releases unfired - entries (with a fake event whose `query` is False — `wait_event` is - a no-op in CPU mode since there's no current stream's wait_event for - a FakeEvent, so we test the release path).""" - _pool, fa, _kv = self._make_full(lazy=True) - a = fa.alloc(4) - - class _FakeEvent: - def __init__(self): - self.waited = False - - def query(self): - return False # never fires - - # Inject ONE batch entry into _pending_reuse keyed by - # Event. Value is `(cpu_list, gpu_tensor)`. The parallel CPU - # set must also be updated. - # (Simulates a prior compaction whose event hasn't fired.) - p = int(fa.virtual_to_physical[int(a[2].item())].item()) - # Clear v2p/p2v so post-drain reuse is safe. - fa.virtual_to_physical[int(a[2].item())] = -1 - fa.physical_to_virtual[p] = -1 - ev = _FakeEvent() - gpu_t = torch.tensor([p], dtype=torch.int64, device=fa.device) - fa._pending_reuse[ev] = ([p], gpu_t) - fa._pending_reuse_pages_cpu.add(p) - # Urgent drain — should release p despite event.query()=False. - # (CPU shim: torch.cuda.current_stream() may not exist; wrap try.) - try: - fa._drain_pending_reuse(urgent=True) - except Exception: - # CPU: wait_event may not work; this test is GPU-only. - self.skipTest("wait_event requires CUDA") - self.assertEqual(len(fa._pending_reuse), 0) - self.assertEqual(len(fa._pending_reuse_pages_cpu), 0) - - def test_lazy_flush_opportunistic_hook(self): - """The public flush_opportunistic method runs the non-urgent path - and is safe to call when no holes exist.""" - _pool, fa, _kv = self._make_full(lazy=True) - # No holes → returns 0 moves, no-op. - self.assertEqual(fa.flush_opportunistic(), 0) - # Create a hole then call flush_opportunistic; latest_event=None - # means src releases immediately. - a = fa.alloc(3) - fa.free(a[0:1].clone()) - moves = fa.flush_opportunistic() - self.assertGreaterEqual(moves, 1) - class TestO3FusedAllocBind(unittest.TestCase): """Fused take_physical_pages + bind_pages. @@ -2465,17 +2250,6 @@ class TestO3FusedAllocBind(unittest.TestCase): ma.bind_peer(fa) return pool, fa, full_kv - def test_helper_exists_and_returns_tensor(self): - """The helper `_alloc_bind_fast_or_slow` is wired and returns a - tensor on success.""" - _pool, fa, _kv = self._make_full(lazy=True) - v_pages = torch.tensor([10, 11, 12], dtype=torch.int64, device="cuda") - phys = fa._alloc_bind_fast_or_slow(v_pages, 3) - self.assertIsNotNone(phys) - self.assertEqual(phys.shape, (3,)) - self.assertEqual(phys.dtype, torch.int64) - self.assertEqual(phys.device.type, "cuda") - def test_fast_path_when_no_holes(self): """When `_free_phys_pages` is empty, the fast path fires. Verifies: watermark advanced, v2p and p2v scattered correctly, diff --git a/test/registered/unit/mem_cache/test_radix_cache_unit.py b/test/registered/unit/mem_cache/test_radix_cache_unit.py index 32afb8840..47b565678 100644 --- a/test/registered/unit/mem_cache/test_radix_cache_unit.py +++ b/test/registered/unit/mem_cache/test_radix_cache_unit.py @@ -17,7 +17,6 @@ Usage: python -m pytest test_radix_cache_unit.py::TestRadixCache::test_insert_basic """ -from sglang.srt.mem_cache.common import available_and_evictable_str from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci # CPU-based unit test, runs quickly on any GPU runner @@ -25,7 +24,6 @@ register_cuda_ci(est_time=15, stage="base-b", runner_config="1-gpu-small") register_amd_ci(est_time=5, suite="stage-b-test-1-gpu-small-amd") import random -import time import unittest import unittest.mock from array import array @@ -49,13 +47,6 @@ DEFAULT_PAGE_SIZE = 4 class TestRadixKey(unittest.TestCase): """Test cases for RadixKey class.""" - def test_init_basic(self): - """Test basic initialization of RadixKey.""" - token_ids = [1, 2, 3, 4] - key = RadixKey(array("q", token_ids)) - self.assertEqual(list(key.token_ids), token_ids) - self.assertIsNone(key.extra_key) - def test_init_with_extra_key(self): """Test initialization with extra_key.""" token_ids = [1, 2, 3] @@ -64,20 +55,6 @@ class TestRadixKey(unittest.TestCase): self.assertEqual(list(key.token_ids), token_ids) self.assertEqual(key.extra_key, extra_key) - def test_len(self): - """Test __len__ method.""" - key = RadixKey(array("q", [1, 2, 3])) - self.assertEqual(len(key), 3) - - empty_key = RadixKey(array("q", [])) - self.assertEqual(len(empty_key), 0) - - def test_iter(self): - """Test __iter__ method.""" - token_ids = [1, 2, 3, 4] - key = RadixKey(array("q", token_ids)) - self.assertEqual(list(key), token_ids) - def test_len_and_iter(self): """Test __len__ and __iter__ methods.""" test_cases = [ @@ -127,21 +104,6 @@ class TestRadixKey(unittest.TestCase): with self.assertRaises(IndexError): _ = key[10] # Out of bounds - def test_repr(self): - """Test __repr__ method.""" - key = RadixKey(array("q", [1, 2, 3]), "test") - repr_str = repr(key) - self.assertIn("RadixKey", repr_str) - self.assertIn("extra_key='test'", repr_str) - self.assertIn("[1, 2, 3]", repr_str) - - def test_repr_long_token_ids(self): - """Test __repr__ with long token_ids.""" - long_tokens = list(range(15)) - key = RadixKey(array("q", long_tokens)) - repr_str = repr(key) - self.assertIn("...", repr_str) # Should be truncated - def _assert_match(self, a, b, page_size, expected, is_bigram=False): key_a = RadixKey(array("q", a), is_bigram=is_bigram) key_b = RadixKey(array("q", b), is_bigram=is_bigram) @@ -225,13 +187,6 @@ class TestTreeNode(unittest.TestCase): node2 = TreeNode() self.assertEqual(node2.id, 1) # Counter was incremented - def test_counter_increment(self): - """Test that counter increments properly.""" - node1 = TreeNode() - node2 = TreeNode() - self.assertEqual(node1.id, 0) - self.assertEqual(node2.id, 1) - def test_evicted_backuped_properties(self): """Test evicted and backuped properties.""" test_cases = [ @@ -313,15 +268,6 @@ class TestTreeNode(unittest.TestCase): n4.hash_value = ["h4"] self.assertEqual(n4.get_prefix_hash_values(n3), ["h1", "h2", "h3"]) - def test_lt_comparison(self): - """Test less than comparison based on last_access_time.""" - node1 = TreeNode() - time.sleep(0.001) # Small delay to ensure different timestamps - node2 = TreeNode() - - self.assertTrue(node1 < node2) - self.assertFalse(node2 < node1) - class TestRadixCache(unittest.TestCase): """Test cases for RadixCache class.""" @@ -677,46 +623,6 @@ class TestRadixCache(unittest.TestCase): match_len = len(result.device_indices) self.assertEqual(match_len % page_size, 0) - def test_pretty_print_basic(self): - """Test pretty_print produces output.""" - cache = RadixCache.create_simulated() - - cache.insert( - InsertParams( - key=RadixKey(array("q", [1, 2, 3])), - value=torch.tensor([10, 20, 30], dtype=torch.int64), - ) - ) - - # Just test that it doesn't crash - try: - cache.pretty_print() - except Exception as e: - self.fail(f"pretty_print raised an exception: {e}") - - def test_all_values_flatten(self): - """Test all_values_flatten method.""" - cache = RadixCache.create_simulated() - - cache.insert( - InsertParams( - key=RadixKey(array("q", [1, 2])), - value=torch.tensor([10, 20], dtype=torch.int64), - ) - ) - cache.insert( - InsertParams( - key=RadixKey(array("q", [3, 4])), - value=torch.tensor([30, 40], dtype=torch.int64), - ) - ) - - all_values = cache.all_values_flatten() - self.assertEqual(len(all_values), 4) - # Values should contain all inserted values (order may vary) - values_set = set(all_values.tolist()) - self.assertEqual(values_set, {10, 20, 30, 40}) - def test_advanced_prefix_match_with_node_splits(self): """Advanced prefix matching: splits inside nodes and across pages.""" for page_size in [1, 2]: @@ -895,14 +801,6 @@ class TestRadixCache(unittest.TestCase): # The cache size should be within reasonable bounds of the actual allocated memory. self.assertLess(torch_allocated, cache_size_bytes * 2) - def test_available_and_evictable_str(self): - mock_allocator = unittest.mock.Mock() - mock_allocator.available_size.return_value = 10 - cache: RadixCache = RadixCache.create_simulated(mock_allocator=mock_allocator) - - print(cache.available_and_evictable_str()) - print(available_and_evictable_str(cache)) - if __name__ == "__main__": unittest.main() diff --git a/test/registered/unit/mem_cache/test_unified_radix_cache_bench.py b/test/registered/unit/mem_cache/test_unified_radix_cache_bench.py index ab4d53a43..03906f7b4 100644 --- a/test/registered/unit/mem_cache/test_unified_radix_cache_bench.py +++ b/test/registered/unit/mem_cache/test_unified_radix_cache_bench.py @@ -740,27 +740,6 @@ _CI_BENCH_CONFIGS = [ num_seqs=5000, kv_size=500_000, ), - dict( - label="FULL_SWA_ps1", - components=(ComponentType.FULL, ComponentType.SWA), - page_size=1, - num_seqs=1000, - kv_size=100_000, - ), - dict( - label="FULL_ps16", - components=(ComponentType.FULL,), - page_size=16, - num_seqs=1000, - kv_size=100_000, - ), - dict( - label="FULL_SWA_ps16", - components=(ComponentType.FULL, ComponentType.SWA), - page_size=16, - num_seqs=1000, - kv_size=100_000, - ), dict( label="FULL_ps128", components=(ComponentType.FULL,), diff --git a/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py b/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py index 77b2f5716..0f724c764 100644 --- a/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py +++ b/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py @@ -26,7 +26,6 @@ from sglang.srt.mem_cache.allocator.swa import SWATokenToKVPoolAllocator from sglang.srt.mem_cache.base_prefix_cache import ( DecLockRefParams, EvictParams, - EvictResult, IncLockRefResult, InitLoadBackParams, InsertParams, @@ -738,22 +737,6 @@ class UnifiedRadixCacheSuite: self.assertEqual(len(m.device_indices), len(base)) cache.sanity_check() - def test_evict_basic(self): - cache, allocator, req_to_token_pool = build_fixture(self.cfg) - seq_a = self._make_seq(1, 2) - seq_b = self._make_seq(500, 2) - - self._insert(cache, allocator, req_to_token_pool, seq_a) - self._insert(cache, allocator, req_to_token_pool, seq_b) - total = len(seq_a) + len(seq_b) - self.assertEqual(cache.full_evictable_size(), total) - - result = cache.evict(EvictParams(num_tokens=len(seq_a))) - self.assertIsInstance(result, EvictResult) - self.assertGreaterEqual(result.num_tokens_evicted, len(seq_a)) - self.assertTrue(cache.full_evictable_size() <= len(seq_b)) - cache.sanity_check() - def test_evict_respects_lock_ref(self): """Lock protects from eviction; unlock allows re-eviction.""" cache, allocator, req_to_token_pool = build_fixture(self.cfg) @@ -793,24 +776,6 @@ class UnifiedRadixCacheSuite: self.assertEqual(result.mamba_num_evicted, 0) cache.sanity_check() - def test_evict_until_empty(self): - cache, allocator, req_to_token_pool = build_fixture(self.cfg) - seqs = [self._make_seq(i * 100, 2) for i in range(5)] - for s in seqs: - self._insert(cache, allocator, req_to_token_pool, s) - total = sum(len(s) for s in seqs) - self.assertEqual(cache.full_evictable_size(), total) - - result = cache.evict(EvictParams(num_tokens=total * 2)) - self.assertGreaterEqual(result.num_tokens_evicted, total) - self.assertEqual(cache.full_evictable_size(), 0) - if self.cfg.has_mamba: - self.assertEqual(cache.mamba_evictable_size(), 0) - - m = cache.match_prefix(MatchPrefixParams(key=RadixKey(array("q", seqs[0])))) - self.assertEqual(len(m.device_indices), 0) - cache.sanity_check() - def test_prev_prefix_len(self): """Three-step test: free overlap, free partial, no free.""" cache, allocator, req_to_token_pool = build_fixture(self.cfg) @@ -860,28 +825,6 @@ class UnifiedRadixCacheSuite: self.assertEqual(allocator.available_size(), avail_before - len(seq_3p)) cache.sanity_check() - def test_node_split_at_boundary(self): - cache, allocator, req_to_token_pool = build_fixture(self.cfg) - base = self._make_seq(1, 3) - self._insert(cache, allocator, req_to_token_pool, base) - - fork_a = base + self._make_seq(100, 1) - fork_b = base + self._make_seq(200, 1) - - self._insert(cache, allocator, req_to_token_pool, fork_a) - result = self._insert(cache, allocator, req_to_token_pool, fork_b) - self.assertEqual(result.prefix_len, len(base)) - - for seq in (fork_a, fork_b): - m = cache.match_prefix(MatchPrefixParams(key=RadixKey(array("q", seq)))) - self.assertEqual(len(m.device_indices), len(seq)) - - m = cache.match_prefix( - MatchPrefixParams(key=RadixKey(array("q", base + self._make_seq(999, 1)))) - ) - self.assertEqual(len(m.device_indices), len(base)) - cache.sanity_check() - def test_cache_finished_req_insert(self): cache, allocator, req_to_token_pool = build_fixture(self.cfg) ps = self.cfg.page_size @@ -1086,34 +1029,6 @@ class UnifiedRadixCacheSuite: cache.pretty_print() cache.sanity_check() - def test_multi_branch_tree(self): - cache, allocator, req_to_token_pool = build_fixture(self.cfg) - base = self._make_seq(1, 2) - self._insert(cache, allocator, req_to_token_pool, base) - - for suffix_start in [100, 200, 300]: - seq = base + self._make_seq(suffix_start, 2) - self._insert(cache, allocator, req_to_token_pool, seq) - - for suffix_start in [100, 200, 300]: - seq = base + self._make_seq(suffix_start, 2) - m = cache.match_prefix(MatchPrefixParams(key=RadixKey(array("q", seq)))) - self.assertEqual(len(m.device_indices), len(seq)) - - m = cache.match_prefix( - MatchPrefixParams(key=RadixKey(array("q", base + self._make_seq(999, 1)))) - ) - self.assertEqual(len(m.device_indices), len(base)) - cache.sanity_check() - - def test_paged_child_key_is_tuple(self): - if self.cfg.page_size == 1: - self.skipTest("page_size > 1 only") - cache, _, _ = build_fixture(self.cfg) - key = RadixKey(array("q", self._make_seq(1, 1))) - child_key = key.child_key(cache.page_size) - self.assertIsInstance(child_key, tuple) - def test_paged_match_truncates_unaligned_key(self): """match_prefix internally aligns keys to page boundary.""" if self.cfg.page_size == 1: @@ -1227,18 +1142,6 @@ class UnifiedRadixCacheSuite: self.assertEqual(len(m.device_indices), 0) cache.sanity_check() - def test_mamba_evict_result_accounting(self): - if not self.cfg.has_mamba: - self.skipTest("requires Mamba component") - cache, allocator, req_to_token_pool = build_fixture(self.cfg) - seq = self._make_seq(1, 3) - self._insert(cache, allocator, req_to_token_pool, seq) - - result = cache.evict(EvictParams(num_tokens=len(seq))) - self.assertGreaterEqual(result.num_tokens_evicted, len(seq)) - self.assertGreaterEqual(result.mamba_num_evicted, 1) - cache.sanity_check() - def test_mamba_evict_cascades_on_full_leaf(self): if not self.cfg.has_mamba: self.skipTest("requires Mamba component") @@ -1276,17 +1179,6 @@ class UnifiedRadixCacheSuite: ) cache.sanity_check() - def test_swa_insert_and_match(self): - if not self.cfg.has_swa: - self.skipTest("requires SWA component") - cache, allocator, req_to_token_pool = build_fixture(self.cfg) - seq = self._make_seq(1, 3) - self._insert(cache, allocator, req_to_token_pool, seq) - - m = cache.match_prefix(MatchPrefixParams(key=RadixKey(array("q", seq)))) - self.assertEqual(len(m.device_indices), len(seq)) - cache.sanity_check() - def test_swa_unfinished_recovery_preserves_locked_full_value(self): if not self.cfg.has_swa or self.cfg.has_mamba: self.skipTest("requires SWA without Mamba") @@ -1391,34 +1283,6 @@ class UnifiedRadixCacheSuite: self.assertIsNone(node.component_data[ComponentType.SWA].value) cache.sanity_check() - def test_swa_evict_cascades(self): - """Evict SWA tokens via swa_num_tokens — cascades to lower-priority components.""" - if not self.cfg.has_swa: - self.skipTest("requires SWA component") - cache, allocator, req_to_token_pool = build_fixture(self.cfg) - seq_short = self._make_seq(1, 2) - seq_long = seq_short + self._make_seq(500, 2) - self._insert(cache, allocator, req_to_token_pool, seq_short) - self._insert(cache, allocator, req_to_token_pool, seq_long) - - result = cache.evict(EvictParams(num_tokens=0, swa_num_tokens=len(seq_short))) - self.assertGreater(result.swa_num_tokens_evicted, 0) - cache.sanity_check() - - def test_swa_evict_cascades_mamba(self): - """SWA eviction on an internal node cascades to Mamba.""" - if not self.cfg.has_swa or not self.cfg.has_mamba: - self.skipTest("requires SWA and Mamba components") - cache, allocator, req_to_token_pool = build_fixture(self.cfg) - seq_short = self._make_seq(1, 3) - seq_long = seq_short + self._make_seq(500, 4) - self._insert(cache, allocator, req_to_token_pool, seq_short) - self._insert(cache, allocator, req_to_token_pool, seq_long) - - result = cache.evict(EvictParams(num_tokens=0, swa_num_tokens=len(seq_short))) - self.assertGreaterEqual(result.swa_num_tokens_evicted, 0) - cache.sanity_check() - def test_leaf_transition_swa_evict_spares_locked_full(self): if not self.cfg.has_swa or not self.cfg.has_mamba: self.skipTest("requires SWA and Mamba components") @@ -1678,46 +1542,6 @@ class UnifiedRadixCacheSuite: self.assertTrue(cache._is_device_leaf(node_a)) cache.sanity_check() - def test_swa_evict_full_leaf_cascades_all(self): - if not self.cfg.has_swa: - self.skipTest("requires SWA component") - cache, allocator, req_to_token_pool = build_fixture(self.cfg) - seq_a = self._make_seq(1, 2) - seq_b = self._make_seq(500, 2) - self._insert(cache, allocator, req_to_token_pool, seq_a) - self._insert(cache, allocator, req_to_token_pool, seq_b) - - result = cache.evict(EvictParams(num_tokens=len(seq_a))) - self.assertGreaterEqual(result.num_tokens_evicted, len(seq_a)) - self.assertGreater(result.swa_num_tokens_evicted, 0) - if self.cfg.has_mamba: - self.assertGreaterEqual(result.mamba_num_evicted, 1) - cache.sanity_check() - - def test_swa_lock_protects_from_eviction(self): - if not self.cfg.has_swa: - self.skipTest("requires SWA component") - cache, allocator, req_to_token_pool = build_fixture(self.cfg) - seq_a = self._make_seq(1, 2) - seq_b = self._make_seq(500, 2) - self._insert(cache, allocator, req_to_token_pool, seq_a) - self._insert(cache, allocator, req_to_token_pool, seq_b) - - m = cache.match_prefix(MatchPrefixParams(key=RadixKey(array("q", seq_a)))) - lock_result = cache.inc_lock_ref(m.last_device_node) - - result = cache.evict(EvictParams(num_tokens=len(seq_a) + len(seq_b))) - self.assertGreaterEqual(result.num_tokens_evicted, len(seq_b)) - - m = cache.match_prefix(MatchPrefixParams(key=RadixKey(array("q", seq_a)))) - self.assertEqual(len(m.device_indices), len(seq_a)) - - cache.dec_lock_ref( - m.last_device_node, - DecLockRefParams(swa_uuid_for_lock=lock_result.swa_uuid_for_lock), - ) - cache.sanity_check() - def test_swa_leaf_capped_to_window_on_insert(self): """A long SWA leaf is split so locking it protects one window of SWA while full attention still protects the whole sequence.""" @@ -1909,44 +1733,6 @@ class UnifiedRadixCacheSuite: ) cache.sanity_check() - def test_swa_lru_cushion_bound_is_sliding_window_plus_page_size(self): - if not self._swa_pinning_cfg_supported(): - self.skipTest("requires SWA-only config with node size >= cushion") - cache, allocator, req_to_token_pool = build_fixture(self.cfg) - - seq_a = self._make_seq(1, 8) - seq_ab = seq_a + self._make_seq(100, 8) - seq_abc = seq_ab + self._make_seq(200, 8) - self._insert(cache, allocator, req_to_token_pool, seq_a) - self._insert(cache, allocator, req_to_token_pool, seq_ab) - self._insert(cache, allocator, req_to_token_pool, seq_abc) - - seq_side = self._make_seq(900, 5) - self._insert(cache, allocator, req_to_token_pool, seq_side) - - pre = self._swa_lru_order(cache) - self.assertEqual(len(pre), 8) - side_node, c_node, b_node, a_node = pre[0], pre[2], pre[4], pre[6] - c_prefix = pre[3] # C's prefix pairs with its tail (c_node) at pre[2:4] - - m = cache.match_prefix(MatchPrefixParams(key=RadixKey(array("q", seq_abc)))) - self.assertEqual(len(m.device_indices), len(seq_abc)) - post = self._swa_lru_order(cache) - - cushion = self.cfg.sliding_window_size + self.cfg.page_size - # Under leaf-cap no single node exceeds the cushion; it spans C's capped - # tail plus its prefix, so both of C's nodes are refreshed to the MRU - # side while B and A keep their relative order below. - self.assertLess(len(c_node.key), cushion) - self.assertIn(c_node, post[:2]) - self.assertIn(c_prefix, post[:2]) - side_pos = post.index(side_node) - b_pos = post.index(b_node) - a_pos = post.index(a_node) - self.assertLess(side_pos, b_pos, "B was below side in pre, must stay below") - self.assertLess(b_pos, a_pos, "A was below B in pre, must stay below") - cache.sanity_check() - def test_swa_eager_eviction_on_unfinished_req(self): if not self.cfg.has_swa or self.cfg.has_mamba: self.skipTest( @@ -2362,51 +2148,6 @@ class UnifiedRadixCacheSuite: self.assertEqual(result.num_tokens_evicted, before - after) cache.sanity_check() - def test_evict_locked_subtree_skipped(self): - """All nodes in a locked path are skipped during eviction.""" - cache, allocator, req_to_token_pool = build_fixture(self.cfg) - seq_a = self._make_seq(1, 3) - seq_b = self._make_seq(500, 2) - self._insert(cache, allocator, req_to_token_pool, seq_a) - self._insert(cache, allocator, req_to_token_pool, seq_b) - - # Lock seq_a - m = cache.match_prefix(MatchPrefixParams(key=RadixKey(array("q", seq_a)))) - lr = cache.inc_lock_ref(m.last_device_node) - - # Try to evict everything - total = cache.full_evictable_size() + cache.full_protected_size() - result = cache.evict(EvictParams(num_tokens=total)) - - # seq_a should still be matchable (protected) - m2 = cache.match_prefix(MatchPrefixParams(key=RadixKey(array("q", seq_a)))) - self.assertEqual(len(m2.device_indices), len(seq_a)) - - cache.dec_lock_ref( - m.last_device_node, - DecLockRefParams(swa_uuid_for_lock=getattr(lr, "swa_uuid_for_lock", None)), - ) - cache.sanity_check() - - def test_mamba_internal_tombstone_evict(self): - """Mamba eviction on internal node tombstones mamba only, keeps Full.""" - if not self.cfg.has_mamba: - self.skipTest("requires Mamba component") - cache, allocator, req_to_token_pool = build_fixture(self.cfg) - # Create internal node with mamba and leaf extending it - seq_short = self._make_seq(1, 2) - seq_long = seq_short + self._make_seq(500, 2) - self._insert(cache, allocator, req_to_token_pool, seq_short) - self._insert(cache, allocator, req_to_token_pool, seq_long) - - # Evict only mamba - result = cache.evict(EvictParams(num_tokens=0, mamba_num=10)) - self.assertEqual(cache.mamba_evictable_size(), 0) - - # Full should still be accessible for at least the long seq base - # (mamba gone breaks match, but full data might still be in tree) - cache.sanity_check() - def test_evict_reinsert_after_full_eviction(self): """After evicting everything, new inserts work correctly.""" cache, allocator, req_to_token_pool = build_fixture(self.cfg) @@ -2990,33 +2731,6 @@ class UnifiedRadixCacheSuite: cache.sanity_check() - def test_hicache_node_states(self): - """Verify device-only to device+host transition after real backup.""" - if self._skip_unsupported_hicache_test(): - return - cache, allocator, req_to_token_pool = self._build_hicache_fixture() - seq = self._make_seq(1, 2) - self._insert(cache, allocator, req_to_token_pool, seq) - - # Find the leaf node - m = cache.match_prefix(MatchPrefixParams(key=RadixKey(array("q", seq)))) - node = m.last_device_node - self.assertIsNot(node, cache.root_node) - - ct = ComponentType.FULL - # S1: device only - self.assertIsNotNone(node.component_data[ct].value) - self.assertIsNone(node.component_data[ct].host_value) - self.assertFalse(node.backuped) - self.assertFalse(node.evicted) - - self._backup_node(cache, node) - self.assertIsNotNone(node.component_data[ct].value) - self.assertIsNotNone(node.component_data[ct].host_value) - self.assertTrue(node.backuped) - self.assertFalse(node.evicted) - cache.sanity_check() - def test_hicache_evict_to_host(self): """Evicting a backed-up device leaf demotes it to host-only state.""" if self._skip_unsupported_hicache_test(): @@ -3130,27 +2844,6 @@ class UnifiedRadixCacheSuite: self.assertTrue(cur.evicted and cur.backuped) cache.sanity_check() - def test_hicache_d_leaf_h_leaf_mutual_exclusion(self): - """D-leaf and H-leaf sets are always disjoint.""" - if self._skip_unsupported_hicache_test(): - return - cache, allocator, req_to_token_pool = self._build_hicache_fixture() - seqs = [self._make_seq(i * 100, 2) for i in range(4)] - for s in seqs: - self._insert(cache, allocator, req_to_token_pool, s) - - for i in range(2): - m = cache.match_prefix(MatchPrefixParams(key=RadixKey(array("q", seqs[i])))) - self._backup_node(cache, m.last_device_node) - - # Evict one backed-up node - cache.evict(EvictParams(num_tokens=len(seqs[0]))) - - # Check mutual exclusion - overlap = cache.evictable_device_leaves & cache.evictable_host_leaves - self.assertEqual(len(overlap), 0) - cache.sanity_check() - def test_hicache_host_leaf_eviction(self): """Evicting a host leaf removes the node from the tree entirely.""" if self._skip_unsupported_hicache_test(): @@ -3713,44 +3406,6 @@ class UnifiedRadixCacheSuite: self.assertGreaterEqual(int(xfer.host_indices.numel()), sw) self.assertEqual(xfer.nodes_to_load, chain[-expected_pages:]) - def test_hicache_swa_host_independent_of_full(self): - """FULL host and SWA host are physically independent. - Freeing one component's host_value must not touch the other. - """ - if not self.cfg.has_swa: - self.skipTest("requires SWA") - - cache, allocator, req_to_token_pool = build_fixture(self.cfg) - seq = self._make_seq(1, 2) - self._insert(cache, allocator, req_to_token_pool, seq) - m = cache.match_prefix(MatchPrefixParams(key=RadixKey(array("q", seq)))) - node = m.last_device_node - - self._simulate_backup(cache, node) - cache.evict(EvictParams(num_tokens=len(seq))) - - cd_full = node.component_data[ComponentType.FULL] - cd_swa = node.component_data[ComponentType.SWA] - self.assertIsNotNone(cd_full.host_value) - self.assertIsNotNone(cd_swa.host_value) - self.assertIn(node, cache.evictable_host_leaves) - self.assertTrue(cache.host_lru_lists[ComponentType.SWA].in_list(node)) - - # Drop FULL host bookkeeping. SWA side must stay intact. - cache.evictable_host_leaves.discard(node) - cd_full.host_value = None - self.assertIsNotNone(cd_swa.host_value) - self.assertTrue(cache.host_lru_lists[ComponentType.SWA].in_list(node)) - self.assertNotIn(node, cache.evictable_host_leaves) - - # Drop SWA host bookkeeping. FULL side (already cleared) stays cleared. - cache.host_lru_lists[ComponentType.SWA].remove_node(node) - cd_swa.host_value = None - self.assertIsNone(cd_full.host_value) - self.assertIsNone(cd_swa.host_value) - self.assertFalse(cache.host_lru_lists[ComponentType.SWA].in_list(node)) - self.assertNotIn(node, cache.evictable_host_leaves) - def _swa_finalize_setup(self): """Build a SWA chain long enough to fill at least the window plus one extra page, and host-back every node so we can flip @@ -4325,22 +3980,6 @@ class UnifiedLRUListBoundedRefreshTest(CustomTestCase): cur = cur.lru_next[pt] return out - def test_bounded_refresh_stops_after_accumulated_meets_window(self): - root, [a, b, c, d] = self._build_chain([2, 2, 2, 2]) - lru = UnifiedLRUList(ComponentType.SWA, self.components) - for n in (a, b, c, d): - lru.insert_mru(n) - self.assertEqual(self._lru_order(lru), [d, c, b, a]) - - # window=5, page_size=1 implicit; nodes are size 2 each - # Walking up from D: visit D(acc=2<5) -> visit C(acc=4<5) -> visit - # B(acc=6>=5, refresh and stop). A is NOT touched. - lru.reset_node_and_window_ancestors_mru( - d, root, window_size=5, should_include=lambda _n: True - ) - # Expected MRU->LRU: D, C, B (refreshed in walk-up order), A (untouched) - self.assertEqual(self._lru_order(lru), [d, c, b, a]) - def test_bounded_refresh_skips_non_included(self): root, [a, b, c, d] = self._build_chain([2, 2, 2, 2]) lru = UnifiedLRUList(ComponentType.SWA, self.components) diff --git a/test/registered/unit/model_executor/test_cuda_graph_buffer_registry.py b/test/registered/unit/model_executor/test_cuda_graph_buffer_registry.py index 13f2bd53c..c8e35ca19 100644 --- a/test/registered/unit/model_executor/test_cuda_graph_buffer_registry.py +++ b/test/registered/unit/model_executor/test_cuda_graph_buffer_registry.py @@ -76,16 +76,6 @@ class TestGraphSlot(unittest.TestCase): axis="garbage", ) - def test_slice_for_before_buffer_alloc_raises(self): - slot = GraphSlot( - name="x", - shape_fn=lambda bs, mt: (bs,), - dtype=torch.int32, - axis="bs", - ) - with self.assertRaises(RuntimeError): - slot.slice_for(padded_bs=1, padded_num_tokens=1) - class TestRegistryRegister(unittest.TestCase): def test_register_allocates_zero_buffer(self): @@ -102,22 +92,6 @@ class TestRegistryRegister(unittest.TestCase): self.assertEqual(slot.buffer.dtype, torch.int64) self.assertTrue(torch.equal(slot.buffer, torch.zeros(16, dtype=torch.int64))) - def test_register_fill_sentinel_init(self): - r = _make_registry() - slot = r.register_slot( - GraphSlot( - name="seq_lens", - shape_fn=lambda bs, mt: (bs,), - dtype=torch.int32, - axis="bs", - padding_policy=PaddingPolicy.FILL_SENTINEL, - pad_value=7, - ) - ) - self.assertTrue( - torch.equal(slot.buffer, torch.full((8,), 7, dtype=torch.int32)) - ) - def test_register_duplicate_raises(self): r = _make_registry() r.register_slot( @@ -153,22 +127,6 @@ class TestRegistryRegister(unittest.TestCase): self.assertFalse(r.has_slot("off")) self.assertNotIn("off", r.slot_names()) - def test_cpu_device_override(self): - r = _make_registry() - slot = r.register_slot( - GraphSlot( - name="seq_lens_cpu", - shape_fn=lambda bs, mt: (bs,), - dtype=torch.int32, - axis="bs", - device=torch.device("cpu"), - padding_policy=PaddingPolicy.FILL_SENTINEL, - pad_value=11, - ) - ) - self.assertEqual(slot.buffer.device.type, "cpu") - self.assertEqual(int(slot.buffer[0].item()), 11) - class TestFillFromAndExtract(unittest.TestCase): """End-to-end exercise: register a representative slot set, fill from @@ -233,35 +191,6 @@ class TestFillFromAndExtract(unittest.TestCase): ) return r - def test_basic_fill_no_padding(self): - r = self._build_registry() - fb = _MiniForwardBatch( - batch_size=4, - input_ids=torch.arange(8, dtype=torch.int64), - req_pool_indices=torch.tensor([3, 1, 4, 2], dtype=torch.int64), - seq_lens=torch.tensor([10, 11, 12, 13], dtype=torch.int32), - out_cache_loc=torch.arange(8, dtype=torch.int64) + 100, - positions=torch.arange(8, dtype=torch.int64), - seq_lens_cpu=torch.tensor([10, 11, 12, 13], dtype=torch.int32), - ) - r.fill_from( - fb, - raw_bs=4, - padded_bs=4, - raw_num_tokens=8, - padded_num_tokens=8, - ) - self.assertTrue(torch.equal(r.get_slot("input_ids").buffer, fb.input_ids)) - self.assertTrue( - torch.equal(r.get_slot("req_pool_indices").buffer, fb.req_pool_indices) - ) - self.assertTrue(torch.equal(r.get_slot("seq_lens").buffer, fb.seq_lens)) - self.assertTrue( - torch.equal(r.get_slot("out_cache_loc").buffer, fb.out_cache_loc) - ) - self.assertTrue(torch.equal(r.get_slot("positions").buffer, fb.positions)) - self.assertTrue(torch.equal(r.get_slot("seq_lens_cpu").buffer, fb.seq_lens_cpu)) - def test_fill_with_padding_resets_zero_and_sentinel(self): r = self._build_registry() # Pre-poison the padded tail to a non-zero value so we can prove @@ -404,39 +333,6 @@ class TestFillFromAndExtract(unittest.TestCase): class TestMissingAndOptionalSlots(unittest.TestCase): - def test_missing_fb_attr_is_skipped(self): - r = _make_registry() - r.register_slot( - GraphSlot( - name="encoder_lens", - shape_fn=lambda bs, mt: (bs,), - dtype=torch.int32, - axis="bs", - padding_policy=PaddingPolicy.FILL_SENTINEL, - pad_value=0, - ) - ) - fb = _MiniForwardBatch( - batch_size=2, - input_ids=torch.arange(4, dtype=torch.int64), - encoder_lens=None, # FB doesn't carry this for this request. - ) - # Should NOT raise; encoder_lens buffer stays at the FILL_SENTINEL - # init value. - r.fill_from( - fb, - raw_bs=2, - padded_bs=4, - raw_num_tokens=4, - padded_num_tokens=8, - ) - self.assertTrue( - torch.equal( - r.get_slot("encoder_lens").buffer, - torch.zeros(8, dtype=torch.int32), - ) - ) - def test_extract_carries_none_for_absent_plain_slot(self): # A plain copy slot absent this iter (mrope on a non-multimodal batch) # must be carried as None, not exposed as the stale/zero buffer. @@ -468,6 +364,35 @@ class TestMissingAndOptionalSlots(unittest.TestCase): ) self.assertIsNone(fb_view.mrope_positions) + def test_plain_slot_with_missing_fb_attr_keeps_sentinel(self): + # A plain copy slot whose FB field is None must be skipped, leaving its + # buffer at the FILL_SENTINEL init value rather than raising. + r = _make_registry() + slot = r.register_slot( + GraphSlot( + "encoder_lens", + lambda bs, mt: (bs,), + torch.int32, + axis="bs", + padding_policy=PaddingPolicy.FILL_SENTINEL, + pad_value=0, + ) + ) + fb = _MiniForwardBatch( + batch_size=2, + input_ids=torch.arange(4, dtype=torch.int64), + encoder_lens=None, + ) + r.fill_from( + fb, + raw_bs=2, + padded_bs=4, + raw_num_tokens=4, + padded_num_tokens=8, + ) + # Buffer untouched by the copy; stays at the sentinel pad value. + self.assertTrue(torch.equal(slot.buffer, torch.zeros_like(slot.buffer))) + def test_extract_exposes_computed_slot_even_when_fb_field_none(self): # A computed slot (copy_from_fb=False) is always exposed, even when its # FB field is None — the None-skip carry applies only to plain copies. @@ -635,40 +560,6 @@ class TestSourceFnSlots(unittest.TestCase): r.fill_from(fb, raw_bs=3, padded_bs=8, raw_num_tokens=3, padded_num_tokens=16) self.assertTrue(torch.all(buf == 7)) # untouched - def test_side_input_source_via_fill_context(self): - r = _make_registry(max_bs=8, max_num_tokens=16) - r.register_slot( - GraphSlot( - name="pp_proxy_tensors.hidden_states", - shape_fn=lambda _bs, mt: (mt,), - dtype=torch.int32, - axis="none", - padding_policy=PaddingPolicy.KEEP_PAD, - source_fn=lambda fb, ctx: ( - None - if ctx.pp_proxy_tensors is None - else ctx.pp_proxy_tensors.tensors["hidden_states"] - ), - ) - ) - buf = r.get_slot("pp_proxy_tensors.hidden_states").buffer - buf.zero_() - fb = _MiniForwardBatch(batch_size=4) - pp = SimpleNamespace( - tensors={"hidden_states": torch.tensor([5, 6, 7, 8], dtype=torch.int32)} - ) - r.fill_from( - fb, - raw_bs=4, - padded_bs=8, - raw_num_tokens=4, - padded_num_tokens=16, - pp_proxy_tensors=pp, - ) - self.assertTrue( - torch.equal(buf[:4], torch.tensor([5, 6, 7, 8], dtype=torch.int32)) - ) - def test_extract_buffer_skips_dotted_slots(self): r = _make_registry(max_bs=8, max_num_tokens=16) r.register_slot( @@ -725,31 +616,6 @@ class TestPoolBackedAlloc(unittest.TestCase): r2.get_slot("ids").buffer.data_ptr(), ) - def test_same_size_shares_one_allocation(self): - a = self._reg(max_num_tokens=16, share_pool=True) - b = self._reg(max_num_tokens=16, share_pool=True) - a.register_slot(self._ids_slot("ids")) - b.register_slot(self._ids_slot("ids")) - # Identical (name, size, dtype, device) -> one shared allocation. - self.assertEqual( - a.get_slot("ids").buffer.data_ptr(), - b.get_slot("ids").buffer.data_ptr(), - ) - - def test_different_sizes_do_not_share(self): - big = self._reg(max_num_tokens=32, share_pool=True) - small = self._reg(max_num_tokens=16, share_pool=True) - big.register_slot(self._ids_slot("ids")) - small.register_slot(self._ids_slot("ids")) - # Different sizes -> different pool keys -> independent storage (no - # aliasing a smaller request onto a larger buffer). - self.assertEqual(tuple(small.get_slot("ids").buffer.shape), (16,)) - self.assertEqual(tuple(big.get_slot("ids").buffer.shape), (32,)) - self.assertNotEqual( - small.get_slot("ids").buffer.data_ptr(), - big.get_slot("ids").buffer.data_ptr(), - ) - def test_sharing_is_independent_of_registration_order(self): from sglang.srt.model_executor import input_buffers diff --git a/test/registered/unit/parser/test_conversation.py b/test/registered/unit/parser/test_conversation.py index e5986b077..da225fea7 100644 --- a/test/registered/unit/parser/test_conversation.py +++ b/test/registered/unit/parser/test_conversation.py @@ -119,19 +119,6 @@ class TestConversationGetPrompt(CustomTestCase): self.assertIn("[USER]Hello\n", prompt) self.assertTrue(prompt.endswith("[ASST]")) - def test_none_message_in_prompt(self): - """Test that None message produces role-only output (no content).""" - conv = Conversation( - name="test", - system_message="", - roles=("User", "Assistant"), - messages=[["User", "Q"], ["Assistant", None]], - sep_style=SeparatorStyle.ADD_COLON_SINGLE, - sep="\n", - ) - prompt = conv.get_prompt() - self.assertTrue(prompt.endswith("Assistant:")) - def test_empty_system_message(self): """Test that empty system message produces empty prefix for LLAMA3.""" conv = Conversation( @@ -189,22 +176,6 @@ class TestConversationGetPrompt(CustomTestCase): self.assertIn("[A]A", prompt) self.assertTrue(prompt.endswith("[U]")) - def test_llama2_with_system(self): - """Test LLAMA2 with system message.""" - conv = Conversation( - name="test", - system_message="<>\nBe helpful\n<>\n\n", - system_template="[INST] {system_message}", - roles=("[INST]", "[/INST]"), - messages=[["[INST]", "Hi"], ["[/INST]", None]], - sep_style=SeparatorStyle.LLAMA2, - sep=" ", - sep2=" ", - ) - prompt = conv.get_prompt() - self.assertIn("Be helpful", prompt) - self.assertIn("Hi ", prompt) - def test_llama2_without_system(self): """Test LLAMA2 without system message falls back to '[INST] ' prefix.""" conv = Conversation( @@ -570,23 +541,6 @@ class TestConversationGetPrompt(CustomTestCase): self.assertIn("USER: Describe this\n", prompt) self.assertIn("ASSISTANT: It shows a cat", prompt) - def test_mpt_with_tuple_message(self): - """Test MPT style extracts first element from tuple messages.""" - conv = Conversation( - name="test", - system_message="<|system|>", - roles=("<|user|>", "<|assistant|>"), - messages=[ - ["<|user|>", ("Hello", "extra1", "extra2")], - ["<|assistant|>", None], - ], - sep_style=SeparatorStyle.MPT, - sep="\n", - ) - prompt = conv.get_prompt() - self.assertIn("<|user|>Hello\n", prompt) - self.assertNotIn("extra1", prompt) - def test_invalid_sep_style_raises(self): """Test that an invalid SeparatorStyle raises ValueError.""" conv = Conversation( @@ -611,100 +565,6 @@ class TestConversationMethods(CustomTestCase): sep="\n", ) - def test_append_message(self): - """Test appending messages to conversation.""" - conv = self._make_conv() - conv.append_message("User", "Hello") - conv.append_message("Assistant", "Hi") - self.assertEqual(len(conv.messages), 2) - self.assertEqual(conv.messages[0], ["User", "Hello"]) - - def test_set_system_message(self): - """Test setting the system message.""" - conv = self._make_conv() - conv.set_system_message("Be helpful") - self.assertEqual(conv.system_message, "Be helpful") - - def test_update_last_message(self): - """Test updating the last message in-place.""" - conv = self._make_conv() - conv.append_message("User", "Q") - conv.append_message("Assistant", None) - conv.update_last_message("Answer") - self.assertEqual(conv.messages[-1][1], "Answer") - - def test_to_openai_api_messages_with_system(self): - """Test conversion to OpenAI format with system message.""" - conv = self._make_conv() - conv.system_message = "Be helpful" - conv.append_message("User", "Hello") - conv.append_message("Assistant", "Hi") - result = conv.to_openai_api_messages() - self.assertEqual(result[0], {"role": "system", "content": "Be helpful"}) - self.assertEqual(result[1], {"role": "user", "content": "Hello"}) - self.assertEqual(result[2], {"role": "assistant", "content": "Hi"}) - - def test_to_openai_api_messages_without_system(self): - """Test conversion to OpenAI format without system message.""" - conv = self._make_conv() - conv.append_message("User", "Hello") - result = conv.to_openai_api_messages() - self.assertEqual(len(result), 1) - self.assertEqual(result[0]["role"], "user") - - def test_to_openai_api_messages_skips_none_assistant(self): - """Test that None assistant message is omitted from OpenAI format.""" - conv = self._make_conv() - conv.append_message("User", "Hello") - conv.append_message("Assistant", None) - result = conv.to_openai_api_messages() - self.assertEqual(len(result), 1) # only user message - - def test_to_gradio_chatbot(self): - """Test conversion to Gradio chatbot format (user/assistant pairs).""" - conv = self._make_conv() - conv.append_message("User", "Q1") - conv.append_message("Assistant", "A1") - conv.append_message("User", "Q2") - conv.append_message("Assistant", "A2") - result = conv.to_gradio_chatbot() - self.assertEqual(len(result), 2) - self.assertEqual(result[0], ["Q1", "A1"]) - self.assertEqual(result[1], ["Q2", "A2"]) - - def test_to_gradio_chatbot_pending_response(self): - """Test Gradio format with pending assistant response (None).""" - conv = self._make_conv() - conv.append_message("User", "Q1") - conv.append_message("Assistant", None) - result = conv.to_gradio_chatbot() - self.assertEqual(result, [["Q1", None]]) - - def test_append_image(self): - """Test appending image data to conversation.""" - conv = self._make_conv() - conv.image_data = [] - conv.append_image("http://example.com/img.jpg", "auto") - self.assertEqual(len(conv.image_data), 1) - self.assertEqual(conv.image_data[0].url, "http://example.com/img.jpg") - self.assertEqual(conv.image_data[0].detail, "auto") - - def test_append_video(self): - """Test appending video data to conversation.""" - conv = self._make_conv() - conv.video_data = [] - conv.append_video("http://example.com/vid.mp4") - self.assertEqual(len(conv.video_data), 1) - self.assertEqual(conv.video_data[0], "http://example.com/vid.mp4") - - def test_append_audio(self): - """Test appending audio data to conversation.""" - conv = self._make_conv() - conv.audio_data = [] - conv.append_audio("http://example.com/audio.wav") - self.assertEqual(len(conv.audio_data), 1) - self.assertEqual(conv.audio_data[0], "http://example.com/audio.wav") - def test_copy_is_independent(self): """Test that copy() creates an independent conversation.""" conv = self._make_conv() @@ -714,15 +574,6 @@ class TestConversationMethods(CustomTestCase): self.assertEqual(len(conv.messages), 1) self.assertEqual(len(copied.messages), 2) - def test_dict_serialization(self): - """Test dict() returns expected keys.""" - conv = self._make_conv() - conv.append_message("User", "Hello") - d = conv.dict() - self.assertEqual(d["template_name"], "test") - self.assertIn("messages", d) - self.assertIn("roles", d) - class TestTemplateRegistry(CustomTestCase): def test_builtin_templates_exist(self): @@ -730,24 +581,6 @@ class TestTemplateRegistry(CustomTestCase): self.assertTrue(chat_template_exists("chatml")) self.assertTrue(chat_template_exists("llama-2")) - def test_unregistered_template_not_found(self): - """Test that non-existent template returns False.""" - self.assertFalse(chat_template_exists("_nonexistent_template_xyz")) - - def test_register_and_lookup(self): - """Test registering and looking up a custom template.""" - t = Conversation( - name="_test_conv_template", - roles=("A", "B"), - messages=[], - sep_style=SeparatorStyle.ADD_COLON_SINGLE, - sep="\n", - ) - register_conv_template(t) - self.assertTrue(chat_template_exists("_test_conv_template")) - # Cleanup - del chat_templates["_test_conv_template"] - def test_register_duplicate_raises(self): """Test that registering a duplicate name without override raises.""" with self.assertRaises(AssertionError): @@ -841,32 +674,6 @@ class TestGenerateEmbeddingConvs(CustomTestCase): self.assertIn("Hello world", convs[0].messages[0][1]) self.assertIsNone(convs[0].messages[1][1]) # assistant placeholder - def test_with_image(self): - """Test generating embedding conversations with image.""" - convs = generate_embedding_convs( - texts=["Describe"], - images=["http://example.com/img.jpg"], - videos=[None], - template_name="chatml", - ) - self.assertEqual(len(convs), 1) - msg = convs[0].messages[0][1] - self.assertIn("", msg) - self.assertIn("Describe", msg) - - def test_with_video(self): - """Test generating embedding conversations with video.""" - convs = generate_embedding_convs( - texts=["Describe"], - images=[None], - videos=["http://example.com/vid.mp4"], - template_name="chatml", - ) - self.assertEqual(len(convs), 1) - msg = convs[0].messages[0][1] - self.assertIn("