diff --git a/.github/ci/README.md b/.github/ci/README.md index f0403512b2f..2fde10963f0 100644 --- a/.github/ci/README.md +++ b/.github/ci/README.md @@ -76,8 +76,17 @@ selection, add caches, or reduce test data sizes. ## Network and workflow retries -Dataset commands have at most three attempts, with a 15-minute timeout per -attempt. `--delete_old False` makes retries noninteractive in these fresh CI +Dataset commands name the existing `v2` release archives explicitly: source +jobs use `qlib_data_simple_cn_1d_latest.zip`, and PyPI jobs use +`qlib_data_cn_1d_latest.zip`. These are the same datasets selected by successful +version-probing downloads, not smaller substitutes. Naming the archive avoids +an unnecessary request for a nonexistent package-version asset: the downloader +can otherwise mistake an HTTP 504 response for confirmation that the file exists. + +Download steps have at most five attempts, with a 15-minute timeout per +attempt and a 60-second pause between failures, allowing transient GitHub +download outages more time to recover than three closely spaced attempts. +`--delete_old False` makes retries noninteractive in these fresh CI directories, including when an earlier attempt already extracted one dataset. Bash command blocks use `set -euo pipefail`, so an earlier failure cannot be hidden by a later successful command. Exhausted retries still fail the job. diff --git a/.github/workflows/test_qlib_from_pip.yml b/.github/workflows/test_qlib_from_pip.yml index cdb3d459452..1d992bd95c6 100644 --- a/.github/workflows/test_qlib_from_pip.yml +++ b/.github/workflows/test_qlib_from_pip.yml @@ -64,13 +64,13 @@ jobs: uses: nick-fields/retry@ad984534de44a9489a53aefd81eb77f87c70dc60 # v4 with: timeout_minutes: 15 - max_attempts: 3 - retry_wait_seconds: 10 + max_attempts: 5 + retry_wait_seconds: 60 shell: bash command: | set -euo pipefail cd "$RUNNER_TEMP" - python -m qlib.cli.data qlib_data --target_dir ~/.qlib/qlib_data/cn_data --region cn --delete_old False + python -m qlib.cli.data download_data --file_name v2/qlib_data_cn_1d_latest.zip --target_dir ~/.qlib/qlib_data/cn_data --delete_old False - name: Test workflow by config uses: nick-fields/retry@ad984534de44a9489a53aefd81eb77f87c70dc60 # v4 diff --git a/.github/workflows/test_qlib_from_source.yml b/.github/workflows/test_qlib_from_source.yml index 78fe587a5c4..0135b252510 100644 --- a/.github/workflows/test_qlib_from_source.yml +++ b/.github/workflows/test_qlib_from_source.yml @@ -108,12 +108,12 @@ jobs: uses: nick-fields/retry@ad984534de44a9489a53aefd81eb77f87c70dc60 # v4 with: timeout_minutes: 15 - max_attempts: 3 - retry_wait_seconds: 10 + max_attempts: 5 + retry_wait_seconds: 60 shell: bash command: | set -euo pipefail - python scripts/get_data.py qlib_data --name qlib_data_simple --target_dir ~/.qlib/qlib_data/cn_data --interval 1d --region cn --delete_old False + python scripts/get_data.py download_data --file_name v2/qlib_data_simple_cn_1d_latest.zip --target_dir ~/.qlib/qlib_data/cn_data --delete_old False python scripts/get_data.py download_data --file_name rl_data.zip --target_dir tests/.data/rl --delete_old False - name: Check Qlib ipynb with nbconvert diff --git a/.github/workflows/test_qlib_from_source_slow.yml b/.github/workflows/test_qlib_from_source_slow.yml index 255eb3fb11a..61356a63ac1 100644 --- a/.github/workflows/test_qlib_from_source_slow.yml +++ b/.github/workflows/test_qlib_from_source_slow.yml @@ -78,10 +78,10 @@ jobs: uses: nick-fields/retry@ad984534de44a9489a53aefd81eb77f87c70dc60 # v4 with: timeout_minutes: 15 - max_attempts: 3 - retry_wait_seconds: 10 + max_attempts: 5 + retry_wait_seconds: 60 shell: bash - command: python scripts/get_data.py qlib_data --name qlib_data_simple --target_dir ~/.qlib/qlib_data/cn_data --interval 1d --region cn --delete_old False + command: python scripts/get_data.py download_data --file_name v2/qlib_data_simple_cn_1d_latest.zip --target_dir ~/.qlib/qlib_data/cn_data --delete_old False - name: Unit tests with Pytest uses: nick-fields/retry@ad984534de44a9489a53aefd81eb77f87c70dc60 # v4 diff --git a/CHANGELOG.md b/CHANGELOG.md index e69de29bb2d..2dfc52f99ab 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -0,0 +1,13 @@ +# Changelog + +## Unreleased + +- **BREAKING:** New source builds restrict recorder artifact loading by default. + Reloading executable artifacts requires verified source/storage and explicit + `trusted=True` (CLI: `--trusted=True`). Supported data-only reads and fresh + in-memory training need no opt-in. See the + [artifact loading migration guide](https://qlib.readthedocs.io/en/latest/start/artifact_migration.html) + for workflow, HIST and high-frequency cache upgrades. +- Merging into `main` affects source installs before a PyPI release. These changes + remain unreleased until included in a tagged release; its versioned upgrade notes + should link to the same guide. diff --git a/CHANGES.rst b/CHANGES.rst index 76aa4829304..190cd282713 100644 --- a/CHANGES.rst +++ b/CHANGES.rst @@ -2,6 +2,18 @@ Changelog ========= Here you can see the full list of changes between each QLib release. +Unreleased +---------- +- **BREAKING:** New source builds restrict recorder artifact loading by default. + Reloading executable artifacts requires verified source/storage and explicit + ``trusted=True`` (CLI: ``--trusted=True``); supported data-only reads and fresh + in-memory training need no opt-in. See the + `artifact loading migration guide `_ + for workflow, HIST and high-frequency cache upgrades. +- Merging into ``main`` affects source installs before a PyPI release. These + changes remain unreleased until included in a tagged release; its versioned + upgrade notes should link to the same guide. + Version 0.1.0 ------------- This is the initial release of QLib library. diff --git a/README.md b/README.md index 42c446e69e5..5cfe55ba48a 100644 --- a/README.md +++ b/README.md @@ -7,6 +7,20 @@ [![License](https://img.shields.io/pypi/l/pyqlib)](LICENSE) [![Join the chat at https://gitter.im/Microsoft/qlib](https://badges.gitter.im/Microsoft/qlib.svg)](https://gitter.im/Microsoft/qlib?utm_source=badge&utm_medium=badge&utm_campaign=pr-badge&utm_content=badge) +## Unreleased breaking change: artifact loading + +New source builds introduce a **backward-incompatible safety default**: recorder +loads are restricted. If you reload saved models/datasets, resume online or delayed +training, or use DDG-DA, verify the artifact writer and storage permissions before +opting in with `trusted=True` (CLI: `--trusted=True`). Fresh in-memory training and +supported data-only reads need no opt-in. HIST mappings and high-frequency cache +paths also have migration requirements. + +Follow the [artifact loading migration guide](https://qlib.readthedocs.io/en/latest/start/artifact_migration.html) +before upgrading. Merging into `main` affects source installs before a new PyPI +release; this change remains **unreleased until the first tagged release containing +it**. That release's versioned upgrade notes should link to the same guide. + ## :newspaper: **What's NEW!**   :sparkling_heart: Recent released features diff --git a/docs/advanced/serial.rst b/docs/advanced/serial.rst index e50ee91ddaf..d2896b24c60 100644 --- a/docs/advanced/serial.rst +++ b/docs/advanced/serial.rst @@ -23,6 +23,13 @@ Example ``Qlib``'s serializable class includes ``DataHandler``, ``DataSet``, ``Processor`` and ``Model``, etc., which are subclass of ``qlib.utils.serial.Serializable``. Specifically, ``qlib.data.dataset.DatasetH`` is one of them. Users can serialize ``DatasetH`` as follows. +.. warning:: + + The direct local ``pickle.load`` below can execute code and must only read a + file from an independently trusted source and store. Recorder/workflow + ``trusted`` settings do not authorize or protect this API. For the separate + recorder loading policy and upgrade steps, see :ref:`artifact_loading_migration`. + .. code-block:: Python ##=============dump dataset============= diff --git a/docs/advanced/task_management.rst b/docs/advanced/task_management.rst index b1cb6c696a5..e59d7f91f2f 100644 --- a/docs/advanced/task_management.rst +++ b/docs/advanced/task_management.rst @@ -79,6 +79,17 @@ Meanwhile, ``Qlib`` provides a module called ``Trainer``. If you do not want to use ``Task Manager`` to manage tasks, then use TrainerR to train a list of tasks generated by ``TaskGen`` is enough. `Here <../reference/api.html#Trainer>`_ are the details about different ``Trainer``. +Delayed training reloads saved task configurations from recorders. For tasks +containing executable objects such as reweighters, first verify their writer and +artifact store, then configure +``DelayTrainerR(trusted=True)`` or +``DelayTrainerRM(trusted=True)``. This consent is forwarded to delayed +training completion, including the ``DelayTrainerRM`` worker. Direct callers can +instead pass ``trusted=True`` to ``end_task_train`` or ``end_train``. +Do not pass this constructor option to ordinary ``TrainerR`` or ``TrainerRM``. +The option concerns recorder artifacts, not the independent trust requirements of +MongoDB task contents. See :ref:`artifact_loading_migration`. + Task Collecting =============== Before collecting model training results, you need to use the ``qlib.init`` to specify the path of mlruns. @@ -87,6 +98,14 @@ To collect the results of ``task`` after training, ``Qlib`` provides `Collector `Collector <../reference/api.html#Collector>`_ can collect objects from everywhere and process them such as merging, grouping, averaging and so on. It has 2 step action including ``collect`` (collect anything in a dict) and ``process_collect`` (process collected dict). +``RecorderCollector`` keeps artifact reads restricted by default. To collect +verified executable objects alongside predictions, use per-artifact loading +options, for example ``artifact_load_kwargs={"model": {"trusted": True}}`` with +``artifacts_path={"model": "params.pkl", "pred": "pred.pkl"}``. Only the named +``model`` entry receives consent; prediction reads remain restricted. Use +``collect(only_exist=False)`` when loading failures must not produce a partial +collection. See :ref:`artifact_loading_migration` for a complete example. + `Group <../reference/api.html#Group>`_ also has 2 steps including ``group`` (can group a set of object based on `group_func` and change them to a dict) and ``reduce`` (can make a dict become an ensemble based on some rule). For example: {(A,B,C1): object, (A,B,C2): object} ---``group``---> {(A,B): {C1: object, C2: object}} ---``reduce``---> {(A,B): object} diff --git a/docs/component/highfreq.rst b/docs/component/highfreq.rst index 655b26b9367..a7effb582c7 100644 --- a/docs/component/highfreq.rst +++ b/docs/component/highfreq.rst @@ -35,6 +35,13 @@ An example of a nested decision execution framework for high-frequency can be fo Besides, the above examples, here are some other related works about high-frequency trading in Qlib. +.. note:: + + New source builds require ``HighFreqProvider`` artifact paths, including derived + cache files, to remain inside ``artifact_root`` (the current directory by + default). Cached pickle contents still require independent trust. See + :ref:`artifact_loading_migration` before reusing existing provider configurations. + - `Prediction with high-frequency data `_ - `Examples `_ to extract features from high-frequency data without fixed frequency. - `A paper `_ for high-frequency trading. diff --git a/docs/component/meta.rst b/docs/component/meta.rst index 2b5f22a18aa..d79388af9cc 100644 --- a/docs/component/meta.rst +++ b/docs/component/meta.rst @@ -66,3 +66,16 @@ which adapts to the market dynamics. 4. Apply guide information to the forecasting models to improve their performances. The `above example `_ can be found in ``examples/benchmarks_dynamic/DDG-DA/workflow.py``. + +DDG-DA uses restricted loading by default for recorder artifacts and local +handler/internal-data pickle caches. To reload executable meta-models, tasks and +caches from a trusted writer and access-controlled storage, configure +``DDGDA(..., trusted=True)`` at the workflow entry point. +Lower-level callers can configure ``MetaDatasetDS`` or ``InternalData.setup`` with +the same option for recorder task reads. Prediction, label and numerical-report reads +remain restricted. DDG-DA's opt-in also authorizes its handler/internal-data +cache reads, so verify ``working_dir``, the configuration directory and any +explicit ``h_path`` as well as the MLflow store. There is no automatic unsafe retry +and the global restricted loader is unchanged. +See :ref:`artifact_loading_migration` +and the example README for CLI commands. diff --git a/docs/component/model.rst b/docs/component/model.rst index 60cfa58fab4..310e94c065a 100644 --- a/docs/component/model.rst +++ b/docs/component/model.rst @@ -110,6 +110,21 @@ Also, the above example has been given in ``examples/train_backtest_analyze.ipyn Technically, the meaning of the model prediction depends on the label setting designed by user. By default, the meaning of the score is normally the rating of the instruments by the forecasting model. The higher the score, the more profit the instruments. +Reloading a saved model +======================= + +Training and prediction with in-memory objects, as above, do not require a recorder +trust grant. Reloading a saved model or dataset may execute Python code, however. +For artifacts from a verified writer in an access-controlled store, explicitly use +``recorder.load_object("params.pkl", trusted=True)`` (or the corresponding ``R`` +facade). Keep prediction and label data reads restricted by omitting ``trusted``. +Do not automatically retry a refused load with unrestricted deserialization. + +Workflows that resume training can configure ``DelayTrainerR`` or +``DelayTrainerRM`` with ``trusted=True`` at construction, rather than +changing internal loads. See :ref:`artifact_loading_migration` for migration details, +version limitations and the separate trust requirements of local pickle files. + Custom Model ============ diff --git a/docs/component/online.rst b/docs/component/online.rst index d7113c19fcf..54af83bc6ff 100644 --- a/docs/component/online.rst +++ b/docs/component/online.rst @@ -26,6 +26,31 @@ The `examples ` Known limitations currently - Currently, the daily updating prediction for the next trading day is supported. But generating orders for the next trading day is not supported due to the `limitations of public data _` +Recorder artifact trust +======================= + +Online updates may reload executable model, dataset and task objects from recorders. +The default is restricted loading. After verifying the artifact writer and the +store's write permissions, opt in with ``trusted=True`` on each +``RollingStrategy`` or on a directly constructed ``OnlineToolR``/updater. +``RollingStrategy`` forwards this setting through its online tool and updater; +prediction, label and numerical-report reads remain restricted. + +``OnlineManager`` does not grant trust globally. Configure newly added strategies +as well as initial strategies, and configure ``DelayTrainerR`` or ``DelayTrainerRM`` +separately if used. A supplied trainer keeps its caller-selected trust policy. +Local serialized manager files must also be independently trusted; restoring one +retains the settings saved in it, rather than applying a new manager-wide grant. +Legacy components without a saved flag default to restricted loading. Explicitly +reconfigure or recreate each strategy, its ``strategy.tool``, and any delayed +trainer as needed; an example constructor flag does not override a subsequently +loaded manager. + +See :ref:`artifact_loading_migration` for supported data, refusal handling and custom +loader migration, and the +`example commands `_ +for the default-off ``--trusted=True`` CLI option. + Online Manager ============== diff --git a/docs/component/recorder.rst b/docs/component/recorder.rst index ca545b75bd1..f4f5a642eb3 100644 --- a/docs/component/recorder.rst +++ b/docs/component/recorder.rst @@ -91,9 +91,23 @@ Here are some important APIs that are not included in the ``QlibRecorder``: For other interfaces such as `save_objects`, `load_object`, please refer to `Recorder API <../reference/api.html#recorder>`_. +.. _artifact_trust_migration: + +Migration: explicit artifact trust +================================== + +The migration guide has moved to :ref:`artifact_loading_migration`. See that guide +for the unreleased restricted-loading default, explicit ``trusted=True`` opt-in, +workflow and custom-loader changes, and HIST/high-frequency cache migrations. + Record Template =============== +``RecordTemp.load(name, trusted=True)`` can explicitly authorize an individual +artifact from a verified source and store, including parent-path lookup. The +default remains restricted, and this does not authorize other loads or the +template's ``generate`` method. See :ref:`artifact_loading_migration` for examples. + The ``RecordTemp`` class is a class that enables generate experiment results such as IC and backtest in a certain format. We have provided three different `Record Template` class: - ``SignalRecord``: This class generates the `prediction` results of the model. @@ -152,3 +166,5 @@ For more information about the APIs, please refer to `Record Template API <../re Known Limitations ================= - The Python objects are saved based on pickle, which may results in issues when the environment dumping objects and loading objects are different. +- Restricted loading is intentionally not compatible with arbitrary Python objects. + See :ref:`artifact_loading_migration` before changing trust settings. diff --git a/docs/component/report.rst b/docs/component/report.rst index 01d3a21234d..45f51fe148e 100644 --- a/docs/component/report.rst +++ b/docs/component/report.rst @@ -23,6 +23,14 @@ Introduction All of the accumulated profit metrics(e.g. return, max drawdown) in Qlib are calculated by summation. This avoids the metrics or the plots being skewed exponentially over time. +.. note:: + + Supported numerical report artifacts load in the restricted default mode. + Saved backtest artifacts containing ``Position`` instances or indicator objects + instead require explicit ``trusted=True`` after verifying their writer and + store. See :ref:`artifact_loading_migration`; a report filename alone does not + establish that its contents are data-only. + Graphical Reports ================= diff --git a/docs/index.rst b/docs/index.rst index 49596234a6c..a49ccc6fb2d 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -26,6 +26,7 @@ Document Structure :caption: FIRST STEPS: Installation + Artifact Loading Migration Initialization Data Retrieval Custom Model Integration diff --git a/docs/start/artifact_migration.rst b/docs/start/artifact_migration.rst new file mode 100644 index 00000000000..94ea068323c --- /dev/null +++ b/docs/start/artifact_migration.rst @@ -0,0 +1,531 @@ +.. _artifact_loading_migration: + +============================================= +Migration: explicit trust for saved artifacts +============================================= + +.. important:: + + **Unreleased backward-incompatible safety-default change.** This guide applies + to new source builds containing the artifact-loading changes. Merging into + ``main`` affects source installs before a new ``pyqlib`` release is published + on PyPI; it does not change an already-installed tagged release. Until the + first tagged release containing these changes, treat them as unreleased. + That release's versioned upgrade notes should link to this guide. + +What changes when upgrading +=========================== + +The built-in MLflow recorder now uses a restricted unpickler by default. Fresh +training and prediction using in-memory objects need no opt-in. Supported data-only +reads, including predictions, labels and numerical reports, also remain available +without consent. + +Saved models, datasets, meta-models and tasks containing executable objects such as +reweighters require explicit ``trusted=True`` when reloaded. Backtest artifacts +containing ``Position`` instances or indicator objects require the same decision; +they are not data-only reports. A ``.pkl`` suffix, a familiar artifact name or a +successful run does not establish safety. + +This also applies to objects saved earlier in the same new workflow. Creating a +model or dataset does not implicitly authorize its later deserialization from +an artifact store. + +Before updating an existing workflow: + +1. Identify which loads read data and which restore executable objects. Leave + supported data-only reads restricted. +2. Verify both the artifact producer and storage write permissions before setting + ``trusted=True`` at each relevant workflow entry point (CLI: ``--trusted=True``). +3. Review separately saved workflow components, local caches and custom loaders + using the sections below. Update HIST mappings and high-frequency cache paths + where applicable. +4. Test representative artifacts in a compatible environment before deployment. + +Ordinary trusted model/dataset artifacts do not need deletion, conversion or a +full retrain merely to adopt explicit consent. Example ``main``/``first_run`` +methods can reset experiments and task pools; do not use them as migration +commands. Preserve existing results and use dedicated stores when trying examples. + +Single-object loads and trust boundaries +======================================== + +Keep data-only reads in the default mode: + +.. code-block:: python + + from qlib.workflow import R + + rec = R.get_recorder() # Select the intended run in your experiment. + predictions = rec.load_object("pred.pkl") + labels = R.load_object("label.pkl") + report = rec.load_object("portfolio_analysis/report_normal_1day.pkl") + +After verifying the writer and store, explicitly authorize executable objects: + +.. code-block:: python + + model = rec.load_object("params.pkl", trusted=True) + dataset = R.load_object("dataset", trusted=True) + positions = rec.load_object( + "portfolio_analysis/positions_normal_1day.pkl", trusted=True + ) + +``Recorder.load_object`` and ``R.load_object`` expose keyword-only +``trusted=False``. All public trust options use this single name and require an +actual Python ``bool``, not strings, integers or other truthy values. The example +CLIs accept ``--trusted=True`` as boolean consent. Unrestricted artifact loads +emit an unsafe-loading warning and may execute code with the loading process's +permissions. The flag does not validate, sanitize or authenticate the contents. + +``MLflowRecorder.load_object(name, unpickler=None, *, trusted=False)`` also accepts +a custom unpickler. That is trusted code, emits an unsafe-loading warning, and is +not a sandbox or a way to make unknown artifacts safe. A custom ``unpickler`` and +``trusted=True`` are mutually exclusive and raise ``ValueError`` together; choose +one deliberate loading policy, not both. + +Verify **both source and store**: who produced the artifact, how it reached the +run, and who can replace it. Restrict write access to MLflow artifact directories, +shared filesystems and remote object stores. Creating a run yourself is not +enough if another user or job can overwrite its artifacts. Prefer dedicated, +access-controlled storage. + +There is no automatic fallback to ordinary pickle. Do not catch a restricted-load +failure and retry with ``trusted=True``. A refusal may indicate an unsupported +representation, a missing dependency or executable content; it does not establish +trust. Inspect the reported type and provenance. Do not expand the global class +allowlist simply to suppress a model-loading error. + +Record templates and mixed artifact collections +------------------------------------------------ + +``RecordTemp.load`` accepts the same keyword-only ``trusted=False``. Consent +applies to that call, including lookup through parent record-template paths; +it does not change subsequent loads or the policy used by ``generate``: + +.. code-block:: python + + from qlib.workflow.record_temp import RecordTemp + + template = RecordTemp(rec) + model = template.load("params.pkl", trusted=True) + predictions = template.load("pred.pkl") + +For a collection containing both executable objects and data, supply +``artifact_load_kwargs`` to ``RecorderCollector``. Keys are the aliases in +``artifacts_path``, not filenames. After verifying the model's source and store: + +.. code-block:: python + + from qlib.workflow.task.collect import RecorderCollector + + collector = RecorderCollector( + experiment=lambda: [rec], + artifacts_path={"model": "params.pkl", "pred": "pred.pkl", "label": "label.pkl"}, + artifact_load_kwargs={"model": {"trusted": True}}, + ) + artifacts = collector.collect(only_exist=False) + +Only the ``model`` entry receives consent; ``pred`` and ``label`` keep restricted +loading. Loading options are copied at construction and forwarded to each +recorder's ``load_object`` method. Backend-specific options such as a custom +unpickler retain that backend's trust requirements. There is no collector-wide +permission or automatic retry with unrestricted loading. + +``only_exist=False`` propagates loading failures when every requested artifact +is required. The existing default, ``only_exist=True``, skips failed artifact +reads with a warning that includes the reason; do not treat a partial collection +as proof that all requested artifacts were loaded. + +Workflow-level entry points +=========================== + +These public entry points use default-off ``trusted=False`` so callers need not +patch internal recorder loads: + +* ``RMDLoader``, ``DSBasedUpdater``, ``PredUpdater`` and ``LabelUpdater``: + model/dataset loads needed for updating. +* ``OnlineToolR`` and ``RollingStrategy``: executable task/model/dataset reads. + A strategy forwards its setting to the online tool it creates, then through the + updater to the loader. +* ``DelayTrainerR`` and ``DelayTrainerRM``: saved task reads when finishing delayed + training. Constructor consent reaches ``end_train`` and the ``DelayTrainerRM`` + worker. Direct callers can use ``end_task_train(..., trusted=True)`` or pass an + explicit ``trusted`` override to ``end_train``/``worker`` for that call. +* ``DDGDA``: recorder-backed meta-model/task reads and the local handler and + ``InternalData`` caches needed by that workflow, described below. +* ``MetaDatasetDS`` and ``InternalData.setup``: lower-level recorder task reads. + Their recorder consent alone does not authorize unrelated local pickle caches. + +For example, after checking this workflow's artifact sources: + +.. code-block:: python + + from qlib.model.trainer import DelayTrainerR + from qlib.workflow.online.manager import OnlineManager + from qlib.workflow.online.strategy import RollingStrategy + + strategy = RollingStrategy( + "my_strategy", + task_template=task_template, + rolling_gen=rolling_gen, + trusted=True, + ) + trainer = DelayTrainerR(trusted=True) + manager = OnlineManager(strategy, trainer=trainer) + +Ordinary ``TrainerR``/``TrainerRM`` and ``OnlineManager`` have no ``trusted`` +constructor option. Configure each strategy, including ones added later, and +configure a caller-supplied delayed trainer separately. A workflow does not +silently change that trainer's policy. + +Consent covers the necessary executable objects, not every artifact in a run. +Prediction, label and numerical-report reads in these workflows remain restricted +even when ``trusted=True``. See :ref:`online_serving` and the +`example commands `_. + +Restored managers and components +--------------------------------- + +Only restore a local serialized ``OnlineManager`` from an independently trusted +source. Its strategies, tools and delayed trainer retain their own saved settings; +previously released objects without a trust field default to ``False``. + +After reviewing each source, explicitly reconfigure or recreate every affected +strategy, its ``strategy.tool``, and any delayed trainer. Changing only +``strategy.trusted`` does not update an already-created tool. A new example +constructor or CLI flag does not overwrite a manager subsequently restored from +disk. ``add_strategy`` applies the current flag to new strategies only. There is +no manager-wide permission or revocation. + +Example: re-authorizing a restored manager +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +Initialize Qlib with your existing market-data and tracking configuration first. +Replace the input filename and ``my_strategy`` below with the saved manager and +strategy you intend to resume. Review that strategy's experiment storage and, +separately, all task stores used by its delayed trainer before granting consent. +A tool or trainer can be shared: its policy changes for every reference to that +same instance, not only for the selected strategy. + +.. warning:: + + ``OnlineManager.load`` itself uses unrestricted pickle/dill and can execute + code before any component settings are changed. Only restore an independently + trusted manager file. Setting component flags later does not make this initial + load safe. + +.. code-block:: python + :name: restored-manager-migration + + from pathlib import Path + + from qlib.model.trainer import DelayTrainerR, DelayTrainerRM + from qlib.workflow.online.manager import OnlineManager + from qlib.workflow.online.strategy import RollingStrategy + from qlib.workflow.online.utils import OnlineToolR + + manager = OnlineManager.load("manager.pkl") + matches = [s for s in manager.strategies if s.name_id == "my_strategy"] + if len(matches) != 1: + raise ValueError("Expected exactly one saved strategy named my_strategy") + strategy = matches[0] + if not isinstance(strategy, RollingStrategy) or not isinstance(strategy.tool, OnlineToolR): + raise TypeError("Adapt this example to your custom strategy/tool's loading policy") + + strategy.trusted = True + strategy.tool.trusted = True + if isinstance(manager.trainer, (DelayTrainerR, DelayTrainerRM)): + manager.trainer.trusted = True + + with Path("manager.migrated.pkl").open("xb") as stream: + manager.get_backend().dump(manager, stream, protocol=4) + +This changes the selected strategy and its existing tool, plus the delayed +trainer when present. Other strategy objects are not automatically opted in. +Ordinary ``TrainerR``/``TrainerRM`` need no trust setting, and custom trainers +require their own loading-policy review. Repeat +the strategy configuration only for other components whose sources you have +reviewed, before resuming routines that use them. + +The new file retains the manager's saved history and selected settings; exclusive +creation (``"xb"``) refuses to overwrite an existing output. After reviewing it, +point your resume command at ``manager.migrated.pkl`` rather than loading the old +file again. This migration does not run ``first_train``, reset experiments, or +retrain models. Prediction and label reads remain restricted. + +DDG-DA caches and exported tasks +================================ + +``DDGDA(..., trusted=True)`` (CLI: ``--trusted=True``) authorizes the necessary +recorder objects and DDG-DA's local handler/``InternalData`` pickle cache reads. +Check ``working_dir``, the configuration directory used for handler caches, any +explicit ``h_path``, and the MLflow store, including all write permissions. +Restricted cache loads refuse executable handlers and ``InternalData`` objects. +Explicit consent uses ordinary pickle with a warning; it neither authenticates +files nor relaxes the global restricted loader. + +Generated tasks keep a lightweight handler-cache **configuration reference**, +including its path and selected ``trusted`` setting, instead of embedding all +market data. Code consuming ``task["dataset"]["kwargs"]["handler"]`` must not +assume it is a ``file://`` string. Treat exported tasks as executable +configurations: their selected cache consent persists when saved and reused, +independently of a newly created workflow's default setting. + +An old exported task may need regeneration using its matching meta-model and +workflow configuration to obtain the current cache-loader reference and policy. +Reloading a recorder task containing a reweighter still requires recorder consent; +that flag alone does not globally permit local caches referenced by the task. +Do not delete experiments or retrain everything just to migrate normal trusted +artifacts. See the +`DDG-DA example `_ +for workflow commands and existing full-flow regression coverage. + +HIST stock-index mapping +======================== + +The bundled mapping is now +``examples/benchmarks/HIST/qlib_csi300_stock_index.json``, preserving all **735 +entries** and their concept-matrix row assignments. Update +``task.model.kwargs.stock_index`` in custom YAML from the old object ``.npy`` path +to this JSON file. The bundled workflow already uses JSON. + +An old saved HIST model also retains its own ``stock_index`` path. After +independently trusting and restoring that model, update ``model.stock_index`` to +the corresponding JSON path before prediction or further fitting; changing YAML +alone does not update a restored instance. Preserve the matching ``stock2concept`` +matrix. It remains a numeric, two-dimensional ``.npy`` file loaded without pickle, +not JSON, and must contain the unknown-stock row **733** (at least **734 rows**). +Every mapped index must also be within the matrix's row bounds. + +For a known-trusted custom mapping, re-export from the original trusted metadata +or producer into a JSON object with instrument strings as keys and non-negative +integer row indices as values (not booleans, floats or strings). Preserve each +instrument's row assignment. Merely renaming a file does not convert it. +Object-pickled ``.npy`` mappings remain refused even with recorder/workflow +``trusted=True``. Do not load an unknown object file to convert it; recover or +regenerate trusted source metadata instead. See the +`HIST example `_. + +High-frequency provider artifact paths +====================================== + +``HighFreqProvider`` confines artifact paths to ``artifact_root``, which defaults +to the current working directory at construction. Choose a dedicated, +access-controlled root and update feature, label and backtest configuration paths +accordingly. Relative paths resolve against this root, not an arbitrary later +working directory. + +All configured and derived paths, including split, per-day and per-stock files, +must remain inside the root after canonical resolution (including symlinks and +``..``). Returned artifact paths are canonical absolute paths; callers should use +them rather than assuming the original relative spelling is preserved. A trust +flag does not bypass path containment. Do not use ``artifact_root="/"`` as a +workaround; move/reconfigure artifacts within the intended dedicated root. + +Containment does not make cache contents safe: these dataset caches still contain +Python pickles and must be independently trusted. Re-create an old serialized +provider that lacks ``artifact_root`` using reviewed configuration and an explicit +root rather than relying on a missing-field fallback. + +Supported data and version compatibility +======================================== + +The restricted loader supports common built-in containers, NumPy arrays/scalars +and pandas ``Series``/``DataFrame`` objects, including typical prediction/label +``MultiIndex`` layouts. Supported reconstruction cases include pickle protocols +4 and 5, NumPy masked arrays, record arrays and record scalars, pandas nullable +integer/float/boolean and Python-backed string arrays, categorical data, +datetime/timedelta data, supported timezone metadata (such as UTC and ``pytz``), +period/interval data and sparse arrays. +``BusinessHour`` and ``CustomBusinessDay`` frequency metadata are also supported, +including custom business hours, weekmasks and holidays, without removing the +frequency or changing the index values. + +Not every dtype or object is supported. Object-dtype cells, custom subclasses, +extension arrays and metadata can introduce executable classes. Arrow-backed +pandas data and ``zoneinfo.ZoneInfo``-backed timezones are not supported by default. +The representation depends on Python, NumPy and pandas versions and dtype +settings. Regenerate supported data in a trusted producer environment rather than +enabling unrestricted loading simply to read predictions or numerical reports. + +Pickle's cross-version limitations still apply. Protocol support does not +guarantee compatibility across Python, NumPy, pandas or model-library versions, +and ``trusted=True`` does not fix missing or renamed classes. Preserve the +producing environment for legacy executable artifacts and test representative +loads before upgrading. + +Example: converting Arrow-backed columns +----------------------------------------- + +Prefer converting the trusted producer's in-memory data before saving it. +The sample below requires pandas 2.x and PyArrow; replace ``original`` with your +actual DataFrame and choose target types from its schema. Do not cast an entire +table to ``float64``: that can lose large integer values or change other columns. +These explicit targets preserve the sample's float width, integer values, +missing-value masks and Python-backed strings. + +.. code-block:: python + :name: arrow-artifact-migration + + import pandas as pd + + original = pd.DataFrame( + { + "score": pd.Series([0.1, None, -0.3], dtype="float32[pyarrow]"), + "count": pd.Series([2**60 + 1, None, 2**60 + 3], dtype="int64[pyarrow]"), + "enabled": pd.Series([True, None, False], dtype="bool[pyarrow]"), + "instrument": pd.Series(["SH600000", None, "SZ000001"], dtype="string[pyarrow]"), + } + ) + compatible = original.astype( + { + "score": "Float32", + "count": "Int64", + "enabled": "boolean", + "instrument": pd.StringDtype(storage="python"), + } + ) + pd.testing.assert_frame_equal(original.isna(), compatible.isna()) + for column in original.columns: + assert original[column].dropna().tolist() == compatible[column].dropna().tolist() + assert compatible.loc[0, "count"] == 2**60 + 1 + assert compatible.loc[2, "count"] == 2**60 + 3 + +This is not a generic converter for Arrow decimals, nested arrays, timestamps or +custom extension types. Review the index, ``attrs`` and object-valued cells too; +converting the listed columns does not necessarily remove every unsupported +object from a real artifact. + +Example: preserving a ZoneInfo index's time semantics +------------------------------------------------------ + +This example requires Python 3.9+ and uses a Qlib-style ``datetime`` / +``instrument`` MultiIndex. It changes the timezone implementation to ``pytz`` +while keeping the same named timezone. The sample crosses a daylight-saving +transition: two different instants have the same local ``01:30`` clock label. + +.. code-block:: python + :name: zoneinfo-artifact-migration + + from zoneinfo import ZoneInfo + + import pandas as pd + import pytz + + times = pd.date_range("2023-11-05 04:30", periods=4, freq="h", tz="UTC") + times = times.tz_convert(ZoneInfo("America/New_York")) + original = pd.DataFrame( + {"score": [0.1, None, 0.3, 0.4]}, + index=pd.MultiIndex.from_arrays( + [times, ["DEMO"] * len(times)], names=["datetime", "instrument"] + ), + ) + compatible = original.tz_convert(pytz.timezone("America/New_York"), level="datetime") + + before = original.index.get_level_values("datetime") + after = compatible.index.get_level_values("datetime") + pd.testing.assert_index_equal(before.tz_convert("UTC"), after.tz_convert("UTC")) + pd.testing.assert_index_equal(before.tz_localize(None), after.tz_localize(None)) + pd.testing.assert_index_equal( + original.index.get_level_values("instrument"), + compatible.index.get_level_values("instrument"), + ) + pd.testing.assert_frame_equal( + original.reset_index(drop=True), compatible.reset_index(drop=True) + ) + +The comparisons check both absolute instants and local clock labels; timezone +removal is used only for the comparison, not on the saved result. Do not simply +strip timezones to bypass a loading error. Converting to UTC preserves instants +but can change local dates, so use it only if your trading-calendar conventions +allow that change. For timezone-aware columns, use ``Series.dt.tz_convert`` with +the intended timezone as well; changing an index does not convert its columns. + +Example: checking and saving the converted artifact +---------------------------------------------------- + +After either conversion above, check the entire result with the restricted +loader before publishing it. The recorder example uses your initialized Qlib +configuration with the built-in MLflow backend and an access-controlled store. +It creates a new run and does not overwrite the original artifact. + +.. code-block:: python + :name: converted-artifact-roundtrip + + import pickle + + import pandas as pd + + from qlib.utils.pickle_utils import restricted_pickle_loads + from qlib.workflow import R + + restored = restricted_pickle_loads(pickle.dumps(compatible, protocol=4)) + pd.testing.assert_frame_equal(compatible, restored) + + with R.start(experiment_name="artifact-migration"): + destination = R.get_recorder() + destination.save_objects(**{"data.migrated.pkl": compatible}) + restored = destination.load_object("data.migrated.pkl") + pd.testing.assert_frame_equal(compatible, restored) + print("Migrated data recorder:", destination.id) + +For real predictions or labels, preserve their original columns, index layout +and business meaning rather than adding the sample columns above. This staging +run contains converted data only, not a complete model/Dataset/task bundle. +Keep the original run backed up and explicitly plan how to publish the verified +data under the ``pred.pkl`` / ``label.pkl`` names expected by your workflow; do +not replace its model recorder with this data-only run. + +If only a legacy pickle remains, recover it only in an independently trusted +producer environment. A deliberately selected MLflow artifact can be read with +``source_rec.load_object("pred.pkl", trusted=True)`` after verifying its writer +and storage, but that unrestricted load can execute code. Never open an unknown +pickle merely to convert it, or automatically retry a restricted-load failure +with consent. Online workflow consent still does not authorize unrestricted +prediction/label reads. + +Custom recorders, loaders and completion callbacks +================================================== + +Custom recorders should implement ``load_object(self, name, *, trusted=False)``, +validate actual boolean consent and enforce restricted loading by default. +Unrestricted loading requires explicit ``trusted=True``. Never ignore the flag or +add an unsafe retry. + +For compatibility, ``R.load_object(name)`` and ``R.load_object(name, trusted=False)`` +call a legacy recorder's ``load_object(name)`` without a new keyword; explicit +``trusted=True`` is forwarded. This preserves default call signatures, **not a +custom backend's security**. A legacy backend using ordinary pickle must implement +the restricted default itself. Adapt its signature before using explicit consent +through ``R``; do not rely on this facade accommodation for direct backend calls. + +An updater's default construction of a custom ``loader_cls`` likewise omits the +new keyword when consent is ``False``. Opted-in construction forwards +``trusted=True``; adapt the loader's constructor, for example +``__init__(self, rec, *, trusted=False)``, and enforce its policy on executable +artifact loads. Data-only reads must remain restricted. + +Custom delayed-training completion callbacks should accept +``end_train_func(rec, experiment_name, *, trusted=False)`` (plus any existing +workflow arguments), validate the boolean, and pass consent only to necessary +executable task reads. Default delayed completion preserves legacy callback calls +without adding the keyword when no consent/override is requested. Opted-in trainers +and explicit per-call overrides forward ``trusted`` to the callback, including in +``DelayTrainerRM`` workers; accepting it without enforcing the policy is not enough. + +Other executable inputs are independent +======================================= + +.. warning:: + + This is a scoped artifact policy, not an all-Qlib sandbox or global + authorization. Apart from DDG-DA's explicitly covered caches, existing local + pickle/model loaders, serialized manager files, handler caches and MongoDB + task stores retain their own trust requirements. YAML/task configurations can + select executable Python components and must also be trusted. + + ``trusted=False`` does not make those inputs safe, and ``trusted=True`` does + not authenticate or globally authorize them. Only open executable inputs + from independently verified sources and access-controlled storage. Never + deserialize unknown files merely to convert or migrate them. diff --git a/docs/start/installation.rst b/docs/start/installation.rst index 451fc0a426d..d48eca659cd 100644 --- a/docs/start/installation.rst +++ b/docs/start/installation.rst @@ -7,6 +7,18 @@ Installation .. currentmodule:: qlib +.. important:: + + **Unreleased upgrade notice for new source builds:** recorder artifact loading + is restricted by default. Reloading executable models, datasets or workflow + objects requires explicit ``trusted=True`` after verifying their source and + storage; supported data-only reads and fresh in-memory training need no opt-in. + Follow :ref:`artifact_loading_migration` before upgrading existing workflows. + Merging into ``main`` affects source installs before a new PyPI release. + This change is unreleased until included in a tagged release, whose upgrade + notes should link to that guide. + + ``Qlib`` Installation ===================== .. note:: diff --git a/examples/README.md b/examples/README.md index 124c34363a2..b881ec47222 100644 --- a/examples/README.md +++ b/examples/README.md @@ -8,3 +8,88 @@ Here is the minimal hardware requirements to run the `workflow_by_code` example. # NOTE The results will slightly vary on different OSs(the variance of annualized return will be less than 2%). The evaluation results in the `README.md` page are from Linux OS. + +# Recorder artifact trust + +Recorder loading is restricted by default. Prediction and label data normally need +no opt-in, but resuming a workflow can require executable model, dataset or task +objects. Only enable `trusted=True` after verifying **both the writer and +the artifact store**, including who can replace files in a shared experiment. +Unrestricted pickle loading can execute code. Do not automatically retry a refused +load with trust enabled. + +See the [artifact loading migration guide](https://qlib.readthedocs.io/en/latest/start/artifact_migration.html) +for supported NumPy/pandas data, version limitations, custom recorder/loader changes +and low-level `load_object(..., trusted=True)` usage. Some notebooks and direct-load +examples already opt in explicitly; their model/task loads assume your own trusted +runs, not arbitrary downloaded or shared artifacts. + +## Online workflows + +The three `online_srv` CLIs expose `--trusted`, defaulting to `False`. +For **your own artifacts in an access-controlled store**, run from `examples/`: + +```bash +python online_srv/update_online_pred.py --trusted=True main +# Subsequent daily updates use the same explicit consent: +python online_srv/update_online_pred.py --trusted=True update_online_pred +``` + +The simulation and rolling-management examples also accept the flag. **Their +`main`/`first_run` methods reset experiment data (and rolling task pools); use +dedicated experiment stores and save anything important before running them.** +With Qlib data and, for rolling management, your MongoDB service configured: + +```bash +python online_srv/online_management_simulate.py --trusted=True --exp_name=my_own_rolling_exp main +python online_srv/rolling_online_management.py --trusted=True --task_url=mongodb://localhost:27017/ first_run +python online_srv/rolling_online_management.py --trusted=True --task_url=mongodb://localhost:27017/ add_strategy +python online_srv/rolling_online_management.py --task_url=mongodb://localhost:27017/ routine +``` + +Consent is passed to each strategy, including newly added strategies, and through +its online tool/updater to executable artifact loads. Prediction, label and +numerical-report reads stay restricted. `OnlineManager` has no trust constructor +or global permission. + +`RollingOnlineExample` creates a `DelayTrainerRM` with the selected consent only +when no trainer is supplied. In Python, configure a caller-supplied delayed trainer +yourself; the example does not change its policy: + +```python +from qlib.model.trainer import DelayTrainerR +from online_srv.rolling_online_management import RollingOnlineExample + +example = RollingOnlineExample( + trainer=DelayTrainerR(trusted=True), + trusted=True, +) +``` + +Ordinary `TrainerR`/`TrainerRM` constructors do not accept `trusted`. +The simulation example leaves those trainers unchanged and configures its strategy. + +Rolling management saves a local `.RollingOnlineExample` pickle. Only restore a +file you independently trust. Restoring it retains its saved strategy, tool and trainer +settings; legacy components without a saved flag default to restricted loading. +An example constructor/CLI flag does not override a subsequently loaded manager. +After reviewing artifact sources, explicitly reconfigure or recreate each +strategy, its `strategy.tool`, and any delayed trainer; changing a strategy's flag +alone does not update its existing tool. `add_strategy` uses the current CLI flag +for **new** strategies only. Remember that this example's `first_run` is destructive. +Ordinary trusted artifacts do not require deleting experiments or full retraining +to migrate; select consent on the actual components that reload them. + +## Other migrations + +- [DDG-DA](benchmarks_dynamic/DDG-DA/README.md#recorder-artifacts-and-local-working-files): + the workflow's opt-in covers necessary recorder and handler/internal-data cache + reads. Verify local `working_dir`, configuration directory and `h_path` contents + as well as MLflow storage. +- [HIST](benchmarks/HIST/README.md#stock-index-mapping-migration): update the + stock-index mapping path to JSON, including on restored models; legacy + object-pickled `.npy` mappings are not accepted. + +These settings cover scoped artifact loads, not every Qlib deserialization API. +Other local model files, handler caches, YAML configurations and task stores +retain their own trust requirements. diff --git a/examples/benchmarks/HIST/README.md b/examples/benchmarks/HIST/README.md index 87f9a35c9b8..512028c8884 100644 --- a/examples/benchmarks/HIST/README.md +++ b/examples/benchmarks/HIST/README.md @@ -1,3 +1,67 @@ # HIST * Code: [https://github.com/Wentao-Xu/HIST](https://github.com/Wentao-Xu/HIST) -* Paper: [HIST: A Graph-based Framework for Stock Trend Forecasting via Mining Concept-Oriented Shared InformationAdaRNN: Adaptive Learning and Forecasting for Time Series](https://arxiv.org/abs/2110.13716). \ No newline at end of file +* Paper: [HIST: A Graph-based Framework for Stock Trend Forecasting via Mining Concept-Oriented Shared InformationAdaRNN: Adaptive Learning and Forecasting for Time Series](https://arxiv.org/abs/2110.13716). + +## Stock-index mapping migration + +The bundled `qlib_csi300_stock_index.npy` object mapping has been replaced by +[`qlib_csi300_stock_index.json`](qlib_csi300_stock_index.json), preserving all +**735 entries**. The [example YAML](workflow_config_hist_Alpha360.yaml) already +uses the new file. Update the same setting in your own workflow YAML: + +```yaml +task: + model: + kwargs: + stock_index: "benchmarks/HIST/qlib_csi300_stock_index.json" +``` + +This path follows the example's convention of running from `examples/`; adjust it +to your working directory. `stock2concept` is a separate numeric matrix and still +uses a two-dimensional `.npy` loaded without pickle. Do not rename or replace that +matrix with the JSON mapping. It must include the unknown-stock row **733** (at +least **734 rows**), and all mapped indices must be within its row bounds. + +A previously saved HIST model retains its own path. After independently trusting +and restoring it, update `model.stock_index` to the matching JSON file before +prediction or further fitting. Changing YAML alone does not update a restored +model. Preserve all 735 bundled entries and their original matrix row assignments. + +For a **known-trusted custom mapping**, re-export it from your original trusted +metadata or producer into a JSON object: + +- Keys must be instrument strings, matching your Qlib dataset. +- Values must be non-negative integers (not strings, floats or booleans) indexing + rows in the corresponding `stock2concept` matrix. +- Preserve each instrument's original row assignment, and verify all indices are + within the matrix's row bounds. +- Point `task.model.kwargs.stock_index` at your new `.json` file and keep the + matching concept matrix. + +If you only have an old object-pickled `.npy`, recover or regenerate the mapping +from the trusted source rather than loading an unknown file to convert it. +Changing the filename extension alone does not convert the contents. HIST +deliberately rejects the old object format because it requires executable pickle +deserialization; there is no fallback or recorder trust flag that re-enables it. + +See the [artifact loading migration guide](https://qlib.readthedocs.io/en/latest/start/artifact_migration.html) +for the separate policy on model/dataset artifacts. The JSON mapping change does +not make pre-existing local model checkpoints or other pickle loaders safe; those +inputs still require independent trust. + +## Full workflow regression + +From the repository root, with the test and model dependencies installed: + +```bash +python -m pytest tests/model/test_hist_workflow.py -m slow -q +``` + +This offline CPU regression runs real Alpha360/DatasetH preparation, one HIST +training epoch, signal analysis and a six-day TopkDropout backtest in an isolated +MLflow store. It checks 48 prediction/label rows, actual optimizer updates, +finite reports and trading activity. Saved model/dataset objects are refused by +default; explicitly trusted reloads reproduce predictions exactly, including in +a fresh Python process. Metadata and concept fixtures are local JSON and numeric +NumPy files, with no downloads. This is functional integration coverage, not a +paper-scale model-quality benchmark. \ No newline at end of file diff --git a/examples/benchmarks/HIST/qlib_csi300_stock_index.json b/examples/benchmarks/HIST/qlib_csi300_stock_index.json new file mode 100644 index 00000000000..065dd8c660c --- /dev/null +++ b/examples/benchmarks/HIST/qlib_csi300_stock_index.json @@ -0,0 +1,737 @@ +{ + "SH600000": 0, + "SH600001": 1, + "SH600003": 307, + "SH600004": 2, + "SH600005": 3, + "SH600006": 4, + "SH600007": 5, + "SH600008": 6, + "SH600009": 7, + "SH600010": 8, + "SH600011": 9, + "SH600012": 10, + "SH600015": 11, + "SH600016": 12, + "SH600017": 308, + "SH600018": 13, + "SH600019": 14, + "SH600020": 15, + "SH600021": 16, + "SH600022": 17, + "SH600023": 538, + "SH600025": 677, + "SH600026": 18, + "SH600027": 19, + "SH600028": 20, + "SH600029": 21, + "SH600030": 22, + "SH600031": 23, + "SH600033": 24, + "SH600035": 25, + "SH600036": 26, + "SH600037": 27, + "SH600038": 556, + "SH600048": 28, + "SH600050": 29, + "SH600058": 30, + "SH600060": 31, + "SH600061": 600, + "SH600062": 32, + "SH600066": 33, + "SH600068": 34, + "SH600073": 35, + "SH600074": 601, + "SH600078": 36, + "SH600079": 526, + "SH600085": 37, + "SH600087": 38, + "SH600088": 39, + "SH600089": 40, + "SH600096": 41, + "SH600098": 42, + "SH600100": 43, + "SH600102": 44, + "SH600104": 45, + "SH600108": 46, + "SH600109": 338, + "SH600110": 47, + "SH600111": 339, + "SH600115": 340, + "SH600117": 48, + "SH600118": 49, + "SH600121": 50, + "SH600123": 51, + "SH600125": 52, + "SH600132": 53, + "SH600143": 54, + "SH600150": 55, + "SH600151": 56, + "SH600153": 57, + "SH600157": 508, + "SH600158": 341, + "SH600160": 58, + "SH600161": 59, + "SH600166": 412, + "SH600169": 363, + "SH600170": 60, + "SH600171": 61, + "SH600176": 391, + "SH600177": 62, + "SH600183": 63, + "SH600188": 64, + "SH600190": 65, + "SH600196": 66, + "SH600198": 67, + "SH600200": 68, + "SH600205": 69, + "SH600208": 309, + "SH600210": 70, + "SH600216": 392, + "SH600219": 310, + "SH600220": 71, + "SH600221": 72, + "SH600231": 73, + "SH600233": 636, + "SH600236": 74, + "SH600239": 413, + "SH600246": 414, + "SH600251": 380, + "SH600252": 480, + "SH600256": 75, + "SH600259": 481, + "SH600266": 76, + "SH600267": 77, + "SH600269": 78, + "SH600270": 79, + "SH600271": 80, + "SH600276": 393, + "SH600277": 539, + "SH600282": 81, + "SH600296": 82, + "SH600297": 613, + "SH600299": 83, + "SH600300": 427, + "SH600307": 84, + "SH600308": 85, + "SH600309": 86, + "SH600312": 87, + "SH600315": 495, + "SH600316": 88, + "SH600317": 577, + "SH600320": 89, + "SH600325": 342, + "SH600331": 90, + "SH600332": 91, + "SH600333": 92, + "SH600339": 678, + "SH600340": 509, + "SH600346": 679, + "SH600348": 93, + "SH600350": 94, + "SH600352": 394, + "SH600357": 95, + "SH600361": 96, + "SH600362": 97, + "SH600369": 415, + "SH600372": 482, + "SH600373": 557, + "SH600376": 364, + "SH600377": 98, + "SH600380": 337, + "SH600383": 99, + "SH600390": 661, + "SH600395": 395, + "SH600398": 558, + "SH600403": 518, + "SH600406": 442, + "SH600410": 100, + "SH600415": 101, + "SH600418": 102, + "SH600426": 381, + "SH600428": 103, + "SH600432": 104, + "SH600436": 519, + "SH600438": 680, + "SH600446": 602, + "SH600456": 105, + "SH600460": 106, + "SH600472": 107, + "SH600481": 443, + "SH600482": 614, + "SH600485": 559, + "SH600487": 681, + "SH600489": 108, + "SH600497": 109, + "SH600498": 110, + "SH600500": 111, + "SH600501": 303, + "SH600508": 112, + "SH600516": 396, + "SH600517": 416, + "SH600518": 397, + "SH600519": 113, + "SH600521": 114, + "SH600522": 637, + "SH600528": 311, + "SH600535": 115, + "SH600546": 444, + "SH600547": 116, + "SH600548": 117, + "SH600549": 118, + "SH600550": 119, + "SH600566": 693, + "SH600569": 120, + "SH600570": 560, + "SH600578": 561, + "SH600582": 398, + "SH600583": 121, + "SH600585": 122, + "SH600588": 123, + "SH600591": 124, + "SH600595": 343, + "SH600596": 125, + "SH600597": 126, + "SH600598": 127, + "SH600600": 128, + "SH600601": 129, + "SH600602": 130, + "SH600606": 603, + "SH600611": 344, + "SH600616": 131, + "SH600627": 132, + "SH600628": 133, + "SH600631": 134, + "SH600633": 540, + "SH600635": 135, + "SH600637": 136, + "SH600638": 345, + "SH600639": 137, + "SH600642": 138, + "SH600643": 139, + "SH600648": 417, + "SH600649": 140, + "SH600653": 141, + "SH600654": 615, + "SH600655": 142, + "SH600657": 418, + "SH600660": 143, + "SH600662": 144, + "SH600663": 145, + "SH600664": 382, + "SH600666": 604, + "SH600674": 383, + "SH600675": 146, + "SH600682": 638, + "SH600685": 312, + "SH600688": 147, + "SH600690": 148, + "SH600694": 149, + "SH600703": 445, + "SH600704": 605, + "SH600705": 527, + "SH600717": 150, + "SH600718": 151, + "SH600724": 152, + "SH600733": 711, + "SH600737": 384, + "SH600739": 153, + "SH600741": 154, + "SH600747": 155, + "SH600748": 156, + "SH600754": 157, + "SH600760": 694, + "SH600761": 158, + "SH600770": 159, + "SH600779": 160, + "SH600780": 161, + "SH600782": 399, + "SH600783": 483, + "SH600786": 162, + "SH600787": 163, + "SH600795": 164, + "SH600797": 165, + "SH600804": 365, + "SH600808": 166, + "SH600809": 167, + "SH600811": 168, + "SH600812": 169, + "SH600816": 400, + "SH600820": 170, + "SH600823": 446, + "SH600827": 479, + "SH600832": 171, + "SH600834": 172, + "SH600835": 173, + "SH600837": 313, + "SH600839": 174, + "SH600848": 723, + "SH600851": 175, + "SH600859": 314, + "SH600863": 176, + "SH600867": 541, + "SH600868": 177, + "SH600871": 178, + "SH600873": 484, + "SH600874": 179, + "SH600875": 180, + "SH600879": 181, + "SH600880": 542, + "SH600881": 182, + "SH600884": 183, + "SH600886": 184, + "SH600887": 185, + "SH600893": 447, + "SH600894": 186, + "SH600895": 187, + "SH600900": 188, + "SH600909": 639, + "SH600919": 640, + "SH600926": 641, + "SH600928": 724, + "SH600958": 578, + "SH600959": 586, + "SH600961": 315, + "SH600968": 725, + "SH600970": 189, + "SH600971": 462, + "SH600977": 642, + "SH600978": 346, + "SH600989": 726, + "SH600997": 190, + "SH600998": 463, + "SH600999": 428, + "SH601001": 191, + "SH601002": 316, + "SH601003": 347, + "SH601005": 348, + "SH601006": 192, + "SH601009": 349, + "SH601012": 662, + "SH601016": 587, + "SH601018": 464, + "SH601021": 579, + "SH601066": 695, + "SH601088": 335, + "SH601098": 465, + "SH601099": 419, + "SH601101": 448, + "SH601106": 449, + "SH601107": 420, + "SH601108": 682, + "SH601111": 193, + "SH601117": 429, + "SH601118": 466, + "SH601127": 616, + "SH601138": 696, + "SH601139": 430, + "SH601155": 617, + "SH601158": 450, + "SH601162": 712, + "SH601163": 643, + "SH601166": 301, + "SH601168": 350, + "SH601169": 351, + "SH601179": 431, + "SH601186": 385, + "SH601198": 588, + "SH601211": 589, + "SH601212": 663, + "SH601216": 485, + "SH601225": 543, + "SH601228": 664, + "SH601229": 644, + "SH601231": 528, + "SH601233": 486, + "SH601236": 727, + "SH601238": 510, + "SH601258": 487, + "SH601268": 451, + "SH601288": 441, + "SH601298": 713, + "SH601299": 432, + "SH601318": 302, + "SH601319": 714, + "SH601328": 306, + "SH601333": 317, + "SH601336": 496, + "SH601360": 683, + "SH601369": 452, + "SH601375": 645, + "SH601377": 467, + "SH601390": 366, + "SH601398": 194, + "SH601519": 468, + "SH601555": 497, + "SH601558": 469, + "SH601566": 498, + "SH601577": 715, + "SH601588": 318, + "SH601600": 304, + "SH601601": 367, + "SH601607": 195, + "SH601608": 590, + "SH601611": 618, + "SH601618": 421, + "SH601628": 300, + "SH601633": 499, + "SH601666": 319, + "SH601668": 422, + "SH601669": 500, + "SH601688": 433, + "SH601698": 728, + "SH601699": 320, + "SH601717": 470, + "SH601718": 471, + "SH601727": 401, + "SH601766": 402, + "SH601788": 434, + "SH601800": 511, + "SH601801": 435, + "SH601808": 352, + "SH601818": 453, + "SH601828": 684, + "SH601838": 685, + "SH601857": 336, + "SH601866": 368, + "SH601872": 321, + "SH601877": 436, + "SH601878": 665, + "SH601881": 646, + "SH601888": 437, + "SH601898": 369, + "SH601899": 386, + "SH601901": 501, + "SH601918": 370, + "SH601919": 353, + "SH601928": 502, + "SH601929": 544, + "SH601933": 472, + "SH601939": 371, + "SH601958": 387, + "SH601966": 647, + "SH601969": 580, + "SH601985": 591, + "SH601988": 196, + "SH601989": 438, + "SH601991": 322, + "SH601992": 488, + "SH601997": 648, + "SH601998": 305, + "SH603000": 529, + "SH603019": 716, + "SH603156": 697, + "SH603160": 649, + "SH603259": 698, + "SH603260": 686, + "SH603288": 562, + "SH603501": 729, + "SH603699": 545, + "SH603799": 666, + "SH603833": 667, + "SH603858": 650, + "SH603885": 592, + "SH603899": 730, + "SH603986": 699, + "SH603993": 520, + "SH999998": 411, + "SH999999": 197, + "SZ000001": 198, + "SZ000002": 199, + "SZ000008": 619, + "SZ000009": 200, + "SZ000012": 201, + "SZ000021": 202, + "SZ000024": 203, + "SZ000027": 204, + "SZ000029": 205, + "SZ000031": 206, + "SZ000036": 207, + "SZ000039": 208, + "SZ000046": 323, + "SZ000059": 209, + "SZ000060": 210, + "SZ000061": 211, + "SZ000063": 212, + "SZ000066": 213, + "SZ000068": 214, + "SZ000069": 215, + "SZ000088": 216, + "SZ000089": 217, + "SZ000099": 218, + "SZ000100": 219, + "SZ000156": 521, + "SZ000157": 220, + "SZ000166": 574, + "SZ000301": 362, + "SZ000333": 525, + "SZ000338": 354, + "SZ000400": 221, + "SZ000401": 222, + "SZ000402": 223, + "SZ000408": 700, + "SZ000410": 224, + "SZ000413": 546, + "SZ000415": 593, + "SZ000422": 225, + "SZ000423": 226, + "SZ000425": 227, + "SZ000488": 228, + "SZ000503": 229, + "SZ000520": 230, + "SZ000527": 231, + "SZ000528": 232, + "SZ000536": 473, + "SZ000538": 233, + "SZ000539": 234, + "SZ000540": 594, + "SZ000541": 235, + "SZ000543": 355, + "SZ000550": 236, + "SZ000553": 701, + "SZ000555": 620, + "SZ000559": 237, + "SZ000562": 238, + "SZ000568": 239, + "SZ000572": 356, + "SZ000581": 240, + "SZ000596": 512, + "SZ000598": 530, + "SZ000601": 241, + "SZ000612": 324, + "SZ000617": 242, + "SZ000623": 243, + "SZ000625": 244, + "SZ000627": 621, + "SZ000629": 245, + "SZ000630": 246, + "SZ000631": 423, + "SZ000636": 247, + "SZ000651": 248, + "SZ000652": 249, + "SZ000656": 522, + "SZ000661": 702, + "SZ000667": 372, + "SZ000671": 622, + "SZ000680": 250, + "SZ000682": 251, + "SZ000685": 403, + "SZ000686": 373, + "SZ000690": 325, + "SZ000698": 252, + "SZ000703": 503, + "SZ000708": 253, + "SZ000709": 254, + "SZ000712": 581, + "SZ000717": 255, + "SZ000718": 357, + "SZ000723": 668, + "SZ000725": 256, + "SZ000728": 374, + "SZ000729": 257, + "SZ000738": 575, + "SZ000750": 513, + "SZ000751": 258, + "SZ000755": 375, + "SZ000758": 259, + "SZ000761": 260, + "SZ000767": 261, + "SZ000768": 262, + "SZ000776": 454, + "SZ000778": 263, + "SZ000780": 424, + "SZ000783": 376, + "SZ000786": 264, + "SZ000792": 265, + "SZ000793": 266, + "SZ000800": 267, + "SZ000807": 268, + "SZ000822": 269, + "SZ000825": 270, + "SZ000826": 531, + "SZ000828": 271, + "SZ000829": 326, + "SZ000831": 532, + "SZ000839": 272, + "SZ000858": 273, + "SZ000869": 274, + "SZ000875": 275, + "SZ000876": 327, + "SZ000878": 276, + "SZ000883": 514, + "SZ000895": 277, + "SZ000897": 278, + "SZ000898": 279, + "SZ000900": 280, + "SZ000912": 281, + "SZ000917": 282, + "SZ000927": 283, + "SZ000930": 284, + "SZ000932": 285, + "SZ000933": 286, + "SZ000937": 287, + "SZ000938": 623, + "SZ000939": 288, + "SZ000951": 358, + "SZ000959": 328, + "SZ000960": 289, + "SZ000961": 439, + "SZ000962": 290, + "SZ000963": 523, + "SZ000968": 388, + "SZ000969": 291, + "SZ000970": 292, + "SZ000977": 606, + "SZ000983": 293, + "SZ000997": 294, + "SZ000999": 404, + "SZ001965": 687, + "SZ001979": 599, + "SZ002001": 389, + "SZ002007": 425, + "SZ002008": 295, + "SZ002010": 717, + "SZ002024": 296, + "SZ002025": 297, + "SZ002027": 607, + "SZ002028": 405, + "SZ002032": 703, + "SZ002038": 406, + "SZ002044": 651, + "SZ002049": 624, + "SZ002050": 688, + "SZ002051": 298, + "SZ002052": 299, + "SZ002065": 533, + "SZ002069": 329, + "SZ002073": 474, + "SZ002074": 625, + "SZ002078": 330, + "SZ002081": 504, + "SZ002083": 331, + "SZ002085": 626, + "SZ002092": 455, + "SZ002097": 332, + "SZ002106": 333, + "SZ002110": 334, + "SZ002120": 704, + "SZ002122": 359, + "SZ002128": 360, + "SZ002129": 534, + "SZ002131": 627, + "SZ002142": 361, + "SZ002146": 377, + "SZ002152": 390, + "SZ002153": 563, + "SZ002155": 378, + "SZ002174": 628, + "SZ002179": 705, + "SZ002183": 608, + "SZ002194": 407, + "SZ002195": 595, + "SZ002202": 379, + "SZ002230": 535, + "SZ002236": 515, + "SZ002241": 505, + "SZ002242": 408, + "SZ002244": 409, + "SZ002252": 547, + "SZ002269": 410, + "SZ002271": 706, + "SZ002275": 426, + "SZ002292": 548, + "SZ002294": 456, + "SZ002299": 489, + "SZ002304": 440, + "SZ002310": 475, + "SZ002311": 707, + "SZ002344": 490, + "SZ002352": 652, + "SZ002353": 506, + "SZ002375": 516, + "SZ002378": 491, + "SZ002385": 457, + "SZ002399": 458, + "SZ002400": 549, + "SZ002405": 459, + "SZ002410": 550, + "SZ002411": 653, + "SZ002415": 460, + "SZ002416": 551, + "SZ002422": 461, + "SZ002424": 609, + "SZ002426": 629, + "SZ002429": 552, + "SZ002431": 492, + "SZ002450": 524, + "SZ002456": 536, + "SZ002460": 669, + "SZ002465": 553, + "SZ002466": 630, + "SZ002468": 670, + "SZ002470": 554, + "SZ002475": 555, + "SZ002493": 476, + "SZ002498": 477, + "SZ002500": 478, + "SZ002508": 654, + "SZ002555": 655, + "SZ002558": 656, + "SZ002568": 610, + "SZ002570": 507, + "SZ002572": 671, + "SZ002594": 493, + "SZ002601": 672, + "SZ002602": 657, + "SZ002603": 494, + "SZ002607": 731, + "SZ002608": 673, + "SZ002624": 674, + "SZ002625": 689, + "SZ002653": 537, + "SZ002673": 517, + "SZ002714": 631, + "SZ002736": 582, + "SZ002739": 596, + "SZ002773": 708, + "SZ002797": 632, + "SZ002831": 658, + "SZ002839": 659, + "SZ002841": 660, + "SZ002916": 732, + "SZ002925": 690, + "SZ002938": 718, + "SZ002939": 719, + "SZ002945": 720, + "SZ002958": 733, + "SZ300002": 583, + "SZ300003": 576, + "SZ300015": 564, + "SZ300017": 565, + "SZ300024": 566, + "SZ300027": 567, + "SZ300033": 633, + "SZ300058": 568, + "SZ300059": 584, + "SZ300070": 569, + "SZ300072": 634, + "SZ300085": 611, + "SZ300104": 585, + "SZ300122": 675, + "SZ300124": 570, + "SZ300133": 571, + "SZ300136": 676, + "SZ300142": 709, + "SZ300144": 597, + "SZ300146": 572, + "SZ300168": 612, + "SZ300182": 635, + "SZ300251": 573, + "SZ300296": 710, + "SZ300315": 598, + "SZ300347": 734, + "SZ300408": 691, + "SZ300413": 721, + "SZ300433": 692, + "SZ300498": 722 +} diff --git a/examples/benchmarks/HIST/qlib_csi300_stock_index.npy b/examples/benchmarks/HIST/qlib_csi300_stock_index.npy deleted file mode 100644 index 36db027df62..00000000000 Binary files a/examples/benchmarks/HIST/qlib_csi300_stock_index.npy and /dev/null differ diff --git a/examples/benchmarks/HIST/workflow_config_hist_Alpha360.yaml b/examples/benchmarks/HIST/workflow_config_hist_Alpha360.yaml index cd50b338793..42370fd31fa 100644 --- a/examples/benchmarks/HIST/workflow_config_hist_Alpha360.yaml +++ b/examples/benchmarks/HIST/workflow_config_hist_Alpha360.yaml @@ -59,7 +59,7 @@ task: base_model: LSTM model_path: "benchmarks/LSTM/model_lstm_csi300.pkl" stock2concept: "benchmarks/HIST/qlib_csi300_stock2concept.npy" - stock_index: "benchmarks/HIST/qlib_csi300_stock_index.npy" + stock_index: "benchmarks/HIST/qlib_csi300_stock_index.json" GPU: 0 dataset: class: DatasetH diff --git a/examples/benchmarks_dynamic/DDG-DA/README.md b/examples/benchmarks_dynamic/DDG-DA/README.md index ac4349d91ec..056fcc28419 100644 --- a/examples/benchmarks_dynamic/DDG-DA/README.md +++ b/examples/benchmarks_dynamic/DDG-DA/README.md @@ -14,16 +14,78 @@ The data in the paper are private. So we conduct experiments on Qlib's public da Though the dataset is different, the conclusion remains the same. By applying `DDG-DA`, users can see rising trends at the test phase both in the proxy models' ICs and the performances of the forecasting models. # Run the Code -Users can try `DDG-DA` by running the following command: +After verifying that the recorder artifacts and local working files are your own +and cannot be replaced by untrusted writers (see below), run from this directory: ```bash - python workflow.py run + python workflow.py --trusted=True run ``` The default forecasting models are `Linear`. Users can choose other forecasting models by changing the `forecast_model` parameter when `DDG-DA` initializes. For example, users can try `LightGBM` forecasting models by running the following command: ```bash - python workflow.py --conf_path=../workflow_config_lightgbm_Alpha158.yaml run + python workflow.py --trusted=True --conf_path=../baseline/workflow_config_lightgbm_Alpha158.yaml run ``` +## Recorder artifacts and local working files + +`workflow.py` exposes the `DDGDA` workflow through the `DDGDABench` Fire entry +point. Its `trusted` option defaults to `False`. Set it explicitly only +for artifacts and caches from a verified writer in access-controlled MLflow and +local storage: unrestricted pickle loading can execute code. Creating a run +yourself is not enough if someone else can overwrite its files. + +The option covers recorder-backed executable meta-model/task loading, including +`InternalData.setup`, and DDG-DA's local handler/internal-data pickle cache reads; +prediction, label and numerical-report artifact loads remain restricted. +Lower-level callers can also pass `trusted=True` to `MetaDatasetDS` or +`InternalData.setup` for recorder task reads, not as a global local-cache grant. +A refused load is not a reason to retry automatically with trust enabled. +See the [artifact loading migration guide](https://qlib.readthedocs.io/en/latest/start/artifact_migration.html) +for data compatibility and migration details. + +This example also saves and reuses **local pickle files in `working_dir`**, which +the benchmark sets to this directory. Handler/internal-data caches default to +restricted loading, which refuses executable objects such as `Alpha158` or +`InternalData`. The explicit opt-in allows these caches to be restored with +ordinary pickle and emits a warning; it does not authenticate their contents. +Protect `working_dir`, the configuration directory (also used for handler +caching), and any supplied `h_path` from untrusted writes. Do not copy unknown +cached handlers, meta-information or models into them. There is no automatic +unsafe retry and no change to the global restricted loader. Other pickle APIs +and workflow YAML retain their own trust requirements; only use trusted +configurations and files. + +Generated tasks keep lightweight handler-cache configuration references, including +the path and chosen `trusted` setting, rather than embedding the full market data. +Do not assume `task["dataset"]["kwargs"]["handler"]` is a `file://` string. +Treat saved task configurations as executable inputs; reusing an opted-in task +also reuses that local-cache consent independently of a new workflow's default. +An old exported task may need regeneration with its matching meta-model and +configuration to obtain the current cache reference and selected policy. Loading +a saved task containing a reweighter through a recorder still requires explicit +recorder consent; that flag alone does not globally authorize local cache reads. + +The Makefile's `clean` target deletes local pickle files and `mlruns`; preserve any +results you need before using it. Deleting experiments or running a full retrain +is not necessary to migrate ordinary trusted artifacts. + +## Full workflow regression + +From the repository root, with the test and model dependencies installed: + +```bash +python -m pytest tests/rolling_tests/test_ddgda.py -m slow -q +``` + +The offline regression uses deterministic local daily market data and an isolated +MLflow store. Both linear and LightGBM similarity models run through feature +selection, seven similarity-training windows, daily rank IC, cache restoration, +30-epoch meta-training, inferred time weights, two rolling training windows, +prediction/label collection and a 40-day portfolio backtest. It checks default +refusal, explicit authorization, restored predictions, delayed replay of both +saved rolling tasks and non-empty numerical results without replacing workflow +stages with mocks. These small integration +cases verify functionality, not paper-scale performance or investment returns. + # Results The results of related methods in Qlib's public dataset can be found [here](../) diff --git a/examples/benchmarks_dynamic/DDG-DA/vis_data.py b/examples/benchmarks_dynamic/DDG-DA/vis_data.py index 18e63a36a54..e5b5fe39afb 100644 --- a/examples/benchmarks_dynamic/DDG-DA/vis_data.py +++ b/examples/benchmarks_dynamic/DDG-DA/vis_data.py @@ -43,7 +43,8 @@ exp = R.get_exp(experiment_name="DDG-DA") meta_rec = exp.list_recorders(rtype="list", max_results=1)[0] -meta_m = meta_rec.load_object("model") +# Only load executable models/tasks from experiments whose storage you trust. +meta_m = meta_rec.load_object("model", trusted=True) pd.DataFrame(meta_m.tn.twm.linear.weight.detach().numpy()).T[0].plot() @@ -86,10 +87,10 @@ def show_linear_weight(exp): coef_df = {} for r in exp.list_recorders("list"): - t = r.load_object("task") + t = r.load_object("task", trusted=True) if None in t["dataset"]["kwargs"]["segments"]["test"]: continue - m = r.load_object("params.pkl") + m = r.load_object("params.pkl", trusted=True) coef_df[t["dataset"]["kwargs"]["segments"]["test"]] = pd.Series(m.coef_) coef_df = pd.concat(coef_df) diff --git a/examples/benchmarks_dynamic/DDG-DA/workflow.py b/examples/benchmarks_dynamic/DDG-DA/workflow.py index 8209e0e9061..988b824ba72 100644 --- a/examples/benchmarks_dynamic/DDG-DA/workflow.py +++ b/examples/benchmarks_dynamic/DDG-DA/workflow.py @@ -23,10 +23,12 @@ class DDGDABench(DDGDA): DEFAULT_CONF = CONF_LIST[0] # Linear by default due to efficiency - def __init__(self, conf_path: Union[str, Path] = DEFAULT_CONF, horizon=20, **kwargs) -> None: + def __init__( + self, conf_path: Union[str, Path] = DEFAULT_CONF, horizon=20, *, trusted: bool = False, **kwargs + ) -> None: # This code is for being compatible with the previous old code conf_path = Path(conf_path) - super().__init__(conf_path=conf_path, horizon=horizon, working_dir=DIRNAME, **kwargs) + super().__init__(conf_path=conf_path, horizon=horizon, working_dir=DIRNAME, trusted=trusted, **kwargs) for f in self.CONF_LIST: if conf_path.samefile(f): diff --git a/examples/model_rolling/task_manager_rolling.py b/examples/model_rolling/task_manager_rolling.py index 2fb7c85b565..b98399f11ef 100644 --- a/examples/model_rolling/task_manager_rolling.py +++ b/examples/model_rolling/task_manager_rolling.py @@ -83,7 +83,7 @@ def task_collecting(self): print("========== task_collecting ==========") def rec_key(recorder): - task_config = recorder.load_object("task") + task_config = recorder.load_object("task", trusted=True) model_key = task_config["model"]["class"] rolling_key = task_config["dataset"]["kwargs"]["segments"]["test"] return model_key, rolling_key diff --git a/examples/online_srv/online_management_simulate.py b/examples/online_srv/online_management_simulate.py index dccc56b6828..dbcd7b0278e 100644 --- a/examples/online_srv/online_management_simulate.py +++ b/examples/online_srv/online_management_simulate.py @@ -35,6 +35,8 @@ def __init__( end_time="2018-10-31", tasks=None, trainer="TrainerR", + *, + trusted=False, ): """ Init OnlineManagerExample. @@ -50,6 +52,7 @@ def __init__( start_time (str, optional): the start time of simulating. Defaults to "2018-09-10". end_time (str, optional): the end time of simulating. Defaults to "2018-10-31". tasks (dict or list[dict]): a set of the task config waiting for rolling and training + trusted (bool): allow executable recorder artifacts from a trusted writer and store. """ if tasks is None: tasks = [CSI100_RECORD_XGBOOST_TASK_CONFIG_ONLINE, CSI100_RECORD_LGB_TASK_CONFIG_ONLINE] @@ -73,7 +76,12 @@ def __init__( # TODO: support all the trainers: TrainerR, TrainerRM, DelayTrainerR raise NotImplementedError(f"This type of input is not supported") self.rolling_online_manager = OnlineManager( - RollingStrategy(exp_name, task_template=tasks, rolling_gen=self.rolling_gen), + RollingStrategy( + exp_name, + task_template=tasks, + rolling_gen=self.rolling_gen, + trusted=trusted, + ), trainer=self.trainer, begin_time=self.start_time, ) @@ -133,5 +141,6 @@ def worker(self): if __name__ == "__main__": ## to run all workflow automatically with your own parameters, use the command below - # python online_management_simulate.py main --experiment_name="your_exp_name" --rolling_step=60 + # Only opt in for your own artifacts in an access-controlled store. main resets the experiment. + # python online_management_simulate.py --trusted=True --exp_name="your_exp_name" --rolling_step=60 main fire.Fire(OnlineSimulationExample) diff --git a/examples/online_srv/rolling_online_management.py b/examples/online_srv/rolling_online_management.py index 6abbbfb0e85..226090b907a 100644 --- a/examples/online_srv/rolling_online_management.py +++ b/examples/online_srv/rolling_online_management.py @@ -14,6 +14,7 @@ import fire import qlib from qlib.model.trainer import DelayTrainerR, DelayTrainerRM, TrainerR, TrainerRM, end_task_train, task_train +from qlib.utils.pickle_utils import validate_trusted from qlib.workflow import R from qlib.workflow.online.strategy import RollingStrategy from qlib.workflow.task.gen import RollingGen @@ -23,16 +24,20 @@ class RollingOnlineExample: + trusted = False + def __init__( self, provider_uri="~/.qlib/qlib_data/cn_data", region="cn", - trainer=DelayTrainerRM(), # you can choose from TrainerR, TrainerRM, DelayTrainerR, DelayTrainerRM + trainer=None, # defaults to DelayTrainerRM; a supplied trainer keeps its own trust policy task_url="mongodb://10.0.0.4:27017/", # not necessary when using TrainerR or DelayTrainerR task_db_name="rolling_db", # not necessary when using TrainerR or DelayTrainerR rolling_step=550, tasks=None, add_tasks=None, + *, + trusted=False, ): if add_tasks is None: add_tasks = [CSI100_RECORD_LGB_TASK_CONFIG_ROLLING] @@ -46,6 +51,7 @@ def __init__( self.tasks = tasks self.add_tasks = add_tasks self.rolling_step = rolling_step + self.trusted = validate_trusted(trusted) strategies = [] for task in tasks: name_id = task["model"]["class"] # NOTE: Assumption: The model class can specify only one strategy @@ -54,9 +60,10 @@ def __init__( name_id, task, RollingGen(step=rolling_step, rtype=RollingGen.ROLL_SD), + trusted=self.trusted, ) ) - self.trainer = trainer + self.trainer = DelayTrainerRM(trusted=trusted) if trainer is None else trainer self.rolling_online_manager = OnlineManager(strategies, trainer=self.trainer) _ROLLING_MANAGER_PATH = ( @@ -119,6 +126,7 @@ def add_strategy(self): name_id, task, RollingGen(step=self.rolling_step, rtype=RollingGen.ROLL_SD), + trusted=self.trusted, ) ) self.rolling_online_manager.add_strategy(strategies=strategies) @@ -134,11 +142,16 @@ def main(self): if __name__ == "__main__": ####### to train the first version's models, use the command below - # python rolling_online_management.py first_run + # Only opt in for artifacts whose writer and store you trust. first_run resets the experiments. + # python rolling_online_management.py --trusted=True first_run ####### to update the models and predictions after the trading time, use the command below + # The saved manager is a separately trusted local pickle and retains its original trust settings. # python rolling_online_management.py routine + ####### to give newly added strategies the same explicit consent + # python rolling_online_management.py --trusted=True add_strategy + ####### to define your own parameters, use `--` - # python rolling_online_management.py first_run --exp_name='your_exp_name' --rolling_step=40 + # python rolling_online_management.py --trusted=True --rolling_step=40 first_run fire.Fire(RollingOnlineExample) diff --git a/examples/online_srv/update_online_pred.py b/examples/online_srv/update_online_pred.py index 369a6d84df2..a562e20acb7 100644 --- a/examples/online_srv/update_online_pred.py +++ b/examples/online_srv/update_online_pred.py @@ -26,11 +26,17 @@ class UpdatePredExample: def __init__( - self, provider_uri="~/.qlib/qlib_data/cn_data", region=REG_CN, experiment_name="online_srv", task_config=task + self, + provider_uri="~/.qlib/qlib_data/cn_data", + region=REG_CN, + experiment_name="online_srv", + task_config=task, + *, + trusted=False, ): qlib.init(provider_uri=provider_uri, region=region) self.experiment_name = experiment_name - self.online_tool = OnlineToolR(self.experiment_name) + self.online_tool = OnlineToolR(self.experiment_name, trusted=trusted) self.task_config = task_config def first_train(self): @@ -49,7 +55,8 @@ def main(self): ## to train a model and set it to online model, use the command below # python update_online_pred.py first_train ## to update online predictions once a day, use the command below - # python update_online_pred.py update_online_pred + # Only opt in for artifacts whose writer and store you trust. + # python update_online_pred.py --trusted=True update_online_pred ## to see the whole process with your own parameters, use the command below - # python update_online_pred.py main --experiment_name="your_exp_name" + # python update_online_pred.py --trusted=True --experiment_name="your_exp_name" main fire.Fire(UpdatePredExample) diff --git a/examples/tutorial/detailed_workflow.ipynb b/examples/tutorial/detailed_workflow.ipynb index 2d70c23f4a7..b1b85b801b8 100644 --- a/examples/tutorial/detailed_workflow.ipynb +++ b/examples/tutorial/detailed_workflow.ipynb @@ -1096,7 +1096,7 @@ "# load previous results\n", "pred_df = recorder.load_object(\"pred.pkl\")\n", "report_normal_df = recorder.load_object(\"portfolio_analysis/report_normal_1day.pkl\")\n", - "positions = recorder.load_object(\"portfolio_analysis/positions_normal_1day.pkl\")\n", + "positions = recorder.load_object(\"portfolio_analysis/positions_normal_1day.pkl\", trusted=True)\n", "analysis_df = recorder.load_object(\"portfolio_analysis/port_analysis_1day.pkl\")" ] }, @@ -1108,7 +1108,7 @@ "outputs": [], "source": [ "# Previous Model can be loaded. but it is not used.\n", - "loaded_model = recorder.load_object(\"trained_model\")\n", + "loaded_model = recorder.load_object(\"trained_model\", trusted=True)\n", "loaded_model" ] }, diff --git a/examples/workflow_by_code.ipynb b/examples/workflow_by_code.ipynb index 2cf3f1fa2a9..bd3170c9e7d 100644 --- a/examples/workflow_by_code.ipynb +++ b/examples/workflow_by_code.ipynb @@ -229,7 +229,7 @@ "# backtest and analysis\n", "with R.start(experiment_name=\"backtest_analysis\"):\n", " recorder = R.get_recorder(recorder_id=rid, experiment_name=\"train_model\")\n", - " model = recorder.load_object(\"trained_model\")\n", + " model = recorder.load_object(\"trained_model\", trusted=True)\n", "\n", " # prediction\n", " recorder = R.get_recorder()\n", @@ -262,7 +262,7 @@ "print(recorder)\n", "pred_df = recorder.load_object(\"pred.pkl\")\n", "report_normal_df = recorder.load_object(\"portfolio_analysis/report_normal_1day.pkl\")\n", - "positions = recorder.load_object(\"portfolio_analysis/positions_normal_1day.pkl\")\n", + "positions = recorder.load_object(\"portfolio_analysis/positions_normal_1day.pkl\", trusted=True)\n", "analysis_df = recorder.load_object(\"portfolio_analysis/port_analysis_1day.pkl\")" ] }, diff --git a/qlib/contrib/data/highfreq_provider.py b/qlib/contrib/data/highfreq_provider.py index 611e30d861f..37aa1009a73 100644 --- a/qlib/contrib/data/highfreq_provider.py +++ b/qlib/contrib/data/highfreq_provider.py @@ -1,6 +1,7 @@ import os import time import datetime +from pathlib import Path from typing import Optional import qlib @@ -8,6 +9,7 @@ from qlib.data import D from qlib.config import REG_CN from qlib.utils import init_instance_by_config +from qlib.utils.pickle_utils import ARTIFACT_MIGRATION_URL from qlib.data.dataset.handler import DataHandlerLP from qlib.data.data import Cal from qlib.contrib.ops.high_freq import get_calendar_day, DayLast, FFillNan, BFillNan, Date, Select, IsNull, IsInf, Cut @@ -16,6 +18,12 @@ class HighFreqProvider: + """Generate and reuse datasets within ``artifact_root`` (the current directory by default). + + Relative paths are resolved against this root. Cached datasets are Python + pickles, so the root and its contents must be controlled by trusted users. + """ + def __init__( self, start_time: str, @@ -29,6 +37,7 @@ def __init__( label_conf: Optional[dict] = None, backtest_conf: dict = None, freq: str = "1min", + artifact_root: Optional[str] = None, **kwargs, ) -> None: self.start_time = start_time @@ -44,6 +53,21 @@ def __init__( self.qlib_conf = qlib_conf self.logger = get_module_logger("HighFreqProvider") self.freq = freq + self.artifact_root = Path(artifact_root or Path.cwd()).expanduser().resolve() + + def _resolve_artifact_path(self, path) -> Path: + """Resolve an artifact path without allowing it to escape artifact_root.""" + path = Path(path).expanduser() + target = path.resolve() if path.is_absolute() else (self.artifact_root / path).resolve() + try: + target.relative_to(self.artifact_root) + except ValueError as exc: + raise ValueError( + f"Artifact path {str(path)!r} escapes artifact_root {str(self.artifact_root)!r}. " + "Place caches under a dedicated trusted artifact_root and resolve configured paths relative to it. " + f"Migration guide: {ARTIFACT_MIGRATION_URL}" + ) from exc + return target def get_pre_datasets(self): """Generate the training, validation and test datasets for prediction @@ -52,15 +76,21 @@ def get_pre_datasets(self): Tuple[BaseDataset, BaseDataset, BaseDataset]: The training and test datasets """ - dict_feature_path = self.feature_conf["path"] - train_feature_path = dict_feature_path[:-4] + "_train.pkl" - valid_feature_path = dict_feature_path[:-4] + "_valid.pkl" - test_feature_path = dict_feature_path[:-4] + "_test.pkl" + dict_feature_path = self._resolve_artifact_path(self.feature_conf["path"]) + train_feature_path = self._resolve_artifact_path( + dict_feature_path.with_name(dict_feature_path.stem + "_train.pkl") + ) + valid_feature_path = self._resolve_artifact_path( + dict_feature_path.with_name(dict_feature_path.stem + "_valid.pkl") + ) + test_feature_path = self._resolve_artifact_path( + dict_feature_path.with_name(dict_feature_path.stem + "_test.pkl") + ) - dict_label_path = self.label_conf["path"] - train_label_path = dict_label_path[:-4] + "_train.pkl" - valid_label_path = dict_label_path[:-4] + "_valid.pkl" - test_label_path = dict_label_path[:-4] + "_test.pkl" + dict_label_path = self._resolve_artifact_path(self.label_conf["path"]) + train_label_path = self._resolve_artifact_path(dict_label_path.with_name(dict_label_path.stem + "_train.pkl")) + valid_label_path = self._resolve_artifact_path(dict_label_path.with_name(dict_label_path.stem + "_valid.pkl")) + test_label_path = self._resolve_artifact_path(dict_label_path.with_name(dict_label_path.stem + "_test.pkl")) if ( not os.path.isfile(train_feature_path) @@ -96,7 +126,7 @@ def get_pre_datasets(self): "test": test_label_path, } - return feature, label + return {key: str(value) for key, value in feature.items()}, {key: str(value) for key, value in label.items()} def get_backtest(self, **kwargs) -> None: self._gen_data(self.backtest_conf) @@ -122,8 +152,9 @@ def _prepare_calender_cache(self): get_calendar_day(freq=self.freq) def _gen_dataframe(self, config, datasets=["train", "valid", "test"]): + config = dict(config) try: - path = config.pop("path") + path = self._resolve_artifact_path(config.pop("path")) except KeyError as e: raise ValueError("Must specify the path to save the dataset.") from e if os.path.isfile(path): @@ -131,7 +162,7 @@ def _gen_dataframe(self, config, datasets=["train", "valid", "test"]): self.logger.info(f"[{__name__}]Dataset exists, load from disk.") # res = dataset.prepare(['train', 'valid', 'test']) - with open(path, "rb") as f: + with path.open("rb") as f: data = pkl.load(f) if isinstance(data, dict): res = [data[i] for i in datasets] @@ -139,8 +170,11 @@ def _gen_dataframe(self, config, datasets=["train", "valid", "test"]): res = data.prepare(datasets) self.logger.info(f"[{__name__}]Data loaded, time cost: {time.time() - start:.2f}") else: - if not os.path.exists(os.path.dirname(path)): - os.makedirs(os.path.dirname(path)) + split_paths = [ + self._resolve_artifact_path(path.with_name(path.stem + split + ".pkl")) + for split in ("train", "valid", "test") + ] + path.parent.mkdir(parents=True, exist_ok=True) self.logger.info(f"[{__name__}]Generating dataset") start_time = time.time() self._prepare_calender_cache() @@ -151,21 +185,19 @@ def _gen_dataframe(self, config, datasets=["train", "valid", "test"]): "valid": validset, "test": testset, } - with open(path, "wb") as f: + with path.open("wb") as f: pkl.dump(data, f) - with open(path[:-4] + "train.pkl", "wb") as f: - pkl.dump(trainset, f) - with open(path[:-4] + "valid.pkl", "wb") as f: - pkl.dump(validset, f) - with open(path[:-4] + "test.pkl", "wb") as f: - pkl.dump(testset, f) + for split_path, split_data in zip(split_paths, (trainset, validset, testset)): + with split_path.open("wb") as f: + pkl.dump(split_data, f) res = [data[i] for i in datasets] self.logger.info(f"[{__name__}]Data generated, time cost: {(time.time() - start_time):.2f}") return res def _gen_data(self, config, datasets=["train", "valid", "test"]): + config = dict(config) try: - path = config.pop("path") + path = self._resolve_artifact_path(config.pop("path")) except KeyError as e: raise ValueError("Must specify the path to save the dataset.") from e if os.path.isfile(path): @@ -173,7 +205,7 @@ def _gen_data(self, config, datasets=["train", "valid", "test"]): self.logger.info(f"[{__name__}]Dataset exists, load from disk.") # res = dataset.prepare(['train', 'valid', 'test']) - with open(path, "rb") as f: + with path.open("rb") as f: data = pkl.load(f) if isinstance(data, dict): res = [data[i] for i in datasets] @@ -181,8 +213,7 @@ def _gen_data(self, config, datasets=["train", "valid", "test"]): res = data.prepare(datasets) self.logger.info(f"[{__name__}]Data loaded, time cost: {time.time() - start:.2f}") else: - if not os.path.exists(os.path.dirname(path)): - os.makedirs(os.path.dirname(path)) + path.parent.mkdir(parents=True, exist_ok=True) self.logger.info(f"[{__name__}]Generating dataset") start_time = time.time() self._prepare_calender_cache() @@ -194,21 +225,21 @@ def _gen_data(self, config, datasets=["train", "valid", "test"]): return res def _gen_dataset(self, config): + config = dict(config) try: - path = config.pop("path") + path = self._resolve_artifact_path(config.pop("path")) except KeyError as e: raise ValueError("Must specify the path to save the dataset.") from e if os.path.isfile(path): start = time.time() self.logger.info(f"[{__name__}]Dataset exists, load from disk.") - with open(path, "rb") as f: + with path.open("rb") as f: dataset = pkl.load(f) self.logger.info(f"[{__name__}]Data loaded, time cost: {time.time() - start:.2f}") else: start = time.time() - if not os.path.exists(os.path.dirname(path)): - os.makedirs(os.path.dirname(path)) + path.parent.mkdir(parents=True, exist_ok=True) self.logger.info(f"[{__name__}]Generating dataset") self._prepare_calender_cache() dataset = init_instance_by_config(config) @@ -220,32 +251,34 @@ def _gen_dataset(self, config): return dataset def _gen_day_dataset(self, config, conf_type): + config = dict(config) try: - path = config.pop("path") + path = self._resolve_artifact_path(config.pop("path")) except KeyError as e: raise ValueError("Must specify the path to save the dataset.") from e - if os.path.isfile(path + "tmp_dataset.pkl"): + tmp_dataset_path = self._resolve_artifact_path(path / "tmp_dataset.pkl") + if tmp_dataset_path.is_file(): start = time.time() self.logger.info(f"[{__name__}]Dataset exists, load from disk.") else: start = time.time() - if not os.path.exists(os.path.dirname(path)): - os.makedirs(os.path.dirname(path)) + path.mkdir(parents=True, exist_ok=True) self.logger.info(f"[{__name__}]Generating dataset") self._prepare_calender_cache() dataset = init_instance_by_config(config) self.logger.info(f"[{__name__}]Dataset init, time cost: {time.time() - start:.2f}") dataset.config(dump_all=False, recursive=True) - dataset.to_pickle(path + "tmp_dataset.pkl") + dataset.to_pickle(tmp_dataset_path) - with open(path + "tmp_dataset.pkl", "rb") as f: + with tmp_dataset_path.open("rb") as f: new_dataset = pkl.load(f) time_list = D.calendar(start_time=self.start_time, end_time=self.end_time, freq=self.freq)[::240] def generate_dataset(times): - if os.path.isfile(path + times.strftime("%Y-%m-%d") + ".pkl"): + output_path = self._resolve_artifact_path(path / f"{times:%Y-%m-%d}.pkl") + if output_path.is_file(): print("exist " + times.strftime("%Y-%m-%d")) return self._init_qlib(self.qlib_conf) @@ -256,31 +289,32 @@ def generate_dataset(times): else: new_dataset.handler.setup_data(init_type=DataHandlerLP.IT_LS) new_dataset.config(dump_all=True, recursive=True) - new_dataset.to_pickle(path + times.strftime("%Y-%m-%d") + ".pkl") + new_dataset.to_pickle(output_path) Parallel(n_jobs=8)(delayed(generate_dataset)(times) for times in time_list) def _gen_stock_dataset(self, config, conf_type): + config = dict(config) try: - path = config.pop("path") + path = self._resolve_artifact_path(config.pop("path")) except KeyError as e: raise ValueError("Must specify the path to save the dataset.") from e - if os.path.isfile(path + "tmp_dataset.pkl"): + tmp_dataset_path = self._resolve_artifact_path(path / "tmp_dataset.pkl") + if tmp_dataset_path.is_file(): start = time.time() self.logger.info(f"[{__name__}]Dataset exists, load from disk.") else: start = time.time() - if not os.path.exists(os.path.dirname(path)): - os.makedirs(os.path.dirname(path)) + path.mkdir(parents=True, exist_ok=True) self.logger.info(f"[{__name__}]Generating dataset") self._prepare_calender_cache() dataset = init_instance_by_config(config) self.logger.info(f"[{__name__}]Dataset init, time cost: {time.time() - start:.2f}") dataset.config(dump_all=False, recursive=True) - dataset.to_pickle(path + "tmp_dataset.pkl") + dataset.to_pickle(tmp_dataset_path) - with open(path + "tmp_dataset.pkl", "rb") as f: + with tmp_dataset_path.open("rb") as f: new_dataset = pkl.load(f) instruments = D.instruments(market="all") @@ -289,7 +323,8 @@ def _gen_stock_dataset(self, config, conf_type): ) def generate_dataset(stock): - if os.path.isfile(path + stock + ".pkl"): + output_path = self._resolve_artifact_path(path / f"{stock}.pkl") + if output_path.is_file(): print("exist " + stock) return self._init_qlib(self.qlib_conf) @@ -299,6 +334,6 @@ def generate_dataset(stock): else: new_dataset.handler.setup_data(init_type=DataHandlerLP.IT_LS) new_dataset.config(dump_all=True, recursive=True) - new_dataset.to_pickle(path + stock + ".pkl") + new_dataset.to_pickle(output_path) Parallel(n_jobs=32)(delayed(generate_dataset)(stock) for stock in stock_list) diff --git a/qlib/contrib/meta/data_selection/dataset.py b/qlib/contrib/meta/data_selection/dataset.py index 61efdd63cfb..1cd8caad60e 100644 --- a/qlib/contrib/meta/data_selection/dataset.py +++ b/qlib/contrib/meta/data_selection/dataset.py @@ -14,6 +14,7 @@ from qlib.log import get_module_logger from qlib.utils import auto_filter_kwargs, get_date_by_shift, init_instance_by_config from qlib.utils.data import deepcopy_basic_type +from qlib.utils.pickle_utils import validate_trusted from qlib.workflow import R from qlib.workflow.task.gen import RollingGen, task_generator from qlib.workflow.task.utils import TimeAdjuster @@ -26,11 +27,13 @@ def __init__(self, task_tpl: dict, step: int, exp_name: str): self.step = step self.exp_name = exp_name - def setup(self, trainer=TrainerR, trainer_kwargs={}): + def setup(self, trainer=TrainerR, trainer_kwargs={}, *, trusted: bool = False): """ after running this function `self.data_ic_df` will become set. Each col represents a data. Each row represents the Timestamp of performance of that data. + Set ``trusted=True`` only for saved tasks from a trusted + source and experiment store. Predictions always use restricted loading. For example, .. code-block:: python @@ -47,6 +50,7 @@ def setup(self, trainer=TrainerR, trainer_kwargs={}): """ + trusted = validate_trusted(trusted) # 1) prepare the prediction of proxy models perf_task_tpl = deepcopy(self.task_tpl) # this task is supposed to contains no complicated objects # The only thing we want to save is the prediction @@ -93,7 +97,7 @@ def setup(self, trainer=TrainerR, trainer_kwargs={}): ic_l = [] for _, rec in tqdm(recorders.items(), desc="calc"): pred = rec.load_object("pred.pkl") - task = rec.load_object("task") + task = rec.load_object("task", trusted=trusted) data_key = task["dataset"]["kwargs"]["segments"]["train"] key_l.append(data_key) ic_l.append(delayed(self._calc_perf)(pred.iloc[:, 0], label_df.iloc[:, 0])) @@ -106,7 +110,7 @@ def setup(self, trainer=TrainerR, trainer_kwargs={}): def _calc_perf(self, pred, label): df = pd.DataFrame({"pred": pred, "label": label}) - df = df.groupby("datetime", group_keys=False).corr(method="spearman") + df = df.groupby("datetime", group_keys=True).corr(method="spearman") corr = df.loc(axis=0)[:, "pred"]["label"].droplevel(axis=0, level=-1) return corr @@ -247,6 +251,7 @@ def __init__( hist_step_n: int = 10, task_mode: str = MetaTask.PROC_MODE_FULL, fill_method: str = "max", + trusted: bool = False, ): """ A dataset for meta model. @@ -283,13 +288,17 @@ def __init__( Number of steps of the data similarity information task_mode : str Please refer to the docs of MetaTask + trusted : bool + Explicitly trust task objects read from ``exp_name`` when preparing + internal data. Not needed for an already prepared InternalData. """ + trusted = validate_trusted(trusted) super().__init__(segments=segments) if isinstance(exp_name, InternalData): self.internal_data = exp_name else: self.internal_data = InternalData(task_tpl, step=step, exp_name=exp_name) - self.internal_data.setup() + self.internal_data.setup(trusted=trusted) self.task_tpl = deepcopy(task_tpl) # FIXME: if the handler is shared, how to avoid the explosion of the memroy. self.trunc_days = trunc_days self.hist_step_n = hist_step_n diff --git a/qlib/contrib/model/gbdt.py b/qlib/contrib/model/gbdt.py index 22c29cd4997..950baa904ba 100644 --- a/qlib/contrib/model/gbdt.py +++ b/qlib/contrib/model/gbdt.py @@ -68,19 +68,20 @@ def fit( evals_result = {} # in case of unsafety of Python default values ds_l = self._prepare_data(dataset, reweighter) ds, names = list(zip(*ds_l)) - early_stopping_callback = lgb.early_stopping( - self.early_stopping_rounds if early_stopping_rounds is None else early_stopping_rounds - ) + stopping_rounds = self.early_stopping_rounds if early_stopping_rounds is None else early_stopping_rounds # NOTE: if you encounter error here. Please upgrade your lightgbm verbose_eval_callback = lgb.log_evaluation(period=verbose_eval) evals_result_callback = lgb.record_evaluation(evals_result) + callbacks = [verbose_eval_callback, evals_result_callback] + if stopping_rounds is not None: + callbacks.append(lgb.early_stopping(stopping_rounds)) self.model = lgb.train( self.params, ds[0], # training dataset num_boost_round=self.num_boost_round if num_boost_round is None else num_boost_round, valid_sets=ds, valid_names=names, - callbacks=[early_stopping_callback, verbose_eval_callback, evals_result_callback], + callbacks=callbacks, **kwargs, ) for k in names: diff --git a/qlib/contrib/model/pytorch_hist.py b/qlib/contrib/model/pytorch_hist.py index 779cde9c859..8752088ec7c 100644 --- a/qlib/contrib/model/pytorch_hist.py +++ b/qlib/contrib/model/pytorch_hist.py @@ -6,8 +6,10 @@ from __future__ import print_function import os +import json import numpy as np import pandas as pd +from pathlib import Path from typing import Text, Union import urllib.request import copy @@ -22,6 +24,38 @@ from ...data.dataset.handler import DataHandlerLP from ...contrib.model.pytorch_lstm import LSTMModel from ...contrib.model.pytorch_gru import GRUModel +from ...utils.pickle_utils import ARTIFACT_MIGRATION_URL + +UNKNOWN_STOCK_INDEX = 733 + + +def _load_stock_index(path, upper_bound=None): + """Load and validate HIST's non-executable stock-index metadata.""" + path = Path(path) + if path.suffix.lower() != ".json": + raise ValueError( + "HIST stock_index must be a JSON file. Object .npy files are no longer " + "supported because loading them requires unsafe pickle deserialization. " + "Export the mapping from trusted source metadata and update stock_index, " + f"including on restored model objects. Migration guide: {ARTIFACT_MIGRATION_URL}" + ) + + with path.open("r", encoding="utf-8") as fp: + stock_index = json.load(fp) + + if not isinstance(stock_index, dict): + raise ValueError("HIST stock_index must be a JSON object") + + validated = {} + for instrument, index in stock_index.items(): + if not isinstance(instrument, str): + raise ValueError("HIST stock_index keys must be strings") + if isinstance(index, bool) or not isinstance(index, int): + raise ValueError(f"Invalid stock index for {instrument!r}: expected an integer") + if index < 0 or (upper_bound is not None and index >= upper_bound): + raise ValueError(f"Stock index out of range for {instrument!r}: {index}") + validated[instrument] = index + return validated class HIST(Model): @@ -188,11 +222,11 @@ def get_daily_inter(self, df, shuffle=False): return daily_index, daily_count def train_epoch(self, x_train, y_train, stock_index): - stock2concept_matrix = np.load(self.stock2concept) + stock2concept_matrix = np.load(self.stock2concept, allow_pickle=False) x_train_values = x_train.values y_train_values = np.squeeze(y_train.values) stock_index = stock_index.values - stock_index[np.isnan(stock_index)] = 733 + stock_index[np.isnan(stock_index)] = UNKNOWN_STOCK_INDEX self.HIST_model.train() # organize the train data into daily batches @@ -213,11 +247,11 @@ def train_epoch(self, x_train, y_train, stock_index): def test_epoch(self, data_x, data_y, stock_index): # prepare training data - stock2concept_matrix = np.load(self.stock2concept) + stock2concept_matrix = np.load(self.stock2concept, allow_pickle=False) x_values = data_x.values y_values = np.squeeze(data_y.values) stock_index = stock_index.values - stock_index[np.isnan(stock_index)] = 733 + stock_index[np.isnan(stock_index)] = UNKNOWN_STOCK_INDEX self.HIST_model.eval() scores = [] @@ -259,10 +293,15 @@ def fit( url = "https://github.com/SunsetWolf/qlib_dataset/releases/download/v0/qlib_csi300_stock2concept.npy" urllib.request.urlretrieve(url, self.stock2concept) - stock_index = np.load(self.stock_index, allow_pickle=True).item() - df_train["stock_index"] = 733 + stock2concept_matrix = np.load(self.stock2concept, allow_pickle=False) + if stock2concept_matrix.ndim != 2 or stock2concept_matrix.dtype.hasobject: + raise ValueError("HIST stock2concept must be a two-dimensional numeric array") + stock_index = _load_stock_index(self.stock_index, upper_bound=len(stock2concept_matrix)) + if UNKNOWN_STOCK_INDEX >= len(stock2concept_matrix): + raise ValueError("HIST unknown stock index is outside the stock2concept matrix") + df_train["stock_index"] = UNKNOWN_STOCK_INDEX df_train["stock_index"] = df_train.index.get_level_values("instrument").map(stock_index) - df_valid["stock_index"] = 733 + df_valid["stock_index"] = UNKNOWN_STOCK_INDEX df_valid["stock_index"] = df_valid.index.get_level_values("instrument").map(stock_index) x_train, y_train, stock_index_train = df_train["feature"], df_train["label"], df_train["stock_index"] @@ -331,13 +370,17 @@ def predict(self, dataset: DatasetH, segment: Union[Text, slice] = "test"): if not self.fitted: raise ValueError("model is not fitted yet!") - stock2concept_matrix = np.load(self.stock2concept) - stock_index = np.load(self.stock_index, allow_pickle=True).item() + stock2concept_matrix = np.load(self.stock2concept, allow_pickle=False) + if stock2concept_matrix.ndim != 2 or stock2concept_matrix.dtype.hasobject: + raise ValueError("HIST stock2concept must be a two-dimensional numeric array") + stock_index = _load_stock_index(self.stock_index, upper_bound=len(stock2concept_matrix)) + if UNKNOWN_STOCK_INDEX >= len(stock2concept_matrix): + raise ValueError("HIST unknown stock index is outside the stock2concept matrix") df_test = dataset.prepare(segment, col_set="feature", data_key=DataHandlerLP.DK_I) - df_test["stock_index"] = 733 + df_test["stock_index"] = UNKNOWN_STOCK_INDEX df_test["stock_index"] = df_test.index.get_level_values("instrument").map(stock_index) stock_index_test = df_test["stock_index"].values - stock_index_test[np.isnan(stock_index_test)] = 733 + stock_index_test[np.isnan(stock_index_test)] = UNKNOWN_STOCK_INDEX stock_index_test = stock_index_test.astype("int") df_test = df_test.drop(["stock_index"], axis=1) index = df_test.index @@ -429,7 +472,7 @@ def cal_cos_similarity(self, x, y): # the 2nd dimension of x and y are the same return cos_similarity def forward(self, x, concept_matrix): - device = torch.device(torch.get_device(x)) + device = x.device x_hidden = x.reshape(len(x), self.d_feat, -1) # [N, F, T] x_hidden = x_hidden.permute(0, 2, 1) # [N, T, F] diff --git a/qlib/contrib/rolling/ddgda.py b/qlib/contrib/rolling/ddgda.py index 0fe01d04550..d559e7a295b 100644 --- a/qlib/contrib/rolling/ddgda.py +++ b/qlib/contrib/rolling/ddgda.py @@ -1,7 +1,9 @@ # Copyright (c) Microsoft Corporation. # Licensed under the MIT License. from pathlib import Path +from copy import deepcopy import pickle +import warnings from typing import Optional, Union import pandas as pd @@ -14,8 +16,13 @@ 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.utils.pickle_utils import ( + ARTIFACT_MIGRATION_URL, + restricted_pickle_load, + validate_trusted, +) from qlib.workflow import R +from qlib.workflow.recorder import UnsafeArtifactWarning from qlib.workflow.task.utils import replace_task_handler_with_cache from .base import Rolling @@ -65,6 +72,29 @@ PROC_ARGS = yaml.load(PROC_ARGS, Loader=yaml.FullLoader) UTIL_MODEL_TYPE = Literal["linear", "gbdt"] +_CACHE_LOADER = "qlib.contrib.rolling.ddgda._load_cache" + + +def _load_cache(path, *, trusted: bool = False): + trusted = validate_trusted(trusted) + with Path(path).open("rb") as stream: + if trusted: + warnings.warn( + "Loading a trusted DDG-DA pickle cache may execute arbitrary code. " + "Only use trusted=True when the cache source and storage are trusted.", + UnsafeArtifactWarning, + stacklevel=2, + ) + return pickle.load(stream) + try: + return restricted_pickle_load(stream) + except pickle.UnpicklingError as error: + guide = "" if ARTIFACT_MIGRATION_URL in str(error) else f" Migration guide: {ARTIFACT_MIGRATION_URL}" + raise pickle.UnpicklingError( + f"Restricted loading of DDG-DA cache {str(path)!r} failed: {error}. " + "Set trusted=True at the DDGDA workflow entry point only when the cache source and storage are trusted. " + f"{guide}" + ) from error class DDGDA(Rolling): @@ -76,6 +106,8 @@ class DDGDA(Rolling): - `rm -r mlruns` """ + trusted = False + def __init__( self, sim_task_model: UTIL_MODEL_TYPE = "gbdt", @@ -87,6 +119,7 @@ def __init__( segments: Union[float, str] = 0.62, hist_step_n: int = 30, working_dir: Optional[Union[str, Path]] = None, + trusted: bool = False, **kwargs, ): """ @@ -109,10 +142,15 @@ def __init__( The ratio of training data in the meta task dataset if segments is a string: it will try its best to put its data in training and ensure that the date `segments` is in the test set + trusted : bool + Explicitly allow executable task/meta-model objects from trusted + MLflow storage and handler/internal-data pickle caches from trusted + local storage. Defaults to False. Predictions remain restricted. """ # NOTE: # the horizon must match the meaning in the base task template self.meta_exp_name = "DDG-DA" + self.trusted = validate_trusted(trusted) self.sim_task_model: UTIL_MODEL_TYPE = sim_task_model # The model to capture the distribution of data. self.alpha = alpha self.meta_1st_train_end = meta_1st_train_end @@ -125,6 +163,29 @@ def __init__( self.segments = segments self.hist_step_n = hist_step_n + def _load_cache(self, path): + return _load_cache(path, trusted=self.trusted) + + def _replace_handler_with_cache(self, task, cache_dir=None): + handler = task["dataset"]["kwargs"]["handler"] + if isinstance(handler, dict) and handler.get("class") == _CACHE_LOADER: + handler["kwargs"]["trusted"] = self.trusted + return task + if cache_dir is None: + task = super()._replace_handler_with_cache(task) + else: + task = replace_task_handler_with_cache(task, cache_dir) + handler = task["dataset"]["kwargs"]["handler"] + if isinstance(handler, str) and handler.startswith("file://"): + handler = Path(handler[len("file://") :]) + if isinstance(handler, Path): + # Keep tasks lightweight and reloadable after training changes the handler's serialization settings. + task["dataset"]["kwargs"]["handler"] = { + "class": _CACHE_LOADER, + "kwargs": {"path": str(handler), "trusted": self.trusted}, + } + return task + def _adjust_task(self, task: dict, astype: UTIL_MODEL_TYPE): """ Base on the original task, we need to do some extra things. @@ -139,17 +200,18 @@ def _adjust_task(self, task: dict, astype: UTIL_MODEL_TYPE): # NOTE: here is just for aligning with previous implementation # It is not necessary for the current implementation handler = task["dataset"].setdefault("kwargs", {}).setdefault("handler", {}) + adjustable_handler = isinstance(handler, dict) and handler.get("class") != _CACHE_LOADER if astype == "gbdt": - task["model"] = LGBM_MODEL - if isinstance(handler, dict): + task["model"] = deepcopy(LGBM_MODEL) + if adjustable_handler: # We don't need preprocessing when using GBDT model for k in ["infer_processors", "learn_processors"]: if k in handler.setdefault("kwargs", {}): handler["kwargs"].pop(k) elif astype == "linear": - task["model"] = LINEAR_MODEL - if isinstance(handler, dict): - handler["kwargs"].update(PROC_ARGS) + task["model"] = deepcopy(LINEAR_MODEL) + if adjustable_handler: + handler["kwargs"].update(deepcopy(PROC_ARGS)) else: self.logger.warning("The handler can't be adjusted.") else: @@ -160,7 +222,7 @@ def _get_feature_importance(self): # this must be lightGBM, because it needs to get the feature importance task = self.basic_task(enable_handler_cache=False) task = self._adjust_task(task, astype="gbdt") - task = replace_task_handler_with_cache(task, self.working_dir) + task = self._replace_handler_with_cache(task, self.working_dir) with R.start(experiment_name="feature_importance"): model = init_instance_by_config(task["model"]) @@ -186,7 +248,7 @@ def _dump_data_for_proxy_model(self): # NOTE: adjusting to `self.sim_task_model` just for aligning with previous implementation. # In previous version. The data for proxy model is using sim_task_model's way for processing task = self._adjust_task(self.basic_task(enable_handler_cache=False), self.sim_task_model) - task = replace_task_handler_with_cache(task, self.working_dir) + task = self._replace_handler_with_cache(task, self.working_dir) # if self.meta_data_proc is not None: # else: # # Otherwise, we don't need futher processing @@ -225,7 +287,7 @@ def _dump_data_for_proxy_model(self): "kwargs": {"config": self.working_dir / "fea_label_df.pkl"}, } ) - handler.to_pickle(self.working_dir / self.proxy_hd, dump_all=True) + handler.to_pickle(self.proxy_hd, dump_all=True) @property def _internal_data_path(self): @@ -238,7 +300,7 @@ def _dump_meta_ipt(self): """ # According to the experiments, the choice of the model type is very important for achieving good results sim_task = self._adjust_task(self.basic_task(enable_handler_cache=False), astype=self.sim_task_model) - sim_task = replace_task_handler_with_cache(sim_task, self.working_dir) + sim_task = self._replace_handler_with_cache(sim_task, self.working_dir) if self.sim_task_model == "gbdt": sim_task["model"].setdefault("kwargs", {}).update({"early_stopping_rounds": None, "num_boost_round": 150}) @@ -246,7 +308,7 @@ def _dump_meta_ipt(self): exp_name_sim = f"data_sim_s{self.step}" internal_data = InternalData(sim_task, self.step, exp_name=exp_name_sim) - internal_data.setup(trainer=TrainerR) + internal_data.setup(trainer=TrainerR, trusted=self.trusted) with self._internal_data_path.open("wb") as f: pickle.dump(internal_data, f) @@ -271,7 +333,7 @@ def _train_meta_model(self, fill_method="max"): "dataset": { "class": "qlib.data.dataset.DatasetH", "kwargs": { - "handler": f"file://{(self.working_dir / self.proxy_hd).absolute()}", + "handler": f"file://{self.proxy_hd.absolute()}", "segments": { "train": (train_start, train_end), "test": (test_start, self.basic_task()["dataset"]["kwargs"]["segments"]["test"][1]), @@ -298,8 +360,7 @@ def _train_meta_model(self, fill_method="max"): # the input of meta model (internal data) are shared between proxy model and final forecasting model # 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 = restricted_pickle_load(f) + internal_data = self._load_cache(self._internal_data_path) md = MetaDatasetDS(exp_name=internal_data, **kwargs) @@ -333,7 +394,7 @@ def get_task_list(self): # 1) get meta model exp = R.get_exp(experiment_name=self.meta_exp_name) rec = exp.list_recorders(rtype=exp.RT_L)[0] - meta_model: MetaModelDS = rec.load_object("model") + meta_model: MetaModelDS = rec.load_object("model", trusted=self.trusted) # 2) # we are transfer to knowledge of meta model to final forecasting tasks. @@ -360,8 +421,7 @@ def get_task_list(self): task_mode=MetaTask.PROC_MODE_TRANSFER, ) - with self._internal_data_path.open("rb") as f: - internal_data = restricted_pickle_load(f) + internal_data = self._load_cache(self._internal_data_path) mds = MetaDatasetDS(exp_name=internal_data, **kwargs) # 3) meta model make inference and get new qlib task diff --git a/qlib/contrib/strategy/signal_strategy.py b/qlib/contrib/strategy/signal_strategy.py index bad19ddfdc9..6ab9cd356c1 100644 --- a/qlib/contrib/strategy/signal_strategy.py +++ b/qlib/contrib/strategy/signal_strategy.py @@ -19,7 +19,6 @@ from qlib.log import get_module_logger from qlib.utils import get_pre_trading_date, load_dataset from qlib.contrib.strategy.order_generator import OrderGenerator, OrderGenWOInteract -from qlib.contrib.strategy.optimizer import EnhancedIndexingOptimizer class BaseSignalStrategy(BaseStrategy, ABC): @@ -414,6 +413,12 @@ def __init__( verbose=False, **kwargs, ): + # Import lazily so unrelated strategies do not load CVXPY's native + # solver extensions during module discovery or test collection. + from qlib.contrib.strategy.optimizer.enhanced_indexing import ( # pylint: disable=C0415 + EnhancedIndexingOptimizer, + ) + super().__init__(**kwargs) self.logger = get_module_logger("EnhancedIndexingStrategy") diff --git a/qlib/model/base.py b/qlib/model/base.py index 009a3bd1441..afe6677ac94 100644 --- a/qlib/model/base.py +++ b/qlib/model/base.py @@ -98,7 +98,7 @@ def finetune(self, dataset: Dataset): # Finetune model based on previous trained model with R.start(experiment_name="finetune model"): recorder = R.get_recorder(recorder_id=rid, experiment_name="init models") - model = recorder.load_object("init_model") + model = recorder.load_object("init_model", trusted=True) model.finetune(dataset, num_boost_round=10) diff --git a/qlib/model/trainer.py b/qlib/model/trainer.py index ce204420f81..03ec3f28c7f 100644 --- a/qlib/model/trainer.py +++ b/qlib/model/trainer.py @@ -28,11 +28,19 @@ init_instance_by_config, ) from qlib.utils.paral import call_in_subproc +from qlib.utils.pickle_utils import validate_trusted from qlib.workflow import R from qlib.workflow.recorder import Recorder from qlib.workflow.task.manage import TaskManager, run_task +def _set_trust_kwargs(kwargs: dict, trusted: bool) -> None: + if "trusted" in kwargs: + validate_trusted(kwargs["trusted"]) + elif trusted is not False: + kwargs["trusted"] = validate_trusted(trusted) + + def _log_task_info(task_config: dict): R.log_params(**flatten_dict(task_config)) R.save_objects(**{"task": task_config}) # keep the original format and datatype @@ -88,19 +96,22 @@ def begin_task_train(task_config: dict, experiment_name: str, recorder_name: str return R.get_recorder() -def end_task_train(rec: Recorder, experiment_name: str) -> Recorder: +def end_task_train(rec: Recorder, experiment_name: str, *, trusted: bool = False) -> Recorder: """ Finish task training with real model fitting and saving. Args: rec (Recorder): the recorder will be resumed experiment_name (str): the name of experiment + trusted (bool): allow executable objects in a saved task only + when its source and artifact store are trusted. Defaults to False. Returns: Recorder: the model recorder """ + trusted = validate_trusted(trusted) with R.start(experiment_name=experiment_name, recorder_id=rec.info["id"], resume=True): - task_config = R.load_object("task") + task_config = R.load_object("task", trusted=trusted) _exe_task(task_config) return rec @@ -295,8 +306,16 @@ class DelayTrainerR(TrainerR): A delayed implementation based on TrainerR, which means `train` method may only do some preparation and `end_train` method can do the real model fitting. """ + trusted = False + def __init__( - self, experiment_name: str = None, train_func=begin_task_train, end_train_func=end_task_train, **kwargs + self, + experiment_name: str = None, + train_func=begin_task_train, + end_train_func=end_task_train, + *, + trusted: bool = False, + **kwargs, ): """ Init TrainerRM. @@ -305,9 +324,12 @@ def __init__( experiment_name (str): the default name of experiment. train_func (Callable, optional): default train method. Defaults to `begin_task_train`. end_train_func (Callable, optional): default end_train method. Defaults to `end_task_train`. + trusted (bool): explicitly trust saved task objects and + their artifact store when resuming training. Defaults to False. """ super().__init__(experiment_name, train_func, **kwargs) self.end_train_func = end_train_func + self.trusted = validate_trusted(trusted) self.delay = True def end_train(self, models, end_train_func=None, experiment_name: str = None, **kwargs) -> List[Recorder]: @@ -330,6 +352,7 @@ def end_train(self, models, end_train_func=None, experiment_name: str = None, ** end_train_func = self.end_train_func if experiment_name is None: experiment_name = self.experiment_name + _set_trust_kwargs(kwargs, self.trusted) for rec in models: if rec.list_tags()[self.STATUS_KEY] == self.STATUS_END: continue @@ -494,6 +517,8 @@ class DelayTrainerRM(TrainerRM): """ + trusted = False + def __init__( self, experiment_name: str = None, @@ -501,6 +526,8 @@ def __init__( train_func=begin_task_train, end_train_func=end_task_train, skip_run_task: bool = False, + *, + trusted: bool = False, **kwargs, ): """ @@ -511,6 +538,8 @@ def __init__( task_pool (str): task pool name in TaskManager. None for use same name as experiment_name. train_func (Callable, optional): default train method. Defaults to `begin_task_train`. end_train_func (Callable, optional): default end_train method. Defaults to `end_task_train`. + trusted (bool): explicitly trust saved task objects and + their artifact store in both end_train and worker. Defaults to False. skip_run_task (bool): If skip_run_task == True: Only run_task in the worker. Otherwise skip run_task. @@ -518,6 +547,7 @@ def __init__( """ super().__init__(experiment_name, task_pool, train_func, **kwargs) self.end_train_func = end_train_func + self.trusted = validate_trusted(trusted) self.delay = True self.skip_run_task = skip_run_task @@ -577,6 +607,7 @@ def end_train(self, recs, end_train_func=None, experiment_name: str = None, **kw _id_list.append(rec.list_tags()[self.TM_ID]) query = {"_id": {"$in": _id_list}} + _set_trust_kwargs(kwargs, self.trusted) if not self.skip_run_task: run_task( end_train_func, @@ -593,13 +624,15 @@ def end_train(self, recs, end_train_func=None, experiment_name: str = None, **kw rec.set_tags(**{self.STATUS_KEY: self.STATUS_END}) return recs - def worker(self, end_train_func=None, experiment_name: str = None): + def worker(self, end_train_func=None, experiment_name: str = None, **kwargs): """ The multiprocessing method for `end_train`. It can share a same task_pool with `end_train` and can run in other progress or other machines. Args: end_train_func (Callable, optional): the end_train method which need at least `recorders` and `experiment_name`. Defaults to None for using self.end_train_func. experiment_name (str): the experiment name, None for use default name. + kwargs: parameters for end_train_func, including an explicit + trusted override. """ if end_train_func is None: end_train_func = self.end_train_func @@ -608,11 +641,13 @@ def worker(self, end_train_func=None, experiment_name: str = None): task_pool = self.task_pool if task_pool is None: task_pool = experiment_name + _set_trust_kwargs(kwargs, self.trusted) run_task( end_train_func, task_pool=task_pool, experiment_name=experiment_name, before_status=TaskManager.STATUS_PART_DONE, + **kwargs, ) def has_worker(self) -> bool: diff --git a/qlib/utils/mod.py b/qlib/utils/mod.py index 5cb2ed3f453..47d4fd67c01 100644 --- a/qlib/utils/mod.py +++ b/qlib/utils/mod.py @@ -199,8 +199,10 @@ def class_casting(obj: object, cls: type): """ orig_cls = obj.__class__ obj.__class__ = cls - yield - obj.__class__ = orig_cls + try: + yield + finally: + obj.__class__ = orig_cls def find_all_classes(module_path: Union[str, ModuleType], cls: type) -> List[type]: diff --git a/qlib/utils/pickle_utils.py b/qlib/utils/pickle_utils.py index 920692f3c89..ec0047e9062 100644 --- a/qlib/utils/pickle_utils.py +++ b/qlib/utils/pickle_utils.py @@ -11,6 +11,15 @@ import pickle from typing import Any, BinaryIO, Set, Tuple +ARTIFACT_MIGRATION_URL = "https://qlib.readthedocs.io/en/latest/start/artifact_migration.html" + + +def validate_trusted(trusted: bool) -> bool: + if not isinstance(trusted, bool): + raise TypeError(f"`trusted` must be a bool. Migration guide: {ARTIFACT_MIGRATION_URL}") + return trusted + + # 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]] = { @@ -49,15 +58,112 @@ ("qlib.data.dataset.handler", "DataHandler"), ("qlib.data.dataset.handler", "DataHandlerLP"), ("qlib.data.dataset.loader", "StaticDataLoader"), + # NumPy reconstruction primitives. Keep this list explicit: trusting the + # whole numpy namespace would also expose functions such as numpy.load. + ("numpy", "ndarray"), + # Record arrays use the ndarray reconstruction/state path, with a record + # scalar type in their structured dtype. NumPy 2 moved the array's path. + ("numpy", "recarray"), + ("numpy.rec", "recarray"), + ("numpy", "record"), + ("numpy", "dtype"), + ("numpy", "scalar"), + ("numpy.core.multiarray", "_reconstruct"), + ("numpy.core.multiarray", "scalar"), + ("numpy._core.multiarray", "_reconstruct"), + ("numpy._core.multiarray", "scalar"), + # Protocol 5 uses _frombuffer instead of _reconstruct for numeric arrays. + ("numpy.core.numeric", "_frombuffer"), + ("numpy._core.numeric", "_frombuffer"), + ("numpy.ma.core", "_mareconstruct"), + # NumPy 1.x and 2.x pickle this class under different module paths. + ("numpy.ma.core", "MaskedArray"), + ("numpy.ma", "MaskedArray"), + # Pandas reconstruction primitives used by Series/DataFrame pickles. + # These entries are deliberately exact. I/O helpers such as + # pandas.read_pickle must never be added here. + ("pandas.core.series", "Series"), + ("pandas.core.frame", "DataFrame"), + ("pandas.core.internals.managers", "BlockManager"), + ("pandas.core.internals.managers", "SingleBlockManager"), + ("pandas.core.internals.blocks", "new_block"), + ("pandas._libs.internals", "_unpickle_block"), + ("pandas.core.indexes.base", "_new_Index"), + ("pandas.core.indexes.base", "Index"), + ("pandas.core.indexes.range", "RangeIndex"), + ("pandas.core.indexes.multi", "MultiIndex"), + ("pandas.core.indexes.datetimes", "_new_DatetimeIndex"), + ("pandas.core.indexes.datetimes", "DatetimeIndex"), + ("pandas.core.indexes.timedeltas", "TimedeltaIndex"), + ("pandas.core.indexes.period", "PeriodIndex"), + ("pandas.core.indexes.interval", "_new_IntervalIndex"), + ("pandas.core.indexes.interval", "IntervalIndex"), + ("pandas._libs.tslibs.timestamps", "_unpickle_timestamp"), + ("pandas._libs.tslibs.timestamps", "Timestamp"), + ("pandas._libs.tslibs.timedeltas", "Timedelta"), + ("pandas._libs.tslibs.period", "Period"), + ("pandas._libs.arrays", "__pyx_unpickle_NDArrayBacked"), + ("pandas.core.arrays.datetimes", "DatetimeArray"), + ("pandas.core.arrays.timedeltas", "TimedeltaArray"), + ("pandas.core.arrays.period", "PeriodArray"), + ("pandas.core.arrays.categorical", "Categorical"), + ("pandas.core.dtypes.dtypes", "CategoricalDtype"), + ("pandas.core.dtypes.dtypes", "PeriodDtype"), + ("pandas.core.dtypes.dtypes", "IntervalDtype"), + ("pandas.core.dtypes.dtypes", "SparseDtype"), + ("pandas.core.arrays.sparse.dtype", "SparseDtype"), + ("pandas.core.arrays.interval", "IntervalArray"), + ("pandas._libs.interval", "__pyx_unpickle_IntervalMixin"), + ("pandas.core.arrays.sparse.array", "SparseArray"), + ("pandas._libs.sparse", "IntIndex"), + ("pandas._libs.sparse", "BlockIndex"), + ("pandas.core.dtypes.dtypes", "DatetimeTZDtype"), + ("pandas._libs.tslibs.nattype", "__nat_unpickle"), + ("pandas._libs.missing", "NA"), + # DatetimeIndex/PeriodIndex retain their frequency and timezone metadata. + ("pandas._libs.tslibs.offsets", "Day"), + ("pandas._libs.tslibs.offsets", "BusinessDay"), + # Reconstruct business schedules from times, weekmasks and date scalars; + # CustomBusinessDay rebuilds its calendar rather than unpickling one. + ("pandas._libs.tslibs.offsets", "BusinessHour"), + ("pandas._libs.tslibs.offsets", "CustomBusinessDay"), + ("pandas._libs.tslibs.offsets", "Week"), + ("pandas._libs.tslibs.offsets", "MonthBegin"), + ("pandas._libs.tslibs.offsets", "MonthEnd"), + ("pandas._libs.tslibs.offsets", "BusinessMonthBegin"), + ("pandas._libs.tslibs.offsets", "BusinessMonthEnd"), + ("pandas._libs.tslibs.offsets", "QuarterBegin"), + ("pandas._libs.tslibs.offsets", "QuarterEnd"), + ("pandas._libs.tslibs.offsets", "YearBegin"), + ("pandas._libs.tslibs.offsets", "YearEnd"), + ("pandas._libs.tslibs.offsets", "Hour"), + ("pandas._libs.tslibs.offsets", "Minute"), + ("pandas._libs.tslibs.offsets", "Second"), + ("pandas._libs.tslibs.offsets", "Milli"), + ("pandas._libs.tslibs.offsets", "Micro"), + ("pandas._libs.tslibs.offsets", "Nano"), + ("pytz", "_UTC"), + ("pytz", "_p"), + # Nullable arrays serialize their masks and dtype objects as well as data. + ("pandas.core.arrays.integer", "IntegerArray"), + ("pandas.core.arrays.integer", "Int8Dtype"), + ("pandas.core.arrays.integer", "Int16Dtype"), + ("pandas.core.arrays.integer", "Int32Dtype"), + ("pandas.core.arrays.integer", "Int64Dtype"), + ("pandas.core.arrays.integer", "UInt8Dtype"), + ("pandas.core.arrays.integer", "UInt16Dtype"), + ("pandas.core.arrays.integer", "UInt32Dtype"), + ("pandas.core.arrays.integer", "UInt64Dtype"), + ("pandas.core.arrays.floating", "FloatingArray"), + ("pandas.core.arrays.floating", "Float32Dtype"), + ("pandas.core.arrays.floating", "Float64Dtype"), + ("pandas.core.arrays.boolean", "BooleanArray"), + ("pandas.core.arrays.boolean", "BooleanDtype"), + ("pandas.core.arrays.string_", "StringArray"), + ("pandas.core.arrays.string_", "StringDtype"), } -TRUSTED_MODULE_PREFIXES = ( - "pandas", - "numpy", -) - - class RestrictedUnpickler(pickle.Unpickler): """Custom unpickler that only allows safe classes to be deserialized. @@ -82,17 +188,14 @@ def find_class(self, module: str, name: str): 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." + f"This is to prevent arbitrary code execution through pickle deserialization. " + f"Migration guide: {ARTIFACT_MIGRATION_URL}" ) diff --git a/qlib/workflow/__init__.py b/qlib/workflow/__init__.py index a29e471c04b..8cde98da280 100644 --- a/qlib/workflow/__init__.py +++ b/qlib/workflow/__init__.py @@ -21,6 +21,7 @@ from .recorder import Recorder from ..utils import Wrapper from ..utils.exceptions import RecorderInitializationError +from ..utils.pickle_utils import validate_trusted class QlibRecorder: @@ -533,11 +534,18 @@ def save_objects(self, local_path=None, artifact_path=None, **kwargs: Dict[Text, ) self.get_exp().get_recorder(start=True).save_objects(local_path, artifact_path, **kwargs) - def load_object(self, name: Text): + def load_object(self, name: Text, *, trusted: bool = False): """ Method for loading an object from artifacts in the experiment in the uri. + + Set ``trusted=True`` only for pickle artifacts whose source and storage + are trusted. Unrestricted pickle loading may execute arbitrary code. """ - return self.get_exp().get_recorder(start=True).load_object(name) + trusted = validate_trusted(trusted) + recorder = self.get_exp().get_recorder(start=True) + if trusted is False: + return recorder.load_object(name) + return recorder.load_object(name, trusted=trusted) def log_params(self, **kwargs): """ diff --git a/qlib/workflow/online/strategy.py b/qlib/workflow/online/strategy.py index d545e4bc9a6..19bbed234e7 100644 --- a/qlib/workflow/online/strategy.py +++ b/qlib/workflow/online/strategy.py @@ -9,6 +9,7 @@ from qlib.log import get_module_logger from qlib.model.ens.group import RollingGroup from qlib.utils import transform_end_date +from qlib.utils.pickle_utils import validate_trusted from qlib.workflow.online.utils import OnlineTool, OnlineToolR from qlib.workflow.recorder import Recorder from qlib.workflow.task.collect import Collector, RecorderCollector @@ -94,11 +95,15 @@ class RollingStrategy(OnlineStrategy): This example strategy always uses the latest rolling model sas online models. """ + trusted = False + def __init__( self, name_id: str, task_template: Union[dict, List[dict]], rolling_gen: RollingGen, + *, + trusted: bool = False, ): """ Init RollingStrategy. @@ -109,15 +114,19 @@ def __init__( name_id (str): a unique name or id. Will be also the name of the Experiment. task_template (Union[dict, List[dict]]): a list of task_template or a single template, which will be used to generate many tasks using rolling_gen. rolling_gen (RollingGen): an instance of RollingGen + trusted (bool): allow unrestricted task/model/dataset + loading for this strategy's trusted experiment. Data artifacts + remain restricted. Defaults to False. """ super().__init__(name_id=name_id) self.exp_name = self.name_id + self.trusted = validate_trusted(trusted) if not isinstance(task_template, list): task_template = [task_template] self.task_template = task_template self.rg = rolling_gen assert issubclass(self.rg.__class__, RollingGen), "The rolling strategy relies on the feature if RollingGen" - self.tool = OnlineToolR(self.exp_name) + self.tool = OnlineToolR(self.exp_name, trusted=trusted) self.ta = TimeAdjuster() def get_collector(self, process_list=[RollingGroup()], rec_key_func=None, rec_filter_func=None, artifacts_key=None): @@ -134,7 +143,7 @@ def get_collector(self, process_list=[RollingGroup()], rec_key_func=None, rec_fi """ def rec_key(recorder): - task_config = recorder.load_object("task") + task_config = recorder.load_object("task", trusted=self.trusted) model_key = task_config["model"]["class"] rolling_key = task_config["dataset"]["kwargs"]["segments"]["test"] return model_key, rolling_key @@ -184,7 +193,7 @@ def prepare_tasks(self, cur_time) -> List[dict]: ) res = [] for rec in latest_records: - task = rec.load_object("task") + task = rec.load_object("task", trusted=self.trusted) res.extend(self.rg.gen_following_tasks(task, calendar_latest)) return res @@ -200,9 +209,11 @@ def _list_latest(self, rec_list: List[Recorder]): """ if len(rec_list) == 0: return rec_list, None - max_test = max(rec.load_object("task")["dataset"]["kwargs"]["segments"]["test"] for rec in rec_list) + max_test = max( + rec.load_object("task", trusted=self.trusted)["dataset"]["kwargs"]["segments"]["test"] for rec in rec_list + ) latest_rec = [] for rec in rec_list: - if rec.load_object("task")["dataset"]["kwargs"]["segments"]["test"] == max_test: + if rec.load_object("task", trusted=self.trusted)["dataset"]["kwargs"]["segments"]["test"] == max_test: latest_rec.append(rec) return latest_rec, max_test diff --git a/qlib/workflow/online/update.py b/qlib/workflow/online/update.py index 5047a1bd25e..832acef4699 100644 --- a/qlib/workflow/online/update.py +++ b/qlib/workflow/online/update.py @@ -14,6 +14,7 @@ from qlib.data.dataset.handler import DataHandlerLP from qlib.model import Model from qlib.utils import get_date_by_shift +from qlib.utils.pickle_utils import validate_trusted from qlib.workflow.recorder import Recorder from qlib.workflow.record_temp import SignalRecord @@ -23,8 +24,12 @@ class RMDLoader: Recorder Model Dataset Loader """ - def __init__(self, rec: Recorder): + trusted = False + + def __init__(self, rec: Recorder, *, trusted: bool = False): + """Only enable ``trusted`` for trusted model/dataset storage.""" self.rec = rec + self.trusted = validate_trusted(trusted) def get_dataset( self, start_time, end_time, segments=None, unprepared_dataset: Optional[DatasetH] = None @@ -52,7 +57,7 @@ def get_dataset( if segments is None: segments = {"test": (start_time, end_time)} if unprepared_dataset is None: - dataset: DatasetH = self.rec.load_object("dataset") + dataset: DatasetH = self.rec.load_object("dataset", trusted=self.trusted) else: dataset = unprepared_dataset dataset.config(handler_kwargs={"start_time": start_time, "end_time": end_time}, segments=segments) @@ -60,7 +65,7 @@ def get_dataset( return dataset def get_model(self) -> Model: - return self.rec.load_object("params.pkl") + return self.rec.load_object("params.pkl", trusted=self.trusted) class RecordUpdater(metaclass=ABCMeta): @@ -101,6 +106,8 @@ class DSBasedUpdater(RecordUpdater, metaclass=ABCMeta): SZ300676 -0.001321 """ + trusted = False + def __init__( self, record: Recorder, @@ -110,6 +117,8 @@ def __init__( freq="day", fname="pred.pkl", loader_cls: type = RMDLoader, + *, + trusted: bool = False, ): """ Init PredUpdater. @@ -144,6 +153,10 @@ def __init__( loader_cls : type the class to load the model and dataset + trusted : bool + Allow unrestricted loading of model/dataset artifacts from a + trusted source and store. Predictions and labels remain restricted. + A custom loader must accept this keyword when it is enabled. """ # TODO: automate this hist_ref in the future. @@ -153,7 +166,11 @@ def __init__( self.hist_ref = hist_ref self.freq = freq self.fname = fname - self.rmdl = loader_cls(rec=record) + self.trusted = validate_trusted(trusted) + if trusted is False: + self.rmdl = loader_cls(rec=record) + else: + self.rmdl = loader_cls(rec=record, trusted=trusted) latest_date = D.calendar(freq=freq)[-1] if to_date is None: @@ -190,7 +207,11 @@ def prepare_data(self, unprepared_dataset: Optional[DatasetH] = None) -> Dataset """ # automatically getting the historical dependency if not specified if self.hist_ref is None: - dataset: DatasetH = self.record.load_object("dataset") if unprepared_dataset is None else unprepared_dataset + dataset: DatasetH = ( + self.record.load_object("dataset", trusted=self.trusted) + if unprepared_dataset is None + else unprepared_dataset + ) # Special treatment of historical dependencies if isinstance(dataset, TSDatasetH): hist_ref = dataset.step_len - 1 @@ -289,8 +310,8 @@ class LabelUpdater(DSBasedUpdater): - The label is generated from record_temp.SignalRecord. """ - def __init__(self, record: Recorder, to_date=None, **kwargs): - super().__init__(record, to_date=to_date, fname="label.pkl", **kwargs) + def __init__(self, record: Recorder, to_date=None, *, trusted: bool = False, **kwargs): + super().__init__(record, to_date=to_date, fname="label.pkl", trusted=trusted, **kwargs) def get_update_data(self, dataset: Dataset) -> pd.DataFrame: new_label = SignalRecord.generate_label(dataset) diff --git a/qlib/workflow/online/utils.py b/qlib/workflow/online/utils.py index c390ca00921..e2fab67171a 100644 --- a/qlib/workflow/online/utils.py +++ b/qlib/workflow/online/utils.py @@ -11,6 +11,7 @@ from qlib.log import get_module_logger from qlib.utils.exceptions import LoadObjectError +from qlib.utils.pickle_utils import validate_trusted from qlib.workflow.online.update import PredUpdater from qlib.workflow.recorder import Recorder from qlib.workflow.task.utils import list_recorders @@ -89,15 +90,20 @@ class OnlineToolR(OnlineTool): The implementation of OnlineTool based on (R)ecorder. """ - def __init__(self, default_exp_name: str = None): + trusted = False + + def __init__(self, default_exp_name: str = None, *, trusted: bool = False): """ Init OnlineToolR. Args: default_exp_name (str): the default experiment name. + trusted (bool): explicitly allow model/dataset pickle loading + from trusted sources and storage when updating predictions. """ super().__init__() self.default_exp_name = default_exp_name + self.trusted = validate_trusted(trusted) def set_online_tag(self, tag, recorder: Union[Recorder, List]): """ @@ -168,7 +174,7 @@ def update_online_pred(self, to_date=None, from_date=None, exp_name: str = None) online_models = self.online_models(exp_name=exp_name) for rec in online_models: try: - updater = PredUpdater(rec, to_date=to_date, from_date=from_date) + updater = PredUpdater(rec, to_date=to_date, from_date=from_date, trusted=self.trusted) except LoadObjectError as e: # skip the recorder without pred self.logger.warn(f"An exception `{str(e)}` happened when load `pred.pkl`, skip it.") diff --git a/qlib/workflow/record_temp.py b/qlib/workflow/record_temp.py index ecd58ec2098..a123f3d48c2 100644 --- a/qlib/workflow/record_temp.py +++ b/qlib/workflow/record_temp.py @@ -20,6 +20,7 @@ from ..utils.time import Freq from ..utils.data import deepcopy_basic_type from ..utils.exceptions import QlibException +from ..utils.pickle_utils import validate_trusted from ..contrib.eva.alpha import calc_ic, calc_long_short_return, calc_long_short_prec logger = get_module_logger("workflow", logging.INFO) @@ -78,7 +79,7 @@ def generate(self, **kwargs): """ raise NotImplementedError(f"Please implement the `generate` method.") - def load(self, name: str, parents: bool = True): + def load(self, name: str, parents: bool = True, *, trusted: bool = False): """ It behaves the same as self.recorder.load_object. But it is an easier interface because users don't have to care about `get_path` and `artifact_path` @@ -93,17 +94,23 @@ def load(self, name: str, parents: bool = True): So parents recursively find the path in parents Sub classes has higher priority + trusted : bool + Explicitly allow unrestricted loading of this artifact, including + parent-path lookup, only when its source and storage are trusted. + This does not change the policy of subsequent loads or generation. + Return ------ The stored records. """ + load_kwargs = {"trusted": True} if validate_trusted(trusted) else {} try: - return self.recorder.load_object(self.get_path(name)) + return self.recorder.load_object(self.get_path(name), **load_kwargs) except LoadObjectError as e: if parents: if self.depend_cls is not None: with class_casting(self, self.depend_cls): - return self.load(name, parents=True) + return self.load(name, parents=True, **load_kwargs) raise e def list(self): diff --git a/qlib/workflow/recorder.py b/qlib/workflow/recorder.py index 5fd99c0769f..de023de8423 100644 --- a/qlib/workflow/recorder.py +++ b/qlib/workflow/recorder.py @@ -7,6 +7,7 @@ import mlflow import shutil import pickle +import warnings import tempfile import subprocess import platform @@ -16,6 +17,7 @@ from qlib.utils.serial import Serializable from qlib.utils.exceptions import LoadObjectError from qlib.utils.paral import AsyncCaller +from qlib.utils.pickle_utils import ARTIFACT_MIGRATION_URL, RestrictedUnpickler, validate_trusted from ..log import TimeInspector, get_module_logger from mlflow.store.artifact.azure_blob_artifact_repo import AzureBlobArtifactRepository @@ -25,6 +27,10 @@ mlflow.utils.validation.MAX_PARAM_VAL_LENGTH = 1000 +class UnsafeArtifactWarning(UserWarning): + """Warning emitted when an artifact is loaded with unrestricted pickle.""" + + class Recorder: """ This is the `Recorder` class for logging the experiments. The API is designed similar to mlflow. @@ -87,7 +93,7 @@ def save_objects(self, local_path=None, artifact_path=None, **kwargs): """ raise NotImplementedError(f"Please implement the `save_objects` method.") - def load_object(self, name): + def load_object(self, name, *, trusted=False): """ Load objects such as prediction file or model checkpoints. @@ -95,6 +101,10 @@ def load_object(self, name): ---------- name : str name of the file to be loaded. + trusted : bool + Whether to allow unrestricted pickle loading. Defaults to False. + Implementations must not silently retry restricted loads with an + unrestricted loader. Only enable this for trusted sources and storage. Returns ------- @@ -410,14 +420,17 @@ def save_objects(self, local_path=None, artifact_path=None, **kwargs): self.client.log_artifact(self.id, temp_dir / name, artifact_path) shutil.rmtree(temp_dir) - def load_object(self, name, unpickler=pickle.Unpickler): + def load_object(self, name, unpickler=None, *, trusted=False): """ Load object such as prediction file or model checkpoint in mlflow. Args: name (str): the object name - unpickler: Supporting using custom unpickler + unpickler: Optional custom unpickler. Custom unpicklers are trusted + code and may execute arbitrary code while loading an artifact. + trusted (bool): Use Python's unrestricted pickle loader. This must + only be enabled for artifacts from a trusted source. Raises: LoadObjectError: if raise some exceptions when load the object @@ -427,12 +440,44 @@ def load_object(self, name, unpickler=pickle.Unpickler): """ assert self.uri is not None, "Please start the experiment and recorder first before using recorder directly." + trusted = validate_trusted(trusted) + if trusted and unpickler is not None: + raise ValueError("`trusted` and `unpickler` cannot be used together") + path = None try: path = self.client.download_artifacts(self.id, name) with Path(path).open("rb") as f: - data = unpickler(f).load() + if trusted: + warnings.warn( + "Loading a trusted pickle artifact may execute arbitrary code. " + "Only use trusted=True when the artifact source and storage are trusted.", + UnsafeArtifactWarning, + stacklevel=2, + ) + loader = pickle.Unpickler(f) + elif unpickler is not None: + warnings.warn( + "A custom artifact unpickler may execute arbitrary code.", + UnsafeArtifactWarning, + stacklevel=2, + ) + loader = unpickler(f) + else: + loader = RestrictedUnpickler(f) + data = loader.load() return data + except pickle.UnpicklingError as e: + if not trusted and unpickler is None: + guide = "" if ARTIFACT_MIGRATION_URL in str(e) else f" Migration guide: {ARTIFACT_MIGRATION_URL}" + raise LoadObjectError( + f"Restricted loading of artifact {name!r} failed: {e}. " + "Use a supported data representation, or explicitly set trusted=True " + "on load_object() or the workflow entry point only when both the artifact " + "source and storage are trusted. Unrestricted pickle may execute arbitrary code." + f"{guide}" + ) from e + raise LoadObjectError(str(e)) from e except Exception as e: raise LoadObjectError(str(e)) from e finally: diff --git a/qlib/workflow/task/collect.py b/qlib/workflow/task/collect.py index bedbd96d201..b60b364f245 100644 --- a/qlib/workflow/task/collect.py +++ b/qlib/workflow/task/collect.py @@ -6,11 +6,13 @@ """ from collections import defaultdict +from collections.abc import Mapping from qlib.log import TimeInspector -from typing import Callable, Dict, Iterable, List +from typing import Any, Callable, Dict, Iterable, List, Optional from qlib.log import get_module_logger from qlib.utils.serial import Serializable from qlib.utils.exceptions import LoadObjectError +from qlib.utils.pickle_utils import validate_trusted from qlib.workflow import R from qlib.workflow.exp import Experiment from qlib.workflow.recorder import Recorder @@ -135,6 +137,7 @@ def collect(self) -> dict: class RecorderCollector(Collector): ART_KEY_RAW = "__raw" + artifact_load_kwargs: Optional[Dict[str, Dict[str, Any]]] = None def __init__( self, @@ -146,6 +149,8 @@ def __init__( artifacts_key=None, list_kwargs={}, status: Iterable = {Recorder.STATUS_FI}, + *, + artifact_load_kwargs: Optional[Dict[str, Dict[str, Any]]] = None, ): """ Init RecorderCollector. @@ -161,7 +166,24 @@ def __init__( artifacts_key (str or List, optional): the artifacts key you want to get. If None, get all artifacts. list_kwargs (str): arguments for list_recorders function. status (Iterable): only collect recorders with specific status. None indicating collecting all the recorders + artifact_load_kwargs (dict, optional): keyword arguments for ``load_object``, + keyed by artifact alias in ``artifacts_path``. For example, + ``{"model": {"trusted": True}}`` authorizes only the model artifact. + Unspecified artifacts retain restricted loading. Only opt in after + verifying the selected artifacts' sources and storage. """ + if artifact_load_kwargs is not None and not isinstance(artifact_load_kwargs, Mapping): + raise TypeError("artifact_load_kwargs must map artifact keys to loading keyword arguments") + self.artifact_load_kwargs = {} + for key, options in (artifact_load_kwargs or {}).items(): + if key == self.ART_KEY_RAW or key not in artifacts_path: + raise ValueError(f"Loading keyword arguments require a configured artifact path: {key!r}") + if not isinstance(options, Mapping): + raise TypeError(f"Loading keyword arguments for {key!r} must be a mapping") + options = dict(options) + if "trusted" in options and not validate_trusted(options["trusted"]): + del options["trusted"] + self.artifact_load_kwargs[key] = options super().__init__(process_list=process_list) if isinstance(experiment, str): experiment = R.get_exp(experiment_name=experiment) @@ -224,6 +246,7 @@ def collect(self, artifacts_key=None, rec_filter_func=None, only_exist=True) -> for r in recs: status_stat[r.status] += 1 logger.info(f"Nubmer of recorders after filter: {status_stat}") + load_kwargs = self.artifact_load_kwargs or {} for rec in recs: rec_key = self.rec_key_func(rec) for key in artifacts_key: @@ -231,11 +254,11 @@ def collect(self, artifacts_key=None, rec_filter_func=None, only_exist=True) -> artifact = rec else: try: - artifact = rec.load_object(self.artifacts_path[key]) + artifact = rec.load_object(self.artifacts_path[key], **load_kwargs.get(key, {})) except LoadObjectError as e: if only_exist: # only collect existing artifact - logger.warning(f"Fail to load {self.artifacts_path[key]} and it is ignored.") + logger.warning(f"Fail to load {self.artifacts_path[key]} and it is ignored: {e}") continue raise e # give user some warning if the values are overridden diff --git a/qlib/workflow/task/gen.py b/qlib/workflow/task/gen.py index cf95e600633..294bfc876b1 100644 --- a/qlib/workflow/task/gen.py +++ b/qlib/workflow/task/gen.py @@ -112,7 +112,9 @@ def handler_mod(task: dict, rolling_gen): 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: + if handler_end_time is not None and ( + 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. diff --git a/scripts/data_collector/pit/README.md b/scripts/data_collector/pit/README.md index cec430628ab..6c1fc8e9f7d 100644 --- a/scripts/data_collector/pit/README.md +++ b/scripts/data_collector/pit/README.md @@ -26,6 +26,11 @@ python collector.py download_data --source_dir ~/.qlib/stock_data/source/pit --s ### Normalize Data + +The normalizer uses report publication dates, or calendar-day offsets of 45 days +(quarterly) / 90 days (annual) when a publication date is missing. It does not +fetch or use a Baostock trading calendar. + ```bash python collector.py normalize_data --interval quarterly --source_dir ~/.qlib/stock_data/source/pit --normalize_dir ~/.qlib/stock_data/source/pit_normalized ``` diff --git a/scripts/data_collector/pit/collector.py b/scripts/data_collector/pit/collector.py index c34b31348dc..a3bb33d91a7 100644 --- a/scripts/data_collector/pit/collector.py +++ b/scripts/data_collector/pit/collector.py @@ -16,7 +16,7 @@ sys.path.append(str(BASE_DIR.parent.parent)) from data_collector.base import BaseCollector, BaseRun, BaseNormalize -from data_collector.utils import get_hs_stock_symbols, get_calendar_list +from data_collector.utils import get_hs_stock_symbols class PitCollector(BaseCollector): @@ -239,7 +239,8 @@ def normalize(self, df: pd.DataFrame) -> pd.DataFrame: return df def _get_calendar_list(self) -> Iterable[pd.Timestamp]: - return get_calendar_list() + # PIT uses report dates and calendar-day offsets, not trading sessions. + return [] class Run(BaseRun): diff --git a/tests/conftest.py b/tests/conftest.py index 7a869da16ee..e7f03d1acc4 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -1,6 +1,8 @@ import os import sys +import pytest + """Ignore RL tests on non-linux platform.""" collect_ignore = [] @@ -8,3 +10,76 @@ for root, dirs, files in os.walk("rl"): for file in files: collect_ignore.append(os.path.join(root, file)) + + +@pytest.fixture +def workflow_context(tmp_path, monkeypatch, request): + """An offline daily market and isolated recorder store for real workflows.""" + from copy import deepcopy + from types import SimpleNamespace + + import numpy as np + import pandas as pd + import qlib + from qlib.config import C + from qlib.data.cache import H + from qlib.workflow import R + + calendar = pd.bdate_range("2020-01-01", periods=260) + instruments = [f"SH{600000 + i:06d}" for i in range(getattr(request, "param", 8))] + provider_uri = tmp_path / "market" + calendars = provider_uri / "calendars" + calendars.mkdir(parents=True) + calendars.joinpath("day.txt").write_text("\n".join(calendar.strftime("%Y-%m-%d")) + "\n") + instruments_dir = provider_uri / "instruments" + instruments_dir.mkdir() + spans = "".join(f"{symbol}\t{calendar[0]:%Y-%m-%d}\t{calendar[-1]:%Y-%m-%d}\n" for symbol in instruments) + instruments_dir.joinpath("csi300.txt").write_text(spans) + instruments_dir.joinpath("all.txt").write_text(spans) + random = np.random.RandomState(42) + market_return = random.normal(0.0004, 0.008, len(calendar)) + for i, symbol in enumerate(instruments + ["SH000300"]): + stock_return = random.normal(0, 0.006, len(calendar)) + for day in range(1, len(calendar)): + stock_return[day] += 0.7 * stock_return[day - 1] + returns = market_return + stock_return + close = (20 + i) * np.exp(np.cumsum(returns)) + opening = close * (1 + random.normal(0, 0.003, len(calendar))) + values = { + "close": close, + "open": opening, + "high": np.maximum(close, opening) * 1.01, + "low": np.minimum(close, opening) * 0.99, + "vwap": (close + opening) / 2, + "volume": random.uniform(1000000, 2000000, len(calendar)), + "factor": np.ones(len(calendar)), + } + directory = provider_uri / "features" / symbol.lower() + directory.mkdir(parents=True) + for field, array in values.items(): + np.concatenate(([0], array)).astype("", "topk": 3, "n_drop": 1}, + }, + "backtest": { + "start_time": context["backtest_dates"][0], + "end_time": context["backtest_dates"][-1], + "account": 1000000, + "benchmark": "SH000300", + "exchange_kwargs": { + "limit_threshold": 0.095, + "deal_price": "close", + "open_cost": 0.0005, + "close_cost": 0.0015, + "min_cost": 5, + }, + }, + } + }, + }, + ], + } + recorder = task_train(task, "hist-workflow") + return recorder.id, recorder.experiment_id + + +def _check_artifacts(context, recorder): + from qlib.contrib.data.handler import Alpha360 + from qlib.data.dataset import DatasetH + from qlib.data.dataset.handler import DataHandlerLP + from qlib.utils.exceptions import LoadObjectError + from qlib.workflow.online.update import RMDLoader + from qlib.workflow.record_temp import SignalRecord + from qlib.workflow.recorder import UnsafeArtifactWarning + + assert recorder.status == "FINISHED" + assert recorder.load_object("task")["model"]["kwargs"]["n_epochs"] == 1 + segments = context["segments"] + for name, class_name in (("params.pkl", "HIST"), ("dataset", "DatasetH")): + with pytest.raises(LoadObjectError, match=class_name): + recorder.load_object(name) + restricted = RMDLoader(recorder) + with pytest.raises(LoadObjectError, match="HIST"): + restricted.get_model() + with pytest.raises(LoadObjectError, match="DatasetH"): + restricted.get_dataset(*segments["test"]) + + loader = RMDLoader(recorder, trusted=True) + with pytest.warns(UnsafeArtifactWarning): + model = loader.get_model() + with pytest.warns(UnsafeArtifactWarning): + dataset = loader.get_dataset(segments["train"][0], segments["test"][1], segments=segments) + assert type(model).__name__ == "HIST" + assert isinstance(dataset, DatasetH) + assert isinstance(dataset.handler, Alpha360) + assert model.fitted and model.n_epochs == 1 + assert model.device == torch.device("cpu") + assert all(parameter.device.type == "cpu" for parameter in model.HIST_model.parameters()) + assert all(torch.isfinite(parameter).all() for parameter in model.HIST_model.parameters()) + + for segment, days in (("train", 12), ("valid", 4), ("test", 6)): + data = dataset.prepare(segment, col_set=["feature", "label"], data_key=DataHandlerLP.DK_L) + expected_index = pd.MultiIndex.from_product( + [pd.bdate_range(*segments[segment]), context["instruments"]], names=["datetime", "instrument"] + ) + pd.testing.assert_index_equal(data.index, expected_index) + assert data["feature"].shape == (days * 8, 360) + assert data["label"].shape == (days * 8, 1) + assert np.isfinite(data.to_numpy()).all() + + optimizer_state = model.train_optimizer.state_dict()["state"] + assert optimizer_state + # Every used parameter must have completed all 12 daily batches of the single epoch. + assert {int(state["step"]) for state in optimizer_state.values()} == {12} + initial = torch.load(Path(context["root"]) / "base_gru.pt", map_location="cpu", weights_only=True) + weight_delta = (model.HIST_model.rnn.weight_ih_l0 - initial["rnn.weight_ih_l0"]).abs().max().item() + assert weight_delta > 0 + + predictions = recorder.load_object("pred.pkl") + labels = recorder.load_object("label.pkl") + expected_index = pd.MultiIndex.from_product( + [pd.to_datetime(context["test_dates"]), context["instruments"]], names=["datetime", "instrument"] + ) + pd.testing.assert_index_equal(predictions.index, expected_index) + pd.testing.assert_index_equal(labels.index, expected_index) + assert predictions.shape == labels.shape == (48, 1) + assert list(predictions.columns) == ["score"] + assert np.isfinite(predictions.to_numpy()).all() + assert np.isfinite(labels.to_numpy()).all() + assert predictions["score"].groupby(level="datetime").std().gt(0).all() + actual = model.predict(dataset).to_frame("score") + pd.testing.assert_frame_equal(actual, predictions, check_exact=True) + pd.testing.assert_frame_equal(SignalRecord.generate_label(dataset), labels, check_exact=True) + + for name in ("ic.pkl", "ric.pkl"): + signal_analysis = recorder.load_object(f"sig_analysis/{name}") + assert signal_analysis.shape == (6,) + np.testing.assert_array_equal(signal_analysis.index, pd.to_datetime(context["test_dates"])) + assert np.isfinite(signal_analysis.to_numpy()).all() + metrics = recorder.list_metrics() + assert np.isfinite([metrics[key] for key in ("IC", "ICIR", "Rank IC", "Rank ICIR")]).all() + report = recorder.load_object("portfolio_analysis/report_normal_1day.pkl") + assert report.shape == (6, 9) + pd.testing.assert_index_equal(report.index, pd.DatetimeIndex(context["backtest_dates"], name="datetime")) + assert np.isfinite(report.to_numpy()).all() + assert report["account"].gt(0).all() and report["value"].gt(0).all() + assert report["turnover"].sum() > 0 and report["cost"].sum() > 0 + risk = recorder.load_object("portfolio_analysis/port_analysis_1day.pkl") + assert risk.shape == (10, 1) + assert np.isfinite(risk.to_numpy()).all() + + # Opting in once must not change subsequent default loading. + with pytest.raises(LoadObjectError, match="HIST"): + recorder.load_object("params.pkl") + with pytest.raises(LoadObjectError, match="DatasetH"): + recorder.load_object("dataset") + return { + "prediction_rows": len(predictions), + "backtest_rows": len(report), + "optimizer_steps": 12, + "weight_delta": weight_delta, + "prediction_max_error": float((actual - predictions).abs().to_numpy().max()), + "turnover": float(report["turnover"].sum()), + } + + +@pytest.mark.slow +def test_hist_full_workflow_artifact_trust(workflow_context): + context = workflow_context + dates = context.calendar.strftime("%Y-%m-%d").tolist() + inputs = { + "root": str(context.root), + "provider_uri": str(context.provider_uri), + "uri": context.uri, + "instruments": context.instruments, + "segments": { + "train": [dates[100], dates[111]], + "valid": [dates[112], dates[115]], + "test": [dates[116], dates[121]], + }, + "test_dates": dates[116:122], + "backtest_dates": dates[117:123], + } + source = Path(__file__).resolve() + environment = dict(os.environ) + environment.update( + { + "HOME": str(context.root), + "USERPROFILE": str(context.root), + "PYTHONUSERBASE": site.getuserbase(), + "PYTHONPATH": os.pathsep.join(filter(None, [str(source.parents[2]), environment.get("PYTHONPATH")])), + "OMP_NUM_THREADS": "1", + "OPENBLAS_NUM_THREADS": "1", + "MLFLOW_ALLOW_FILE_STORE": "true", + } + ) + # Isolate Torch RNG/threads and HIST's default ~/tmp checkpoint without replacing any workflow stages. + results = [] + for stage in ("train", "reload"): + result = subprocess.run( + [sys.executable, str(source), stage, json.dumps(inputs)], + cwd=context.root, + env=environment, + capture_output=True, + text=True, + timeout=180, + check=False, + ) + assert result.returncode == 0, f"{stage} failed:\n{result.stdout}\n{result.stderr}" + results.append(json.loads((context.root / f"{stage}_result.json").read_text(encoding="utf-8"))) + assert results[0] == results[1] + assert results[1]["prediction_rows"] == 48 + assert results[1]["backtest_rows"] == 6 + assert results[1]["optimizer_steps"] == 12 + assert results[1]["prediction_max_error"] == 0 + + +if __name__ == "__main__": + from qlib.workflow import R + + stage, encoded_context = sys.argv[1:] + context = json.loads(encoded_context) + root = Path(context["root"]) + _init_workflow(context) + if stage == "train": + identifiers = _train_workflow(context) + (root / "recorder.json").write_text(json.dumps(identifiers), encoding="utf-8") + else: + assert stage == "reload" + assert "qlib.contrib.model.pytorch_hist" not in sys.modules + identifiers = json.loads((root / "recorder.json").read_text(encoding="utf-8")) + recorder = R.get_recorder(recorder_id=identifiers[0], experiment_id=identifiers[1]) + result = _check_artifacts(context, recorder) + assert "qlib.contrib.model.pytorch_hist" in sys.modules + (root / f"{stage}_result.json").write_text(json.dumps(result), encoding="utf-8") diff --git a/tests/rolling_tests/test_ddgda.py b/tests/rolling_tests/test_ddgda.py new file mode 100644 index 00000000000..5004258f308 --- /dev/null +++ b/tests/rolling_tests/test_ddgda.py @@ -0,0 +1,344 @@ +import copy +import pickle +from pathlib import Path + +import numpy as np +import pandas as pd +import pytest +import torch +import yaml +from joblib import parallel_backend + +from qlib.contrib.meta.data_selection.dataset import InternalData +from qlib.contrib.meta.data_selection.model import MetaModelDS, TimeReweighter +from qlib.contrib.rolling.ddgda import DDGDA +from qlib.model.trainer import DelayTrainerR +from qlib.utils import init_instance_by_config +from qlib.utils.exceptions import LoadObjectError +from qlib.workflow import R +from qlib.workflow.online.update import RMDLoader +from qlib.workflow.recorder import UnsafeArtifactWarning +from qlib.workflow.task.gen import RollingGen, handler_mod + + +def _task(context): + dates = context.calendar.strftime("%Y-%m-%d") + return { + "model": { + "class": "LinearModel", + "module_path": "qlib.contrib.model.linear", + "kwargs": {"estimator": "ridge", "alpha": 0.05}, + }, + "dataset": { + "class": "DatasetH", + "module_path": "qlib.data.dataset", + "kwargs": { + "handler": { + "class": "Alpha158", + "module_path": "qlib.contrib.data.handler", + "kwargs": { + "start_time": dates[60], + "end_time": dates[219], + "fit_start_time": dates[60], + "fit_end_time": dates[119], + "instruments": context.instruments, + }, + }, + "segments": { + "train": [dates[60], dates[119]], + "valid": [dates[120], dates[139]], + "test": [dates[160], dates[199]], + }, + }, + }, + "record": [ + "qlib.workflow.record_temp.SignalRecord", + "qlib.workflow.record_temp.SigAnaRecord", + { + "class": "PortAnaRecord", + "module_path": "qlib.workflow.record_temp", + "kwargs": { + "config": { + "strategy": { + "class": "TopkDropoutStrategy", + "module_path": "qlib.contrib.strategy", + "kwargs": {"signal": "", "topk": 3, "n_drop": 1}, + }, + "backtest": { + "start_time": dates[160], + "end_time": dates[199], + "account": 1000000, + "benchmark": "SH000300", + "exchange_kwargs": { + "limit_threshold": 0.095, + "deal_price": "close", + "open_cost": 0.0005, + "close_cost": 0.0015, + "min_cost": 5, + }, + }, + } + }, + }, + ], + } + + +@pytest.mark.slow +@pytest.mark.parametrize( + "workflow_context,sim_task_model", [(8, "linear"), (192, "gbdt")], indirect=["workflow_context"] +) +def test_ddgda_full_workflow(workflow_context, sim_task_model, monkeypatch): + context = workflow_context + monkeypatch.chdir(context.root) + config = context.root / "workflow.yaml" + config.write_text(yaml.safe_dump({"task": _task(context)})) + work = context.root / "work" + work.mkdir() + workflow = DDGDA( + conf_path=config, + exp_name="ddgda-result", + rolling_exp="ddgda-rolling", + working_dir="work", + sim_task_model=sim_task_model, + train_start=str(context.calendar[60].date()), + meta_1st_train_end=str(context.calendar[119].date()), + horizon=1, + step=20, + hist_step_n=2, + loss_skip_thresh=2, + fea_imp_n=4, + segments=0.5, + ) + previous_threads = torch.get_num_threads() + previous_grad = torch.is_grad_enabled() + previous_rng = torch.get_rng_state() + try: + torch.set_num_threads(1) + with pytest.raises(pickle.UnpicklingError, match="trusted=True"): + workflow.run() + assert not workflow._internal_data_path.exists() + workflow.trusted = True + with parallel_backend("threading"), pytest.warns(UnsafeArtifactWarning): + workflow.run() + finally: + torch.set_num_threads(previous_threads) + torch.set_grad_enabled(previous_grad) + torch.set_rng_state(previous_rng) + + with pytest.warns(UnsafeArtifactWarning): + internal_data = workflow._load_cache(workflow._internal_data_path) + assert isinstance(internal_data, InternalData) + assert internal_data.data_ic_df.shape == (160, 7) + assert internal_data.data_ic_df.notna().any().all() + similarity_recorders = R.list_recorders(experiment_name=internal_data.exp_name) + assert len(similarity_recorders) == 7 + if sim_task_model == "gbdt": + for similarity_recorder in similarity_recorders.values(): + with pytest.warns(UnsafeArtifactWarning): + similarity_model = RMDLoader(similarity_recorder, trusted=True).get_model() + assert similarity_model.early_stopping_rounds is None + assert similarity_model.num_boost_round == 150 + assert similarity_model.model.num_trees() > 1 + assert (work / "handler_proxy.pkl").is_file() + assert (work / "fea_label_df.pkl").is_file() + assert workflow._task_path.is_file() + + meta_recorders = R.list_recorders(experiment_name=workflow.meta_exp_name) + assert len(meta_recorders) == 1 + meta_recorder = next(iter(meta_recorders.values())) + with pytest.raises(LoadObjectError, match="MetaModelDS"): + meta_recorder.load_object("model") + with pytest.warns(UnsafeArtifactWarning): + meta_model = meta_recorder.load_object("model", trusted=True) + assert isinstance(meta_model, MetaModelDS) + assert meta_model.fitted + assert meta_model.max_epoch == 30 + for name in ("loss/train", "loss/test", "ic/train", "ic/test"): + assert np.isfinite(meta_recorder.list_metrics()[name]) + history = meta_recorder.client.get_metric_history(meta_recorder.id, "loss/train") + assert {entry.step for entry in history} == set(range(30)) + assert np.isfinite([entry.value for entry in history]).all() + assert np.ptp([entry.value for entry in history]) > 1e-8 + assert all(torch.isfinite(parameter).all() for parameter in meta_model.tn.parameters()) + + records = R.list_recorders(experiment_name=workflow.rolling_exp) + assert len(records) == 2 + rolling_predictions = {} + for recorder in records.values(): + with pytest.raises(LoadObjectError, match="LinearModel"): + RMDLoader(recorder).get_model() + with pytest.raises(LoadObjectError, match="DatasetH"): + recorder.load_object("dataset") + with pytest.raises(LoadObjectError, match="Forbidden class"): + recorder.load_object("task") + with pytest.warns(UnsafeArtifactWarning): + task = recorder.load_object("task", trusted=True) + model = RMDLoader(recorder, trusted=True).get_model() + dataset = recorder.load_object("dataset", trusted=True) + assert isinstance(task["reweighter"], TimeReweighter) + assert np.isfinite(task["reweighter"].time_weight).all() + assert (task["reweighter"].time_weight > 0).all() + assert task["reweighter"].time_weight.std() > 0 + dataset.setup_data(handler_kwargs={"init_type": "load_state"}) + actual = model.predict(dataset) + expected = recorder.load_object("pred.pkl").iloc[:, 0] + pd.testing.assert_series_equal(actual, expected, check_names=False, check_exact=True) + rolling_predictions[tuple(task["dataset"]["kwargs"]["segments"]["test"])] = expected + + assert workflow._task_path.stat().st_size < 10000 + with pytest.warns(UnsafeArtifactWarning): + saved_tasks = workflow._load_cache(workflow._task_path) + assert len(saved_tasks) == 2 + assert all(isinstance(task["dataset"]["kwargs"]["handler"], dict) for task in saved_tasks) + replay_trainer = DelayTrainerR(experiment_name="ddgda-replay") + replay_records = replay_trainer.train(saved_tasks) + with pytest.raises(LoadObjectError, match="TimeReweighter"): + replay_trainer.end_train(replay_records) + with pytest.warns(UnsafeArtifactWarning): + replay_trainer.end_train(replay_records, trusted=True) + for task, replay_recorder in zip(saved_tasks, replay_records): + replay_prediction = replay_recorder.load_object("pred.pkl").iloc[:, 0] + expected = rolling_predictions[tuple(task["dataset"]["kwargs"]["segments"]["test"])] + pd.testing.assert_series_equal(replay_prediction, expected, check_exact=True) + + recorder = R.get_recorder(experiment_name=workflow.exp_name, recorder_id=workflow._rid) + predictions = recorder.load_object("pred.pkl") + labels = recorder.load_object("label.pkl") + expected_index = pd.MultiIndex.from_product( + [context.calendar[160:200], context.instruments], names=["datetime", "instrument"] + ) + pd.testing.assert_index_equal(predictions.index, expected_index) + pd.testing.assert_index_equal(labels.index, expected_index) + assert np.isfinite(predictions.values).all() + assert np.isfinite(labels.values).all() + ic = recorder.load_object("sig_analysis/ic.pkl") + assert len(ic) == 40 + assert np.isfinite(ic).all() + report = recorder.load_object("portfolio_analysis/report_normal_1day.pkl") + pd.testing.assert_index_equal(report.index, context.calendar[160:200].rename("datetime"), check_names=False) + assert np.isfinite(report[["return", "cost", "bench", "account"]].values).all() + assert report["cost"].sum() > 0 + assert report["turnover"].sum() > 0 + + +def test_internal_data_daily_rank_ic(): + dates = pd.DatetimeIndex(["2020-01-03", "2020-01-06", "2020-01-07"], name="datetime") + index = pd.MultiIndex.from_product([dates, ["A", "B", "C"]], names=["datetime", "instrument"]) + pred = pd.Series([1.0, 2.0, 3.0, 3.0, 2.0, 1.0, 1.0, 1.0, 1.0], index=index) + label = pd.Series([1.0, 2.0, 3.0] * 3, index=index) + result = InternalData({}, 20, "unused")._calc_perf(pred, label) + pd.testing.assert_series_equal(result, pd.Series([1.0, -1.0, np.nan], index=dates, name="label")) + + +def test_ddgda_cache_requires_explicit_trust(tmp_path): + path = tmp_path / "internal.pkl" + expected = InternalData({}, 20, "unused") + path.write_bytes(pickle.dumps(expected)) + workflow = object.__new__(DDGDA) + with pytest.raises(pickle.UnpicklingError, match="trusted=True"): + workflow._load_cache(path) + workflow.trusted = True + with pytest.warns(UnsafeArtifactWarning, match="cache source and storage"): + actual = workflow._load_cache(path) + assert isinstance(actual, InternalData) + assert actual.__dict__ == expected.__dict__ + workflow.trusted = False + with pytest.raises(pickle.UnpicklingError, match="trusted=True"): + workflow._load_cache(path) + + +def test_ddgda_external_cache_keeps_a_reloadable_reference(tmp_path): + path = tmp_path / "handler.pkl" + expected = {"data": [1, 2, 3]} + path.write_bytes(pickle.dumps(expected)) + workflow = DDGDA(conf_path=tmp_path / "unused.yaml", h_path=path, trusted=True) + task = {"dataset": {"kwargs": {"handler": "replaced-by-h-path"}}} + task = workflow._replace_handler_with_cache(task) + for model_type in ("linear", "gbdt"): + workflow._adjust_task(task, model_type) + task = workflow._replace_handler_with_cache(task, tmp_path / "unused") + handler = task["dataset"]["kwargs"]["handler"] + assert handler["kwargs"] == {"path": str(path), "trusted": True} + with pytest.warns(UnsafeArtifactWarning): + assert init_instance_by_config(handler) == expected + assert not (tmp_path / "unused").exists() + workflow.trusted = False + task = workflow._replace_handler_with_cache(task) + assert task["dataset"]["kwargs"]["handler"]["kwargs"]["trusted"] is False + assert init_instance_by_config(task["dataset"]["kwargs"]["handler"]) == expected + + +@pytest.mark.parametrize("value", ["false", 0, 1, None, np.bool_(True)]) +def test_ddgda_cache_rejects_non_boolean_trust(tmp_path, value): + workflow = object.__new__(DDGDA) + workflow.trusted = value + with pytest.raises(TypeError, match="must be a bool"): + workflow._load_cache(tmp_path / "not-opened.pkl") + + +def test_ddgda_fresh_cache_loads_in_open_ended_rolling_tail(workflow_context): + context = workflow_context + dates = context.calendar + task = _task(context) + dataset_kwargs = task["dataset"]["kwargs"] + dataset_kwargs["handler"]["kwargs"]["end_time"] = str(dates[-1].date()) + dataset_kwargs["segments"]["test"] = [str(dates[-10].date()), str(dates[-1].date())] + cache_dir = context.root / "handler-cache" + cache_dir.mkdir() + workflow = DDGDA(conf_path=context.root / "unused.yaml", trusted=True) + task = workflow._replace_handler_with_cache(task, cache_dir) + reference = copy.deepcopy(task["dataset"]["kwargs"]["handler"]) + assert Path(reference["kwargs"]["path"]).is_file() + + with pytest.warns(UnsafeArtifactWarning): + expected = init_instance_by_config(task["dataset"]).prepare("test", col_set=["feature", "label"]) + assert len(expected) == 10 * len(context.instruments) + rolling_tasks = RollingGen(step=20).generate(task) + assert len(rolling_tasks) == 1 + dataset_config = rolling_tasks[0]["dataset"] + assert dataset_config["kwargs"]["segments"]["test"][1] is None + assert dataset_config["kwargs"]["handler"] == reference + assert task["dataset"]["kwargs"]["handler"] == reference + with pytest.warns(UnsafeArtifactWarning): + actual = init_instance_by_config(dataset_config).prepare("test", col_set=["feature", "label"]) + pd.testing.assert_frame_equal(actual, expected, check_exact=True) + + restricted = copy.deepcopy(dataset_config) + restricted["kwargs"]["handler"]["kwargs"]["trusted"] = False + with pytest.raises(pickle.UnpicklingError, match="Alpha158"): + init_instance_by_config(restricted) + + +@pytest.mark.parametrize( + "handler_end,test_end,expected_end", + [(20, 30, 30), (40, 30, 40), (20, None, None), (None, 30, None), (None, None, None)], +) +def test_rolling_handler_mod_extends_only_bounded_handlers(workflow_context, handler_end, test_end, expected_end): + dates = workflow_context.calendar + handler_kwargs = {"end_time": None if handler_end is None else dates[handler_end]} + task = { + "dataset": { + "kwargs": { + "handler": {"kwargs": handler_kwargs}, + "segments": {"test": (dates[10], None if test_end is None else dates[test_end])}, + } + } + } + handler_mod(task, RollingGen(step=20)) + assert handler_kwargs["end_time"] == (None if expected_end is None else dates[expected_end]) + + +def test_ddgda_adjusted_task_does_not_mutate_defaults(): + from qlib.contrib.rolling.ddgda import LGBM_MODEL, LINEAR_MODEL, PROC_ARGS + + workflow = object.__new__(DDGDA) + expected = copy.deepcopy((LGBM_MODEL, LINEAR_MODEL, PROC_ARGS)) + for kind in ("gbdt", "linear"): + task = {"dataset": {"kwargs": {"handler": {"kwargs": {}}}}} + workflow._adjust_task(task, kind) + task["model"]["kwargs"]["num_boost_round"] = 150 + if kind == "linear": + task["dataset"]["kwargs"]["handler"]["kwargs"]["infer_processors"][0]["kwargs"]["clip_outlier"] = False + task["dataset"]["kwargs"]["handler"]["kwargs"].clear() + assert (LGBM_MODEL, LINEAR_MODEL, PROC_ARGS) == expected diff --git a/tests/rolling_tests/test_update_pred.py b/tests/rolling_tests/test_update_pred.py index b3ca2e0368f..8d8d3fb1f4d 100644 --- a/tests/rolling_tests/test_update_pred.py +++ b/tests/rolling_tests/test_update_pred.py @@ -10,8 +10,9 @@ from qlib.model.trainer import task_train from qlib.tests import TestAutoData from qlib.tests.config import CSI300_GBDT_TASK +from qlib.utils.exceptions import LoadObjectError from qlib.workflow.online.utils import OnlineToolR -from qlib.workflow.online.update import LabelUpdater +from qlib.workflow.online.update import LabelUpdater, PredUpdater class TestRolling(TestAutoData): @@ -49,7 +50,11 @@ def test_update_pred(self): pred = rec.load_object("pred.pkl") - online_tool = OnlineToolR(exp_name) + with pytest.raises(LoadObjectError, match="trusted=True"): + PredUpdater(rec, from_date=latest_date - pd.Timedelta(days=20)).update() + pd.testing.assert_frame_equal(rec.load_object("pred.pkl"), pred) + + online_tool = OnlineToolR(exp_name, trusted=True) online_tool.reset_online_tag(rec) # set to online model online_tool.update_online_pred(to_date=latest_date + pd.Timedelta(days=10)) @@ -111,8 +116,11 @@ def test_update_label(self): pred = rec.load_object("pred.pkl") - online_tool = OnlineToolR(exp_name) + online_tool = OnlineToolR(exp_name, trusted=True) online_tool.reset_online_tag(rec) # set to online model + with pytest.raises(LoadObjectError, match="trusted=True"): + OnlineToolR(exp_name).update_online_pred() + pd.testing.assert_frame_equal(rec.load_object("pred.pkl"), pred) online_tool.update_online_pred() new_pred = rec.load_object("pred.pkl") @@ -124,7 +132,10 @@ def test_update_label(self): self.assertTrue(label_date < pred_date) # Update label now - lu = LabelUpdater(rec) + with pytest.raises(LoadObjectError, match="trusted=True"): + LabelUpdater(rec).update() + pd.testing.assert_frame_equal(rec.load_object("label.pkl"), label) + lu = LabelUpdater(rec, trusted=True) lu.update() new_label = rec.load_object("label.pkl") new_label_date = new_label.index.get_level_values("datetime").max() diff --git a/tests/security/test_artifact_reader_trust.py b/tests/security/test_artifact_reader_trust.py new file mode 100644 index 00000000000..0c24ede14d2 --- /dev/null +++ b/tests/security/test_artifact_reader_trust.py @@ -0,0 +1,201 @@ +import pickle +from types import SimpleNamespace +from unittest.mock import Mock + +import numpy as np +import pandas as pd +import pytest + +from qlib.contrib.model.linear import LinearModel +from qlib.data.dataset import DatasetH +from qlib.data.dataset.handler import DataHandlerLP +from qlib.data.dataset.loader import StaticDataLoader +from qlib.utils.exceptions import LoadObjectError +from qlib.workflow import R +from qlib.workflow.record_temp import RecordTemp +from qlib.workflow.recorder import Recorder, UnsafeArtifactWarning +from qlib.workflow.task.collect import RecorderCollector + + +class NestedRecord(RecordTemp): + artifact_path = "nested" + depend_cls = RecordTemp + + +@pytest.fixture +def reader_artifacts(workflow_context): + dates = workflow_context.calendar[:8] + index = pd.MultiIndex.from_product([dates, ["SH600000"]], names=["datetime", "instrument"]) + x = np.arange(8, dtype=float) + data = pd.DataFrame( + np.column_stack([x, x**2, 2 * x - 0.3 * x**2 + 1]), + index=index, + columns=pd.MultiIndex.from_tuples([("feature", "x"), ("feature", "x2"), ("label", "y")]), + ) + handler = DataHandlerLP(instruments=None, data_loader=StaticDataLoader(data)) + handler.config(dump_all=True) + dataset = DatasetH(handler=handler, segments={"train": (dates[0], dates[-1]), "test": (dates[0], dates[-1])}) + model = LinearModel(estimator="ridge", alpha=0.1, fit_intercept=True) + model.fit(dataset) + prediction = model.predict(dataset) + assert prediction.nunique() > 1 + with R.start(experiment_name="artifact-readers"): + recorder = R.get_recorder() + recorder.save_objects( + **{ + "params.pkl": model, + "dataset": dataset, + "pred.pkl": prediction.to_frame("score"), + "not-data.pkl": model, + } + ) + return SimpleNamespace(recorder=recorder, dataset=dataset, prediction=prediction) + + +@pytest.mark.parametrize("nested", [False, True]) +def test_record_template_authorizes_only_the_requested_load(reader_artifacts, nested): + artifacts = reader_artifacts + record = NestedRecord(artifacts.recorder) if nested else RecordTemp(artifacts.recorder) + with pytest.raises(LoadObjectError, match="LinearModel"): + record.load("params.pkl") + assert type(record) is (NestedRecord if nested else RecordTemp) + with pytest.warns(UnsafeArtifactWarning): + model = record.load("params.pkl", trusted=True) + with pytest.warns(UnsafeArtifactWarning): + dataset = record.load("dataset", trusted=True) + pd.testing.assert_series_equal(model.predict(dataset), artifacts.prediction, check_exact=True) + pd.testing.assert_frame_equal(record.load("pred.pkl"), artifacts.prediction.to_frame("score"), check_exact=True) + with pytest.raises(LoadObjectError, match="LinearModel"): + record.load("not-data.pkl") + assert type(record) is (NestedRecord if nested else RecordTemp) + + +def test_record_template_can_disable_parent_lookup(reader_artifacts): + record = NestedRecord(reader_artifacts.recorder) + with pytest.raises(LoadObjectError): + record.load("params.pkl", parents=False, trusted=True) + assert type(record) is NestedRecord + + +def test_record_template_keeps_its_class_when_parent_artifact_is_missing(reader_artifacts): + record = NestedRecord(reader_artifacts.recorder) + with pytest.raises(LoadObjectError): + record.load("missing.pkl", trusted=True) + assert type(record) is NestedRecord + + +@pytest.mark.parametrize("value", ["false", 0, 1, None, np.bool_(True)]) +def test_record_template_rejects_invalid_consent_before_reading(value): + recorder = Mock() + with pytest.raises(TypeError, match="must be a bool"): + RecordTemp(recorder).load("unused.pkl", trusted=value) + recorder.load_object.assert_not_called() + + +def test_collector_authorizes_selected_artifacts_without_widening_data_reads(reader_artifacts): + artifacts = reader_artifacts + paths = {"model": "params.pkl", "pred": "pred.pkl"} + options = {"model": {"trusted": True}} + collector = RecorderCollector(lambda: [artifacts.recorder], artifacts_path=paths, artifact_load_kwargs=options) + assert options == {"model": {"trusted": True}} + options["model"]["trusted"] = False + options["pred"] = {"trusted": True} + + with pytest.warns(UnsafeArtifactWarning) as caught: + collected = collector.collect(only_exist=False) + assert len(caught) == 1 + recorder_id = artifacts.recorder.info["id"] + pd.testing.assert_series_equal( + collected["model"][recorder_id].predict(artifacts.dataset), artifacts.prediction, check_exact=True + ) + pd.testing.assert_frame_equal( + collected["pred"][recorder_id], artifacts.prediction.to_frame("score"), check_exact=True + ) + + collector.artifacts_path["pred"] = "not-data.pkl" + with pytest.warns(UnsafeArtifactWarning), pytest.raises(LoadObjectError, match="LinearModel"): + collector.collect(only_exist=False) + + +def test_collector_default_refusal_reports_why_an_artifact_is_skipped(reader_artifacts, monkeypatch): + from qlib.workflow.task import collect + + logger = Mock() + monkeypatch.setattr(collect, "get_module_logger", lambda *args: logger) + collector = RecorderCollector(lambda: [reader_artifacts.recorder], artifacts_path={"model": "params.pkl"}) + with pytest.raises(LoadObjectError, match="LinearModel"): + collector.collect(only_exist=False) + assert collector.collect() == {} + assert "Forbidden class:" in logger.warning.call_args.args[0] + assert "LinearModel" in logger.warning.call_args.args[0] + + +def test_collector_preserves_backend_unpickler_policy(reader_artifacts): + artifacts = reader_artifacts + collector = RecorderCollector( + lambda: [artifacts.recorder], + artifacts_path={"model": "params.pkl", "pred": "pred.pkl"}, + artifact_load_kwargs={"model": {"unpickler": pickle.Unpickler}}, + ) + with pytest.warns(UnsafeArtifactWarning) as caught: + collected = collector.collect(only_exist=False) + assert len(caught) == 1 + recorder_id = artifacts.recorder.info["id"] + pd.testing.assert_series_equal( + collected["model"][recorder_id].predict(artifacts.dataset), artifacts.prediction, check_exact=True + ) + pd.testing.assert_frame_equal( + collected["pred"][recorder_id], artifacts.prediction.to_frame("score"), check_exact=True + ) + + collector.artifact_load_kwargs["model"]["trusted"] = True + with pytest.raises(ValueError, match="unpickler"): + collector.collect(only_exist=False) + + +@pytest.mark.parametrize("value", ["false", 0, 1, None, np.bool_(True)]) +def test_collector_rejects_invalid_consent_before_opening_an_experiment(monkeypatch, value): + get_exp = Mock() + monkeypatch.setattr(R, "get_exp", get_exp) + with pytest.raises(TypeError, match="must be a bool"): + RecorderCollector("unused", artifact_load_kwargs={"pred": {"trusted": value}}) + get_exp.assert_not_called() + + +@pytest.mark.parametrize( + "options,error", + [ + ([], TypeError), + (False, TypeError), + ({"pred": None}, TypeError), + ({"pred": []}, TypeError), + ({"unknown": {"trusted": True}}, ValueError), + ({"__raw": {"trusted": True}}, ValueError), + ], +) +def test_collector_rejects_invalid_loading_options(options, error): + with pytest.raises(error): + RecorderCollector(lambda: [], artifact_load_kwargs=options) + + +@pytest.mark.parametrize("options", [None, {"pred": {"trusted": False}}]) +def test_artifact_readers_preserve_legacy_default_recorder_signature(options): + expected = pd.DataFrame({"score": [0.5]}) + + class LegacyRecorder: + info = {"id": "legacy-reader"} + status = Recorder.STATUS_FI + + def load_object(self, name): + assert name == "pred.pkl" + return expected + + recorder = LegacyRecorder() + assert RecordTemp(recorder).load("pred.pkl", trusted=False) is expected + collector = RecorderCollector(lambda: [recorder], artifact_load_kwargs=options) + assert collector.collect(only_exist=False)["pred"]["legacy-reader"] is expected + + +def test_collector_raw_records_do_not_require_deserialization(reader_artifacts): + collector = RecorderCollector(lambda: [reader_artifacts.recorder], artifacts_key="__raw") + assert collector.collect()["__raw"][reader_artifacts.recorder.info["id"]] is reader_artifacts.recorder diff --git a/tests/security/test_artifact_trust_api.py b/tests/security/test_artifact_trust_api.py new file mode 100644 index 00000000000..d0014b5b98e --- /dev/null +++ b/tests/security/test_artifact_trust_api.py @@ -0,0 +1,147 @@ +import inspect +import pickle +import runpy +import warnings +from importlib import import_module +from pathlib import Path + +import numpy as np +import pytest +from fire.helptext import HelpText + +from qlib.utils.pickle_utils import ARTIFACT_MIGRATION_URL + + +WORKFLOWS = [ + ("qlib.workflow.online.strategy", "RollingStrategy"), + ("qlib.workflow.online.utils", "OnlineToolR"), + ("qlib.workflow.online.update", "RMDLoader"), + ("qlib.workflow.online.update", "PredUpdater"), + ("qlib.workflow.online.update", "LabelUpdater"), + ("qlib.model.trainer", "DelayTrainerR"), + ("qlib.model.trainer", "DelayTrainerRM"), + ("qlib.contrib.rolling.ddgda", "DDGDA"), +] + + +@pytest.fixture(params=WORKFLOWS, ids=[name for _, name in WORKFLOWS]) +def workflow_cls(request): + module, name = request.param + return getattr(import_module(module), name) + + +@pytest.mark.parametrize("consent", [False, True]) +def test_saved_consent_survives_restoring_a_workflow(workflow_cls, consent): + original = object.__new__(workflow_cls) + original.trusted = consent + original.marker = "preserved" + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + restored = pickle.loads(pickle.dumps(original)) + assert not caught + assert restored.trusted is consent + assert restored.__dict__ == {"trusted": consent, "marker": "preserved"} + + +def test_saved_workflow_without_consent_does_not_gain_it(workflow_cls): + original = object.__new__(workflow_cls) + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + restored = pickle.loads(pickle.dumps(original)) + assert not caught + assert restored.trusted is False + assert "trusted" not in restored.__dict__ + + +def test_old_online_manager_requires_consent_on_each_component(): + from qlib.model.trainer import DelayTrainerR + from qlib.workflow.online.manager import OnlineManager + from qlib.workflow.online.strategy import RollingStrategy + from qlib.workflow.online.utils import OnlineToolR + + manager = object.__new__(OnlineManager) + strategy = object.__new__(RollingStrategy) + strategy.tool = object.__new__(OnlineToolR) + manager.strategies = [strategy] + manager.trainer = object.__new__(DelayTrainerR) + + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + restored = pickle.loads(pickle.dumps(manager)) + + assert not caught + assert restored.strategies[0].trusted is False + assert restored.strategies[0].tool.trusted is False + assert restored.trainer.trusted is False + + restored.strategies[0].trusted = True + reloaded = pickle.loads(pickle.dumps(restored)) + assert reloaded.strategies[0].trusted is True + assert reloaded.strategies[0].tool.trusted is False + assert reloaded.trainer.trusted is False + + +def test_public_workflow_constructors_default_to_restricted_loading(workflow_cls): + parameters = inspect.signature(workflow_cls).parameters + assert parameters["trusted"].default is False + + +def _entry_points(): + from qlib.contrib.meta.data_selection.dataset import InternalData, MetaDatasetDS + from qlib.contrib.rolling.ddgda import DDGDA + from qlib.model.trainer import DelayTrainerR, DelayTrainerRM, end_task_train + from qlib.workflow import QlibRecorder + from qlib.workflow.online.strategy import RollingStrategy + from qlib.workflow.online.update import LabelUpdater, PredUpdater, RMDLoader + from qlib.workflow.online.utils import OnlineToolR + from qlib.workflow.recorder import MLflowRecorder, Recorder + from qlib.workflow.record_temp import RecordTemp + + return [ + (RMDLoader, {"rec": None}), + (PredUpdater, {"record": None}), + (LabelUpdater, {"record": None}), + (OnlineToolR, {}), + (RollingStrategy, {"name_id": "unused", "task_template": {}, "rolling_gen": None}), + (DelayTrainerR, {}), + (DelayTrainerRM, {}), + (DDGDA, {"conf_path": "unused.yaml"}), + (MetaDatasetDS, {"task_tpl": [], "step": 20, "exp_name": "unused", "segments": 0.5}), + (InternalData({}, 20, "unused").setup, {}), + (end_task_train, {"rec": None, "experiment_name": "unused"}), + (QlibRecorder(None).load_object, {"name": "unused"}), + (RecordTemp(None).load, {"name": "unused"}), + (MLflowRecorder.load_object, {}), + (Recorder.load_object, {}), + ] + + +def test_all_entry_points_use_the_same_public_keyword(): + for entry, _ in _entry_points(): + parameters = inspect.signature(entry).parameters + assert parameters["trusted"].default is False, entry + + +@pytest.mark.parametrize("value", ["false", 0, 1, None, np.bool_(True)]) +def test_workflow_entries_reject_non_boolean_consent_before_loading(value): + for entry, kwargs in _entry_points()[:-2]: + with pytest.raises(TypeError, match="`trusted` must be a bool") as caught: + entry(**kwargs, trusted=value) + assert ARTIFACT_MIGRATION_URL in str(caught.value) + + +@pytest.mark.parametrize( + "path,class_name", + [ + ("examples/online_srv/update_online_pred.py", "UpdatePredExample"), + ("examples/online_srv/online_management_simulate.py", "OnlineSimulationExample"), + ("examples/online_srv/rolling_online_management.py", "RollingOnlineExample"), + ("examples/benchmarks_dynamic/DDG-DA/workflow.py", "DDGDABench"), + ], +) +def test_example_cli_help_exposes_the_same_trust_flag(path, class_name): + namespace = runpy.run_path(str(Path(__file__).resolve().parents[2] / path)) + example = namespace[class_name] + assert inspect.signature(example).parameters["trusted"].default is False + text = HelpText(example) + assert "--trusted" in text diff --git a/tests/security/test_highfreq_artifact_paths.py b/tests/security/test_highfreq_artifact_paths.py new file mode 100644 index 00000000000..a9d88d493f1 --- /dev/null +++ b/tests/security/test_highfreq_artifact_paths.py @@ -0,0 +1,134 @@ +from pathlib import Path +import pickle +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest + +from qlib.contrib.data.highfreq_provider import HighFreqProvider +from qlib.utils.pickle_utils import ARTIFACT_MIGRATION_URL + + +def _provider(root): + provider = object.__new__(HighFreqProvider) + provider.artifact_root = Path(root).resolve() + return provider + + +def test_highfreq_artifact_path_stays_under_root(tmp_path): + provider = _provider(tmp_path) + assert provider._resolve_artifact_path("data/features.pkl") == tmp_path / "data/features.pkl" + + +def test_highfreq_artifact_path_rejects_parent_traversal(tmp_path): + provider = _provider(tmp_path / "artifacts") + with pytest.raises(ValueError, match="escapes artifact_root") as caught: + provider._resolve_artifact_path("../outside.pkl") + assert ARTIFACT_MIGRATION_URL in str(caught.value) + assert "dedicated trusted artifact_root" in str(caught.value) + + +def test_highfreq_artifact_path_rejects_absolute_path(tmp_path): + provider = _provider(tmp_path / "artifacts") + with pytest.raises(ValueError, match="escapes artifact_root"): + provider._resolve_artifact_path(tmp_path / "outside.pkl") + + +def _symlink(link, target): + try: + link.symlink_to(target) + except (OSError, NotImplementedError): + pytest.skip("Symlink creation is unavailable") + + +@pytest.mark.parametrize( + "method", ["_gen_data", "_gen_dataframe", "_gen_dataset", "_gen_day_dataset", "_gen_stock_dataset"] +) +def test_generation_rejects_escaping_path_without_mutating_config(tmp_path, method): + provider = _provider(tmp_path / "artifacts") + config = {"path": "../outside.pkl"} + args = (config, "feature") if method in ("_gen_day_dataset", "_gen_stock_dataset") else (config,) + with pytest.raises(ValueError, match="escapes artifact_root"): + getattr(provider, method)(*args) + assert config == {"path": "../outside.pkl"} + assert not (tmp_path / "outside.pkl").exists() + + +@pytest.mark.parametrize("method", ["_gen_day_dataset", "_gen_stock_dataset"]) +def test_temporary_dataset_rejects_symlink_escape(tmp_path, method): + root = tmp_path / "artifacts" + root.mkdir() + outside = tmp_path / "outside.pkl" + outside.write_bytes(b"must not be read") + _symlink(root / "tmp_dataset.pkl", outside) + with pytest.raises(ValueError, match="escapes artifact_root"): + getattr(_provider(root), method)({"path": "."}, "feature") + assert outside.read_bytes() == b"must not be read" + + +@pytest.mark.parametrize( + "method, filename", [("_gen_dataframe", "featurestrain.pkl"), ("get_pre_datasets", "features_train.pkl")] +) +def test_split_dataset_rejects_symlink_escape(tmp_path, method, filename): + root = tmp_path / "artifacts" + root.mkdir() + outside = tmp_path / "outside.pkl" + outside.write_bytes(b"must not be overwritten") + _symlink(root / filename, outside) + provider = _provider(root) + provider.feature_conf = {"path": "features.pkl"} + provider.label_conf = {"path": "labels.pkl"} + with pytest.raises(ValueError, match="escapes artifact_root"): + if method == "get_pre_datasets": + provider.get_pre_datasets() + else: + provider._gen_dataframe(provider.feature_conf) + assert outside.read_bytes() == b"must not be overwritten" + + +def test_cached_data_load_preserves_input_config(tmp_path): + expected = {"train": [1], "valid": [2], "test": [3]} + (tmp_path / "features.pkl").write_bytes(pickle.dumps(expected)) + provider = _provider(tmp_path) + provider.logger = Mock() + config = {"path": "features.pkl"} + assert provider._gen_data(config) == [[1], [2], [3]] + assert config == {"path": "features.pkl"} + + +@pytest.mark.parametrize( + "method, filename", + [ + ("_gen_day_dataset", "2024-01-01.pkl"), + ("_gen_stock_dataset", "SH600000.pkl"), + ("_gen_stock_dataset", "../outside.pkl"), + ], +) +def test_generated_dataset_rejects_escaping_filename(tmp_path, monkeypatch, method, filename): + import pandas as pd + from qlib.contrib.data import highfreq_provider as module + + root = tmp_path / "artifacts" + root.mkdir() + outside = tmp_path / "outside.pkl" + outside.write_bytes(b"must not be overwritten") + if filename != "../outside.pkl": + _symlink(root / filename, outside) + (root / "tmp_dataset.pkl").write_bytes(pickle.dumps(None)) + provider = _provider(root) + provider.logger = Mock() + provider.start_time = provider.end_time = "2024-01-01" + provider.freq = "1min" + monkeypatch.setattr( + module, + "D", + SimpleNamespace( + calendar=lambda **kwargs: [pd.Timestamp("2024-01-01")], + instruments=lambda **kwargs: [], + list_instruments=lambda **kwargs: [filename[:-4]], + ), + ) + monkeypatch.setattr(module, "Parallel", lambda **kwargs: lambda jobs: [fun(*args, **kw) for fun, args, kw in jobs]) + with pytest.raises(ValueError, match="escapes artifact_root"): + getattr(provider, method)({"path": "."}, "feature") + assert outside.read_bytes() == b"must not be overwritten" diff --git a/tests/security/test_hist_artifacts.py b/tests/security/test_hist_artifacts.py new file mode 100644 index 00000000000..f7a07669eb0 --- /dev/null +++ b/tests/security/test_hist_artifacts.py @@ -0,0 +1,63 @@ +import json +from unittest.mock import Mock + +import numpy as np +import pandas as pd +import pytest + +pytest.importorskip("torch") + +from qlib.contrib.model.pytorch_hist import HIST, _load_stock_index +from qlib.utils.pickle_utils import ARTIFACT_MIGRATION_URL + + +def test_load_stock_index_from_json(tmp_path): + path = tmp_path / "stock_index.json" + path.write_text(json.dumps({"SH600000": 0, "SZ000001": 1}), encoding="utf-8") + + assert _load_stock_index(path, upper_bound=2) == {"SH600000": 0, "SZ000001": 1} + + +@pytest.mark.parametrize("value", [-1, 2, 1.5, True, None]) +def test_load_stock_index_rejects_invalid_values(tmp_path, value): + path = tmp_path / "stock_index.json" + path.write_text(json.dumps({"SH600000": value}), encoding="utf-8") + + with pytest.raises(ValueError): + _load_stock_index(path, upper_bound=2) + + +def test_load_stock_index_rejects_object_npy(tmp_path): + path = tmp_path / "stock_index.npy" + np.save(path, {"SH600000": 0}, allow_pickle=True) + + with pytest.raises(ValueError, match="must be a JSON file") as caught: + _load_stock_index(path) + assert ARTIFACT_MIGRATION_URL in str(caught.value) + assert "restored model objects" in str(caught.value) + + +class _MarkerPayload: + def __init__(self, path): + self.path = str(path) + + def __reduce__(self): + return eval, (f"open({self.path!r}, 'w').write('executed')",) + + +@pytest.mark.parametrize("method", ["fit", "predict"]) +def test_hist_rejects_malicious_metadata_before_execution(tmp_path, method): + marker = tmp_path / "executed.txt" + metadata = tmp_path / "index.npy" + np.save(metadata, {"SH600000": _MarkerPayload(marker)}, allow_pickle=True) + concepts = tmp_path / "concepts.npy" + np.save(concepts, np.zeros((734, 1)), allow_pickle=False) + model = object.__new__(HIST) + model.stock_index = metadata + model.stock2concept = concepts + model.fitted = True + dataset = Mock() + dataset.prepare.return_value = [pd.DataFrame({"value": [1]})] * 3 + with pytest.raises(ValueError, match="must be a JSON file"): + getattr(model, method)(dataset) + assert not marker.exists() diff --git a/tests/security/test_mlflow_artifacts.py b/tests/security/test_mlflow_artifacts.py new file mode 100644 index 00000000000..ec7dfe5de48 --- /dev/null +++ b/tests/security/test_mlflow_artifacts.py @@ -0,0 +1,509 @@ +import os +import pickle +import warnings +from contextlib import nullcontext +from datetime import timedelta +from types import SimpleNamespace +from unittest.mock import Mock + +import pandas as pd +import numpy as np +import pytest + +from qlib.utils.exceptions import LoadObjectError +from qlib.utils.pickle_utils import ARTIFACT_MIGRATION_URL +from qlib.workflow.recorder import MLflowRecorder, Recorder, UnsafeArtifactWarning + + +class _TrackingClient: + @staticmethod + def _get_artifact_repo(_run_id): + return object() + + +class _ArtifactClient: + _tracking_client = _TrackingClient() + + def __init__(self, path): + self.path = path + + def download_artifacts(self, _run_id, _name): + return str(self.path) + + +class _CustomArtifact: + def __init__(self, value=42): + self.value = value + + +class _ExecutableArtifact: + def __reduce__(self): + return os.system, ("echo vulnerable",) + + +def _recorder(path): + recorder = object.__new__(MLflowRecorder) + recorder._uri = "file:///unused" + recorder.id = "run-id" + recorder.client = _ArtifactClient(path) + return recorder + + +@pytest.mark.parametrize("options", [{}, {"trusted": False}]) +def test_mlflow_artifact_uses_restricted_loading_without_fallback(tmp_path, monkeypatch, options): + from qlib.workflow import recorder as recorder_module + + path = tmp_path / "payload.pkl" + path.write_bytes(pickle.dumps(_CustomArtifact())) + unrestricted = Mock(side_effect=AssertionError("Unrestricted loading must not be attempted")) + monkeypatch.setattr(recorder_module.pickle, "Unpickler", unrestricted) + + with warnings.catch_warnings(record=True) as caught: + with pytest.raises(LoadObjectError, match="Forbidden class") as error: + _recorder(path).load_object("payload.pkl", **options) + + assert "payload.pkl" in str(error.value) + assert "trusted=True" in str(error.value) + assert "workflow entry point" in str(error.value) + assert ARTIFACT_MIGRATION_URL in str(error.value) + unrestricted.assert_not_called() + assert not any(issubclass(warning.category, UnsafeArtifactWarning) for warning in caught) + + +@pytest.mark.parametrize("protocol", [4, 5]) +def test_mlflow_artifact_loads_common_data_without_trusted_flag(tmp_path, protocol): + path = tmp_path / "frame.pkl" + expected = pd.DataFrame({"value": [1, 2]}, index=pd.date_range("2024-01-01", periods=2)) + path.write_bytes(pickle.dumps(expected, protocol=protocol)) + + actual = _recorder(path).load_object("frame.pkl") + + pd.testing.assert_frame_equal(actual, expected) + + +def test_mlflow_artifact_requires_explicit_trust_for_arbitrary_pickle(tmp_path): + path = tmp_path / "payload.pkl" + path.write_bytes(pickle.dumps(_CustomArtifact())) + recorder = _recorder(path) + + with pytest.warns(UnsafeArtifactWarning, match="may execute arbitrary code"): + actual = recorder.load_object("payload.pkl", trusted=True) + assert isinstance(actual, _CustomArtifact) + assert actual.value == 42 + with pytest.raises(LoadObjectError, match="Forbidden class"): + recorder.load_object("payload.pkl") + + +def test_custom_artifact_unpickler_requires_explicit_selection_and_warns(tmp_path): + path = tmp_path / "payload.pkl" + path.write_bytes(pickle.dumps(_CustomArtifact())) + with pytest.warns(UnsafeArtifactWarning, match="custom artifact unpickler"): + actual = _recorder(path).load_object("payload.pkl", unpickler=pickle.Unpickler) + assert isinstance(actual, _CustomArtifact) + assert actual.value == 42 + + +def test_artifact_loader_rejects_ambiguous_trust_options(tmp_path): + with pytest.raises(ValueError, match="cannot be used together"): + _recorder(tmp_path / "unused.pkl").load_object("unused.pkl", unpickler=pickle.Unpickler, trusted=True) + + +@pytest.mark.parametrize("trusted", [None, 0, 1, "False", "True", np.bool_(True)]) +def test_mlflow_artifact_requires_an_actual_boolean_before_download(tmp_path, trusted): + recorder = _recorder(tmp_path / "unused.pkl") + recorder.client.download_artifacts = Mock() + + with pytest.raises(TypeError, match="bool"): + recorder.load_object("unused.pkl", trusted=trusted) + + recorder.client.download_artifacts.assert_not_called() + + +def _facade(recorder): + from qlib.workflow import QlibRecorder + + experiment = SimpleNamespace(get_recorder=Mock(return_value=recorder)) + facade = object.__new__(QlibRecorder) + facade.get_exp = Mock(return_value=experiment) + return facade + + +@pytest.mark.parametrize("options", [{}, {"trusted": False}]) +def test_recorder_facade_preserves_legacy_subclass_signature(options): + class LegacyRecorder(Recorder): + def load_object(self, name): + return {"name": name} + + recorder = Mock(wraps=LegacyRecorder("experiment", "legacy")) + + assert _facade(recorder).load_object("data.pkl", **options) == {"name": "data.pkl"} + recorder.load_object.assert_called_once_with("data.pkl") + + +def test_recorder_facade_forwards_explicit_trust_to_modern_subclass(): + class ModernRecorder(Recorder): + def load_object(self, name, *, trusted=False): + return {"name": name, "trusted": trusted} + + recorder = Mock(wraps=ModernRecorder("experiment", "modern")) + + assert _facade(recorder).load_object("model.pkl", trusted=True) == {"name": "model.pkl", "trusted": True} + recorder.load_object.assert_called_once_with("model.pkl", trusted=True) + + +@pytest.fixture +def mlflow_recorders(tmp_path): + from mlflow.tracking import MlflowClient + + uri = (tmp_path / "mlruns").as_uri() + client = MlflowClient(tracking_uri=uri) + experiment_id = client.create_experiment("artifact-trust") + run = client.create_run(experiment_id, tags={"mlflow.runName": "roundtrip"}) + writer = MLflowRecorder(experiment_id, uri, mlflow_run=run) + reader = MLflowRecorder(experiment_id, uri, mlflow_run=client.get_run(run.info.run_id)) + try: + yield writer, reader + finally: + client.set_terminated(run.info.run_id) + + +@pytest.mark.parametrize( + "value", + [ + pd.Series([1, 2], index=pd.period_range("2024-01", periods=2, freq="M")), + pd.Series([1, 2], index=pd.IntervalIndex.from_breaks([0, 1, 2])), + pd.Series([0.0, 1.0, 0.0], dtype=pd.SparseDtype("float64", 0)), + np.ma.array([1, 2, 3], mask=[False, True, False]), + ], +) +def test_real_mlflow_store_roundtrips_data_artifacts(mlflow_recorders, value): + writer, reader = mlflow_recorders + writer.save_objects(**{"data.pkl": value}) + # A new recorder/client must download and deserialize the stored artifact. + actual = reader.load_object("data.pkl") + if isinstance(value, pd.Series): + pd.testing.assert_series_equal(actual, value) + else: + np.testing.assert_array_equal(actual.data, value.data) + np.testing.assert_array_equal(actual.mask, value.mask) + + +@pytest.mark.parametrize("protocol", [4, 5]) +@pytest.mark.parametrize("multi_index", [False, True], ids=["datetime-index", "multi-index"]) +@pytest.mark.parametrize( + "offset", + [ + pytest.param( + pd.offsets.BusinessHour( + n=2, start=["08:30", "13:00"], end=["11:30", "16:00"], offset=timedelta(minutes=15) + ), + id="split-business-hours", + ), + pytest.param( + pd.offsets.CustomBusinessDay( + n=2, + weekmask="Mon Tue Thu Fri", + holidays=["2024-01-04", "2024-01-15"], + offset=timedelta(hours=1, minutes=15), + ), + id="custom-business-days", + ), + ], +) +def test_real_mlflow_store_roundtrips_business_frequency_predictions( + mlflow_recorders, monkeypatch, offset, multi_index, protocol +): + from qlib.config import C + + monkeypatch.setitem(C, "dump_protocol_version", protocol) + dates = pd.date_range("2024-01-01 08:30", periods=4, freq=offset, name="datetime") + index = ( + pd.MultiIndex.from_product([dates, ["SH600000", "SH600004"]], names=["datetime", "instrument"]) + if multi_index + else dates + ) + expected = pd.DataFrame( + { + "score": np.resize([1.25, np.nan], len(index)), + "volume": pd.array(np.resize([100, None], len(index)), dtype="Int64"), + }, + index=index, + ) + original = expected.copy(deep=True) + frequency_args = offset.__reduce__()[1] + writer, reader = mlflow_recorders + + writer.save_objects(**{"pred.pkl": expected}) + actual = reader.load_object("pred.pkl") + + pd.testing.assert_frame_equal(actual, original, check_exact=True, check_freq=True) + actual_dates = actual.index.levels[0] if multi_index else actual.index + pd.testing.assert_index_equal(actual_dates, dates, exact=True, check_exact=True) + assert type(actual_dates.freq) is type(offset) + assert actual_dates.freq.__reduce__()[1] == frequency_args + actual.iloc[0, 0] = -100.0 + pd.testing.assert_frame_equal(expected, original, check_exact=True, check_freq=True) + assert offset.__reduce__()[1] == frequency_args + + +@pytest.mark.parametrize("protocol", [4, 5]) +@pytest.mark.parametrize("record_index", [None, 0, 1], ids=["recarray", "record", "missing-record"]) +def test_real_mlflow_store_roundtrips_records(mlflow_recorders, monkeypatch, protocol, record_index): + from qlib.config import C + + monkeypatch.setitem(C, "dump_protocol_version", protocol) + records = np.array( + [("SH600000", 1.25, 100, "2024-01-01"), ("SH600004", np.nan, 0, "NaT")], + dtype=[("instrument", "U8"), ("score", " 1: self.assertIn("set -euo pipefail", commands) + def test_downloads_use_explicit_release_archives_without_version_probes(self): + expected = { + "test_qlib_from_source.yml": ["v2/qlib_data_simple_cn_1d_latest.zip", "rl_data.zip"], + "test_qlib_from_source_slow.yml": ["v2/qlib_data_simple_cn_1d_latest.zip"], + "test_qlib_from_pip.yml": ["v2/qlib_data_cn_1d_latest.zip"], + } + for name, workflow in self.workflows.items(): + with self.subTest(workflow=name): + archives = [] + for step in workflow["jobs"]["build"]["steps"]: + if ( + "data download" not in step["name"].lower() + and "downloads dependencies" not in step["name"].lower() + ): + continue + for line in step["with"]["command"].splitlines(): + args = shlex.split(line) + self.assertNotIn("qlib_data", args) + if "download_data" in args: + archives.append(args[args.index("--file_name") + 1]) + target = args[args.index("--target_dir") + 1] + self.assertEqual( + target, + "tests/.data/rl" if archives[-1] == "rl_data.zip" else "~/.qlib/qlib_data/cn_data", + ) + self.assertEqual(archives, expected[name]) + def test_pypi_workflow_does_not_import_the_checkout(self): steps = self.workflows["test_qlib_from_pip.yml"]["jobs"]["build"]["steps"] for name in ("Downloads dependencies data", "Test workflow by config"): diff --git a/tests/test_pit.py b/tests/test_pit.py index 1c2228f4733..886dcaee8c9 100644 --- a/tests/test_pit.py +++ b/tests/test_pit.py @@ -9,6 +9,7 @@ import pytest import pandas as pd from pathlib import Path +from unittest.mock import Mock from qlib.data import D from qlib.tests.data import GetData @@ -17,7 +18,7 @@ from dump_pit import DumpPitData sys.path.append(str(Path(__file__).resolve().parent.parent.joinpath("scripts/data_collector/pit"))) -from collector import Run +from collector import PitNormalize, Run pd.set_option("display.width", 1000) pd.set_option("display.max_columns", None) @@ -29,6 +30,24 @@ QLIB_DIR.mkdir(exist_ok=True, parents=True) +@pytest.mark.parametrize( + ("interval", "period", "missing_date"), + [ + ("quarterly", 202304, "2024-02-14"), + ("annual", 2023, "2024-03-30"), + ], +) +def test_pit_normalization_does_not_fetch_trading_calendar(monkeypatch, interval, period, missing_date): + login = Mock(side_effect=AssertionError("PIT normalization must not request an online trading calendar")) + monkeypatch.setattr("data_collector.utils.bs.login", login) + normalizer = PitNormalize(interval=interval) + source = pd.DataFrame({"date": [None, "2024-01-20"], "period": ["2023-12-31", "2023-12-31"], "value": [0.1, 0.2]}) + expected = pd.DataFrame({"date": [missing_date, "2024-01-20"], "period": [period, period], "value": [0.1, 0.2]}) + + pd.testing.assert_frame_equal(normalizer.normalize(source), expected) + login.assert_not_called() + + class TestPIT(unittest.TestCase): @classmethod def tearDownClass(cls) -> None: