aquin 0.0.1__py3-none-any.whl

aquin/__init__.py

@@ -0,0 +1,3 @@
+from .run import Run, init
+
+__all__ = ["Run", "init"]

aquin/cli.py

@@ -0,0 +1,197 @@
+from __future__ import annotations
+
+import json
+import os
+import sys
+import tarfile
+from pathlib import Path
+
+_CONFIG_PATH = Path.home() / ".aquin" / "config.json"
+_BASE_URL = "https://api.aquin.app"
+_PACKAGE_PATH = Path.cwd() / "aquin_run.tar.gz"
+
+
+def _load_config() -> dict:
+    if _CONFIG_PATH.exists():
+        with open(_CONFIG_PATH) as f:
+            return json.load(f)
+    return {}
+
+
+def _save_config(data: dict) -> None:
+    _CONFIG_PATH.parent.mkdir(parents=True, exist_ok=True)
+    with open(_CONFIG_PATH, "w") as f:
+        json.dump(data, f, indent=2)
+
+
+def _get_api_key(explicit: str | None = None) -> str:
+    if explicit:
+        return explicit
+    cfg = _load_config()
+    key = cfg.get("api_key") or os.environ.get("AQUIN_API_KEY") or ""
+    if not key:
+        print("No API key found. Run:  aquin login")
+        sys.exit(1)
+    return key
+
+
+def _package_path() -> Path:
+    return Path.cwd() / "aquin_run.tar.gz"
+
+
+def _read_meta() -> dict:
+    run_dir = Path.cwd() / "aquin_run"
+    if not run_dir.exists():
+        print("No aquin_run/ directory found. Run your training script first.")
+        sys.exit(1)
+    meta_path = run_dir / "meta.json"
+    if not meta_path.exists():
+        print("aquin_run/meta.json not found. Run your training script first.")
+        sys.exit(1)
+    with open(meta_path) as f:
+        return json.load(f)
+
+
+def cmd_login(args: list[str]) -> None:
+    """aquin login  — save your API key"""
+    key = input("Paste your Aquin API key: ").strip()
+    if not key:
+        print("No key entered.")
+        sys.exit(1)
+    cfg = _load_config()
+    cfg["api_key"] = key
+    _save_config(cfg)
+    print(f"API key saved to {_CONFIG_PATH}")
+
+
+def cmd_package(args: list[str]) -> None:
+    """aquin package  — bundle aquin_run/ into aquin_run.tar.gz"""
+    meta = _read_meta()
+    run_name = meta.get("run_name", meta["run_id"])
+    out = _package_path()
+
+    print(f"[aquin] Packaging run '{run_name}'...")
+    run_dir = Path.cwd() / "aquin_run"
+    with tarfile.open(out, "w:gz") as tar:
+        tar.add(run_dir, arcname="aquin_run")
+
+    size_mb = round(out.stat().st_size / 1_000_000, 1)
+    print(f"[aquin] Package ready: {out.name} ({size_mb} MB)")
+    print(f"[aquin] Run  aquin push  to send it to Aquin.")
+
+
+def cmd_push(args: list[str]) -> None:
+    """aquin push  — push aquin_run.tar.gz to Aquin"""
+    import requests
+
+    explicit_key = None
+    for i, a in enumerate(args):
+        if a == "--key" and i + 1 < len(args):
+            explicit_key = args[i + 1]
+
+    pkg = _package_path()
+    if not pkg.exists():
+        print("No aquin_run.tar.gz found. Run  aquin package  first.")
+        sys.exit(1)
+
+    meta = _read_meta()
+    run_id = meta["run_id"]
+    run_name = meta.get("run_name", run_id)
+    api_key = _get_api_key(explicit_key)
+    base_url = os.environ.get("AQUIN_BASE_URL", _BASE_URL).rstrip("/")
+
+    size_mb = round(pkg.stat().st_size / 1_000_000, 1)
+    print(f"[aquin] Pushing '{run_name}' ({size_mb} MB)...")
+
+    try:
+        # Step 1: get the VM upload URL from Next.js (auth only, no body)
+        auth_resp = requests.get(
+            f"{base_url}/api/sdk/push",
+            params={
+                "run_id": run_id,
+                "run_name": run_name,
+                "base_model": meta.get("base_model", ""),
+            },
+            headers={"Authorization": f"Bearer {api_key}"},
+            timeout=15,
+        )
+        auth_resp.raise_for_status()
+        upload = auth_resp.json()
+
+        # Step 2: push tar.gz directly to the VM (bypasses Vercel body limit)
+        file_size = pkg.stat().st_size
+        with open(pkg, "rb") as f:
+            resp = requests.post(
+                upload["url"],
+                data=f,
+                headers={**upload["headers"], "Content-Length": str(file_size)},
+                timeout=600,
+            )
+        resp.raise_for_status()
+        print("[aquin] Run pushed successfully.")
+        print("[aquin] Open Aquin, your run will appear as a notification.")
+    except Exception as exc:
+        print(f"[aquin] Push failed: {exc}")
+        sys.exit(1)
+
+
+def cmd_whoami(args: list[str]) -> None:
+    """aquin whoami  — show the currently logged in user"""
+    import requests
+    key = _get_api_key()
+    base_url = os.environ.get("AQUIN_BASE_URL", _BASE_URL).rstrip("/")
+    try:
+        resp = requests.get(
+            f"{base_url}/api/sdk/whoami",
+            headers={"Authorization": f"Bearer {key}"},
+            timeout=10,
+        )
+        resp.raise_for_status()
+        data = resp.json()
+        print(f"Logged in as: {data.get('email', data.get('user_id', '?'))}")
+    except Exception as exc:
+        print(f"Error: {exc}")
+
+
+def cmd_help(args: list[str]) -> None:
+    print("Aquin CLI")
+    print("")
+    print("Usage: aquin <command>")
+    print("")
+    print("Commands:")
+    print("  login       Save your Aquin API key")
+    print("  package     Bundle aquin_run/ into aquin_run.tar.gz")
+    print("  push        Push aquin_run.tar.gz to Aquin")
+    print("  whoami      Show the currently logged in account")
+    print("  help        Show this help message")
+    print("")
+    print("Typical workflow:")
+    print("  1. Add aquin.init() / run.log() / run.finish() to your training script")
+    print("  2. Run your training script")
+    print("  3. aquin package")
+    print("  4. aquin push")
+
+
+def main() -> None:
+    args = sys.argv[1:]
+    if not args:
+        cmd_help([])
+        return
+
+    cmd = args[0]
+    rest = args[1:]
+
+    if cmd == "login":
+        cmd_login(rest)
+    elif cmd == "package":
+        cmd_package(rest)
+    elif cmd == "push":
+        cmd_push(rest)
+    elif cmd == "whoami":
+        cmd_whoami(rest)
+    elif cmd in ("help", "--help", "-h"):
+        cmd_help(rest)
+    else:
+        print(f"Unknown command: {cmd}")
+        print("Run  aquin help  to see available commands.")
+        sys.exit(1)

