Upload folder using huggingface_hub

scripts/cache_dataset.py

@@ -12,6 +12,8 @@ from tqdm import tqdm
 from dotenv import load_dotenv
 import huggingface_hub
 import logging
+from concurrent.futures import ProcessPoolExecutor, as_completed
+import multiprocessing as mp
 
 # Suppress noisy libraries
 logging.getLogger("httpx").setLevel(logging.WARNING)
@@ -21,8 +23,6 @@ logging.getLogger("huggingface_hub").setLevel(logging.WARNING)
 # Add parent directory to path to import modules
 sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
 
-from data.data_loader import OracleDataset
-from data.data_fetcher import DataFetcher
 from scripts.analyze_distribution import get_return_class_map
 # Import quality score calculator
 from scripts.compute_quality_score import get_token_quality_scores, fetch_token_metrics, _bucket_id, _midrank_percentiles, EPS
@@ -30,18 +30,161 @@ from scripts.compute_quality_score import get_token_quality_scores, fetch_token_
 from clickhouse_driver import Client as ClickHouseClient
 from neo4j import GraphDatabase
 
+# Global variables for worker processes (initialized per-worker)
+_worker_dataset = None
+_worker_return_class_map = None
+_worker_quality_scores_map = None
+
+
+def _init_worker(db_config, dataset_config, return_class_map, quality_scores_map):
+    """Initialize worker process with its own DB connections and dataset."""
+    global _worker_dataset, _worker_return_class_map, _worker_quality_scores_map
+
+    from data.data_loader import OracleDataset
+    from data.data_fetcher import DataFetcher
+
+    # Each worker gets its own DB connections
+    clickhouse_client = ClickHouseClient(
+        host=db_config['clickhouse_host'],
+        port=db_config['clickhouse_port']
+    )
+    neo4j_driver = GraphDatabase.driver(
+        db_config['neo4j_uri'],
+        auth=(db_config['neo4j_user'], db_config['neo4j_password'])
+    )
+
+    data_fetcher = DataFetcher(clickhouse_client=clickhouse_client, neo4j_driver=neo4j_driver)
+
+    _worker_dataset = OracleDataset(
+        data_fetcher=data_fetcher,
+        max_samples=dataset_config['max_samples'],
+        start_date=dataset_config['start_date'],
+        ohlc_stats_path=dataset_config['ohlc_stats_path'],
+        horizons_seconds=dataset_config['horizons_seconds'],
+        quantiles=dataset_config['quantiles'],
+        min_trade_usd=dataset_config['min_trade_usd'],
+        max_seq_len=dataset_config['max_seq_len']
+    )
+
+    # Set the filtered mints
+    _worker_dataset.sampled_mints = dataset_config['sampled_mints']
+
+    _worker_return_class_map = return_class_map
+    _worker_quality_scores_map = quality_scores_map
+
+
+def _process_single_token_context(args):
+    """Worker function to process a single token in context mode."""
+    idx, mint_addr, samples_per_token, output_dir = args
+
+    global _worker_dataset, _worker_return_class_map, _worker_quality_scores_map
+
+    try:
+        class_id = _worker_return_class_map.get(mint_addr)
+        if class_id is None:
+            return {'status': 'skipped', 'reason': 'not in class map', 'mint': mint_addr}
+
+        # Generate contexts
+        contexts = _worker_dataset.__cacheitem_context__(idx, num_samples_per_token=samples_per_token)
+
+        if not contexts:
+            return {'status': 'skipped', 'reason': 'no valid contexts', 'mint': mint_addr}
+
+        q_score = _worker_quality_scores_map.get(mint_addr)
+        if q_score is None:
+            return {'status': 'skipped', 'reason': 'no quality score', 'mint': mint_addr}
+
+        # Save contexts - use mint_addr hash for unique filenames
+        saved_files = []
+        for ctx_idx, ctx in enumerate(contexts):
+            ctx["quality_score"] = q_score
+            ctx["class_id"] = class_id
+            ctx["source_token"] = mint_addr
+            ctx["cache_mode"] = "context"
+
+            # Use hash-based filename to avoid conflicts
+            filename = f"sample_{mint_addr[:16]}_{ctx_idx}.pt"
+            output_path = Path(output_dir) / filename
+            torch.save(ctx, output_path)
+            saved_files.append(filename)
+
+        return {
+            'status': 'success',
+            'mint': mint_addr,
+            'class_id': class_id,
+            'q_score': q_score,
+            'n_contexts': len(contexts),
+            'n_events': len(contexts[0].get('event_sequence', [])) if contexts else 0,
+            'files': saved_files
+        }
+
+    except Exception as e:
+        import traceback
+        return {
+            'status': 'error',
+            'mint': mint_addr,
+            'error': str(e),
+            'traceback': traceback.format_exc()
+        }
+
+
+def _process_single_token_raw(args):
+    """Worker function to process a single token in raw mode."""
+    idx, mint_addr, output_dir = args
+
+    global _worker_dataset, _worker_return_class_map, _worker_quality_scores_map
+
+    try:
+        class_id = _worker_return_class_map.get(mint_addr)
+        if class_id is None:
+            return {'status': 'skipped', 'reason': 'not in class map', 'mint': mint_addr}
+
+        item = _worker_dataset.__cacheitem__(idx)
+
+        if item is None:
+            return {'status': 'skipped', 'reason': 'cacheitem returned None', 'mint': mint_addr}
+
+        q_score = _worker_quality_scores_map.get(mint_addr)
+        if q_score is None:
+            return {'status': 'skipped', 'reason': 'no quality score', 'mint': mint_addr}
+
+        item["quality_score"] = q_score
+        item["class_id"] = class_id
+        item["cache_mode"] = "raw"
+
+        filename = f"sample_{mint_addr[:16]}.pt"
+        output_path = Path(output_dir) / filename
+        torch.save(item, output_path)
+
+        return {
+            'status': 'success',
+            'mint': mint_addr,
+            'class_id': class_id,
+            'q_score': q_score,
+            'n_trades': len(item.get('trades', [])),
+            'files': [filename]
+        }
+
+    except Exception as e:
+        import traceback
+        return {
+            'status': 'error',
+            'mint': mint_addr,
+            'error': str(e),
+            'traceback': traceback.format_exc()
+        }
+
+
 def compute_save_ohlc_stats(client: ClickHouseClient, output_path: str):
     """
     Computes global mean/std for price/volume from ClickHouse and saves to .npz
     This allows the dataset loader to normalize inputs correctly.
     """
     print(f"INFO: Computing OHLC stats (mean/std) from ClickHouse...")
