adds template files , adds non destructive dataset updates

requirements/requirements_core.txt

@@ -6,7 +6,7 @@ accelerate>=0.20.0
 peft>=0.17.0  # Updated for GPT-OSS LoRA support
 trl>=0.20.0  # Updated for GPT-OSS compatibility
 kernels
-
+openai-harmony
 # Hugging Face Hub for model and space management
 huggingface_hub>=0.19.0
 

scripts/trackio_tonic/app.py

@@ -26,6 +26,7 @@ class TrackioSpace:
         self.experiments = {}
         self.current_experiment = None
         self.backup_mode = False
+        self.dataset_manager = None
         
         # Get dataset repository and HF token from parameters or environment variables
         # Respect explicit values; avoid hardcoded defaults that might point to test repos
@@ -38,6 +39,17 @@ class TrackioSpace:
         if not self.hf_token:
             logger.warning("⚠️ HF_TOKEN not found. Some features may not work.")
         
+        # Initialize dataset manager for safe, non-destructive operations
+        try:
+            import sys
+            sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', '..', 'src'))
+            from dataset_utils import TrackioDatasetManager  # type: ignore
+            if self.hf_token and self.dataset_repo:
+                self.dataset_manager = TrackioDatasetManager(self.dataset_repo, self.hf_token)
+                logger.info("✅ Dataset manager initialized (data preservation enabled)")
+        except Exception as e:
+            logger.warning(f"⚠️ Dataset manager not available, using legacy save mode: {e}")
+        
         self._load_experiments()
         
     def _load_experiments(self):
@@ -314,13 +326,45 @@ class TrackioSpace:
         self.current_experiment = 'exp_20250720_134319'
         logger.info(f"✅ Loaded {len(backup_experiments)} backup experiments")
     
+    def _upsert_experiment(self, experiment_id: str):
+        """Non-destructive upsert of a single experiment to the dataset if manager available."""
+        try:
+            if not self.dataset_manager or not self.hf_token:
+                # Fallback to legacy save method
+                self._save_experiments()
+                return
+            exp = self.experiments.get(experiment_id)
+            if not exp:
+                return
+            # Build dataset row with JSON-encoded fields
+            payload = {
+                'experiment_id': experiment_id,
+                'name': exp.get('name', ''),
+                'description': exp.get('description', ''),
+                'created_at': exp.get('created_at', ''),
+                'status': exp.get('status', 'running'),
+                'metrics': json.dumps(exp.get('metrics', []), default=str),
+                'parameters': json.dumps(exp.get('parameters', {}), default=str),
+                'artifacts': json.dumps(exp.get('artifacts', []), default=str),
+                'logs': json.dumps(exp.get('logs', []), default=str),
+                'last_updated': datetime.now().isoformat()
+            }
+            self.dataset_manager.upsert_experiment(payload)
+        except Exception as e:
+            logger.warning(f"⚠️ Upsert failed, falling back to legacy save: {e}")
+            self._save_experiments()
+    
     def _save_experiments(self):
-        """Save experiments to HF Dataset"""
+        """Save experiments to HF Dataset (legacy fallback).
+
+        Prefer using dataset manager upserts in per-operation paths. This method is
+        retained as a fallback when the manager isn't available.
+        """
         try:
             if self.backup_mode:
                 logger.warning("⚠️ Backup mode active; skipping dataset save to avoid overwriting real data with demo values")
                 return
-            if self.hf_token:
+            if self.hf_token and not self.dataset_manager:
                 from datasets import Dataset
                 from huggingface_hub import HfApi
                 
@@ -351,10 +395,10 @@ class TrackioSpace:
                     private=True  # Make it private for security
                 )
                 
-                logger.info(f"✅ Saved {len(dataset_data)} experiments to {self.dataset_repo}")
+                logger.info(f"✅ Saved {len(dataset_data)} experiments to {self.dataset_repo} (legacy mode)")
                 
             else:
-                logger.warning("⚠️ No HF_TOKEN available, experiments not saved to dataset")
+                logger.warning("⚠️ No dataset manager and/or HF_TOKEN available, experiments not saved to dataset")
                 
         except Exception as e:
             logger.error(f"Failed to save experiments to dataset: {e}")
@@ -389,7 +433,8 @@ class TrackioSpace:
         
         self.experiments[experiment_id] = experiment
         self.current_experiment = experiment_id