aquin/run.py

@@ -0,0 +1,197 @@
+from __future__ import annotations
+
+import json
+import os
+import threading
+import time
+import uuid
+from pathlib import Path
+from typing import Any
+
+
+def init(
+    base_model: str | None = None,
+    run_name: str | None = None,
+    config: dict[str, Any] | None = None,
+) -> "Run":
+    """Start recording a training run. Call this before your training loop."""
+    return Run(base_model=base_model, run_name=run_name, config=config)
+
+
+class Run:
+    """
+    Records metrics and checkpoints locally during training.
+    Call aquin.init() to create one.
+
+    After training, use the CLI:
+        aquin package  # then: aquin push
+
+    Example
+    -------
+    run = aquin.init(
+        base_model="meta-llama/Llama-3.2-1B-Instruct",
+        run_name="my-run",
+        config={"lr": 2e-4, "epochs": 3, "rank": 16, "lora_alpha": 32,
+                "method": "qlora", "per_device_train_batch_size": 2,
+                "gradient_accumulation_steps": 8, "dataset": "data.jsonl"},
+    )
+    for step, batch in enumerate(dataloader):
+        loss = train_step(batch)
+        run.log(step, loss=loss.item(), learning_rate=scheduler.get_last_lr()[0],
+                grad_norm=grad_norm, epoch=epoch)
+    run.checkpoint(model, step=step)
+    run.finish()
+    """
+
+    def __init__(
+        self,
+        base_model: str | None = None,
+        run_name: str | None = None,
+        config: dict[str, Any] | None = None,
+    ) -> None:
+        self._run_id = str(uuid.uuid4())
+        self._run_name = run_name or f"run-{self._run_id[:8]}"
+        self._base_model = base_model or os.environ.get("AQUIN_BASE_MODEL") or ""
+        self._started_at = time.time()
+
+        self._run_dir = Path.cwd() / "aquin_run"
+        self._run_dir.mkdir(exist_ok=True)
+        (self._run_dir / "checkpoints").mkdir(exist_ok=True)
+
+        self._metrics: list[dict[str, Any]] = []
+        self._metrics_lock = threading.Lock()
+
+        self._write_meta()
+        if config:
+            self._write_config(config)
+
+        print(f"[aquin] Recording run '{self._run_name}' (id: {self._run_id})")
+        print(f"[aquin] Run data will be saved to {self._run_dir}")
+
+    # ── Public API ──────────────────────────────────────────────────────────────
+
+    def log(
+        self,
+        step: int,
+        *,
+        loss: float,
+        learning_rate: float | None = None,
+        grad_norm: float | None = None,
+        momentum_norm: float | None = None,
+        epoch: int | None = None,
+        batch: int | None = None,
+        total_batches: int | None = None,
+        step_ms: float | None = None,
+        **extra: float,
+    ) -> None:
+        """
+        Record metrics for one training step. Call this every step inside your loop.
+
+        Required
+        --------
+        step          : global training step index (0-based or 1-based, be consistent)
+        loss          : scalar training loss for this step
+
+        Recommended
+        -----------
+        learning_rate : current LR from your scheduler
+        grad_norm     : gradient norm — compute with torch.nn.utils.clip_grad_norm_()
+        momentum_norm : optimizer momentum norm — enables Momentum chart in the UI
+                        for AdamW: sum(p.exp_avg.norm() for p in optimizer.state.values())
+        epoch         : current epoch number — enables epoch summary table in the UI
+        batch         : batch index within the epoch
+        total_batches : total batches per epoch
+        step_ms       : wall-clock time for this step in milliseconds
+
+        Any additional keyword arguments are stored alongside the standard metrics.
+        """
+        entry: dict[str, Any] = {"step": step, "metrics": {"loss": loss}}
+        if learning_rate is not None:
+            entry["metrics"]["learning_rate"] = learning_rate
+        if grad_norm is not None:
+            entry["metrics"]["grad_norm"] = grad_norm
+        if momentum_norm is not None:
+            entry["metrics"]["momentum_norm"] = momentum_norm
+        if step_ms is not None:
+            entry["metrics"]["step_ms"] = step_ms
+        if extra:
+            entry["metrics"].update(extra)
+        if epoch is not None:
+            entry["epoch"] = epoch
+        if batch is not None:
+            entry["batch"] = batch
+        if total_batches is not None:
+            entry["total_batches"] = total_batches
+
+        with self._metrics_lock:
+            self._metrics.append(entry)
+
+    def checkpoint(self, model: Any, step: int) -> None:
+        """
+        Save the final model checkpoint. Call once at the end of training.
+        One checkpoint per run — always replaces the previous save.
+
+        Args
+        ----
+        model: a PyTorch nn.Module
+        step : the training step this checkpoint corresponds to
+        """
+        try:
+            import torch
+        except ImportError:
+            raise ImportError("torch is required to save checkpoints.")
+
+        ckpt_path = self._run_dir / "checkpoints" / "checkpoint.pt"
+        print(f"[aquin] Saving checkpoint (step {step})...")
+        torch.save({"step": step, "state_dict": model.state_dict()}, ckpt_path)
+        self._write_metrics()
+        print(f"[aquin] Checkpoint saved (step {step}).")
+
+    def finish(self, config: dict[str, Any] | None = None) -> None:
+        """
+        Finalise the run. Flushes all metrics and optional config to disk.
+
+        If you passed config to aquin.init() you don't need to pass it again here.
+        After this, run:  aquin package  # then:  aquin push
+
+        Args
+        ----
+        config: training hyperparameters to store alongside the run.
+                Recognised keys (all optional):
+                  lr, epochs, rank, lora_alpha, method,
+                  per_device_train_batch_size, gradient_accumulation_steps,
+                  dataset, max_seq_len, optimizer, scheduler,
+                  dropout, weight_decay, grad_clip
+        """
+        self._write_metrics()
+        if config:
+            self._write_config(config)
+        elapsed = round(time.time() - self._started_at)
+        print(f"[aquin] Run finished. {len(self._metrics)} steps recorded in {elapsed}s.")
+        print(f"[aquin] Push to Aquin with:  aquin package  # then: aquin push")
+
+    @property
+    def run_id(self) -> str:
+        return self._run_id
+
+    # ── Internal ────────────────────────────────────────────────────────────────
+
+    def _write_meta(self) -> None:
+        meta = {
+            "run_id": self._run_id,
+            "run_name": self._run_name,
+            "base_model": self._base_model,
+            "started_at": self._started_at,
+        }
+        with open(self._run_dir / "meta.json", "w") as f:
+            json.dump(meta, f, indent=2)
+
+    def _write_metrics(self) -> None:
+        with self._metrics_lock:
+            data = list(self._metrics)
+        with open(self._run_dir / "metrics.json", "w") as f:
+            json.dump(data, f)
+
+    def _write_config(self, config: dict[str, Any]) -> None:
+        with open(self._run_dir / "config.json", "w") as f:
+            json.dump(config, f, indent=2)