-    
-    # Query matching preprocess_distribution.py logic
-    # We use hardcoded min_price/vol filters to avoid skewing stats with dust
+
     min_price = 0.0
     min_vol = 0.0
-    
+
     query = """
         SELECT
             AVG(t.price_usd)         AS mean_price_usd,
@@ -53,9 +196,9 @@ def compute_save_ohlc_stats(client: ClickHouseClient, output_path: str):
         FROM trades AS t
         WHERE t.price_usd > %(min_price)s AND t.total_usd > %(min_vol)s
     """
-    
+
     params = {"min_price": min_price, "min_vol": min_vol}
-    
+
     try:
         result = client.execute(query, params=params)
         if not result or not result[0]:
@@ -67,10 +210,9 @@ def compute_save_ohlc_stats(client: ClickHouseClient, output_path: str):
             }
         else:
             row = result[0]
-            # Handle potential None values if DB is empty
             def safe_float(x, default=0.0):
                 return float(x) if x is not None else default
-            
+
             def safe_std(x):
                 val = safe_float(x, 1.0)
                 return val if val > 1e-9 else 1.0
@@ -83,29 +225,24 @@ def compute_save_ohlc_stats(client: ClickHouseClient, output_path: str):
                 "mean_trade_value_usd": safe_float(row[4]),
                 "std_trade_value_usd": safe_std(row[5]),
             }
-            
-        # Save to NPZ
+
         out_p = Path(output_path)
         out_p.parent.mkdir(parents=True, exist_ok=True)
         np.savez(out_p, **stats)
-        
+
         print(f"INFO: Saved OHLC stats to {out_p}")
         for k, v in stats.items():
             print(f"  {k}: {v:.4f}")
