heulistic 0.1.0__py3-none-any.whl

heulistic/__main__.py

@@ -0,0 +1,11 @@
+"""Entry point for `python -m heulistic`.
+
+The `heulistic` console script is the normal way in; this makes the module form
+work too, which is how the CLI gets run inside a virtualenv that has not been
+activated.
+"""
+
+from .main import main
+
+if __name__ == "__main__":
+    main()

heulistic/client.py

@@ -0,0 +1,161 @@
+"""Talk to the local Heulistic API.
+
+Standard library only, so `pip install heulistic` stays light — which is why
+multipart bodies are assembled here rather than handed to a library.
+
+Local mode accepts a fixed bearer token in place of a Cognito JWT. When cloud
+support lands this is the seam where a real credential goes.
+"""
+
+from __future__ import annotations
+
+import json
+import mimetypes
+import secrets
+import urllib.error
+import urllib.request
+from pathlib import Path
+from typing import Iterator
+
+__all__ = [
+    "ApiError",
+    "encode_multipart",
+    "submit",
+    "get_job",
+    "download_url",
+    "stream_logs",
+]
+
+DEV_TOKEN = "dev-localuser"
+_TIMEOUT = 30
+# Training runs are long; the log stream must not time out mid-job.
+_STREAM_TIMEOUT = 60 * 60 * 24
+
+
+class ApiError(Exception):
+    """The API refused the request, carrying its `detail` where there is one."""
+
+
+def _headers(content_type: str | None = None) -> dict[str, str]:
+    headers = {"Authorization": f"Bearer {DEV_TOKEN}"}
+    if content_type:
+        headers["Content-Type"] = content_type
+    return headers
+
+
+def _raise_for_error(exc: urllib.error.HTTPError) -> None:
+    try:
+        detail = json.loads(exc.read()).get("detail", "")
+    except Exception:
+        detail = ""
+    raise ApiError(detail or f"HTTP {exc.code}") from exc
+
+
+def _request(
+    method: str,
+    url: str,
+    *,
+    body: bytes | None = None,
+    content_type: str | None = None,
+    stream: bool = False,
+):
+    """Perform one HTTP request. Tests replace this wholesale."""
+    request = urllib.request.Request(
+        url, data=body, method=method, headers=_headers(content_type)
+    )
+    try:
+        response = urllib.request.urlopen(
+            request, timeout=_STREAM_TIMEOUT if stream else _TIMEOUT
+        )
+    except urllib.error.HTTPError as exc:
+        _raise_for_error(exc)
+    except urllib.error.URLError as exc:
+        raise ApiError(
+            f"Could not reach the Heulistic API at {url} — is it running? "
+            "Start it with: heulistic start"
+        ) from exc
+    return response if stream else response.read()
+
+
+def encode_multipart(
+    fields: dict[str, str | None], files: dict[str, tuple[str, bytes]]
+) -> tuple[str, bytes]:
+    """Build a multipart/form-data body. Returns (content_type, body)."""
+    boundary = secrets.token_hex(16)
+    parts: list[bytes] = []
+
+    for name, value in fields.items():
+        if value in (None, ""):
+            continue  # an empty optional field must not be sent at all
+        parts.append(
+            f"--{boundary}\r\n"
+            f'Content-Disposition: form-data; name="{name}"\r\n\r\n'
+            f"{value}\r\n".encode()
+        )
+
+    for name, (filename, content) in files.items():
+        mime = mimetypes.guess_type(filename)[0] or "application/octet-stream"
+        parts.append(
+            f"--{boundary}\r\n"
+            f'Content-Disposition: form-data; name="{name}"; filename="{filename}"\r\n'
+            f"Content-Type: {mime}\r\n\r\n".encode()
+        )
+        parts.append(content)
+        parts.append(b"\r\n")
+
+    parts.append(f"--{boundary}--\r\n".encode())
+    return f"multipart/form-data; boundary={boundary}", b"".join(parts)
+
+
+def submit(
+    api_url: str,
+    config: Path,
+    dataset: Path | None = None,
+    hf_token: str | None = None,
+) -> dict:
+    """Submit a training job. Returns the API's job record."""
+    files = {"config": (config.name, config.read_bytes())}
+    if dataset is not None:
+        files["dataset"] = (dataset.name, dataset.read_bytes())
+
+    content_type, body = encode_multipart({"hf_token": hf_token}, files)
+    raw = _request(
+        "POST", f"{api_url}/api/jobs", body=body, content_type=content_type
+    )
+    return json.loads(raw)
+
+
+def get_job(api_url: str, job_id: str) -> dict:
+    return json.loads(_request("GET", f"{api_url}/api/jobs/{job_id}"))
+
+
+def download_url(api_url: str, job_id: str) -> str:
+    raw = _request("GET", f"{api_url}/api/jobs/{job_id}/download")
+    return json.loads(raw)["url"]
+
+
+def _stream_lines(url: str) -> Iterator[bytes]:
+    response = _request("GET", url, stream=True)
+    for raw_line in response:
+        yield raw_line.rstrip(b"\n")
+
+
+def stream_logs(api_url: str, job_id: str) -> Iterator[tuple[str, str]]:
+    """Yield ("line", text) for each log line, then ("done", status) at the end.
+
+    Server-sent events; anything that is not a well-formed `data:` frame — a
+    keepalive, a blank separator — is skipped rather than surfaced.
+    """
+    url = f"{api_url}/api/jobs/{job_id}/logs"
+    for raw_line in _stream_lines(url):
+        if not raw_line.startswith(b"data: "):
+            continue
+        try:
+            payload = json.loads(raw_line[len(b"data: "):])
+        except json.JSONDecodeError:
+            continue
+        if payload.get("event") == "done":
+            yield "done", payload.get("status", "")
+            return
+        if "line" in payload:
+            yield "line", payload["line"]

