synth-ai 0.2.8.dev11__py3-none-any.whl → 0.2.8.dev13__py3-none-any.whl

synth_ai/api/train/__init__.py

@@ -0,0 +1,5 @@
+from __future__ import annotations
+
+from .cli import register, train_command
+
+__all__ = ["register", "train_command"]

synth_ai/api/train/builders.py

@@ -0,0 +1,165 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any
+
+import click
+
+from .utils import ensure_api_base, load_toml, TrainError
+
+
+@dataclass(slots=True)
+class RLBuildResult:
+    payload: dict[str, Any]
+    task_url: str
+    idempotency: str | None
+
+
+@dataclass(slots=True)
+class SFTBuildResult:
+    payload: dict[str, Any]
+    train_file: Path
+    validation_file: Path | None
+
+
+def build_rl_payload(
+    *,
+    config_path: Path,
+    task_url: str,
+    overrides: dict[str, Any],
+    idempotency: str | None,
+) -> RLBuildResult:
+    data = load_toml(config_path)
+    services = data.get("services") if isinstance(data.get("services"), dict) else {}
+    model_cfg = data.get("model") if isinstance(data.get("model"), dict) else {}
+
+    final_task_url = (overrides.get("task_url") or task_url or services.get("task_url") or "").strip()
+    if not final_task_url:
+        raise click.ClickException("Task app URL required (provide --task-url or set services.task_url in TOML)")
+
+    model_source = (model_cfg.get("source") or "").strip()
+    model_base = (model_cfg.get("base") or "").strip()
+    override_model = (overrides.get("model") or "").strip()
+    if override_model:
+        model_source = override_model
+        model_base = ""
+    if bool(model_source) == bool(model_base):
+        raise click.ClickException("Model section must specify exactly one of [model].source or [model].base")
+
+    # Force TOML services.task_url to the effective endpoint to avoid split URLs
+    try:
+        if isinstance(data.get("services"), dict):
+            data["services"]["task_url"] = final_task_url
+        else:
+            data["services"] = {"task_url": final_task_url}
+    except Exception:
+        pass
+
+    payload: dict[str, Any] = {
+        "job_type": "rl",
+        "compute": data.get("compute", {}),
+        "data": {
+            "endpoint_base_url": final_task_url.rstrip("/"),
+            "config": data,
+        },
+        "tags": {"source": "train-cli"},
+    }
+    if model_source:
+        payload["data"]["model"] = model_source
+    if model_base:
+        payload["data"]["base_model"] = model_base
+
+    backend = overrides.get("backend")
+    if backend:
+        payload.setdefault("metadata", {})["backend_base_url"] = ensure_api_base(str(backend))
+
+    return RLBuildResult(payload=payload, task_url=final_task_url, idempotency=idempotency)
+
+
+def build_sft_payload(
+    *,
+    config_path: Path,
+    dataset_override: Path | None,
+) -> SFTBuildResult:
+    data = load_toml(config_path)
+    job_cfg = data.get("job") if isinstance(data.get("job"), dict) else {}
+    data_cfg = data.get("data") if isinstance(data.get("data"), dict) else {}
+    hp_cfg = data.get("hyperparameters") if isinstance(data.get("hyperparameters"), dict) else {}
+    train_cfg = data.get("training") if isinstance(data.get("training"), dict) else {}
+    compute_cfg = data.get("compute") if isinstance(data.get("compute"), dict) else {}
+
+    raw_dataset = dataset_override or job_cfg.get("data") or job_cfg.get("data_path")
+    if not raw_dataset:
+        raise TrainError("Dataset not specified; pass --dataset or set [job].data")
+    dataset_path = Path(raw_dataset)
+    dataset_path = (dataset_path if dataset_path.is_absolute() else (config_path.parent / dataset_path)).resolve()
+    if not dataset_path.exists():
+        raise TrainError(f"Dataset not found: {dataset_path}")
+
+    validation_path = data_cfg.get("validation_path") if isinstance(data_cfg.get("validation_path"), str) else None
+    validation_file = None
+    if validation_path:
+        vpath = Path(validation_path)
+        vpath = (vpath if vpath.is_absolute() else (config_path.parent / vpath)).resolve()
+        if not vpath.exists():
+            click.echo(f"[WARN] Validation dataset {vpath} missing; continuing without validation")
+        else:
+            validation_file = vpath
+
+    hp_block: dict[str, Any] = {
+        "n_epochs": int(hp_cfg.get("n_epochs", 1)),
+    }
+    for key in (
+        "batch_size",
+        "global_batch",
+        "per_device_batch",
+        "gradient_accumulation_steps",
+        "sequence_length",
+        "learning_rate",
+        "warmup_ratio",
+        "train_kind",
+    ):
+        if key in hp_cfg:
+            hp_block[key] = hp_cfg[key]
+    if isinstance(hp_cfg.get("parallelism"), dict):
+        hp_block["parallelism"] = hp_cfg["parallelism"]
+
+    compute_block = {k: compute_cfg[k] for k in ("gpu_type", "gpu_count", "nodes") if k in compute_cfg}
+
+    effective = {
+        "compute": compute_block,
+        "data": {"topology": data_cfg.get("topology", {}) if isinstance(data_cfg.get("topology"), dict) else {}},
+        "training": {k: v for k, v in train_cfg.items() if k in ("mode", "use_qlora")},
+    }
+
+    validation_cfg = train_cfg.get("validation") if isinstance(train_cfg.get("validation"), dict) else None
+    if isinstance(validation_cfg, dict):
+        hp_block.update(
+            {
+                "evaluation_strategy": validation_cfg.get("evaluation_strategy", "steps"),
+                "eval_steps": int(validation_cfg.get("eval_steps", 0) or 0),
+                "save_best_model_at_end": bool(validation_cfg.get("save_best_model_at_end", True)),
+                "metric_for_best_model": validation_cfg.get("metric_for_best_model", "val.loss"),
+                "greater_is_better": bool(validation_cfg.get("greater_is_better", False)),
+            }
+        )
+        effective.setdefault("training", {})["validation"] = {"enabled": bool(validation_cfg.get("enabled", True))}
+
+    payload = {
+        "model": job_cfg.get("model") or data.get("model"),
+        "training_file_id": None,  # populated after upload
+        "training_type": "sft_offline",
+        "hyperparameters": hp_block,
+        "metadata": {"effective_config": effective},
+    }
+
+    return SFTBuildResult(payload=payload, train_file=dataset_path, validation_file=validation_file)
+
+
+__all__ = [
+    "RLBuildResult",
+    "SFTBuildResult",
+    "build_rl_payload",
+    "build_sft_payload",
+]

