Blame for RSS_AI_Engine/Integration_Layer/engine_orchestrator.py

"""Integration layer: the ONLY place that wires the modules together.

It imports exactly four public interfaces - data_api, hardware_api, training_api, serving_api - and never
touches anything inside an `internal/` folder.

Purge policy (PURGE_OLD_MODEL_ON_START=true or --fresh-model):
  train_all / pretrain  start a NEW training lineage  -> old artifacts are purged and the base model re-downloaded
  clean / finetune / serve                              never purge (finetune needs the stage-1 checkpoint,
                                                        serve needs the trained model) unless --fresh-model is given
"""
from __future__ import annotations

import json
import os
import sys
import time
from pathlib import Path
from typing import Dict, List, Optional

ROOT = Path(__file__).resolve().parent.parent
if str(ROOT) not in sys.path:
    sys.path.insert(0, str(ROOT))

from Data_Module import data_api                      # noqa: E402
from Hardware_Module import hardware_api             # noqa: E402
from Serving_Module import serving_api               # noqa: E402
from Training_Module import training_api             # noqa: E402

DEFAULT_SOURCES = [{"repo": "Biraboro/Youtop-trenscprit", "limit": 150}]


# ------------------------------------------------------------------ configuration
def _path(root: Path, value: str) -> Path:
    p = Path(value).expanduser()
    return p if p.is_absolute() else (root / p).resolve()


def _parse_sources(raw: str) -> List[Dict]:
    raw = (raw or "").strip()
    if not raw:
        return list(DEFAULT_SOURCES)
    try:
        data = json.loads(raw)
    except ValueError:
        data = [{"repo": part.strip(), "limit": 150} for part in raw.split(",") if part.strip()]
    if isinstance(data, dict):
        data = [data]
    out = []
    for item in data if isinstance(data, list) else []:
        if isinstance(item, str):
            out.append({"repo": item, "limit": 150})
        elif isinstance(item, dict) and item.get("repo"):
            out.append({"repo": str(item["repo"]), "limit": item.get("limit", 150)})
    if not out:
        print("[Config] SOURCE_REPOS could not be parsed -> using the default repo")
        return list(DEFAULT_SOURCES)
    return out


def _flag(value: str, default: bool = False) -> bool:
    return str(value if value is not None else default).strip().lower() in ("1", "true", "yes", "y", "on")


def _float(value: str, default: float) -> float:
    try:
        return float(value)
    except (TypeError, ValueError):
        return default


def build_config(args=None) -> Dict:
    env_file = ROOT / ".env"
    try:
        from dotenv import load_dotenv
        load_dotenv(env_file)
    except ImportError:                                       # minimal fallback parser
        if env_file.exists():
            for line in env_file.read_text(encoding="utf-8").splitlines():
                if "=" in line and not line.strip().startswith("#"):
                    key, _, val = line.partition("=")
                    os.environ.setdefault(key.strip(), val.strip().strip("'\""))
    env = os.environ.get
    tune: Dict = {}
    raw_tune = env("TUNE_JSON", "").strip()
    if raw_tune:
        try:
            tune = json.loads(raw_tune)
        except ValueError:
            print("[Config] TUNE_JSON is not valid JSON - ignored")
    extra = [e.strip() for e in (env("EXTRA_BLOCKED_EXTENSIONS", "") or "").replace(";", ",").split(",") if e.strip()]
    cfg: Dict = {
        "root": str(ROOT),
        "hf_token": (env("HF_TOKEN", "") or "").strip(),
        "telegram_token": (env("TELEGRAM_BOT_TOKEN", "") or "").strip(),
        "telegram_chat_id": (env("TELEGRAM_CHAT_ID", "") or "").strip(),
        "base_model_repo": (env("BASE_MODEL_REPO", "Qwen/Qwen2.5-0.5B-Instruct") or "").strip(),
        "output_hf_repo": (env("OUTPUT_HF_REPO", "Biraboro/rss-ai-trained") or "").strip(),
        "source_repos": _parse_sources(env("SOURCE_REPOS", "")),
        "chat_file": str(_path(ROOT, env("LOCAL_CHAT_FILE", "./chat.jsonl") or "./chat.jsonl")),
        "json_quota": _float(env("JSON_PRETRAIN_QUOTA", "0.15"), 0.15),
        "conv_ratio": _float(env("PRETRAIN_CONV_RATIO", "0.25"), 0.25),
        "purge_old_model": _flag(env("PURGE_OLD_MODEL_ON_START", "true"), True),
        "extra_blocked_extensions": extra,
        "tune": tune,
        "force_tier": None,
        "cpu_threads": None,
        "limit_override": None,
    }
    if args is not None:
        if getattr(args, "limit", None):
            cfg["limit_override"] = args.limit
        if getattr(args, "tier", None):
            cfg["force_tier"] = args.tier
        if getattr(args, "threads", None):
            cfg["cpu_threads"] = args.threads
    return cfg


def _dirs(cfg: Dict) -> Dict[str, Path]:
    root = Path(cfg["root"])
    return {"stage1": root / "checkpoints" / "stage1", "final": root / "trained_model", "root": root}