heulistic/download_cmd.py

@@ -0,0 +1,92 @@
+"""`heulistic download` — fetch a finished model.
+
+Locally the artifact sits in MinIO inside Docker, so without this the output of
+a run is hard to reach. The API hands back a URL; this follows it.
+"""
+
+from __future__ import annotations
+
+import shutil
+import urllib.error
+import urllib.request
+from pathlib import Path
+from typing import Optional
+
+import typer
+
+from . import client, stack
+
+_SUFFIX = "-model.tar.gz"
+
+
+def _fetch(url: str, dest: Path) -> Path:
+    """Stream the artifact to disk. Replaced in tests."""
+    with urllib.request.urlopen(url, timeout=300) as response, open(dest, "wb") as out:
+        shutil.copyfileobj(response, out)
+    return dest
+
+
+def _resolve_destination(output: Optional[Path], job_id: str) -> Path:
+    default_name = f"{job_id}{_SUFFIX}"
+    if output is None:
+        return Path.cwd() / default_name
+    if output.is_dir():
+        return output / default_name
+    return output
+
+
+def download(
+    job_id: str = typer.Argument(..., help="Job ID to download the model for."),
+    output: Optional[Path] = typer.Option(
+        None, "--output", "-o", help="File or directory to write to."
+    ),
+    force: bool = typer.Option(False, "--force", "-f", help="Overwrite an existing file."),
+    url_only: bool = typer.Option(
+        False, "--url-only", help="Print the download URL instead of fetching it."
+    ),
+) -> None:
+    """Download the trained model for a completed job."""
+    api_url = stack.api_url()
+    if not api_url:
+        typer.secho("Heulistic is not running.", fg=typer.colors.RED)
+        typer.echo("Start it with: heulistic start")
+        raise typer.Exit(code=1)
+
+    try:
+        job = client.get_job(api_url, job_id)
+    except client.ApiError as exc:
+        typer.secho(str(exc), fg=typer.colors.RED)
+        raise typer.Exit(code=1)
+
+    status = job.get("status", "unknown")
+    if status != "completed":
+        typer.secho(f"Job {job_id} is '{status}' — a model is only available once "
+                    "the job has completed.", fg=typer.colors.RED)
+        raise typer.Exit(code=1)
+
+    try:
+        url = client.download_url(api_url, job_id)
+    except client.ApiError as exc:
+        typer.secho(str(exc), fg=typer.colors.RED)
+        raise typer.Exit(code=1)
+
+    if url_only:
+        typer.echo(url)
+        raise typer.Exit(code=0)
+
+    dest = _resolve_destination(output, job_id)
+    if dest.exists() and not force:
+        typer.secho(f"{dest} already exists — pass --force to overwrite.",
+                    fg=typer.colors.RED)
+        raise typer.Exit(code=1)
+
+    typer.echo(f"Downloading {job_id} → {dest}")
+    try:
+        _fetch(url, dest)
+    except (urllib.error.URLError, OSError) as exc:
+        typer.secho(f"Download failed: {exc}", fg=typer.colors.RED)
+        raise typer.Exit(code=1)
+
+    size_mb = dest.stat().st_size / (1024 * 1024)
+    typer.secho(f"Saved {dest} ({size_mb:.1f} MB)", fg=typer.colors.GREEN)
+    typer.echo(f"  tar xzf {dest.name}")

heulistic/gpu.py

@@ -0,0 +1,90 @@
+"""Local GPU detection via nvidia-smi.
+
+Used to check whether a config fits the machine it is about to run on. Every
+path here degrades to "no GPU found" rather than raising — running
+`validate-config` on a laptop without an NVIDIA card is ordinary, not an error.
+"""
+
+from __future__ import annotations
+
+import shutil
+import subprocess
+from dataclasses import dataclass
+
+from heulistic_axolotl import Hardware
+
+__all__ = ["DetectedGpu", "detect_gpus", "detect_hardware"]
+
+_MB_PER_GB = 1024
+_QUERY = ["--query-gpu=name,memory.total", "--format=csv,noheader,nounits"]
+
+
+@dataclass(frozen=True)
+class DetectedGpu:
+    name: str
+    vram_mb: int
+
+
+def _nvidia_smi_output() -> str | None:
+    """Raw `nvidia-smi` CSV, or None when unavailable.
+
+    Isolated so tests can supply canned output without mocking subprocess.
+    """
+    if shutil.which("nvidia-smi") is None:
+        return None
+    try:
+        result = subprocess.run(
+            ["nvidia-smi", *_QUERY],
+            capture_output=True,
+            text=True,
+            timeout=10,
+        )
+    except (OSError, subprocess.SubprocessError):
+        return None
+    return result.stdout if result.returncode == 0 else None
+
+
+def detect_gpus() -> list[DetectedGpu]:
+    """Every NVIDIA GPU on this machine. Empty when none can be found."""
+    output = _nvidia_smi_output()
+    if not output:
+        return []
+
+    cards: list[DetectedGpu] = []
+    for line in output.splitlines():
+        line = line.strip()
+        if not line or "," not in line:
+            continue
+        name, _, memory = line.rpartition(",")
+        try:
+            vram_mb = int(memory.strip())
+        except ValueError:
+            continue  # a driver error lands in this column instead of a number
+        cards.append(DetectedGpu(name=name.strip(), vram_mb=vram_mb))
+    return cards
+
+
+def _label(cards: list[DetectedGpu]) -> str:
+    names = {c.name for c in cards}
+    if len(cards) == 1:
+        return cards[0].name
+    if len(names) == 1:
+        return f"{len(cards)}x {cards[0].name}"
+    return f"{len(cards)} GPUs"
+
+
+def detect_hardware() -> Hardware | None:
+    """The local GPUs as a validation target, or None if there are none usable."""
+    cards = detect_gpus()
+    if not cards:
+        return None
+    total_gb = sum(c.vram_mb for c in cards) / _MB_PER_GB
+    try:
+        return Hardware(
+            total_vram_gb=round(total_gb, 2),
+            gpu_count=len(cards),
+            label=_label(cards),
+        )
+    except ValueError:
+        # A device reporting zero memory is not something we can size against.
+        return None

heulistic/init_cmd.py

