checkpointing

Author: mitzimorris

cmdstanpy/install_cmdstan.py

@@ -32,8 +32,6 @@
 
 from cmdstanpy import _DOT_CMDSTAN, _DOT_CMDSTANPY
 
-from cmdstanpy.progress import allow_show_progress, disable_progress
-
 from cmdstanpy.utils import (
     cmdstan_path,
     do_command,
@@ -43,6 +41,8 @@
     wrap_url_progress_hook,
 )
 
+from . import progress as progbar
+
 MAKE = os.getenv(
     'MAKE', 'make' if platform.system() != 'Windows' else 'mingw32-make'
 )
@@ -112,7 +112,7 @@ def build(verbose: bool = False, progress: bool = True) -> None:
     try:
         if verbose:
             do_command(cmd)
-        elif progress and allow_show_progress():
+        elif progress and progbar.allow_show_progress():
             progress_hook: Any = _wrap_build_progress_hook()
             do_command(cmd, fd_out=None, pbar=progress_hook)
         else:
@@ -146,43 +146,33 @@ def build(verbose: bool = False, progress: bool = True) -> None:
         )
 
 
-# pylint: disable=no-self-use
+@progbar.wrap_callback
 def _wrap_build_progress_hook() -> Optional[Callable[[str], None]]:
     """Sets up tqdm callback for CmdStan sampler console msgs."""
     pad = ' ' * 20
     msgs_expected = 150  # hack: 2.27 make build send ~140 msgs to console
-    try:
-        pbar: tqdm = tqdm(
-            total=msgs_expected,
-            bar_format="{desc} ({elapsed}) | {bar} | {postfix[0][value]}",
-            postfix=[dict(value=f'Building CmdStan {pad}')],
-            colour='blue',
-            desc='',
-            position=0,
-        )
-    # pylint: disable=broad-except
-    except Exception as e:
-        disable_progress(e)
-
-        # pylint: disable=unused-argument
-        def build_progress_hook(line: str) -> None:
-            return
-
-    else:
+    pbar: tqdm = tqdm(
+        total=msgs_expected,
+        bar_format="{desc} ({elapsed}) | {bar} | {postfix[0][value]}",
+        postfix=[dict(value=f'Building CmdStan {pad}')],
+        colour='blue',
+        desc='',
+        position=0,
+    )
 
-        def build_progress_hook(line: str) -> None:
-            if line.startswith('--- CmdStan'):
-                pbar.set_description('Done')
+    def build_progress_hook(line: str) -> None:
+        if line.startswith('--- CmdStan'):
+            pbar.set_description('Done')
+            pbar.postfix[0]["value"] = line
+            pbar.update(msgs_expected - pbar.n)
+            pbar.close()
+        else:
+            if line.startswith('--'):
                 pbar.postfix[0]["value"] = line
-                pbar.update(msgs_expected - pbar.n)
-                pbar.close()
             else:
-                if line.startswith('--'):
-                    pbar.postfix[0]["value"] = line
-                else:
-                    pbar.postfix[0]["value"] = f'{line[:8]} ... {line[-20:]}'
-                    pbar.set_description('Compiling')
-                    pbar.update(1)
+                pbar.postfix[0]["value"] = f'{line[:8]} ... {line[-20:]}'
+                pbar.set_description('Compiling')
+                pbar.update(1)
 
     return build_progress_hook
 
@@ -336,7 +326,7 @@ def retrieve_version(version: str, progress: bool = True) -> None:
     )
     for i in range(6):  # always retry to allow for transient URLErrors
         try:
-            if progress and allow_show_progress():
+            if progress and progbar.allow_show_progress():
                 progress_hook: Optional[
                     Callable[[int, int, int], None]
                 ] = wrap_url_progress_hook()
@@ -386,7 +376,7 @@ def retrieve_version(version: str, progress: bool = True) -> None:
             # fixes long-path limitation on Windows
             target = r'\\?\{}'.format(target)
 