def _valid_stage1_checkpoint(folder: Path) -> bool:
    """Validate Stage-1 metadata AND the actual weight payload before any resume decision."""
    folder = Path(folder)
    meta = training_api.read_meta(folder)
    if not meta or int(meta.get("stage", 0)) != 1:
        return False
    kind = meta.get("kind")
    if kind == "full":
        if not (folder / "config.json").is_file():
            return False
        single = folder / "model.safetensors"
        index = folder / "model.safetensors.index.json"
        if single.is_file() and single.stat().st_size > 0:
            return True
        if index.is_file():
            try:
                weight_map = json.loads(index.read_text(encoding="utf-8")).get("weight_map", {})
                shards = {str(name) for name in weight_map.values()}
                return bool(shards) and all((folder / name).is_file() and (folder / name).stat().st_size > 0
                                             for name in shards)
            except Exception:
                return False
        return False
    if kind == "adapter":
        adapter = folder / (meta.get("adapter_dir") or "adapter")
        if not (adapter / "adapter_config.json").is_file():
            return False
        return any(path.is_file() and path.stat().st_size > 0 for path in
                   (adapter / "adapter_model.safetensors", adapter / "adapter_model.bin"))
    return False


def _banner(text: str) -> None:
    print("\n" + "=" * 72 + f"\n {text}\n" + "=" * 72)


# ------------------------------------------------------------------ building blocks
def _should_purge(cfg: Dict, args, mode: str) -> bool:
    # Stage-1 is a protected lineage checkpoint. Never let the old default PURGE=true,
    # --fresh-model, or train_all remove ./checkpoints/stage1.
    return False


def _load(model_dir: str, cfg: Dict, recipe: Optional[Dict] = None):
    """profile -> plan -> tokenizer surgery -> load -> finalize. Returns (model, tokenizer, plan, surgery)."""
    prof = hardware_api.profile_system(model_dir, cfg)
    plan = hardware_api.plan_execution(prof, cfg)
    tokenizer, surgery = training_api.load_tokenizer(model_dir, recipe=recipe)
    model = hardware_api.load_model(model_dir, plan, cfg, row_recipe=surgery["recipe"] or None)
    training_api.finalize_model(model, tokenizer, surgery)
    return model, tokenizer, plan, surgery


def _ensure_clean(cfg: Dict) -> None:
    if not data_api.cleaned_data_ready(cfg):
        print("[Orchestrator] cleaned data not found -> running the clean step first")
        data_api.run_clean(cfg)


# ------------------------------------------------------------------ modes
def mode_clean(cfg: Dict, args=None) -> Dict:
    _banner("MODE clean  -  stream + clean data")
    data_api.ensure_chat_file(cfg)
    return data_api.run_clean(cfg)


def mode_pretrain(cfg: Dict, args=None, model_dir: Optional[str] = None) -> Dict:
    _banner("MODE pretrain  -  Stage 1 continual pre-training")
    dirs = _dirs(cfg)
    if _valid_stage1_checkpoint(dirs["stage1"]):
        meta = training_api.read_meta(dirs["stage1"]) or {}
        print(f"[Orchestrator] valid Stage-1 checkpoint found -> SKIP Stage 1; resume from {dirs['stage1']}",
              flush=True)
        return {"kind": meta.get("kind"), "checkpoint": str(dirs["stage1"]), "skipped": True}
    if dirs["stage1"].exists():
        raise RuntimeError(f"Stage-1 directory exists but failed validation; refusing to overwrite/delete it: {dirs['stage1']}")
    data_api.ensure_chat_file(cfg)
    _ensure_clean(cfg)
    dirs = _dirs(cfg)
    if model_dir is None:
        model_dir = hardware_api.prepare_base_model(cfg, fresh=False)
    paths = data_api.cleaned_paths(cfg)
    model, tokenizer, plan, surgery = _load(model_dir, cfg)
    try:
        result = training_api.run_pretrain(model, tokenizer, plan, cfg, surgery, paths["pretrain"], dirs["stage1"])
    finally:
        del model
        hardware_api.release_memory()
    print(f"[Orchestrator] Stage 1 done (loss {result['loss']:.4f}) -> {result['checkpoint']}")
    return result


