adds local and remote training monitors to config

scripts/training/train_gpt_oss.py

@@ -19,6 +19,11 @@ except Exception:  # pragma: no cover - optional import depending on TRL version
     DPOTrainer = None
 from datasets import load_dataset
 from pathlib import Path
+# Import monitoring utilities from project src for persistent logging
+try:
+    from src.monitoring import create_monitor_from_config  # type: ignore
+except Exception:
+    create_monitor_from_config = None  # type: ignore
 
 # Ensure project root and config package are importable for configs that do `from config...` imports
 project_root = Path(__file__).resolve().parents[2]
@@ -876,6 +881,23 @@ def train_gpt_oss(config_path, experiment_name, output_dir, trackio_url, trainer
     # Setup Trackio tracking
     trackio_client = setup_trackio_tracking(config)
     
+    # Initialize project monitor (HF Datasets + Trackio Space if configured)
+    monitor = None
+    monitor_callback = None
+    if create_monitor_from_config is not None:
+        try:
+            monitor = create_monitor_from_config(config, experiment_name=experiment_name)
+            # Persist configuration immediately
+            try:
+                cfg_dict = {k: v for k, v in config.__dict__.items() if not k.startswith('_')}
+                monitor.log_config(cfg_dict)
+            except Exception:
+                pass
+            # Create callback for SFTTrainer
+            monitor_callback = monitor.create_monitoring_callback()
+        except Exception:
+            monitor = None
+    
     # Create SFT configuration
     sft_config = create_sft_config(config, output_dir)
     
@@ -949,6 +971,10 @@ def train_gpt_oss(config_path, experiment_name, output_dir, trackio_url, trainer
         if "packing" in sft_params:
             sft_kwargs["packing"] = getattr(config, 'packing', False)
 
+        # Attach monitoring callback if supported
+        if "callbacks" in sft_params:
+            sft_kwargs["callbacks"] = ([monitor_callback] if monitor_callback is not None else [])
+
         # Remove any None values
         sft_kwargs = {k: v for k, v in sft_kwargs.items() if v is not None}
 
@@ -959,7 +985,15 @@ def train_gpt_oss(config_path, experiment_name, output_dir, trackio_url, trainer
     
     # Start training
     print("Starting GPT-OSS training...")
-    trainer.train()
+    try:
+        trainer.train()
+    finally:
+        # Ensure periodic metrics are flushed at the end even if interrupted
+        try:
+            if monitor is not None:
+                monitor._save_to_hf_dataset({'status': 'running'})
+        except Exception:
+            pass
     
     # Save model
     print("Saving trained model...")
@@ -970,6 +1004,18 @@ def train_gpt_oss(config_path, experiment_name, output_dir, trackio_url, trainer
         print("Pushing model to Hugging Face Hub...")
         trainer.push_to_hub(dataset_name="HuggingFaceH4/Multilingual-Thinking")
     
+    # Log training summary and close monitor
+    try:
+        if monitor is not None:
+            summary = {
+                'output_dir': output_dir,
+                'model_name': getattr(config, 'model_name', 'unknown'),
+            }
+            monitor.log_training_summary(summary)
+            monitor.close()
+    except Exception:
+        pass
+
     print("GPT-OSS training completed successfully!")
     
     return trainer

src/monitoring.py

@@ -50,6 +50,11 @@ class SmolLM3Monitor:
         self.log_artifacts = log_artifacts
         self.log_metrics_enabled = log_metrics  # Rename to avoid conflict
         self.log_config_enabled = log_config  # Rename to avoid conflict
+        # Flush interval for dataset persistence (metrics)
+        try:
+            self.flush_interval = int(os.environ.get('TRACKIO_FLUSH_INTERVAL', '10'))
+        except Exception:
+            self.flush_interval = 10
         
         # HF Datasets configuration
         self.hf_token = hf_token or os.environ.get('HF_TOKEN')
@@ -343,12 +348,12 @@ class SmolLM3Monitor:
     
     def log_configuration(self, config: Dict[str, Any]):
         """Log experiment configuration"""
-        if not self.enable_tracking or not self.log_config_enabled:
+        if not self.log_config_enabled:
             return
         
         try:
             # Log configuration as parameters
-            if self.trackio_client:
+            if self.enable_tracking and self.trackio_client:
                 try:
                     result = self.trackio_client.log_parameters(
                         experiment_id=self.experiment_id,
@@ -390,7 +395,7 @@ class SmolLM3Monitor:
         - throughput, step_time, batch_size, seq_len
         - token_acc, train/gate_ortho, train/center, etc.
         """
-        if not self.enable_tracking or not self.log_metrics_enabled:
+        if not self.log_metrics_enabled:
             return
         
         try:
@@ -400,7 +405,7 @@ class SmolLM3Monitor:
                 metrics['step'] = step
             
             # Log to Trackio (if available)
-            if self.trackio_client:
+            if self.enable_tracking and self.trackio_client:
                 try:
                     result = self.trackio_client.log_metrics(
                         experiment_id=self.experiment_id,
@@ -418,8 +423,8 @@ class SmolLM3Monitor:
             # Store locally
             self.metrics_history.append(metrics)
             
-            # Save to HF Dataset periodically
-            if len(self.metrics_history) % 10 == 0:  # Save every 10 metrics
+            # Save to HF Dataset periodically (configurable)
+            if self.flush_interval > 0 and (len(self.metrics_history) % self.flush_interval == 0):
                 self._save_to_hf_dataset({'metrics': self.metrics_history})
             
             logger.debug("Metrics logged: %s", metrics)
@@ -429,7 +434,7 @@ class SmolLM3Monitor:
     
     def log_model_checkpoint(self, checkpoint_path: str, step: Optional[int] = None):
         """Log model checkpoint"""
-        if not self.enable_tracking or not self.log_artifacts:
+        if not self.log_artifacts:
             return
         
         try:
@@ -441,7 +446,7 @@ class SmolLM3Monitor:
                 "checkpoint_size": os.path.getsize(checkpoint_path) if os.path.exists(checkpoint_path) else 0
             }
             
-            if self.trackio_client:
+            if self.enable_tracking and self.trackio_client:
                 result = self.trackio_client.log_parameters(
                     experiment_id=self.experiment_id,
                     parameters=checkpoint_info
@@ -453,6 +458,11 @@ class SmolLM3Monitor:
                     logger.error("Failed to log checkpoint to Trackio: %s", result)
             
             self.artifacts.append(checkpoint_path)
+            # Also preserve checkpoint info in HF dataset
+            try:
+                self._save_to_hf_dataset({'artifacts': [checkpoint_path], **checkpoint_info})
+            except Exception:
+                pass
             logger.info("Checkpoint logged: %s", checkpoint_path)
             
         except Exception as e:
@@ -460,9 +470,6 @@ class SmolLM3Monitor:
     
     def log_evaluation_results(self, results: Dict[str, Any], step: Optional[int] = None):
         """Log evaluation results"""
-        if not self.enable_tracking:
-            return
-        
         try:
             # Add evaluation prefix to metrics
             eval_metrics = {f"eval_{k}": v for k, v in results.items()}
@@ -485,9 +492,6 @@ class SmolLM3Monitor:
     
     def log_system_metrics(self, step: Optional[int] = None):
         """Log system metrics (GPU, memory, etc.)"""
-        if not self.enable_tracking:
-            return
-        
         try:
             system_metrics = {}
             
@@ -513,9 +517,6 @@ class SmolLM3Monitor:
     
     def log_training_summary(self, summary: Dict[str, Any]):
         """Log training summary at the end"""
-        if not self.enable_tracking:
-            return
-        
         try:
             # Add experiment duration
             end_time = datetime.now()
@@ -524,7 +525,7 @@ class SmolLM3Monitor:
             summary['experiment_duration_hours'] = duration / 3600
             
             # Log final summary to Trackio
-            if self.trackio_client:
+            if self.enable_tracking and self.trackio_client:
                 result = self.trackio_client.log_parameters(
                     experiment_id=self.experiment_id,
                     parameters=summary

src/trainer.py

@@ -78,6 +78,7 @@ class SmolLM3Trainer:
         # Add simple console callback for basic monitoring
         from transformers import TrainerCallback
         
+        outer = self
         class SimpleConsoleCallback(TrainerCallback):
             def on_init_end(self, args, state, control, **kwargs):
                 """Called when training initialization is complete"""
@@ -99,6 +100,16 @@ class SmolLM3Trainer:
                     else:
                         lr_str = str(lr)
                     print(f"Step {step}: loss={loss_str}, lr={lr_str}")
+
+                    # Persist metrics via our monitor when Trackio callback isn't active
+                    try:
+                        if outer.monitor:
+                            # Avoid double logging when Trackio callback is used
+                            if not outer.monitor.enable_tracking:
+                                outer.monitor.log_metrics(dict(logs), step if isinstance(step, int) else None)
+                                outer.monitor.log_system_metrics(step if isinstance(step, int) else None)
+                    except Exception as e:
+                        logger.warning("SimpleConsoleCallback metrics persistence failed: %s", e)
             
             def on_train_begin(self, args, state, control, **kwargs):
                 print("🚀 Training started!")
@@ -109,28 +120,40 @@ class SmolLM3Trainer:
             def on_save(self, args, state, control, **kwargs):
                 step = state.global_step if hasattr(state, 'global_step') else 'unknown'
                 print(f"💾 Checkpoint saved at step {step}")
+                try:
+                    if outer.monitor and not outer.monitor.enable_tracking:
+                        checkpoint_path = os.path.join(args.output_dir, f"checkpoint-{step}")
+                        if os.path.exists(checkpoint_path):
+                            outer.monitor.log_model_checkpoint(checkpoint_path, step if isinstance(step, int) else None)
+                except Exception as e:
+                    logger.warning("SimpleConsoleCallback checkpoint persistence failed: %s", e)
             
             def on_evaluate(self, args, state, control, metrics=None, **kwargs):
                 if metrics and isinstance(metrics, dict):
                     step = state.global_step if hasattr(state, 'global_step') else 'unknown'
                     eval_loss = metrics.get('eval_loss', 'N/A')
                     print(f"📊 Evaluation at step {step}: eval_loss={eval_loss}")
+                    try:
+                        if outer.monitor and not outer.monitor.enable_tracking:
+                            outer.monitor.log_evaluation_results(dict(metrics), step if isinstance(step, int) else None)
+                    except Exception as e:
+                        logger.warning("SimpleConsoleCallback eval persistence failed: %s", e)
         
         # Add console callback
         callbacks.append(SimpleConsoleCallback())
         logger.info("Added simple console monitoring callback")
         
-        # Add Trackio callback if available
-        if self.monitor and self.monitor.enable_tracking:
+        # Add monitoring callback if available (always attach; it persists to dataset even if Trackio is disabled)
+        if self.monitor:
             try:
                 trackio_callback = self.monitor.create_monitoring_callback()
                 if trackio_callback:
                     callbacks.append(trackio_callback)
-                    logger.info("Added Trackio monitoring callback")
+                    logger.info("Added monitoring callback")
                 else:
-                    logger.warning("Failed to create Trackio callback")
+                    logger.warning("Failed to create monitoring callback")
             except Exception as e:
-                logger.error("Error creating Trackio callback: %s", e)
+                logger.error("Error creating monitoring callback: %s", e)
                 logger.info("Continuing with console monitoring only")
         
         logger.info("Total callbacks: %d", len(callbacks))
@@ -220,16 +243,20 @@ class SmolLM3Trainer:
         """Start training"""
         logger.info("Starting training")
         
-        # Log configuration to Trackio
-        if self.monitor and self.monitor.enable_tracking:
-            config_dict = {k: v for k, v in self.config.__dict__.items() 
-                          if not k.startswith('_')}
-            self.monitor.log_config(config_dict)
-            
-            # Log experiment URL
-            experiment_url = self.monitor.get_experiment_url()
-            if experiment_url:
-                logger.info("Trackio experiment URL: %s", experiment_url)
+        # Log configuration (always persist to dataset; Trackio if enabled)
+        if self.monitor:
+            try:
+                config_dict = {k: v for k, v in self.config.__dict__.items() if not k.startswith('_')}
+                self.monitor.log_config(config_dict)
+            except Exception as e:
+                logger.warning("Failed to log configuration: %s", e)
+            # Log experiment URL only if available
+            try:
+                experiment_url = self.monitor.get_experiment_url()
+                if experiment_url:
+                    logger.info("Trackio experiment URL: %s", experiment_url)
+            except Exception:
+                pass
         
         # Load checkpoint if resuming
         if self.init_from == "resume":
@@ -251,17 +278,20 @@ class SmolLM3Trainer:
             with open(os.path.join(self.output_dir, "train_results.json"), "w") as f:
                 json.dump(train_result.metrics, f, indent=2)
             
-            # Log training summary to Trackio
-            if self.monitor and self.monitor.enable_tracking:
-                summary = {
-                    'final_loss': train_result.metrics.get('train_loss', 0),
-                    'total_steps': train_result.metrics.get('train_runtime', 0),
-                    'training_time': train_result.metrics.get('train_runtime', 0),
-                    'output_dir': self.output_dir,
-                    'model_name': getattr(self.config, 'model_name', 'unknown'),
-                }
-                self.monitor.log_training_summary(summary)
-                self.monitor.close()
+            # Log training summary (always persist to dataset; Trackio if enabled)
+            if self.monitor:
+                try:
+                    summary = {
+                        'final_loss': train_result.metrics.get('train_loss', 0),
+                        'total_steps': train_result.metrics.get('train_runtime', 0),
+                        'training_time': train_result.metrics.get('train_runtime', 0),
+                        'output_dir': self.output_dir,
+                        'model_name': getattr(self.config, 'model_name', 'unknown'),
+                    }
+                    self.monitor.log_training_summary(summary)
+                    self.monitor.close()
+                except Exception as e:
+                    logger.warning("Failed to log training summary: %s", e)
             
             # Finish trackio experiment
             try:
@@ -276,9 +306,12 @@ class SmolLM3Trainer:
             
         except Exception as e:
             logger.error("Training failed: %s", e)
-            # Close monitoring on error
-            if self.monitor and self.monitor.enable_tracking:
-                self.monitor.close()
+            # Close monitoring on error (still persist final status to dataset)
+            if self.monitor:
+                try:
+                    self.monitor.close(final_status="failed")
+                except Exception:
+                    pass
             
             # Finish trackio experiment on error
             try: