Skip to content

update: add interruptible execution of loops decorator#272

Open
VsevolodX wants to merge 15 commits intomainfrom
feature/SOF-7852
Open

update: add interruptible execution of loops decorator#272
VsevolodX wants to merge 15 commits intomainfrom
feature/SOF-7852

Conversation

@VsevolodX
Copy link
Member

@VsevolodX VsevolodX commented Feb 27, 2026

Summary by CodeRabbit

  • Chores

    • Updated dependency versioning for broader compatibility.
  • Refactor

    • Converted job monitoring from synchronous to asynchronous polling for improved responsiveness.
    • Enhanced status output formatting for job monitoring operations.

@coderabbitai
Copy link

coderabbitai bot commented Feb 27, 2026

📝 Walkthrough

Walkthrough

Converts 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

Cohort / File(s) Summary
Async Job Polling Refactor
utils/api.py
Function wait_for_jobs_to_finish renamed to wait_for_jobs_to_finish_async, decorated with @interruptible_polling_loop(), and reimplemented: removes iterative polling loop in favor of single-pass status check; returns bool instead of None; uses Counter for status aggregation and pretty_print for formatted output; removes poll_interval parameter and dependencies on time and tabulate.
Async Adoption
other/materials_designer/workflows/band_gap.ipynb
Updated to import wait_for_jobs_to_finish_async instead of wait_for_jobs_to_finish and changed invocation to use await syntax.
Dependency Update
pyproject.toml
Relaxed mat3ra-utils dependency constraint from pinned minimum version >=2024.6.11.post0 to unpinned mat3ra-utils.

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~20 minutes

Poem

🐰 A polling loop transformed to a single glance,
With counters and pretty prints joining the dance,
Async awaits now the jobs' grand show,
While dependencies loose their version grip, let them flow! ✨

🚥 Pre-merge checks | ✅ 3
✅ Passed checks (3 passed)
Check name Status Explanation
Description Check ✅ Passed Check skipped - CodeRabbit’s high-level summary is enabled.
Title check ✅ Passed The title 'update: add interruptible execution of loops decorator' accurately describes the main change: introducing an interruptible_polling_loop decorator to enable async/await-based polling with interruption capability.
Docstring Coverage ✅ Passed Docstring coverage is 100.00% which is sufficient. The required threshold is 80.00%.

✏️ Tip: You can configure your own custom pre-merge checks in the settings.

✨ Finishing Touches
  • 📝 Generate docstrings (stacked PR)
  • 📝 Generate docstrings (commit on current branch)
🧪 Generate unit tests (beta)
  • Create PR with unit tests
  • Post copyable unit tests in a comment
  • Commit unit tests in branch feature/SOF-7852

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.

❤️ Share

Comment @coderabbitai help to get the list of available commands and usage tips.

@review-notebook-app
Copy link

Check out this pull request on  ReviewNB

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:
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

show_controls: bool = True,
):
"""
Decorator for single-iteration polling functions.
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is not clear

show_controls: bool = True,
) -> None:
"""
Minimal wrapper.
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

?

from functools import wraps
from typing import Any, Awaitable, Callable

from mat3ra.utils.jupyterlite.environment import ENVIRONMENT, EnvironmentsEnum
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Let's put this to utils



@dataclass
class BroadcastChannelAbortController:
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Use pydantic

Base automatically changed from feature/SOF-7843 to main March 10, 2026 06:16
Copy link

@coderabbitai coderabbitai bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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 bool indicating 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

📥 Commits

Reviewing files that changed from the base of the PR and between 509aa7e and 48c176f.

📒 Files selected for processing (3)
  • other/materials_designer/workflows/band_gap.ipynb
  • pyproject.toml
  • utils/api.py

Comment on lines +592 to +594
"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)"
Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🔴 Critical

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",
Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🔴 Critical

🧩 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 py

Repository: Exabyte-io/api-examples

Length of output: 2314


🏁 Script executed:

# Locate utils.api module
fd -e py -e ipynb "utils" | grep -i api | head -20

Repository: 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 -20

Repository: 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 -100

Repository: 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.ipynb
  • examples/reproducing_publications/band_structure_for_interface_bilayer_twisted_molybdenum_disulfide.ipynb
  • examples/job/run-simulations-and-extract-properties.ipynb
  • examples/job/ml-train-model-predict-properties.ipynb
  • examples/job/get-file-from-job.ipynb
  • examples/workflow/qe_scf_calculation.ipynb
  • examples/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",
Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟠 Major

🧩 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 -20

Repository: Exabyte-io/api-examples

Length of output: 79


🏁 Script executed:

# Search for interruptible_polling_loop imports in the codebase
rg "interruptible_polling_loop" --type py

Repository: 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.

Suggested change
"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.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants