openadapt-ml 0.1.0__py3-none-any.whl → 0.2.0__py3-none-any.whl

openadapt_ml/training/trainer.py

@@ -6,13 +6,7 @@ from dataclasses import dataclass, field
 from pathlib import Path
 from typing import Any, Callable, Dict, List, Optional
 
-import torch
-from torch.optim import Optimizer
-from torch.optim.lr_scheduler import LambdaLR
-from torch.utils.data import DataLoader, Dataset
-
-from openadapt_ml.models.base_adapter import BaseVLMAdapter
-from openadapt_ml.schemas.sessions import Episode
+from openadapt_ml.schema import Episode, Step, Action, ActionType
 from openadapt_ml.training.shared_ui import (
     get_shared_header_css as _get_shared_header_css,
     generate_shared_header_html as _generate_shared_header_html,
@@ -21,6 +15,10 @@ from openadapt_ml.training.shared_ui import (
 from openadapt_ml.training.viewer import (
     generate_unified_viewer_from_output_dir,
 )
+from openadapt_ml.training.benchmark_viewer import (
+    _get_azure_jobs_panel_css,
+    _get_azure_jobs_panel_html,
+)
 
 
 def setup_job_directory(base_dir: str | Path, job_id: str) -> Path:
@@ -116,6 +114,11 @@ class TrainingState:
     capture_path: str = ""
     config_path: str = ""
     goal: str = ""  # Task goal/description for the training run
+    # Model configuration
+    model_name: str = ""  # e.g. "Qwen/Qwen3-VL-2B-Instruct"
+    lora_r: int = 0  # LoRA rank
+    lora_alpha: int = 0  # LoRA alpha
+    load_in_4bit: bool = False  # Quantization
     # Training progress
     epoch: int = 0
     step: int = 0
@@ -185,6 +188,11 @@ class TrainingState:
             "capture_path": self.capture_path,
             "config_path": self.config_path,
             "goal": self.goal,
+            # Model configuration
+            "model_name": self.model_name,
+            "lora_r": self.lora_r,
+            "lora_alpha": self.lora_alpha,
+            "load_in_4bit": self.load_in_4bit,
             "instance_type": self.instance_type,
             "instance_ip": self.instance_ip,
             "started_at": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime(self.start_time)),
@@ -227,6 +235,11 @@ class TrainingLogger:
         cloud_dashboard_url: str = "",
         cloud_instance_id: str = "",
         job_id: str = "",
+        # Model configuration
+        model_name: str = "",
+        lora_r: int = 0,
+        lora_alpha: int = 0,
+        load_in_4bit: bool = False,
     ):
         # Generate job_id if not provided
         if not job_id:
@@ -242,6 +255,10 @@ class TrainingLogger:
             capture_path=capture_path,
             config_path=config_path,
             goal=goal,
+            model_name=model_name,
+            lora_r=lora_r,
+            lora_alpha=lora_alpha,
+            load_in_4bit=load_in_4bit,
             instance_ip=instance_ip,
             instance_type=instance_type,
             total_epochs=config.num_train_epochs,
@@ -428,14 +445,18 @@ def generate_training_dashboard(state: TrainingState, config: TrainingConfig) ->
                 # Generate comparison data with null predictions (shows "— No prediction")
                 comparison_data = []
                 for i, step in enumerate(episode.steps):
