feat: wandb local converter table (#1469)

* try

* refactor

* fix

Author: Zeyi-Lin

swanlab/converter/wb/wb_local_converter.py

@@ -34,6 +34,7 @@
 
 import yaml
 import swanlab
+from swanlab import echarts
 from swanlab.log import swanlog as swl
 from swanlab.data.porter import DataPorter
 from swanlab.env import create_time
@@ -143,6 +144,28 @@ def _proto_items_to_dict(self, items) -> dict:
                 swl.warning(f"Could not decode json for key '{key}': {item.value_json}")
         return mapping
 
+    def _validate_path(self, base_dir: str, file_path: str) -> Optional[str]:
+        """Validate that file_path stays within base_dir to prevent path traversal."""
+        if not file_path:
+            return None
+        abs_base = os.path.abspath(base_dir)
+        abs_path = os.path.abspath(os.path.join(base_dir, file_path))
+        if not abs_path.startswith(abs_base + os.sep) and abs_path != abs_base:
+            swl.warning(f"Path traversal attempt detected: {file_path}")
+            return None
+        return abs_path
+
+    def _filter_text_columns(self, columns, data):
+        """Filter out non-text columns (images, media) from table data."""
+        non_text_indices = {
+            i for row in data for i, cell in enumerate(row) if isinstance(cell, dict) and "_type" in cell
+        }
+        text_indices = [i for i in range(len(columns)) if i not in non_text_indices]
+        return (
+            [columns[i] for i in text_indices],
+            [[row[i] for i in text_indices if i < len(row)] for row in data]
+        )
+
     def _parse_run(self, run_dir: str):
         """Parses a single wandb run directory and converts it to a SwanLab run."""
         wandb_files = glob.glob(os.path.join(run_dir, "*.wandb"))
@@ -166,38 +189,38 @@ def _parse_run(self, run_dir: str):
 
         # Direct-write state (bypasses swanlab_run.log() overhead for scalar metrics)
         porter = None           # set after swanlab.init()
-        _conv_column_kids = {}  # key -> kid str
-        _conv_epoch_counters = {}  # key -> cumulative epoch count
-        _upload_pre_count = [0]  # items uploaded before upload progress bar appears
-        _total_scalars_logged = [0]  # total scalar data points logged
+        column_kids = {}  # key -> kid str
+        epoch_counters = {}  # key -> cumulative epoch count
+        upload_pre_count = 0  # items uploaded before upload progress bar appears
+        total_scalars_logged = 0  # total scalar data points logged
 
         def _log_scalars_direct(scalars, step):
             """Write float scalars directly to the porter, bypassing log() overhead."""
+            nonlocal total_scalars_logged
             if not scalars or porter is None:
                 return
             for key in scalars:
-                if key not in _conv_column_kids:
-                    kid = len(_conv_column_kids)
-                    _conv_column_kids[key] = str(kid)
+                if key not in column_kids:
+                    kid = len(column_kids)
+                    column_kids[key] = str(kid)
                     split_key = key.split("/")
                     sname = split_key[0] if len(split_key) > 1 and split_key[0] else None
                     porter.trace_column(ColumnInfo(
                         key=key, kid=str(kid), name=key, cls='CUSTOM',
                         chart_type=ChartType.LINE, chart_reference='STEP',
                         section_name=sname, section_type="PUBLIC",
                     ))
-            for key in scalars:
-                _conv_epoch_counters[key] = _conv_epoch_counters.get(key, 0) + 1
-            _total_scalars_logged[0] += len(scalars)
-            porter.trace_scalars_step(step, scalars, dict(_conv_epoch_counters), create_time())
+                epoch_counters[key] = epoch_counters.get(key, 0) + 1
+            total_scalars_logged += len(scalars)
+            porter.trace_scalars_step(step, scalars, dict(epoch_counters), create_time())
 
         def _finish_with_progress():
             """Run swanlab_run.finish() while showing a Rich upload progress bar."""
             _pool = porter._pool if porter is not None else None
             if _pool is None:
                 swanlab_run.finish()
                 return
-            total = len(_conv_column_kids) + _total_scalars_logged[0]
+            total = len(column_kids) + total_scalars_logged
             up = Progress(
                 TextColumn("[bold green]{task.description}"),
                 BarColumn(bar_width=40),
@@ -206,32 +229,34 @@ def _finish_with_progress():
                 TimeRemainingColumn(),
             )
             up.start()
-            t = up.add_task("Uploading to SwanLab", total=total, completed=_upload_pre_count[0])
+            t = up.add_task("Uploading to SwanLab", total=total, completed=upload_pre_count)
 
-            _last_completed = [_upload_pre_count[0]]
-            _stall_check_time = [time.time()]
+            last_completed = upload_pre_count
+            stall_check_time = time.time()
 
             def _upload_cb(n):
-                _last_completed[0] += n
-                _stall_check_time[0] = time.time()
+                nonlocal last_completed, stall_check_time
+                last_completed += n
+                stall_check_time = time.time()
                 up.update(t, advance=n)
 
             _pool.collector.upload_callback = _upload_cb
 
             # Monitor for stalls and update description
             import threading
-            _stop_monitor = [False]
+            stop_monitor = False
             def _monitor_stalls():
-                while not _stop_monitor[0]:
+                nonlocal stop_monitor
+                while not stop_monitor:
                     time.sleep(1)
-                    if _stop_monitor[0]:
+                    if stop_monitor:
                         break
-                    elapsed = time.time() - _stall_check_time[0]
-                    if elapsed > 5 and _last_completed[0] < total:
-                        remaining = total - _last_completed[0]
+                    elapsed = time.time() - stall_check_time
+                    if elapsed > 5 and last_completed < total:
+                        remaining = total - last_completed
                         batch_info = f"batch ~{min(1000, remaining)} items" if remaining > 0 else "final batch"
                         up.update(t, description=f"[bold yellow]Uploading to SwanLab (processing {batch_info}, {int(elapsed)}s)")
-                    elif _last_completed[0] < total:
+                    elif last_completed < total:
                         up.update(t, description="[bold green]Uploading to SwanLab")
 
             monitor_thread = threading.Thread(target=_monitor_stalls, daemon=True)
@@ -244,7 +269,7 @@ def _monitor_stalls():
                 swanlab_run.finish()
             finally:
                 swl.info = _orig_info
-                _stop_monitor[0] = True
+                stop_monitor = True
             up.update(t, completed=total, description="[bold green]Uploading to SwanLab")
             up.stop()
 
@@ -283,7 +308,8 @@ def initialize_swanlab_run_if_needed():
             # Set pre-count callback: track items uploaded during parse before the upload bar appears
             if porter is not None and porter._pool is not None:
                 def _pre_upload_cb(n):
-                    _upload_pre_count[0] += n
+                    nonlocal upload_pre_count
+                    upload_pre_count += n
                 porter._pool.collector.upload_callback = _pre_upload_cb
 
             # 恢复进度条
@@ -369,39 +395,79 @@ def _pre_upload_cb(n):
                 initialize_swanlab_run_if_needed()
                 scalar_dict = {}
                 media_dict = {}
+                grouped_items = {}
                 step = 0
-                # Single-pass: parse proto items directly, fast-path float() for scalars
+                # First pass: group items by base key
                 for item in record_pb.history.item:
                     key = item.key or '/'.join(item.nested_key)
                     if not key:
                         continue
-                    vj = item.value_json
+                    value_json = item.value_json
                     if key == '_step':
                         try:
-                            step = int(float(vj))
+                            step = int(float(value_json))
                         except (ValueError, TypeError):
                             pass
                         continue
                     if key.startswith('_'):
                         continue
-                    # Fast path: direct float conversion (avoids full JSON parse for scalars)
-                    try:
-                        scalar_dict[key] = float(vj)
-                        continue
-                    except (ValueError, TypeError):
-                        pass
-                    # Slow path: full JSON parse for complex types (media, etc.)
-                    try:
-                        value = _json_loads(vj)
-                    except (ValueError, Exception):
-                        continue
-                    if isinstance(value, int):
-                        scalar_dict[key] = float(value)
-                    elif isinstance(value, dict) and "_type" in value:
-                        media_type = value["_type"]
-                        path = os.path.join(files_root_dir, value.get("path", ""))
-                        if os.path.exists(path) and media_type == "image-file":
-                            media_dict[key] = swanlab.Image(path)
+                    # Check if key has nested structure (e.g., "table/_type")
+                    if '/' in key:
+                        base_key, sub_key = key.split('/', 1)
+                        if base_key not in grouped_items:
+                            grouped_items[base_key] = {}
+                        try:
+                            grouped_items[base_key][sub_key] = _json_loads(value_json)
+                        except (ValueError, Exception):
+                            grouped_items[base_key][sub_key] = value_json
+                    else:
+                        # Fast path: direct float conversion for scalars
+                        try:
+                            scalar_dict[key] = float(value_json)
+                            continue
+                        except (ValueError, TypeError):
+                            pass
+                        # Slow path: full JSON parse
+                        try:
+                            value = _json_loads(value_json)
+                            if isinstance(value, int):
+                                scalar_dict[key] = float(value)
+                            elif isinstance(value, dict) and "path" in value:
+                                validated_path = self._validate_path(files_root_dir, value["path"])
+                                if validated_path and os.path.exists(validated_path):
+                                    if value.get("_type") == "image-file":
+                                        media_dict[key] = swanlab.Image(validated_path)
+                                    elif value.get("_type") == "audio-file":
+                                        media_dict[key] = swanlab.Audio(validated_path)
+                        except (ValueError, Exception):
+                            pass
+
+                # Second pass: process grouped items for tables and media
+                for base_key, props in grouped_items.items():
+                    if props.get('_type') == 'table-file' and 'path' in props:
+                        validated_path = self._validate_path(files_root_dir, props['path'])
+                        if validated_path and os.path.exists(validated_path):
+                            try:
+                                with open(validated_path, 'r', encoding='utf-8') as f:
+                                    table_data = _json_loads(f.read())
+                                columns = table_data.get("columns", [])
+                                data = table_data.get("data", [])
+                                # Filter text-only columns
+                                filtered_cols, filtered_data = self._filter_text_columns(columns, data)
+                                if filtered_cols:
+                                    table = echarts.Table()
+                                    table.add(filtered_cols, filtered_data)
+                                    media_dict[base_key] = table
+                            except Exception as e:
+                                swl.warning(f"Failed to parse table from {validated_path}: {e}")
+                    elif props.get('_type') == 'image-file' and 'path' in props:
+                        validated_path = self._validate_path(files_root_dir, props['path'])
+                        if validated_path and os.path.exists(validated_path):
+                            media_dict[base_key] = swanlab.Image(validated_path)
+                    elif props.get('_type') == 'audio-file' and 'path' in props:
+                        validated_path = self._validate_path(files_root_dir, props['path'])
+                        if validated_path and os.path.exists(validated_path):
+                            media_dict[base_key] = swanlab.Audio(validated_path)
 
                 if scalar_dict or media_dict:
                     if scalar_dict:
@@ -416,17 +482,13 @@ def _pre_upload_cb(n):
                 # would result in N redundant log calls (N = number of steps).
                 for item in record_pb.summary.update:
                     key = item.key or '/'.join(item.nested_key)
-                    if not key or key.startswith('_'):
+                    if not key or key.startswith('_') or '/' in key:
                         continue
                     try:
                         last_summary[key] = float(item.value_json)
                     except (ValueError, TypeError):
                         pass
 
-            # 清理公共变量，释放内存
-            del record_pb
-            del record_bin
-
             # GC every GC_INTERVAL records to reduce overhead
             record_count += 1
             if record_count % GC_INTERVAL == 0:
@@ -436,12 +498,6 @@ def _pre_upload_cb(n):
         if progress is not None:
             progress.stop()
 
-        # Log 最终的 summary 数据（只 log 一次，避免对每条 summary record 都 log）
-        if last_summary:
-            initialize_swanlab_run_if_needed()
-            if swanlab_run:
-                _log_scalars_direct(last_summary, last_summary_step)
-
         if swanlab_run:
             swl.info(f"Finished Parsing run: {run_metadata['name']}")
             _finish_with_progress()