orchestrator.py

# orchestrator.py
# Phase 4: Batch orchestration with training triggers, evaluation, and rollback.
#
# Features:
# - Run sessions in batches with configurable delays
# - Track per-session and rolling metrics
# - Trigger training on a sensible cadence
# - Post-train evaluation with rollback on regression
# - Graceful error handling

import os
import sys
import json
import time
import signal
import shutil
import argparse
from datetime import datetime
from typing import Optional, List
from contextlib import contextmanager

from config import (
    LOG_DIR,
    MODELS_DIR,
    CHECKPOINT_DIR,
    SESSIONS_PER_BATCH,
    INTER_SESSION_DELAY_SEC,
    TRAIN_EVERY_GRADED_TURNS,
    get_train_every_graded_turns,
    get_train_every_usable_turns,
    MIN_SESSIONS_BEFORE_TRAIN,
    TRAIN_ONLY_IF_PROGRESS_SIGNAL_RATE_AT_LEAST,
    TRAIN_PROGRESS_SIGNAL_WINDOW,
    EVAL_SESSIONS_AFTER_TRAIN,
    ROLLBACK_IF_SCORE_DROP_PCT,
    ROLLBACK_IF_COMPLIANCE_DROP_ABS,
    EVAL_COMPARE_WINDOW,
    get_basil_model_name,
    DEFAULT_WORKERS,
)
from auto_session import run_session
from metrics_manager import (
    load_rolling_metrics,
    save_rolling_metrics,
    compute_baseline,
    compute_eval_metrics,
    get_progress_signal_rate,
    record_training_attempt,
    get_recent_session_metrics,
)
from memory_manager import load_basil_assessment
from identity_probe import run_identity_probe
from curriculum_manager import clear_used_lessons
from prompts.storytime.storytime_session import run_storytime_session, clear_used_stories
from prompts.howitworks.howitworks_session import run_howitworks_session, clear_used_topics
from whychain_session import run_whychain_session, clear_used_questions


# =============================================================================
# GRACEFUL SHUTDOWN
# =============================================================================

_shutdown_requested = False
_original_sigint = signal.getsignal(signal.SIGINT)


def _sigint_handler(signum, frame):
    """
    First Ctrl+C: request graceful shutdown (finish current operation, then exit).
    Second Ctrl+C: force exit immediately.
    """
    global _shutdown_requested
    if _shutdown_requested:
        # Second Ctrl+C -- force exit
        print("\n\n[Orchestrator] Force quit! Exiting immediately.")
        sys.exit(1)
    
    _shutdown_requested = True
    print("\n\n[Orchestrator] Graceful shutdown requested. "
          "Will exit after the current session/operation finishes.\n"
          "[Orchestrator] Press Ctrl+C again to force quit.\n")


def shutdown_requested() -> bool:
    """Check if graceful shutdown has been requested."""
    return _shutdown_requested


def install_signal_handler():
    """Install the graceful shutdown signal handler."""
    signal.signal(signal.SIGINT, _sigint_handler)


def reset_signal_handler():
    """Restore the original signal handler (e.g., for training subprocess)."""
    signal.signal(signal.SIGINT, _original_sigint)


def _debug_log(run_id: str, hypothesis_id: str, location: str, message: str, data: dict):
    # region agent log
    try:
        payload = {
            "id": f"log_{int(time.time() * 1000)}_{os.getpid()}",
            "timestamp": int(time.time() * 1000),
            "runId": run_id,
            "hypothesisId": hypothesis_id,
            "location": location,
            "message": message,
            "data": data,
        }
        with open("/home/ubuntu/bootstrap-basil/bootstrap-basil/.cursor/debug.log", "a", encoding="utf-8") as f:
            f.write(json.dumps(payload, ensure_ascii=True) + "\n")
    except Exception:
        pass
    # endregion


# =============================================================================
# TRAINING LOG CAPTURE
# =============================================================================

class _TeeWriter:
    """Write to both a file and the original stream (stdout or stderr)."""
    
    def __init__(self, file, original):
        self.file = file
        self.original = original
    
    def write(self, data):
        self.original.write(data)
        self.file.write(data)
        self.file.flush()
    
    def flush(self):
        self.original.flush()
        self.file.flush()
    
    # Forward any other attribute lookups to the original stream so that
    # code checking e.g. sys.stdout.isatty() doesn't break.
    def __getattr__(self, name):
        return getattr(self.original, name)


