diff --git a/justfile b/justfile
index 0a82d616..d8da7690 100644
--- a/justfile
+++ b/justfile
@@ -1,3 +1,5 @@
+export NIX_CONFIG := "extra-experimental-features = nix-command flakes"
+
fmt:
nix fmt
diff --git a/src/exo/master/api.py b/src/exo/master/api.py
index 30f87e2a..c24ab75b 100644
--- a/src/exo/master/api.py
+++ b/src/exo/master/api.py
@@ -13,12 +13,6 @@ from hypercorn.asyncio import serve # pyright: ignore[reportUnknownVariableType
from hypercorn.config import Config
from hypercorn.typing import ASGIFramework
from loguru import logger
-from openai_harmony import ( # pyright: ignore[reportMissingTypeStubs]
- HarmonyEncodingName,
- Role,
- StreamableParser,
- load_harmony_encoding,
-)
from exo.master.placement import place_instance as get_instance_placements
from exo.shared.apply import apply
@@ -67,8 +61,6 @@ from exo.utils.channels import Receiver, Sender, channel
from exo.utils.dashboard_path import find_dashboard
from exo.utils.event_buffer import OrderedBuffer
-encoding = load_harmony_encoding(HarmonyEncodingName.HARMONY_GPT_OSS)
-
def chunk_to_response(
chunk: TokenChunk, command_id: CommandId
@@ -381,35 +373,8 @@ class API:
instance_id=instance_id,
)
- async def _process_gpt_oss(self, token_chunks: Receiver[TokenChunk]):
- stream = StreamableParser(encoding, role=Role.ASSISTANT)
- thinking = False
-
- async for chunk in token_chunks:
- stream.process(chunk.token_id)
-
- delta = stream.last_content_delta
- ch = stream.current_channel
-
- if ch == "analysis" and not thinking:
- thinking = True
- yield chunk.model_copy(update={"text": ""})
-
- if ch != "analysis" and thinking:
- thinking = False
- yield chunk.model_copy(update={"text": ""})
-
- if delta:
- yield chunk.model_copy(update={"text": delta})
-
- if chunk.finish_reason is not None:
- if thinking:
- yield chunk.model_copy(update={"text": ""})
- yield chunk
- break
-
async def _chat_chunk_stream(
- self, command_id: CommandId, parse_gpt_oss: bool
+ self, command_id: CommandId
) -> AsyncGenerator[TokenChunk, None]:
"""Yield `TokenChunk`s for a given command until completion."""
@@ -417,16 +382,10 @@ class API:
self._chat_completion_queues[command_id], recv = channel[TokenChunk]()
with recv as token_chunks:
- if parse_gpt_oss:
- async for chunk in self._process_gpt_oss(token_chunks):
- yield chunk
- if chunk.finish_reason is not None:
- break
- else:
- async for chunk in token_chunks:
- yield chunk
- if chunk.finish_reason is not None:
- break
+ async for chunk in token_chunks:
+ yield chunk
+ if chunk.finish_reason is not None:
+ break
except anyio.get_cancelled_exc_class():
# TODO: TaskCancelled
@@ -442,11 +401,11 @@ class API:
del self._chat_completion_queues[command_id]
async def _generate_chat_stream(
- self, command_id: CommandId, parse_gpt_oss: bool
+ self, command_id: CommandId
) -> AsyncGenerator[str, None]:
"""Generate chat completion stream as JSON strings."""
- async for chunk in self._chat_chunk_stream(command_id, parse_gpt_oss):
+ async for chunk in self._chat_chunk_stream(command_id):
chunk_response: ChatCompletionResponse = chunk_to_response(
chunk, command_id
)
@@ -458,7 +417,7 @@ class API:
yield "data: [DONE]\n\n"
async def _collect_chat_completion(
- self, command_id: CommandId, parse_gpt_oss: bool
+ self, command_id: CommandId
) -> ChatCompletionResponse:
"""Collect all token chunks for a chat completion and return a single response."""
@@ -466,7 +425,7 @@ class API:
model: str | None = None
finish_reason: FinishReason | None = None
- async for chunk in self._chat_chunk_stream(command_id, parse_gpt_oss):
+ async for chunk in self._chat_chunk_stream(command_id):
if model is None:
model = chunk.model
@@ -495,7 +454,7 @@ class API:
)
async def _collect_chat_completion_with_stats(
- self, command_id: CommandId, parse_gpt_oss: bool
+ self, command_id: CommandId
) -> BenchChatCompletionResponse:
text_parts: list[str] = []
model: str | None = None
@@ -503,7 +462,7 @@ class API:
stats: GenerationStats | None = None
- async for chunk in self._chat_chunk_stream(command_id, parse_gpt_oss):
+ async for chunk in self._chat_chunk_stream(command_id):
if model is None:
model = chunk.model
@@ -544,8 +503,6 @@ class API:
"""Handle chat completions, supporting both streaming and non-streaming responses."""
model_meta = await resolve_model_meta(payload.model)
payload.model = model_meta.model_id
- parse_gpt_oss = "gpt-oss" in model_meta.model_id.lower()
- logger.info(f"{parse_gpt_oss=}")
if not any(
instance.shard_assignments.model_id == payload.model
@@ -562,17 +519,16 @@ class API:
await self._send(command)
if payload.stream:
return StreamingResponse(
- self._generate_chat_stream(command.command_id, parse_gpt_oss),
+ self._generate_chat_stream(command.command_id),
media_type="text/event-stream",
)
- return await self._collect_chat_completion(command.command_id, parse_gpt_oss)
+ return await self._collect_chat_completion(command.command_id)
async def bench_chat_completions(
self, payload: BenchChatCompletionTaskParams
) -> BenchChatCompletionResponse:
model_meta = await resolve_model_meta(payload.model)
- parse_gpt_oss = "gpt-oss" in model_meta.model_id.lower()
payload.model = model_meta.model_id
if not any(
@@ -589,10 +545,7 @@ class API:
command = ChatCompletion(request_params=payload)
await self._send(command)
- response = await self._collect_chat_completion_with_stats(
- command.command_id,
- parse_gpt_oss,
- )
+ response = await self._collect_chat_completion_with_stats(command.command_id)
return response
def _calculate_total_available_memory(self) -> Memory:
diff --git a/src/exo/worker/engines/mlx/utils_mlx.py b/src/exo/worker/engines/mlx/utils_mlx.py
index 64807eb6..81fc68aa 100644
--- a/src/exo/worker/engines/mlx/utils_mlx.py
+++ b/src/exo/worker/engines/mlx/utils_mlx.py
@@ -366,6 +366,8 @@ def apply_chat_template(
tools=chat_task_data.tools,
)
+ logger.info(prompt)
+
return prompt
diff --git a/src/exo/worker/runner/runner.py b/src/exo/worker/runner/runner.py
index 029bbf4c..fae4acf8 100644
--- a/src/exo/worker/runner/runner.py
+++ b/src/exo/worker/runner/runner.py
@@ -1,6 +1,15 @@
import time
+from collections.abc import Generator
+from functools import cache
import mlx.core as mx
+from mlx_lm.models.gpt_oss import Model as GptOssModel
+from openai_harmony import ( # pyright: ignore[reportMissingTypeStubs]
+ HarmonyEncodingName,
+ Role,
+ StreamableParser,
+ load_harmony_encoding,
+)
from exo.shared.types.api import ChatCompletionMessageText
from exo.shared.types.chunks import TokenChunk
@@ -153,11 +162,19 @@ def main(
_check_for_debug_prompts(task_params.messages[0].content)
# Generate responses using the actual MLX generation
- for response in mlx_generate(
+ mlx_generator = mlx_generate(
model=model,
tokenizer=tokenizer,
task=task_params,
- ):
+ )
+
+ # GPT-OSS specific parsing to match other model formats.
+ if isinstance(model, GptOssModel):
+ mlx_generator = parse_gpt_oss(mlx_generator)
+
+ # TODO: Add tool call parser here
+
+ for response in mlx_generator:
match response:
case GenerationResponse():
if shard_metadata.device_rank == 0:
@@ -207,6 +224,43 @@ def main(
break
+@cache
+def get_gpt_oss_encoding():
+ encoding = load_harmony_encoding(HarmonyEncodingName.HARMONY_GPT_OSS)
+ return encoding
+
+
+def parse_gpt_oss(
+ responses: Generator[GenerationResponse],
+) -> Generator[GenerationResponse]:
+ encoding = get_gpt_oss_encoding()
+ stream = StreamableParser(encoding, role=Role.ASSISTANT)
+ thinking = False
+
+ for response in responses:
+ stream.process(response.token)
+
+ delta = stream.last_content_delta
+ ch = stream.current_channel
+
+ if ch == "analysis" and not thinking:
+ thinking = True
+ yield response.model_copy(update={"text": ""})
+
+ if ch != "analysis" and thinking:
+ thinking = False
+ yield response.model_copy(update={"text": ""})
+
+ if delta:
+ yield response.model_copy(update={"text": delta})
+
+ if response.finish_reason is not None:
+ if thinking:
+ yield response.model_copy(update={"text": ""})
+ yield response
+ break
+
+
EXO_RUNNER_MUST_FAIL = "EXO RUNNER MUST FAIL"
EXO_RUNNER_MUST_OOM = "EXO RUNNER MUST OOM"
EXO_RUNNER_MUST_TIMEOUT = "EXO RUNNER MUST TIMEOUT"