-            
+
     except Exception as e:
         print(f"ERROR: Failed to compute OHLC stats: {e}")
-        # Don't crash, let it try to proceed (though dataset might complain if file missing)
+
 
 def build_quality_missing_reason_map(client: ClickHouseClient, max_ret: float = 1e9):
-    """
-    Build a map: token_address -> reason string for why a quality score is missing.
-    This mirrors compute_quality_scores filtering and feature availability.
-    """
+    """Build a map: token_address -> reason string for why a quality score is missing."""
     data = fetch_token_metrics(client)
     metrics_by_token = {d.get("token_address"): d for d in data if d.get("token_address")}
 
-    # Build buckets with the same return filtering as compute_quality_scores
     buckets = {}
     for d in data:
         ret_val = d.get("ret")
@@ -117,7 +254,6 @@ def build_quality_missing_reason_map(client: ClickHouseClient, max_ret: float =
         d["bucket_id"] = b
         buckets.setdefault(b, []).append(d)
 
-    # Same feature definitions as compute_quality_scores
     feature_defs = [
         ("fees_log", lambda d: math.log1p(d["fees_sol"]) if d.get("fees_sol") is not None else None, True),
         ("volume_log", lambda d: math.log1p(d["volume_usd"]) if d.get("volume_usd") is not None else None, True),
@@ -132,7 +268,6 @@ def build_quality_missing_reason_map(client: ClickHouseClient, max_ret: float =
         ("dev_hold_pct", lambda d: d.get("dev_hold_pct"), True),
     ]
 
-    # Precompute percentiles per bucket + feature
     bucket_feature_percentiles = {}
     for b, items in buckets.items():
         feature_percentiles = {}
@@ -149,10 +284,10 @@ def build_quality_missing_reason_map(client: ClickHouseClient, max_ret: float =
     def _reason_for(token_address: str) -> str:
         d = metrics_by_token.get(token_address)
         if not d:
-            return "no metrics found (missing from token_metrics/trades/mints joins)"
+            return "no metrics found"
         ret_val = d.get("ret")
         if ret_val is None:
-            return "ret is None (missing ATH/launch metrics)"
+            return "ret is None"
         if ret_val <= 0:
             return f"ret <= 0 ({ret_val})"
         if ret_val > max_ret:
@@ -160,27 +295,17 @@ def build_quality_missing_reason_map(client: ClickHouseClient, max_ret: float =
         b = _bucket_id(ret_val)
         if b == -1:
             return f"ret {ret_val} not in RETURN_THRESHOLDS"
-        items = buckets.get(b, [])
-        if not items:
-            return f"bucket {b} empty after filtering"
-        feature_percentiles = bucket_feature_percentiles.get(b, {})
-        has_any = False
-        missing_features = []
-        for fname, _fget, _pos in feature_defs:
-            if feature_percentiles.get(fname, {}).get(token_address) is None:
-                missing_features.append(fname)
-            else:
-                has_any = True
-        if not has_any:
-            return "no valid feature percentiles for token (all features missing/invalid)"
-        return f"unexpected: has feature percentiles but no score; missing features={','.join(missing_features)}"
+        return "unknown"
 
     return _reason_for
 
+
 def main():
     load_dotenv()
-    
-    # Explicit Login
+
+    # Use spawn method for multiprocessing (safer with CUDA/DB connections)
+    mp.set_start_method('spawn', force=True)
+
     hf_token = os.getenv("HF_TOKEN")
     if hf_token:
         print(f"INFO: Logging in to Hugging Face with token starting with: {hf_token[:4]}...")
@@ -195,15 +320,19 @@ def main():
     parser.add_argument("--ohlc_stats_path", type=str, default="data/ohlc_stats.npz")
     parser.add_argument("--min_trade_usd", type=float, default=0.0)
 
-    # NEW: Context caching mode args
+    # Context caching mode args
     parser.add_argument("--cache_mode", type=str, default="raw", choices=["raw", "context"],
-                        help="Cache mode: 'raw' caches raw token data (old behavior), 'context' caches fully processed training contexts (new behavior)")
+                        help="Cache mode: 'raw' or 'context'")
     parser.add_argument("--context_length", type=int, default=8192,
-                        help="Max sequence length for context caching mode. Triggers H/B/H dynamic sampling when events exceed this limit.")
+                        help="Max sequence length for H/B/H threshold")
     parser.add_argument("--min_trades", type=int, default=10,
-                        help="Minimum number of trades required for T_cutoff sampling. Tokens with fewer trades are skipped.")
+                        help="Minimum trades for T_cutoff sampling")
     parser.add_argument("--samples_per_token", type=int, default=1,
-                        help="Number of different T_cutoff samples to generate per token in context mode.")
+                        help="Number of T_cutoff samples per token")
+
+    # Parallelization args
+    parser.add_argument("--num_workers", type=int, default=1,
+                        help="Number of parallel workers (default: 1, use 0 for auto-detect)")
 
     # DB Args
     parser.add_argument("--clickhouse_host", type=str, default=os.getenv("CLICKHOUSE_HOST", "localhost"))
@@ -213,36 +342,40 @@ def main():
     parser.add_argument("--neo4j_password", type=str, default=os.getenv("NEO4J_PASSWORD", "password"))
 
     args = parser.parse_args()
-    
+
+    # Auto-detect workers if set to 0
+    if args.num_workers == 0:
+        args.num_workers = max(1, mp.cpu_count() - 4)
+
     output_dir = Path(args.output_dir)
     output_dir.mkdir(parents=True, exist_ok=True)
-    
+
     start_date_dt = None
     if args.start_date:
         start_date_dt = datetime.datetime.strptime(args.start_date, "%Y-%m-%d")
-        
+
     print(f"INFO: Initializing DB Connections...")
     clickhouse_client = ClickHouseClient(host=args.clickhouse_host, port=args.clickhouse_port)
     neo4j_driver = GraphDatabase.driver(args.neo4j_uri, auth=(args.neo4j_user, args.neo4j_password))
-    
+
     try:
         # --- 1. Compute OHLC Stats (Global) ---
         compute_save_ohlc_stats(clickhouse_client, args.ohlc_stats_path)
 
-        # --- 2. Initialize DataFetcher and OracleDataset ---
+        # --- 2. Initialize DataFetcher and OracleDataset (main process) ---
+        from data.data_loader import OracleDataset
+        from data.data_fetcher import DataFetcher
+
         data_fetcher = DataFetcher(clickhouse_client=clickhouse_client, neo4j_driver=neo4j_driver)
-        
-        # Pre-fetch the Return Class Map
+
         print("INFO: Fetching Return Classification Map...")
         return_class_map, thresholds = get_return_class_map(clickhouse_client)
         print(f"INFO: Loaded {len(return_class_map)} valid classified tokens.")
 
-        # Pre-fetch Quality Scores
         print("INFO: Fetching Token Quality Scores...")
         quality_scores_map = get_token_quality_scores(clickhouse_client)
-        quality_missing_reason = build_quality_missing_reason_map(clickhouse_client, max_ret=1e9)
         print(f"INFO: Loaded {len(quality_scores_map)} quality scores.")
-    
+
         dataset = OracleDataset(
             data_fetcher=data_fetcher,
             max_samples=args.max_samples,
@@ -251,19 +384,18 @@ def main():
             horizons_seconds=[60, 180, 300, 600, 1800, 3600, 7200],
             quantiles=[0.5],
             min_trade_usd=args.min_trade_usd,
-            max_seq_len=args.context_length  # Pass context_length for H/B/H threshold
+            max_seq_len=args.context_length
         )
-    
+
         if len(dataset) == 0:
             print("WARNING: Dataset initialization resulted in 0 samples. Nothing to cache.")
             return
 
-        # --- FILTER DATASET BY CLASS MAP ---
-        # Only keep mints that are classified (valid return, sufficient data)
+        # Filter dataset by class map
         original_size = len(dataset)
         print(f"INFO: Filtering dataset... Original size: {original_size}")
         dataset.sampled_mints = [
-            m for m in dataset.sampled_mints 
+            m for m in dataset.sampled_mints
             if m['mint_address'] in return_class_map
         ]
         filtered_size = len(dataset)
@@ -274,156 +406,101 @@ def main():
             print("WARNING: No tokens remain after filtering by return_class_map.")
             return
 
