diff --git a/pilot/model/model_adapter.py b/pilot/model/model_adapter.py index b8f4f0982..57e3cf251 100644 --- a/pilot/model/model_adapter.py +++ b/pilot/model/model_adapter.py @@ -51,6 +51,12 @@ _OLD_MODELS = [ class LLMModelAdaper: """New Adapter for DB-GPT LLM models""" + def use_fast_tokenizer(self) -> bool: + """Whether use a [fast Rust-based tokenizer](https://huggingface.co/docs/tokenizers/index) if it is supported + for a given model. + """ + return False + def model_type(self) -> str: return ModelType.HF @@ -169,6 +175,9 @@ class OldLLMModelAdaperWrapper(LLMModelAdaper): self._adapter = adapter self._chat_adapter = chat_adapter + def use_fast_tokenizer(self) -> bool: + return self._adapter.use_fast_tokenizer() + def model_type(self) -> str: return self._adapter.model_type() @@ -200,6 +209,9 @@ class FastChatLLMModelAdaperWrapper(LLMModelAdaper): def __init__(self, adapter: "BaseModelAdapter") -> None: self._adapter = adapter + def use_fast_tokenizer(self) -> bool: + return self._adapter.use_fast_tokenizer + def load(self, model_path: str, from_pretrained_kwargs: dict): return self._adapter.load_model(model_path, from_pretrained_kwargs) diff --git a/pilot/utils/tracer/span_storage.py b/pilot/utils/tracer/span_storage.py index 8967f9ee5..914aa0126 100644 --- a/pilot/utils/tracer/span_storage.py +++ b/pilot/utils/tracer/span_storage.py @@ -34,7 +34,9 @@ class FileSpanStorage(SpanStorage): self.flush_signal_queue = queue.Queue() if not os.path.exists(filename): - with open(filename, "w") as _: + # New file if not exist + os.makedirs(os.path.dirname(filename), exist_ok=True) + with open(filename, "a"): pass self.flush_thread = threading.Thread(target=self._flush_to_file, daemon=True) self.flush_thread.start() diff --git a/pilot/utils/tracer/tests/test_base.py b/pilot/utils/tracer/tests/test_base.py index 4a8061cb1..445120e8a 100644 --- a/pilot/utils/tracer/tests/test_base.py +++ b/pilot/utils/tracer/tests/test_base.py @@ -1,7 +1,7 @@ from typing import Dict from pilot.component import SystemApp -from pilot.utils.tracer import Span, SpanStorage, Tracer +from pilot.utils.tracer import Span, SpanType, SpanStorage, Tracer # Mock implementations @@ -31,7 +31,9 @@ class MockTracer(Tracer): self._new_uuid() if parent_span_id is None else parent_span_id.split(":")[0] ) span_id = f"{trace_id}:{self._new_uuid()}" - span = Span(trace_id, span_id, parent_span_id, operation_name, metadata) + span = Span( + trace_id, span_id, SpanType.BASE, parent_span_id, operation_name, metadata + ) self.current_span = span return span @@ -50,7 +52,14 @@ class MockTracer(Tracer): def test_span_creation(): - span = Span("trace_id", "span_id", "parent_span_id", "operation", {"key": "value"}) + span = Span( + "trace_id", + "span_id", + SpanType.BASE, + "parent_span_id", + "operation", + {"key": "value"}, + ) assert span.trace_id == "trace_id" assert span.span_id == "span_id" assert span.parent_span_id == "parent_span_id" diff --git a/pilot/utils/tracer/tests/test_span_storage.py b/pilot/utils/tracer/tests/test_span_storage.py index 5ad518b15..0c63992a6 100644 --- a/pilot/utils/tracer/tests/test_span_storage.py +++ b/pilot/utils/tracer/tests/test_span_storage.py @@ -5,7 +5,7 @@ import json import tempfile import time -from pilot.utils.tracer import SpanStorage, FileSpanStorage, Span +from pilot.utils.tracer import SpanStorage, FileSpanStorage, Span, SpanType @pytest.fixture @@ -44,7 +44,7 @@ def read_spans_from_file(filename): "storage", [{"batch_size": 1, "flush_interval": 5}], indirect=True ) def test_write_span(storage: SpanStorage): - span = Span("1", "a", "b", "op1") + span = Span("1", "a", SpanType.BASE, "b", "op1") storage.append_span(span) time.sleep(0.1) @@ -57,8 +57,8 @@ def test_write_span(storage: SpanStorage): "storage", [{"batch_size": 1, "flush_interval": 5}], indirect=True ) def test_incremental_write(storage: SpanStorage): - span1 = Span("1", "a", "b", "op1") - span2 = Span("2", "c", "d", "op2") + span1 = Span("1", "a", SpanType.BASE, "b", "op1") + span2 = Span("2", "c", SpanType.BASE, "d", "op2") storage.append_span(span1) storage.append_span(span2) @@ -72,7 +72,7 @@ def test_incremental_write(storage: SpanStorage): "storage", [{"batch_size": 2, "flush_interval": 5}], indirect=True ) def test_sync_and_async_append(storage: SpanStorage): - span = Span("1", "a", "b", "op1") + span = Span("1", "a", SpanType.BASE, "b", "op1") storage.append_span(span) @@ -88,7 +88,7 @@ def test_sync_and_async_append(storage: SpanStorage): @pytest.mark.asyncio async def test_flush_policy(storage: SpanStorage): - span = Span("1", "a", "b", "op1") + span = Span("1", "a", SpanType.BASE, "b", "op1") for _ in range(storage.batch_size - 1): storage.append_span(span) @@ -108,8 +108,8 @@ async def test_flush_policy(storage: SpanStorage): "storage", [{"batch_size": 2, "file_does_not_exist": True}], indirect=True ) def test_non_existent_file(storage: SpanStorage): - span = Span("1", "a", "b", "op1") - span2 = Span("2", "c", "d", "op2") + span = Span("1", "a", SpanType.BASE, "b", "op1") + span2 = Span("2", "c", SpanType.BASE, "d", "op2") storage.append_span(span) time.sleep(0.1)