update: add interruptible execution of loops decorator#272
update: add interruptible execution of loops decorator#272
Conversation
📝 WalkthroughWalkthroughConverts a synchronous job-polling function to an asynchronous variant using a decorator-based approach with single-pass polling and status counting. Updates the call site in a notebook to use the new async function with await syntax. Relaxes dependency version constraint in pyproject.toml. Changes
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~20 minutes Poem
🚥 Pre-merge checks | ✅ 3✅ Passed checks (3 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
|
Check out this pull request on See visual diffs & provide feedback on Jupyter Notebooks. Powered by ReviewNB |
utils/api.py
Outdated
|
|
||
|
|
||
| def wait_for_jobs_to_finish(endpoint: JobEndpoints, job_ids: list, poll_interval: int = 10) -> None: | ||
| def _print_jobs_status_table(counts: Counter) -> None: |
There was a problem hiding this comment.
utils/interrupts.py
Outdated
| show_controls: bool = True, | ||
| ): | ||
| """ | ||
| Decorator for single-iteration polling functions. |
utils/interrupts.py
Outdated
| show_controls: bool = True, | ||
| ) -> None: | ||
| """ | ||
| Minimal wrapper. |
utils/interrupts.py
Outdated
| from functools import wraps | ||
| from typing import Any, Awaitable, Callable | ||
|
|
||
| from mat3ra.utils.jupyterlite.environment import ENVIRONMENT, EnvironmentsEnum |
utils/interrupts.py
Outdated
|
|
||
|
|
||
| @dataclass | ||
| class BroadcastChannelAbortController: |
There was a problem hiding this comment.
Actionable comments posted: 3
🧹 Nitpick comments (1)
utils/api.py (1)
63-71: Update docstring to document the return value.The function now returns a
boolindicating whether jobs are still active, but the docstring doesn't document this. This is important for understanding the polling semantics.📝 Suggested docstring update
""" Waits for jobs to finish and prints their statuses. A job is considered finished if it is not in "pre-submission", "submitted", or, "active" status. Args: endpoint (JobEndpoints): Job endpoint object from the Exabyte API Client job_ids (list): list of job IDs to wait for + + Returns: + bool: True if any jobs are still active (polling should continue), False if all jobs have finished. """🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@utils/api.py` around lines 63 - 71, Update the docstring for wait_for_jobs_to_finish_async to include a Returns section that documents the boolean return value: state that the function returns a bool indicating whether any of the supplied jobs are still active after polling (True if at least one job remains in an active/submitted/pre-submission state, False if all jobs are finished), and mention the polling semantics/timing if relevant; edit the docstring immediately above the function definition for clarity.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@other/materials_designer/workflows/band_gap.ipynb`:
- Around line 592-594: The call to wait_for_jobs_to_finish_async is passing an
unsupported keyword argument poll_interval; update the call in the notebook to
call wait_for_jobs_to_finish_async(client.jobs, [job_id]) (remove
poll_interval=POLL_INTERVAL) and remove the now-unused POLL_INTERVAL variable
declaration earlier in the notebook; target the invocation of
wait_for_jobs_to_finish_async and the POLL_INTERVAL definition to make these
changes.
- Line 592: Replace usages of the removed synchronous wait_for_jobs_to_finish
with the new async function wait_for_jobs_to_finish_async from utils.api: update
imports to "from utils.api import wait_for_jobs_to_finish_async" and change all
calls to either "await wait_for_jobs_to_finish_async(...)" inside async
cells/functions or wrap calls in a top-level runner like "import asyncio;
asyncio.run(wait_for_jobs_to_finish_async(...))" in notebook cells that must
remain synchronous; update every notebook that imports or calls
wait_for_jobs_to_finish (search for that symbol) and ensure the call site uses
await or asyncio.run and any surrounding functions are marked async if needed.
In `@pyproject.toml`:
- Line 15: Update the pyproject.toml dependency entry for mat3ra-utils to
include a minimum version that provides the interruptible_polling_loop decorator
(used in utils/api.py); change the existing bare "mat3ra-utils" requirement to a
constrained requirement for the release that added interruptible_polling_loop,
then run a quick install/test to confirm the import in utils/api.py succeeds.
---
Nitpick comments:
In `@utils/api.py`:
- Around line 63-71: Update the docstring for wait_for_jobs_to_finish_async to
include a Returns section that documents the boolean return value: state that
the function returns a bool indicating whether any of the supplied jobs are
still active after polling (True if at least one job remains in an
active/submitted/pre-submission state, False if all jobs are finished), and
mention the polling semantics/timing if relevant; edit the docstring immediately
above the function definition for clarity.
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 00b05956-1ea1-4325-8aca-14ac2b510f66
📒 Files selected for processing (3)
other/materials_designer/workflows/band_gap.ipynbpyproject.tomlutils/api.py
| "from utils.api import wait_for_jobs_to_finish_async\n", | ||
| "\n", | ||
| "wait_for_jobs_to_finish(client.jobs, [job_id], poll_interval=POLL_INTERVAL)" | ||
| "await wait_for_jobs_to_finish_async(client.jobs, [job_id], poll_interval=POLL_INTERVAL)" |
There was a problem hiding this comment.
TypeError: wait_for_jobs_to_finish_async does not accept poll_interval argument.
The function call passes poll_interval=POLL_INTERVAL, but the new function signature in utils/api.py is:
def wait_for_jobs_to_finish_async(endpoint: JobEndpoints, job_ids: List[str]) -> bool:This will raise TypeError: wait_for_jobs_to_finish_async() got an unexpected keyword argument 'poll_interval' at runtime.
🐛 Proposed fix
- "await wait_for_jobs_to_finish_async(client.jobs, [job_id], poll_interval=POLL_INTERVAL)"
+ "await wait_for_jobs_to_finish_async(client.jobs, [job_id])"Additionally, the POLL_INTERVAL variable defined on line 63 becomes unused and can be removed.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@other/materials_designer/workflows/band_gap.ipynb` around lines 592 - 594,
The call to wait_for_jobs_to_finish_async is passing an unsupported keyword
argument poll_interval; update the call in the notebook to call
wait_for_jobs_to_finish_async(client.jobs, [job_id]) (remove
poll_interval=POLL_INTERVAL) and remove the now-unused POLL_INTERVAL variable
declaration earlier in the notebook; target the invocation of
wait_for_jobs_to_finish_async and the POLL_INTERVAL definition to make these
changes.
| "outputs": [], | ||
| "source": [ | ||
| "from utils.api import wait_for_jobs_to_finish\n", | ||
| "from utils.api import wait_for_jobs_to_finish_async\n", |
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Search for other usages of the old wait_for_jobs_to_finish function
rg -n "wait_for_jobs_to_finish[^_]" --type-add 'notebook:*.ipynb' -t notebook -t pyRepository: Exabyte-io/api-examples
Length of output: 2314
🏁 Script executed:
# Locate utils.api module
fd -e py -e ipynb "utils" | grep -i api | head -20Repository: Exabyte-io/api-examples
Length of output: 49
🏁 Script executed:
# Find files containing wait_for_jobs_to_finish definition (not just imports)
rg -n "def wait_for_jobs_to_finish" --type py
# Also search for the utils directory structure
git ls-files | grep -E "(utils|api)" | head -20Repository: Exabyte-io/api-examples
Length of output: 651
🏁 Script executed:
# Check the utils/api.py file for function definitions
cat -n utils/api.py | head -100Repository: Exabyte-io/api-examples
Length of output: 4301
Update all notebooks importing the removed wait_for_jobs_to_finish function.
The synchronous function wait_for_jobs_to_finish has been removed from utils/api.py and replaced with wait_for_jobs_to_finish_async. However, at least 8 notebooks still import or call the old function and will fail with ImportError:
other/jarvis/run_job_using_material_from_jarvis_db.ipynbexamples/reproducing_publications/band_structure_for_interface_bilayer_twisted_molybdenum_disulfide.ipynbexamples/job/run-simulations-and-extract-properties.ipynbexamples/job/ml-train-model-predict-properties.ipynbexamples/job/get-file-from-job.ipynbexamples/workflow/qe_scf_calculation.ipynbexamples/reproducing_publications/band_gaps_for_interface_bilayer_twisted_molybdenum_disulfide.ipynb
Update all affected notebooks to use wait_for_jobs_to_finish_async before merging.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@other/materials_designer/workflows/band_gap.ipynb` at line 592, Replace
usages of the removed synchronous wait_for_jobs_to_finish with the new async
function wait_for_jobs_to_finish_async from utils.api: update imports to "from
utils.api import wait_for_jobs_to_finish_async" and change all calls to either
"await wait_for_jobs_to_finish_async(...)" inside async cells/functions or wrap
calls in a top-level runner like "import asyncio;
asyncio.run(wait_for_jobs_to_finish_async(...))" in notebook cells that must
remain synchronous; update every notebook that imports or calls
wait_for_jobs_to_finish (search for that symbol) and ensure the call site uses
await or asyncio.run and any surrounding functions are marked async if needed.
| "pymatgen==2024.4.13", | ||
| "mat3ra-made>=2025.8.13.post0", | ||
| "mat3ra-utils>=2024.6.11.post0", | ||
| "mat3ra-utils", |
There was a problem hiding this comment.
🧩 Analysis chain
🌐 Web query:
mat3ra-utils interruptible_polling_loop release version
💡 Result:
interruptible_polling_loop is present in the Exabyte-io/utils repo at tag 2026.3.6-0 (which corresponds to the PyPI release mat3ra-utils 2026.3.6.post0, released Mar 6, 2026). [1] [2]
🏁 Script executed:
# Check if utils/api.py exists and imports interruptible_polling_loop
fd -t f api.py | head -20Repository: Exabyte-io/api-examples
Length of output: 79
🏁 Script executed:
# Search for interruptible_polling_loop imports in the codebase
rg "interruptible_polling_loop" --type pyRepository: Exabyte-io/api-examples
Length of output: 197
🏁 Script executed:
# Check current pyproject.toml for mat3ra-utils dependency
cat pyproject.toml | grep -A 2 -B 2 "mat3ra-utils"Repository: Exabyte-io/api-examples
Length of output: 178
Retain a minimum version constraint for mat3ra-utils.
Removing the version constraint allows older versions without the interruptible_polling_loop decorator (imported in utils/api.py) to be installed, causing runtime ImportError. Specify a minimum version to ensure the required feature is available.
- "mat3ra-utils",
+ "mat3ra-utils>=2026.3.6.post0",📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| "mat3ra-utils", | |
| "mat3ra-utils>=2026.3.6.post0", |
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@pyproject.toml` at line 15, Update the pyproject.toml dependency entry for
mat3ra-utils to include a minimum version that provides the
interruptible_polling_loop decorator (used in utils/api.py); change the existing
bare "mat3ra-utils" requirement to a constrained requirement for the release
that added interruptible_polling_loop, then run a quick install/test to confirm
the import in utils/api.py succeeds.
Summary by CodeRabbit
Chores
Refactor