-        # --- 3. Iterate and cache based on mode ---
+        # --- 3. Parallel caching ---
         print(f"INFO: Cache mode: {args.cache_mode}")
-        print(f"INFO: Starting to generate and cache from {len(dataset)} tokens...")
+        print(f"INFO: Number of workers: {args.num_workers}")
+        print(f"INFO: Starting to cache {len(dataset)} tokens...")
+
+        # Prepare configs for workers
+        db_config = {
+            'clickhouse_host': args.clickhouse_host,
+            'clickhouse_port': args.clickhouse_port,
+            'neo4j_uri': args.neo4j_uri,
+            'neo4j_user': args.neo4j_user,
+            'neo4j_password': args.neo4j_password,
+        }
+
+        dataset_config = {
+            'max_samples': args.max_samples,
+            'start_date': start_date_dt,
+            'ohlc_stats_path': args.ohlc_stats_path,
+            'horizons_seconds': [60, 180, 300, 600, 1800, 3600, 7200],
+            'quantiles': [0.5],
+            'min_trade_usd': args.min_trade_usd,
+            'max_seq_len': args.context_length,
+            'sampled_mints': dataset.sampled_mints,  # Pass filtered mints
+        }
+
+        # Prepare task list
+        tasks = []
+        for i in range(len(dataset)):
+            mint_addr = dataset.sampled_mints[i]['mint_address']
+            if args.cache_mode == "context":
+                tasks.append((i, mint_addr, args.samples_per_token, str(output_dir)))
+            else:
+                tasks.append((i, mint_addr, str(output_dir)))
 
+        # Track results
+        success_count = 0
         skipped_count = 0
-        cached_count = 0
-        global_sample_idx = 0  # Global counter for unique sample filenames
-
-        # Track class distribution for balanced sampling metadata
+        error_count = 0
         class_distribution = {}
 
-        if args.cache_mode == "context":
-            # =========================================================================
-            # CONTEXT MODE: Cache fully processed training contexts
-            # - Samples T_cutoff during caching (non-deterministic moved to cache time)
-            # - Applies H/B/H dynamic sampling based on context_length
-            # - Avoids caching tokens that won't be seen (garbage filtered out)
-            # - Training becomes fully deterministic (just loads cached contexts)
-            # =========================================================================
-            print(f"INFO: Context mode settings:")
-            print(f"  - context_length (H/B/H threshold): {args.context_length}")
-            print(f"  - min_trades (T_cutoff threshold): {args.min_trades}")
-            print(f"  - samples_per_token: {args.samples_per_token}")
-
-            for i in tqdm(range(len(dataset)), desc="Caching contexts"):
-                mint_addr = dataset.sampled_mints[i]['mint_address']
-                class_id = return_class_map[mint_addr]
-
-                try:
-                    # Generate multiple training contexts per token
-                    contexts = dataset.__cacheitem_context__(i, num_samples_per_token=args.samples_per_token)
-
-                    if not contexts:
-                        skipped_count += 1
-                        continue
-
-                    # Require quality score
-                    if mint_addr not in quality_scores_map:
-                        reason = quality_missing_reason(mint_addr)
-                        raise RuntimeError(
-                            f"Missing quality score for mint {mint_addr}. Reason: {reason}."
-                        )
-                    q_score = quality_scores_map[mint_addr]
-
-                    # Save each context as a separate sample
-                    for ctx in contexts:
-                        ctx["quality_score"] = q_score
-                        ctx["class_id"] = class_id
-                        ctx["source_token"] = mint_addr  # Track origin for debugging
-                        ctx["cache_mode"] = "context"
-
-                        filename = f"sample_{global_sample_idx}.pt"
-                        output_path = output_dir / filename
-                        torch.save(ctx, output_path)
-
-                        # Track class distribution
-                        class_distribution[class_id] = class_distribution.get(class_id, 0) + 1
-
-                        global_sample_idx += 1
-                        cached_count += 1
-
-                    n_events = len(contexts[0].get("event_sequence", [])) if contexts else 0
-                    tqdm.write(
-                        f"  + Cached {len(contexts)} contexts: {mint_addr} | Class: {class_id} | Q: {q_score:.4f} | Events: {n_events}"
-                    )
-
-                except Exception as e:
-                    error_msg = str(e)
-                    if "FATAL" in error_msg or "AuthenticationRateLimit" in error_msg:
-                        print(f"\nCRITICAL: Fatal error processing sample {i}. Stopping.\nError: {e}", file=sys.stderr)
-                        sys.exit(1)
-
-                    print(f"\nERROR: Failed to cache contexts for {mint_addr}. Error: {e}", file=sys.stderr)
-                    import traceback
-                    traceback.print_exc()
+        if args.num_workers == 1:
+            # Single-threaded mode (no multiprocessing overhead)
+            print("INFO: Running in single-threaded mode...")
+            _init_worker(db_config, dataset_config, return_class_map, quality_scores_map)
+
+            process_fn = _process_single_token_context if args.cache_mode == "context" else _process_single_token_raw
+
+            for task in tqdm(tasks, desc="Caching"):
+                result = process_fn(task)
+
+                if result['status'] == 'success':
+                    success_count += 1
+                    cid = result['class_id']
+                    class_distribution[cid] = class_distribution.get(cid, 0) + 1
+                    if args.cache_mode == "context":
+                        tqdm.write(f"  + {result['mint'][:16]} | Class: {cid} | Q: {result['q_score']:.4f} | Contexts: {result['n_contexts']} | Events: {result['n_events']}")
+                    else:
+                        tqdm.write(f"  + {result['mint'][:16]} | Class: {cid} | Q: {result['q_score']:.4f} | Trades: {result['n_trades']}")
+                elif result['status'] == 'skipped':
                     skipped_count += 1
