Add timeout parameter.

Author: tillahoffmann

cmdstanpy/model.py

@@ -13,6 +13,7 @@
 from io import StringIO
 from multiprocessing import cpu_count
 from pathlib import Path
+import threading
 from typing import Any, Callable, Dict, Iterable, List, Mapping, Optional, Union
 
 import ujson as json
@@ -593,6 +594,7 @@ def optimize(
         show_console: bool = False,
         refresh: Optional[int] = None,
         time_fmt: str = "%Y%m%d%H%M%S",
+        timeout: Optional[float] = None,
     ) -> CmdStanMLE:
         """
         Run the specified CmdStan optimize algorithm to produce a
@@ -692,6 +694,8 @@ def optimize(
             :meth:`~datetime.datetime.strftime` to decide the file names for
             output CSVs. Defaults to "%Y%m%d%H%M%S"
 
+        :param timeout: Duration at which optimization times out in seconds.
+
         :return: CmdStanMLE object
         """
         optimize_args = OptimizeArgs(
@@ -723,7 +727,8 @@ def optimize(
             )
             dummy_chain_id = 0
             runset = RunSet(args=args, chains=1, time_fmt=time_fmt)
-            self._run_cmdstan(runset, dummy_chain_id, show_console=show_console)
+            self._run_cmdstan(runset, dummy_chain_id, show_console=show_console,
+                              timeout=timeout)
 
         if not runset._check_retcodes():
             msg = 'Error during optimization: {}'.format(runset.get_err_msgs())
@@ -767,6 +772,7 @@ def sample(
         show_console: bool = False,
         refresh: Optional[int] = None,
         time_fmt: str = "%Y%m%d%H%M%S",
+        timeout: Optional[float] = None,
         *,
         force_one_process_per_chain: Optional[bool] = None,
     ) -> CmdStanMCMC:
@@ -964,6 +970,8 @@ def sample(
             model was compiled with STAN_THREADS=True, and utilize the
             parallel chain functionality if those conditions are met.
 
+        :param timeout: Duration at which sampling times out in seconds.
+
         :return: CmdStanMCMC object
         """
         if fixed_param is None:
@@ -1139,6 +1147,7 @@ def sample(
                         show_progress=show_progress,
                         show_console=show_console,
                         progress_hook=progress_hook,
+                        timeout=timeout,
                     )
             if show_progress and progress_hook is not None:
                 progress_hook("Done", -1)  # -1 == all chains finished
@@ -1209,6 +1218,7 @@ def generate_quantities(
         show_console: bool = False,
         refresh: Optional[int] = None,
         time_fmt: str = "%Y%m%d%H%M%S",
+        timeout: Optional[float] = None,
     ) -> CmdStanGQ:
         """
         Run CmdStan's generate_quantities method which runs the generated
@@ -1267,6 +1277,8 @@ def generate_quantities(
             :meth:`~datetime.datetime.strftime` to decide the file names for
             output CSVs. Defaults to "%Y%m%d%H%M%S"
 
+        :param timeout: Duration at which generation times out in seconds.
+
         :return: CmdStanGQ object
         """
         if isinstance(mcmc_sample, CmdStanMCMC):
@@ -1329,6 +1341,7 @@ def generate_quantities(
                         runset,
                         i,
                         show_console=show_console,
+                        timeout=timeout,
                     )
 
             errors = runset.get_err_msgs()
@@ -1366,6 +1379,7 @@ def variational(
         show_console: bool = False,
         refresh: Optional[int] = None,
         time_fmt: str = "%Y%m%d%H%M%S",
+        timeout: Optional[float] = None,
     ) -> CmdStanVB:
         """
         Run CmdStan's variational inference algorithm to approximate
@@ -1458,6 +1472,9 @@ def variational(
             :meth:`~datetime.datetime.strftime` to decide the file names for
             output CSVs. Defaults to "%Y%m%d%H%M%S"
 
+        :param timeout: Duration at which variational Bayesian inference times
+            out in seconds.
+
         :return: CmdStanVB object
         """
         variational_args = VariationalArgs(
@@ -1491,7 +1508,8 @@ def variational(
 
             dummy_chain_id = 0
             runset = RunSet(args=args, chains=1, time_fmt=time_fmt)
-            self._run_cmdstan(runset, dummy_chain_id, show_console=show_console)
+            self._run_cmdstan(runset, dummy_chain_id, show_console=show_console,
+                              timeout=timeout)
 
         # treat failure to converge as failure
         transcript_file = runset.stdout_files[dummy_chain_id]
@@ -1541,6 +1559,7 @@ def _run_cmdstan(
         show_progress: bool = False,
         show_console: bool = False,
         progress_hook: Optional[Callable[[str, int], None]] = None,
+        timeout: Optional[float] = None,
     ) -> None:
         """
         Helper function which encapsulates call to CmdStan.
@@ -1577,6 +1596,12 @@ def _run_cmdstan(
                 env=os.environ,
                 universal_newlines=True,
             )
+            if timeout:
+                timer = threading.Timer(timeout, proc.terminate)
+                timer.setDaemon(True)
+                timer.start()
+            else:
+                timer = None
             while proc.poll() is None:
                 if proc.stdout is not None:
                     line = proc.stdout.readline()
@@ -1589,6 +1614,8 @@ def _run_cmdstan(
 
             stdout, _ = proc.communicate()
             retcode = proc.returncode
+            if timer and retcode == -15:
+                retcode = 60
             runset._set_retcode(idx, retcode)
 
             if stdout:

cmdstanpy/stanfit/runset.py

@@ -234,6 +234,8 @@ def get_err_msgs(self) -> str:
         """Checks console messages for each CmdStan run."""
         msgs = []
         for i in range(self._num_procs):
+            if self._retcodes[i] == 60:
+                msgs.append("processing timed out")
             if (
                 os.path.exists(self._stdout_files[i])
                 and os.stat(self._stdout_files[i]).st_size > 0

test/data/timeout.stan

@@ -0,0 +1,21 @@
+data {
+    // Indicator for endless looping.
+    int loop;
+}
+
+transformed data {
+    // Maybe loop forever so the model times out.
+    real y = 1;
+    while(loop && y) {
+        y += 1;
+    }
+}
+
+parameters {
+    real x;
+}
+
+model {
+    // A nice model so we can get a fit for the `generated_quantities` call.
+    x ~ normal(0, 1);
+}

test/test_generate_quantities.py

@@ -476,6 +476,16 @@ def test_attrs(self):
         with self.assertRaisesRegex(AttributeError, 'Unknown variable name:'):
             dummy = fit.c
 
+    def test_timeout(self):
+        stan = os.path.join(DATAFILES_PATH, 'timeout.stan')
+        timeout_model = CmdStanModel(stan_file=stan)
+        fit = timeout_model.sample(data={'loop': 0}, chains=1, iter_sampling=10)
+        self.assertRaisesRegex(
+            RuntimeError, 'processing timed out',
+            timeout_model.generate_quantities, timeout=0.1,
+            mcmc_sample=fit, data={'loop': 1},
+        )
+
 
 if __name__ == '__main__':
     unittest.main()

test/test_optimize.py

@@ -634,6 +634,14 @@ def test_attrs(self):
         with self.assertRaisesRegex(AttributeError, 'Unknown variable name:'):
             dummy = fit.c
 
+    def test_timeout(self):
+        stan = os.path.join(DATAFILES_PATH, 'timeout.stan')
+        timeout_model = CmdStanModel(stan_file=stan)
+        self.assertRaisesRegex(
+            RuntimeError, 'processing timed out', timeout_model.optimize,
+            data={'loop': 1}, timeout=0.1,
+        )
+
 
 if __name__ == '__main__':
     unittest.main()

test/test_sample.py

@@ -1913,6 +1913,14 @@ def test_diagnostics(self):
         self.assertEqual(fit.max_treedepths, None)
         self.assertEqual(fit.divergences, None)
 
+    def test_timeout(self):
+        stan = os.path.join(DATAFILES_PATH, 'timeout.stan')
+        timeout_model = CmdStanModel(stan_file=stan)
+        self.assertRaisesRegex(
+            RuntimeError, 'processing timed out', timeout_model.sample,
+            timeout=0.1, chains=1, data={'loop': 1},
+        )
+
 
 if __name__ == '__main__':
     unittest.main()

test/test_variational.py

@@ -291,6 +291,14 @@ def test_attrs(self):
         with self.assertRaisesRegex(AttributeError, 'Unknown variable name:'):
             dummy = fit.c
 
+    def test_timeout(self):
+        stan = os.path.join(DATAFILES_PATH, 'timeout.stan')
+        timeout_model = CmdStanModel(stan_file=stan)
+        self.assertRaisesRegex(
+            RuntimeError, 'processing timed out', timeout_model.variational,
+            timeout=0.1, data={'loop': 1}, show_console=True,
+        )
+
 
 if __name__ == '__main__':
     unittest.main()