-        self._save_experiments()
+        # Prefer non-destructive upsert
+        self._upsert_experiment(experiment_id)
         
         logger.info(f"Created experiment: {experiment_id} - {name}")
         return experiment
@@ -406,7 +451,7 @@ class TrackioSpace:
         }
         
         self.experiments[experiment_id]['metrics'].append(metric_entry)
-        self._save_experiments()
+        self._upsert_experiment(experiment_id)
         logger.info(f"Logged metrics for experiment {experiment_id}: {metrics}")
     
     def log_parameters(self, experiment_id: str, parameters: Dict[str, Any]):
@@ -415,7 +460,7 @@ class TrackioSpace:
             raise ValueError(f"Experiment {experiment_id} not found")
         
         self.experiments[experiment_id]['parameters'].update(parameters)
-        self._save_experiments()
+        self._upsert_experiment(experiment_id)
         logger.info(f"Logged parameters for experiment {experiment_id}: {parameters}")
     
     def log_artifact(self, experiment_id: str, artifact_name: str, artifact_data: str):
@@ -430,7 +475,7 @@ class TrackioSpace:
         }
         
         self.experiments[experiment_id]['artifacts'].append(artifact_entry)
-        self._save_experiments()
+        self._upsert_experiment(experiment_id)
         logger.info(f"Logged artifact for experiment {experiment_id}: {artifact_name}")
     
     def get_experiment(self, experiment_id: str) -> Optional[Dict[str, Any]]:
@@ -449,7 +494,7 @@ class TrackioSpace:
         """Update experiment status"""
         if experiment_id in self.experiments:
             self.experiments[experiment_id]['status'] = status
-            self._save_experiments()
+            self._upsert_experiment(experiment_id)
             logger.info(f"Updated experiment {experiment_id} status to {status}")
     
     def get_metrics_dataframe(self, experiment_id: str) -> pd.DataFrame:

scripts/trackio_tonic/dataset_utils.py