@@ -0,0 +1,152 @@
+"""`heulistic init` — write a starter config sized to this machine.
+
+A generic sample config is a poor first experience: the usual failure is an
+OOM twenty minutes in, because the example assumed a bigger card. Since the
+CLI can already see the local GPU and estimate VRAM, the starter config picks
+a model that actually fits.
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Optional
+
+import typer
+
+from heulistic_axolotl import validate
+from heulistic_axolotl.vram import vram_required_gb
+
+from . import gpu
+
+# QLoRA-friendly models, smallest first. The largest entry that fits the card
+# wins.
+#
+# The ladder deliberately stops at 8B rather than climbing to 70B. This is a
+# starter config: the goal is a run that finishes quickly and works, not one
+# that fills the GPU. Someone with an 80 GB card can pass --model.
+_LADDER = [
+    "Qwen/Qwen2.5-0.5B",
+    "Qwen/Qwen2.5-1.5B",
+    "Qwen/Qwen2.5-3B",
+    "Qwen/Qwen2.5-7B",
+]
+# Every entry must be ungated. The ladder previously led with Llama-3.2, which
+# is gated on the Hub — so the first `heulistic train` after `init` was rejected
+# for want of an approved HuggingFace token, which is precisely the wall `init`
+# exists to avoid. Qwen2.5 is Apache-2.0 and needs no token.
+_FALLBACK_MODEL = "Qwen/Qwen2.5-0.5B"
+
+# Leave room for activations, fragmentation and the desktop's own usage.
+_USABLE_FRACTION = 0.75
+
+
+def _template(base_model: str) -> dict:
+    """A complete, warning-free QLoRA config.
+
+    Every field the validator would otherwise flag as missing is set
+    explicitly, so `heulistic validate-config` on a freshly generated file is
+    silent — a starter config that immediately complains teaches the wrong
+    thing.
+    """
+    return {
+        "base_model": base_model,
+        "adapter": "qlora",
+        "load_in_4bit": True,
+        "datasets": [{"path": "tatsu-lab/alpaca", "type": "alpaca"}],
+        "sequence_len": 1024,
+        "micro_batch_size": 1,
+        "gradient_accumulation_steps": 8,
+        "num_epochs": 1,
+        "learning_rate": 0.0002,
+        "lora_r": 16,
+        "lora_alpha": 32,
+        "lora_dropout": 0.05,
+        "lora_target_linear": True,
+        "bf16": "auto",
+        "gradient_checkpointing": True,
+        "save_steps": 100,
+        "val_set_size": 0.05,
+        "output_dir": "./outputs",
+    }
+
+
+def _pick_model(budget_gb: float | None) -> str:
+    """Largest model from the ladder that fits the budget."""
+    if budget_gb is None:
+        return _FALLBACK_MODEL
+    chosen = None
+    for model in _LADDER:
+        if vram_required_gb(_template(model)) <= budget_gb:
+            chosen = model
+    return chosen or _LADDER[0]
+
+
+def _render(config: dict) -> str:
+    """Emit YAML by hand, so the file is commented and ordered for a human."""
+    import yaml
+
+    header = (
+        "# Generated by `heulistic init`.\n"
+        "# Check it before training:  heulistic validate-config config.yaml\n"
+        "# Then run it:               heulistic train config.yaml\n\n"
+    )
+    return header + yaml.safe_dump(config, sort_keys=False, default_flow_style=False)
+
+
+def init(
+    output: Path = typer.Option(
+        Path("config.yaml"), "--output", "-o", help="Where to write the config."
+    ),
+    model: Optional[str] = typer.Option(
+        None, "--model", "-m", help="Base model to use instead of the auto-picked one."
+    ),
+    force: bool = typer.Option(False, "--force", "-f", help="Overwrite an existing file."),
+) -> None:
+    """Write a starter Axolotl config, sized to the local GPU."""
+    if output.exists() and not force:
+        typer.secho(f"{output} already exists — pass --force to overwrite.",
+                    fg=typer.colors.RED)
+        raise typer.Exit(code=1)
+
+    hardware = gpu.detect_hardware()
+    if hardware is None:
+        typer.secho(
+            "No NVIDIA GPU detected — writing a small starter config. "
+            "Use --model to choose another.",
+            fg=typer.colors.YELLOW,
+        )
+        budget = None
+    else:
+        budget = hardware.total_vram_gb * _USABLE_FRACTION
+        typer.echo(
+            f"Detected {hardware.label} ({hardware.total_vram_gb:g} GB) — "
+            f"sizing for roughly {budget:.0f} GB of usable VRAM."
+        )
+
+    base_model = model or _pick_model(budget)
+    config = _template(base_model)
+    required = vram_required_gb(config)
+
+    if model and hardware is not None and required > hardware.total_vram_gb:
+        typer.secho(
+            f"  Warning: {model} needs roughly {required:.0f} GB of VRAM but "
+            f"{hardware.label} has {hardware.total_vram_gb:g} GB.",
+            fg=typer.colors.YELLOW,
+        )
+
+    output.write_text(_render(config))
+
+    typer.secho(f"Wrote {output}", fg=typer.colors.GREEN)
+    typer.echo(f"  model: {base_model}  (~{required:.1f} GB VRAM as QLoRA)")
+    typer.echo(f"  data:  tatsu-lab/alpaca — swap in your own under 'datasets'")
+    typer.echo()
+    typer.echo(f"Next:  heulistic train {output}")
+
+    # Never ship a starter file our own validator would complain about.
+    problems = [d for d in validate(output.read_text()) if d.level in ("error", "warning")]
+    if problems:  # pragma: no cover - guards against template drift
+        typer.secho(
+            f"  ({len(problems)} validator finding(s) in the generated file — "
+            "please report this)",
+            fg=typer.colors.YELLOW,
+        )

heulistic/local_compose.yml

@@ -0,0 +1,101 @@
+name: heulistic
+
+services:
+  db:
+    image: postgres:15
+    environment:
+      POSTGRES_DB: heulistic
+      POSTGRES_USER: heulistic_admin
+      POSTGRES_PASSWORD: password
+    volumes:
+      - postgres_data:/var/lib/postgresql/data
+    healthcheck:
+      test: ["CMD-SHELL", "pg_isready -U heulistic_admin -d heulistic"]
+      interval: 5s
+      timeout: 5s
+      retries: 5
+
+  minio:
+    image: minio/minio
+    command: server /data --console-address ":9001"
+    environment:
+      MINIO_ROOT_USER: minioadmin
+      MINIO_ROOT_PASSWORD: minioadmin
+    ports:
+      - "9000:9000"
+      - "9001:9001"
+    volumes:
+      - minio_data:/data
+
+  minio-init:
+    image: minio/mc
+    depends_on:
+      - minio
+    entrypoint: >
+      /bin/sh -c "
+        until mc alias set local http://minio:9000 minioadmin minioadmin; do sleep 1; done;
+        mc mb local/heulistic-local --ignore-existing;
+        exit 0"
+
+  api:
+    image: heulistic/api:latest
+    environment:
+      DATABASE_URL: postgresql+asyncpg://heulistic_admin:password@db:5432/heulistic
+      LOCAL_MODE: "true"
+      S3_BUCKET: heulistic-local
+      S3_ENDPOINT_URL: http://minio:9000
+      MINIO_ACCESS_KEY: minioadmin
+      MINIO_SECRET_KEY: minioadmin
+      FRONTEND_URL: http://localhost:${PORT:-3000}
+      # Download links are built from this, so it must match the published port.
+      API_PUBLIC_URL: http://localhost:${API_PORT:-8000}
+      AWS_REGION: us-east-1
+    ports:
+      - "${API_PORT:-8000}:8000"
+    depends_on:
+      db:
+        condition: service_healthy
+      minio-init:
+        condition: service_completed_successfully
+
+  docker-proxy:
+    image: tecnativa/docker-socket-proxy:edge
+    environment:
+      CONTAINERS: 1
+      IMAGES: 1
+      POST: 1
+    volumes:
+      - /var/run/docker.sock:/var/run/docker.sock:ro
+    restart: unless-stopped
+
+  scheduler:
+    image: heulistic/scheduler:latest
+    environment:
+      DATABASE_URL: postgresql://heulistic_admin:password@db:5432/heulistic
+      LOCAL_MODE: "true"
+      API_INTERNAL_URL: http://api:8000
+      LOCAL_TRAINING_IMAGE: heulistic/training:latest
+      DOCKER_NETWORK: heulistic_default
+      LOCAL_CACHE_VOLUME: heulistic-cache
+      LOCAL_GPU: "${LOCAL_GPU:-true}"
+      MAX_CONCURRENT_JOBS: "1"
+      HF_TOKEN: "${HF_TOKEN:-}"
+      AWS_REGION: us-east-1
+      DOCKER_SOCKET: tcp://docker-proxy:2375
+    depends_on:
+      - api
+      - docker-proxy
+
+  frontend:
+    image: heulistic/frontend-local:latest
+    environment:
+      API_URL: http://api:8000
+    ports:
+      - "${PORT:-3000}:3000"
+    depends_on:
+      - api
+
+volumes:
+  postgres_data:
+  minio_data:
+  heulistic-cache: