ins-pricing 0.3.3__py3-none-any.whl → 0.4.0__py3-none-any.whl

ins_pricing/frontend/ft_workflow.py

@@ -0,0 +1,316 @@
+"""
+FT-Transformer Two-Step Workflow Helper
+Automates the FT → XGB/ResN two-step training process.
+"""
+
+import json
+import copy
+from pathlib import Path
+from typing import Dict, Any, List, Tuple, Optional
+import pandas as pd
+
+
+class FTWorkflowHelper:
+    """
+    Helper for FT-Transformer two-step workflow.
+
+    Step 1: Train FT as unsupervised embedding generator
+    Step 2: Merge embeddings with raw data and train XGB/ResN
+    """
+
+    def __init__(self):
+        self.step1_config = None
+        self.step2_configs = {}
+
+    def prepare_step1_config(
+        self,
+        base_config: Dict[str, Any],
+        output_dir: str = "./ResultsFTUnsupervisedDDP",
+        ft_feature_prefix: str = "ft_emb",
+        use_ddp: bool = True,
+        nproc_per_node: int = 2,
+    ) -> Dict[str, Any]:
+        """
+        Prepare configuration for Step 1: FT unsupervised embedding.
+
+        Args:
+            base_config: Base configuration dictionary
+            output_dir: Output directory for FT embeddings
+            ft_feature_prefix: Prefix for embedding column names
+            use_ddp: Whether to use DDP for FT training
+            nproc_per_node: Number of processes for DDP
+
+        Returns:
+            Step 1 configuration
+        """
+        config = copy.deepcopy(base_config)
+
+        # Set FT role to unsupervised embedding
+        config['ft_role'] = 'unsupervised_embedding'
+        config['ft_feature_prefix'] = ft_feature_prefix
+        config['output_dir'] = output_dir
+        config['cache_predictions'] = True
+        config['prediction_cache_format'] = 'csv'
+
+        # Disable other models in step 1
+        config['stack_model_keys'] = []
+
+        # DDP settings
+        config['use_ft_ddp'] = use_ddp
+        config['use_resn_ddp'] = False
+        config['use_gnn_ddp'] = False
+        config['use_ft_data_parallel'] = False
+        config['use_resn_data_parallel'] = False
+        config['use_gnn_data_parallel'] = False
+
+        # Optuna storage
+        config['optuna_storage'] = f"{output_dir}/optuna/bayesopt.sqlite3"
+        config['optuna_study_prefix'] = 'pricing_ft_unsup'
+
+        # Runner config
+        runner = config.get('runner', {})
+        runner['mode'] = 'entry'
+        runner['model_keys'] = ['ft']
+        runner['nproc_per_node'] = nproc_per_node if use_ddp else 1
+        runner['plot_curves'] = False
+        config['runner'] = runner
+
+        # Disable plotting
+        config['plot_curves'] = False
+        plot_cfg = config.get('plot', {})
+        plot_cfg['enable'] = False
+        config['plot'] = plot_cfg
+
+        self.step1_config = config
+        return config
+
+    def generate_step2_configs(
+        self,
+        step1_config_path: str,
+        target_models: List[str] = None,
+        augmented_data_dir: str = "./DataFTUnsupervised"
+    ) -> Tuple[Optional[Dict[str, Any]], Optional[Dict[str, Any]]]:
+        """
+        Generate Step 2 configurations for XGB and/or ResN.
+
+        This requires that Step 1 has completed and embeddings are cached.
+
+        Args:
+            step1_config_path: Path to the Step 1 config file
+            target_models: Models to train in step 2 (e.g., ['xgb', 'resn'])
+            augmented_data_dir: Directory to save augmented data with embeddings
+
+        Returns:
+            Tuple of (xgb_config, resn_config) - None if not in target_models
+        """
+        if target_models is None:
+            target_models = ['xgb', 'resn']
+
+        # Load step 1 config
+        cfg_path = Path(step1_config_path)
+        with open(cfg_path, 'r', encoding='utf-8') as f:
+            cfg = json.load(f)
+
+        # Read raw data and split
+        model_name = f"{cfg['model_list'][0]}_{cfg['model_categories'][0]}"
+        data_dir = (cfg_path.parent / cfg["data_dir"]).resolve()
+        raw_path = data_dir / f"{model_name}.csv"
+
+        if not raw_path.exists():
+            raise FileNotFoundError(f"Data file not found: {raw_path}")
+
+        raw = pd.read_csv(raw_path)
+
+        # Import split function
+        try:
+            from ins_pricing.cli.utils.cli_common import split_train_test
+        except ImportError:
+            raise ImportError("Cannot import split_train_test. Ensure ins_pricing is installed.")
+
+        # Split data using same settings as step 1
+        holdout_ratio = cfg.get("holdout_ratio", cfg.get("prop_test", 0.25))
+        split_strategy = cfg.get("split_strategy", "random")
+        split_group_col = cfg.get("split_group_col")
+        split_time_col = cfg.get("split_time_col")
+        split_time_ascending = cfg.get("split_time_ascending", True)
+        rand_seed = cfg.get("rand_seed", 13)
+
+        train_df, test_df = split_train_test(
+            raw,
+            holdout_ratio=holdout_ratio,
+            strategy=split_strategy,
+            group_col=split_group_col,
+            time_col=split_time_col,
+            time_ascending=split_time_ascending,
+            rand_seed=rand_seed,
+            reset_index_mode="time_group",
+            ratio_label="holdout_ratio",
+        )
+
+        # Load cached embeddings
+        out_root = (cfg_path.parent / cfg["output_dir"]).resolve()
+        pred_prefix = cfg.get("ft_feature_prefix", "ft_emb")
+        pred_dir = out_root / "Results" / "predictions"
+
+        train_pred_path = pred_dir / f"{model_name}_{pred_prefix}_train.csv"
+        test_pred_path = pred_dir / f"{model_name}_{pred_prefix}_test.csv"
+
+        if not train_pred_path.exists() or not test_pred_path.exists():
+            raise FileNotFoundError(
+                f"Embedding files not found. Run Step 1 first.\n"
+                f"Expected: {train_pred_path} and {test_pred_path}"
+            )
+
+        pred_train = pd.read_csv(train_pred_path)
+        pred_test = pd.read_csv(test_pred_path)
+
+        if len(pred_train) != len(train_df) or len(pred_test) != len(test_df):
+            raise ValueError(
+                "Prediction rows do not match split sizes; check split settings.")
+
+        # Merge embeddings with raw data
+        aug = raw.copy()
+        aug.loc[train_df.index, pred_train.columns] = pred_train.values
+        aug.loc[test_df.index, pred_test.columns] = pred_test.values
+
+        # Save augmented data
+        data_out_dir = cfg_path.parent / augmented_data_dir
+        data_out_dir.mkdir(parents=True, exist_ok=True)
+        aug_path = data_out_dir / f"{model_name}.csv"
+        aug.to_csv(aug_path, index=False)
+
+        # Get embedding column names
+        embed_cols = list(pred_train.columns)
+
+        # Generate configs
+        xgb_config = None
+        resn_config = None
+
+        if 'xgb' in target_models:
+            xgb_config = self._build_xgb_config(cfg, cfg_path, embed_cols, augmented_data_dir)
+            self.step2_configs['xgb'] = xgb_config
+
+        if 'resn' in target_models:
+            resn_config = self._build_resn_config(cfg, cfg_path, embed_cols, augmented_data_dir)
+            self.step2_configs['resn'] = resn_config
+
+        return xgb_config, resn_config
+
+    def _build_xgb_config(
+        self,
+        base_cfg: Dict[str, Any],
+        cfg_path: Path,
+        embed_cols: List[str],
+        data_dir: str
+    ) -> Dict[str, Any]:
+        """Build XGB config for Step 2."""
+        xgb_cfg = copy.deepcopy(base_cfg)
+
+        xgb_cfg["data_dir"] = str(data_dir)
+        xgb_cfg["feature_list"] = base_cfg["feature_list"] + embed_cols
+        xgb_cfg["ft_role"] = "model"
+        xgb_cfg["stack_model_keys"] = ["xgb"]
+        xgb_cfg["cache_predictions"] = False
+
+        # Disable DDP for XGB
+        xgb_cfg["use_resn_ddp"] = False
+        xgb_cfg["use_ft_ddp"] = False
+        xgb_cfg["use_gnn_ddp"] = False
+        xgb_cfg["use_resn_data_parallel"] = False
+        xgb_cfg["use_ft_data_parallel"] = False
+        xgb_cfg["use_gnn_data_parallel"] = False
+
+        xgb_cfg["output_dir"] = "./ResultsXGBFromFTUnsupervised"
+        xgb_cfg["optuna_storage"] = "./ResultsXGBFromFTUnsupervised/optuna/bayesopt.sqlite3"
+        xgb_cfg["optuna_study_prefix"] = "pricing_ft_unsup_xgb"
+        xgb_cfg["loss_name"] = "mse"
+
+        runner_cfg = xgb_cfg.get("runner", {})
+        runner_cfg["model_keys"] = ["xgb"]
+        runner_cfg["nproc_per_node"] = 1
+        runner_cfg["plot_curves"] = False
+        xgb_cfg["runner"] = runner_cfg
+
+        xgb_cfg["plot_curves"] = False
+        plot_cfg = xgb_cfg.get("plot", {})
+        plot_cfg["enable"] = False
+        xgb_cfg["plot"] = plot_cfg
+
+        return xgb_cfg
+
+    def _build_resn_config(
+        self,
+        base_cfg: Dict[str, Any],
+        cfg_path: Path,
+        embed_cols: List[str],
+        data_dir: str
+    ) -> Dict[str, Any]:
+        """Build ResNet config for Step 2."""
+        resn_cfg = copy.deepcopy(base_cfg)
+
+        resn_cfg["data_dir"] = str(data_dir)
+        resn_cfg["feature_list"] = base_cfg["feature_list"] + embed_cols
+        resn_cfg["ft_role"] = "model"
+        resn_cfg["stack_model_keys"] = ["resn"]
+        resn_cfg["cache_predictions"] = False
+
+        # Enable DDP for ResNet
+        resn_cfg["use_resn_ddp"] = True
+        resn_cfg["use_ft_ddp"] = False
+        resn_cfg["use_gnn_ddp"] = False
+        resn_cfg["use_resn_data_parallel"] = False
+        resn_cfg["use_ft_data_parallel"] = False
+        resn_cfg["use_gnn_data_parallel"] = False
+
+        resn_cfg["output_dir"] = "./ResultsResNFromFTUnsupervised"
+        resn_cfg["optuna_storage"] = "./ResultsResNFromFTUnsupervised/optuna/bayesopt.sqlite3"
+        resn_cfg["optuna_study_prefix"] = "pricing_ft_unsup_resn_ddp"
+        resn_cfg["loss_name"] = "mse"
+
+        runner_cfg = resn_cfg.get("runner", {})
+        runner_cfg["model_keys"] = ["resn"]
+        runner_cfg["nproc_per_node"] = 2
+        runner_cfg["plot_curves"] = False
+        resn_cfg["runner"] = runner_cfg
+
+        resn_cfg["plot_curves"] = False
+        plot_cfg = resn_cfg.get("plot", {})
+        plot_cfg["enable"] = False
+        resn_cfg["plot"] = plot_cfg
+
+        return resn_cfg
+
+    def save_configs(self, output_dir: str = ".") -> Dict[str, str]:
+        """
+        Save generated configs to files.
+
+        Args:
+            output_dir: Directory to save config files
+
+        Returns:
+            Dictionary mapping model names to saved file paths
+        """
+        output_path = Path(output_dir)
+        output_path.mkdir(parents=True, exist_ok=True)
+
+        saved_files = {}
+
+        if self.step1_config:
+            step1_path = output_path / "config_ft_step1_unsupervised.json"
+            with open(step1_path, 'w', encoding='utf-8') as f:
+                json.dump(self.step1_config, f, indent=2)
+            saved_files['ft_step1'] = str(step1_path)
+
+        if 'xgb' in self.step2_configs:
+            xgb_path = output_path / "config_xgb_from_ft_step2.json"
+            with open(xgb_path, 'w', encoding='utf-8') as f:
+                json.dump(self.step2_configs['xgb'], f, indent=2)
+            saved_files['xgb_step2'] = str(xgb_path)
+
+        if 'resn' in self.step2_configs:
+            resn_path = output_path / "config_resn_from_ft_step2.json"
+            with open(resn_path, 'w', encoding='utf-8') as f:
+                json.dump(self.step2_configs['resn'], f, indent=2)
+            saved_files['resn_step2'] = str(resn_path)
+
+        return saved_files

