Merge branch 'fix_moe' of https://github.com/modelscope/twinkle into fix_moe

Author: Yunnglin

cookbook/client/tinker/self_congnition.py

@@ -8,7 +8,7 @@
 # The server must be running first (see server.py and server_config.yaml).
 import numpy as np
 import os
-from modelscope import AutoTokenizer
+from tqdm import tqdm
 from tinker import types
 from twinkle_client import init_tinker_compat_client
 from twinkle.data_format import Message, Trajectory
@@ -125,5 +125,5 @@ def eval():
 
 
 if __name__ == '__main__':
-    # train()   # Uncomment to run training
-    eval()      # Run evaluation / inference
+    train()   # Uncomment to run training
+    # eval()      # Run evaluation / inference

src/twinkle/model/megatron/megatron.py

@@ -848,13 +848,13 @@ def load(self, name: str, output_dir: Optional[str] = None, **kwargs):
         Args:
             name: Checkpoint name or HuggingFace Hub model id.
             output_dir: Parent directory that contains the checkpoint folder.
-                If None **and** ``resume`` is False, downloads from Hub.
-            resume: If True, restore optimizer, lr_scheduler and RNG state
+                If None **and** ``load_optimizer`` is False, downloads from Hub.
+            load_optimizer: If True, restore optimizer, lr_scheduler and RNG state
                 from the mcore sub-checkpoint for training resumption.
             **kwargs: Additional arguments (``adapter_name``, ``no_load_optim``,
                 ``no_load_rng``, etc.).
         """
-        resume = kwargs.pop('resume', False)
+        resume = kwargs.pop('load_optimizer', False)
         if output_dir is None and not resume:
             # Load from hub
             token = kwargs.pop('token', None)

src/twinkle/server/utils/state.py

@@ -31,7 +31,8 @@ class ServerState:
     def __init__(
             self,
             expiration_timeout: float = 86400.0,  # 24 hours in seconds
-            cleanup_interval: float = 3600.0) -> None:  # 1 hour in seconds
+            cleanup_interval: float = 3600.0,
+            **kwargs) -> None:  # 1 hour in seconds
         # Session tracking
         self.sessions: dict[str, dict[str, Any]] = {}
         # Model registration

src/twinkle/server/utils/task_queue.py

@@ -222,7 +222,7 @@ async def _queue_worker(self) -> None:
         Selection policy: round-robin across queue keys. If a task is rate-limited
         at execution time, it is requeued and the worker tries other queues.
         """
-        print('[TaskQueue] Worker started')
+        logger.debug('[TaskQueue] Worker started')
         while True:
             try:
                 # Wait until there is at least one queue with a task
@@ -470,7 +470,7 @@ async def schedule_task(
         if self._event_loop is None:
             self._event_loop = asyncio.get_running_loop()
 
-        print(
+        logger.debug(
             f'[TaskQueue] Scheduling task {request_id}, rps_limit={self._task_queue_config.rps_limit}, enabled={self._task_queue_config.enabled}'  # noqa: E501
         )
 
@@ -487,7 +487,7 @@ async def schedule_task(
 
         # 5. Put task in queue and update status
         q = self._task_queues[queue_key]
-        print(
+        logger.debug(
             f'[TaskQueue] Adding task {request_id} to queue key={queue_key} (current size: {q.qsize()}) type={task_type}'  # noqa: E501
         )
         await q.put(
@@ -502,7 +502,7 @@ async def schedule_task(
             ))
         self.state.store_future_status(
             request_id, TaskStatus.QUEUED.value, model_id, queue_state=QueueState.ACTIVE.value)
-        print(f'[TaskQueue] Task {request_id} queued, new queue size: {q.qsize()} key={queue_key}')
+        logger.debug(f'[TaskQueue] Task {request_id} queued, new queue size: {q.qsize()} key={queue_key}')
 
         self._new_task_event.set()
 
@@ -567,4 +567,4 @@ async def shutdown_task_queue(self) -> None:
         self._task_queues.clear()
         self._queue_order.clear()
 
-        print('[TaskQueue] Task queue shutdown complete')
+        logger.debug('[TaskQueue] Task queue shutdown complete')

src/twinkle_client/__init__.py

@@ -21,7 +21,7 @@ def init_tinker_compat_client(base_url: str | None = None, api_key: str | None =
     # Apply patch to bypass tinker:// prefix validation
     patch_tinker()
 
-    if api_key is None:
+    if not api_key:
         api_key = get_api_key()
 
     if base_url and not base_url.startswith(('http://', 'https://')):