Qionk/a-share-quant
0
1"""快速训练 worker — 独立进程,不受 Streamlit 干扰"""2import os, sys3 4os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3"5os.environ["CUDA_VISIBLE_DEVICES"] = ""6 7# 必须在 import akshare/pandas 之前先 import TF,否则线程池冲突导致死锁8import tensorflow # noqa: F4019 10import pickle11 12if __name__ == "__main__":13 input_path = sys.argv[1]14 output_path = sys.argv[2]15 16 sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))17 18 with open(input_path, "rb") as f:19 args = pickle.load(f)20 21 from src.predict.training import train_all_models22 23 results = train_all_models(24 args["df"], args["selected_models"], args["config"],25 forecast_days=args["forecast_days"], quick=True)26 27 for r in results.values():28 r.model_object = None29 r.scaler = None30 31 with open(output_path, "wb") as f:32 pickle.dump(results, f)33 