synth_ai/api/train/cli.py

@@ -0,0 +1,429 @@
+from __future__ import annotations
+
+import os
+from pathlib import Path
+from typing import Any, Dict
+
+import click
+
+from .builders import RLBuildResult, SFTBuildResult, build_rl_payload, build_sft_payload
+from .config_finder import discover_configs, prompt_for_config
+from .env_resolver import KeySpec, resolve_env
+from .pollers import RLJobPoller, SFTJobPoller
+from .task_app import check_task_app_health
+from .utils import (
+    TrainError,
+    REPO_ROOT,
+    ensure_api_base,
+    http_post,
+    http_get,
+    limit_jsonl_examples,
+    mask_value,
+    post_multipart,
+    preview_json,
+    sleep,
+    validate_sft_jsonl,
+)
+
+
+def _discover_dataset_candidates(config_path: Path, limit: int = 50) -> list[Path]:
+    search_dirs: list[Path] = [
+        config_path.parent,
+        config_path.parent / "datasets",
+        REPO_ROOT / "traces",
+        REPO_ROOT / "datasets",
+    ]
+
+    candidates: list[Path] = []
+    seen: set[Path] = set()
+    for directory in search_dirs:
+        if not directory.exists() or not directory.is_dir():
+            continue
+        for path in directory.rglob("*.jsonl"):
+            try:
+                resolved = path.resolve()
+            except OSError:
+                continue
+            if resolved in seen:
+                continue
+            seen.add(resolved)
+            if resolved.stat().st_size == 0:
+                continue
+            candidates.append(resolved)
+            if len(candidates) >= limit:
+                return candidates
+    return candidates
+
+
+def prompt_for_dataset(config_path: Path) -> Path:
+    candidates = _discover_dataset_candidates(config_path)
+    while True:
+        if candidates:
+            click.echo("Select dataset JSONL file:")
+            for idx, candidate in enumerate(candidates, start=1):
+                click.echo(f"  {idx}) {candidate}")
+            click.echo("  m) Enter path manually")
+            click.echo("  0) Abort")
+            choice = click.prompt("Choice", default="m").strip().lower()
+            if choice == "0":
+                raise click.ClickException("Aborted by user")
+            if choice in {"m", "manual"}:
+                selected = _prompt_manual_dataset()
+            else:
+                try:
+                    idx = int(choice)
+                except ValueError:
+                    click.echo("Invalid selection; try again")
+                    continue
+                if idx < 1 or idx > len(candidates):
+                    click.echo("Invalid selection; try again")
+                    continue
+                selected = candidates[idx - 1]
+        else:
+            selected = _prompt_manual_dataset()
+
+        if selected.exists() and selected.suffix == ".jsonl":
+            return selected.resolve()
+        click.echo("File not found or not a .jsonl; please try again.")
+
+
+def _prompt_manual_dataset() -> Path:
+    manual = click.prompt("Enter dataset JSONL path", type=str).strip()
+    return Path(manual).expanduser()
+
+
+@click.command("train")
+@click.option("--config", "config_paths", multiple=True, type=click.Path(), help="Path to training TOML (repeatable)")
+@click.option("--type", "train_type", type=click.Choice(["auto", "rl", "sft"]), default="auto")
+@click.option("--env-file", "env_files", multiple=True, type=click.Path(), help=".env file(s) to preload (skips selection prompt)")
+@click.option("--task-url", default=None, help="Override task app base URL (RL only)")
+@click.option("--dataset", "dataset_path", type=click.Path(), default=None, help="Override dataset JSONL path (SFT)")
+@click.option("--backend", default=lambda: os.getenv("BACKEND_BASE_URL", "http://localhost:8000/api"), help="Backend base URL")
+@click.option("--model", default=None, help="Override model identifier")
+@click.option("--idempotency", default=None, help="Idempotency-Key header for job creation")
+@click.option("--dry-run", is_flag=True, help="Preview payload without submitting")
+@click.option("--poll/--no-poll", default=True, help="Poll job status until terminal state")
+@click.option("--poll-timeout", default=3600.0, type=float, help="Maximum seconds to poll before timing out")
+@click.option("--poll-interval", default=5.0, type=float, help="Seconds between poll attempts")
+@click.option("--examples", "examples_limit", type=int, default=None, help="Limit SFT training to the first N examples")
+def train_command(
+    config_paths: tuple[str, ...],
+    train_type: str,
+    env_files: tuple[str, ...],
+    task_url: str | None,
+    dataset_path: str | None,
+    backend: str,
+    model: str | None,
+    idempotency: str | None,
+    dry_run: bool,
+    poll: bool,
+    poll_timeout: float,
+    poll_interval: float,
+    examples_limit: int | None,
+) -> None:
+    """Interactive launcher for RL / SFT jobs."""
+
+    candidates = discover_configs(list(config_paths), requested_type=train_type if train_type != "auto" else None)
+    selection = prompt_for_config(candidates, requested_type=train_type if train_type != "auto" else None)
+
+    effective_type = train_type if train_type != "auto" else selection.train_type
+    if effective_type not in {"rl", "sft"}:
+        effective_type = click.prompt("Detected config type is ambiguous. Enter type", type=click.Choice(["rl", "sft"]))
+
+    cfg_path = selection.path
+    click.echo(f"Using config: {cfg_path} ({effective_type})")
+
+    required_keys: list[KeySpec] = []
+    if effective_type == "rl":
+        required_keys.append(KeySpec("SYNTH_API_KEY", "Synth API key for backend"))
+        required_keys.append(
+            KeySpec(
+                "ENVIRONMENT_API_KEY",
+                "Environment API key for task app",
+                allow_modal_secret=True,
+                modal_secret_pattern="env",
+            )
+        )
+        required_keys.append(
+            KeySpec(
+                "TASK_APP_URL",
+                "Task app base URL",
+                secret=False,
+                allow_modal_app=True,
+                optional=bool(task_url),
+            )
+        )
+    else:  # sft
+        required_keys.append(KeySpec("SYNTH_API_KEY", "Synth API key for backend"))
+
+    env_path, env_values = resolve_env(
+        config_path=cfg_path,
+        explicit_env_paths=env_files,
+        required_keys=required_keys,
+    )
+    click.echo(f"Using env file: {env_path}")
+
+    synth_key = env_values.get("SYNTH_API_KEY") or os.environ.get("SYNTH_API_KEY")
+    if not synth_key:
+        raise click.ClickException("SYNTH_API_KEY required")
+
+    backend_base = ensure_api_base(backend)
+    click.echo(f"Backend base: {backend_base} (key {mask_value(synth_key)})")
+
+    if effective_type == "rl":
+        handle_rl(
+            cfg_path=cfg_path,
+            backend_base=backend_base,
+            synth_key=synth_key,
+            task_url_override=task_url,
+            model_override=model,
+            idempotency=idempotency,
+            dry_run=dry_run,
+            poll=poll,
+            poll_timeout=poll_timeout,
+            poll_interval=poll_interval,
+        )
+    else:
+        dataset_override_path = Path(dataset_path).expanduser().resolve() if dataset_path else None
+        handle_sft(
+            cfg_path=cfg_path,
+            backend_base=backend_base,
+            synth_key=synth_key,
+            dataset_override=dataset_override_path,
+            dry_run=dry_run,
+            poll=poll,
+            poll_timeout=poll_timeout,
+            poll_interval=poll_interval,
+            examples_limit=examples_limit,
+        )
+
+
+def _wait_for_training_file(backend_base: str, api_key: str, file_id: str, *, timeout: float = 120.0) -> None:
+    url = f"{backend_base}/learning/files/{file_id}"
+    headers = {"Authorization": f"Bearer {api_key}"}
+    elapsed = 0.0
+    interval = 2.0
+    while True:
+        resp = http_get(url, headers=headers, timeout=30.0)
+        if resp.status_code == 200:
+            try:
+                data = resp.json()
+            except Exception:
+                data = {}
+            status = str(data.get("status") or data.get("state") or data.get("storage_state") or "ready").lower()
+            if status in {"ready", "uploaded", "stored", "complete"}:
+                return
+        elif resp.status_code == 404:
+            # Keep polling; object may not be visible yet
+            pass
+        else:
+            click.echo(f"[WARN] Unexpected response while checking training file {file_id}: {resp.status_code}")
+
+        if elapsed >= timeout:
+            raise click.ClickException(f"Training file {file_id} not ready after {timeout:.0f}s")
+        sleep(interval)
+        elapsed += interval
+
+
+def handle_rl(
+    *,
+    cfg_path: Path,
+    backend_base: str,
+    synth_key: str,
+    task_url_override: str | None,
+    model_override: str | None,
+    idempotency: str | None,
+    dry_run: bool,
+    poll: bool,
+    poll_timeout: float,
+    poll_interval: float,
+) -> None:
+    overrides: Dict[str, Any] = {"backend": backend_base, "task_url": task_url_override, "model": model_override}
+    build = build_rl_payload(
+        config_path=cfg_path,
+        task_url=task_url_override or os.environ.get("TASK_APP_URL", ""),
+        overrides=overrides,
+        idempotency=idempotency,
+    )
+
+    # Backend-side verification: try ALL org environment keys against /health and /task_info
+    verify_url = f"{backend_base}/rl/verify_task_app"
+    verify_headers = {"Authorization": f"Bearer {synth_key}", "Content-Type": "application/json"}
+    try:
+        vresp = http_post(verify_url, headers=verify_headers, json_body={"endpoint_base_url": build.task_url})
+        try:
+            vjs = vresp.json()
+        except Exception:
+            vjs = {"status": vresp.status_code, "text": (vresp.text or "")[:400]}
+    except Exception as _ve:
+        raise click.ClickException(f"Task app verification call failed: {type(_ve).__name__}: {_ve}") from _ve
+    if vresp.status_code >= 400:
+        click.echo("Task app verification error:\n" + preview_json(vjs, limit=800))
+        raise click.ClickException(f"Verification failed with status {vresp.status_code}")
+    if not bool(vjs.get("any_ok")):
+        click.echo("Task app verification failed; no auth combination succeeded. Full report:")
+        click.echo(preview_json(vjs, limit=1200))
+        raise click.ClickException("Task app verification failed (auth)")
+    else:
+        # Print concise summary
+        try:
+            cands = vjs.get("candidates_first15") or []
+            attempts = vjs.get("attempts") or []
+            statuses = [a.get("status") for a in attempts]
+            click.echo(f"Verification OK (candidates={cands}, statuses={statuses})")
+        except Exception:
+            pass
+
+    env_key = os.environ.get("ENVIRONMENT_API_KEY")
+    if not env_key:
+        raise click.ClickException("ENVIRONMENT_API_KEY required for RL flow")
+
+    click.echo("Performing task app health check…")
+    health = check_task_app_health(build.task_url, env_key)
+    if not health.ok:
+        click.echo(f"Task app health check failed: {health.detail}")
+        raise click.ClickException("Aborting due to failing health check")
+    else:
+        click.echo("Task app healthy")
+
+    create_url = f"{backend_base}/rl/jobs"
+    headers = {"Authorization": f"Bearer {synth_key}", "Content-Type": "application/json"}
+    if build.idempotency:
+        headers["Idempotency-Key"] = build.idempotency
+
+    click.echo(f"POST {create_url}")
+    click.echo("Payload preview:\n" + preview_json(build.payload, limit=800))
+    if dry_run:
+        click.echo("Dry run enabled; skipping submission")
+        return
+
+    resp = http_post(create_url, headers=headers, json_body=build.payload)
+    try:
+        js = resp.json()
+    except Exception:
+        js = {"status": resp.status_code, "text": resp.text[:400]}
+    click.echo(f"Response {resp.status_code}: {preview_json(js, limit=400)}")
+    if resp.status_code not in (200, 201):
+        raise click.ClickException("Job creation failed")
+    job_id = js.get("job_id") or js.get("id")
+    if not job_id:
+        raise click.ClickException("Response missing job id")
+
+    if not poll:
+        click.echo(f"Created job {job_id} (polling disabled)")
+        return
+
+    poller = RLJobPoller(backend_base, synth_key, interval=poll_interval, timeout=poll_timeout)
+    outcome = poller.poll_job(job_id)
+    click.echo(f"Final status: {outcome.status}")
+    click.echo(preview_json(outcome.payload, limit=600))
+
+
+def handle_sft(
+    *,
+    cfg_path: Path,
+    backend_base: str,
+    synth_key: str,
+    dataset_override: Path | None,
+    dry_run: bool,
+    poll: bool,
+    poll_timeout: float,
+    poll_interval: float,
+    examples_limit: int | None,
+) -> None:
+    dataset_path = dataset_override
+
+    while True:
+        try:
+            build = build_sft_payload(config_path=cfg_path, dataset_override=dataset_path)
+            break
+        except TrainError as exc:
+            click.echo(str(exc))
+            dataset_path = prompt_for_dataset(cfg_path)
+
+    limited_path: Path | None = None
+
+    try:
+        if examples_limit is not None:
+            limited_path = limit_jsonl_examples(build.train_file, examples_limit)
+            click.echo(
+                f"Using first {examples_limit} examples from {build.train_file} -> {limited_path}"
+            )
+            build.train_file = limited_path
+
+        click.echo("Validating training dataset…")
+        validate_sft_jsonl(build.train_file)
+        if build.validation_file and build.validation_file.suffix == ".jsonl":
+            click.echo("Validating validation dataset…")
+            validate_sft_jsonl(build.validation_file)
+
+        upload_url = f"{backend_base}/learning/files"
+        click.echo(f"Uploading dataset {build.train_file}")
+        if dry_run:
+            click.echo("Dry run: skipping upload")
+            train_file_id = "dry-run-train"
+            val_file_id = None
+        else:
+            resp = post_multipart(upload_url, api_key=synth_key, file_field="file", file_path=build.train_file)
+            js = resp.json() if resp.headers.get("content-type", "").startswith("application/json") else {}
+            if resp.status_code >= 400 or "id" not in js:
+                raise click.ClickException(f"Training file upload failed ({resp.status_code}): {js or resp.text[:200]}")
+            train_file_id = js["id"]
+            val_file_id = None
+            if build.validation_file:
+                click.echo(f"Uploading validation dataset {build.validation_file}")
+                vresp = post_multipart(upload_url, api_key=synth_key, file_field="file", file_path=build.validation_file)
+                vjs = vresp.json() if vresp.headers.get("content-type", "").startswith("application/json") else {}
+                if vresp.status_code < 400 and "id" in vjs:
+                    val_file_id = vjs["id"]
+                else:
+                    click.echo(f"[WARN] Validation upload failed: {vresp.status_code} {vjs or vresp.text[:200]}")
+        payload = dict(build.payload)
+        payload["training_file_id"] = train_file_id
+        if val_file_id:
+            payload.setdefault("metadata", {}).setdefault("effective_config", {}).setdefault("data", {})["validation_files"] = [val_file_id]
+
+        try:
+            _wait_for_training_file(backend_base, synth_key, train_file_id)
+        except click.ClickException as exc:
+            raise click.ClickException(f"Training file {train_file_id} not ready: {exc}") from exc
+
+        click.echo("FFT job payload:\n" + preview_json(payload, limit=800))
+        if dry_run:
+            click.echo("Dry run: skipping job submission")
+            return
+
+        create_url = f"{backend_base}/learning/jobs"
+        headers = {"Authorization": f"Bearer {synth_key}", "Content-Type": "application/json"}
+        resp = http_post(create_url, headers=headers, json_body=payload)
+        js = resp.json() if resp.headers.get("content-type", "").startswith("application/json") else {}
+        click.echo(f"Response {resp.status_code}: {preview_json(js, limit=400)}")
+        if resp.status_code not in (200, 201):
+            raise click.ClickException("Failed to create learning job")
+        job_id = js.get("job_id") or js.get("id")
+        if not job_id:
+            raise click.ClickException("Response missing job id")
+
+        start_url = f"{backend_base}/learning/jobs/{job_id}/start"
+        click.echo(f"POST {start_url} (start)")
+        _ = http_post(start_url, headers=headers, json_body={})
+
+        if not poll:
+            click.echo(f"Started job {job_id} (polling disabled)")
+            return
+
+        poller = SFTJobPoller(backend_base, synth_key, interval=poll_interval, timeout=poll_timeout)
+        outcome = poller.poll_job(job_id)
+        click.echo(f"Final status: {outcome.status}")
+        click.echo(preview_json(outcome.payload, limit=600))
+    finally:
+        if limited_path is not None:
+            try:
+                limited_path.unlink(missing_ok=True)
+                limited_path.parent.rmdir()
+            except Exception:
+                pass
+
+
+def register(cli: click.Group) -> None:
+    cli.add_command(train_command)