ins_pricing/frontend/runner.py

@@ -0,0 +1,388 @@
+"""
+Unified Task Runner with Real-time Logging
+Executes model training, explanation, plotting, and other tasks based on config.
+"""
+
+import sys
+import threading
+import queue
+import time
+import json
+import subprocess
+from pathlib import Path
+from typing import Generator, Optional, Dict, Any, List, Sequence, Tuple
+import logging
+
+
+class LogCapture:
+    """Capture stdout and stderr for real-time display."""
+
+    def __init__(self):
+        self.log_queue = queue.Queue()
+        self.stop_flag = threading.Event()
+
+    def write(self, text: str):
+        """Write method for capturing output."""
+        if text and text.strip():
+            self.log_queue.put(text)
+
+    def flush(self):
+        """Flush method (required for file-like objects)."""
+        pass
+
+
+class TaskRunner:
+    """
+    Run model tasks (training, explain, plotting, etc.) and capture logs.
+
+    Supports all task modes defined in config.runner.mode:
+    - entry: Standard model training
+    - explain: Model explanation (permutation, SHAP, etc.)
+    - incremental: Incremental training
+    - watchdog: Watchdog mode for monitoring
+    """
+
+    def __init__(self):
+        self.task_thread = None
+        self.log_capture = None
+
+    def _detect_task_mode(self, config_path: str) -> str:
+        """
+        Detect the task mode from config file.
+
+        Args:
+            config_path: Path to configuration JSON file
+
+        Returns:
+            Task mode string (e.g., 'entry', 'explain', 'incremental', 'watchdog')
+        """
+        try:
+            with open(config_path, 'r', encoding='utf-8') as f:
+                config = json.load(f)
+
+            runner_config = config.get('runner', {})
+            mode = runner_config.get('mode', 'entry')
+            return str(mode).lower()
+
+        except Exception as e:
+            print(f"Warning: Could not detect task mode, defaulting to 'entry': {e}")
+            return 'entry'
+
+    def _build_cmd_from_config(self, config_path: str) -> Tuple[List[str], str]:
+        """
+        Build the command to execute based on config.runner.mode, mirroring
+        ins_pricing.cli.utils.notebook_utils.run_from_config behavior.
+        """
+        from ins_pricing.cli.utils.cli_config import set_env
+        from ins_pricing.cli.utils.notebook_utils import (
+            build_bayesopt_entry_cmd,
+            build_incremental_cmd,
+            build_explain_cmd,
+            wrap_with_watchdog,
+        )
+
+        cfg_path = Path(config_path).resolve()
+        raw = json.loads(cfg_path.read_text(encoding="utf-8", errors="replace"))
+        set_env(raw.get("env", {}))
+        runner = dict(raw.get("runner") or {})
+
+        mode = str(runner.get("mode") or "entry").strip().lower()
+        use_watchdog = bool(runner.get("use_watchdog", False))
+        if mode == "watchdog":
+            use_watchdog = True
+            mode = "entry"
+
+        idle_seconds = int(runner.get("idle_seconds", 7200))
+        max_restarts = int(runner.get("max_restarts", 50))
+        restart_delay_seconds = int(runner.get("restart_delay_seconds", 10))
+
+        if mode == "incremental":
+            inc_args = runner.get("incremental_args") or []
+            if not isinstance(inc_args, list):
+                raise ValueError("config.runner.incremental_args must be a list of strings.")
+            cmd = build_incremental_cmd(cfg_path, extra_args=[str(x) for x in inc_args])
+            if use_watchdog:
+                cmd = wrap_with_watchdog(
+                    cmd,
+                    idle_seconds=idle_seconds,
+                    max_restarts=max_restarts,
+                    restart_delay_seconds=restart_delay_seconds,
+                )
+            return cmd, "incremental"
+
+        if mode == "explain":
+            exp_args = runner.get("explain_args") or []
+            if not isinstance(exp_args, list):
+                raise ValueError("config.runner.explain_args must be a list of strings.")
+            cmd = build_explain_cmd(cfg_path, extra_args=[str(x) for x in exp_args])
+            if use_watchdog:
+                cmd = wrap_with_watchdog(
+                    cmd,
+                    idle_seconds=idle_seconds,
+                    max_restarts=max_restarts,
+                    restart_delay_seconds=restart_delay_seconds,
+                )
+            return cmd, "explain"
+
+        if mode != "entry":
+            raise ValueError(
+                f"Unsupported runner.mode={mode!r}, expected 'entry', 'incremental', or 'explain'."
+            )
+
+        model_keys = runner.get("model_keys") or raw.get("model_keys") or ["ft"]
+        if not isinstance(model_keys, list):
+            raise ValueError("runner.model_keys must be a list of strings.")
+        nproc_per_node = int(runner.get("nproc_per_node", 1))
+        max_evals = int(runner.get("max_evals", raw.get("max_evals", 50)))
+        plot_curves = bool(runner.get("plot_curves", raw.get("plot_curves", True)))
+        ft_role = runner.get("ft_role", raw.get("ft_role"))
+
+        extra_args: List[str] = ["--max-evals", str(max_evals)]
+        if plot_curves:
+            extra_args.append("--plot-curves")
+        if ft_role:
+            extra_args += ["--ft-role", str(ft_role)]
+
+        cmd = build_bayesopt_entry_cmd(
+            cfg_path,
+            model_keys=[str(x) for x in model_keys],
+            nproc_per_node=nproc_per_node,
+            extra_args=extra_args,
+        )
+        if use_watchdog:
+            cmd = wrap_with_watchdog(
+                cmd,
+                idle_seconds=idle_seconds,
+                max_restarts=max_restarts,
+                restart_delay_seconds=restart_delay_seconds,
+            )
+        return cmd, "entry"
+
+    def run_task(self, config_path: str) -> Generator[str, None, None]:
+        """
+        Run task based on config file with real-time log capture.
+
+        This method automatically detects the task mode from the config file
+        (training, explain, plotting, etc.) and runs the appropriate task.
+
+        Args:
+            config_path: Path to configuration JSON file
+
+        Yields:
+            Log lines as they are generated
+        """
+        self.log_capture = LogCapture()
+
+        # Configure logging to capture both file and stream output
+        log_handler = logging.StreamHandler(self.log_capture)
+        log_handler.setLevel(logging.INFO)
+        formatter = logging.Formatter('%(asctime)s - %(levelname)s - %(message)s')
+        log_handler.setFormatter(formatter)
+
+        # Add handler to root logger
+        root_logger = logging.getLogger()
+        original_handlers = root_logger.handlers.copy()
+        root_logger.addHandler(log_handler)
+
+        # Store original stdout/stderr
+        original_stdout = sys.stdout
+        original_stderr = sys.stderr
+
+        try:
+            # Detect task mode
+            task_mode = self._detect_task_mode(config_path)
+
+            # Start task in separate thread
+            exception_holder = []
+
+            def task_worker():
+                try:
+                    sys.stdout = self.log_capture
+                    sys.stderr = self.log_capture
+
+                    # Log start
+                    cmd, task_mode = self._build_cmd_from_config(config_path)
+                    print(f"Starting task [{task_mode}] with config: {config_path}")
+                    print("=" * 80)
+
+                    # Run subprocess with streamed output
+                    proc = subprocess.Popen(
+                        cmd,
+                        stdout=subprocess.PIPE,
+                        stderr=subprocess.STDOUT,
+                        text=True,
+                        bufsize=1,
+                        cwd=str(Path(config_path).resolve().parent),
+                    )
+                    if proc.stdout is not None:
+                        for line in proc.stdout:
+                            print(line.rstrip())
+                    return_code = proc.wait()
+                    if return_code != 0:
+                        raise RuntimeError(f"Task exited with code {return_code}")
+
+                    print("=" * 80)
+                    print(f"Task [{task_mode}] completed successfully!")
+
+                except Exception as e:
+                    exception_holder.append(e)
+                    print(f"Error during task execution: {str(e)}")
+                    import traceback
+                    print(traceback.format_exc())
+
+                finally:
+                    sys.stdout = original_stdout
+                    sys.stderr = original_stderr
+
+            self.task_thread = threading.Thread(target=task_worker, daemon=True)
+            self.task_thread.start()
+
+            # Yield log lines as they come in
+            last_update = time.time()
+            while self.task_thread.is_alive() or not self.log_capture.log_queue.empty():
+                try:
+                    # Try to get log with timeout
+                    log_line = self.log_capture.log_queue.get(timeout=0.1)
+                    yield log_line
+                    last_update = time.time()
+
+                except queue.Empty:
+                    # Send heartbeat every 5 seconds
+                    if time.time() - last_update > 5:
+                        yield "."
+                        last_update = time.time()
+                    continue
+
+            # Wait for thread to complete
+            self.task_thread.join(timeout=1)
+
+            # Check for exceptions
+            if exception_holder:
+                raise exception_holder[0]
+
+        finally:
+            # Restore original stdout/stderr
+            sys.stdout = original_stdout
+            sys.stderr = original_stderr
+
+            # Restore original logging handlers
+            root_logger.handlers = original_handlers
+
+    def run_callable(self, func, *args, **kwargs) -> Generator[str, None, None]:
+        """Run an in-process callable and stream stdout/stderr."""
+        self.log_capture = LogCapture()
+
+        log_handler = logging.StreamHandler(self.log_capture)
+        log_handler.setLevel(logging.INFO)
+        formatter = logging.Formatter('%(asctime)s - %(levelname)s - %(message)s')
+        log_handler.setFormatter(formatter)
+
+        root_logger = logging.getLogger()
+        original_handlers = root_logger.handlers.copy()
+        root_logger.addHandler(log_handler)
+
+        original_stdout = sys.stdout
+        original_stderr = sys.stderr
+
+        try:
+            exception_holder = []
+
+            def task_worker():
+                try:
+                    sys.stdout = self.log_capture
+                    sys.stderr = self.log_capture
+                    func(*args, **kwargs)
+                except Exception as e:
+                    exception_holder.append(e)
+                    print(f"Error during task execution: {str(e)}")
+                    import traceback
+                    print(traceback.format_exc())
+                finally:
+                    sys.stdout = original_stdout
+                    sys.stderr = original_stderr
+
+            self.task_thread = threading.Thread(target=task_worker, daemon=True)
+            self.task_thread.start()
+
+            last_update = time.time()
+            while self.task_thread.is_alive() or not self.log_capture.log_queue.empty():
+                try:
+                    log_line = self.log_capture.log_queue.get(timeout=0.1)
+                    yield log_line
+                    last_update = time.time()
+                except queue.Empty:
+                    if time.time() - last_update > 5:
+                        yield "."
+                        last_update = time.time()
+                    continue
+
+            self.task_thread.join(timeout=1)
+            if exception_holder:
+                raise exception_holder[0]
+        finally:
+            sys.stdout = original_stdout
+            sys.stderr = original_stderr
+            root_logger.handlers = original_handlers
+
+    def stop_task(self):
+        """Stop the current task process."""
+        if self.log_capture:
+            self.log_capture.stop_flag.set()
+
+        if self.task_thread and self.task_thread.is_alive():
+            # Note: Thread.join() will wait for completion
+            # For forceful termination, you may need to use process-based approach
+            self.task_thread.join(timeout=5)
+
+
+# Backward compatibility aliases
+TrainingRunner = TaskRunner
+
+
+class StreamToLogger:
+    """
+    Fake file-like stream object that redirects writes to a logger instance.
+    """
+
+    def __init__(self, logger, log_level=logging.INFO):
+        self.logger = logger
+        self.log_level = log_level
+        self.linebuf = ''
+
+    def write(self, buf):
+        for line in buf.rstrip().splitlines():
+            self.logger.log(self.log_level, line.rstrip())
+
+    def flush(self):
+        pass
+
+
+def setup_logger(name: str = "task") -> logging.Logger:
+    """
+    Set up a logger for task execution.
+
+    Args:
+        name: Logger name
+
+    Returns:
+        Configured logger instance
+    """
+    logger = logging.getLogger(name)
+    logger.setLevel(logging.INFO)
+
+    # Create console handler
+    console_handler = logging.StreamHandler()
+    console_handler.setLevel(logging.INFO)
+
+    # Create formatter
+    formatter = logging.Formatter(
+        '%(asctime)s - %(name)s - %(levelname)s - %(message)s',
+        datefmt='%Y-%m-%d %H:%M:%S'
+    )
+    console_handler.setFormatter(formatter)
+
+    # Add handler to logger
+    if not logger.handlers:
+        logger.addHandler(console_handler)
+
+    return logger