@contextmanager
def _training_log(label: str = "train"):
    """
    Context manager that tees all stdout and stderr to a timestamped log file
    in LOG_DIR for the duration of the block.  Output still appears on the
    terminal in real time.
    
    Usage:
        with _training_log("train"):
            run_train_with_checkpoint_and_eval()
    
    Creates: logs/train_YYYYMMDD_HHMMSS.log
    """
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    log_path = os.path.join(LOG_DIR, f"{label}_{timestamp}.log")
    os.makedirs(LOG_DIR, exist_ok=True)
    
    log_file = open(log_path, "w")
    
    # Write a header
    log_file.write(f"{'='*60}\n")
    log_file.write(f"Training run log: {label}_{timestamp}\n")
    log_file.write(f"Started: {datetime.now().isoformat()}\n")
    log_file.write(f"{'='*60}\n\n")
    log_file.flush()
    
    old_stdout = sys.stdout
    old_stderr = sys.stderr
    sys.stdout = _TeeWriter(log_file, old_stdout)
    sys.stderr = _TeeWriter(log_file, old_stderr)
    
    try:
        yield log_path
    finally:
        sys.stdout = old_stdout
        sys.stderr = old_stderr
        
        log_file.write(f"\n{'='*60}\n")
        log_file.write(f"Finished: {datetime.now().isoformat()}\n")
        log_file.write(f"{'='*60}\n")
        log_file.close()
        
        print(f"[Orchestrator] Training log saved to {log_path}")


# =============================================================================
# TRAINING TRIGGER
# =============================================================================

def should_train(metrics: dict = None) -> tuple:
    """
    Check if we should trigger a training run.
    
    Returns:
        (should_train: bool, reason: str)
    """
    if metrics is None:
        metrics = load_rolling_metrics()
    
    total_sessions = metrics.get("total_sessions", 0)
    usable_since_train = metrics.get("usable_turns_since_last_train", 0)
    
    # Check minimum sessions
    if total_sessions < MIN_SESSIONS_BEFORE_TRAIN:
        return False, f"not enough sessions ({total_sessions} < {MIN_SESSIONS_BEFORE_TRAIN})"
    
    # Get dynamic training threshold based on current age_band (usable turns)
    assessment = load_basil_assessment()
    age_band = assessment.get("age_band", 0)
    train_threshold = get_train_every_usable_turns(age_band)
    
    # Check usable turns threshold
    if usable_since_train < train_threshold:
        return False, f"not enough usable turns ({usable_since_train} < {train_threshold}, age_band={age_band})"
    
    # Check progress signal rate
    progress_rate = get_progress_signal_rate(TRAIN_PROGRESS_SIGNAL_WINDOW)
    if progress_rate < TRAIN_ONLY_IF_PROGRESS_SIGNAL_RATE_AT_LEAST:
        return False, f"progress signal rate too low ({progress_rate:.1%} < {TRAIN_ONLY_IF_PROGRESS_SIGNAL_RATE_AT_LEAST:.0%})"
    
    return True, f"graded_turns={graded_since_train}, progress_rate={progress_rate:.1%}"


# =============================================================================
# CHECKPOINTING
# =============================================================================

def ensure_checkpoint_dir():
    """Ensure checkpoint directory exists."""
    os.makedirs(CHECKPOINT_DIR, exist_ok=True)


def save_checkpoint(label: str) -> str:
    """
    Save a checkpoint of the current Basil model.
    
    Args:
        label: Checkpoint label (e.g., "pretrain", "posttrain")
    
    Returns:
        Path to the checkpoint directory
    """
    ensure_checkpoint_dir()
    
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    checkpoint_name = f"{label}_{timestamp}"
    checkpoint_path = os.path.join(CHECKPOINT_DIR, checkpoint_name)
    
    try:
        model_path = get_basil_model_name()
        shutil.copytree(model_path, checkpoint_path)
        print(f"[Orchestrator] Saved checkpoint to {checkpoint_path}")
        return checkpoint_path
    except Exception as e:
        print(f"[Orchestrator] Failed to save checkpoint: {e}")
        return None


