Fix TQDM_DISABLE=1 compatibility issue with cmdstanpy#888
Fix TQDM_DISABLE=1 compatibility issue with cmdstanpy#888mohsinm-dev wants to merge 1 commit intouber:devfrom
Conversation
- Add cmdstanpy_compat.py with patch_tqdm_progress_hook() function - Patch cmdstanpy progress hooks to safely handle disabled tqdm objects - Apply patch in stan_estimator.py when TQDM_DISABLE=1 is set - Add comprehensive test suite covering all scenarios Fixes uber#887
|
|
There was a problem hiding this comment.
Pull request overview
This PR fixes issue #887 where setting TQDM_DISABLE=1 before importing orbit would cause an AttributeError crash because cmdstanpy's progress hook assumes tqdm progress bars have a postfix attribute, which disabled tqdm objects don't provide.
Changes:
- Adds
orbit/utils/cmdstanpy_compat.pywith apatch_tqdm_progress_hook()function that monkey-patches cmdstanpy's_wrap_sampler_progress_hookmethod to safely guardpostfixattribute accesses - Applies the compatibility patch at module load time in
orbit/estimators/stan_estimator.py, and removes a now-unusedfrom sys import platform, version_infoimport - Adds
tests/orbit/utils/test_cmdstanpy_compat.pywith unit and integration test coverage for the new utility
Reviewed changes
Copilot reviewed 3 out of 3 changed files in this pull request and generated 6 comments.
| File | Description |
|---|---|
orbit/utils/cmdstanpy_compat.py |
New module containing the tqdm compatibility patch for cmdstanpy |
orbit/estimators/stan_estimator.py |
Applies the compatibility patch at import time; removes unused sys import |
tests/orbit/utils/test_cmdstanpy_compat.py |
New test file covering patch application, no-op scenarios, idempotency, and error handling |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
|
|
||
| try: | ||
| import cmdstanpy.model | ||
| import re |
There was a problem hiding this comment.
The re module is imported inside the try block at line 33, and is captured by the safe_wrap_sampler_progress_hook closure. While this works correctly, re is a standard library module. Importing it at the top of the file (next to the import os) would be cleaner and more consistent with Python conventions and with the rest of the codebase.
| try: | ||
| from tqdm import tqdm | ||
|
|
||
| pat = re.compile(r"Chain \[(\d*)\] (Iteration.*)") | ||
| pbars: Dict[int, tqdm] = { | ||
| chain_id: tqdm( | ||
| total=total, | ||
| bar_format="{desc} |{bar}| {elapsed} {postfix[0][value]}", | ||
| postfix=[{"value": "Status"}], | ||
| desc=f"chain {chain_id}", | ||
| colour="yellow", | ||
| ) | ||
| for chain_id in chain_ids | ||
| } | ||
|
|
||
| def progress_hook(line: str, idx: int) -> None: | ||
| if line == "Done": | ||
| for pbar in pbars.values(): | ||
| # safe postfix access | ||
| if hasattr(pbar, "postfix") and pbar.postfix: | ||
| try: | ||
| pbar.postfix[0]["value"] = "Sampling completed" | ||
| except (AttributeError, KeyError, IndexError): | ||
| pass | ||
| pbar.update(total - pbar.n) | ||
| pbar.close() | ||
| else: | ||
| match = pat.match(line) | ||
| if match: | ||
| idx = int(match.group(1)) | ||
| mline = match.group(2).strip() | ||
| elif line.startswith("Iteration"): | ||
| mline = line | ||
| idx = chain_ids[idx] | ||
| else: | ||
| return | ||
|
|
||
| if idx in pbars: | ||
| if "Sampling" in mline and hasattr(pbars[idx], "colour"): | ||
| pbars[idx].colour = "blue" | ||
| pbars[idx].update(1) | ||
|
|
||
| # safe postfix access | ||
| if hasattr(pbars[idx], "postfix") and pbars[idx].postfix: | ||
| try: | ||
| pbars[idx].postfix[0]["value"] = mline | ||
| except (AttributeError, KeyError, IndexError): | ||
| pass | ||
|
|
||
| return progress_hook | ||
|
|
||
| except Exception as e: | ||
| logger.warning( | ||
| f"Progress bar setup failed: {e}. Disabling progress bars." | ||
| ) | ||
| return None |
There was a problem hiding this comment.
When TQDM_DISABLE=1, the purpose is to disable all progress output. The safe_wrap_sampler_progress_hook function still creates tqdm progress bar instances (lines 55–64) even in this case — they are silently disabled by tqdm, but creating them adds unnecessary overhead. A simpler and more direct solution would be to return a no-op callable from safe_wrap_sampler_progress_hook when TQDM_DISABLE=1 is set, rather than creating progress bar objects with safe guards. The no-op callable would satisfy cmdstanpy's requirement for a callable while completely skipping progress bar logic:
def progress_hook(line: str, idx: int) -> None:
pass
return progress_hookThis makes the intent clearer and avoids any future tqdm API compatibility issues.
| # Only patch if TQDM_DISABLE is set | ||
| if os.environ.get("TQDM_DISABLE") != "1": |
There was a problem hiding this comment.
The guard on line 28 only triggers the patch when TQDM_DISABLE is exactly "1". However, tqdm treats several other values as truthy for disabling progress bars (e.g., "true", "yes", "True"). Users who set TQDM_DISABLE=true or TQDM_DISABLE=yes would still encounter the original crash, as the patch would not be applied. The condition should be broadened to cover all values that tqdm considers as "disabled", for example by checking os.environ.get("TQDM_DISABLE", "0").lower() not in ("0", "false", "", "no"), or by checking whether tqdm.tqdm.is_disabled() returns True.
| # Only patch if TQDM_DISABLE is set | |
| if os.environ.get("TQDM_DISABLE") != "1": | |
| # Only patch when tqdm is effectively disabled via TQDM_DISABLE | |
| env_value = os.environ.get("TQDM_DISABLE", "0") | |
| tqdm_disabled = env_value.lower() not in ("0", "false", "", "no") | |
| if not tqdm_disabled: |
| # Apply cmdstanpy compatibility patches | ||
| patch_tqdm_progress_hook() |
There was a problem hiding this comment.
The patch_tqdm_progress_hook() function is called unconditionally at module load time (line 20), which means the patch is only applied if TQDM_DISABLE=1 was set before orbit/estimators/stan_estimator.py was first imported. If a user sets TQDM_DISABLE=1 after importing orbit, the patch will not be applied, and the same crash will occur. While this matches the specific scenario described in issue #887, this limitation is not documented, which could lead to user confusion. Consider adding a note to the docstring or a logged warning if TQDM_DISABLE=1 is set but the patch has already been skipped.
| finally: | ||
| # Clean up | ||
| if hasattr(cmdstanpy.model.CmdStanModel, "_orbit_tqdm_patched"): | ||
| delattr(cmdstanpy.model.CmdStanModel, "_orbit_tqdm_patched") | ||
| if original_method is not None: | ||
| cmdstanpy.model.CmdStanModel._wrap_sampler_progress_hook = ( | ||
| original_method | ||
| ) |
There was a problem hiding this comment.
The test_integration_with_actual_cmdstanpy test has a cleanup issue: the finally block (lines 97–103) unconditionally removes _orbit_tqdm_patched and restores the original method. If patch_tqdm_progress_hook() was already called during the import of stan_estimator.py (because TQDM_DISABLE=1 was set at import time), this cleanup will remove the module-level patch that was applied at import time. The method will be restored while _orbit_tqdm_patched is deleted, leaving the class in a partially inconsistent state for subsequent test runs. The cleanup should check whether the patch was applied by this test's invocation specifically, or only clean up if the test itself applied the patch (e.g., by checking whether the flag was absent before the call).
| @staticmethod | ||
| def safe_wrap_sampler_progress_hook( |
There was a problem hiding this comment.
The @staticmethod decorator is applied to a function defined inside another function (line 45). This is an unusual pattern that creates a staticmethod descriptor object as a local variable. Assigning this descriptor to the class attribute at line 109 works correctly through Python's descriptor protocol, but this pattern can be confusing for maintainers who may not be familiar with it. A clearer equivalent would be to define the inner function without the @staticmethod decorator, and then use staticmethod(safe_wrap_sampler_progress_hook) when assigning to the class attribute on line 109. This is functionally identical but makes the intent clearer.
Summary
Fixes AttributeError when TQDM_DISABLE=1 is set before importing orbit. The issue occurs because cmdstanpy assumes tqdm progress bars have a 'postfix' attribute, but disabled tqdm objects don't have this attribute.
Changes
orbit/utils/cmdstanpy_compat.pywith safe progress hook implementationorbit/estimators/stan_estimator.pywhen TQDM_DISABLE=1Test Plan
Closes #887