""" ⚡ ULTRON Systems // Streaming & Sharded Parquet Dataset Loader Safely streams the 90GB HuggingFace OHLCV dataset month-by-month. Guarantees O(1) constant disk and RAM usage to prevent Colab quota violations. """ import os import gc import glob import shutil import urllib.request import numpy as np import pandas as pd import pyarrow.parquet as pq import torch from torch.utils.data import Dataset, DataLoader from typing import Generator, List, Tuple, Optional from huggingface_hub import HfFileSystem, hf_hub_download from config import DatasetConfig from features import compute_stationary_features class ShardedOHLCVDataset(Dataset): """ In-memory dataset representing windows from a single monthly shard. Efficiently created from vectorized sliding window views. """ def __init__(self, windows: np.ndarray, targets: np.ndarray): # windows: [N_samples, seq_len, in_features] # targets: [N_samples, num_horizons] self.windows = torch.from_numpy(windows).float() self.targets = torch.from_numpy(targets).float() def __len__(self): return len(self.windows) def __getitem__(self, idx: int) -> Tuple[torch.Tensor, torch.Tensor]: return self.windows[idx], self.targets[idx] class StreamingShardManager: """ Orchestrates downloading, windowing, and memory pruning for monthly Parquet files. """ def __init__(self, cfg: DatasetConfig): self.cfg = cfg os.makedirs(self.cfg.cache_dir, exist_ok=True) self.parquet_files = self._discover_remote_shards() print(f"⚡ [ULTRON] Discovered {len(self.parquet_files)} target monthly shards between {cfg.start_year} and {cfg.end_year}.") def _discover_remote_shards(self) -> List[str]: """Discovers Parquet file paths in HF repo matching year filters.""" try: fs = HfFileSystem(token=self.cfg.hf_token) all_files = fs.ls(f"datasets/{self.cfg.repo_id}/data", detail=False) except Exception as e: print(f"⚠️ [ULTRON] Direct HfFileSystem query failed ({e}). Constructing standard filenames...") all_files = [] for y in range(self.cfg.start_year, self.cfg.end_year + 1): for m in range(1, 13): all_files.append(f"data/ohlcv_{y}-{m:02d}.parquet") # Filter matching year range and format filtered = [] for f in all_files: fname = os.path.basename(f) if fname.startswith("ohlcv_") and fname.endswith(".parquet"): # Extract year try: parts = fname.replace("ohlcv_", "").replace(".parquet", "").split("-") year = int(parts[0]) if self.cfg.start_year <= year <= self.cfg.end_year: filtered.append(f"data/{fname}") except Exception: continue return sorted(list(set(filtered))) def stream_monthly_shards(self) -> Generator[ShardedOHLCVDataset, None, None]: """ Yields a ShardedOHLCVDataset for each month, cleaning up disk and RAM immediately after. """ for rel_path in self.parquet_files: filename = os.path.basename(rel_path) local_path = os.path.join(self.cfg.cache_dir, filename) print(f"\n⚡ [ULTRON] Fetching shard: {filename} ...") try: # Download single shard (~300MB) directly from HuggingFace with authentication downloaded_file = hf_hub_download( repo_id=self.cfg.repo_id, repo_type="dataset", filename=rel_path, local_dir=self.cfg.cache_dir, token=self.cfg.hf_token, force_download=False ) except Exception as e: print(f"⚠️ [ULTRON] Failed to download {rel_path} via HF Hub: {e}. Trying raw URL fallback...") url = f"https://huggingface.co/datasets/{self.cfg.repo_id}/resolve/main/{rel_path}" try: urllib.request.urlretrieve(url, local_path) downloaded_file = local_path except Exception as ex2: print(f"❌ [ULTRON] Shard {filename} download failed: {ex2}. Skipping.") continue # Read and process shard dataset = self._process_parquet_to_dataset(downloaded_file) # Remove file from disk immediately to preserve Colab storage quota! try: if os.path.exists(downloaded_file): os.remove(downloaded_file) except Exception: pass if dataset is not None and len(dataset) > 0: print(f"✅ [ULTRON] Shard {filename} ingested: Generated {len(dataset)} overlapping windows.") yield dataset else: print(f"⚠️ [ULTRON] Shard {filename} yielded no valid samples. Proceeding to next shard.") # Trigger garbage collection gc.collect() def _process_parquet_to_dataset(self, file_path: str) -> Optional[ShardedOHLCVDataset]: """Reads parquet, computes stationary features, and extracts sliding windows.""" try: # Read only required columns table = pq.read_table( file_path, columns=["timestamp", "open", "high", "low", "close", "volume", "ticker"] ) df = table.to_pandas() del table gc.collect() except Exception as e: print(f"❌ Error reading parquet file {file_path}: {e}") return None # Filter tickers if specified if self.cfg.liquid_tickers: df = df[df["ticker"].isin(self.cfg.liquid_tickers)] if len(df) == 0: del df return None all_windows = [] all_targets = [] # Group by ticker to prevent cross-ticker leakage for ticker, ticker_df in df.groupby("ticker"): feats, targets = compute_stationary_features(ticker_df, self.cfg.pred_horizons) n_rows = len(feats) if n_rows < self.cfg.seq_len: continue # Vectorized overlapping sliding window slicing: # Stride over sequence with window_stride (e.g. 16 minutes) num_windows = (n_rows - self.cfg.seq_len) // self.cfg.window_stride + 1 if num_windows <= 0: continue for i in range(0, n_rows - self.cfg.seq_len + 1, self.cfg.window_stride): w_feats = feats[i : i + self.cfg.seq_len] # Target corresponds to the future return from the final bar of this window w_target = targets[i + self.cfg.seq_len - 1] # Check for NaNs or Inf if not np.isnan(w_target).any() and not np.isnan(w_feats).any(): all_windows.append(w_feats) all_targets.append(w_target) del df gc.collect() if not all_windows: return None windows_arr = np.array(all_windows, dtype=np.float32) targets_arr = np.array(all_targets, dtype=np.float32) return ShardedOHLCVDataset(windows_arr, targets_arr)