def restore_checkpoint(checkpoint_path: str) -> bool:
    """
    Restore a checkpoint by making it the latest model.
    
    This copies the checkpoint back to models/ with a new version number.
    
    Args:
        checkpoint_path: Path to the checkpoint to restore
    
    Returns:
        True if successful, False otherwise
    """
    if not os.path.exists(checkpoint_path):
        print(f"[Orchestrator] Checkpoint not found: {checkpoint_path}")
        return False
    
    try:
        # Find next version number
        existing_versions = sorted([
            int(name.split("_v")[-1])
            for name in os.listdir(MODELS_DIR)
            if name.startswith("basil_v") and name.split("_v")[-1].isdigit()
        ])
        next_version = max(existing_versions, default=1) + 1
        new_model_dir = os.path.join(MODELS_DIR, f"basil_v{next_version:03}")
        
        # Copy checkpoint as new model
        shutil.copytree(checkpoint_path, new_model_dir)
        print(f"[Orchestrator] Restored checkpoint to {new_model_dir}")
        return True
    except Exception as e:
        print(f"[Orchestrator] Failed to restore checkpoint: {e}")
        return False


# =============================================================================
# TRAINING
# =============================================================================

def run_training() -> bool:
    """
    Run the training script.
    
    Returns:
        True if training succeeded, False otherwise
    """
    print("\n[Orchestrator] Starting training...")
    
    try:
        # Import and run training directly
        from train_basil_v2 import train
        train(mode="mixed")  # Dual-objective: alternating world/trunk + basil/LoRA
        return True
    except Exception as e:
        print(f"[Orchestrator] Training failed: {e}")
        return False


def run_train_with_checkpoint_and_eval() -> str:
    """
    Full training workflow with checkpointing, eval, and rollback.
    
    1. Save pre-train checkpoint
    2. Compute baseline metrics
    3. Run training
    4. Run post-train eval sessions
    5. Compare and decide keep/rollback
    
    Returns:
        Result: "kept", "rolled_back", or "train_failed"
    """
    print(f"\n{'#'*60}")
    print("TRAINING WORKFLOW")
    print(f"{'#'*60}\n")
    
    # 0. Capture the most recent session ID as the training-run boundary.
    #    All sessions up to (and including) this one were generated by the
    #    current model version. After training, the new model produces a new run.
    metrics = load_rolling_metrics()
    recent_ids = metrics.get("last_n_session_ids", [])
    boundary_session_id = recent_ids[-1] if recent_ids else datetime.now().strftime("%Y%m%d_%H%M%S")
    
    # 1. Save pre-train checkpoint
    pretrain_checkpoint = save_checkpoint("pretrain")
    if not pretrain_checkpoint:
        print("[Orchestrator] Failed to save pre-train checkpoint, aborting training")
        record_training_attempt("train_failed", session_id=boundary_session_id)
        return "train_failed"
    
    # 2. Compute baseline metrics
    baseline = compute_baseline(EVAL_COMPARE_WINDOW)
    print(f"[Orchestrator] Baseline (last {EVAL_COMPARE_WINDOW} normal sessions):")
    print(f"  - avg_score: {baseline['avg_score']:.2f}")
    print(f"  - avg_compliance: {baseline['avg_compliance']:.1%}")
    print(f"  - sessions_count: {baseline['sessions_count']}")
    
    # 3. Run training
    training_success = run_training()
    if not training_success:
        print("[Orchestrator] Training failed, keeping current model")
        record_training_attempt("train_failed", session_id=boundary_session_id)
        return "train_failed"
    
    # 4. Save post-train checkpoint
    posttrain_checkpoint = save_checkpoint("posttrain")
    
    # 5. Run post-train eval sessions
    print(f"\n[Orchestrator] Running {EVAL_SESSIONS_AFTER_TRAIN} post-train eval sessions...")
    eval_session_ids = []
    
    for i in range(EVAL_SESSIONS_AFTER_TRAIN):
        if shutdown_requested():
            print(f"\n[Orchestrator] Shutdown requested, skipping remaining eval sessions")
            break
        
        print(f"\n[Orchestrator] Eval session {i+1}/{EVAL_SESSIONS_AFTER_TRAIN}")
        try:
            result = run_session(
                training_phase="posttrain_eval",
                verbose=True,
            )
            eval_session_ids.append(result["session_id"])
        except KeyboardInterrupt:
            print(f"\n[Orchestrator] Eval session interrupted")
            break
        except Exception as e:
            print(f"[Orchestrator] Eval session failed: {e}")
        
        if i < EVAL_SESSIONS_AFTER_TRAIN - 1:
            time.sleep(INTER_SESSION_DELAY_SEC)
    
    # 6. Compute eval metrics
    eval_metrics = compute_eval_metrics(eval_session_ids)
    print(f"\n[Orchestrator] Eval metrics ({len(eval_session_ids)} sessions):")
    print(f"  - avg_score: {eval_metrics['avg_score']:.2f}")
    print(f"  - avg_compliance: {eval_metrics['avg_compliance']:.1%}")
    
    # 7. Decide keep or rollback
    result = decide_rollback(baseline, eval_metrics, pretrain_checkpoint)
    
    # 8. Record training attempt (with boundary for recency weighting)
    record_training_attempt(result, session_id=boundary_session_id)
    
    # 9. Clear used lessons and stories lists (reset for next training run's data collection)
    clear_used_lessons()
    clear_used_stories()
    clear_used_topics()
    clear_used_questions()
    print("[Orchestrator] Cleared used lessons, stories, topics, and questions for next training run")
    
    # 10. Run identity probe on the new (or restored) model
    try:
        model_version = os.path.basename(get_basil_model_name())
        run_identity_probe(label=f"post_train_{model_version}_{result}")
    except Exception as e:
        print(f"[Orchestrator] Identity probe failed: {e}")
    
    print(f"\n{'#'*60}")
    print(f"TRAINING WORKFLOW COMPLETE: {result}")
    print(f"{'#'*60}\n")
    
    return result


