change print to logger.log

Author: binary-husky

ajet/backbone/main_verl.py

@@ -22,6 +22,7 @@
 import hydra
 import ray
 from beast_logger import print_dict
+from loguru import logger
 from omegaconf import OmegaConf
 from verl.trainer.ppo.reward import load_reward_manager
 from verl.utils.device import is_cuda_available
@@ -112,7 +113,7 @@ def run(self, config):
         from omegaconf import OmegaConf
         from verl.utils.fs import copy_to_local
 
-        print(f"TaskRunner hostname: {socket.gethostname()}, PID: {os.getpid()}")
+        logger.info(f"TaskRunner hostname: {socket.gethostname()}, PID: {os.getpid()}")
         pprint(OmegaConf.to_container(config, resolve=True))
         OmegaConf.resolve(config)
 
@@ -148,8 +149,6 @@ def run(self, config):
                 from verl.workers.fsdp_workers import CriticWorker
             elif use_legacy_worker_impl == "disable":
                 from verl.workers.roles import CriticWorker
-
-                print("Using new worker implementation")
             else:
                 raise ValueError(f"Invalid use_legacy_worker_impl: {use_legacy_worker_impl}")
 

ajet/backbone/main_vllm.py

@@ -10,6 +10,7 @@
 from ajet.utils.launch_utils import set_loguru_default_color
 from ajet.schema.logprob import TokenAndProb
 from ajet.utils.core_env_vars import get_runtime_env
+from loguru import logger
 
 set_loguru_default_color()
 
@@ -116,12 +117,11 @@ def run(config):
         config.ajet.task_reader,
     )
     tasks = task_reader.get_validation_tasks()
-    print(tasks[:2])
+    logger.info(tasks[:n_task])
     ctx_tracker = parallel_env.rollout(
         tasks=tasks[:n_task], mode="sample", epoch="1"
     )  # "sample" or "validate"
     _ = parallel_env.to_dataproto(ctx_tracker)
-    print("Generated batch output")
 
 
 @hydra.main(
@@ -133,7 +133,6 @@ def main(config):
     from omegaconf import OmegaConf
 
     OmegaConf.resolve(config)
-    print("*" * 20)
 
     runtime_env = get_runtime_env()
     os.environ.update(runtime_env["env_vars"])
@@ -147,12 +146,12 @@ def companion_launch():
 
         from ajet.utils.smart_daemon import LaunchCommandWhenAbsent
 
-        print("Launching companion process for async LLM server...")
+        logger.info("Launching companion process for async LLM server...")
         model_path = config.ajet.model.path
         tensor_parallel_size = config.ajet.debug.debug_tensor_parallel_size
         n_avail_gpus = torch.cuda.device_count()
         if tensor_parallel_size > n_avail_gpus:
-            print(
+            logger.info(
                 f"Warning: tensor_parallel_size {tensor_parallel_size} is greater than available GPUs {n_avail_gpus}. Setting tensor_parallel_size to {n_avail_gpus}."
             )
             tensor_parallel_size = n_avail_gpus

ajet/backbone/trainer_verl.py

@@ -293,15 +293,15 @@ def check_mutually_exclusive(mbs, mbs_per_gpu, name: str):
             )
 
         if self.config.algorithm.use_kl_in_reward and config.actor_rollout_ref.actor.use_kl_loss:
-            print("NOTICE: You have both enabled in-reward kl and kl loss.")
+            logger.warning("NOTICE: You have both enabled in-reward kl and kl loss.")
 
         # critic
         if self.use_critic:
             critic_config = omega_conf_to_dataclass(config.critic)
             critic_config.validate(n_gpus, config.ajet.data.train_batch_size)
 
         if config.data.get("val_batch_size", None) is not None:
-            print(
+            logger.warning(
                 "WARNING: val_batch_size is deprecated."
                 + " Validation datasets are sent to inference engines as a whole batch,"
                 + " which will schedule the memory themselves."
@@ -313,7 +313,7 @@ def check_mutually_exclusive(mbs, mbs_per_gpu, name: str):
                 config.ajet.rollout.temperature > 0
             ), "validation gen temperature should be greater than 0 when enabling do_sample"
 
-        print("[validate_config] All configuration checks passed successfully!")
+        logger.success("[validate_config] All configuration checks passed successfully!")
 
     def init_workers(self):
         """Initialize distributed training workers using Ray backend.
@@ -787,7 +787,7 @@ def fit(self):  # noqa: C901
                         or esi_close_to_expiration
                     ):
                         if esi_close_to_expiration:
-                            print("Force saving checkpoint: ESI instance expiration approaching.")
+                            logger.info("Force saving checkpoint: ESI instance expiration approaching.")
                         with marked_timer("save_checkpoint", timing_raw, color="green"):
                             self._save_checkpoint()
 

ajet/context_tracker/basic_tracker.py

@@ -1,8 +1,8 @@
+import torch
 import copy
 from collections import defaultdict
 from typing import List, Tuple
-
-import torch
+from loguru import logger
 
 from ajet.context_tracker.base_tracker import (
     BaseTracker,
@@ -231,7 +231,7 @@ def group_tokenize_multi_group(self):
             sample_arr += [sample]
 
         if len(sample_arr) > max_num_group:
-            print(f"Warning: allow {max_num_group} groups, but got {len(sample_arr)} groups")
+            logger.warning(f"Warning: allow {max_num_group} groups, but got {len(sample_arr)} groups")
             import random
 
             sample_arr = random.sample(sample_arr, max_num_group)  # preserve max_num_group groups

ajet/default_config/ajet_default.yaml

@@ -7,7 +7,7 @@ ajet:
 
 
   # the experimental reverse proxy feature that allows `tuner.as_oai_baseurl_apikey` feature
-  enable_experimental_reverse_proxy: True
+  enable_experimental_reverse_proxy: False
 
   model:
     # which model should be trained

ajet/schema/convertion.py

@@ -46,7 +46,9 @@ def convert_llm_proxy_response_to_oai_response(llm_proxy_response):
         usage=usage,
     )
 
-# copied from AgentScope's DashScopeChatModule
+
+
+# modified from AgentScope's DashScopeChatModule
 def convert_llm_proxy_response_to_agentscope_response(
     message,
     structured_model: Type[BaseModel] | None = None,
@@ -105,91 +107,3 @@ def convert_llm_proxy_response_to_agentscope_response(
 
     return parsed_response
 
-
-
-def test_convert_llm_proxy_response_to_oai_response():
-    """Test the conversion from llm_proxy_response to OpenAI ChatCompletion format."""
-
-    from ajet.schema.logprob import TokenAndProb
-    # Test case 1: Basic response with content only
-    llm_proxy_response_basic = {
-        "role": "assistant",
-        "request_id": "req-123456",
-        "content": "Hello, how can I help you today?",
-        "tool_calls": None,
-        "tokens": [
-            TokenAndProb(
-                token_id=123,
-                logprob=-0.5,
-                decoded_string="Hello",
-            ),
-            TokenAndProb(
-                token_id=456,
-                logprob=-0.3,
-                decoded_string=",",
-            ),
-        ],
-    }
-
-    result = convert_llm_proxy_response_to_oai_response(llm_proxy_response_basic)
-
-    assert result.id == "req-123456"
-    assert result.object == "chat.completion"
-    assert len(result.choices) == 1
-    assert result.choices[0].message.role == "assistant"
-    assert result.choices[0].message.content == "Hello, how can I help you today?"
-    assert result.choices[0].message.tool_calls is None
-    assert result.choices[0].finish_reason == "stop"
-    assert result.usage is not None
-    assert result.usage.completion_tokens == 2
-    assert result.usage.total_tokens == 2
-
-    print("✓ Test case 1 passed: Basic response with content")
-
-    # Test case 2: Response with tool calls
-    llm_proxy_response_with_tools = {
-        "role": "assistant",
-        "request_id": "req-789012",
-        "content": "",
-        "tool_calls": [
-            {
-                "id": "call_abc123",
-                "type": "function",
-                "function": {
-                    "name": "get_weather",
-                    "arguments": '{"location": "San Francisco"}'
-                }
-            }
-        ],
-        "tokens": [],
-    }
-
-    result2 = convert_llm_proxy_response_to_oai_response(llm_proxy_response_with_tools)
-
-    assert result2.id == "req-789012"
-    assert result2.choices[0].message.content == ""
-    assert result2.choices[0].message.tool_calls is not None
-    assert len(result2.choices[0].message.tool_calls) == 1
-    assert result2.usage is None  # No tokens provided
-
-    print("✓ Test case 2 passed: Response with tool calls")
-
-    # Test case 3: Minimal response with defaults
-    llm_proxy_response_minimal = {
-        "content": "Test response"
-    }
-
-    result3 = convert_llm_proxy_response_to_oai_response(llm_proxy_response_minimal)
-
-    assert result3.id == "chatcmpl-default"
-    assert result3.choices[0].message.role == "assistant"
-    assert result3.choices[0].message.content == "Test response"
-    assert result3.model == "unknown"
-
-    print("✓ Test case 3 passed: Minimal response with defaults")
-
-    print("\n✅ All tests passed!")
-
-
-if __name__ == "__main__":
-    test_convert_llm_proxy_response_to_oai_response()

ajet/schema/trajectory.py

@@ -2,6 +2,7 @@
 
 import numpy as np
 from pydantic import BaseModel, Field
+from loguru import logger
 
 
 class Reward(BaseModel):
@@ -31,7 +32,6 @@ def performance_reward(self):
         # this reward is NOT used in training
         if (self.step_reward_arr is not None) and len(self.step_reward_arr) > 0:
             res = np.mean(self.step_reward_arr)
-            # print(f"Performance reward computed as mean of step_reward_arr: {res}")
             return res
         else:
             return self.raw_reward
@@ -146,13 +146,13 @@ def truncate_output_ids(self) -> None:
 
         if len(self.response_ids) > self.max_response_len:
             truncate_any = True
-            print(
+            logger.warning(
                 "-------------------------------------------------------------------------------------------------------"
             )
-            print(
+            logger.warning(
                 f"Warning: response_ids length {len(self.response_ids)} exceeds max_response_len {self.max_response_len}, truncating."
             )
-            print(
+            logger.warning(
                 "-------------------------------------------------------------------------------------------------------"
             )
             self.response_ids = self.response_ids[: self.max_response_len]

ajet/task_reader/tracing_reader/__init__.py

@@ -31,7 +31,6 @@ def __init__(
 
         super().__init__(reader_config)
         # config patch
-        # print("*********", config, "**********")
         self.reader_config = reader_config.feedback_tracing
 
         logger.info(

ajet/task_rollout/dashscope_llm_bridge.py

@@ -71,8 +71,8 @@ def external_llm_chat_fn(messages, sampling_params_override={}, request_id=""):
                 return {"role": message["role"], "content": message["content"]}
             except Exception as e:
                 logger.bind(exception=True).exception(f"Error calling alien llm: {e}")
+                logger.warning(f"Error calling alien llm: {e}, retrying...")
                 time.sleep(5)
-                print(f"Error calling alien llm: {e}, retrying...")
         raise RuntimeError(f"Failed to get response from alien llm after {max_try} attempts")
 
     return external_llm_chat_fn

ajet/tuner.py

@@ -101,6 +101,7 @@ def as_oai_baseurl_apikey(
             ```
         """
 
+        assert self.config.ajet.enable_experimental_reverse_proxy, "Please enable `ajet.enable_experimental_reverse_proxy` in yaml config to use `as_oai_baseurl_apikey` feature."
         baseurl_apikey_model = OpenaiClientBaseUrlTuner(
             config=self.config,
             context_tracker=self.context_tracker,