-        if progress and allow_show_progress():
+        if progress and progbar.allow_show_progress():
             for member in tqdm(
                 iterable=tar.getmembers(),
                 total=len(tar.getmembers()),

cmdstanpy/model.py

@@ -16,23 +16,22 @@
 
 from tqdm.auto import tqdm  # type: ignore
 
-
 from cmdstanpy import (
     _CMDSTAN_REFRESH,
     _CMDSTAN_SAMPLING,
     _CMDSTAN_WARMUP,
 )
 
-from cmdstanpy.progress import allow_show_progress, disable_progress
-
 from cmdstanpy.cmdstan_args import (
     CmdStanArgs,
     GenerateQuantitiesArgs,
     OptimizeArgs,
     SamplerArgs,
     VariationalArgs,
 )
+
 from cmdstanpy.compiler_opts import CompilerOptions
+
 from cmdstanpy.stanfit import (
     CmdStanGQ,
     CmdStanMCMC,
@@ -41,6 +40,7 @@
     RunSet,
     from_csv,
 )
+
 from cmdstanpy.utils import (
     EXTENSION,
     MaybeDictToFilePath,
@@ -51,6 +51,8 @@
     returncode_msg,
 )
 
+from . import progress as progbar
+
 
 class CmdStanModel:
     # overview, omitted from doc comment in order to improve Sphinx docs.
@@ -905,7 +907,7 @@ def sample(
             if show_console:
                 show_progress = False
             else:
-                show_progress = show_progress and allow_show_progress()
+                show_progress = show_progress and progbar.allow_show_progress()
 
             get_logger().info('sampling: %s', runset.cmds[0])
             with ThreadPoolExecutor(max_workers=parallel_chains) as executor:
@@ -918,7 +920,7 @@ def sample(
                         show_console,
                         iter_total,
                     )
-            if show_progress and allow_show_progress():
+            if show_progress and progbar.allow_show_progress():
                 # advance terminal window cursor past progress bars
                 term_size: os.terminal_size = shutil.get_terminal_size(
                     fallback=(80, 24)
@@ -1283,7 +1285,7 @@ def _run_cmdstan(
         get_logger().debug(
             'threads: %s', str(os.environ.get('STAN_NUM_THREADS'))
         )
-        if show_progress and allow_show_progress():
+        if show_progress and progbar.allow_show_progress():
             progress_hook: Optional[
                 Callable[[str], None]
             ] = self._wrap_sampler_progress_hook(idx + 1, iter_total)
@@ -1352,37 +1354,28 @@ def _run_cmdstan(
                     sampler_args.fixed_param = True
 
     # pylint: disable=no-self-use
+    @progbar.wrap_callback
     def _wrap_sampler_progress_hook(
         self, chain_id: int, total: int
     ) -> Optional[Callable[[str], None]]:
         """Sets up tqdm callback for CmdStan sampler console msgs."""
-        try:
-            pbar: tqdm = tqdm(
-                total=total,
-                bar_format="{desc} |{bar}| {elapsed} {postfix[0][value]}",
-                postfix=[dict(value="Status")],
-                desc=f'chain {chain_id}',
-                colour='yellow',
-            )
-        # pylint: disable=broad-except
-        except Exception as e:
-            disable_progress(e)
-
-            # pylint: disable=unused-argument
-            def sampler_progress_hook(line: str) -> None:
-                return
-
-        else:
+        pbar: tqdm = tqdm(
+            total=total,
+            bar_format="{desc} |{bar}| {elapsed} {postfix[0][value]}",
+            postfix=[dict(value="Status")],
+            desc=f'chain {chain_id}',
+            colour='yellow',
+        )
 
-            def sampler_progress_hook(line: str) -> None:
-                if line == "Done":
-                    pbar.postfix[0]["value"] = 'Sampling completed'
-                    pbar.update(total - pbar.n)
-                    pbar.close()
-                elif line.startswith("Iteration"):
-                    if 'Sampling' in line:
-                        pbar.colour = 'blue'
-                    pbar.update(1)
-                    pbar.postfix[0]["value"] = line
+        def sampler_progress_hook(line: str) -> None:
+            if line == "Done":
+                pbar.postfix[0]["value"] = 'Sampling completed'
+                pbar.update(total - pbar.n)
+                pbar.close()
+            elif line.startswith("Iteration"):
+                if 'Sampling' in line:
+                    pbar.colour = 'blue'
+                pbar.update(1)
+                pbar.postfix[0]["value"] = line
 
         return sampler_progress_hook

cmdstanpy/progress.py

@@ -1,25 +1,49 @@
 """
 Record tqdm progress bar fail during session
 """
-
+import functools
 import logging
 
-SHOW_PROGRESS: bool = True
+_SHOW_PROGRESS: bool = True
 
 
 def allow_show_progress() -> bool:
-    return SHOW_PROGRESS
+    """Return False if any progressbar errors have occurred this session"""
+    return _SHOW_PROGRESS
 
 
-def disable_progress(e: Exception) -> None:
-    print("DISABLE")
+def _disable_progress(e: Exception) -> None:
+    """Print an exception and disable progress bars for this session"""
     # pylint: disable=global-statement
-    global SHOW_PROGRESS
-    if SHOW_PROGRESS:
+    global _SHOW_PROGRESS
+    if _SHOW_PROGRESS:
+        _SHOW_PROGRESS = False
         logging.getLogger('cmdstanpy').error(
             'Error in progress bar initialization:\n'
             '\t%s\n'
             'Disabling progress bars for this session',
             str(e),
         )
-    SHOW_PROGRESS = False
+
+
+def wrap_callback(func):  # type: ignore
+    """Wrap a callback generator so it fails safely"""
+
+    @functools.wraps(func)
+    def safe_progress(*args, **kwargs):  # type: ignore
+        # pylint: disable=unused-argument
+        def callback(*args, **kwargs):  # type: ignore
+            # totally empty callback
+            return None
+
+        if not allow_show_progress():
+            return callback
+
+        try:
+            return func(*args, **kwargs)
+        # pylint: disable=broad-except
+        except Exception as e:
+            _disable_progress(e)
+            return callback
+
+    return safe_progress

cmdstanpy/utils.py

@@ -42,8 +42,7 @@
     _TMPDIR,
 )
 
-from cmdstanpy.progress import disable_progress
-
+from . import progress as progbar
 
 EXTENSION = '.exe' if platform.system() == 'Windows' else ''
 
@@ -1210,40 +1209,27 @@ def install_cmdstan(
     return True
 
 
+@progbar.wrap_callback
 def wrap_url_progress_hook() -> Optional[Callable[[int, int, int], None]]:
     """Sets up tqdm callback for url downloads."""
-    try:
-        pbar: tqdm = tqdm(
-            unit='B',
-            unit_scale=True,
-            unit_divisor=1024,
-            colour='blue',
-            leave=False,
-        )  # type: ignore
-    # pylint: disable=broad-except
-    except Exception as e:
-        disable_progress(e)
-
-        def download_progress_hook(
-            # pylint: disable=unused-argument
-            count: int,
-            block_size: int,
-            total_size: int,
-        ) -> None:
-            return
-
-    else:
+    pbar: tqdm = tqdm(
+        unit='B',
+        unit_scale=True,
+        unit_divisor=1024,
+        colour='blue',
+        leave=False,
+    )
 
-        def download_progress_hook(
-            count: int, block_size: int, total_size: int
-        ) -> None:
-            if pbar.total is None:
-                pbar.total = total_size
-                pbar.reset()
-            downloaded_size = count * block_size
-            pbar.update(downloaded_size - pbar.n)
-            if pbar.n >= total_size:
-                pbar.close()
+    def download_progress_hook(
+        count: int, block_size: int, total_size: int
+    ) -> None:
+        if pbar.total is None:
+            pbar.total = total_size
+            pbar.reset()
+        downloaded_size = count * block_size
+        pbar.update(downloaded_size - pbar.n)
+        if pbar.n >= total_size:
+            pbar.close()
 
     return download_progress_hook