def decide_rollback(baseline: dict, eval_metrics: dict, pretrain_checkpoint: str) -> str:
    """
    Decide whether to keep new model or rollback to checkpoint.
    
    Rollback if BOTH conditions are met:
    - Score dropped by >= ROLLBACK_IF_SCORE_DROP_PCT
    - Compliance dropped by >= ROLLBACK_IF_COMPLIANCE_DROP_ABS
    
    Exception: If baseline score is extremely low (<1.0), require both signals.
    
    Args:
        baseline: Baseline metrics dict
        eval_metrics: Post-train eval metrics dict
        pretrain_checkpoint: Path to pre-train checkpoint
    
    Returns:
        "kept" or "rolled_back"
    """
    baseline_score = baseline.get("avg_score", 0)
    baseline_compliance = baseline.get("avg_compliance", 0)
    eval_score = eval_metrics.get("avg_score", 0)
    eval_compliance = eval_metrics.get("avg_compliance", 0)
    
    # Calculate drops
    if baseline_score > 1e-6:
        score_drop_pct = (baseline_score - eval_score) / baseline_score
    else:
        score_drop_pct = 0  # Can't compute drop from near-zero baseline
    
    compliance_drop_abs = baseline_compliance - eval_compliance
    
    print(f"\n[Orchestrator] Regression check:")
    print(f"  - Score drop: {score_drop_pct:.1%} (threshold: {ROLLBACK_IF_SCORE_DROP_PCT:.0%})")
    print(f"  - Compliance drop: {compliance_drop_abs:.2f} (threshold: {ROLLBACK_IF_COMPLIANCE_DROP_ABS:.2f})")
    
    # Rollback requires BOTH conditions to be met
    score_regressed = score_drop_pct >= ROLLBACK_IF_SCORE_DROP_PCT
    compliance_regressed = compliance_drop_abs >= ROLLBACK_IF_COMPLIANCE_DROP_ABS
    
    # Special case: if baseline is extremely low, be more lenient
    if baseline_score < 1.0:
        print(f"  - Baseline score very low ({baseline_score:.2f}), requiring both signals for rollback")
        should_rollback = score_regressed and compliance_regressed
    else:
        # Normal case: rollback if both conditions met
        should_rollback = score_regressed and compliance_regressed
    
    if should_rollback:
        print(f"\n[Orchestrator] REGRESSION DETECTED - Rolling back to {pretrain_checkpoint}")
        if restore_checkpoint(pretrain_checkpoint):
            return "rolled_back"
        else:
            print("[Orchestrator] Rollback failed, keeping new model anyway")
            return "kept"
    else:
        print("\n[Orchestrator] No significant regression - keeping new model")
        return "kept"


# =============================================================================
# BATCH RUNNER
# =============================================================================