+                    # Extract normalized coordinates if available
+                    action_x, action_y = None, None
+                    if step.action.normalized_coordinates:
+                        action_x, action_y = step.action.normalized_coordinates
                     step_data = {
                         "index": i,
-                        "time": step.t,
-                        "image_path": step.observation.image_path,
+                        "time": step.step_index,
+                        "image_path": step.observation.screenshot_path,
                         "human_action": {
-                            "type": step.action.type,
-                            "x": step.action.x,
-                            "y": step.action.y,
+                            "type": step.action.type.value if isinstance(step.action.type, ActionType) else step.action.type,
+                            "x": action_x,
+                            "y": action_y,
                             "text": step.action.text,
                         },
                         "predicted_action": None,  # Shows "— No prediction" in viewer
@@ -596,6 +617,42 @@ def generate_training_dashboard(state: TrainingState, config: TrainingConfig) ->
         .setup-log-line.current {{
             color: var(--accent);
         }}
+        .config-panel {{
+            background: var(--bg-secondary);
+            border: 1px solid var(--border-color);
+            border-radius: 12px;
+            padding: 16px 20px;
+            margin-bottom: 24px;
+        }}
+        .config-grid {{
+            display: grid;
+            grid-template-columns: repeat(auto-fit, minmax(200px, 1fr));
+            gap: 16px;
+        }}
+        .config-item {{
+            display: flex;
+            flex-direction: column;
+            gap: 4px;
+        }}
+        .config-label {{
+            font-size: 0.7rem;
+            color: var(--text-secondary);
+            text-transform: uppercase;
+            letter-spacing: 0.5px;
+        }}
+        .config-value {{
+            font-family: "SF Mono", Monaco, monospace;
+            font-size: 0.85rem;
+            color: var(--text-primary);
+        }}
+        .config-value.model {{
+            color: var(--accent);
+        }}
+        .config-value.goal {{
+            font-family: -apple-system, BlinkMacSystemFont, "Inter", sans-serif;
+            font-size: 0.8rem;
+            opacity: 0.9;
+        }}
         .status {{
             display: flex;
             align-items: center;
@@ -754,6 +811,8 @@ def generate_training_dashboard(state: TrainingState, config: TrainingConfig) ->
         }}
         /* Shared header styles (injected from _get_shared_header_css) */
         {_get_shared_header_css()}
+        /* Azure ML Jobs panel styles */
+        {_get_azure_jobs_panel_css()}
         .eval-panel {{
             background: var(--bg-secondary);
             border: 1px solid var(--border-color);
@@ -1095,6 +1154,33 @@ def generate_training_dashboard(state: TrainingState, config: TrainingConfig) ->
             </div>
         </div>
 
+        {_get_azure_jobs_panel_html()}
+
+        <div class="config-panel" id="config-panel">
+            <div class="config-grid">
+                <div class="config-item">
+                    <span class="config-label">Model</span>
+                    <span class="config-value model" id="config-model">{state.model_name or 'Not specified'}</span>
+                </div>
+                <div class="config-item">
+                    <span class="config-label">Goal</span>
+                    <span class="config-value goal" id="config-goal">{state.goal or 'Not specified'}</span>
+                </div>
+                <div class="config-item">
+                    <span class="config-label">LoRA</span>
+                    <span class="config-value" id="config-lora">{f'r={state.lora_r}, α={state.lora_alpha}' if state.lora_r else 'Not specified'}</span>
+                </div>
+                <div class="config-item">
+                    <span class="config-label">Quantization</span>
+                    <span class="config-value" id="config-quant">{'4-bit' if state.load_in_4bit else 'None'}</span>
+                </div>
+                <div class="config-item">
+                    <span class="config-label">Config</span>
+                    <span class="config-value" id="config-path">{state.config_path or 'Not specified'}</span>
+                </div>
+            </div>
+        </div>
+
         <div class="stats-grid">
             <div class="stat-card" id="card-epoch">
                 <div class="stat-label">Epoch Progress</div>
@@ -1519,6 +1605,28 @@ def generate_training_dashboard(state: TrainingState, config: TrainingConfig) ->
                     }}
                 }}
 
