mlx-knife/docs/2.0-IMPLEMENTATION-GUIDE.md

2.0 Server/Run Implementation Guide

Purpose: Step-by-step guide for Sonnet sessions implementing server/run functionality
Created: 2025-09-10
Target: 2.0.0-beta.1-local through beta.3 (public)

Quick Reference for Sonnet

What You're Building

• Port server/run functionality from 1.x (main branch) to 2.0 (feature/2.0.0-alpha.1)
• Preserve 2.0's modular architecture (mlxk2/core/, mlxk2/operations/, mlxk2/output/)
• Test-first approach using specifications in docs/2.0-TEST-SPECIFICATIONS.md

Key Files to Reference

# 1.x source files (use git show to view)
git show main:mlx_knife/server.py          # FastAPI server implementation
git show main:mlx_knife/mlx_runner.py      # MLX execution engine
git show main:mlx_knife/reasoning_utils.py # Reasoning model support
git show main:mlx_knife/cli.py             # CLI command definitions

# 2.0 existing structure
mlxk2/core/cache.py      # Extend with model detection
mlxk2/operations/*.py    # Add run.py, serve.py, chat.py
mlxk2/output/*.py        # Extend for streaming support
mlxk2/cli.py            # Add new commands

Implementation Steps

Step 1.0: Core Runner Implementation

File: mlxk2/core/runner.py

# Key components to port from mlx_runner.py:
class MLXRunner:
    """Core MLX model execution engine"""
    
    def __init__(self, model_name_or_path):
        # Model loading logic
        # Memory tracking
        
    def __enter__(self):
        # Context manager entry
        return self
        
    def __exit__(self, exc_type, exc_val, exc_tb):
        # CRITICAL: Cleanup even on exception
        
    def generate_streaming(self, prompt, **kwargs):
        # Generator for token-by-token output
        yield from self._generate_tokens(prompt, **kwargs)
        
    def generate_batch(self, prompt, **kwargs):
        # Complete generation at once
        return "".join(self.generate_streaming(prompt, **kwargs))

Critical Requirements:

1. Context manager pattern for memory safety
2. Separate streaming vs batch generation
3. Stop token filtering (CHAT_STOP_TOKENS)
4. Dynamic token limits based on model context

Step 1.1: Complete Run Command

File: mlxk2/operations/run.py

from mlxk2.core.runner import MLXRunner

def run_model(
    model_spec: str,
    prompt: Optional[str] = None,
    stream: bool = True,
    max_tokens: Optional[int] = None,
    temperature: float = 0.7,
    top_p: float = 0.9,
    **kwargs
):
    """Execute model with prompt - supports both single-shot and interactive modes.
    
    Args:
        model_spec: Model specification
        prompt: Input prompt (None = interactive mode)
        stream: Enable streaming output
        max_tokens: Maximum tokens (None = full model context)
        temperature: Sampling temperature
        top_p: Top-p sampling parameter
    """
    with MLXRunner(model_spec) as runner:
        # Interactive mode: no prompt provided
        if prompt is None:
            interactive_chat(runner, stream=stream, max_tokens=max_tokens, **kwargs)
        else:
            # Single-shot mode: prompt provided
            single_shot_generation(runner, prompt, stream=stream, max_tokens=max_tokens, **kwargs)

def interactive_chat(runner, stream=True, **kwargs):
    """Interactive conversation mode with history tracking."""
    print("Starting interactive chat. Type 'exit' or 'quit' to end.\n")
    
    conversation_history = []
    
    while True:
        try:
            user_input = input("You: ").strip()
            
            if user_input.lower() in ['exit', 'quit', 'q']:
                print("\nGoodbye!")
                break
                
            if not user_input:
                continue
                
            # Add user message to conversation history
            conversation_history.append({"role": "user", "content": user_input})
            
            # Format conversation using chat template
            formatted_prompt = runner._format_conversation(conversation_history)
            
            # Generate response
            print("\nAssistant: ", end="", flush=True)
            
            if stream:
                # Streaming mode
                response_tokens = []
                for token in runner.generate_streaming(formatted_prompt, use_chat_template=False, **kwargs):
                    print(token, end="", flush=True)
                    response_tokens.append(token)
                response = "".join(response_tokens).strip()
            else:
                # Batch mode
                response = runner.generate_batch(formatted_prompt, use_chat_template=False, **kwargs)
                print(response)
            
            # Add assistant response to history
            conversation_history.append({"role": "assistant", "content": response})
            print()  # Newline after response
            
        except KeyboardInterrupt:
            print("\n\nChat interrupted. Goodbye!")
            break
        except Exception as e:
            print(f"\n[ERROR] {e}")
            continue

def single_shot_generation(runner, prompt, stream=True, **kwargs):
    """Single prompt generation."""
    if stream:
        for token in runner.generate_streaming(prompt, **kwargs):
            print(token, end="", flush=True)
        print()  # Final newline
    else:
        result = runner.generate_batch(prompt, **kwargs)
        print(result)

CLI Integration (mlxk2/cli.py):

# Run command parser
run_parser = subparsers.add_parser("run", help="Run model with prompt")
run_parser.add_argument("model", help="Model name to run")
run_parser.add_argument("prompt", nargs="?", help="Input prompt (optional - triggers interactive mode if omitted)")
run_parser.add_argument("--max-tokens", type=int, help="Maximum tokens to generate (default: full model context)")
run_parser.add_argument("--temperature", type=float, default=0.7, help="Sampling temperature")
run_parser.add_argument("--top-p", type=float, default=0.9, help="Top-p sampling parameter")
run_parser.add_argument("--no-stream", action="store_true", help="Disable streaming output (batch mode)")
run_parser.add_argument("--json", action="store_true", help="Output in JSON format")
run_parser.add_argument("--verbose", action="store_true", help="Show detailed output")

# Usage examples:
# mlxk2 run model "prompt"                    # Single-shot streaming
# mlxk2 run model "prompt" --no-stream        # Single-shot batch
# mlxk2 run model                             # Interactive streaming  
# mlxk2 run model --no-stream                 # Interactive batch

Key Changes from Basic to Complete:

• ✅ Interactive mode: prompt parameter is now optional
• ✅ Conversation history: Tracks full chat context
• ✅ Stream control: --no-stream works in both modes
• ✅ Full context tokens: No arbitrary limits for run command
• ✅ Chat template integration: Uses model's native conversation format

Step 1.2: Beta.1 Completion

Complete the remaining Beta.1 requirements:

1.2.1: Full Context Token Limits

File: mlxk2/core/runner.py

def _calculate_dynamic_max_tokens(self, server_mode: bool = False) -> int:
    """Calculate dynamic max tokens based on model context and usage mode."""
    if not self._context_length:
        return 2048
    
    if server_mode:
        # Server: half context for DoS protection
        return self._context_length // 2
    else:
        # Run command: full context (user's own machine, be generous)
        return self._context_length

# Update generate_streaming and generate_batch to use:
effective_max_tokens = max_tokens if max_tokens is not None else self._calculate_dynamic_max_tokens(server_mode=False)

1.2.2: Ctrl-C Handling

Already implemented in our MLXRunner: ✅

• Signal handler in __init__
• _interrupted flag checking during generation
• Graceful interruption with user message

1.2.3: Interactive Mode Implementation

Server Model Caching (Hot‑Swap, kein Reload pro Prompt)

Ziel: Die UX‑Verbesserung aus 1.1.1 beibehalten – der Server lädt Modelle nicht für jeden Prompt neu.

• Mechanik:
  • In mlxk2/core/server_base.py existiert ein globaler Runner‑Cache:
    • _model_cache: Dict[str, MLXRunner] und _current_model_path: Optional[str].
    • get_or_load_model(model_spec): gibt einen bestehenden MLXRunner zurück, falls bereits geladen; lädt nur bei Modellwechsel neu.
    • Beim Wechsel wird der alte Runner unter Lock bereinigt (runner.cleanup()), dann der neue geladen (Hot‑Swap).
  • Für den Server wird MLXRunner(..., install_signal_handlers=False) verwendet (keine Signal‑Handler‑Konflikte).
• Verhalten:
  • Gleiches Modell über mehrere Requests → kein Reload → zügige Antworten, stabile UX.
  • Anderes Modell → altes Modell freigeben, neues laden (Hot‑Swap), weiterhin kein Reload pro Prompt.
• Kontextlänge (Erinnerung):
  • Run‑Command nutzt volle Kontextlänge; Server nutzt halbe Kontextlänge als DoS‑Schutz (get_effective_max_tokens(..., server_mode=True)).

File: mlxk2/operations/run.py - Add missing methods:

def _format_conversation(self, messages: List[Dict[str, str]]) -> str:
    """Format conversation history into a prompt using chat template."""
    if hasattr(self.tokenizer, 'chat_template') and self.tokenizer.chat_template:
        try:
            return self.tokenizer.apply_chat_template(
                messages,
                tokenize=False,
                add_generation_prompt=True
            )
        except Exception:
            # Fall back to legacy format
            pass
    
    # Legacy Human:/Assistant: format
    formatted_parts = []
    for msg in messages:
        role = msg["role"]
        content = msg["content"]
        if role == "system":
            formatted_parts.append(f"System: {content}")
        elif role == "user":
            formatted_parts.append(f"Human: {content}")
        elif role == "assistant":
            formatted_parts.append(f"Assistant: {content}")
    
    return "\n\n".join(formatted_parts) + "\n\nAssistant: "

1.2.4: Update CLI for Interactive Mode

File: mlxk2/cli.py

# Update run command argument parser
run_parser.add_argument("prompt", nargs="?", help="Input prompt (optional - triggers interactive mode if omitted)")

# Update run command handler
elif args.command == "run":
    result_text = run_model_enhanced(
        model_spec=args.model,
        prompt=args.prompt,  # Can be None for interactive mode
        stream=not args.no_stream,
        # ... other parameters
    )

1.2.5: Beta.1 Test Coverage

Files: Complete test implementation for:

• tests_2.0/test_run_complete.py - All run command scenarios
• tests_2.0/test_interactive_mode.py - Conversation history and chat templates
• tests_2.0/test_token_limits.py - Full context vs server context
• tests_2.0/test_ctrl_c_handling.py - Interruption scenarios

Coverage Target: 80% for run command functionality

Step 2.0: Server Implementation (Beta.2-local Core)

File: mlxk2/core/server_base.py

from fastapi import FastAPI
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel

# OpenAI-compatible request/response models
class ChatCompletionRequest(BaseModel):
    model: str
    messages: List[Dict[str, str]]
    stream: Optional[bool] = False
    max_tokens: Optional[int] = None
    
class ChatCompletionResponse(BaseModel):
    choices: List[Dict]
    model: str
    usage: Dict

File: mlxk2/operations/serve.py

def start_server(model=None, port=8000, host="127.0.0.1"):
    """Start OpenAI-compatible API server"""
    # 1. Create FastAPI app
    # 2. Setup endpoints (/v1/chat/completions, /v1/models)
    # 3. Handle streaming vs non-streaming with SSE
    # 4. Model hot-swapping support
    # 5. Half context token limits (DoS protection)

Step 2.1: Beta.2 Parity Features

2.1.1: Reasoning Support (GPT-OSS/MXFP4)

CRITICAL: This is already implemented in 1.1.1-beta.3 and must be ported for parity!

File: mlxk2/core/reasoning.py

# Port from mlx_knife/reasoning_utils.py (1.x main branch)
class ReasoningExtractor:
    """Extract reasoning from GPT-OSS/MXFP4 models"""
    
    PATTERNS = {
        'gpt-oss': {
            'reasoning': r'<\|channel\|>analysis<\|message\|>(.*?)<\|end\|>',
            'final': r'<\|channel\|>final<\|message\|>(.*?)(?:<\|return\|>|$)',
        }
    }
    
class StreamingReasoningParser:
    """Parse reasoning tokens in real-time"""
    # Real-time token classification
    # Format as **[Reasoning]** / **[Answer]**

Integration:

• Runner detects MXFP4/GPT-OSS models via _is_reasoning_model()
• Formats output as [Reasoning] ... --- [Answer]
• Server API includes reasoning in response metadata (optional)

2.1.2: Issue #30 - Gated Models Preflight

File: mlxk2/operations/pull.py

def preflight_repo_access(model_spec):
    """Check repository access before download."""
    try:
        HfApi().model_info(repo_id, token=os.getenv("HUGGINGFACE_HUB_TOKEN"))
    except HTTPError as e:
        if e.response.status_code in [401, 403]:
            return {"error": "Model requires authentication"}
    return {"status": "accessible"}

Testing Strategy

Test Organization

tests_2.0/
├── test_runner_core.py         # Core MLXRunner tests
├── test_run_command.py         # CLI run tests
├── test_server_api.py          # OpenAI API compliance
├── test_reasoning.py           # GPT-OSS reasoning
└── test_chat_mode.py          # Interactive chat

Test Fixtures to Use

# From tests_2.0/conftest.py
@pytest.fixture
def temp_cache_dir():
    """Isolated cache for testing"""
    
@pytest.fixture
def mock_tiny_model():
    """Minimal model for fast tests"""

CRITICAL NOTES FOR SONNET

⚠️ Open Issues to Fix During Port

Issue #30: Gated Models Preflight Check [Beta.2]

Problem: Pull von gated models startet Download, dann 403 → Cache pollution Target: 2.0.0-beta.2-local Solution für 2.0:

# In mlxk2/operations/pull.py
def preflight_repo_access(model_spec):
    try:
        HfApi().model_info(repo_id, token=os.getenv("HUGGINGFACE_HUB_TOKEN"))
    except HTTPError as e:
        if e.response.status_code in [401, 403]:
            # Fail fast BEVOR Download
            return {"error": "Model requires authentication. Please accept terms and set HUGGINGFACE_HUB_TOKEN"}

Ctrl-C Handling [Beta.1] (Nicht als Issue dokumentiert)

Problem: Run/Server blockiert während Model-Generation, Ctrl-C funktioniert nicht Target: 2.0.0-beta.1-local (Core functionality!) Solution für 2.0:

import signal
import threading

class MLXRunner:
    def __init__(self):
        self._interrupted = False
        signal.signal(signal.SIGINT, self._handle_interrupt)
    
    def _handle_interrupt(self, signum, frame):
        self._interrupted = True
        # Generation-Loop checkt self._interrupted
    
    def generate_streaming(self):
        for token in model.generate():
            if self._interrupted:
                yield "\n[Generation interrupted by user]"
                break
            yield token

⚠️ Model Loading & Caching

WICHTIG: Der Server in 1.x cached Modelle im Memory. In 2.0:

• Model-Cache global in mlxk2/core/server_base.py
• NICHT bei jedem Request neu laden!
• Hot-swapping = nur wenn anderes Modell requested

⚠️ JSON vs Human Output (CLI-Ebene)

WICHTIG: 2.0 hat BEIDE Output-Modi auf CLI-Ebene:

• Default ohne --json: Human-readable output (wie 1.x)
• Mit --json: JSON output auf stdout
• Server API: Immer OpenAI-JSON Format (unabhängig von CLI)
• Streaming: Technisch separate Implementierung (SSE für Server, direktes Token-Streaming für CLI)

⚠️ Stop Tokens & Code-Sharing

DESIGN-PRINZIP: Server baut maximal auf run-Funktionalität auf!

# Runner implementiert die Core-Logik
CHAT_STOP_TOKENS = ["\nHuman:", "\nAssistant:", "\nUser:", "\nYou:"]

# Server nutzt Runner - KEINE Duplikation
from mlxk2.core.runner import MLXRunner
# Server ruft runner.generate_streaming() oder runner.generate_batch()

VORTEIL: Einmal richtig implementiert, überall korrekt

⚠️ Test Models & RAM-aware Filtering

LOKALE TESTS: RAM-aware Filtering aus 1.x BEIBEHALTEN!

# Aus 1.x TESTING.md - diese Logik portieren:
- 8GB Mac: Nur tiny models
- 16GB Mac: Bis zu 7B models  
- 32GB+ Mac: Alle models möglich

GitHub CI: Nicht möglich (keine Apple Silicon Runner)

• Docs müssen klar sagen: "Lokale Tests only"
• Badge "166/166 tests" bezieht sich auf lokale Ausführung

Common Pitfalls & Solutions

1. Memory Leaks & Process Monitoring

Problem: Model stays in memory after error / Zombie processes Solution:

• Context manager mit garantiertem cleanup in __exit__
• Portiere Process-Monitoring aus 1.x beta.2:
  • test_server_functionality.py: Server lifecycle tests
  • Process guards gegen orphaned Python processes
  • Automatic cleanup on Ctrl-C/SIGTERM

2. Streaming vs Batch Inconsistency

Problem: Different output between modes
Solution: Filter stop tokens in BOTH paths

3. Token Limits

Problem: Hardcoded limits truncate output
Solution: Dynamic limits aus 1.x (funktioniert gut!)

# Von 1.x beibehalten:
- max_tokens=None → Dynamische Limits basierend auf Model-Context
- Explicit max_tokens → Respektieren
- Formel aus 1.x mlx_runner.py übernehmen

Mögliche Verbesserung: Config-basierte Overrides für spezielle Modelle

4. Model Path Resolution

Problem: Can't find models in cache
Solution: Use existing mlxk2/core/cache.py resolution

Version Milestones

2.0.0-beta.1-local

Step 1.0: ✅ MLXRunner core engine
Step 1.1: ✅ Complete run command (single-shot + interactive)
Step 1.2: 🔄 Beta.1 completion

• Full context token limits (no DoS protection)
• Interactive mode implementation
• CLI integration for interactive mode
• 80% test coverage
• Ctrl-C handling (already implemented)

2.0.0-beta.2-local

Goal: 1.1.1-beta.3 parity + core stability Step 2.0: 🔄 Server implementation
Step 2.1: 🔄 Parity features (required for 1.x compatibility)

• OpenAI-compatible API server
• Half context token limits for server (DoS protection)
• Model hot-swapping support
• SSE streaming endpoints
• Reasoning models (GPT-OSS/MXFP4) ← ALREADY IN 1.1.1-beta.3!
• Issue #30: Gated models preflight
• Enhanced error handling and logging
• Server lifecycle management (Ctrl-C, cleanup)
• 90% test coverage

2.0.0-beta.3 (public)

Goal: Production-ready with 1.1.1-beta.3 complete parity

• All core features stable and battle-tested
• Performance optimized
• Documentation complete
• 95%+ test coverage
• Integration testing with real-world scenarios

Beyond 2.0.0-beta.3 (Future Releases)

New features for post-beta.3 versions:

• System Prompt CLI Support (--system parameter) - not yet specified
• Advanced reasoning model support (DeepSeek R1, QwQ, etc.)
• Custom reasoning token markers (--reasoning-start, --reasoning-end)
• Enhanced chat template system

Push Function Notes

The push operation (experimental in alpha.3) remains functional throughout beta phases:

• May receive fixes between beta versions
• Minor enhancements possible
• Not blocking for server/run implementation
• Already working with user's workflow

Quick Commands for Development

# View 1.x implementation
git show main:mlx_knife/server.py | less

# Run 2.0 tests
pytest tests_2.0/

# Test specific functionality
pytest tests_2.0/test_runner_core.py -v

# Check coverage
pytest tests_2.0/ --cov=mlxk2 --cov-report=term-missing

# Create local beta tag (not pushed)
git tag -a 2.0.0-beta.1-local -m "Initial server/run port"

# Run local 2.0 version
python -m mlxk2.cli run model "prompt"

References for Each Step

Step 1.0 (Runner Core)

• Source: git show main:mlx_knife/mlx_runner.py
• Tests: git show main:tests/unit/test_mlx_runner_memory.py

Step 1.1 (Run Command)

• Source: git show main:mlx_knife/cli.py (run_model function)
• Tests: git show main:tests/integration/test_run_command_advanced.py

Step 2.0 (Server)

• Source: git show main:mlx_knife/server.py
• Tests: git show main:tests/integration/test_server_functionality.py

Step 3.0 (Reasoning)

• Source: git show main:mlx_knife/reasoning_utils.py
• Context: CLAUDE.md reasoning architecture section

Step 3.1 (Chat)

• Source: Search for "interactive_chat" in main branch
• Tests: Look for chat-related tests in integration

Success Criteria

Each Sonnet session should:

1. Write tests first (TDD)
2. Implement minimal working version
3. Verify tests pass
4. Document any deviations from 1.x

Remember: The goal is feature parity with 1.1.1-beta.3, not innovation. Port conservatively.