def run_batch(
    n_sessions: int = None,
    include_sophie: bool = True,
    verbose: bool = True,
    clear_tracking: bool = False,
) -> dict:
    """
    Run a batch of sessions.
    
    All session logs are consolidated into 3 batch-level files:
      - batch_<id>_graded.jsonl  (graded turns - consumed by training)
      - batch_<id>_sessions.jsonl (transcripts, episodes, debug data)
      - batch_<id>_meta.jsonl    (per-session metrics, one line per session)
    
    Args:
        n_sessions: Number of sessions (default: SESSIONS_PER_BATCH)
        include_sophie: Include Sophie in conversations
        verbose: Print progress
        clear_tracking: If True, clear used lessons/stories/topics before
            starting. Used by standalone `batch` CLI command. The orchestrator
            loop does NOT set this -- dedup state persists across batches and
            is only cleared after training completes.
    
    Returns:
        Batch results summary
    """
    n = n_sessions or SESSIONS_PER_BATCH
    batch_id = datetime.now().strftime("%Y%m%d_%H%M%S")
    results = []
    
    # Create batch-level file paths
    os.makedirs(LOG_DIR, exist_ok=True)
    batch_graded_path = os.path.join(LOG_DIR, f"batch_{batch_id}_graded.jsonl")
    batch_sessions_path = os.path.join(LOG_DIR, f"batch_{batch_id}_sessions.jsonl")
    batch_meta_path = os.path.join(LOG_DIR, f"batch_{batch_id}_meta.jsonl")
    
    # Only clear used lists when explicitly requested (standalone batch command).
    # In the orchestrator loop, dedup state persists across batches and is
    # cleared only after training completes (in run_train_with_checkpoint_and_eval).
    if clear_tracking:
        clear_used_lessons()
        clear_used_stories()
        clear_used_topics()
        clear_used_questions()
    
    print(f"\n{'#'*60}")
    print(f"Starting Batch {batch_id}: {n} sessions")
    print(f"  graded  -> {batch_graded_path}")
    print(f"  sessions -> {batch_sessions_path}")
    print(f"  meta     -> {batch_meta_path}")
    if clear_tracking:
        print(f"  (Cleared used lessons/stories/topics/questions)")
    print(f"{'#'*60}\n")
    
    for i in range(n):
        # Check for graceful shutdown before starting next session
        if shutdown_requested():
            print(f"\n[Batch] Graceful shutdown: stopping after {i} of {n} sessions")
            break
        
        # Rotate between four session types: regular, howitworks, whychain, storytime
        session_type = i % 4  # 0=regular, 1=howitworks, 2=whychain, 3=storytime
        
        if session_type == 0:
            print(f"\n[Batch] Session {i+1}/{n} (Regular)")
        elif session_type == 1:
            print(f"\n[Batch] Session {i+1}/{n} (How It Works)")
        elif session_type == 2:
            print(f"\n[Batch] Session {i+1}/{n} (WhyChain)")
        else:
            print(f"\n[Batch] Session {i+1}/{n} (Storytime)")
        
        try:
            if session_type == 1:
                result = run_howitworks_session(
                    verbose=verbose,
                    training_phase="howitworks",
                    batch_graded_path=batch_graded_path,
                    batch_sessions_path=batch_sessions_path,
                    batch_meta_path=batch_meta_path,
                )
            elif session_type == 2:
                result = run_whychain_session(
                    verbose=verbose,
                    training_phase="whychain",
                    batch_graded_path=batch_graded_path,
                    batch_sessions_path=batch_sessions_path,
                    batch_meta_path=batch_meta_path,
                )
            elif session_type == 3:
                result = run_storytime_session(
                    verbose=verbose,
                    training_phase="storytime",
                    batch_graded_path=batch_graded_path,
                    batch_sessions_path=batch_sessions_path,
                    batch_meta_path=batch_meta_path,
                )
            else:
                result = run_session(
                    include_sophie=include_sophie,
                    verbose=verbose,
                    training_phase="normal",
                    batch_graded_path=batch_graded_path,
                    batch_sessions_path=batch_sessions_path,
                    batch_meta_path=batch_meta_path,
                )
            results.append(result)
        except KeyboardInterrupt:
            # Session interrupted -- treat as shutdown request
            print(f"\n[Batch] Session interrupted, stopping batch")
            break
        except Exception as e:
            print(f"[Batch] Session failed with exception: {e}")
            # Record a failed session
            failed_result = {
                "session_id": datetime.now().strftime("%Y%m%d_%H%M%S"),
                "session_metrics": {
                    "early_stopped": True,
                    "stop_reason": "exception",
                    "graded_turns_count": 0,
                    "avg_score_session": 0,
                    "compliance_rate_session": 0,
                },
            }
            results.append(failed_result)
        
        # Sleep between sessions (interruptible)
        if i < n - 1 and not shutdown_requested():
            time.sleep(INTER_SESSION_DELAY_SEC)
    
    # Calculate batch summary
    metrics_list = [r.get("session_metrics", {}) for r in results]
    total_graded = sum(m.get("graded_turns_count", 0) for m in metrics_list)
    scores = [m.get("avg_score_session", 0) for m in metrics_list if m.get("graded_turns_count", 0) > 0]
    avg_score = sum(scores) / len(scores) if scores else 0
    
    # Count session types
    regular_sessions = sum(1 for r in results if r.get("session_metrics", {}).get("training_phase") == "normal")
    howitworks_sessions = sum(1 for r in results if r.get("session_metrics", {}).get("training_phase") == "howitworks")
    whychain_sessions = sum(1 for r in results if r.get("session_metrics", {}).get("training_phase") == "whychain")
    storytime_sessions = sum(1 for r in results if r.get("session_metrics", {}).get("training_phase") == "storytime")
    
    summary = {
        "batch_id": batch_id,
        "sessions_completed": len(results),
        "regular_sessions": regular_sessions,
        "howitworks_sessions": howitworks_sessions,
        "whychain_sessions": whychain_sessions,
        "storytime_sessions": storytime_sessions,
        "total_graded_turns": total_graded,
        "average_score": avg_score,
        "session_ids": [r.get("session_id") for r in results],
        "batch_graded_path": batch_graded_path,
        "batch_sessions_path": batch_sessions_path,
        "batch_meta_path": batch_meta_path,
    }
    
    # Save batch summary (small JSON, replaces old per-batch file)
    batch_log_path = os.path.join(LOG_DIR, f"batch_{batch_id}_summary.json")
    with open(batch_log_path, "w") as f:
        json.dump(summary, f, indent=2)
    
    print(f"\n{'#'*60}")
    print(f"Batch Complete")
    print(f"  Sessions: {len(results)} ({regular_sessions} regular, {howitworks_sessions} howitworks, {whychain_sessions} whychain, {storytime_sessions} storytime)")
    print(f"  Graded turns: {total_graded}")
    print(f"  Average score: {avg_score:.2f}/7")
    print(f"  Files: 3 batch logs + 1 summary")
    print(f"{'#'*60}\n")
    
    return summary


