Skip to content

feat: local MLX training and inference behind arkor start/dev --local (ENG-1065) - #228

Open
k-taro56 wants to merge 39 commits into
mainfrom
eng-1065
Open

feat: local MLX training and inference behind arkor start/dev --local (ENG-1065)#228
k-taro56 wants to merge 39 commits into
mainfrom
eng-1065

Conversation

@k-taro56

@k-taro56 k-taro56 commented Aug 16, 2026

Copy link
Copy Markdown
Contributor

Implements local training and inference for ENG-1065: arkor start --local and arkor dev --local run the whole loop on an Apple Silicon Mac via MLX, with no cloud account and nothing leaving the machine except model/dataset downloads.

Architecture

  • New published package @arkor/local (lockstep-versioned, ~60 KB tarball). Users opt in with pnpm add -D @arkor/local; the arkor package gains zero runtime dependencies and stays cloud-lean. The CLI resolves the runtime from the user's project (same resolution model as the build artifact resolving arkor) with a protocol-version handshake that turns incompatible pairings into actionable upgrade errors.
  • Local training server (loopback Hono, ephemeral port, per-launch bearer token with the same threat model as Studio's CSRF token) implements the cloud API subset the SDK and Studio actually call: POST/GET /v1/jobs, cancel, the SSE event stream, and POST /v1/inference/chat. Because the contract matches, trainer.wait(), the Studio jobs list, the live loss chart, and the Playground work unchanged.
  • Durable job store under .arkor/local/jobs/<id>/: append-only events.jsonl with per-job monotonic sequence ids (true Last-Event-ID replay across restarts), atomic job.json writes, size-capped console capture, and pid-guarded orphan reconciliation on boot.
  • Backend extensibility: MLX is the first implementation of a LocalTrainingBackend interface (preflight, config validation, run/inference spec building) behind a registry. All platform knowledge lives in each backend's preflight; a future CUDA or ROCm backend is one backends/<id>.ts, one Python shim speaking the same JSON-line protocol, and one registry entry.
  • Python shim (bundled MIT code, executed via uv run --with "mlx-lm[train]==0.31.3"): prepares datasets (HuggingFace or blob URL; chatml/sharegpt/alpaca/text/prompt_completion with columnMapping; automatic validation holdout), drives mlx-lm LoRA fine-tuning through its TrainingCallback, normalises checkpoints into adapters/step-<N>/ and adapters/final/, and reports structured progress as @arkor -marked JSON lines (non-finite floats sanitised).
  • Local inference: a lazily spawned OpenAI-compatible mlx_lm server child keyed by (model, adapter), with readiness polling, in-flight refcounting (idle shutdown never kills a mid-download or mid-generation child), crash respawn, and streaming passthrough. onCheckpoint.infer() and the Playground work against trained adapters.
  • Isolation guarantees: local mode never reads or writes ~/.arkor/credentials.json, never bootstraps projects, and ignores .arkor/state.json, so a local run can never poison later cloud use.

CI

  • New local-training-smoke.yaml: runs on every push and pull request (plus nightly and manual dispatch) on macos-15 (Apple Silicon), training SmolLM2-135M-Instruct for 4 real steps through the full stack and asserting on events, adapter artifacts, and checkpoint inference.
  • Release workflows pack, attest, and publish the third package (version reconciliation, not-yet-published guards, tarball globs, postflight verification updated).
  • ci.yaml uploads arkor-local coverage and test results under their own flags.

Testing

  • 89 unit tests in @arkor/local (fake Node children for every process seam), plus local-mode suites in arkor (trainer bypass with fs spies proving zero ~/.arkor access, Studio proxy repointing, loader handshake).
  • e2e: arkor start --local end-to-end through the real bin with a fake uv (Apple Silicon CI leg; the not-installed error path runs on every OS), and a Playwright spec driving arkor dev --local in Chromium against an unroutable cloud URL.
  • Full repo sweep green: 27 turbo tasks (lint, typecheck, test) plus format:check.
  • Two self-review passes fixed 10 real defects before this PR (cancel/spawn and SSE replay races, pipeline poisoning hang, unhandled stream errors, NaN protocol lines, idle-timer kills, a raw NUL byte, double-terminal races).

Docs

local-training.mdx added on both axes and both languages with nav registration; dev / build-and-start / overview / models / roadmap / quickstart updated (EN + JA in this PR).

Deliberately not implemented (documented)

Deployments/endpoints locally, tools/responseFormat/structuredOutputs/embeddings in local inference, pretokenized datasets, Playground free-text base-model UI, and the CUDA/ROCm backend bodies (seams only).


Summary by cubic

Implements ENG-1065 by adding opt-in local MLX training and inference through @arkor/local. Training was previously cloud-only; arkor start --local and arkor dev --local now run the full loop on Apple Silicon while cloud behavior remains unchanged.

New Features

  • Dynamically loads @arkor/local from the project with protocol checks and actionable install or upgrade errors.
  • Adds a loopback API server with durable jobs, SSE replay, cancellation, crash recovery, and local artifacts under .arkor/local/; startup reconciliation skips unreadable job dirs instead of aborting the launch.
  • Runs pinned mlx-lm LoRA training through uv (--no-config --no-project so user config can't break it), including dataset preparation, progress events, checkpoints, and arbitrary Hugging Face model IDs.
  • Manages local inference servers with readiness polling, streaming, crash recovery, reference counting, and prompt abort handling.
  • Wires --local and --backend <id> through the CLI and Studio; local mode skips cloud credentials and supports job cancellation.
  • Validates local configuration strictly (rejects unknown nested keys, non-string columns, unsafe blob URLs, unsupported optim options, and https-to-http redirects; warns on unknown top-level keys), protects local files, and prevents duplicate terminal events or orphaned processes across cancellation and restart races.
  • Keeps local dry-run jobs and unavailable deployments out of Studio actions, while cloud dry runs retain their existing behavior.
  • Publishes @arkor/local in lockstep and adds Apple Silicon CI and real MLX smoke coverage.

Migration

  • Requires an Apple Silicon Mac, uv on PATH, and @arkor/local installed with pnpm add -D @arkor/local@alpha.
  • Local mode uses the exact loopback hand-off http://127.0.0.1:<port> and stores models, datasets, logs, and adapters locally; only downloads leave the machine.
  • Local deployments and advanced inference features such as embeddings, tools, and structured outputs remain unsupported. Python SBOM and cross-process SSE cancel notices are known limitations.

Written for commit 344c511. Summary will update on new commits.

Review in cubic

Summary by CodeRabbit

  • New Features

    • Added optional local MLX training and inference on Apple Silicon through @arkor/local.
    • Added arkor start --local and arkor dev --local, with backend selection and Studio integration.
    • Added durable job tracking, progress streaming, cancellation, recovery, and local chat inference.
    • Local runs support arbitrary Hugging Face model IDs without cloud credentials.
    • Added Studio controls for cancelling queued or running jobs.
  • Bug Fixes

    • Dry-run jobs no longer appear as playable adapters in Studio.
    • Studio now indicates when cloud deployments are unavailable in local mode.
  • Documentation

    • Added comprehensive English and Japanese guides for local workflows.

Loading
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants