# ChronosML

[![License](https://img.shields.io/badge/License-Apache_2.0-blue.svg)](https://opensource.org/licenses/Apache-2.0)
[![Python 3.10+](https://img.shields.io/badge/python-3.10%2B-blue.svg)](https://www.python.org/downloads/)
[![MCP Compatible](https://img.shields.io/badge/MCP-%3C200%20tokens-success.svg)](https://modelcontextprotocol.io)

**ChronosML** is a resilient, zero-guesswork ML runtime, memory, and throughput estimation engine that enables any local or frontier AI agent (via **MCP**, **CLI**, and **Python SDK**) to accurately calculate wall-clock execution time, peak memory footprint, OOM risk, and optimal batching/parallelism **before** launching ML experiments, training runs, or inference benchmarks.

---

## Why ChronosML?

AI coding and research agents routinely misestimate ML experiment runtimes by **10x–100x** because they guess or use theoretical peak FLOPs instead of accounting for:
1. **Actual execution backend & silent traps**: Container/sandbox CUDA mismatches (`torch.cuda.is_available() == False`), CPU `bfloat16` software emulation, missing `flash-attn` / `causal-conv1d` / `fla` fast paths on hybrid linear-attention/SSM models.
2. **Memory-bandwidth Roofline bottlenecks**: Autoregressive token decoding ($B=1$) and pointwise activation streaming are memory-bandwidth bound, whereas prefill and large-batch training are compute Tensor Core bound.
3. **Padding waste, KV-cache growth, & cold-start I/O**: Sequence-length padding waste, dynamic KV-cache footprint across MHA/GQA/MoE/Hybrid layers, and `safetensors` disk `mmap` load latency.

---

## 4-Layer Architecture

```mermaid
flowchart LR
    A["Layer 1: Hardware Profiler\n(chronosml.hardware)"] --> B["Layer 2A: Analytical Roofline\n(chronosml.roofline)"]
    A --> C["Layer 2B: 2-Step Micro-Pilot\n(chronosml.pilot)"]
    B --> D["Layer 2: Hybrid Estimator\n(chronosml.estimator)"]
    C --> D
    D --> E["Layer 3: Bottleneck Advisor\n(chronosml.advisor)"]
    E --> F["Layer 4: Ultra-Compact MCP & CLI\n(chronosml.mcp / chronosml.ledger)"]
    F -. "Online Bayesian Residual" .-> D
```

1. **Layer 1: Live Hardware & Backend Telemetry Profiler (`chronosml.hardware`)**
   - Detects CPU architecture (`ARM64` / `x86_64`, physical/logical cores), system RAM & Unified Memory Architecture (UMA systems where `nvidia-smi` reports `[N/A]` for dedicated VRAM and OS page cache is dynamically reclaimable), GPU compute capability/SMs, and disk `mmap` read speed.
   - Detects silent performance traps (`CUDA_SANDBOX_MISMATCH`, `CPU_BF16_SOFTWARE_EMULATION`, `MISSING_FLASH_ATTENTION`, `MISSING_SSM_LINEAR_ATTN_KERNELS`, `SINGLE_THREAD_CPU_BOTTLENECK`, `HIGH_SWAP_OR_MEMORY_PRESSURE`).
   - Maintains `~/.cache/chronosml/hardware_profile.json` via a `<2s` calibration micro-benchmark suite (`chronosml calibrate`).

2. **Layer 2: Analytical Roofline + Micro-Pilot Hybrid Estimator (`chronosml.estimator`)**
   - **Analytical Roofline (`chronosml.roofline`)**: Parses HuggingFace `config.json`, zero-copy `safetensors` headers, PyTorch `nn.Module` objects, or presets across Dense MHA, GQA, MQA, MoE ($P_{\text{act}} \ll P$), and Linear-Attention/SSM/Hybrid (Mamba, Gated DeltaNet) architectures:
     - Cold-start load time: $T_{\text{load}} = \text{ModelBytes} / \text{BW}_{\text{disk/mmap}}$
     - Prefill time: $\max(\text{FLOPs}_{\text{prefill}} / \text{TFLOPs}_{\text{compute}}, \text{Bytes}_{\text{prefill}} / \text{BW}_{\text{mem}}) + \text{Bytes}_{\text{pointwise}} / \text{BW}_{\text{mem}}$
     - Decode step time per token: $\max(2 P_{\text{act}} B / \text{TFLOPs}_{\text{batch}}, (P_{\text{act}} \cdot b + \text{KV\_Bytes}(S)) / \text{BW}_{\text{mem}})$
     - Training step time: Forward + Backward ($\approx 3\times$ forward FLOPs for full training, $\approx 2.15\times$ for LoRA, plus activation checkpointing and AdamW optimizer state traffic).
   - **Empirical Micro-Pilot Probe (`chronosml.pilot`)**: Executes `1` warmup iteration + `2` timed micro-iterations (`pilot_probe`), measures step latency variance and $\Delta\text{RSS}$ / $\Delta\text{VRAM}$, extrapolates total runtime for $N$ items with a **95% prediction interval** $(T_{\text{p50}}, T_{\text{p90}}, T_{\text{p99}})$, and flags OOM risk if projected peak memory exceeds **85%** of available memory.

3. **Layer 3: Auto-Optimizer & Bottleneck Advisor (`chronosml.advisor`)**
   - Generates a structured `OptimizationPlan` with expected speedup (`speedup_if_optimized`) and concrete configuration recommendations: optimal `batch_size`, `dtype`, `torch.set_num_threads()`, `stop_strings` / `max_new_tokens` cap, left-padding length bucketing (`bucket_by_length`), activation checkpointing, and vectorized `llm_logprob` scoring vs. autoregressive decoding.

4. **Layer 4: Ultra-Compact MCP Server (`chronosml.mcp`) & Online Residual Ledger (`chronosml.ledger`)**
   - Exposes **at most 2 minimal MCP tools** (`estimate_ml_task` and `log_completed_run`) keeping total schema context overhead under **200 tokens** (<300 token budget).
   - Continuously fits per-hardware damped Bayesian residual multipliers in `~/.cache/chronosml/calibration_ledger.sqlite` & `.jsonl`.

---

## Installation

```bash
git clone https://github.com/kardashevproject/chronosml.git
cd chronosml
pip install -e .

# Run the <2-second hardware calibration suite
chronosml calibrate
```

---

## Plug-and-Play MCP Server Setup (Zero Context Tax)

ChronosML's MCP server (`chronosml-mcp` or `python3 -m chronosml.mcp_server`) speaks standard JSON-RPC 2.0 over `stdio` and adds `<200 tokens` to your agent's context window.

### Claude Code
```bash
claude mcp add chronosml -- chronosml-mcp
```

### Cursor / Claude Desktop / Windsurf (`mcp.json`)
```json
{
  "mcpServers": {
    "chronosml": {
      "command": "chronosml-mcp",
      "args": []
    }
  }
}
```

### Gemini CLI (`~/.gemini/settings.json`)
```json
{
  "mcpServers": {
    "chronosml": {
      "command": "chronosml-mcp",
      "args": []
    }
  }
}
```

### OpenAI Codex CLI (`~/.codex/config.toml`)
```toml
[mcp_servers.chronosml]
command = "chronosml-mcp"
args = []
```

---

## CLI Usage

```bash
# 1. Calibrate hardware & detect silent performance traps (<2s)
chronosml calibrate

# 2. Estimate runtime, peak memory, OOM risk & optimization plan (with 2-step micro-pilot)
chronosml estimate \
  --model qwen3.5-0.8b \
  --task-type llm_generate \
  --num-samples 200 \
  --batch-size 1 \
  --seq-len 512 \
  --max-new-tokens 128 \
  --pilot --compact

# 3. Generate a full Layer 3 OptimizationPlan with code hints
chronosml advise --model llama-3-8b --task-type llm_generate --num-samples 128 --batch-size 1 --dtype fp32

# 4. Probe any custom training or inference shell command (1 warmup + 2 micro-steps)
chronosml pilot --cmd "python3 -c 'import torch; a=torch.randn(1024,1024); _=a@a'" --num-samples 100 --batch-size 4

# 5. Log a completed run for online Bayesian residual calibration
chronosml log-run --task-type llm_generate --predicted-sec 14.2 --actual-sec 15.1 --backend cuda --dtype bf16

# 6. Benchmark predictions against real PyTorch runs on your machine and report MAPE
chronosml benchmark
```

---

## Python SDK Usage

```python
import torch
import torch.nn as nn
from chronosml import calibrate_hardware, estimate_task, pilot_probe, log_completed_run

# 1. Inspect calibrated hardware & silent performance traps
hw = calibrate_hardware(force=False)
print(f"Backend: {hw.primary_backend}, Free Memory: {hw.effective_free_memory_gb:.1f} GB")

# 2. Hybrid Roofline + Micro-Pilot estimation
report = estimate_task(
    model_path_or_config="qwen3.5-2b",  # or path to HF config.json / .safetensors / nn.Module
    task_type="llm_generate",
    num_samples=200,
    batch_size=1,
    seq_len=512,
    max_new_tokens=128,
    run_micro_pilot=True,
)
print(report.to_compact_dict())
# -> {'p50_sec': ..., 'p90_sec': ..., 'peak_mem_gb': ..., 'oom_risk': False,
#     'bottleneck': 'memory_bandwidth_bound', 'recommended_config': {...}, 'speedup_if_optimized': ...}

# 3. Direct 2-step Micro-Pilot Probe on a custom PyTorch training step
model = nn.Linear(2048, 2048, device=hw.primary_backend, dtype=torch.bfloat16)
x = torch.randn(8, 256, 2048, device=hw.primary_backend, dtype=torch.bfloat16)

probe = pilot_probe(
    lambda: model(x).sum().backward() or model.zero_grad(set_to_none=True),
    n_warmup=1,
    n_steps=2,
    num_samples=160,
    batch_size=8,
)
print(f"Extrapolated p50={probe.p50_sec:.4f}s, 95% PI=[{probe.ci95_low_sec:.4f}s, {probe.ci95_high_sec:.4f}s]")
```

---

## Empirical Verification Benchmarks

Evaluated via `chronosml benchmark` on an ARM64 Unified Memory GPU system (`92.87 TFLOP/s` `bf16` GEMM, `197.15 GB/s` sustained memory bandwidth):

| Workload | Task Type | Samples ($N$) | Batch ($B$) | Micro-Pilot $\hat{T}_{\text{p50}}$ (s) | Analytical Roofline $\hat{T}$ (s) | Actual Wall-Clock $T_{\text{actual}}$ (s) | Micro-Pilot APE (%) | Roofline APE (%) |
| :--- | :--- | :---: | :---: | :---: | :---: | :---: | :---: | :---: |
| `prefill_logprob_b8_s256` | `llm_logprob` | 160 | 8 | `0.19056` | `0.18485` | `0.20041` | **4.92%** | 7.76% |
| `autoregressive_decode_b4_t16` | `llm_generate` | 64 | 4 | `0.47441` | `0.31561` | `0.44145` | **7.46%** | 28.51% |
| `full_train_step_b4_s256` | `full_train` | 64 | 4 | `0.16031` | `0.19232` | `0.15852` | **1.13%** | 21.33% |
| `large_batch_embedding_b16_s512` | `embedding` | 320 | 16 | `1.69368` | `1.39222` | `1.65705` | **2.21%** | 15.98% |
| **Overall Mean (MAPE)** | — | — | — | — | — | — | **3.93%** | **18.39%** |

---

## License

Licensed under the [Apache License, Version 2.0](LICENSE).