mirror of
https://github.com/microsoft/qlib.git
synced 2026-07-19 02:14:33 +08:00
Compare commits
11 Commits
release-pl
...
8355990ac5
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
8355990ac5 | ||
|
|
16acb76aba | ||
|
|
4e0f5d5ec9 | ||
|
|
50c32ac15f | ||
|
|
80982f8904 | ||
|
|
477160e4ac | ||
|
|
3472e82d5c | ||
|
|
cb285bccac | ||
|
|
2e9a00a9f7 | ||
|
|
d631b4450b | ||
|
|
0826879481 |
1
.gitignore
vendored
1
.gitignore
vendored
@@ -22,6 +22,7 @@ dist/
|
||||
qlib/VERSION.txt
|
||||
qlib/data/_libs/expanding.cpp
|
||||
qlib/data/_libs/rolling.cpp
|
||||
qlib/_version.py
|
||||
examples/estimator/estimator_example/
|
||||
examples/rl/data/
|
||||
examples/rl/checkpoints/
|
||||
|
||||
14
CHANGELOG.md
14
CHANGELOG.md
@@ -1,14 +0,0 @@
|
||||
# Changelog
|
||||
|
||||
## [0.9.8](https://github.com/microsoft/qlib/compare/v0.9.7...v0.9.8) (2025-11-13)
|
||||
|
||||
|
||||
### Bug Fixes
|
||||
|
||||
* download orderbook data error ([#1990](https://github.com/microsoft/qlib/issues/1990)) ([136b2dd](https://github.com/microsoft/qlib/commit/136b2ddf9a16e4106d62b8d1336a56273a8abef0))
|
||||
* **gbdt:** correct dtrain assignment in finetune() to use Dataset instead of tuple ([#2049](https://github.com/microsoft/qlib/issues/2049)) ([2b41782](https://github.com/microsoft/qlib/commit/2b41782f0cfb81e8cc065f2915b215758a7838ef))
|
||||
* **macd:** remove extra division by close in DEA calculation to ensure dimension consistency ([#2046](https://github.com/microsoft/qlib/issues/2046)) ([66c3622](https://github.com/microsoft/qlib/commit/66c36226aafceabe497e5967f67921e5d3c9d497))
|
||||
* replace deprecated pandas fillna(method=) with ffill()/bfill() ([#1987](https://github.com/microsoft/qlib/issues/1987)) ([7095e75](https://github.com/microsoft/qlib/commit/7095e755fa57e011f0483d24b45fc5bd5a4deaf8))
|
||||
* spelling errors ([#1996](https://github.com/microsoft/qlib/issues/1996)) ([f26b341](https://github.com/microsoft/qlib/commit/f26b3417363410531dbbb39e425bce6cf05528a1))
|
||||
* the bug when auto_mount=True ([#2009](https://github.com/microsoft/qlib/issues/2009)) ([213eb6c](https://github.com/microsoft/qlib/commit/213eb6c2cd12342b6ec98f21300217e1659f3d58))
|
||||
* typo in integration documentation: 'userd' -> 'used' ([#2034](https://github.com/microsoft/qlib/issues/2034)) ([3dc5a7d](https://github.com/microsoft/qlib/commit/3dc5a7d299074f0fa45a4b7bb50ab446a8824a32))
|
||||
|
||||
23
Makefile
23
Makefile
@@ -74,34 +74,37 @@ prerequisite:
|
||||
|
||||
# Install the package in editable mode.
|
||||
dependencies:
|
||||
python -m pip install -e .
|
||||
python -m pip install --no-cache-dir -e .
|
||||
|
||||
lightgbm:
|
||||
python -m pip install lightgbm --prefer-binary
|
||||
python -m pip install --no-cache-dir lightgbm --prefer-binary
|
||||
|
||||
rl:
|
||||
python -m pip install -e .[rl]
|
||||
python -m pip install --no-cache-dir -e .[rl]
|
||||
|
||||
develop:
|
||||
python -m pip install -e .[dev]
|
||||
python -m pip install --no-cache-dir -e .[dev]
|
||||
|
||||
lint:
|
||||
python -m pip install -e .[lint]
|
||||
python -m pip install --no-cache-dir -e .[lint]
|
||||
|
||||
docs:
|
||||
python -m pip install -e .[docs]
|
||||
python -m pip install --no-cache-dir -e .[docs]
|
||||
|
||||
package:
|
||||
python -m pip install -e .[package]
|
||||
python -m pip install --no-cache-dir -e .[package]
|
||||
|
||||
test:
|
||||
python -m pip install -e .[test]
|
||||
python -m pip install --no-cache-dir -e .[test]
|
||||
|
||||
analysis:
|
||||
python -m pip install -e .[analysis]
|
||||
python -m pip install --no-cache-dir -e .[analysis]
|
||||
|
||||
client:
|
||||
python -m pip install --no-cache-dir -e .[client]
|
||||
|
||||
all:
|
||||
python -m pip install -e .[pywinpty,dev,lint,docs,package,test,analysis,rl]
|
||||
python -m pip install --no-cache-dir -e .[pywinpty,dev,lint,docs,package,test,analysis,rl]
|
||||
|
||||
install: prerequisite dependencies
|
||||
|
||||
|
||||
@@ -69,8 +69,10 @@ rl = [
|
||||
"torch",
|
||||
"numpy<2.0.0",
|
||||
]
|
||||
# We exclude black version 26.1.0 due to known issues with nbqa when formatting Jupyter notebooks,
|
||||
# which can cause false-positive --check results and inconsistent notebook formatting.
|
||||
lint = [
|
||||
"black",
|
||||
"black!=26.1.0",
|
||||
"pylint",
|
||||
"mypy<1.5.0",
|
||||
"flake8",
|
||||
@@ -101,6 +103,10 @@ analysis = [
|
||||
"plotly",
|
||||
"statsmodels",
|
||||
]
|
||||
client = [
|
||||
"python-socketio<6",
|
||||
"tables",
|
||||
]
|
||||
|
||||
# In the process of releasing a new version, when checking the manylinux package with twine, an error is reported:
|
||||
# InvalidDistribution: Invalid distribution metadata: unrecognized or malformed field 'license-file'
|
||||
|
||||
@@ -87,7 +87,7 @@ def workflow(config_path, experiment_name="workflow", uri_folder="mlruns"):
|
||||
"""
|
||||
This is a Qlib CLI entrance.
|
||||
User can run the whole Quant research workflow defined by a configure file
|
||||
- the code is located here ``qlib/cli/run.py`
|
||||
- the code is located here ``qlib/cli/run.py``
|
||||
|
||||
User can specify a base_config file in your workflow.yml file by adding "BASE_CONFIG_PATH".
|
||||
Qlib will load the configuration in BASE_CONFIG_PATH first, and the user only needs to update the custom fields
|
||||
|
||||
@@ -10,6 +10,7 @@ import os
|
||||
import gc
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
from packaging import version
|
||||
from typing import Callable, Optional, Text, Union
|
||||
from sklearn.metrics import roc_auc_score, mean_squared_error
|
||||
|
||||
@@ -148,7 +149,7 @@ class DNNModelPytorch(Model):
|
||||
if scheduler == "default":
|
||||
# In torch version 2.7.0, the verbose parameter has been removed. Reference Link:
|
||||
# https://github.com/pytorch/pytorch/pull/147301/files#diff-036a7470d5307f13c9a6a51c3a65dd014f00ca02f476c545488cd856bea9bcf2L1313
|
||||
if str(torch.__version__).split("+", maxsplit=1)[0] <= "2.6.0":
|
||||
if version.parse(str(torch.__version__).split("+", maxsplit=1)[0]) <= version.parse("2.6.0"):
|
||||
# Reduce learning rate when loss has stopped decrease
|
||||
self.scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau( # pylint: disable=E1123
|
||||
self.train_optimizer,
|
||||
|
||||
@@ -14,6 +14,7 @@ from qlib.model.meta.task import MetaTask
|
||||
from qlib.model.trainer import TrainerR
|
||||
from qlib.typehint import Literal
|
||||
from qlib.utils import init_instance_by_config
|
||||
from qlib.utils.pickle_utils import restricted_pickle_load
|
||||
from qlib.workflow import R
|
||||
from qlib.workflow.task.utils import replace_task_handler_with_cache
|
||||
|
||||
@@ -298,7 +299,7 @@ class DDGDA(Rolling):
|
||||
# but their task test segment are not aligned! It worked in my previous experiment.
|
||||
# So the misalignment will not affect the effectiveness of the method.
|
||||
with self._internal_data_path.open("rb") as f:
|
||||
internal_data = pickle.load(f)
|
||||
internal_data = restricted_pickle_load(f)
|
||||
|
||||
md = MetaDatasetDS(exp_name=internal_data, **kwargs)
|
||||
|
||||
@@ -360,7 +361,7 @@ class DDGDA(Rolling):
|
||||
)
|
||||
|
||||
with self._internal_data_path.open("rb") as f:
|
||||
internal_data = pickle.load(f)
|
||||
internal_data = restricted_pickle_load(f)
|
||||
mds = MetaDatasetDS(exp_name=internal_data, **kwargs)
|
||||
|
||||
# 3) meta model make inference and get new qlib task
|
||||
|
||||
@@ -8,7 +8,6 @@ import os
|
||||
import yaml
|
||||
import json
|
||||
import copy
|
||||
import pickle
|
||||
import logging
|
||||
import importlib
|
||||
import subprocess
|
||||
@@ -18,6 +17,7 @@ import numpy as np
|
||||
from abc import abstractmethod
|
||||
|
||||
from ...log import get_module_logger, TimeInspector
|
||||
from ...utils.pickle_utils import restricted_pickle_load
|
||||
from hyperopt import fmin, tpe
|
||||
from hyperopt import STATUS_OK, STATUS_FAIL
|
||||
|
||||
@@ -136,7 +136,7 @@ class QLibTuner(Tuner):
|
||||
exp_result_dir = os.path.join(self.ex_dir, QLibTuner.EXP_RESULT_DIR.format(estimator_ex_id))
|
||||
exp_result_path = os.path.join(exp_result_dir, QLibTuner.EXP_RESULT_NAME)
|
||||
with open(exp_result_path, "rb") as fp:
|
||||
analysis_df = pickle.load(fp)
|
||||
analysis_df = restricted_pickle_load(fp)
|
||||
|
||||
# 4. Get the backtest factor which user want to optimize, if user want to maximize the factor, then reverse the result
|
||||
res = analysis_df.loc[self.optim_config.report_type].loc[self.optim_config.report_factor]
|
||||
|
||||
@@ -30,6 +30,7 @@ from ..utils import (
|
||||
normalize_cache_fields,
|
||||
normalize_cache_instruments,
|
||||
)
|
||||
from ..utils.pickle_utils import restricted_pickle_load
|
||||
|
||||
from ..log import get_module_logger
|
||||
from .base import Feature
|
||||
@@ -225,7 +226,7 @@ class CacheUtils:
|
||||
cache_path = Path(cache_path)
|
||||
meta_path = cache_path.with_suffix(".meta")
|
||||
with meta_path.open("rb") as f:
|
||||
d = pickle.load(f)
|
||||
d = restricted_pickle_load(f)
|
||||
with meta_path.open("wb") as f:
|
||||
try:
|
||||
d["meta"]["last_visit"] = str(time.time())
|
||||
@@ -592,7 +593,7 @@ class DiskExpressionCache(ExpressionCache):
|
||||
|
||||
with CacheUtils.writer_lock(self.r, f"{str(C.dpm.get_data_uri())}:expression-{cache_uri}"):
|
||||
with meta_path.open("rb") as f:
|
||||
d = pickle.load(f)
|
||||
d = restricted_pickle_load(f)
|
||||
instrument = d["info"]["instrument"]
|
||||
field = d["info"]["field"]
|
||||
freq = d["info"]["freq"]
|
||||
@@ -959,7 +960,7 @@ class DiskDatasetCache(DatasetCache):
|
||||
im = DiskDatasetCache.IndexManager(cp_cache_uri)
|
||||
with CacheUtils.writer_lock(self.r, f"{str(C.dpm.get_data_uri())}:dataset-{cache_uri}"):
|
||||
with meta_path.open("rb") as f:
|
||||
d = pickle.load(f)
|
||||
d = restricted_pickle_load(f)
|
||||
instruments = d["info"]["instruments"]
|
||||
fields = d["info"]["fields"]
|
||||
freq = d["info"]["freq"]
|
||||
|
||||
@@ -2,15 +2,15 @@
|
||||
# Licensed under the MIT License.
|
||||
|
||||
|
||||
from __future__ import division
|
||||
from __future__ import print_function
|
||||
from __future__ import division, print_function
|
||||
|
||||
import json
|
||||
|
||||
import socketio
|
||||
|
||||
import qlib
|
||||
from ..config import C
|
||||
|
||||
from ..log import get_module_logger
|
||||
import pickle
|
||||
|
||||
|
||||
class Client:
|
||||
@@ -96,7 +96,7 @@ class Client:
|
||||
self.logger.debug("connected")
|
||||
# The pickle is for passing some parameters with special type(such as
|
||||
# pd.Timestamp)
|
||||
request_content = {"head": head_info, "body": pickle.dumps(request_content, protocol=C.dump_protocol_version)}
|
||||
request_content = {"head": head_info, "body": json.dumps(request_content, default=str)}
|
||||
self.sio.on(request_type + "_response", request_callback)
|
||||
self.logger.debug("try sending")
|
||||
self.sio.emit(request_type + "_request", request_content)
|
||||
|
||||
@@ -2,7 +2,6 @@
|
||||
# Licensed under the MIT License.
|
||||
|
||||
import abc
|
||||
import pickle
|
||||
from pathlib import Path
|
||||
import warnings
|
||||
import pandas as pd
|
||||
@@ -11,6 +10,7 @@ from typing import Tuple, Union, List, Dict
|
||||
|
||||
from qlib.data import D
|
||||
from qlib.utils import load_dataset, init_instance_by_config, time_to_slc_point
|
||||
from qlib.utils.pickle_utils import restricted_pickle_load
|
||||
from qlib.log import get_module_logger
|
||||
from qlib.utils.serial import Serializable
|
||||
|
||||
@@ -283,7 +283,7 @@ class StaticDataLoader(DataLoader, Serializable):
|
||||
self._data = pd.read_parquet(self._config, engine="pyarrow")
|
||||
else:
|
||||
with Path(self._config).open("rb") as f:
|
||||
self._data = pickle.load(f)
|
||||
self._data = restricted_pickle_load(f)
|
||||
elif isinstance(self._config, pd.DataFrame):
|
||||
self._data = self._config
|
||||
|
||||
|
||||
@@ -168,7 +168,7 @@ class SeriesDFilter(BaseDFilter):
|
||||
for _ts, _bool in timestamp_series.items():
|
||||
# there is likely to be NAN when the filter series don't have the
|
||||
# bool value, so we just change the NAN into False
|
||||
if _bool == np.nan:
|
||||
if np.isnan(_bool):
|
||||
_bool = False
|
||||
if _lbool is None:
|
||||
_cur_start = _ts
|
||||
|
||||
@@ -2,17 +2,18 @@
|
||||
# Licensed under the MIT License.
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
from pathlib import Path
|
||||
from typing import cast, List
|
||||
from typing import List, cast
|
||||
|
||||
import cachetools
|
||||
import pandas as pd
|
||||
import pickle
|
||||
import os
|
||||
|
||||
from qlib.backtest import Exchange, Order
|
||||
from qlib.backtest.decision import TradeRange, TradeRangeByTime
|
||||
from qlib.constant import EPS_T
|
||||
from qlib.utils.pickle_utils import restricted_pickle_load
|
||||
|
||||
from .base import BaseIntradayBacktestData, BaseIntradayProcessedData, ProcessedDataProvider
|
||||
|
||||
|
||||
@@ -162,7 +163,7 @@ class HandlerIntradayProcessedData(BaseIntradayProcessedData):
|
||||
path = os.path.join(data_dir, "backtest" if backtest else "feature", f"{stock_id}.pkl")
|
||||
start_time, end_time = date.replace(hour=0, minute=0, second=0), date.replace(hour=23, minute=59, second=59)
|
||||
with open(path, "rb") as fstream:
|
||||
dataset = pickle.load(fstream)
|
||||
dataset = restricted_pickle_load(fstream)
|
||||
data = dataset.handler.fetch(pd.IndexSlice[stock_id, start_time:end_time], level=None)
|
||||
|
||||
if index_only:
|
||||
|
||||
@@ -11,7 +11,6 @@ import contextlib
|
||||
import importlib
|
||||
import os
|
||||
from pathlib import Path
|
||||
import pickle
|
||||
import pkgutil
|
||||
import re
|
||||
import sys
|
||||
@@ -20,6 +19,7 @@ from typing import Any, Dict, List, Tuple, Union
|
||||
from urllib.parse import urlparse
|
||||
|
||||
from qlib.typehint import InstConf
|
||||
from qlib.utils.pickle_utils import restricted_pickle_load
|
||||
|
||||
|
||||
def get_module_by_module_path(module_path: Union[str, ModuleType]):
|
||||
@@ -168,10 +168,10 @@ def init_instance_by_config(
|
||||
|
||||
pr_path = os.path.join(pr.netloc, path) if bool(pr.path) else pr.netloc
|
||||
with open(os.path.normpath(pr_path), "rb") as f:
|
||||
return pickle.load(f)
|
||||
return restricted_pickle_load(f)
|
||||
else:
|
||||
with config.open("rb") as f:
|
||||
return pickle.load(f)
|
||||
return restricted_pickle_load(f)
|
||||
|
||||
klass, cls_kwargs = get_callable_kwargs(config, default_module=default_module)
|
||||
|
||||
|
||||
@@ -6,6 +6,7 @@ import tempfile
|
||||
from pathlib import Path
|
||||
|
||||
from qlib.config import C
|
||||
from qlib.utils.pickle_utils import restricted_pickle_load
|
||||
|
||||
|
||||
class ObjManager:
|
||||
@@ -116,7 +117,7 @@ class FileManager(ObjManager):
|
||||
|
||||
def load_obj(self, name):
|
||||
with (self.path / name).open("rb") as f:
|
||||
return pickle.load(f)
|
||||
return restricted_pickle_load(f)
|
||||
|
||||
def exists(self, name):
|
||||
return (self.path / name).exists()
|
||||
|
||||
171
qlib/utils/pickle_utils.py
Normal file
171
qlib/utils/pickle_utils.py
Normal file
@@ -0,0 +1,171 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# Licensed under the MIT License.
|
||||
"""
|
||||
Secure pickle utilities to prevent arbitrary code execution through deserialization.
|
||||
|
||||
This module provides a secure alternative to pickle.load() and pickle.loads()
|
||||
that restricts deserialization to a whitelist of safe classes.
|
||||
"""
|
||||
|
||||
import io
|
||||
import pickle
|
||||
from typing import Any, BinaryIO, Set, Tuple
|
||||
|
||||
# Whitelist of safe classes that are allowed to be unpickled
|
||||
# These are common data types used in qlib that should be safe to deserialize
|
||||
SAFE_PICKLE_CLASSES: Set[Tuple[str, str]] = {
|
||||
# python builtins
|
||||
("builtins", "slice"),
|
||||
("builtins", "range"),
|
||||
("builtins", "dict"),
|
||||
("builtins", "list"),
|
||||
("builtins", "tuple"),
|
||||
("builtins", "set"),
|
||||
("builtins", "frozenset"),
|
||||
("builtins", "bytearray"),
|
||||
("builtins", "bytes"),
|
||||
("builtins", "str"),
|
||||
("builtins", "int"),
|
||||
("builtins", "float"),
|
||||
("builtins", "bool"),
|
||||
("builtins", "complex"),
|
||||
("builtins", "type"),
|
||||
("builtins", "property"),
|
||||
# common utility classes
|
||||
("datetime", "datetime"),
|
||||
("datetime", "date"),
|
||||
("datetime", "time"),
|
||||
("datetime", "timedelta"),
|
||||
("datetime", "timezone"),
|
||||
("decimal", "Decimal"),
|
||||
("collections", "OrderedDict"),
|
||||
("collections", "defaultdict"),
|
||||
("collections", "Counter"),
|
||||
("collections", "namedtuple"),
|
||||
("enum", "Enum"),
|
||||
("pathlib", "Path"),
|
||||
("pathlib", "PosixPath"),
|
||||
("pathlib", "WindowsPath"),
|
||||
("qlib.data.dataset.handler", "DataHandler"),
|
||||
("qlib.data.dataset.handler", "DataHandlerLP"),
|
||||
("qlib.data.dataset.loader", "StaticDataLoader"),
|
||||
}
|
||||
|
||||
|
||||
TRUSTED_MODULE_PREFIXES = (
|
||||
"pandas",
|
||||
"numpy",
|
||||
)
|
||||
|
||||
|
||||
class RestrictedUnpickler(pickle.Unpickler):
|
||||
"""Custom unpickler that only allows safe classes to be deserialized.
|
||||
|
||||
This prevents arbitrary code execution through malicious pickle files by
|
||||
restricting deserialization to a whitelist of safe classes.
|
||||
|
||||
Example:
|
||||
>>> with open("data.pkl", "rb") as f:
|
||||
... data = RestrictedUnpickler(f).load()
|
||||
"""
|
||||
|
||||
def find_class(self, module: str, name: str):
|
||||
"""Override find_class to restrict allowed classes.
|
||||
|
||||
Args:
|
||||
module: Module name of the class
|
||||
name: Class name
|
||||
|
||||
Returns:
|
||||
The class object if it's in the whitelist
|
||||
|
||||
Raises:
|
||||
pickle.UnpicklingError: If the class is not in the whitelist
|
||||
"""
|
||||
if module.startswith(TRUSTED_MODULE_PREFIXES):
|
||||
return super().find_class(module, name)
|
||||
|
||||
# 2. explicit whitelist (qlib internal)
|
||||
if (module, name) in SAFE_PICKLE_CLASSES:
|
||||
return super().find_class(module, name)
|
||||
|
||||
raise pickle.UnpicklingError(
|
||||
f"Forbidden class: {module}.{name}. "
|
||||
f"Only whitelisted classes are allowed for security reasons. "
|
||||
f"This is to prevent arbitrary code execution through pickle deserialization."
|
||||
)
|
||||
|
||||
|
||||
def restricted_pickle_load(file: BinaryIO) -> Any:
|
||||
"""Safely load a pickle file with restricted classes.
|
||||
|
||||
This is a drop-in replacement for pickle.load() that prevents
|
||||
arbitrary code execution by only allowing whitelisted classes.
|
||||
|
||||
Args:
|
||||
file: An opened file object in binary mode
|
||||
|
||||
Returns:
|
||||
The unpickled Python object
|
||||
|
||||
Raises:
|
||||
pickle.UnpicklingError: If the pickle contains forbidden classes
|
||||
|
||||
Example:
|
||||
>>> with open("data.pkl", "rb") as f:
|
||||
... data = restricted_pickle_load(f)
|
||||
"""
|
||||
return RestrictedUnpickler(file).load()
|
||||
|
||||
|
||||
def restricted_pickle_loads(data: bytes) -> Any:
|
||||
"""Safely load a pickle from bytes with restricted classes.
|
||||
|
||||
This is a drop-in replacement for pickle.loads() that prevents
|
||||
arbitrary code execution by only allowing whitelisted classes.
|
||||
|
||||
Args:
|
||||
data: Bytes object containing pickled data
|
||||
|
||||
Returns:
|
||||
The unpickled Python object
|
||||
|
||||
Raises:
|
||||
pickle.UnpicklingError: If the pickle contains forbidden classes
|
||||
|
||||
Example:
|
||||
>>> data = b'\\x80\\x04\\x95...'
|
||||
>>> obj = restricted_pickle_loads(data)
|
||||
"""
|
||||
file_like = io.BytesIO(data)
|
||||
return RestrictedUnpickler(file_like).load()
|
||||
|
||||
|
||||
def add_safe_class(module: str, name: str) -> None:
|
||||
"""Add a class to the whitelist of safe classes for unpickling.
|
||||
|
||||
Use this function to extend the whitelist if your code needs to deserialize
|
||||
additional classes. However, be very careful when adding classes, as this
|
||||
could potentially introduce security vulnerabilities.
|
||||
|
||||
Args:
|
||||
module: Module name of the class (e.g., 'my_package.my_module')
|
||||
name: Class name (e.g., 'MyClass')
|
||||
|
||||
Warning:
|
||||
Only add classes that you fully control and trust. Adding arbitrary
|
||||
classes from external packages could introduce security risks.
|
||||
|
||||
Example:
|
||||
>>> add_safe_class('my_package.models', 'CustomModel')
|
||||
"""
|
||||
SAFE_PICKLE_CLASSES.add((module, name))
|
||||
|
||||
|
||||
def get_safe_classes() -> Set[Tuple[str, str]]:
|
||||
"""Get a copy of the current whitelist of safe classes.
|
||||
|
||||
Returns:
|
||||
A set of (module, name) tuples representing allowed classes
|
||||
"""
|
||||
return SAFE_PICKLE_CLASSES.copy()
|
||||
@@ -106,15 +106,13 @@ def handler_mod(task: dict, rolling_gen):
|
||||
rg (RollingGen): an instance of RollingGen
|
||||
"""
|
||||
try:
|
||||
interval = rolling_gen.ta.cal_interval(
|
||||
task["dataset"]["kwargs"]["handler"]["kwargs"]["end_time"],
|
||||
task["dataset"]["kwargs"]["segments"][rolling_gen.test_key][1],
|
||||
)
|
||||
# if end_time < the end of test_segments, then change end_time to allow load more data
|
||||
if interval < 0:
|
||||
task["dataset"]["kwargs"]["handler"]["kwargs"]["end_time"] = copy.deepcopy(
|
||||
task["dataset"]["kwargs"]["segments"][rolling_gen.test_key][1]
|
||||
)
|
||||
handler_kwargs = task["dataset"]["kwargs"]["handler"]["kwargs"]
|
||||
handler_end_time = handler_kwargs.get("end_time")
|
||||
test_seg_end_time = task["dataset"]["kwargs"]["segments"][rolling_gen.test_key][1]
|
||||
# if the end of test_segments is None (open-ended segment, i.e., "until now") or end_time < the end of test_segments,
|
||||
# then change end_time to allow load more data
|
||||
if test_seg_end_time is None or rolling_gen.ta.cal_interval(handler_end_time, test_seg_end_time) < 0:
|
||||
handler_kwargs["end_time"] = copy.deepcopy(test_seg_end_time)
|
||||
except KeyError:
|
||||
# Maybe dataset do not have handler, then do nothing.
|
||||
pass
|
||||
|
||||
@@ -28,6 +28,7 @@ from tqdm.cli import tqdm
|
||||
|
||||
from .utils import get_mongodb
|
||||
from ...config import C
|
||||
from ...utils.pickle_utils import restricted_pickle_loads
|
||||
|
||||
|
||||
class TaskManager:
|
||||
@@ -131,7 +132,7 @@ class TaskManager:
|
||||
for prefix in self.ENCODE_FIELDS_PREFIX:
|
||||
for k in list(task.keys()):
|
||||
if k.startswith(prefix):
|
||||
task[k] = pickle.loads(task[k])
|
||||
task[k] = restricted_pickle_loads(task[k])
|
||||
return task
|
||||
|
||||
def _dict_to_str(self, flt):
|
||||
|
||||
@@ -1,12 +1,12 @@
|
||||
from loguru import logger
|
||||
import os
|
||||
from typing import Optional
|
||||
|
||||
import fire
|
||||
import pandas as pd
|
||||
import qlib
|
||||
from loguru import logger
|
||||
from tqdm import tqdm
|
||||
|
||||
import qlib
|
||||
from qlib.data import D
|
||||
|
||||
|
||||
@@ -36,6 +36,7 @@ class DataHealthChecker:
|
||||
self.large_step_threshold_price = large_step_threshold_price
|
||||
self.large_step_threshold_volume = large_step_threshold_volume
|
||||
self.missing_data_num = missing_data_num
|
||||
self.qlib_dir = os.path.abspath(os.path.expanduser(qlib_dir))
|
||||
|
||||
if csv_path:
|
||||
assert os.path.isdir(csv_path), f"{csv_path} should be a directory."
|
||||
@@ -68,6 +69,43 @@ class DataHealthChecker:
|
||||
self.data[instrument] = df
|
||||
print(df)
|
||||
|
||||
# NOTE:
|
||||
# This check is added due to a known issue in Qlib where feature paths
|
||||
# are constructed using lowercased instrument names. On case-sensitive
|
||||
# file systems (e.g. Linux), uppercase directory names under `features/`
|
||||
# will cause data loading failures.
|
||||
#
|
||||
# See: https://github.com/microsoft/qlib/issues/2053
|
||||
def check_features_dir_lowercase(self) -> Optional[pd.DataFrame]:
|
||||
"""
|
||||
Check whether all subdirectories under `<qlib_dir>/features` are named in lowercase.
|
||||
|
||||
This validation helps prevent data loading issues on case-sensitive
|
||||
file systems caused by uppercase instrument directory names.
|
||||
"""
|
||||
if not self.qlib_dir:
|
||||
return None
|
||||
|
||||
features_dir = os.path.join(self.qlib_dir, "features")
|
||||
if not os.path.isdir(features_dir):
|
||||
logger.warning(f"`features` directory not found under {self.qlib_dir}")
|
||||
return None
|
||||
|
||||
bad_dirs = []
|
||||
for name in os.listdir(features_dir):
|
||||
full_path = os.path.join(features_dir, name)
|
||||
if os.path.isdir(full_path) and name != name.lower():
|
||||
bad_dirs.append(name)
|
||||
|
||||
if bad_dirs:
|
||||
result_df = pd.DataFrame({"non_lowercase_dir": bad_dirs})
|
||||
return result_df
|
||||
else:
|
||||
logger.info(
|
||||
f"✅ All subdirectories under `{os.path.join(self.qlib_dir, 'features')}` are named in lowercase."
|
||||
)
|
||||
return None
|
||||
|
||||
def check_missing_data(self) -> Optional[pd.DataFrame]:
|
||||
"""Check if any data is missing in the DataFrame."""
|
||||
result_dict = {
|
||||
@@ -177,11 +215,13 @@ class DataHealthChecker:
|
||||
check_large_step_changes_result = self.check_large_step_changes()
|
||||
check_required_columns_result = self.check_required_columns()
|
||||
check_missing_factor_result = self.check_missing_factor()
|
||||
check_features_dir_case_result = self.check_features_dir_lowercase()
|
||||
if (
|
||||
check_large_step_changes_result is not None
|
||||
or check_large_step_changes_result is not None
|
||||
or check_required_columns_result is not None
|
||||
or check_missing_factor_result is not None
|
||||
or check_features_dir_case_result is not None
|
||||
):
|
||||
print(f"\nSummary of data health check ({len(self.data)} files checked):")
|
||||
print("-------------------------------------------------")
|
||||
@@ -197,6 +237,11 @@ class DataHealthChecker:
|
||||
if isinstance(check_missing_factor_result, pd.DataFrame):
|
||||
logger.warning(f"The factor column does not exist or is empty")
|
||||
print(check_missing_factor_result)
|
||||
if isinstance(check_features_dir_case_result, pd.DataFrame):
|
||||
logger.warning(
|
||||
f"Some subdirectories under `{os.path.join(self.qlib_dir, 'features')}` contain uppercase letters, please rename them to lowercase manually."
|
||||
)
|
||||
print(check_features_dir_case_result)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
||||
@@ -7,12 +7,14 @@ import sys
|
||||
from pathlib import Path
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from typing import List
|
||||
from io import StringIO
|
||||
|
||||
import fire
|
||||
import requests
|
||||
import pandas as pd
|
||||
from tqdm import tqdm
|
||||
from loguru import logger
|
||||
from fake_useragent import UserAgent
|
||||
|
||||
|
||||
CUR_DIR = Path(__file__).resolve().parent
|
||||
@@ -51,6 +53,7 @@ class WIKIIndex(IndexBase):
|
||||
)
|
||||
|
||||
self._target_url = f"{WIKI_URL}/{WIKI_INDEX_NAME_MAP[self.index_name.upper()]}"
|
||||
self._ua = UserAgent()
|
||||
|
||||
@property
|
||||
@abc.abstractmethod
|
||||
@@ -112,7 +115,8 @@ class WIKIIndex(IndexBase):
|
||||
return _calendar_list
|
||||
|
||||
def _request_new_companies(self) -> requests.Response:
|
||||
resp = requests.get(self._target_url, timeout=None)
|
||||
headers = {"User-Agent": self._ua.random}
|
||||
resp = requests.get(self._target_url, timeout=None, headers=headers)
|
||||
if resp.status_code != 200:
|
||||
raise ValueError(f"request error: {self._target_url}")
|
||||
|
||||
@@ -128,7 +132,7 @@ class WIKIIndex(IndexBase):
|
||||
def get_new_companies(self):
|
||||
logger.info(f"get new companies {self.index_name} ......")
|
||||
_data = deco_retry(retry=self._request_retry, retry_sleep=self._retry_sleep)(self._request_new_companies)()
|
||||
df_list = pd.read_html(_data.text)
|
||||
df_list = pd.read_html(StringIO(_data.text))
|
||||
for _df in df_list:
|
||||
_df = self.filter_df(_df)
|
||||
if (_df is not None) and (not _df.empty):
|
||||
@@ -226,7 +230,11 @@ class SP500Index(WIKIIndex):
|
||||
def get_changes(self) -> pd.DataFrame:
|
||||
logger.info(f"get sp500 history changes......")
|
||||
# NOTE: may update the index of the table
|
||||
changes_df = pd.read_html(self.WIKISP500_CHANGES_URL)[-1]
|
||||
# Add headers to avoid 403 Forbidden error from Wikipedia
|
||||
headers = {"User-Agent": self._ua.random}
|
||||
response = requests.get(self.WIKISP500_CHANGES_URL, headers=headers, timeout=None)
|
||||
response.raise_for_status()
|
||||
changes_df = pd.read_html(StringIO(response.text))[-1]
|
||||
changes_df = changes_df.iloc[:, [0, 1, 3]]
|
||||
changes_df.columns = [self.DATE_FIELD_NAME, self.ADD, self.REMOVE]
|
||||
changes_df[self.DATE_FIELD_NAME] = pd.to_datetime(changes_df[self.DATE_FIELD_NAME])
|
||||
|
||||
@@ -3,3 +3,4 @@ requests
|
||||
pandas
|
||||
lxml
|
||||
loguru
|
||||
fake-useragent
|
||||
|
||||
@@ -7,7 +7,6 @@ import importlib
|
||||
import time
|
||||
import bisect
|
||||
import pickle
|
||||
import random
|
||||
import requests
|
||||
import functools
|
||||
from pathlib import Path
|
||||
@@ -80,28 +79,14 @@ def get_calendar_list(bench_code="CSI300") -> List[pd.Timestamp]:
|
||||
calendar = df.index.get_level_values(level="date").map(pd.Timestamp).unique().tolist()
|
||||
else:
|
||||
if bench_code.upper() == "ALL":
|
||||
import akshare as ak # pylint: disable=C0415
|
||||
|
||||
@deco_retry
|
||||
def _get_calendar_from_month(month):
|
||||
_cal = []
|
||||
try:
|
||||
resp = requests.get(
|
||||
SZSE_CALENDAR_URL.format(month=month, random=random.random), timeout=None
|
||||
).json()
|
||||
for _r in resp["data"]:
|
||||
if int(_r["jybz"]):
|
||||
_cal.append(pd.Timestamp(_r["jyrq"]))
|
||||
except Exception as e:
|
||||
raise ValueError(f"{month}-->{e}") from e
|
||||
return _cal
|
||||
|
||||
month_range = pd.date_range(start="2000-01", end=pd.Timestamp.now() + pd.Timedelta(days=31), freq="M")
|
||||
calendar = []
|
||||
for _m in month_range:
|
||||
cal = _get_calendar_from_month(_m.strftime("%Y-%m"))
|
||||
if cal:
|
||||
calendar += cal
|
||||
calendar = list(filter(lambda x: x <= pd.Timestamp.now(), calendar))
|
||||
trade_date_df = ak.tool_trade_date_hist_sina()
|
||||
trade_date_list = trade_date_df["trade_date"].tolist()
|
||||
trade_date_list = [pd.Timestamp(d) for d in trade_date_list]
|
||||
dates = pd.DatetimeIndex(trade_date_list)
|
||||
filtered_dates = dates[(dates >= "2000-01-04") & (dates <= pd.Timestamp.today().normalize())]
|
||||
calendar = filtered_dates.tolist()
|
||||
else:
|
||||
calendar = _get_calendar(CALENDAR_BENCH_URL_MAP[bench_code])
|
||||
_CALENDAR_MAP[bench_code] = calendar
|
||||
|
||||
@@ -10,3 +10,4 @@ joblib
|
||||
beautifulsoup4
|
||||
bs4
|
||||
soupsieve
|
||||
akshare
|
||||
@@ -1,10 +1,10 @@
|
||||
import os
|
||||
import pickle
|
||||
import shutil
|
||||
import unittest
|
||||
from qlib.tests import TestAutoData
|
||||
|
||||
from qlib.data import D
|
||||
from qlib.data.dataset.handler import DataHandlerLP
|
||||
from qlib.tests import TestAutoData
|
||||
from qlib.utils.pickle_utils import restricted_pickle_load
|
||||
|
||||
|
||||
class HandlerTests(TestAutoData):
|
||||
@@ -23,7 +23,7 @@ class HandlerTests(TestAutoData):
|
||||
dh.to_pickle(fname, dump_all=True)
|
||||
|
||||
with open(fname, "rb") as f:
|
||||
dh_d = pickle.load(f)
|
||||
dh_d = restricted_pickle_load(f)
|
||||
|
||||
self.assertTrue(dh_d._data.equals(df))
|
||||
self.assertTrue(dh_d._infer is dh_d._data)
|
||||
|
||||
Reference in New Issue
Block a user