+                // Update config panel
+                const configModel = document.getElementById('config-model');
+                const configGoal = document.getElementById('config-goal');
+                const configLora = document.getElementById('config-lora');
+                const configQuant = document.getElementById('config-quant');
+                const configPath = document.getElementById('config-path');
+                if (configModel && data.model_name) {{
+                    configModel.textContent = data.model_name;
+                }}
+                if (configGoal && data.goal) {{
+                    configGoal.textContent = data.goal;
+                }}
+                if (configLora && (data.lora_r || data.lora_alpha)) {{
+                    configLora.textContent = `r=${{data.lora_r || 0}}, α=${{data.lora_alpha || 0}}`;
+                }}
+                if (configQuant) {{
+                    configQuant.textContent = data.load_in_4bit ? '4-bit' : 'None';
+                }}
+                if (configPath && data.config_path) {{
+                    configPath.textContent = data.config_path;
+                }}
+
                 // Update setup panel if setup logs present
                 if (data.setup_logs && data.setup_logs.length > 0) {{
                     const setupPanel = document.getElementById('setup-panel');
@@ -2093,354 +2201,3 @@ def regenerate_local_dashboard(
     print(f"Regenerated dashboard: {dashboard_path}")
 
     return dashboard_path
-
-
-def run_epoch_evaluation(
-    adapter: BaseVLMAdapter,
-    episode: Episode,
-    epoch: int,
-    config: TrainingConfig,
-    logger: "TrainingLogger",
-    sample_indices: Optional[List[int]] = None,
-) -> Path:
-    """Run inference evaluation on sample steps after an epoch.
-
-    This generates a comparison_epoch{N}.html file showing human vs predicted actions.
-
-    Args:
-        adapter: Trained adapter to use for inference
-        episode: Episode with steps to evaluate
-        epoch: Current epoch number
-        config: Training configuration
-        logger: Training logger for state tracking
-        sample_indices: Specific step indices to evaluate (default: evenly spaced)
-
-    Returns:
-        Path to generated comparison HTML file
-    """
-    from openadapt_ml.scripts.compare import generate_comparison_html, predict_action, format_action
-
-    output_dir = Path(config.output_dir)
-    output_dir.mkdir(parents=True, exist_ok=True)
-
-    # Select sample indices if not provided
-    num_samples = min(config.eval_samples, len(episode.steps))
-    if sample_indices is None:
-        if num_samples >= len(episode.steps):
-            sample_indices = list(range(len(episode.steps)))
-        else:
-            # Evenly space samples across the episode
-            step_size = len(episode.steps) // num_samples
-            sample_indices = [i * step_size for i in range(num_samples)]
-
-    print(f"  Running inference on {len(sample_indices)} sample steps...")
-
-    # Switch adapter to eval mode
-    adapter.eval()
-
-    comparison_data = []
-    action_history: List[str] = []
-    total_steps = len(episode.steps)
-
-    for i, step in enumerate(episode.steps):
-        step_data = {
-            "index": i,
-            "time": step.t,
-            "image_path": step.observation.image_path,
-            "human_action": {
-                "type": step.action.type,
-                "x": step.action.x,
-                "y": step.action.y,
-                "text": step.action.text,
-            },
-            "predicted_action": None,
-            "match": None,
-        }
-
-        # Only run inference on selected samples (for speed)
-        if i in sample_indices and step.observation.image_path:
-            try:
-                predicted = predict_action(
-                    adapter,
-                    step.observation.image_path,
-                    episode.goal,
-                    step_index=i,
-                    total_steps=total_steps,
-                    action_history=action_history.copy(),
-                )
-                step_data["predicted_action"] = predicted
-
-                # Check match and calculate distance
-                if predicted and predicted.get("type") == step.action.type:
-                    step_data["match"] = True
-
-                    # Calculate distance for click actions
-                    if step.action.type == "click":
-                        hx, hy = step.action.x or 0, step.action.y or 0
-                        px, py = predicted.get("x", 0), predicted.get("y", 0)
-                        distance = ((hx - px) ** 2 + (hy - py) ** 2) ** 0.5
-
-                        # Log evaluation to training state
-                        logger.state.log_evaluation(
-                            epoch=epoch,
-                            sample_idx=i,
-                            image_path=step.observation.image_path,
-                            human_action=step_data["human_action"],
-                            predicted_action=predicted,
-                        )
-                else:
-                    step_data["match"] = False
-
-                print(f"    Step {i}: {step.action.type} -> {predicted.get('type') if predicted else 'none'}")
-
-            except Exception as e:
-                print(f"    Step {i}: inference failed - {e}")
-
-        # Build action history for context
-        action_history.append(format_action(step.action, use_som=False))
-        comparison_data.append(step_data)
-
-    # Switch back to train mode
-    adapter.train()
-
-    # Generate comparison HTML
-    output_path = output_dir / f"comparison_epoch{epoch}.html"
-    capture_path = Path(logger.state.capture_path) if logger.state.capture_path else Path(".")
-
-    generate_comparison_html(capture_path, episode, comparison_data, output_path)
-    print(f"  Comparison saved: {output_path}")
-
-    # Also regenerate all dashboards to update navigation
-    regenerate_all_dashboards(output_dir)
-
-    return output_path
-
-
-def _create_dataloader(dataset: Dataset, batch_size: int) -> DataLoader:
-    # Use an identity collate_fn so that each batch is a List[Dict], matching
-    # the expectations of adapters that operate on SFT-style samples.
-    return DataLoader(dataset, batch_size=batch_size, shuffle=True, collate_fn=lambda x: x)
-
-
-def _create_lr_scheduler(
-    optimizer: Optimizer,
-    config: TrainingConfig,
-    num_training_steps: int,
-) -> Optional[LambdaLR]:
-    """Create learning rate scheduler based on config.
-
-    Args:
-        optimizer: The optimizer to schedule.
-        config: Training configuration with lr_scheduler_type and warmup_ratio.
-        num_training_steps: Total number of training steps.
-
-    Returns:
-        LambdaLR scheduler or None if scheduler_type is "none" or "constant".
-    """
-    scheduler_type = config.lr_scheduler_type.lower()
-
-    if scheduler_type in ("none", "constant"):
-        return None
-
-    num_warmup_steps = int(num_training_steps * config.warmup_ratio)
-
-    if scheduler_type == "linear":
-        def lr_lambda(current_step: int) -> float:
-            if current_step < num_warmup_steps:
-                # Linear warmup
-                return float(current_step) / float(max(1, num_warmup_steps))
-            # Linear decay
-            return max(
-                0.0,
-                float(num_training_steps - current_step) / float(max(1, num_training_steps - num_warmup_steps))
-            )
-    elif scheduler_type == "cosine":
-        import math
-        def lr_lambda(current_step: int) -> float:
-            if current_step < num_warmup_steps:
-                # Linear warmup
-                return float(current_step) / float(max(1, num_warmup_steps))
-            # Cosine decay
-            progress = float(current_step - num_warmup_steps) / float(max(1, num_training_steps - num_warmup_steps))
-            return max(0.0, 0.5 * (1.0 + math.cos(math.pi * progress)))
-    else:
-        raise ValueError(f"Unknown lr_scheduler_type: {scheduler_type}. Use 'linear', 'cosine', 'constant', or 'none'.")
-
-    return LambdaLR(optimizer, lr_lambda)
-
-
-def train_supervised(
-    adapter: BaseVLMAdapter,
-    dataset: Dataset,
-    config: TrainingConfig,
-    optimizer: Optional[Optimizer] = None,
-    logger: Optional[TrainingLogger] = None,
-    episode: Optional[Episode] = None,
-) -> bool:
-    """Minimal supervised training loop skeleton.
-
-    This assumes that `adapter.prepare_inputs` and `adapter.compute_loss` are
-    implemented. It will raise if those methods are not implemented.
-
-    Args:
-        adapter: VLM adapter to train.
-        dataset: Training dataset.
-        config: Training configuration.
-        optimizer: Optional optimizer (default: AdamW).
-        logger: Optional training logger for visualization.
-        episode: Optional episode for periodic evaluation (generates comparison_epoch{N}.html).
-
-    Returns:
-        True if training completed successfully, False if aborted due to NaN/Inf loss.
-    """
-
-    device = adapter.device  # type: ignore[attr-defined]
-    dataloader = _create_dataloader(dataset, batch_size=config.per_device_train_batch_size)
-
-    if optimizer is None:
-        optimizer = torch.optim.AdamW(
-            adapter.model.parameters(),  # type: ignore[arg-type]
-            lr=config.learning_rate,
-            weight_decay=config.weight_decay,
-        )
-
-    # Create logger if not provided
-    if logger is None:
-        logger = TrainingLogger(config.output_dir, config)
-
-    # Calculate total training steps for scheduler
-    num_training_steps = len(dataloader) * config.num_train_epochs // config.gradient_accumulation_steps
-
-    # Create learning rate scheduler
-    lr_scheduler = _create_lr_scheduler(optimizer, config, num_training_steps)
-
-    total_steps = 0
-    adapter.train()
-
-    # Early stopping tracking
-    consecutive_low_loss = 0
-    early_stopped = False
-    user_stopped = False
-
-    for epoch in range(config.num_train_epochs):
-        if early_stopped or user_stopped:
-            break
-
-        for _, batch in enumerate(dataloader):
-            # Check for stop signal from dashboard
-            stop_file = Path(config.output_dir) / "STOP_TRAINING"
-            if stop_file.exists():
-                msg = "Stop signal received from dashboard. Stopping training..."
-                print(msg)
-                logger._log_to_terminal(msg)
-                # Set termination status for dashboard
-                logger.state.termination_status = "user_stop"
-                logger.state.termination_message = "Training stopped by user via dashboard"
-                logger.save()
-                user_stopped = True
-                stop_file.unlink()  # Remove signal file
-                break
-
-            # Batch is a List[Dict[str, Any]] of SFT-style samples; adapter is
-            # responsible for converting it into model inputs.
-            samples: List[Dict[str, Any]] = batch
-
-            inputs = adapter.prepare_inputs(samples)
-            inputs = {k: v.to(device) if isinstance(v, torch.Tensor) else v for k, v in inputs.items()}
-
-            loss = adapter.compute_loss(inputs)
-
-            # Guard against invalid losses to avoid propagating NaNs/Infs
-            if torch.isnan(loss) or torch.isinf(loss):
-                msg = f"Encountered invalid loss at epoch={epoch} step={total_steps + 1}: {loss.item()}"
-                print(msg)
-                logger._log_to_terminal(msg)
-                logger.on_train_end()
-                return False
-
-            loss.backward()
-
-            if (total_steps + 1) % config.gradient_accumulation_steps == 0:
-                torch.nn.utils.clip_grad_norm_(adapter.model.parameters(), config.max_grad_norm)  # type: ignore[arg-type]
-                optimizer.step()
-                if lr_scheduler is not None:
-                    lr_scheduler.step()
-                optimizer.zero_grad()
-
-            total_steps += 1
-            loss_val = loss.item()
-
-            # Get current learning rate from optimizer
-            current_lr = optimizer.param_groups[0]['lr']
-
-            # Log step
-            logger.on_step(epoch, total_steps, loss_val, current_lr)
-
-            if config.logging_steps and total_steps % config.logging_steps == 0:
-                msg = f"epoch={epoch} step={total_steps} loss={loss_val:.4f} lr={current_lr:.6f}"
-                print(msg)
-                logger._log_to_terminal(msg)
-
-            # Early stopping check
-            if loss_val < config.early_stop_loss:
-                consecutive_low_loss += 1
-                if consecutive_low_loss >= config.early_stop_patience:
-                    msg = (
-                        f"Early stopping: loss ({loss_val:.6f}) below threshold "
-                        f"({config.early_stop_loss}) for {config.early_stop_patience} consecutive steps"
-                    )
-                    print(msg)
-                    logger._log_to_terminal(msg)
-                    # Set termination status for dashboard
-                    logger.state.termination_status = "auto_low_loss"
-                    logger.state.termination_message = (
-                        f"Loss reached {loss_val:.6f} (< {config.early_stop_loss}) "
-                        f"for {config.early_stop_patience} consecutive steps"
-                    )
-                    logger.save()
-                    early_stopped = True
-                    break
-            else:
-                consecutive_low_loss = 0
-
-        # End of epoch
-        logger.on_epoch_end(epoch)
-
-        # Save checkpoint at end of each epoch
-        if config.save_checkpoint_every_epoch:
-            checkpoint_path = Path(config.checkpoint_dir) / f"epoch_{epoch}"
-            checkpoint_path.mkdir(parents=True, exist_ok=True)
-            try:
-                adapter.save_checkpoint(str(checkpoint_path))
-                msg = f"Checkpoint saved to {checkpoint_path}"
-                print(msg)
-                logger._log_to_terminal(msg)
-            except Exception as e:
-                msg = f"Warning: Failed to save checkpoint: {e}"
-                print(msg)
-                logger._log_to_terminal(msg)
-
-        # Run evaluation after each epoch (generates comparison_epoch{N}.html)
-        if config.eval_every_epoch and episode is not None:
-            try:
-                print(f"Running epoch {epoch} evaluation...")
-                run_epoch_evaluation(
-                    adapter=adapter,
-                    episode=episode,
-                    epoch=epoch,
-                    config=config,
-                    logger=logger,
-                )
-            except Exception as e:
-                print(f"Warning: Epoch evaluation failed: {e}")
-                import traceback
-                traceback.print_exc()
-
-    # Set termination status if not already set (normal completion)
-    if not logger.state.termination_status:
-        logger.state.termination_status = "auto_complete"
-        logger.state.termination_message = f"Training completed all {config.num_train_epochs} epochs"
-        logger.save()
-
-    logger.on_train_end()
-    return True