def run_orchestrator_loop(
    max_batches: int = None,
    max_sessions: int = None,
    include_sophie: bool = True,
    verbose: bool = True,
):
    """
    Main orchestrator loop.
    
    Runs sessions in batches, triggers training when appropriate.
    
    Args:
        max_batches: Maximum batches to run (None = infinite)
        max_sessions: Maximum total sessions to run (None = infinite)
        include_sophie: Include Sophie in sessions
        verbose: Print progress
    """
    # Get current age_band for dynamic threshold display
    assessment = load_basil_assessment()
    age_band = assessment.get("age_band", 0)
    train_threshold = get_train_every_usable_turns(age_band)
    
    print(f"\n{'='*60}")
    print("Bootstrap Basil Orchestrator")
    print(f"{'='*60}")
    print(f"Max batches: {max_batches or 'infinite'}")
    print(f"Max sessions: {max_sessions or 'infinite'}")
    print(f"Sessions per batch: {SESSIONS_PER_BATCH}")
    print(f"Train every: {train_threshold} usable turns (age_band={age_band}, score>{age_band})")
    print(f"{'='*60}\n")
    
    batch_count = 0
    session_count = 0
    
    install_signal_handler()
    
    while True:
        # Check for graceful shutdown
        if shutdown_requested():
            print(f"\n[Orchestrator] Shutting down gracefully after {session_count} sessions, {batch_count} batches.")
            break
        
        # Check batch limit
        if max_batches is not None and batch_count >= max_batches:
            print(f"\n[Orchestrator] Reached max batches ({max_batches}), stopping")
            break
        
        # Determine sessions for this batch
        if max_sessions is not None:
            remaining = max_sessions - session_count
            if remaining <= 0:
                print(f"\n[Orchestrator] Reached max sessions ({max_sessions}), stopping")
                break
            batch_size = min(SESSIONS_PER_BATCH, remaining)
        else:
            batch_size = SESSIONS_PER_BATCH
        
        # Run batch
        batch_result = run_batch(
            n_sessions=batch_size,
            include_sophie=include_sophie,
            verbose=verbose,
        )
        
        batch_count += 1
        session_count += batch_result.get("sessions_completed", 0)
        
        # Check for shutdown before training/memory updates
        if shutdown_requested():
            print(f"\n[Orchestrator] Shutting down gracefully (skipping training check).")
            break
        
        # Check training trigger
        metrics = load_rolling_metrics()
        usable_since = metrics.get("usable_turns_since_last_train", 0)
        graded_since = metrics.get("graded_turns_since_last_train", 0)
        current_age_band = load_basil_assessment().get("age_band", 0)
        usable_target = get_train_every_usable_turns(current_age_band)
        print(f"[Progress] Usable: {usable_since:,} / {usable_target:,} | Total graded: {graded_since:,} (age_band={current_age_band})")
        
        trigger, reason = should_train(metrics)
        
        if trigger:
            print(f"\n[Orchestrator] Training trigger: {reason}")
            with _training_log("train"):
                run_train_with_checkpoint_and_eval()
        else:
            print(f"[Orchestrator] Training not triggered: {reason}")
        
    # Run identity probe on exit (snapshot of current model's identity)
    try:
        model_version = os.path.basename(get_basil_model_name())
        run_identity_probe(label=f"exit_{model_version}")
    except Exception as e:
        print(f"[Orchestrator] Exit identity probe failed: {e}")
    
    print(f"[Orchestrator] Exited cleanly.")