-                    continue
-
+                else:
+                    error_count += 1
+                    tqdm.write(f"  ERROR: {result['mint'][:16]} - {result['error']}")
         else:
-            # =========================================================================
-            # RAW MODE: Cache raw token data (original behavior)
-            # - T_cutoff sampling happens at runtime
-            # - H/B/H applied at runtime
-            # - Non-deterministic training
-            # =========================================================================
-            for i in tqdm(range(len(dataset)), desc="Caching raw samples"):
-                mint_addr = dataset.sampled_mints[i]['mint_address']
-                class_id = return_class_map[mint_addr]
-
-                try:
-                    item = dataset.__cacheitem__(i)
-                    if item is None:
-                        skipped_count += 1
-                        continue
-
-                    if mint_addr not in quality_scores_map:
-                        reason = quality_missing_reason(mint_addr)
-                        raise RuntimeError(
-                            f"Missing quality score for mint {mint_addr}. Reason: {reason}."
-                        )
-                    q_score = quality_scores_map[mint_addr]
-
-                    item["quality_score"] = q_score
-                    item["class_id"] = class_id
-                    item["cache_mode"] = "raw"
-
-                    filename = f"sample_{i}.pt"
-                    output_path = output_dir / filename
-                    torch.save(item, output_path)
-
-                    # Track class distribution
-                    class_distribution[class_id] = class_distribution.get(class_id, 0) + 1
-
-                    cached_count += 1
-
-                    n_trades = len(item.get("trades", []))
-                    n_transfers = len(item.get("transfers", []))
-                    n_pool_creations = len(item.get("pool_creations", []))
-                    n_liquidity_changes = len(item.get("liquidity_changes", []))
-                    n_fee_collections = len(item.get("fee_collections", []))
-                    n_burns = len(item.get("burns", []))
-                    n_supply_locks = len(item.get("supply_locks", []))
-                    n_migrations = len(item.get("migrations", []))
-                    n_mints = 1 if item.get("mint_timestamp") else 0
-                    n_ohlc = len(item.get("ohlc_1s", [])) if item.get("ohlc_1s") is not None else 0
-                    n_snapshots_5m = len(item.get("snapshots_5m", []))
-                    n_holders = len(item.get("holder_snapshots_list", []))
-
-                    tqdm.write(
-                        f"  + Cached: {mint_addr} | Class: {class_id} | Q: {q_score:.4f} | "
-                        f"Events: Mint {n_mints}, Trades {n_trades}, Transfers {n_transfers}, Pool Creations {n_pool_creations}, "
-                        f"Liquidity Changes {n_liquidity_changes}, Fee Collections {n_fee_collections}, "
-                        f"Burns {n_burns}, Supply Locks {n_supply_locks}, Migrations {n_migrations} | "
-                        f"Derived: Ohlc 1s {n_ohlc}, Snapshots 5m {n_snapshots_5m}, Holder Snapshots {n_holders}"
-                    )
-
-                except Exception as e:
-                    error_msg = str(e)
-                    if "FATAL" in error_msg or "AuthenticationRateLimit" in error_msg:
-                        print(f"\nCRITICAL: Fatal error processing sample {i}. Stopping.\nError: {e}", file=sys.stderr)
-                        sys.exit(1)
-
-                    print(f"\nERROR: Failed to cache sample {i} for {mint_addr}. Error: {e}", file=sys.stderr)
-                    import traceback
-                    traceback.print_exc()
-                    skipped_count += 1
-                    continue
-
-        # --- Save class metadata for balanced sampling ---
-        # Build file_class_map for the metadata cache
+            # Multi-process mode
+            print(f"INFO: Running with {args.num_workers} parallel workers...")
+
+            process_fn = _process_single_token_context if args.cache_mode == "context" else _process_single_token_raw
+
+            with ProcessPoolExecutor(
+                max_workers=args.num_workers,
+                initializer=_init_worker,
+                initargs=(db_config, dataset_config, return_class_map, quality_scores_map)
+            ) as executor:
+                futures = {executor.submit(process_fn, task): task for task in tasks}
+
+                for future in tqdm(as_completed(futures), total=len(futures), desc="Caching"):
+                    try:
+                        result = future.result(timeout=300)  # 5 min timeout per token
+
+                        if result['status'] == 'success':
+                            success_count += 1
+                            cid = result['class_id']
+                            class_distribution[cid] = class_distribution.get(cid, 0) + 1
+                        elif result['status'] == 'skipped':
+                            skipped_count += 1
+                        else:
+                            error_count += 1
+                            tqdm.write(f"  ERROR: {result.get('mint', 'unknown')[:16]} - {result.get('error', 'unknown')}")
+                    except Exception as e:
+                        error_count += 1
+                        tqdm.write(f"  WORKER ERROR: {e}")
+
+        # --- 4. Build metadata ---
+        print("INFO: Building class metadata...")
         file_class_map = {}
         for sample_file in sorted(output_dir.glob("sample_*.pt")):
             try:
