import asyncio import os import unittest import ray import torch from openai import BadRequestError from parameterized import parameterized_class from transformers import AutoTokenizer from tests.tools import ( RayUnittestBaseAsync, get_api_model_path, get_model_path, get_template_config, ) from trinity.common.config import Config from trinity.common.models import create_inference_models from trinity.common.models.model import ModelWrapper from trinity.common.models.utils import ( tokenize_and_mask_messages_default, tokenize_and_mask_messages_hf, ) from trinity.manager.synchronizer import Synchronizer DEBUG = False def print_debug(*args): if DEBUG: print(*args) async def prepare_engines(engines, auxiliary_engines): prepare_model_refs = [] for engine in engines: prepare_model_refs.append(engine.prepare.remote()) for engines in auxiliary_engines: for engine in engines: prepare_model_refs.append(engine.prepare.remote()) await asyncio.gather(*prepare_model_refs) # Qwen2.5 chat template with {% generation %} mark CHAT_TEMPLATE = r""" {%- if tools %} {{- '<|im_start|>system\n' }} {%- if messages[0]['role'] == 'system' %} {{- messages[0]['content'] }} {%- else %} {{- 'You are Qwen, created by Alibaba Cloud. You are a helpful assistant.' }} {%- endif %} {{- "\n\n# Tools\n\nYou may call one or more functions to assist with the user query.\n\nYou are provided with function signatures within XML tags:\n" }} {%- for tool in tools %} {{- "\n" }} {{- tool | tojson }} {%- endfor %} {{- "\n\n\nFor each function call, return a json object with function name and arguments within XML tags:\n\n{\"name\": , \"arguments\": }\n<|im_end|>\n" }} {%- else %} {%- if messages[0]['role'] == 'system' %} {{- '<|im_start|>system\n' + messages[0]['content'] + '<|im_end|>\n' }} {%- else %} {{- '<|im_start|>system\nYou are Qwen, created by Alibaba Cloud. You are a helpful assistant.<|im_end|>\n' }} {%- endif %} {%- endif %} {%- for message in messages %} {%- if (message.role == "user") or (message.role == "system" and not loop.first) %} {{- '<|im_start|>' + message.role + '\n' + message.content + '<|im_end|>' + '\n' }} {%- elif (message.role == "assistant" and not message.tool_calls) %} {{- '<|im_start|>' + message.role + '\n'}}{% generation %}{{- message.content + '<|im_end|>' + '\n' }}{% endgeneration %} {%- elif message.role == "assistant" %} {{- '<|im_start|>' + message.role + '\n'}}{% generation %} {%- if message.content %} {{- message.content }} {%- endif %} {%- for tool_call in message.tool_calls %} {%- if tool_call.function is defined %} {%- set tool_call = tool_call.function %} {%- endif %} {{- '\n{"name": "' }} {{- tool_call.name }} {{- '", "arguments": ' }} {{- tool_call.arguments | tojson }} {{- '}\n\n' }} {%- endfor %} {{- '<|im_end|>\n' }}{% endgeneration %} {%- elif message.role == "tool" %} {%- if (loop.index0 == 0) or (messages[loop.index0 - 1].role != "tool") %} {{- '<|im_start|>user' }} {%- endif %} {{- '\n\n' }} {{- message.content }} {{- '\n' }} {%- if loop.last or (messages[loop.index0 + 1].role != "tool") %} {{- '<|im_end|>\n' }} {%- endif %} {%- endif %} {%- endfor %} {%- if add_generation_prompt %} {{- '<|im_start|>assistant\n' }} {%- endif %} """ @parameterized_class( ( "tensor_parallel_size", "engine_num", "repeat_times", "enable_history", "use_async", ), [ (2, 2, 2, True, False), (1, 2, 1, False, True), (2, 1, 3, True, True), ], ) class ModelWrapperTest(RayUnittestBaseAsync): def setUp(self): # configure the model self.config = get_template_config() self.config.mode = "explore" self.config.model.model_path = get_model_path() self.config.model.custom_chat_template = CHAT_TEMPLATE self.config.explorer.rollout_model.engine_num = self.engine_num self.config.explorer.rollout_model.tensor_parallel_size = self.tensor_parallel_size self.config.algorithm.repeat_times = self.repeat_times self.config.explorer.rollout_model.enable_history = self.enable_history self.config.check_and_update() self.engines, self.auxiliary_engines = create_inference_models(self.config) self.model_wrapper = ModelWrapper( self.engines[0], engine_type="vllm", enable_history=self.enable_history ) async def test_generate(self): await prepare_engines(self.engines, self.auxiliary_engines) await self.model_wrapper.prepare() self.assertEqual(self.model_wrapper.model_path, self.config.model.model_path) prompts = ["Hello, world!", "Hello, my name is"] n = self.config.algorithm.repeat_times if self.use_async: generate_results = await self.model_wrapper.generate_async( prompts, n=n, temperature=1.0 ) else: generate_results = self.model_wrapper.generate(prompts, n=n, temperature=1.0) self.assertEqual(len(generate_results), len(prompts) * n) if self.config.explorer.rollout_model.enable_history: history_experiences = self.model_wrapper.extract_experience_from_history( clear_history=False ) self.assertEqual(len(history_experiences), len(generate_results)) for exp, history_exp in zip(generate_results, history_experiences): self.assertEqual(exp.response_text, history_exp.response_text) self.assertEqual(exp.tokens.tolist(), history_exp.tokens.tolist()) self.assertEqual(exp.prompt_length, history_exp.prompt_length) self.assertEqual(exp.logprobs.tolist(), history_exp.logprobs.tolist()) else: with self.assertRaises(ValueError): self.model_wrapper.extract_experience_from_history(clear_history=False) messages = [ {"role": "system", "content": "You are a helpful assistant."}, {"role": "user", "content": "What's the weather like today?"}, { "role": "assistant", "content": "I'm sorry, but as an AI language model, I don't have access to real-time weather information. To get accurate weather information for your location, you can check a weather website or app, or look outside if possible.", }, {"role": "user", "content": "OK, thanks!"}, ] if self.use_async: results = await self.model_wrapper.chat_async(messages, n=n, temperature=1.0) else: results = self.model_wrapper.chat(messages, n=n, temperature=1.0) self.assertEqual(len(results), n) if self.config.explorer.rollout_model.enable_history: history_experiences = self.model_wrapper.extract_experience_from_history() self.assertEqual(len(history_experiences) - len(generate_results), len(results)) for exp, history_exp in zip(results, history_experiences[len(generate_results) :]): self.assertEqual(exp.response_text, history_exp.response_text) self.assertEqual(exp.tokens.tolist(), history_exp.tokens.tolist()) self.assertEqual(exp.prompt_length, history_exp.prompt_length) self.assertEqual(exp.logprobs.tolist(), history_exp.logprobs.tolist()) for result in results: self.assertTrue(torch.any(result.logprobs != 0)) if self.use_async: logprobs = await self.model_wrapper.logprobs_async(results[0].tokens.tolist()) else: logprobs = self.model_wrapper.logprobs(results[0].tokens.tolist()) self.assertEqual(logprobs.shape[0], results[0].tokens.shape[0] - 1) if self.config.explorer.rollout_model.enable_history: history_experiences = self.model_wrapper.extract_experience_from_history() self.assertTrue(len(history_experiences) == 0) messages.append( { "role": "assistant", "content": results[0].response_text, } ) exp = self.model_wrapper.convert_messages_to_experience(messages) tokenizer = AutoTokenizer.from_pretrained(self.config.model.model_path) result_dict = tokenizer.apply_chat_template( messages, chat_template=CHAT_TEMPLATE, add_generation_prompt=False, padding=False, truncation=True, return_tensors="pt", add_special_tokens=False, return_assistant_tokens_mask=True, return_dict=True, ) prompt_length = torch.argmax(result_dict["assistant_masks"][0]).item() self.assertTrue( torch.equal(result_dict["assistant_masks"][0][prompt_length:], exp.action_mask) ) self.assertTrue(exp.logprobs.shape[0] == exp.tokens.shape[0] - prompt_length) self.assertTrue(torch.equal(result_dict["input_ids"][0], exp.tokens)) self.assertRaises(ValueError, self.model_wrapper.get_openai_client) if self.config.explorer.rollout_model.enable_history: history_experiences = self.model_wrapper.extract_experience_from_history() self.assertTrue(len(history_experiences) == 0) @parameterized_class( ( "max_model_len", "max_prompt_tokens", "max_response_tokens", ), [ (20, 19, None), (20, None, 1), (20, 5, 15), ], ) class TestModelLen(RayUnittestBaseAsync): def setUp(self): self.config = get_template_config() self.config.mode = "explore" self.config.model.model_path = get_model_path() self.config.model.max_model_len = self.max_model_len self.config.model.max_prompt_tokens = self.max_prompt_tokens self.config.model.max_response_tokens = self.max_response_tokens self.config.model.enable_prompt_truncation = True self.config.explorer.rollout_model.enable_openai_api = True self.config.check_and_update() self.engines, self.auxiliary_engines = create_inference_models(self.config) self.model_wrapper = ModelWrapper(self.engines[0], engine_type="vllm", enable_history=True) self.tokenizer = AutoTokenizer.from_pretrained(self.config.model.model_path) async def test_model_len(self): await prepare_engines(self.engines, self.auxiliary_engines) await self.model_wrapper.prepare() messages = [ {"role": "system", "content": "You are a helpful assistant."}, {"role": "user", "content": "What's the weather like today?"}, ] def _check_experience(exp): # check prompt content and length encoded_prompt = self.tokenizer.encode(exp.prompt_text, add_special_tokens=False) self.assertEqual(len(encoded_prompt), exp.prompt_length) self.assertLessEqual(exp.prompt_length, self.config.model.max_prompt_tokens) # check response content and length encoded_response = self.tokenizer.encode(exp.response_text, add_special_tokens=False) self.assertEqual(len(encoded_response), len(exp.tokens) - exp.prompt_length) self.assertLessEqual( len(exp.tokens) - exp.prompt_length, self.config.model.max_response_tokens ) # check full sequence self.assertLessEqual(len(exp.tokens), self.config.model.max_model_len) # For vllm engine, max_prompt_tokens and max_response_tokens work response = self.model_wrapper.chat(messages) self.assertEqual(len(response), 1) if self.max_prompt_tokens == 5: self.assertEqual(response[0].truncate_status, "prompt_truncated") _check_experience(response[0]) exps = self.model_wrapper.extract_experience_from_history() self.assertEqual(len(exps), 1) _check_experience(exps[0]) # For openai api, max_prompt_tokens and max_response_tokens do not work openai_client = self.model_wrapper.get_openai_client() model_id = openai_client.models.list().data[0].id with self.assertRaises(BadRequestError): # the prompt is longer than max_model_len openai_client.chat.completions.create(model=model_id, messages=messages, n=1) exps = self.model_wrapper.extract_experience_from_history() self.assertEqual(len(exps), 0) response = openai_client.chat.completions.create(model=model_id, messages=messages[1:], n=1) self.assertEqual(len(response.choices), 1) exps = self.model_wrapper.extract_experience_from_history() self.assertEqual(len(exps), 1) # only generate max_response_tokens tokens self.assertLessEqual( len(exps[0].tokens) - response.usage.prompt_tokens, self.config.model.max_response_tokens, ) class TestModelLenWithoutPromptTruncation(RayUnittestBaseAsync): def setUp(self): self.config = get_template_config() self.config.mode = "explore" self.config.model.model_path = get_model_path() self.config.model.max_model_len = 20 self.config.model.max_prompt_tokens = 1 self.config.model.max_response_tokens = None self.config.model.enable_prompt_truncation = False self.config.explorer.rollout_model.enable_openai_api = True self.config.check_and_update() self.engines, self.auxiliary_engines = create_inference_models(self.config) self.model_wrapper = ModelWrapper(self.engines[0], engine_type="vllm", enable_history=True) async def test_model_len(self): await prepare_engines(self.engines, self.auxiliary_engines) await self.model_wrapper.prepare() messages = [ {"role": "user", "content": "How are you?"}, ] # For vllm engine, max_prompt_tokens and max_response_tokens work response = self.model_wrapper.chat(messages) self.assertEqual(len(response), 1) self.assertLessEqual( len(response[0].tokens) - response[0].prompt_length, self.config.model.max_response_tokens, ) exps = self.model_wrapper.extract_experience_from_history() self.assertEqual(len(exps), 1) self.assertLessEqual( len(exps[0].tokens) - exps[0].prompt_length, self.config.model.max_response_tokens, ) # For openai api openai_client = self.model_wrapper.get_openai_client() model_id = openai_client.models.list().data[0].id response = openai_client.chat.completions.create(model=model_id, messages=messages, n=1) self.assertEqual(len(response.choices), 1) exps = self.model_wrapper.extract_experience_from_history() self.assertEqual(len(exps), 1) self.assertLessEqual( len(exps[0].tokens) - response.usage.prompt_tokens, self.config.model.max_response_tokens, ) class TestMessageProcess(RayUnittestBaseAsync): def setUp(self): self.config = get_template_config() self.config.mode = "explore" self.config.model.model_path = get_model_path() self.config.model.max_model_len = 100 self.config.model.max_prompt_tokens = 50 self.config.model.max_response_tokens = 50 self.config.model.enable_prompt_truncation = True self.config.explorer.rollout_model.enable_openai_api = True self.config.check_and_update() self.engines, self.auxiliary_engines = create_inference_models(self.config) self.model_wrapper = ModelWrapper(self.engines[0], engine_type="vllm", enable_history=True) async def test_truncation_status(self): """Test truncation status for multi-turn conversations.""" await prepare_engines(self.engines, self.auxiliary_engines) await self.model_wrapper.prepare() # Case: "prompt_truncated" messages = [ {"role": "system", "content": "You are helpful."}, {"role": "user", "content": "A very long prompt." * 20}, {"role": "assistant", "content": "OK"}, ] converted_experience = self.model_wrapper.convert_messages_to_experience( messages, ) self._check_experience(converted_experience, "prompt_truncated") # Case: No truncation messages = [ {"role": "system", "content": "You are helpful."}, {"role": "user", "content": "Tell me about weather."}, {"role": "assistant", "content": "OK"}, ] converted_experience = self.model_wrapper.convert_messages_to_experience( messages, ) self._check_experience(converted_experience, None) async def test_no_prompt_truncation(self): """Test truncation status for multi-turn conversations in workflow.""" self.config.model.enable_prompt_truncation = False self.config.check_and_update() await prepare_engines(self.engines, self.auxiliary_engines) await self.model_wrapper.prepare() # Case: No truncation messages = [ {"role": "system", "content": "You are helpful."}, {"role": "user", "content": "Tell me about weather."}, ] converted_experience = self.model_wrapper.convert_messages_to_experience(messages) self._check_experience(converted_experience, None) # Case: "response_truncated" messages = [ {"role": "system", "content": "You are helpful."}, {"role": "user", "content": "Tell me about weather."}, {"role": "assistant", "content": "A very long response" * 20}, ] converted_experience = self.model_wrapper.convert_messages_to_experience(messages) self._check_experience(converted_experience, "response_truncated") def _check_experience(self, exp, target_truncate_status): self.assertIsNotNone(exp) model_len = len(exp.tokens) prompt_length = exp.prompt_length self.assertEqual(exp.truncate_status, target_truncate_status) self.assertLessEqual(prompt_length, self.config.model.max_prompt_tokens) self.assertLessEqual(model_len, self.config.model.max_model_len) class TestAPIServer(RayUnittestBaseAsync): def setUp(self): self.config = get_template_config() self.config.mode = "explore" self.config.model.model_path = get_model_path() self.config.explorer.rollout_model.engine_type = "vllm" self.config.explorer.rollout_model.engine_num = 1 self.config.explorer.rollout_model.tensor_parallel_size = 1 self.config.explorer.rollout_model.chat_template = CHAT_TEMPLATE self.config.explorer.rollout_model.enable_openai_api = True self.config.check_and_update() self.engines, self.auxiliary_engines = create_inference_models(self.config) self.model_wrapper = ModelWrapper(self.engines[0], engine_type="vllm", enable_history=True) self.model_wrapper_no_history = ModelWrapper( self.engines[0], engine_type="vllm", enable_history=False ) async def test_api(self): await prepare_engines(self.engines, self.auxiliary_engines) await self.model_wrapper.prepare() await self.model_wrapper_no_history.prepare() openai_client = self.model_wrapper.get_openai_client() messages = [ {"role": "system", "content": "You are a helpful assistant."}, {"role": "user", "content": "What is your name?"}, ] model_id = openai_client.models.list().data[0].id response = openai_client.chat.completions.create(model=model_id, messages=messages, n=1) self.assertEqual(1, len(response.choices)) self.assertTrue(len(response.choices[0].message.content) > 0) response = openai_client.chat.completions.create( model=model_id, messages=messages, n=2, temperature=0.5, logprobs=True, top_logprobs=0, ) self.assertEqual(2, len(response.choices)) self.assertTrue(response.choices[0].logprobs is not None) self.assertEqual(0, len(response.choices[0].logprobs.content[2].top_logprobs)) # here we check the 3rd token logprob, because the first two tokens (``,`\n` usually have zero logprob) self.assertTrue(response.choices[0].logprobs.content[2].logprob < 0) self.assertTrue(hasattr(response, "prompt_token_ids")) self.assertTrue(len(response.prompt_token_ids) > 0) self.assertTrue(hasattr(response.choices[0], "token_ids")) self.assertTrue(len(response.choices[0].token_ids) > 0) exps = self.model_wrapper.extract_experience_from_history() self.assertEqual(len(exps), 3) response = openai_client.chat.completions.create( model=model_id, messages=messages, n=4, temperature=0.5, logprobs=True, top_logprobs=0, ) exps = self.model_wrapper.extract_experience_from_history() self.assertEqual(len(exps), 4) for exp in exps: self.assertTrue(len(exp.tokens) > 0) self.assertTrue(len(exp.logprobs) > 0) self.assertTrue(exp.prompt_length + len(exp.logprobs) == len(exp.tokens)) self.assertEqual(len(self.model_wrapper.extract_experience_from_history()), 0) response = openai_client.chat.completions.create( model=model_id, messages=messages, ) exps = self.model_wrapper.extract_experience_from_history() self.assertEqual(len(exps), 1) self.assertTrue(len(exps[0].tokens) > 0) self.assertTrue(len(exps[0].logprobs) > 0) self.assertTrue(exps[0].prompt_length + len(exps[0].logprobs) == len(exps[0].tokens)) response = openai_client.chat.completions.create( model=model_id, messages=messages, logprobs=False, ) exps = self.model_wrapper.extract_experience_from_history() self.assertEqual(len(exps), 1) self.assertTrue(len(exps[0].logprobs) == 0) response = self.model_wrapper_no_history.get_openai_client().chat.completions.create( model=model_id, messages=messages, n=2 ) self.assertEqual(2, len(response.choices)) self.assertTrue(hasattr(response.choices[0], "token_ids")) self.assertTrue(response.choices[0].token_ids is None) with self.assertRaises(ValueError): self.model_wrapper_no_history.extract_experience_from_history() self.assertEqual(len(self.model_wrapper_no_history.history), 0) SYSTEM_PROMPT = """ You are Qwen, created by Alibaba Cloud. You are a helpful assistant. You are walking on a frozen lake. FrozenLake Quick Guide Goal: Reach the goal (G). Player (P) and Goal (G) must overlap. Symbols: _ Frozen | O Hole | G Goal | P Player Rules: 1. Avoid falling into holes (O). 2. Frozen tiles are slippery, you may move perpendicular to your intended direction. Valid Action (separated by | ): Up | Down | Left | Right Rewards: Fall into hole: 0 Reach goal: +1.0 You will be provided the current observation, please decide on the next Action. You should show your thought process and then input the final action in ``` ```. You should only output the NEXT ACTION at each iteration in the ``` ```. For example, if you want to move up, you should output ```Up```. You should plan ahead and need to achieve it in minimum number of steps. You should be aware that frozen tiles can be slippery, but the chance is small and you should not overthink it. Please show your thinking process and put the final action in ``` ```. In every turn, the final action MUST be one of Up, Down, Left, Right. """ USER_PROMPT = """Current Observation (0): _ G _ _ _ _ P O O You have not achieved the goal, P has not reached G yet. Please give the next action. The maximum number of steps remaining is 10. """ class TestLogprobs(RayUnittestBaseAsync): def setUp(self): self.config = get_template_config() self.config.mode = "explore" self.config.model.model_path = get_model_path() self.config.explorer.rollout_model.engine_type = "vllm" self.config.explorer.rollout_model.engine_num = 1 self.config.explorer.rollout_model.tensor_parallel_size = 1 self.config.explorer.rollout_model.chat_template = CHAT_TEMPLATE self.config.explorer.rollout_model.enable_openai_api = True self.config.explorer.rollout_model.enable_log_requests = True self.config.check_and_update() self.engines, self.auxiliary_engines = create_inference_models(self.config) self.model_wrapper = ModelWrapper(self.engines[0], engine_type="vllm", enable_history=True) async def test_logprobs_api(self): await prepare_engines(self.engines, self.auxiliary_engines) await self.model_wrapper.prepare() messages = [ {"role": "system", "content": SYSTEM_PROMPT}, {"role": "user", "content": USER_PROMPT}, ] # Test openai api logprobs with different temperature self.model_client = self.model_wrapper.get_openai_async_client() _ = await self.model_client.chat.completions.create( model=self.model_client.model_path, messages=messages, n=1, temperature=1.0, logprobs=True, max_tokens=15, ) response_1 = self.model_wrapper.extract_experience_from_history()[0] _ = await self.model_client.chat.completions.create( model=self.model_client.model_path, messages=messages, n=1, temperature=0.8, logprobs=True, max_tokens=15, ) response_2 = self.model_wrapper.extract_experience_from_history()[0] self.assertTrue(response_1.logprobs is not None) self.assertTrue(len(response_1.logprobs) > 0) self.assertTrue(response_2.logprobs is not None) self.assertTrue(len(response_2.logprobs) > 0) logprobs_1 = self.model_wrapper.logprobs(response_1.tokens.tolist(), temperature=1.0) logprobs_2 = self.model_wrapper.logprobs(response_1.tokens.tolist(), temperature=0.8) logprobs_3 = self.model_wrapper.logprobs(response_2.tokens.tolist(), temperature=1.0) logprobs_4 = self.model_wrapper.logprobs(response_2.tokens.tolist(), temperature=0.8) self.assertEqual(logprobs_1.shape, logprobs_2.shape) self.assertEqual(logprobs_3.shape, logprobs_4.shape) self.assertFalse(torch.allclose(logprobs_1, logprobs_2, rtol=0.3, atol=1e-3)) self.assertFalse(torch.allclose(logprobs_3, logprobs_4, rtol=0.3, atol=1e-3)) logprobs_1_prompt = logprobs_1[: response_1.prompt_length - 1] logprobs_2_prompt = logprobs_2[: response_1.prompt_length - 1] logprobs_3_prompt = logprobs_3[: response_2.prompt_length - 1] logprobs_4_prompt = logprobs_4[: response_2.prompt_length - 1] self.assertEqual(logprobs_1_prompt.shape, logprobs_2_prompt.shape) self.assertFalse(torch.allclose(logprobs_1_prompt, logprobs_2_prompt, rtol=0.3, atol=1e-3)) self.assertFalse(torch.allclose(logprobs_3_prompt, logprobs_4_prompt, rtol=0.3, atol=1e-3)) self.assertTrue(torch.allclose(logprobs_1_prompt, logprobs_3_prompt, rtol=0.3, atol=1e-3)) self.assertTrue(torch.allclose(logprobs_2_prompt, logprobs_4_prompt, rtol=0.3, atol=1e-3)) logprobs_1_response = logprobs_1[response_1.prompt_length - 1 :] logprobs_2_response = logprobs_2[response_1.prompt_length - 1 :] logprobs_3_response = logprobs_3[response_2.prompt_length - 1 :] logprobs_4_response = logprobs_4[response_2.prompt_length - 1 :] self.assertEqual(logprobs_1_response.shape, logprobs_2_response.shape) self.assertEqual(logprobs_3_response.shape, logprobs_4_response.shape) self.assertEqual(logprobs_1_response.shape, logprobs_2_response.shape) self.assertEqual(response_1.logprobs.shape, logprobs_1_response.shape) self.assertTrue( torch.allclose(response_1.logprobs, logprobs_1_response, rtol=0.3, atol=1e-3) ) self.assertFalse( torch.allclose(response_1.logprobs, logprobs_2_response, rtol=0.3, atol=1e-3) ) self.assertTrue( torch.allclose(response_2.logprobs, logprobs_4_response, rtol=0.5, atol=1e-2) ) self.assertFalse( torch.allclose(response_2.logprobs, logprobs_3_response, rtol=0.3, atol=1e-3) ) # test vllm engine logprobs with different temperature response_1 = self.model_wrapper.chat( messages, n=1, temperature=1.0, logprobs=True, max_tokens=15 )[0] response_2 = self.model_wrapper.chat( messages, n=1, temperature=0.8, logprobs=True, max_tokens=15 )[0] self.assertTrue(response_1.logprobs is not None) self.assertTrue(len(response_1.logprobs) > 0) self.assertTrue(response_2.logprobs is not None) self.assertTrue(len(response_2.logprobs) > 0) logprobs_1 = self.model_wrapper.logprobs(response_1.tokens.tolist(), temperature=1.0) logprobs_2 = self.model_wrapper.logprobs(response_1.tokens.tolist(), temperature=0.8) logprobs_3 = self.model_wrapper.logprobs(response_2.tokens.tolist(), temperature=1.0) logprobs_4 = self.model_wrapper.logprobs(response_2.tokens.tolist(), temperature=0.8) self.assertEqual(logprobs_1.shape, logprobs_2.shape) self.assertEqual(logprobs_3.shape, logprobs_4.shape) self.assertFalse(torch.allclose(logprobs_1, logprobs_2, rtol=0.3, atol=1e-3)) self.assertFalse(torch.allclose(logprobs_3, logprobs_4, rtol=0.3, atol=1e-3)) logprobs_1_prompt = logprobs_1[: response_1.prompt_length - 1] logprobs_2_prompt = logprobs_2[: response_1.prompt_length - 1] logprobs_3_prompt = logprobs_3[: response_2.prompt_length - 1] logprobs_4_prompt = logprobs_4[: response_2.prompt_length - 1] self.assertEqual(logprobs_1_prompt.shape, logprobs_2_prompt.shape) self.assertFalse(torch.allclose(logprobs_1_prompt, logprobs_2_prompt, rtol=0.3, atol=1e-3)) self.assertFalse(torch.allclose(logprobs_3_prompt, logprobs_4_prompt, rtol=0.3, atol=1e-3)) self.assertTrue(torch.allclose(logprobs_1_prompt, logprobs_3_prompt, rtol=0.3, atol=1e-3)) self.assertTrue(torch.allclose(logprobs_2_prompt, logprobs_4_prompt, rtol=0.3, atol=1e-3)) logprobs_1_response = logprobs_1[response_1.prompt_length - 1 :] logprobs_2_response = logprobs_2[response_1.prompt_length - 1 :] logprobs_3_response = logprobs_3[response_2.prompt_length - 1 :] logprobs_4_response = logprobs_4[response_2.prompt_length - 1 :] self.assertEqual(logprobs_1_response.shape, logprobs_2_response.shape) self.assertEqual(logprobs_3_response.shape, logprobs_4_response.shape) self.assertEqual(logprobs_1_response.shape, logprobs_2_response.shape) self.assertEqual(response_1.logprobs.shape, logprobs_1_response.shape) self.assertTrue( torch.allclose(response_1.logprobs, logprobs_1_response, rtol=0.3, atol=1e-3) ) self.assertFalse( torch.allclose(response_1.logprobs, logprobs_2_response, rtol=0.3, atol=1e-3) ) self.assertTrue( torch.allclose(response_2.logprobs, logprobs_4_response, rtol=0.5, atol=1e-2) ) self.assertFalse( torch.allclose(response_2.logprobs, logprobs_3_response, rtol=0.3, atol=1e-3) ) # test openai api and vllm engine logprobs consistency await self.model_wrapper.clean_workflow_state() _ = await self.model_client.chat.completions.create( model=self.model_client.model_path, messages=messages, n=1, temperature=1.0, logprobs=0, max_tokens=1, ) response_openai_1 = self.model_wrapper.extract_experience_from_history()[0] _ = await self.model_client.chat.completions.create( model=self.model_client.model_path, messages=messages, n=1, temperature=0.8, logprobs=0, max_tokens=1, ) response_openai_2 = self.model_wrapper.extract_experience_from_history()[0] response_vllm_1 = self.model_wrapper.chat( messages, n=1, temperature=1.0, logprobs=0, max_tokens=1, )[0] response_vllm_2 = self.model_wrapper.chat( messages, n=1, temperature=0.8, logprobs=0, max_tokens=1, )[0] self.assertEqual(len(response_openai_1.tokens), len(response_vllm_1.tokens)) self.assertTrue( torch.allclose( response_openai_1.logprobs, response_vllm_1.logprobs, rtol=0.1, ) ) self.assertTrue( torch.allclose( response_openai_2.logprobs, response_vllm_2.logprobs, rtol=0.1, ) ) class TestAsyncAPIServer(RayUnittestBaseAsync): engine_type: str = "vllm" model_path: str = get_model_path() async def asyncSetUp(self): self.config = get_template_config() self._update_config() await self._setup_engines() def _update_config(self): self.config.mode = "explore" self.config.model.model_path = self.model_path self.config.explorer.rollout_model.engine_type = "vllm" self.config.explorer.rollout_model.engine_num = 1 self.config.explorer.rollout_model.tensor_parallel_size = 1 self.config.explorer.rollout_model.chat_template = CHAT_TEMPLATE self.config.explorer.rollout_model.enable_openai_api = True self.config.check_and_update() async def _setup_engines(self): self.engines, self.auxiliary_engines = create_inference_models(self.config) self.model_wrapper = ModelWrapper( self.engines[0], engine_type=self.engine_type, enable_history=True ) self.model_wrapper_no_history = ModelWrapper( self.engines[0], engine_type=self.engine_type, enable_history=False ) async def test_api_async(self): await prepare_engines(self.engines, self.auxiliary_engines) await self.model_wrapper.prepare() await self.model_wrapper_no_history.prepare() openai_client = self.model_wrapper.get_openai_async_client() messages = [ {"role": "system", "content": "You are a helpful assistant."}, {"role": "user", "content": "What is your name?"}, ] model_id = openai_client.model_path response = await openai_client.chat.completions.create( model=model_id, messages=messages, n=1 ) self.assertEqual(1, len(response.choices)) self.assertTrue(len(response.choices[0].message.content) > 0) response = await openai_client.chat.completions.create( model=model_id, messages=messages, n=2, temperature=0.5, logprobs=True, top_logprobs=0, ) self.assertEqual(2, len(response.choices)) self.assertTrue(response.choices[0].logprobs is not None) self.assertEqual(0, len(response.choices[0].logprobs.content[2].top_logprobs)) # here we check the 3rd token logprob, because the first two tokens (``,`\n` usually have zero logprob) if "Instruct" not in self.model_path: self.assertTrue(response.choices[0].logprobs.content[2].logprob < 0) self.assertTrue(hasattr(response, "prompt_token_ids")) self.assertTrue(len(response.prompt_token_ids) > 0) self.assertTrue(hasattr(response.choices[0], "token_ids")) self.assertTrue(len(response.choices[0].token_ids) > 0) exps = self.model_wrapper.extract_experience_from_history() self.assertEqual(len(exps), 3) response = await openai_client.chat.completions.create( model=model_id, messages=messages, n=4, temperature=0.5, logprobs=True, top_logprobs=0, ) exps = self.model_wrapper.extract_experience_from_history() self.assertEqual(len(exps), 4) for exp in exps: self.assertTrue(len(exp.tokens) > 0) self.assertTrue(len(exp.logprobs) > 0) self.assertTrue(exp.prompt_length + len(exp.logprobs) == len(exp.tokens)) self.assertEqual(len(self.model_wrapper.extract_experience_from_history()), 0) response = await openai_client.chat.completions.create( model=model_id, messages=messages, ) exps = self.model_wrapper.extract_experience_from_history() self.assertEqual(len(exps), 1) self.assertTrue(len(exps[0].tokens) > 0) self.assertTrue(len(exps[0].logprobs) > 0) self.assertTrue(exps[0].prompt_length + len(exps[0].logprobs) == len(exps[0].tokens)) response = await openai_client.chat.completions.create( model=model_id, messages=messages, logprobs=False, ) exps = self.model_wrapper.extract_experience_from_history() self.assertEqual(len(exps), 1) self.assertTrue(len(exps[0].logprobs) == 0) response = ( await self.model_wrapper_no_history.get_openai_async_client().chat.completions.create( model=model_id, messages=messages, n=2 ) ) self.assertEqual(2, len(response.choices)) self.assertTrue(hasattr(response.choices[0], "token_ids")) self.assertTrue(response.choices[0].token_ids is None) with self.assertRaises(ValueError): self.model_wrapper_no_history.extract_experience_from_history() self.assertEqual(len(self.model_wrapper_no_history.history), 0) @unittest.skipIf("TINKER_API_KEY" not in os.environ, "TINKER_API_KEY is not set") class TestTinkerAsyncAPIServer(TestAsyncAPIServer): engine_type: str = "tinker" model_path: str = "Qwen/Qwen3-4B-Instruct-2507" # llama model in Tinker does not support chat template def _update_config(self): self.config.model.tinker.enable = True self.config.algorithm.algorithm_type = "grpo" super()._update_config() async def _setup_engines(self): @ray.remote class FakeTrainer: def __init__(self, config: Config): self.config = config self.synchronizer = Synchronizer.get_actor(config) async def is_alive(self): return True fake_trainer = FakeTrainer.remote(self.config) await fake_trainer.__ray_ready__.remote() await super()._setup_engines() async def test_api_async(self): await super().test_api_async() class TestTokenizer(unittest.TestCase): def test_action_mask(self): messages = [ {"role": "system", "content": "You are a helpful assistant."}, {"role": "user", "content": "What's the weather like today?"}, { "role": "assistant", "content": "I'm sorry, but as an AI language model, I don't have access to real-time weather information. To get accurate weather information for your location, you can check a weather website or app, or look outside if possible.", }, {"role": "user", "content": "OK, thanks!"}, { "role": "assistant", "content": "You're welcome! If you have any other questions, feel free to ask.", }, ] tokenizer = AutoTokenizer.from_pretrained(get_model_path()) token_ids, action_mask, prompt_length = tokenize_and_mask_messages_default( tokenizer=tokenizer, messages=messages, chat_template=CHAT_TEMPLATE, ) token_ids_hf, action_mask_hf, prompt_length_hf = tokenize_and_mask_messages_hf( tokenizer=tokenizer, messages=messages, chat_template=CHAT_TEMPLATE, ) self.assertEqual(token_ids.shape, token_ids_hf.shape) self.assertEqual(action_mask.shape, action_mask_hf.shape) self.assertTrue(torch.equal(token_ids, token_ids_hf)) self.assertTrue(torch.equal(action_mask, action_mask_hf)) self.assertEqual(prompt_length, prompt_length_hf) def test_action_mask_with_tools(self): messages = [ { "role": "system", "content": "You are a helpful assistant with access to various tools. Use them when needed to help users.", }, {"role": "user", "content": "What's the weather like in Beijing today?"}, { "role": "assistant", "content": "Let me get the weather for you.", "tool_calls": [ { "id": "call_abc123", "type": "function", "function": { "name": "get_weather", "arguments": '{"location": "Beijing", "unit": "celsius"}', }, } ], }, { "role": "tool", "content": '{"temperature": 22, "condition": "sunny", "humidity": 45}', "tool_call_id": "call_abc123", }, { "role": "assistant", "content": "The weather in Beijing today is sunny with a temperature of 22°C and humidity at 45%. It's a pleasant day!", }, ] tools = [ { "type": "function", "function": { "name": "get_weather", "description": "Get the current weather in a given location", "parameters": { "type": "object", "properties": { "location": { "type": "string", "description": "The city and state, e.g. San Francisco, CA", }, "unit": { "type": "string", "enum": ["celsius", "fahrenheit"], "description": "The temperature unit", }, }, "required": ["location"], }, }, } ] tokenizer = AutoTokenizer.from_pretrained(get_model_path()) token_ids, action_mask, prompt_length = tokenize_and_mask_messages_default( tokenizer=tokenizer, messages=messages, tools=tools, chat_template=CHAT_TEMPLATE, ) token_ids_hf, action_mask_hf, prompt_length_hf = tokenize_and_mask_messages_hf( tokenizer=tokenizer, messages=messages, tools=tools, chat_template=CHAT_TEMPLATE, ) self.assertEqual(token_ids.shape, token_ids_hf.shape) self.assertEqual(action_mask.shape, action_mask_hf.shape) self.assertTrue(torch.equal(token_ids, token_ids_hf)) self.assertTrue(torch.equal(action_mask, action_mask_hf)) self.assertEqual(prompt_length, prompt_length_hf) @parameterized_class( ("enable_thinking", "reasoning_parser"), [ (True, "deepseek_r1"), (False, None), ], ) class TestAPIServerToolCall(RayUnittestBaseAsync): def setUp(self): self.config = get_template_config() self.config.mode = "explore" self.config.model.model_path = get_api_model_path() self.config.explorer.rollout_model.engine_type = "vllm" self.config.explorer.rollout_model.engine_num = 1 self.config.explorer.rollout_model.tensor_parallel_size = 1 self.config.explorer.rollout_model.chat_template = CHAT_TEMPLATE self.config.explorer.rollout_model.enable_openai_api = True # added for toolcalls self.config.explorer.rollout_model.enable_auto_tool_choice = True self.config.explorer.rollout_model.tool_call_parser = "hermes" self.config.explorer.rollout_model.enable_thinking = self.enable_thinking self.config.explorer.rollout_model.reasoning_parser = self.reasoning_parser self.config.check_and_update() self.engines, self.auxiliary_engines = create_inference_models(self.config) self.model_wrapper = ModelWrapper(self.engines[0], engine_type="vllm", enable_history=True) self.model_wrapper_no_history = ModelWrapper( self.engines[0], engine_type="vllm", enable_history=False ) async def test_api_tool_calls(self): """Tests the full conversation flow of a tool call via the OpenAI API. Note: This test require a model that supports tool calls and thinking mode, e.g. Qwen3-1.7B. """ import json import time await prepare_engines(self.engines, self.auxiliary_engines) await self.model_wrapper.prepare() await self.model_wrapper_no_history.prepare() tokenizer = AutoTokenizer.from_pretrained(get_api_model_path()) print_debug("\n\n" + "=" * 30 + " Running test_api_tool_calls " + "=" * 30) start_time = time.time() # --- Step 0: Get OpenAI Client --- print_debug(f"[{time.time() - start_time:.2f}s] Getting OpenAI client...") openai_client = self.model_wrapper.get_openai_client() model_id = openai_client.models.list().data[0].id print_debug( f"[{time.time() - start_time:.2f}s] Successfully got client. Model ID: {model_id}" ) # --- Step 1: Define Tools and Messages --- print_debug(f"[{time.time() - start_time:.2f}s] Defining tools and initial message...") tools = [ { "type": "function", "function": { "name": "get_current_weather", "description": "Get the current weather in a given location", "parameters": { "type": "object", "properties": { "location": { "type": "string", "description": "The city and state, e.g. San Francisco, CA", }, "unit": {"type": "string", "enum": ["celsius", "fahrenheit"]}, }, "required": ["location"], }, }, } ] messages = [{"role": "user", "content": "What's the weather like in Boston?"}] print_debug( f"[{time.time() - start_time:.2f}s] Initial user message: {messages[0]['content']}" ) print_debug("-" * 80) # --- Step 2: First API Call (Expecting a tool call) --- print_debug(f"[{time.time() - start_time:.2f}s] Making first API call to the model...") response = openai_client.chat.completions.create( model=model_id, messages=messages, tools=tools, tool_choice="auto", extra_body={ "repetition_penalty": 1.05, "chat_template_kwargs": { "enable_thinking": self.enable_thinking }, # default to True }, ) print_debug(f"[{time.time() - start_time:.2f}s] First API call completed.") # --- Step 3: Assert and Print the Tool Call Response --- print_debug(f"[{time.time() - start_time:.2f}s] Asserting response is a tool call...") self.assertEqual(len(response.choices), 1) choice = response.choices[0] print_debug(f" > Finish Reason: {choice.finish_reason}") self.assertEqual(choice.finish_reason, "tool_calls") if self.enable_thinking: self.assertIsNotNone(choice.message.reasoning_content) self.assertIsNotNone(choice.message.tool_calls) self.assertEqual(len(choice.message.tool_calls), 1) tool_call = choice.message.tool_calls[0] print_debug(f" > Tool Call ID: {tool_call.id}") print_debug(f" > Function Name: {tool_call.function.name}") print_debug(f" > Function Arguments: {tool_call.function.arguments}") self.assertEqual(tool_call.type, "function") self.assertEqual(tool_call.function.name, "get_current_weather") self.assertIn("Boston", tool_call.function.arguments) print_debug(f"[{time.time() - start_time:.2f}s] Assertions for tool call passed.") print_debug("-" * 80) # --- Step 4: Check Experience History --- print_debug(f"[{time.time() - start_time:.2f}s] Checking experience history...") exps = self.model_wrapper.extract_experience_from_history() self.assertEqual(len(exps), 1) # The response text in the experience should contain the tool call info print_debug(f" > Recorded experience response_text: {exps[0].response_text}") print_debug(f" > Recorded experience: {exps[0]}") print_debug(f" > message: {choice.message}") exp = exps[0] print_debug("\n" + "-" * 15 + " Decoding Experience Tokens " + "-" * 15) full_decoded_text = tokenizer.decode(exp.tokens, skip_special_tokens=False) print_debug( f" > Full Decoded Text ({len(exp.tokens)} tokens):\n---\n{full_decoded_text}\n---" ) prompt_length = exp.prompt_length prompt_tokens = exp.tokens[:prompt_length] response_tokens = exp.tokens[prompt_length:] prompt_decoded_text = tokenizer.decode(prompt_tokens, skip_special_tokens=False) response_decoded_text = tokenizer.decode(response_tokens, skip_special_tokens=False) print_debug( f" > Decoded Prompt Part ({len(prompt_tokens)} tokens):\n---\n{prompt_decoded_text}\n---" ) print_debug( f" > Decoded Response Part ({len(response_tokens)} tokens):\n---\n{response_decoded_text}\n---" ) action_mask = getattr(exp, "action_mask", None) if action_mask is not None: print_debug(f"\n > Action Mask (Length: {len(action_mask)}):") masked_tokens_info = [] for i, token_id in enumerate(response_tokens): token_text = tokenizer.decode([token_id]) mask_value = action_mask[i] if i < len(action_mask) else "N/A" masked_tokens_info.append(f"({repr(token_text)}, Mask: {mask_value})") print_debug(" " + " ".join(masked_tokens_info)) self.assertTrue( abs(len(action_mask) - len(response_tokens)) <= 1, f"Length of action_mask ({len(action_mask)}) does not match " f"length of response_tokens ({len(response_tokens)})", ) else: print_debug(" > Action Mask: Not found in experience.") print_debug("-" * 52 + "\n") # pass this part # self.assertIn("get_current_weather", exps[0].response_text) self.assertEqual( len(self.model_wrapper.extract_experience_from_history()), 0 ) # Verify cleared print_debug(f"[{time.time() - start_time:.2f}s] Experience history check passed.") print_debug("-" * 80) # --- Step 5: Second API Call (Providing tool result) --- print_debug( f"[{time.time() - start_time:.2f}s] Preparing for the second API call with tool result..." ) messages.append(response.choices[0].message) # Add assistant's tool call message # Mock the result of our tool tool_response_content = json.dumps( {"location": "Boston", "temperature": "72", "unit": "fahrenheit"} ) messages.append( { "role": "tool", "tool_call_id": tool_call.id, "content": tool_response_content, } ) print_debug(f"[{time.time() - start_time:.2f}s] Full message list for second call:") for msg in messages: print_debug(f" - {msg}") print_debug(f"[{time.time() - start_time:.2f}s] Making second API call...") second_response = openai_client.chat.completions.create( model=model_id, messages=messages, tools=tools, extra_body={ "repetition_penalty": 1.05, "chat_template_kwargs": { "enable_thinking": self.enable_thinking }, # default to True }, ) print_debug(f"[{time.time() - start_time:.2f}s] Second API call completed.") # --- Step 6: Assert and Print the Final Response --- print_debug( f"[{time.time() - start_time:.2f}s] Asserting final natural language response..." ) self.assertEqual(len(second_response.choices), 1) final_choice = second_response.choices[0] print_debug(f" > Final Finish Reason: {final_choice.finish_reason}") print_debug(f" > Final Message Content: {final_choice.message.content}") print_debug(f" > Final Message: {final_choice.message}") self.assertEqual(final_choice.finish_reason, "stop") # self.assertIsNone(final_choice.message.tool_calls) self.assertEqual(final_choice.message.tool_calls, []) self.assertIsNotNone(final_choice.message.content) # Check if the model used the information from the tool response self.assertIn("72", final_choice.message.content) self.assertIn("Boston", final_choice.message.content) print_debug(f"[{time.time() - start_time:.2f}s] Assertions for final response passed.") print_debug("-" * 80) # --- Step 7: Check Final Experience History --- print_debug(f"[{time.time() - start_time:.2f}s] Checking final experience history...") final_exps = self.model_wrapper.extract_experience_from_history() self.assertEqual(len(final_exps), 1) print_debug(f" > Final recorded experience response_text: {final_exps[0].response_text}") self.assertEqual(final_exps[0].response_text, final_choice.message.content) print_debug(f"[{time.time() - start_time:.2f}s] Final experience history check passed.") exp = final_exps[0] print_debug("\n" + "-" * 15 + " Decoding Experience Tokens " + "-" * 15) full_decoded_text = tokenizer.decode(exp.tokens, skip_special_tokens=False) print_debug( f" > Full Decoded Text ({len(exp.tokens)} tokens):\n---\n{full_decoded_text}\n---" ) prompt_length = exp.prompt_length prompt_tokens = exp.tokens[:prompt_length] response_tokens = exp.tokens[prompt_length:] prompt_decoded_text = tokenizer.decode(prompt_tokens, skip_special_tokens=False) response_decoded_text = tokenizer.decode(response_tokens, skip_special_tokens=False) print_debug( f" > Decoded Prompt Part ({len(prompt_tokens)} tokens):\n---\n{prompt_decoded_text}\n---" ) print_debug( f" > Decoded Response Part ({len(response_tokens)} tokens):\n---\n{response_decoded_text}\n---" ) action_mask = getattr(exp, "action_mask", None) if action_mask is not None: print_debug(f"\n > Action Mask (Length: {len(action_mask)}):") masked_tokens_info = [] for i, token_id in enumerate(response_tokens): token_text = tokenizer.decode([token_id]) mask_value = action_mask[i] if i < len(action_mask) else "N/A" masked_tokens_info.append(f"({repr(token_text)}, Mask: {mask_value})") print_debug(" " + " ".join(masked_tokens_info)) self.assertTrue( abs(len(action_mask) - len(response_tokens)) <= 1, f"Length of action_mask ({len(action_mask)}) does not match " f"length of response_tokens ({len(response_tokens)})", ) else: print_debug(" > Action Mask: Not found in experience.") total_time = time.time() - start_time print_debug( "\n" + "=" * 28 + f" test_api_tool_calls PASSED in {total_time:.2f}s " + "=" * 28 + "\n" ) class TestSuperLongGeneration(RayUnittestBaseAsync): def setUp(self): self.config = get_template_config() self.config.mode = "explore" self.config.model.model_path = get_model_path() self.config.model.max_model_len = 81920 self.config.model.max_prompt_tokens = 61440 self.config.model.max_response_tokens = 20480 self.config.model.rope_scaling = { "rope_type": "yarn", "factor": 2.0, "original_max_position_embeddings": 40960, } self.config.explorer.rollout_model.engine_type = "vllm" self.config.explorer.rollout_model.engine_num = 1 self.config.explorer.rollout_model.tensor_parallel_size = 1 self.config.explorer.rollout_model.chat_template = CHAT_TEMPLATE self.config.check_and_update() self.engines, self.auxiliary_engines = create_inference_models(self.config) self.model_wrapper = ModelWrapper(self.engines[0], engine_type="vllm", enable_history=True) async def test_generate(self): base_dir = os.path.dirname(__file__) target_dir = os.path.join(base_dir, "..", "..", "trinity", "trainer", "verl") with open(os.path.join(target_dir, "fsdp_workers.py")) as f: fsdp_code = f.read() with open(os.path.join(target_dir, "megatron_workers.py")) as f: megatron_code = f.read() target_dir = os.path.join(base_dir, "..", "..", "trinity", "common") with open(os.path.join(target_dir, "config.py")) as f: config_code = f.read() target_dir = os.path.join(base_dir, "..", "..", "trinity", "manager") with open(os.path.join(target_dir, "config_manager.py")) as f: config_manager_code = f.read() messages = [ {"role": "system", "content": "You are a helpful assistant."}, { "role": "user", "content": """# Please add comments and documentation for these following code, """ """make sure the code is well-structured and easy to read, """ """and the complete code must be shown, do not omit any parts.\n""" f"""## fsdp_workers.py\n{fsdp_code}\n""" f"""## megatron_workers.py\n{megatron_code}\n""" f"""## config.py\n{config_code}\n""" f"""## config_manager.py\n{config_manager_code}\n""", }, ] response = self.model_wrapper.chat(messages, n=1, temperature=0.7, logprobs=True)[0] self.assertGreater( response.prompt_length, 40960 ) # If not long enough, please add more files to prompt self.assertGreater(response.logprobs.shape[0], 1000) class TestTinkerAPI(RayUnittestBaseAsync): """Test the Tinker API integration with the vLLM engine.""" def setUp(self): self.config = get_template_config() self.config.mode = "explore" self.config.model.model_path = get_model_path() self.config.explorer.rollout_model.engine_type = "vllm" self.config.explorer.rollout_model.engine_num = 1 self.config.explorer.rollout_model.tensor_parallel_size = 1 self.config.explorer.rollout_model.chat_template = CHAT_TEMPLATE self.config.explorer.rollout_model.enable_openai_api = True self.config.explorer.rollout_model.enable_lora = True self.config.check_and_update() self.engines, self.auxiliary_engines = create_inference_models(self.config) self.model_wrapper = ModelWrapper(self.engines[0], engine_type="vllm", enable_history=True) async def test_tinker_api(self): from tinker import types from transformers import AutoTokenizer engine = self.engines[0] tokenizer = AutoTokenizer.from_pretrained(self.config.model.model_path) messages = [ {"role": "system", "content": "You are a helpful assistant."}, {"role": "user", "content": "What is your name?"}, ] result_dict = tokenizer.apply_chat_template( messages, chat_template=CHAT_TEMPLATE, add_generation_prompt=False, padding=False, truncation=True, return_tensors="pt", add_special_tokens=False, return_assistant_tokens_mask=True, return_dict=True, ) prompt = types.ModelInput.from_ints( result_dict["input_ids"][0].tolist(), ) # sample api without prompt logprobs num_samples = 4 response = await engine.sample.remote( prompt=prompt, num_samples=num_samples, sampling_params=types.SamplingParams(temperature=0.7), # no limit on length ) self.assertEqual(len(response.sequences), num_samples) for sequence in response.sequences: self.assertEqual(len(sequence.tokens), len(sequence.logprobs)) self.assertEqual(sequence.stop_reason, "stop") self.assertIsNone(response.prompt_logprobs) self.assertIsNone(response.topk_prompt_logprobs) # sample api with prompt logprobs num_samples = 2 topk_prompt_logprobs = 3 response = await engine.sample.remote( prompt=prompt, num_samples=num_samples, sampling_params=types.SamplingParams(temperature=0.7, max_tokens=8), include_prompt_logprobs=True, topk_prompt_logprobs=topk_prompt_logprobs, ) self.assertEqual(len(response.sequences), num_samples) for sequence in response.sequences: self.assertEqual(len(sequence.tokens), len(sequence.logprobs)) self.assertEqual(sequence.stop_reason, "length") self.assertEqual(len(response.prompt_logprobs), len(prompt.to_ints())) self.assertIsNone(response.prompt_logprobs[0]) self.assertEqual(len(response.topk_prompt_logprobs), len(prompt.to_ints())) self.assertIsNone(response.topk_prompt_logprobs[0]) for topk_logprobs in response.topk_prompt_logprobs[1:]: self.assertIsNotNone(topk_logprobs) self.assertEqual(len(topk_logprobs), topk_prompt_logprobs) # compute_logprob api response = await engine.sample.remote( prompt=prompt, num_samples=1, sampling_params=types.SamplingParams(max_tokens=1), include_prompt_logprobs=True, ) self.assertEqual(len(response.sequences), 1) self.assertEqual(response.sequences[0].stop_reason, "length") self.assertEqual(len(prompt.to_ints()), len(response.prompt_logprobs)) self.assertIsNone(response.topk_prompt_logprobs) # test add remove lora from vllm.lora.request import LoRARequest # create a dummy lora adapter with all zero weights lora_path_1 = os.path.join(self.config.checkpoint_job_dir, "adapter_1") lora_path_2 = os.path.join(self.config.checkpoint_job_dir, "adapter_2") _create_adapter(self.config.model.model_path, lora_path_1, "adapter_1") _create_adapter(self.config.model.model_path, lora_path_2, "adapter_2") lora_1 = LoRARequest( lora_name="test_adapter_1", lora_int_id=1, lora_path=os.path.join(lora_path_1, "adapter_1"), ) lora_2 = LoRARequest( lora_name="test_adapter_2", lora_int_id=2, lora_path=os.path.join(lora_path_2, "adapter_2"), ) response = await engine.sample.remote( prompt=prompt, num_samples=1, sampling_params=types.SamplingParams(max_tokens=1), include_prompt_logprobs=True, lora_request=lora_1, ) ids = await engine.list_lora_adapters.remote() self.assertEqual(ids, [1]) self.assertEqual(len(response.sequences), 1) self.assertEqual(response.sequences[0].stop_reason, "length") self.assertEqual(len(prompt.to_ints()), len(response.prompt_logprobs)) self.assertIsNone(response.topk_prompt_logprobs) response = await engine.sample.remote( prompt=prompt, num_samples=1, sampling_params=types.SamplingParams(max_tokens=1), include_prompt_logprobs=True, lora_request=lora_2, ) self.assertEqual(len(response.sequences), 1) self.assertEqual(response.sequences[0].stop_reason, "length") self.assertEqual(len(prompt.to_ints()), len(response.prompt_logprobs)) self.assertIsNone(response.topk_prompt_logprobs) await engine.remove_lora_adapter.remote(lora_id=1) await engine.remove_lora_adapter.remote(lora_id=2) ids = await engine.list_lora_adapters.remote() self.assertEqual(ids, []) def _create_adapter(model_path: str, lora_path: str, name: str): from peft import LoraConfig, get_peft_model from transformers import AutoModelForCausalLM model = AutoModelForCausalLM.from_pretrained( model_path, device_map="cpu", ) lora_config = LoraConfig( r=8, lora_alpha=8, target_modules=["gate_proj", "up_proj", "down_proj"], lora_dropout=0.1, ) lora_model = get_peft_model(model, lora_config, adapter_name=name) lora_model.save_pretrained(lora_path)