@@ -0,0 +1,328 @@
+#!/usr/bin/env python3
+"""
+Dataset utilities for Trackio experiment data management
+Provides functions for safe dataset operations with data preservation
+"""
+
+import json
+import logging
+from datetime import datetime
+from typing import Dict, Any, List, Optional, Union
+from datasets import Dataset, load_dataset
+
+logger = logging.getLogger(__name__)
+
+class TrackioDatasetManager:
+    """
+    Manager class for Trackio experiment datasets with data preservation.
+    
+    This class ensures that existing experiment data is always preserved
+    when adding new experiments or updating existing ones.
+    """
+    
+    def __init__(self, dataset_repo: str, hf_token: str):
+        """
+        Initialize the dataset manager.
+        
+        Args:
+            dataset_repo (str): HF dataset repository ID (e.g., "username/dataset-name")
+            hf_token (str): Hugging Face token for authentication
+        """
+        self.dataset_repo = dataset_repo
+        self.hf_token = hf_token
+        self._validate_repo_format()
+    
+    def _validate_repo_format(self):
+        """Validate dataset repository format"""
+        if not self.dataset_repo or '/' not in self.dataset_repo:
+            raise ValueError(f"Invalid dataset repository format: {self.dataset_repo}")
+    
+    def check_dataset_exists(self) -> bool:
+        """
+        Check if the dataset repository exists and is accessible.
+        
+        Returns:
+            bool: True if dataset exists and is accessible, False otherwise
+        """
+        try:
+            load_dataset(self.dataset_repo, token=self.hf_token)
+            logger.info(f"✅ Dataset {self.dataset_repo} exists and is accessible")
+            return True
+        except Exception as e:
+            logger.info(f"📊 Dataset {self.dataset_repo} doesn't exist or isn't accessible: {e}")
+            return False
+    
+    def load_existing_experiments(self) -> List[Dict[str, Any]]:
+        """
+        Load all existing experiments from the dataset.
+        
+        Returns:
+            List[Dict[str, Any]]: List of existing experiment dictionaries
+        """
+        try:
+            if not self.check_dataset_exists():
+                logger.info("📊 No existing dataset found, returning empty list")
+                return []
+            
+            dataset = load_dataset(self.dataset_repo, token=self.hf_token)
+            
+            if 'train' not in dataset:
+                logger.info("📊 No 'train' split found in dataset")
+                return []
+            
+            experiments = list(dataset['train'])
+            logger.info(f"📊 Loaded {len(experiments)} existing experiments")
+            
+            # Validate experiment structure
+            valid_experiments = []
+            for exp in experiments:
+                if self._validate_experiment_structure(exp):
+                    valid_experiments.append(exp)
+                else:
+                    logger.warning(f"⚠️ Skipping invalid experiment: {exp.get('experiment_id', 'unknown')}")
+            
+            logger.info(f"📊 {len(valid_experiments)} valid experiments loaded")
+            return valid_experiments
+            
+        except Exception as e:
+            logger.error(f"❌ Failed to load existing experiments: {e}")
+            return []
+    
+    def _validate_experiment_structure(self, experiment: Dict[str, Any]) -> bool:
+        """
+        Validate that an experiment has the required structure.
+        
+        Args:
+            experiment (Dict[str, Any]): Experiment dictionary to validate
+            
+        Returns:
+            bool: True if experiment structure is valid
+        """
+        required_fields = [
+            'experiment_id', 'name', 'description', 'created_at', 
+            'status', 'metrics', 'parameters', 'artifacts', 'logs'
+        ]
+        
+        for field in required_fields:
+            if field not in experiment:
+                logger.warning(f"⚠️ Missing required field '{field}' in experiment")
+                return False
+        
+        # Validate JSON fields
+        json_fields = ['metrics', 'parameters', 'artifacts', 'logs']
+        for field in json_fields:
+            if isinstance(experiment[field], str):
+                try:
+                    json.loads(experiment[field])
+                except json.JSONDecodeError:
+                    logger.warning(f"⚠️ Invalid JSON in field '{field}' for experiment {experiment.get('experiment_id')}")
+                    return False
+        
+        return True
+    
+    def save_experiments(self, experiments: List[Dict[str, Any]], commit_message: Optional[str] = None) -> bool:
+        """
+        Save a list of experiments to the dataset, preserving data integrity.
+        
+        Args:
+            experiments (List[Dict[str, Any]]): List of experiment dictionaries
+            commit_message (Optional[str]): Custom commit message
+            
+        Returns:
+            bool: True if save was successful, False otherwise
+        """
+        try:
+            if not experiments:
+                logger.warning("⚠️ No experiments to save")
+                return False
+            
+            # Validate all experiments before saving
+            valid_experiments = []
+            for exp in experiments:
+                if self._validate_experiment_structure(exp):
+                    # Ensure last_updated is set
+                    if 'last_updated' not in exp:
+                        exp['last_updated'] = datetime.now().isoformat()
+                    valid_experiments.append(exp)
+                else:
+                    logger.error(f"❌ Invalid experiment structure: {exp.get('experiment_id', 'unknown')}")
+                    return False
+            
+            # Create dataset
+            dataset = Dataset.from_list(valid_experiments)
+            
+            # Generate commit message if not provided
+            if not commit_message:
+                commit_message = f"Update dataset with {len(valid_experiments)} experiments ({datetime.now().isoformat()})"
+            
+            # Push to hub
+            dataset.push_to_hub(
+                self.dataset_repo,
+                token=self.hf_token,
+                private=True,
+                commit_message=commit_message
+            )
+            
+            logger.info(f"✅ Successfully saved {len(valid_experiments)} experiments to {self.dataset_repo}")
+            return True
+            
+        except Exception as e:
+            logger.error(f"❌ Failed to save experiments to dataset: {e}")
+            return False
+    
+    def upsert_experiment(self, experiment: Dict[str, Any]) -> bool:
+        """
+        Insert a new experiment or update an existing one, preserving all other data.
+        
+        Args:
+            experiment (Dict[str, Any]): Experiment dictionary to upsert
+            
+        Returns:
+            bool: True if operation was successful, False otherwise
+        """
+        try:
+            # Validate the experiment structure
+            if not self._validate_experiment_structure(experiment):
+                logger.error(f"❌ Invalid experiment structure for {experiment.get('experiment_id', 'unknown')}")
+                return False
+            
+            # Load existing experiments
+            existing_experiments = self.load_existing_experiments()
+            
+            # Find if experiment already exists
+            experiment_id = experiment['experiment_id']
+            experiment_found = False
+            updated_experiments = []
+            
+            for existing_exp in existing_experiments:
+                if existing_exp.get('experiment_id') == experiment_id:
+                    # Update existing experiment
+                    logger.info(f"🔄 Updating existing experiment: {experiment_id}")
+                    experiment['last_updated'] = datetime.now().isoformat()
+                    updated_experiments.append(experiment)
+                    experiment_found = True
+                else:
+                    # Preserve existing experiment
+                    updated_experiments.append(existing_exp)
+            
+            # If experiment doesn't exist, add it
+            if not experiment_found:
+                logger.info(f"➕ Adding new experiment: {experiment_id}")
+                experiment['last_updated'] = datetime.now().isoformat()
+                updated_experiments.append(experiment)
+            
+            # Save all experiments
+            commit_message = f"{'Update' if experiment_found else 'Add'} experiment {experiment_id} (preserving {len(existing_experiments)} existing experiments)"
+            
+            return self.save_experiments(updated_experiments, commit_message)
+            
+        except Exception as e:
+            logger.error(f"❌ Failed to upsert experiment: {e}")
+            return False
+    
+    def get_experiment_by_id(self, experiment_id: str) -> Optional[Dict[str, Any]]:
+        """
+        Retrieve a specific experiment by its ID.
+        
+        Args:
+            experiment_id (str): The experiment ID to search for
+            
+        Returns:
+            Optional[Dict[str, Any]]: The experiment dictionary if found, None otherwise
+        """
+        try:
+            experiments = self.load_existing_experiments()
+            
+            for exp in experiments:
+                if exp.get('experiment_id') == experiment_id:
+                    logger.info(f"✅ Found experiment: {experiment_id}")
+                    return exp
+            
+            logger.info(f"📊 Experiment not found: {experiment_id}")
+            return None
+            
+        except Exception as e:
+            logger.error(f"❌ Failed to get experiment {experiment_id}: {e}")
+            return None
+    
+    def list_experiments(self, status_filter: Optional[str] = None) -> List[Dict[str, Any]]:
+        """
+        List all experiments, optionally filtered by status.
+        
+        Args:
+            status_filter (Optional[str]): Filter by experiment status (running, completed, failed, paused)
+            
+        Returns:
+            List[Dict[str, Any]]: List of experiments matching the filter
+        """
+        try:
+            experiments = self.load_existing_experiments()
+            
+            if status_filter:
+                filtered_experiments = [exp for exp in experiments if exp.get('status') == status_filter]
+                logger.info(f"📊 Found {len(filtered_experiments)} experiments with status '{status_filter}'")
+                return filtered_experiments
+            
+            logger.info(f"📊 Found {len(experiments)} total experiments")
+            return experiments
+            
+        except Exception as e:
+            logger.error(f"❌ Failed to list experiments: {e}")
+            return []
+    
+    def backup_dataset(self, backup_suffix: Optional[str] = None) -> str:
+        """
+        Create a backup of the current dataset.
+        
+        Args:
+            backup_suffix (Optional[str]): Optional suffix for backup repo name
+            
+        Returns:
+            str: Backup repository name if successful, empty string otherwise
+        """
+        try:
+            if not backup_suffix:
+                backup_suffix = datetime.now().strftime('%Y%m%d_%H%M%S')
+            
+            backup_repo = f"{self.dataset_repo}-backup-{backup_suffix}"
+            
+            # Load current experiments
+            experiments = self.load_existing_experiments()
+            
+            if not experiments:
+                logger.warning("⚠️ No experiments to backup")
+                return ""
+            
+            # Create backup dataset manager
+            backup_manager = TrackioDatasetManager(backup_repo, self.hf_token)
+            
+            # Save to backup
+            success = backup_manager.save_experiments(
+                experiments, 
+                f"Backup of {self.dataset_repo} created on {datetime.now().isoformat()}"
+            )
+            
+            if success:
+                logger.info(f"✅ Backup created: {backup_repo}")
+                return backup_repo
+            else:
+                logger.error("❌ Failed to create backup")
+                return ""
+                
+        except Exception as e:
+            logger.error(f"❌ Failed to create backup: {e}")
+            return ""
+
+
+def create_dataset_manager(dataset_repo: str, hf_token: str) -> TrackioDatasetManager:
+    """
+    Factory function to create a TrackioDatasetManager instance.
+    
+    Args:
+        dataset_repo (str): HF dataset repository ID
+        hf_token (str): Hugging Face token
+        
+    Returns:
+        TrackioDatasetManager: Configured dataset manager instance
+    """
+    return TrackioDatasetManager(dataset_repo, hf_token)

