Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions .github/workflows/pattern-check.yml
Original file line number Diff line number Diff line change
Expand Up @@ -51,8 +51,8 @@ jobs:
'\.gluon/'
'gluon-runner'
'gluon-app\[bot\]'
# Internal package registries
'nexus\.alm\.europe\.cloudcenter\.corp'
# Internal package registries / internal TLD suffixes (generic)
'\.corp\b'
'artifactory\.santander\.'
# Internal IPv4 ranges (RFC1918)
'\b10\.[0-9]+\.[0-9]+\.[0-9]+\b'
Expand Down
6 changes: 6 additions & 0 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -244,6 +244,12 @@ Training writes these files into `run.output_dir`:
- `network.json`
- `model.pkl`

> **Security note:** `model.pkl` is a Python
> [pickle](https://docs.python.org/3/library/pickle.html) file. Loading a
> pickle can execute arbitrary code, so call `AutoBayesModel.load()` only on
> model directories you created yourself or obtained from a source you fully
> trust.

## Limits

- Each table can have at most one parent relation.
Expand Down
14 changes: 13 additions & 1 deletion src/auto_bayesian/model.py
Original file line number Diff line number Diff line change
Expand Up @@ -148,6 +148,11 @@ def save(self, output_dir: str | Path) -> None:

Writes ``model.pkl``, ``metrics.json``, ``network.json``, and (when
available) ``materialized.parquet`` into ``output_dir``.

.. note::
``model.pkl`` is a :mod:`pickle` file. Treat saved model
directories as trusted artifacts and share them only through
channels where they cannot be tampered with.
"""
output_path = Path(output_dir)
output_path.mkdir(parents=True, exist_ok=True)
Expand Down Expand Up @@ -208,7 +213,14 @@ def save(self, output_dir: str | Path) -> None:

@classmethod
def load(cls, output_dir: str | Path) -> "AutoBayesModel":
"""Load a model previously written by :meth:`save`."""
"""Load a model previously written by :meth:`save`.

.. warning::
``model.pkl`` is deserialized with :mod:`pickle`, which can
execute arbitrary code embedded in a malicious file. Only load
model directories that you created yourself or that come from a
source you fully trust.
"""
output_path = Path(output_dir)
with (output_path / "model.pkl").open("rb") as handle:
payload = pickle.load(handle)
Expand Down
Loading