aquin-0.0.1.dist-info/METADATA

@@ -0,0 +1,136 @@
+Metadata-Version: 2.4
+Name: aquin
+Version: 0.0.1
+Summary: Record training runs locally and push to Aquin for post-hoc inspection — loss curves, SAE diffs, model diffs, and more.
+License: MIT
+Project-URL: Homepage, https://aquin.app
+Requires-Python: >=3.9
+Description-Content-Type: text/markdown
+Requires-Dist: requests>=2.28
+
+# Aquin SDK
+
+Record your training runs locally and push them to [Aquin](https://aquin.app) for post-hoc inspection — loss curves, learning rate, grad norm, epoch summaries, SAE feature diffs, model behaviour diffs, and more.
+
+## Install
+
+```bash
+pip install aquin
+```
+
+## Quickstart
+
+```python
+import aquin
+
+run = aquin.init(
+    base_model="meta-llama/Llama-3.2-1B-Instruct",
+    run_name="my-lora-run",
+    config={
+        "lr": 2e-4, "epochs": 3, "rank": 16, "lora_alpha": 32,
+        "method": "qlora", "per_device_train_batch_size": 2,
+        "gradient_accumulation_steps": 8, "dataset": "data.jsonl",
+    },
+)
+
+for epoch in range(3):
+    for step, batch in enumerate(dataloader):
+        loss = train_step(batch)
+        grad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0).item()
+        run.log(
+            step,
+            loss=loss.item(),
+            learning_rate=scheduler.get_last_lr()[0],
+            grad_norm=grad_norm,
+            epoch=epoch,
+        )
+
+run.checkpoint(model, step=step)
+run.finish()
+```
+
+Then push:
+
+```bash
+aquin package
+aquin push
+```
+
+Your run appears in the Aquin dashboard under **CLI runs** with the full inspection suite.
+
+## API
+
+### `aquin.init(base_model, run_name, config)`
+
+Starts a new run. Creates `aquin_run/` in the current directory.
+
+| Param | Description |
+|---|---|
+| `base_model` | HuggingFace model ID, e.g. `"meta-llama/Llama-3.2-1B-Instruct"` |
+| `run_name` | Display name for the run |
+| `config` | Dict of training hyperparameters (optional, can also pass to `finish()`) |
+
+### `run.log(step, *, loss, ...)`
+
+Record metrics for one training step. Call every step inside your loop.
+
+| Param | Description |
+|---|---|
+| `step` | Global training step (required) |
+| `loss` | Scalar training loss (required) |
+| `learning_rate` | Current LR — enables LR chart |
+| `grad_norm` | Gradient norm — enables grad norm chart |
+| `epoch` | Current epoch — enables epoch summary table |
+| `momentum_norm` | Optimizer momentum norm — enables momentum chart |
+| `step_ms` | Wall-clock time for this step in ms |
+
+### `run.checkpoint(model, step)`
+
+Saves the model checkpoint locally. One checkpoint per run — always replaces the previous save. Call once at the end of training. The checkpoint is included in the push and used for SAE diff and model diff analysis.
+
+### `run.finish(config)`
+
+Flushes all metrics to disk. Pass `config` here if you didn't pass it to `aquin.init()`.
+
+## CLI
+
+```bash
+aquin login       # save your API key
+aquin package     # bundle aquin_run/ into aquin_run.tar.gz
+aquin push        # push to Aquin
+aquin whoami      # check which account you're logged in as
+```
+
+## Using with HuggingFace Trainer / TRL
+
+Use a `TrainerCallback` to hook into the training loop:
+
+```python
+import time
+from transformers import TrainerCallback
+
+class AquinCallback(TrainerCallback):
+    def __init__(self, run):
+        self.run = run
+        self._step_start = 0.0
+
+    def on_step_begin(self, args, state, control, **kwargs):
+        self._step_start = time.time()
+
+    def on_log(self, args, state, control, logs=None, **kwargs):
+        if not logs or "loss" not in logs:
+            return
+        self.run.log(
+            step=state.global_step,
+            loss=float(logs["loss"]),
+            learning_rate=float(logs["learning_rate"]) if "learning_rate" in logs else None,
+            grad_norm=float(logs["grad_norm"]) if "grad_norm" in logs else None,
+            epoch=int(state.epoch) if state.epoch is not None else None,
+            step_ms=round((time.time() - self._step_start) * 1000),
+        )
+
+    def on_train_end(self, args, state, control, **kwargs):
+        model = kwargs.get("model")
+        if model:
+            self.run.checkpoint(model, step=state.global_step)
+```

aquin-0.0.1.dist-info/RECORD

@@ -0,0 +1,8 @@
+aquin/__init__.py,sha256=FraX6aNKjZXE6nlP7kcpYnIsgZ7ybKmVgaBpdF_8_y4,54
+aquin/cli.py,sha256=kW2sWBLyFF9yZQJIg3rEeN5JSVQwPczBHUyao60mWis,5901
+aquin/run.py,sha256=2TtOIWKTe1yT8TG3BZ-MQRbwnBRiG_d3xrPvZtt8_vI,7317
+aquin-0.0.1.dist-info/METADATA,sha256=tWZG32rvIfchvkWYxjPGO2RiTqpUmQwreDBy3nRvU2M,4251
+aquin-0.0.1.dist-info/WHEEL,sha256=aeYiig01lYGDzBgS8HxWXOg3uV61G9ijOsup-k9o1sk,91
+aquin-0.0.1.dist-info/entry_points.txt,sha256=G2OIiZIvNTHNXeuMnkTooZLLDDIEdYsABFP3eQEfHQ0,41
+aquin-0.0.1.dist-info/top_level.txt,sha256=_z0uFSWkloDciYFBw0MXB_sIID7Q1EsVGU234HDrZps,6
+aquin-0.0.1.dist-info/RECORD,,

aquin-0.0.1.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (82.0.1)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

aquin-0.0.1.dist-info/entry_points.txt

@@ -0,0 +1,2 @@
+[console_scripts]
+aquin = aquin.cli:main

aquin-0.0.1.dist-info/top_level.txt

@@ -0,0 +1 @@
+aquin