Repository: exo-explore/exo
File: src/exo/worker/runner/runner.py (Python)

Code before:
                        elif isinstance(inference_model, GptOssModel):
                            mlx_generator = parse_gpt_oss(mlx_generator)

                        if tokenizer.has_tool_calling and not isinstance(
                            inference_model, GptOssModel
                        ):
                            assert tokenizer.tool_call_start
                            assert tokenizer.tool_call_end
                            assert tokenizer.tool_parser  # pyright: ignore[reportAny]
                            mlx_generator = parse_tool_calls(
                                mlx_generator,
                                tokenizer.tool_call_start,
                                tokenizer.tool_call_end,
                                tokenizer.tool_parser,  # pyright: ignore[reportAny]
                            )

                        last_checked = time.perf_counter()
                        for response in mlx_generator:
                            if (t := time.perf_counter()) - last_checked > 0.1:
                                last_checked = t
                                cancelled_tasks.update(cancel_receiver.collect())
                                want_to_cancel = (task.task_id in cancelled_tasks) or (
                                    TaskId("CANCEL_CURRENT_TASK") in cancelled_tasks
                                )
                                if mx_any(want_to_cancel, group):
                                    break

                            match response:
                                case GenerationResponse():
                                    if (
                                        device_rank == 0
                                        and response.finish_reason == "error"
                                    ):
                                        event_sender.send(
                                            ChunkGenerated(
                                                command_id=command_id,
                                                chunk=ErrorChunk(
                                                    error_message=response.text,
                                                    model=shard_metadata.model_card.model_id,
                                                ),
                                            )
                                        )

                                    elif device_rank == 0:
                                        assert response.finish_reason not in (
                                            "error",
                                            "tool_calls",
                                            "function_call",
                                        )
                                        event_sender.send(
                                            ChunkGenerated(

Review comment (performance):
Could you check performance against exo_bench? 

For example, I don't see any issue with, for example, putting this every second instead of 0.1s. Or the fact we are using a different stream will be enough, but perhaps we need to explicitly set a new stream?

Code after:
                        assert image_model
                        image = warmup_image_generator(model=image_model)
                        if image is not None:
                            logger.info(f"warmed up by generating {image.size} image")
                        else:
                            logger.info("warmup completed (non-primary node)")

                    current_status = RunnerReady()
                    logger.info("runner ready")
                case TextGeneration(task_params=task_params, command_id=command_id) if (
                    isinstance(current_status, RunnerReady)
                ):
                    logger.info(f"received chat request: {task}")
                    current_status = RunnerRunning()
                    logger.info("runner running")
                    event_sender.send(
                        RunnerStatusUpdated(
                            runner_id=runner_id, runner_status=current_status
                        )
                    )
                    event_sender.send(TaskAcknowledged(task_id=task.task_id))
                    assert inference_model
                    assert tokenizer
                    assert check_for_cancel_every

                    try:
                        _check_for_debug_prompts(task_params)

                        # Build prompt once - used for both generation and thinking detection
                        prompt = apply_chat_template(tokenizer, task_params)

                        # Generate responses using the actual MLX generation
                        mlx_generator = mlx_generate(
                            model=inference_model,
                            tokenizer=tokenizer,
                            task=task_params,
                            prompt=prompt,
                            kv_prefix_cache=kv_prefix_cache,
                            group=group,
                        )

                        # For other thinking models (GLM, etc.), check if we need to
                        # prepend the thinking tag that was consumed by the chat template
                        if detect_thinking_prompt_suffix(prompt, tokenizer):
                            mlx_generator = parse_thinking_models(
                                mlx_generator, tokenizer
                            )

                        # Kimi-K2 has tool call sections - we don't care about them
                        if "kimi" in shard_metadata.model_card.model_id.lower():
                            mlx_generator = filter_kimi_tokens(mlx_generator)