def mode_finetune(cfg: Dict, args=None, model_dir: Optional[str] = None) -> Dict:
    _banner("MODE finetune  -  Stage 2 rational SFT")
    data_api.ensure_chat_file(cfg)
    _ensure_clean(cfg)
    dirs = _dirs(cfg)
    paths = data_api.cleaned_paths(cfg)
    meta = training_api.read_meta(dirs["stage1"])
    if dirs["stage1"].exists() and not _valid_stage1_checkpoint(dirs["stage1"]):
        raise RuntimeError(f"Stage-1 directory exists but checkpoint metadata/weights are missing/invalid; refusing fallback or overwrite: {dirs['stage1']}")
    adapter_dir: Optional[str] = None
    recipe = None
    if meta is None:
        print("[Orchestrator] no Stage-1 checkpoint found -> fine-tuning directly from the base model")
        source = model_dir or hardware_api.prepare_base_model(cfg, fresh=False)
    elif meta["kind"] == "adapter":
        source = model_dir or hardware_api.prepare_base_model(cfg, fresh=False)
        adapter_dir = str(dirs["stage1"] / meta["adapter_dir"])
        recipe = meta.get("row_recipe") or None
        print(f"[Orchestrator] continuing the Stage-1 LoRA adapter on top of {cfg['base_model_repo']}")
    else:
        source = str(dirs["stage1"])
        recipe = meta.get("row_recipe") or None
        print(f"[Orchestrator] loading the Stage-1 checkpoint {source}")
    model, tokenizer, plan, surgery = _load(source, cfg, recipe=recipe)
    try:
        result = training_api.run_finetune(model, tokenizer, plan, cfg, surgery, paths["sft"], dirs["final"],
                                           adapter_dir=adapter_dir)
    finally:
        del model
        hardware_api.release_memory()
    print(f"[Orchestrator] Stage 2 done (loss {result['loss']:.4f}) -> {result['final_dir']}")
    return result


def mode_upload(cfg: Dict, args=None) -> bool:
    final = _dirs(cfg)["final"]
    ok = training_api.upload_model(final, cfg)
    if ok:
        training_api.clear_cache(cfg["root"])
    return ok


def mode_train_all(cfg: Dict, args=None) -> Dict:
    _banner("MODE train_all  -  checkpoint-aware clean -> pretrain/skip -> finetune -> upload")
    started = time.time()
    data_api.ensure_chat_file(cfg)
    dirs = _dirs(cfg)
    stage1_valid = _valid_stage1_checkpoint(dirs["stage1"])
    if dirs["stage1"].exists() and not stage1_valid:
        raise RuntimeError(f"Stage-1 directory exists but is not a valid checkpoint; preserving it and refusing overwrite: {dirs['stage1']}")

    if stage1_valid:
        stage1_meta = training_api.read_meta(dirs["stage1"]) or {}
        print(f"[1/5] Stage-1 checkpoint verified ({stage1_meta.get('kind')}) -> preserving; no purge/download", flush=True)
        model_dir = None
        if stage1_meta.get("kind") == "adapter":
            model_dir = hardware_api.prepare_base_model(cfg, fresh=False)
    else:
        print("[1/5] base model (reuse verified local copy; destructive purge disabled)", flush=True)
        model_dir = hardware_api.prepare_base_model(cfg, fresh=False)

    print("[2/5] cleaning data", flush=True)
    if data_api.cleaned_data_ready(cfg):
        print("[Data] cleaned SFT/pretrain files already exist -> reuse", flush=True)
    else:
        data_api.run_clean(cfg)

    if stage1_valid:
        print(f"[3/5] stage 1 -> SKIPPED (verified checkpoint {dirs['stage1']})", flush=True)
    else:
        print("[3/5] stage 1 -> train once (checkpoint path is write-once)", flush=True)
        mode_pretrain(cfg, args, model_dir=model_dir)

    print("[4/5] stage 2 -> rational SFT", flush=True)
    result = mode_finetune(cfg, args, model_dir=model_dir)
    uploaded = False
    if getattr(args, "skip_upload", False):
        print("[5/5] upload skipped (--skip-upload)", flush=True)
    else:
        print("[5/5] uploading to Hugging Face", flush=True)
        uploaded = mode_upload(cfg, args)
    result["uploaded"] = uploaded
    print(f"\n[Orchestrator] train_all finished in {(time.time() - started) / 60:.1f} min  "
          f"(uploaded={uploaded}).  Chat with it:  python RSS_AI_Engine/main.py --mode serve", flush=True)
    return result


def mode_serve(cfg: Dict, args=None) -> str:
    _banner("MODE serve  -  load trained model and start the chat front-end")
    data_api.ensure_chat_file(cfg)
    trained = serving_api.resolve_trained_dir(cfg, fetch_remote=training_api.download_trained)
    meta = training_api.read_meta(trained) or {}
    if meta.get("has_adapter") and not meta.get("base_is_self"):
        weights_dir = hardware_api.prepare_base_model(cfg, fresh=False)
    else:
        weights_dir = str(trained)
    recipe = meta.get("row_recipe") or None
    cfg["active_vocab_file"] = str(trained / "active_vocab_ids.json")
    model, tokenizer, plan, surgery = _load(weights_dir, cfg, recipe=recipe)
    if weights_dir != str(trained):                            # adapter-only: tokenizer must be the trained one
        tokenizer, surgery = training_api.load_tokenizer(str(trained), recipe=recipe)
    model = serving_api.prepare_model(model, trained, meta)
    engine = serving_api.build_engine(model, tokenizer, cfg)
    return serving_api.start(engine, cfg)


def run(mode: str, cfg: Dict, args=None):
    table = {"clean": mode_clean, "pretrain": mode_pretrain, "finetune": mode_finetune,
             "train_all": mode_train_all, "serve": mode_serve}
    if getattr(args, "upload_only", False):
        return mode_upload(cfg, args)
    return table[mode](cfg, args)