[object Object]

← back to Exo

Fix tests broken in recent commits (#1239)

a354aaa3e5873b7ec938623adcd12facf6542ab6 · 2026-01-21 18:32:49 +0000 · rltakashige

We'll have good CI soon...

## Test Plan

### Automated Testing
Wroks

Files touched

Diff

commit a354aaa3e5873b7ec938623adcd12facf6542ab6
Author: rltakashige <rl.takashige@gmail.com>
Date:   Wed Jan 21 18:32:49 2026 +0000

    Fix tests broken in recent commits (#1239)
    
    We'll have good CI soon...
    
    ## Test Plan
    
    ### Automated Testing
    Wroks
---
 src/exo/worker/tests/unittests/test_mlx/conftest.py        | 12 ++++++++----
 src/exo/worker/tests/unittests/test_mlx/test_tokenizers.py |  6 ++----
 2 files changed, 10 insertions(+), 8 deletions(-)

diff --git a/src/exo/worker/tests/unittests/test_mlx/conftest.py b/src/exo/worker/tests/unittests/test_mlx/conftest.py
index 87ecce7d..77b3fc11 100644
--- a/src/exo/worker/tests/unittests/test_mlx/conftest.py
+++ b/src/exo/worker/tests/unittests/test_mlx/conftest.py
@@ -19,7 +19,7 @@ from exo.shared.types.tasks import ChatCompletionTaskParams
 from exo.shared.types.worker.shards import PipelineShardMetadata, TensorShardMetadata
 from exo.worker.engines.mlx import Model
 from exo.worker.engines.mlx.generator.generate import mlx_generate
-from exo.worker.engines.mlx.utils_mlx import shard_and_load
+from exo.worker.engines.mlx.utils_mlx import apply_chat_template, shard_and_load
 
 
 class MockLayer(nn.Module):
@@ -119,11 +119,12 @@ def run_gpt_oss_pipeline_device(
             max_tokens=max_tokens,
         )
 
+        prompt = apply_chat_template(tokenizer, task)
+
         generated_text = ""
+
         for response in mlx_generate(
-            model=model,
-            tokenizer=tokenizer,
-            task=task,
+            model=model, tokenizer=tokenizer, task=task, prompt=prompt
         ):
             generated_text += response.text
             if response.finish_reason is not None:
@@ -186,11 +187,14 @@ def run_gpt_oss_tensor_parallel_device(
             max_tokens=max_tokens,
         )
 
+        prompt = apply_chat_template(tokenizer, task)
+
         generated_text = ""
         for response in mlx_generate(
             model=model,
             tokenizer=tokenizer,
             task=task,
+            prompt=prompt,
         ):
             generated_text += response.text
             if response.finish_reason is not None:
diff --git a/src/exo/worker/tests/unittests/test_mlx/test_tokenizers.py b/src/exo/worker/tests/unittests/test_mlx/test_tokenizers.py
index 6ba684e9..d0b6707b 100644
--- a/src/exo/worker/tests/unittests/test_mlx/test_tokenizers.py
+++ b/src/exo/worker/tests/unittests/test_mlx/test_tokenizers.py
@@ -105,7 +105,7 @@ def event_loop():
     TEST_MODELS,
 )
 @pytest.mark.asyncio
-async def test_tokenizer_encode_decode(short_id: str, model_card: ModelCard) -> None:
+async def test_tokenizer_encode_decode(model_card: ModelCard) -> None:
     """Test that tokenizer can encode and decode text correctly."""
     model_id = model_card.model_id
 
@@ -170,9 +170,7 @@ async def test_tokenizer_encode_decode(short_id: str, model_card: ModelCard) ->
     TEST_MODELS,
 )
 @pytest.mark.asyncio
-async def test_tokenizer_has_required_attributes(
-    short_id: str, model_card: ModelCard
-) -> None:
+async def test_tokenizer_has_required_attributes(model_card: ModelCard) -> None:
     """Test that tokenizer has required attributes for inference."""
     model_id = model_card.model_id

← 307f454b feat: initial image generation support (#1095)  ·  back to Exo  ·  treefmt: add Svelte file formatting a7e205e4 →