scripts/trackio_tonic/deploy_trackio_space.py

@@ -199,10 +199,14 @@ class TrackioSpaceDeployer:
             templates_dir = project_root / "templates" / "spaces" / "trackio"
             
             # Files to copy from templates/spaces/trackio
+            # Include dataset_utils.py and trackio_api_client.py to enable
+            # safe dataset upserts and local API client usage on the Space.
             files_to_copy = [
                 "app.py",
                 "requirements.txt",
-                "README.md"
+                "README.md",
+                "dataset_utils.py",
+                "trackio_api_client.py",
             ]
             
             # Copy files from templates/spaces/trackio to temp directory

scripts/training/train_gpt_oss.py

@@ -61,13 +61,20 @@ def load_gpt_oss_model_and_tokenizer(config):
         # No quantization
         quantization_config = None
     
-    # Model kwargs as per tutorial
-    model_kwargs = {
+    # Build model kwargs with sensible defaults and allow config overrides
+    default_model_kwargs = {
         "attn_implementation": "eager",
         "torch_dtype": torch.bfloat16,
         "use_cache": False,
         "device_map": "auto",
     }
+
+    cfg_model_kwargs = getattr(config, "model_kwargs", None)
+    if isinstance(cfg_model_kwargs, dict):
+        # Config overrides defaults (e.g., attn_implementation="kernels-community/vllm-flash-attn3")
+        model_kwargs = {**default_model_kwargs, **cfg_model_kwargs}
+    else:
+        model_kwargs = default_model_kwargs.copy()
     
     # Only add quantization_config if it's not None
     if quantization_config is not None:

src/monitoring.py

@@ -204,36 +204,128 @@ class SmolLM3Monitor:
             self.experiment_id = f"exp_{timestamp}"
     
     def _save_to_hf_dataset(self, experiment_data: Dict[str, Any]):
-        """Save experiment data to HF Dataset with data preservation using dataset manager"""
+        """Save experiment data to HF Dataset with data preservation using dataset manager.
+
+        This method MERGES with any existing experiment entry to avoid overwriting data:
+        - If experiment_data contains a 'metrics' list, append new metric entries (with de-dup by step+timestamp)
+          and store using the nested structure expected by the Trackio Space (each entry has
+          {timestamp, step, metrics: {...}}).
+        - Otherwise, treat experiment_data as a parameters update and dict-merge it into existing parameters.
+        - Artifacts are merged and de-duplicated by their string value.
+        """
         if not self.dataset_manager:
             logger.warning("⚠️ Dataset manager not available")
             return False
-        
+
         try:
-            # Prepare current experiment data with standardized structure
+            experiment_id = self.experiment_id or f"exp_{datetime.now().strftime('%Y%m%d_%H%M%S')}"
+
+            # Load existing experiment (if any)
+            existing = self.dataset_manager.get_experiment_by_id(experiment_id) or {}
+
+            # Helper to safely parse JSON fields from existing
+            def _parse_json_field(value, default):
+                try:
+                    if value is None:
+                        return default
+                    if isinstance(value, str):
+                        return json.loads(value) if value else default
+                    return value
+                except Exception:
+                    return default
+
+            existing_metrics = _parse_json_field(existing.get('metrics'), [])
+            existing_parameters = _parse_json_field(existing.get('parameters'), {})
+            existing_artifacts = _parse_json_field(existing.get('artifacts'), [])
+            existing_logs = _parse_json_field(existing.get('logs'), [])
+
+            # Start from existing fields
+            merged_metrics = list(existing_metrics) if isinstance(existing_metrics, list) else []
+            merged_parameters = dict(existing_parameters) if isinstance(existing_parameters, dict) else {}
+            merged_artifacts = list(existing_artifacts) if isinstance(existing_artifacts, list) else []
+
+            # Merge incoming data
+            if 'metrics' in experiment_data:
+                # Accept either a list of metric dicts or a single metrics dict
+                incoming_metrics = experiment_data.get('metrics')
+
+                # Build a set of (step, timestamp) to deduplicate
+                def _entry_key(entry: Dict[str, Any]):
+                    return (entry.get('step'), entry.get('timestamp'))
+
+                existing_keys = set()
+                for entry in merged_metrics:
+                    # Support both nested and flat formats in existing data
+                    if isinstance(entry, dict) and 'metrics' in entry:
+                        existing_keys.add(_entry_key(entry))
+                    elif isinstance(entry, dict):
+                        existing_keys.add((entry.get('step'), entry.get('timestamp')))
+
+                def _to_nested_entry(metric: Dict[str, Any]) -> Dict[str, Any]:
+                    # If already nested, return as-is
+                    if isinstance(metric, dict) and 'metrics' in metric:
+                        return metric
+                    # Convert flat dict into nested format expected by the Space
+                    step_val = metric.get('step')
+                    ts_val = metric.get('timestamp')
+                    metrics_only = {k: v for k, v in metric.items() if k not in ('step', 'timestamp')}
+                    return {
+                        'timestamp': ts_val,
+                        'step': step_val,
+                        'metrics': metrics_only
+                    }
+
+                if isinstance(incoming_metrics, list):
+                    for m in incoming_metrics:
+                        nested = _to_nested_entry(m if isinstance(m, dict) else {})
+                        if _entry_key(nested) not in existing_keys:
+                            merged_metrics.append(nested)
+                            existing_keys.add(_entry_key(nested))
+                elif isinstance(incoming_metrics, dict):
+                    nested = _to_nested_entry(incoming_metrics)
+                    if _entry_key(nested) not in existing_keys:
+                        merged_metrics.append(nested)
+                # else: ignore invalid metrics payload
+            else:
+                # Treat as parameters update; merge dict
+                try:
+                    if isinstance(experiment_data, dict):
+                        merged_parameters.update(experiment_data)
+                except Exception:
+                    pass
+
+            # Merge artifacts if provided
+            if 'artifacts' in experiment_data and isinstance(experiment_data['artifacts'], list):
+                # De-duplicate while preserving order
+                seen = set(merged_artifacts)
+                for a in experiment_data['artifacts']:
+                    if a not in seen:
+                        merged_artifacts.append(a)
+                        seen.add(a)
+
+            # Build the experiment payload to upsert
             current_experiment = {
-                'experiment_id': self.experiment_id or f"exp_{datetime.now().strftime('%Y%m%d_%H%M%S')}",
-                'name': self.experiment_name,
-                'description': "SmolLM3 fine-tuning experiment",
-                'created_at': self.start_time.isoformat(),
-                'status': 'running',
-                'metrics': json.dumps(self.metrics_history, default=str),
-                'parameters': json.dumps(experiment_data, default=str),
-                'artifacts': json.dumps(self.artifacts, default=str),
-                'logs': json.dumps([], default=str),
+                'experiment_id': experiment_id,
+                'name': existing.get('name') or self.experiment_name,
+                'description': existing.get('description') or "SmolLM3 fine-tuning experiment",
+                'created_at': existing.get('created_at') or self.start_time.isoformat(),
+                'status': existing.get('status') or 'running',
+                'metrics': json.dumps(merged_metrics, default=str),
+                'parameters': json.dumps(merged_parameters, default=str),
+                'artifacts': json.dumps(merged_artifacts, default=str),
+                'logs': json.dumps(existing_logs, default=str),
                 'last_updated': datetime.now().isoformat()
             }
-            
-            # Use dataset manager to safely upsert the experiment
+
             success = self.dataset_manager.upsert_experiment(current_experiment)
-            
+
             if success:
                 logger.info(f"✅ Experiment data saved to HF Dataset: {self.dataset_repo}")
                 return True
             else:
-                logger.error(f"❌ Failed to save experiment data to HF Dataset")
+                logger.error("❌ Failed to save experiment data to HF Dataset")
                 return False
-                
+
         except Exception as e:
             logger.error(f"❌ Failed to save to HF Dataset: {e}")
             return False

templates/spaces/trackio/README.md

@@ -6,7 +6,7 @@ colorTo: yellow
 sdk: gradio
 sdk_version: 5.38.0
 app_file: app.py
-pinned: true
+pinned: false
 license: mit
 short_description: trackio for training monitoring
 tags:

templates/spaces/trackio/app.py

@@ -37,16 +37,20 @@ class TrackioSpace:
         self.dataset_manager = None
         if self.hf_token and self.dataset_repo:
             try:
-                # Import dataset manager
-                import sys
-                sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', '..', 'src'))
-                from dataset_utils import TrackioDatasetManager
+                # Prefer local dataset_utils in Space repo
+                from dataset_utils import TrackioDatasetManager  # type: ignore
                 self.dataset_manager = TrackioDatasetManager(self.dataset_repo, self.hf_token)
-                logger.info("✅ Dataset manager initialized for safe operations")
-            except ImportError:
-                logger.warning("⚠️ Dataset manager not available, using legacy data handling")
-            except Exception as e:
-                logger.warning(f"⚠️ Failed to initialize dataset manager: {e}")
+                logger.info("✅ Dataset manager initialized for safe operations (local)")
+            except Exception as local_e:
+                try:
+                    # Fallback: try project src layout if present
+                    import sys
+                    sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', '..', 'src'))
+                    from dataset_utils import TrackioDatasetManager  # type: ignore
+                    self.dataset_manager = TrackioDatasetManager(self.dataset_repo, self.hf_token)
+                    logger.info("✅ Dataset manager initialized for safe operations (fallback src)")
+                except Exception as e:
+                    logger.warning(f"⚠️ Dataset manager not available, using legacy data handling: {local_e or e}")
         
         logger.info(f"🔧 Using dataset repository: {self.dataset_repo}")
         
@@ -426,7 +430,11 @@ class TrackioSpace:
         logger.info(f"✅ Loaded {len(backup_experiments)} backup experiments")
     
     def _save_experiments(self):
-        """Save experiments to HF Dataset with data preservation"""
+        """Save experiments to HF Dataset with data preservation
+
+        Note: This saves the full in-memory set. Prefer per-operation upsert via
+        dataset manager when available to reduce overwrite risk.
+        """
         try:
             if self.using_backup_data:
                 logger.warning("⚠️ Using backup data; skip saving to dataset to avoid overwriting with demo values")
@@ -474,6 +482,33 @@ class TrackioSpace:
             logger.error(f"❌ Failed to save experiments: {e}")
             # Fallback to legacy method
             self._save_experiments_legacy()
+
+    def _upsert_experiment(self, experiment_id: str):
+        """Non-destructive upsert of a single experiment when dataset manager is available."""
+        try:
+            if not self.dataset_manager:
+                # Fallback to legacy save of full set
+                self._save_experiments()
+                return
+            exp = self.experiments.get(experiment_id)
+            if not exp:
+                return
+            payload = {
+                'experiment_id': experiment_id,
+                'name': exp.get('name', ''),
+                'description': exp.get('description', ''),
+                'created_at': exp.get('created_at', ''),
+                'status': exp.get('status', 'running'),
+                'metrics': json.dumps(exp.get('metrics', []), default=str),
+                'parameters': json.dumps(exp.get('parameters', {}), default=str),
+                'artifacts': json.dumps(exp.get('artifacts', []), default=str),
+                'logs': json.dumps(exp.get('logs', []), default=str),
+                'last_updated': datetime.now().isoformat()
+            }
+            self.dataset_manager.upsert_experiment(payload)
+        except Exception as e:
+            logger.warning(f"⚠️ Upsert failed, falling back to legacy save: {e}")
+            self._save_experiments()
     
     def _save_experiments_legacy(self):
         """Legacy save method without data preservation (fallback only)"""
@@ -550,7 +585,7 @@ class TrackioSpace:
         
         self.experiments[experiment_id] = experiment
         self.current_experiment = experiment_id
-        self._save_experiments()
+        self._upsert_experiment(experiment_id)
         
         logger.info(f"Created experiment: {experiment_id} - {name}")
         return experiment
@@ -567,7 +602,7 @@ class TrackioSpace:
         }
         
         self.experiments[experiment_id]['metrics'].append(metric_entry)
-        self._save_experiments()
+        self._upsert_experiment(experiment_id)
         logger.info(f"Logged metrics for experiment {experiment_id}: {metrics}")
     
     def log_parameters(self, experiment_id: str, parameters: Dict[str, Any]):
@@ -576,7 +611,7 @@ class TrackioSpace:
             raise ValueError(f"Experiment {experiment_id} not found")
         
         self.experiments[experiment_id]['parameters'].update(parameters)
-        self._save_experiments()
+        self._upsert_experiment(experiment_id)
         logger.info(f"Logged parameters for experiment {experiment_id}: {parameters}")
     
     def log_artifact(self, experiment_id: str, artifact_name: str, artifact_data: str):
@@ -610,7 +645,7 @@ class TrackioSpace:
         """Update experiment status"""
         if experiment_id in self.experiments:
             self.experiments[experiment_id]['status'] = status
-            self._save_experiments()
+            self._upsert_experiment(experiment_id)
             logger.info(f"Updated experiment {experiment_id} status to {status}")
     
     def get_metrics_dataframe(self, experiment_id: str) -> pd.DataFrame: