update init

Author: Yunnglin

src/twinkle/server/__init__.py

@@ -1,15 +1,7 @@
 # Copyright (c) ModelScope Contributors. All rights reserved.
 from .launcher import ServerLauncher, launch_server
-from .twinkle.model import build_model_app
-from .twinkle.processor import build_processor_app
-from .twinkle.sampler import build_sampler_app
-from .twinkle.server import build_server_app
 
 __all__ = [
-    'build_model_app',
-    'build_processor_app',
-    'build_sampler_app',
-    'build_server_app',
     'ServerLauncher',
     'launch_server',
 ]

src/twinkle/server/launcher.py

@@ -101,7 +101,7 @@ def _get_builders(self) -> dict[str, Callable]:
                 'build_sampler_app': build_sampler_app,
             }
         else:  # twinkle
-            from twinkle.server import build_model_app, build_processor_app, build_sampler_app, build_server_app
+            from twinkle.server.twinkle import build_model_app, build_processor_app, build_sampler_app, build_server_app
             self._builders = {
                 'build_server_app': build_server_app,
                 'build_model_app': build_model_app,

src/twinkle/server/tinker/__init__.py

@@ -1,15 +1,18 @@
 # Copyright (c) ModelScope Contributors. All rights reserved.
+import sys
+from typing import TYPE_CHECKING
 
-from ..utils import wrap_builder_with_device_group_env
-from .model import build_model_app as _build_model_app
-from .sampler import build_sampler_app as _build_sampler_app
-from .server import build_server_app
+from twinkle.utils.import_utils import _LazyModule
 
-build_model_app = wrap_builder_with_device_group_env(_build_model_app)
-build_sampler_app = wrap_builder_with_device_group_env(_build_sampler_app)
+_import_structure = {
+    'model': ['build_model_app'],
+    'sampler': ['build_sampler_app'],
+    'server': ['build_server_app'],
+}
 
-__all__ = [
-    'build_model_app',
-    'build_sampler_app',
-    'build_server_app',
-]
+if TYPE_CHECKING:
+    from .model import build_model_app
+    from .sampler import build_sampler_app
+    from .server import build_server_app
+else:
+    sys.modules[__name__] = _LazyModule(__name__, __file__, _import_structure, module_spec=__spec__)

src/twinkle/server/tinker/common/datum.py

@@ -3,7 +3,6 @@
 import numpy as np
 from collections import defaultdict
 from tinker import types
-from typing import List, Union
 
 from twinkle.data_format.input_feature import InputFeature
 from twinkle.template import Template
@@ -92,6 +91,8 @@ def input_feature_to_datum(input_feature: InputFeature) -> types.Datum:
         labels_raw = input_feature['labels']
         if isinstance(labels_raw, np.ndarray):
             labels_arr = labels_raw.astype(np.int64)
+        elif isinstance(labels_raw, list):
+            labels_arr = np.asarray(labels_raw, dtype=np.int64)
         else:
             labels_arr = np.asarray(labels_raw.cpu(), dtype=np.int64)
 

src/twinkle/server/tinker/model.py

@@ -24,6 +24,7 @@
 from twinkle.server.utils.task_queue import TaskQueueConfig, TaskQueueMixin
 from twinkle.server.utils.validation import get_token_from_request, verify_request_token
 from twinkle.utils.logger import get_logger
+from ..utils import wrap_builder_with_device_group_env
 from .common.io_utils import create_checkpoint_manager, create_training_run_manager
 from .common.router import StickyLoraRequestRouter
 
@@ -653,3 +654,6 @@ async def _do_load():
 
     return ModelManagement.options(**deploy_options).bind(nproc_per_node, device_group, device_mesh, use_megatron,
                                                           queue_config, **kwargs)
+
+
+build_model_app = wrap_builder_with_device_group_env(build_model_app)

src/twinkle/server/tinker/sampler.py

@@ -23,6 +23,7 @@
 from twinkle.server.utils.task_queue import TaskQueueConfig, TaskQueueMixin
 from twinkle.server.utils.validation import get_token_from_request, verify_request_token
 from twinkle.utils.logger import get_logger
+from ..utils import wrap_builder_with_device_group_env
 from .common.io_utils import create_checkpoint_manager
 
 logger = get_logger()
@@ -245,3 +246,6 @@ async def _do_sample():
 
     return SamplerManagement.options(**deploy_options).bind(nproc_per_node, device_group, device_mesh, sampler_type,
                                                             engine_args, queue_config, **kwargs)
+
+
+build_sampler_app = wrap_builder_with_device_group_env(build_sampler_app)

src/twinkle/server/tinker/server.py

@@ -13,7 +13,6 @@
 
 import asyncio
 import httpx
-import logging
 import os
 from fastapi import FastAPI, HTTPException, Request, Response
 from ray import serve
@@ -24,9 +23,10 @@
 from twinkle.server.utils.state import get_server_state
 from twinkle.server.utils.task_queue import QueueState
 from twinkle.server.utils.validation import get_token_from_request, verify_request_token
+from twinkle.utils.logger import get_logger
 from .common.io_utils import create_checkpoint_manager, create_training_run_manager
 
-logger = logging.getLogger(__name__)
+logger = get_logger()
 
 
 def build_server_app(deploy_options: dict[str, Any],

src/twinkle/server/twinkle/__init__.py

@@ -1,5 +1,20 @@
 # Copyright (c) ModelScope Contributors. All rights reserved.
-from .model import build_model_app
-from .processor import build_processor_app
-from .sampler import build_sampler_app
-from .server import build_server_app
+import sys
+from typing import TYPE_CHECKING
+
+from twinkle.utils.import_utils import _LazyModule
+
+_import_structure = {
+    'model': ['build_model_app'],
+    'processor': ['build_processor_app'],
+    'sampler': ['build_sampler_app'],
+    'server': ['build_server_app'],
+}
+
+if TYPE_CHECKING:
+    from .model import build_model_app
+    from .processor import build_processor_app
+    from .sampler import build_sampler_app
+    from .server import build_server_app
+else:
+    sys.modules[__name__] = _LazyModule(__name__, __file__, _import_structure, module_spec=__spec__)