Raise TimeoutError.

Author: tillahoffmann

cmdstanpy/model.py

@@ -733,6 +733,7 @@ def optimize(
                 show_console=show_console,
                 timeout=timeout,
             )
+        runset.raise_for_timeouts()
 
         if not runset._check_retcodes():
             msg = 'Error during optimization: {}'.format(runset.get_err_msgs())
@@ -1167,6 +1168,8 @@ def sample(
                 sys.stdout.write('\n')
                 get_logger().info('CmdStan done processing.')
 
+            runset.raise_for_timeouts()
+
             get_logger().debug('runset\n%s', repr(runset))
 
             # hack needed to parse CSV files if model has no params
@@ -1348,6 +1351,7 @@ def generate_quantities(
                         timeout=timeout,
                     )
 
+            runset.raise_for_timeouts()
             errors = runset.get_err_msgs()
             if errors:
                 msg = (
@@ -1518,6 +1522,7 @@ def variational(
                 show_console=show_console,
                 timeout=timeout,
             )
+        runset.raise_for_timeouts()
 
         # treat failure to converge as failure
         transcript_file = runset.stdout_files[dummy_chain_id]
@@ -1606,7 +1611,15 @@ def _run_cmdstan(
                 universal_newlines=True,
             )
             if timeout:
-                timer = threading.Timer(timeout, proc.terminate)
+
+                def _timer_target():
+                    # Abort if the process has already terminated.
+                    if proc.poll() is not None:
+                        return
+                    proc.terminate()
+                    runset._set_timeout_flag(idx, True)
+
+                timer = threading.Timer(timeout, _timer_target)
                 timer.setDaemon(True)
                 timer.start()
             else:
@@ -1623,11 +1636,9 @@ def _run_cmdstan(
 
             stdout, _ = proc.communicate()
             retcode = proc.returncode
+            runset._set_retcode(idx, retcode)
             if timer:
                 timer.cancel()
-                if retcode == -15:
-                    retcode = 60
-            runset._set_retcode(idx, retcode)
 
             if stdout:
                 fd_out.write(stdout)

cmdstanpy/stanfit/runset.py

@@ -44,6 +44,7 @@ def __init__(
         else:
             self._num_procs = 1
         self._retcodes = [-1 for _ in range(self._num_procs)]
+        self._timeout_flags = [False for _ in range(self._num_procs)]
         if chain_ids is None:
             chain_ids = [i + 1 for i in range(chains)]
         self._chain_ids = chain_ids
@@ -230,12 +231,14 @@ def _set_retcode(self, idx: int, val: int) -> None:
         """Set retcode at process[idx] to val."""
         self._retcodes[idx] = val
 
+    def _set_timeout_flag(self, idx: int, val: bool) -> None:
+        """Set timeout_flag at process[idx] to val."""
+        self._timeout_flags[idx] = val
+
     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
@@ -296,3 +299,10 @@ def save_csvfiles(self, dir: Optional[str] = None) -> None:
                 raise ValueError(
                     'Cannot save to file: {}'.format(to_path)
                 ) from e
+
+    def raise_for_timeouts(self) -> None:
+        if any(self._timeout_flags):
+            raise TimeoutError(
+                f"{sum(self._timeout_flags)} of {self.num_procs} processes "
+                "timed out"
+            )

test/test_generate_quantities.py

@@ -480,7 +480,7 @@ 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)
-        with self.assertRaisesRegex(RuntimeError, 'processing timed out'):
+        with self.assertRaises(TimeoutError):
             timeout_model.generate_quantities(
                 timeout=0.1, mcmc_sample=fit, data={'loop': 1}
             )

test/test_optimize.py

@@ -637,7 +637,7 @@ def test_attrs(self):
     def test_timeout(self):
         stan = os.path.join(DATAFILES_PATH, 'timeout.stan')
         timeout_model = CmdStanModel(stan_file=stan)
-        with self.assertRaisesRegex(RuntimeError, 'processing timed out'):
+        with self.assertRaises(TimeoutError):
             timeout_model.optimize(data={'loop': 1}, timeout=0.1)
 
 

test/test_sample.py

@@ -1916,7 +1916,7 @@ def test_diagnostics(self):
     def test_timeout(self):
         stan = os.path.join(DATAFILES_PATH, 'timeout.stan')
         timeout_model = CmdStanModel(stan_file=stan)
-        with self.assertRaisesRegex(RuntimeError, 'processing timed out'):
+        with self.assertRaises(TimeoutError):
             timeout_model.sample(timeout=0.1, chains=1, data={'loop': 1})
 
 

test/test_variational.py

@@ -294,7 +294,7 @@ def test_attrs(self):
     def test_timeout(self):
         stan = os.path.join(DATAFILES_PATH, 'timeout.stan')
         timeout_model = CmdStanModel(stan_file=stan)
-        with self.assertRaisesRegex(RuntimeError, 'processing timed out'):
+        with self.assertRaises(TimeoutError):
             timeout_model.variational(timeout=0.1, data={'loop': 1})