Add warnings and protections around loading pickles

Author: ae-foster

README.md

@@ -1,5 +1,7 @@
 ![header](envelope.png)
 
+**Warning: Orbformer checkpoints are stored using `pickle`. Never read a checkpoint from an untrusted source.**
+
 # OneQMC
 
 This package provides an implementation of the [Orbformer wave function foundation model](https://arxiv.org/abs/2506.19960).
@@ -61,6 +63,9 @@ python scripts/transferable.py -d <subdirectory of ./data> -n <number of trainin
 We recommend using distinct output directories for every training run.
 Regarding other optional arguments, run `python scripts/transferable.py -h` for more information.
 
+**Note**: Checkpoints are stored using `pickle`. This means that opening a checkpoint from an untrusted source is a major security risk. Only ever open checkpoint files from trusted sources. To prevent reading untrusted pickle files, checkpoint reading is disabled by default and can be
+re-enabled by settings the environment variables `ORBFORMER_PICKLE_LOADING=1`.
+
 ### Preparing new structure data for fine-tuning
 
 To create a new dataset for fine-tuning, we recommend using [qcelemental](https://github.com/MolSSI/QCElemental) format.

src/oneqmc/entrypoint.py

@@ -18,6 +18,8 @@
 
 
 def load_chkpt_file(chkpt: str, discard_sampler_state: bool) -> Tuple[TrainState, int]:
+    if not os.environ["ORBFORMER_PICKLE_LOADING"] == "1":
+        raise PermissionError("Loading pickle files is disable for security. Set ORBFORMER_PICKLE_LOADING=1 to allow")
     with open(chkpt, "rb") as chkpt_file:
         init_step, (smpl_state, param_state, opt_state) = pickle.load(chkpt_file)
         if discard_sampler_state:
@@ -28,6 +30,8 @@ def load_chkpt_file(chkpt: str, discard_sampler_state: bool) -> Tuple[TrainState
 
 
 def load_density_chkpt_file(chkpt: str, discard_sampler_state: bool) -> Tuple[Tuple, int]:
+    if not os.environ["ORBFORMER_PICKLE_LOADING"] == "1":
+        raise PermissionError("Loading pickle files is disable for security. Set ORBFORMER_PICKLE_LOADING=1 to allow")
     with open(chkpt, "rb") as chkpt_file:
         init_step, (param_state, opt_state) = pickle.load(chkpt_file)
     return (param_state, opt_state), init_step

src/oneqmc/log.py

@@ -133,6 +133,8 @@ def close(self):
     def last(self):
         step_fast, step_slow = -1, -1  # account for the case where a queue is not initialized
         while self.fast_chkpts:
+            if not os.environ["ORBFORMER_PICKLE_LOADING"] == "1":
+                raise PermissionError("Loading pickle files is disable for security. Set ORBFORMER_PICKLE_LOADING=1 to allow")
             with self.fast_chkpts.pop(-1).path.open("rb") as f:
                 step_fast, last_chkpt_fast = pickle.load(f)
             if not jax.tree.reduce(
@@ -142,6 +144,8 @@ def last(self):
             ):
                 break
         while self.slow_chkpts:
+            if not os.environ["ORBFORMER_PICKLE_LOADING"] == "1":
+                raise PermissionError("Loading pickle files is disable for security. Set ORBFORMER_PICKLE_LOADING=1 to allow")
             with self.slow_chkpts.pop(-1).path.open("rb") as f:
                 step_slow, last_chkpt_slow = pickle.load(f)
             if not jax.tree.reduce(

tests/integration_tests/test_scripts.py

@@ -52,6 +52,7 @@ def runner(extra_args):
             ],
             cwd=project_root,
             capture_output=True,
+            env={"ORBFORMER_PICKLE_LOADING": "1"}
         )
         if result.returncode != 0:
             raise OneQMCProcessError(result.stderr.decode())
@@ -90,6 +91,7 @@ def runner(extra_args):
             ],
             cwd=project_root,
             capture_output=True,
+            env={"ORBFORMER_PICKLE_LOADING": "1"}
         )
         if result.returncode != 0:
             raise OneQMCProcessError(result.stderr.decode())