consciousness/training/apollo_plugin/training_worker.py

"""Training subprocess - handles Apollo training and checkpoint sync.

Long-lived process that:
1. Loads IPC handles from vLLM's exported weights
2. Creates HF model with views into vLLM's GPU memory
3. Handles training requests via ZMQ
4. Handles checkpoint sync requests
5. Persists Apollo optimizer state between calls

Communicates with the API server's /train endpoint via ZMQ REP socket.
"""

import logging
import os
import signal
import sys
from pathlib import Path
from typing import Any

# Handle running as script vs module
if __name__ == '__main__' and __package__ is None:
    # Running as script - add parent to path for imports
    sys.path.insert(0, str(Path(__file__).parent.parent))
    __package__ = 'apollo_plugin'

import torch
import torch.nn as nn
import zmq

from .checkpoint_sync import checkpoint_sync
from .optimizer import Apollo
from .weight_mapping import load_hf_model_with_vllm_weights

logger = logging.getLogger(__name__)

DEFAULT_RANK = 64
DEFAULT_ZMQ_ADDR = "ipc:///tmp/apollo_training.sock"
HANDLE_PATH = "/tmp/vllm_weight_handles.pt"
OPTIMIZER_STATE_PATH = "/tmp/apollo_optimizer_state.pt"


class TrainingWorker:
    """Long-lived training worker process."""

    def __init__(self, zmq_addr: str = DEFAULT_ZMQ_ADDR):
        self.zmq_addr = zmq_addr
        self.model: nn.Module | None = None
        self.optimizer: Apollo | None = None
        self.model_path: str | None = None
        self._running = True

    def _create_model_wrapper(self) -> nn.Module:
        """Create HF model wrapper with views into vLLM's GPU memory."""
        if not os.path.exists(HANDLE_PATH):
            raise FileNotFoundError(
                f"Weight handles not found: {HANDLE_PATH}. "
                "Is vLLM running with the export hook?"
            )

        handles = torch.load(HANDLE_PATH, weights_only=False)

        # Extract metadata
        metadata = handles.pop('__metadata__', {})
        self.model_path = metadata.get('model_path') or os.environ.get('APOLLO_MODEL_PATH')
        if not self.model_path:
            raise ValueError(
                "Model path not found in handles metadata or APOLLO_MODEL_PATH env var"
            )

        # Reconstruct tensors from IPC handles
        vllm_params = {}
        for name, info in handles.items():
            func, args = info['handle']
            vllm_params[name] = func(*args)

        model = load_hf_model_with_vllm_weights(vllm_params, self.model_path)
        model.train()
        return model

    def _get_or_create_optimizer(self, config: dict[str, Any]) -> Apollo:
        """Get existing optimizer or create new one."""
        if self.optimizer is not None:
            return self.optimizer

        # Build parameter groups (Apollo for 2D+, standard Adam for small/1D)
        apollo_params, standard_params = [], []
        for p in self.model.parameters():
            if p.requires_grad:
                if p.ndim >= 2 and min(p.shape) >= DEFAULT_RANK:
                    apollo_params.append(p)
                else:
                    standard_params.append(p)

        groups = []
        if apollo_params:
            groups.append({'params': apollo_params})
        if standard_params:
            groups.append({'params': standard_params})

        if not groups:
            raise ValueError("No trainable parameters found")

        self.optimizer = Apollo(
            groups,
            lr=config.get('lr', 1e-5),
            rank=config.get('rank', DEFAULT_RANK),
            betas=tuple(config.get('betas', (0.9, 0.999))),
            eps=config.get('eps', 1e-8),
            weight_decay=config.get('weight_decay', 0.01),
            warmup_steps=config.get('warmup_steps', 0),
            scale=config.get('scale'),
            proj_refresh=config.get('proj_refresh', 200),
            norm_growth_limit=config.get('norm_growth_limit', 1.01),
        )

        # Restore state if exists
        if os.path.exists(OPTIMIZER_STATE_PATH):
            try:
                state = torch.load(OPTIMIZER_STATE_PATH, weights_only=False)
                self.optimizer.load_state_dict(state)
                logger.info(f"Restored optimizer state from {OPTIMIZER_STATE_PATH}")
            except Exception as e:
                logger.warning(f"Could not restore optimizer state: {e}")

        logger.info(
            f"Optimizer: {len(apollo_params)} apollo params, "
            f"{len(standard_params)} standard, "
            f"state={self.optimizer.state_size_bytes()/1e6:.1f}MB"
        )

        return self.optimizer

    def _save_optimizer_state(self):
        """Save optimizer state for persistence."""
        if self.optimizer is not None:
            torch.save(self.optimizer.state_dict(), OPTIMIZER_STATE_PATH)
            logger.info(f"Saved optimizer state to {OPTIMIZER_STATE_PATH}")

    def _run_training(
        self,
        samples: list[dict[str, Any]],
        config: dict[str, Any],
    ) -> list[float]:
        """Run Apollo training on the given samples."""
        optimizer = self._get_or_create_optimizer(config)

        loss_history = []

        for i, sample in enumerate(samples):
            ctx_ids = sample['context_ids']
            cont_ids = sample['continuation_ids']
            all_ids = ctx_ids + cont_ids
            context_len = len(ctx_ids)

            input_ids = torch.tensor([all_ids], device='cuda:0')

            optimizer.zero_grad()

            # Context-frozen forward pass
            with torch.no_grad():
                outputs = self.model(input_ids[:, :context_len], use_cache=True)
                past_kv = outputs.past_key_values

            # Decision tokens with gradients
            with torch.enable_grad():
                outputs = self.model(
                    input_ids[:, context_len:],
                    past_key_values=past_kv,
                    use_cache=False,
                )
                logits = outputs.logits

                # Shift: predict next token from each position
                shift_logits = logits[:, :-1].contiguous()
                shift_labels = input_ids[:, context_len + 1:].contiguous()

                loss = nn.functional.cross_entropy(
                    shift_logits.view(-1, shift_logits.size(-1)),
                    shift_labels.view(-1),
                )

            loss.backward()
            optimizer.step()

            loss_val = loss.item()
            loss_history.append(loss_val)
            logger.info(
                f"Step {i+1}/{len(samples)}: loss={loss_val:.4f} "
                f"(ctx={context_len}, cont={len(cont_ids)} tokens)"
            )

        return loss_history

    def _handle_train(self, request: dict[str, Any]) -> dict[str, Any]:
        """Handle a training request."""
        samples = request.get('samples', [])
        config = request.get('config', {})

        if not samples:
            return {'error': 'No training samples provided'}

        try:
            loss_history = self._run_training(samples, config)
            return {
                'status': 'completed',
                'training_samples': len(samples),
                'loss_history': loss_history,
            }
        except Exception as e:
            logger.exception(f"Training failed: {e}")
            return {'error': str(e)}

    def _handle_checkpoint(self, request: dict[str, Any]) -> dict[str, Any]:
        """Handle a checkpoint sync request."""
        if not self.model_path:
            return {'error': 'Model path not set'}

        try:
            self._save_optimizer_state()
            result = checkpoint_sync(self.model_path)
            return {
                'status': 'completed',
                'total_changed': result['total_changed'],
                'files_changed': result['files_changed'],
            }
        except Exception as e:
            logger.exception(f"Checkpoint sync failed: {e}")
            return {'error': str(e)}

    def _handle_status(self, request: dict[str, Any]) -> dict[str, Any]:
        """Handle a status request."""
        return {
            'status': 'ready',
            'model_loaded': self.model is not None,
            'optimizer_loaded': self.optimizer is not None,
            'model_path': self.model_path,
            'optimizer_state_mb': (
                self.optimizer.state_size_bytes() / 1e6
                if self.optimizer else 0
            ),
        }

    def run(self):
        """Main loop - listen for requests and handle them."""
        # Set up signal handlers
        def handle_signal(signum, frame):
            logger.info(f"Received signal {signum}, shutting down...")
            self._running = False

        signal.signal(signal.SIGTERM, handle_signal)
        signal.signal(signal.SIGINT, handle_signal)

        # Set up ZMQ socket first so API server can connect
        context = zmq.Context()
        socket = context.socket(zmq.REP)
        socket.bind(self.zmq_addr)
        logger.info(f"Training worker listening on {self.zmq_addr}")

        # Create HF model wrapper with views into vLLM's GPU memory
        logger.info("Connecting to vLLM weights via IPC handles...")
        try:
            self.model = self._create_model_wrapper()
            logger.info("HF model wrapper ready (views into vLLM GPU memory)")
        except Exception as e:
            logger.error(f"Failed to connect to vLLM weights: {e}")
            logger.info("Will retry on first training request")

        # Set socket timeout so we can check _running flag
        socket.setsockopt(zmq.RCVTIMEO, 1000)  # 1 second timeout

        while self._running:
            try:
                message = socket.recv_json()
            except zmq.Again:
                # Timeout, check _running and continue
                continue

            request_type = message.get('type', 'train')
            logger.info(f"Received {request_type} request")

            # Ensure model is loaded
            if self.model is None and request_type != 'status':
                try:
                    self.model = self._create_model_wrapper()
                except Exception as e:
                    socket.send_json({'error': f'Model not loaded: {e}'})
                    continue

            # Dispatch request
            if request_type == 'train':
                response = self._handle_train(message)
            elif request_type == 'checkpoint':
                response = self._handle_checkpoint(message)
            elif request_type == 'status':
                response = self._handle_status(message)
            else:
                response = {'error': f'Unknown request type: {request_type}'}

            socket.send_json(response)

        # Cleanup
        logger.info("Saving optimizer state before shutdown...")
        self._save_optimizer_state()
        socket.close()
        context.term()
        logger.info("Training worker shut down")


def main():
    """Entry point for running as a subprocess."""
    logging.basicConfig(
        level=logging.INFO,
        format='[apollo-worker] %(asctime)s %(levelname)s %(message)s',
        datefmt='%H:%M:%S',
    )

    zmq_addr = os.environ.get('APOLLO_ZMQ_ADDR', DEFAULT_ZMQ_ADDR)
    worker = TrainingWorker(zmq_addr)
    worker.run()


if __name__ == '__main__':
    main()