run.py

# -*- coding: utf-8 -*-
# pyinstaller --onefile run.py
import os
import argparse
import builtins
import random
import numpy as np
import torch
import torch.distributed as dist
from datetime import datetime

from exp.exp_forecast_pvpf import Exp_Forecast_PVPF
from exp.exp_chronos2 import Exp_Chronos2
from exp.exp_chronos2_rag import Exp_Chronos2_RAG
from exp.exp_chronos2_rag_fut import Exp_Chronos2_RAG_fut
from exp.exp_chronos2_rag_fut_spec import Exp_Chronos2_RAG_fut as Exp_Chronos2_RAG_fut_Spec


def _enable_timestamped_print():
    original_print = builtins.print

    def timestamped_print(*args, **kwargs):
        prefix = datetime.now().strftime("[%Y-%m-%d %H:%M:%S]")
        if args:
            original_print(prefix, *args, **kwargs)
        else:
            original_print(prefix, **kwargs)

    builtins.print = timestamped_print

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='PVPFLab')

    # basic config
    parser.add_argument('--task_name', type=str, required=True, default='forecast',
                        help='task name, options:[forecast|evf|pipeline]')
    parser.add_argument('--is_training', type=int, required=True, default=1, help='status')
    parser.add_argument('--model_id', type=str, required=True, default='test', help='model id')
    parser.add_argument('--model', type=str, required=True, default='timer', help='model name, options: [timer]')
    parser.add_argument('--seed', type=int, default=2021, help='seed')

    # data loader
    parser.add_argument('--data', type=str, required=True, default='ETTh1', help='dataset type')
    parser.add_argument('--train_data', type=str, default=None, help='dataset type')
    parser.add_argument('--test_data', type=str, default=None, help='dataset type')
    parser.add_argument('--root_path', type=str, default='./dataset/ETT-small/', help='root path of the data file')
    parser.add_argument('--data_path', type=str, default='ETTh1.csv', help='data file')
    parser.add_argument('--checkpoints', type=str, default='./checkpoints/', help='location of model checkpoints')
    parser.add_argument('--data_list_path', type=str, default='./dataset_list.json', help='data list path')
    parser.add_argument('--test_data_list_path', type=str, default=None, help='data list path')
    parser.add_argument('--resume_dir', type=str, default='', help='resume dir')
    parser.add_argument('--scale', action='store_true', help='scale data')
    parser.add_argument('--max_files', type=int, default=None, help='max number of CSV files to load (for MultiXiexinPVPFCV), None means load all')

    # forecasting task
    parser.add_argument('--seq_len', type=int, default=672, help='input sequence length')
    parser.add_argument('--input_token_len', type=int, default=576, help='input token length')
    parser.add_argument('--max_output_token_len', type=int, default=96, help='max output token length')
    parser.add_argument('--future_cov_observed', type=int, default=96, help='number of observed points in future covariate window')

    # test
    parser.add_argument('--test_seq_len', type=int, default=672, help='test seq len')
    parser.add_argument('--test_pred_len', type=int, default=96, help='test pred len')
    parser.add_argument('--test_dir', type=str, default='./test', help='test dir')
    parser.add_argument('--test_with_revin', action='store_true', help='test with revin', default=False)
    parser.add_argument('--test_n_sample', type=int, default=500, help='test n sample')

    # finetune
    parser.add_argument('--adaptation', type=str, help='adaptation method description', default=None)
    parser.add_argument('--pretrain_model_path', type=str, default='pretrain_model.pth', help='pretrain model path')

    # optimization
    parser.add_argument('--num_workers', type=int, default=32, help='data loader num workers')
    parser.add_argument('--train_epochs', type=int, default=10, help='train epochs')
    parser.add_argument('--batch_size', type=int, default=32, help='batch size of train input data')
    parser.add_argument('--val_batch_size', type=int, default=32, help='batch size of val input data')
    parser.add_argument('--patience', type=int, default=3, help='early stopping patience')
    parser.add_argument('--learning_rate', type=float, default=0.0001, help='optimizer learning rate')
    parser.add_argument('--des', type=str, default='test', help='exp description')
    parser.add_argument('--weight_decay', type=float, default=0)
    parser.add_argument('--warmup', action='store_true', help='use warmup', default=False)
    parser.add_argument('--num_warmup_steps', type=int, default=10000, help='num warmup steps')
    parser.add_argument('--num_training_steps', type=int, default=100000, help='num training steps')
    parser.add_argument('--iter_per_epoch', type=int, default=5000, help='iter per epoch')
    parser.add_argument('--train_augment_margin', type=int, default=0, help='augment timestamp margin for rollout dataloaders')
    parser.add_argument('--target_hour', type=int, default=0, help='target hour for price forecasting')
    parser.add_argument('--nofilt', action='store_true', help='no filter data samples by target hour', default=False)
    parser.add_argument('--covariate_shift', action='store_true', help='use covariate shift for multivariate rollout', default=False)

    # GPU
    parser.add_argument('--gpu', type=int, default=0, help='gpu')
    parser.add_argument('--ddp', action='store_true', help='Distributed Data Parallel', default=False)
    parser.add_argument('--devices', type=str, default='0,1,2,3', help='device ids of multile gpus')

    # adaptation
    parser.add_argument("--adapter", type=str, default="adapter_linear", help="adapter name")
    parser.add_argument("--input_channel", type=int, default=1, help="input channel")
    parser.add_argument("--output_channel", type=int, default=1, help="output channel")
    parser.add_argument("--fit_batch_size", type=int, default=128)
    parser.add_argument("--fit_learning_rate", type=float, default=1e-3)
    parser.add_argument("--use_amp", action="store_true", help="enable mixed precision training")
    parser.add_argument("--accum_steps", type=int, default=32)
    parser.add_argument("--finetune_head_path", type=str, default='')
    parser.add_argument("--finetune_ltm", type=str, default='')
    parser.add_argument("--zero_init", action="store_true", help="use zero initialization DW", default=False)
    parser.add_argument("--finetune_mode", type=str, default="full", help="Chronos2 finetune mode, options: [full|lora]")
    parser.add_argument("--chronos2_finetune_input_mode", type=int, default=3, choices=[1, 2, 3],
        help="Chronos2 finetune input mode: 1=target only, 2=target+past_cov, 3=target+past_cov+future_cov")

    # LoRA 配置参数
    parser.add_argument("--lora_r", type=int, default=8, help="LoRA rank (r)")
    parser.add_argument("--lora_alpha", type=int, default=16, help="LoRA alpha scaling factor")
    parser.add_argument("--lora_dropout", type=float, default=0.0, help="LoRA dropout rate")
    parser.add_argument("--lora_target_modules", type=str, default=None,
                        help="LoRA target modules (comma-separated), e.g., 'self_attention.q,self_attention.v'")

    # ===== pipeline switches (cov-only v1; tsfm later) =====
    parser.add_argument('--enable_tsfm', action='store_true', default=False,help='enable TSFM rollout predictions in pipeline')
    parser.add_argument('--tsfm_model_name', type=str, default='sundial',help='TSFM model name for rollout (pipeline)')
    parser.add_argument('--save_dir', type=str, default='./artifacts/pipeline_min', help='pipeline artifacts dir')

    # ===== LGBM params (pipeline) =====
    parser.add_argument('--num_boost_round', type=int, default=5000)
    parser.add_argument('--early_stopping_rounds', type=int, default=200)
    parser.add_argument('--lgbm_learning_rate', type=float, default=0.05)
    parser.add_argument('--num_leaves', type=int, default=63)
    parser.add_argument('--feature_fraction', type=float, default=0.9)
    parser.add_argument('--bagging_fraction', type=float, default=0.8)
    parser.add_argument('--min_gain_to_split', type=float, default=0.1)
    parser.add_argument("--max_depth", type=int, default=-1)
    parser.add_argument("--min_data_in_leaf", type=int, default=20)
    parser.add_argument("--min_sum_hessian_in_leaf", type=float, default=1e-3)
    parser.add_argument('--bagging_freq', type=int, default=5)
    parser.add_argument('--reg_alpha', type=float, default=0.0)
    parser.add_argument('--reg_lambda', type=float, default=0.0)
    parser.add_argument('--n_jobs', type=int, default=-1)
    parser.add_argument('--log_period', type=int, default=200)

    # ===== SHAP explainability (pipeline, LightGBM) =====
    parser.add_argument('--disable_shap', action='store_true', default=False, help='disable SHAP explainability step')
    parser.add_argument('--shap_no_sample', action='store_true', default=False, help='disable SHAP subsampling (use all rows)')
    parser.add_argument('--shap_max_samples', type=int, default=5000, help='max rows used to compute SHAP (subsampled)')
    parser.add_argument('--shap_seed', type=int, default=2021, help='seed for SHAP subsampling')
    parser.add_argument('--shap_topk', type=int, default=50, help='max displayed features in SHAP plots')

    # experiment tracing
    parser.add_argument('--swanlab_key', type=str, default=None, help='use swanlab experiment tracing')
    parser.add_argument('--trace_exp_name', type=str, default="default trace experiment name", help='trace experiment name')
    parser.add_argument('--trace_exp_des', type=str, default="default trace experiment description", help='trace experiment description')

    # ===== RAG (zero-shot inference) =====
    parser.add_argument('--rag_k', type=int, default=3, help='top-k retrieval for RAG')
    parser.add_argument('--rag_k_mode', type=str, default='auto', choices=['auto', 'fixed'],
                        help="RAG k selection mode: 'auto' picks best k per-sample; 'fixed' uses --rag_k")
    parser.add_argument('--rag_max_k', type=int, default=12, help='max k candidates considered when rag_k_mode=auto')
    parser.add_argument('--rag_strategy', type=str, default='coarse_wcos_fine_weuc',
                        help='retrieval strategy in RAG/retriever.py')
    parser.add_argument('--rag_max_kb', type=int, default=5000, help='max KB windows built from train set')
    parser.add_argument('--rag_kb_stride', type=int, default=1, help='stride when scanning train windows for KB')
    parser.add_argument('--rag_alignment', action='store_true', default=False, help='align retrieved context values to query mean')
    parser.add_argument('--rag_context_length', type=int, default=0,
                        help='optionally truncate retrieved-only RAG context to this length; 0 keeps the full retrieved context')
    parser.add_argument('--rag_query_keep', type=int, default=0,
                        help='deprecated in future-only RAG: query history is not appended to Chronos input')
    parser.add_argument('--rag_use_metadata', type=int, default=1, choices=[0, 1],
                        help='use device/series metadata to filter RAG candidates in *_spec experiments; default 1 preserves original *_spec behavior')
    parser.add_argument('--rag_meta_ranking', type=int, default=1, choices=[0, 1],
                        help='rank RAG candidates by static metadata before numeric retrieval in *_spec experiments')
    parser.add_argument('--rag_meta_candidate_cap', type=int, default=0,
                        help='metadata-ranked candidate pool size per query in *_spec experiments; <=0 auto uses max(1024, retrieve_k*64)')
    parser.add_argument('--rag_meta_geo_weight', type=float, default=1.0,
                        help='weight for lat/lng distance in metadata candidate ranking')
    parser.add_argument('--rag_meta_capacity_weight', type=float, default=2.0,
                        help='weight for capacity mismatch in metadata candidate ranking')
    parser.add_argument('--rag_meta_component_weight', type=float, default=1.0,
                        help='weight for component_num mismatch in metadata candidate ranking')
    parser.add_argument('--rag_meta_roof_weight', type=float, default=1.0,
                        help='weight for roof_type_name mismatch in metadata candidate ranking')
    parser.add_argument('--rag_meta_power_weight', type=float, default=2.0,
                        help='weight for recent target-power/load mismatch in metadata candidate ranking')
    parser.add_argument('--rag_final_meta_rerank', type=int, default=1, choices=[0, 1],
                        help='rerank over-fetched retrieved items by metadata/history-power similarity in *_spec experiments')
    parser.add_argument('--rag_retrieve_overfetch', type=int, default=4,
                        help='retrieve this multiple of top-k before final metadata reranking in *_spec experiments')
    parser.add_argument('--rag_final_meta_weight', type=float, default=0.1,
                        help='light metadata identity weight in final retrieved-item reranking')
    parser.add_argument('--rag_final_power_weight', type=float, default=20.0,
                        help='weight for recent history-power/load mismatch in final retrieved-item reranking')
    parser.add_argument('--rag_zero_history_guard', type=int, default=1, choices=[0, 1],
                        help='penalize high-power retrieved histories when query recent history is near-zero in *_spec experiments')
    parser.add_argument('--rag_zero_hist_load_mean_threshold', type=float, default=0.01,
                        help='query history load-mean threshold for near-zero guard')
    parser.add_argument('--rag_zero_hist_load_max_threshold', type=float, default=0.05,
                        help='query history load-max threshold for near-zero guard')
    parser.add_argument('--rag_high_hist_load_mean_threshold', type=float, default=0.05,
                        help='retrieved history load-mean threshold treated as high-power under near-zero guard')
    parser.add_argument('--rag_high_hist_load_max_threshold', type=float, default=0.25,
                        help='retrieved history load-max threshold treated as high-power under near-zero guard')
    parser.add_argument('--rag_zero_high_power_penalty', type=float, default=100.0,
                        help='large final rerank penalty for high-power histories under near-zero guard')
    parser.add_argument('--plt_rag', action='store_true', default=False,
                        help='save RAG target-only visualization figures for uniformly sampled test windows')
    parser.add_argument('--rag_log_interval', type=int, default=100,
                        help='log RAG eval progress every N processed samples; <=0 disables progress logs')
    parser.add_argument('--rag_log_first_n', type=int, default=3,
                        help='log detailed per-sample RAG info for the first N samples of each split')
    parser.add_argument('--rag_debug_topn', type=int, default=8,
                        help='log top-N per-sample p3 MSE outliers with target/pred/cov stats; <=0 disables')
    parser.add_argument('--no_rag_log_metadata', action='store_true', default=False,
                        help='disable detailed RAG metadata/variable summary logs')
    parser.add_argument('--no_log_timestamp', action='store_true', default=False,
                        help='disable timestamp prefixes on stdout logs')

    args = parser.parse_args()
    if not getattr(args, "no_log_timestamp", False):
        _enable_timestamped_print()

    # Chronos2 的 model._compute_loss 会读取 os.environ["DUAL_WEAVER"]，未设置会 KeyError
    os.environ.setdefault("DUAL_WEAVER", "0")

    fix_seed = args.seed
    random.seed(fix_seed)
    torch.manual_seed(fix_seed)
    np.random.seed(fix_seed)

    if args.ddp:
        ip = os.environ.get("MASTER_ADDR", "127.0.0.1")
        port = os.environ.get("MASTER_PORT", "64209")
        hosts = int(os.environ.get("WORLD_SIZE", "8"))
        rank = int(os.environ.get("RANK", "0"))
        local_rank = int(os.environ.get("LOCAL_RANK", "0"))
        gpus = torch.cuda.device_count()
        args.local_rank = local_rank
        if args.local_rank == 0:
            print(ip, port, hosts, rank, local_rank, gpus)
        dist.init_process_group(
            backend="nccl",
            init_method=f"tcp://{ip}:{port}",
            world_size=hosts,
            rank=rank,
        )
        torch.cuda.set_device(local_rank)
    else:
        args.local_rank = 0

    # select Exp
    if args.task_name == 'forecast':
        Exp = Exp_Forecast
    elif args.task_name == 'evf':
        Exp = Exp_EVF
    elif args.task_name == 'pipeline':
        Exp = Exp_Pipeline
    elif args.task_name == 'chronos2':
        Exp = Exp_Chronos2
    elif args.task_name == 'chronos2_rag_fut':
        Exp = Exp_Chronos2_RAG_fut
    elif args.task_name == 'chronos2_rag_fut_spec':
        Exp = Exp_Chronos2_RAG_fut_Spec
    elif args.task_name in ('chronos2_rag', 'chronos2_rag_hist_fut'):
        Exp = Exp_Chronos2_RAG
    elif args.task_name == 'forecast_pvpf':
        Exp = Exp_Forecast_PVPF
    else:
        Exp = Exp_Forecast

    data_path = args.data_list_path.split('/')[-1].split('.')[0]

    if args.is_training:
        exp = Exp(args)
        setting = '{}_{}_sl{}_it{}_ot{}_lr{}_bt{}_wd{}_{}'.format(
            data_path,
            args.model_id,
            args.model,
            args.seq_len,
            args.input_token_len,
            args.max_output_token_len,
            args.learning_rate,
            args.batch_size,
            args.weight_decay
        )
        setting = datetime.now().strftime("%y-%m-%d_%H-%M-%S") + setting

        if args.local_rank == 0:
            print('>>>>>>>start training : {}>>>>>>>>>>>>>>>>>>>>>>>>>>'.format(setting))
        exp.train(setting)
        torch.cuda.empty_cache()
        if args.ddp:
            dist.destroy_process_group()
    else:
        setting = '{}_{}_sl{}_it{}_ot{}_lr{}_bt{}_wd{}_{}'.format(
            data_path,
            args.model_id,
            args.model,
            args.seq_len,
            args.input_token_len,
            args.max_output_token_len,
            args.learning_rate,
            args.batch_size,
            args.weight_decay
        )
        setting = datetime.now().strftime("%y-%m-%d_%H-%M-%S") + setting
        exp = Exp(args)
        exp.test(setting, test=1)
        torch.cuda.empty_cache()
        if args.ddp:
            dist.destroy_process_group()