@@ -437,11 +514,12 @@ def main():
             with open(metadata_path, 'w') as f:
                 json.dump({
                     'file_class_map': file_class_map,
-                    'class_distribution': class_distribution,
+                    'class_distribution': {str(k): v for k, v in class_distribution.items()},
                     'cache_mode': args.cache_mode,
                     'context_length': args.context_length if args.cache_mode == "context" else None,
                     'min_trades': args.min_trades if args.cache_mode == "context" else None,
                     'samples_per_token': args.samples_per_token if args.cache_mode == "context" else None,
+                    'num_workers': args.num_workers,
                 }, f, indent=2)
             print(f"INFO: Saved class metadata to {metadata_path}")
         except Exception as e:
@@ -449,16 +527,19 @@ def main():
 
         print(f"\n--- Caching Complete ---")
         print(f"Cache mode: {args.cache_mode}")
-        print(f"Successfully cached: {cached_count} samples.")
-        print(f"Filtered (Invalid/High Return): {filtered_count} tokens.")
-        print(f"Skipped (Errors/Empty): {skipped_count} tokens.")
+        print(f"Workers used: {args.num_workers}")
+        print(f"Successfully cached: {success_count} tokens")
+        print(f"Total files: {len(file_class_map)}")
+        print(f"Filtered: {filtered_count} tokens")
+        print(f"Skipped: {skipped_count} tokens")
+        print(f"Errors: {error_count} tokens")
         print(f"Class distribution: {class_distribution}")
         print(f"Cache location: {output_dir.resolve()}")
 
     finally:
-        # --- 4. Close connections ---
         clickhouse_client.disconnect()
         neo4j_driver.close()
 
+
 if __name__ == "__main__":
     main()