# =============================================================================
# STATUS
# =============================================================================

def print_status():
    """Print current system status."""
    metrics = load_rolling_metrics()
    
    print("\n" + "="*60)
    print("Bootstrap Basil Status")
    print("="*60)
    print(f"Total sessions: {metrics.get('total_sessions', 0)}")
    print(f"Total graded turns: {metrics.get('total_graded_turns', 0)}")
    print(f"Graded turns since last train: {metrics.get('graded_turns_since_last_train', 0)}")
    print(f"Usable turns since last train: {metrics.get('usable_turns_since_last_train', 0)}")
    print(f"Total training runs: {metrics.get('total_training_runs', 0)}")
    print(f"Total rollbacks: {metrics.get('total_rollbacks', 0)}")
    print(f"Last training: {metrics.get('last_train_timestamp', 'Never')}")
    print(f"Last train result: {metrics.get('last_train_result', 'N/A')}")
    
    # EWMA
    print(f"\nEWMA Metrics:")
    print(f"  EWMA score: {metrics.get('ewma_score', 'N/A')}")
    print(f"  EWMA compliance: {metrics.get('ewma_compliance', 'N/A')}")
    
    # Recent summary
    summary = metrics.get("recent_summary", {})
    print(f"\nRecent Summary (last 10 normal sessions):")
    print(f"  Avg score: {summary.get('avg_score', 'N/A')}")
    print(f"  Avg compliance: {summary.get('avg_compliance', 'N/A')}")
    print(f"  Progress signal rate: {summary.get('progress_signal_rate', 'N/A')}")
    
    # Training trigger check
    print("\nTraining status:")
    trigger, reason = should_train(metrics)
    print(f"  Ready to train: {trigger}")
    print(f"  Reason: {reason}")
    
    # Progress signal rate
    progress_rate = get_progress_signal_rate()
    print(f"  Progress signal rate (last {TRAIN_PROGRESS_SIGNAL_WINDOW}): {progress_rate:.1%}")
    
    print("="*60 + "\n")


# =============================================================================
# CLI
# =============================================================================

def main():
    parser = argparse.ArgumentParser(
        description="Bootstrap Basil Orchestrator",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  python orchestrator.py run                    # Run forever
  python orchestrator.py run --batches 1        # Run one batch
  python orchestrator.py run --sessions 3       # Run 3 sessions
  python orchestrator.py status                 # Print status
  python orchestrator.py train                  # Force training
  python orchestrator.py train --force          # Force training (skip checks)
"""
    )
    
    subparsers = parser.add_subparsers(dest="command", help="Command to run")
    
    # Run command
    run_parser = subparsers.add_parser("run", help="Run orchestrator loop")
    run_parser.add_argument("--batches", type=int, default=None,
                           help="Max batches to run (default: infinite)")
    run_parser.add_argument("--sessions", type=int, default=None,
                           help="Max sessions to run (default: infinite)")
    run_parser.add_argument("--workers", type=int, default=1,
                           help=f"Number of parallel workers (default: 1, use {DEFAULT_WORKERS} for parallel mode)")
    run_parser.add_argument("--no-sophie", action="store_true",
                           help="Disable Sophie")
    run_parser.add_argument("--quiet", action="store_true",
                           help="Reduce output")
    
    # Status command
    subparsers.add_parser("status", help="Print current status")
    
    # Train command
    train_parser = subparsers.add_parser("train", help="Trigger training")
    train_parser.add_argument("--force", action="store_true",
                             help="Force training even if trigger conditions not met")
    
    # Batch command (legacy compatibility)
    batch_parser = subparsers.add_parser("batch", help="Run a single batch (legacy)")
    batch_parser.add_argument("--sessions", type=int, default=SESSIONS_PER_BATCH,
                             help="Number of sessions")
    batch_parser.add_argument("--no-sophie", action="store_true",
                             help="Disable Sophie")
    batch_parser.add_argument("--quiet", action="store_true",
                             help="Reduce output")
    
    # Train-check command (legacy compatibility)
    subparsers.add_parser("train-check", help="Check if training should be triggered")
    
    args = parser.parse_args()
    
    if args.command == "run" or args.command is None:
        # Default to run if no command specified
        batches = getattr(args, "batches", None)
        sessions = getattr(args, "sessions", None)
        workers = getattr(args, "workers", 1)
        no_sophie = getattr(args, "no_sophie", False)
        quiet = getattr(args, "quiet", False)
        run_id = f"orch_{datetime.now().strftime('%Y%m%d_%H%M%S')}_{os.getpid()}"
        # region agent log
        _debug_log(
            run_id=run_id,
            hypothesis_id="H1_cli_route",
            location="orchestrator.py:run_command_entry",
            message="Parsed orchestrator run args",
            data={
                "batches": batches,
                "sessions": sessions,
                "workers": workers,
                "no_sophie": bool(no_sophie),
                "quiet": bool(quiet),
            },
        )
        # endregion
        
        if workers > 1:
            # Parallel mode: delegate to parallel_generate
            from parallel_generate import run_parallel_generation
            if sessions:
                # Explicit --sessions: estimate ~6 graded turns per session
                target = sessions * 6
                # region agent log
                _debug_log(
                    run_id=run_id,
                    hypothesis_id="H2_target_source",
                    location="orchestrator.py:parallel_target_from_sessions",
                    message="Target derived from sessions*6",
                    data={"sessions": sessions, "target_turns": target},
                )
                # endregion
            else:
                # Default: generate enough usable turns to reach training threshold
                assessment = load_basil_assessment()
                age_band = assessment.get("age_band", 0)
                metrics = load_rolling_metrics()
                usable_since = metrics.get("usable_turns_since_last_train", 0)
                target = max(0, get_train_every_usable_turns(age_band) - usable_since)
                # region agent log
                _debug_log(
                    run_id=run_id,
                    hypothesis_id="H2_target_source",
                    location="orchestrator.py:parallel_target_from_threshold",
                    message="Target derived from usable threshold",
                    data={
                        "age_band": age_band,
                        "usable_since_last_train": usable_since,
                        "threshold": get_train_every_usable_turns(age_band),
                        "target_turns": target,
                    },
                )
                # endregion

            if target <= 0 and not sessions:
                # Already have enough usable turns (e.g. aborted training left counters uncleared)
                print(f"\n[Orchestrator] Already at training threshold "
                      f"(usable={usable_since:,} >= {get_train_every_usable_turns(age_band):,}, "
                      f"age_band={age_band}). Running training directly.\n")
                with _training_log("train"):
                    result = run_train_with_checkpoint_and_eval()
                print(f"\n[Orchestrator] Training result: {result}")
            else:
                # region agent log
                _debug_log(
                    run_id=run_id,
                    hypothesis_id="H1_cli_route",
                    location="orchestrator.py:before_run_parallel_generation",
                    message="Calling run_parallel_generation",
                    data={"target_turns": target, "num_workers": workers, "do_train": True},
                )
                # endregion
                run_parallel_generation(
                    target_turns=target,
                    num_workers=workers,
                    do_train=True,
                )
        else:
            run_orchestrator_loop(
                max_batches=batches,
                max_sessions=sessions,
                include_sophie=not no_sophie,
                verbose=not quiet,
            )
    
    elif args.command == "status":
        print_status()
    
    elif args.command == "train":
        if args.force:
            print("[Orchestrator] Forcing training...")
            with _training_log("train"):
                run_train_with_checkpoint_and_eval()
        else:
            trigger, reason = should_train()
            if trigger:
                print(f"[Orchestrator] Training triggered: {reason}")
                with _training_log("train"):
                    run_train_with_checkpoint_and_eval()
            else:
                print(f"[Orchestrator] Training not triggered: {reason}")
                print("Use --force to override")
    
    elif args.command == "batch":
        run_batch(
            n_sessions=args.sessions,
            include_sophie=not args.no_sophie,
            verbose=not args.quiet,
            clear_tracking=True,  # Standalone batch starts fresh
        )
    
    elif args.command == "train-check":
        trigger, reason = should_train()
        if trigger:
            print(f"Training should be triggered: {reason}")
        else:
            print(f"Training not ready: {reason}")


if __name__ == "__main__":
    main()