fix: wait for jobs to finish async

[VsevolodX]

Summary by CodeRabbit

• Chores
  • Updated example notebooks to use asynchronous job completion monitoring across six notebook files in examples and other directories.

[review-notebook-app]

Check out this pull request on 

See visual diffs & provide feedback on Jupyter Notebooks.

---

Powered by ReviewNB

[coderabbitai]

📝 Walkthrough

Walkthrough

Six example notebooks are updated to use asynchronous job-waiting utilities. Imports are changed from wait_for_jobs_to_finish to wait_for_jobs_to_finish_async, and synchronous function calls are replaced with await wait_for_jobs_to_finish_async(...) expressions to enable async/await execution patterns.

Changes

Cohort / File(s)	Summary
Job Examples examples/job/get-file-from-job.ipynb, examples/job/ml-train-model-predict-properties.ipynb, examples/job/run-simulations-and-extract-properties.ipynb	Import and usage updated from synchronous wait_for_jobs_to_finish to asynchronous wait_for_jobs_to_finish_async with await syntax.
Publication Examples examples/reproducing_publications/band_gaps_for_interface_bilayer_twisted_molybdenum_disulfide.ipynb, examples/reproducing_publications/band_structure_for_interface_bilayer_twisted_molybdenum_disulfide.ipynb	Job completion wait calls converted from synchronous to awaited async variant, maintaining existing poll interval parameters.
Workflow & JARVIS Examples examples/workflow/qe_scf_calculation.ipynb, other/jarvis/run_job_using_material_from_jarvis_db.ipynb	Replaced synchronous job-wait function calls with async/await pattern to align with async execution model.

Estimated code review effort

🎯 2 (Simple) | ⏱️ ~10 minutes

Possibly related PRs

• update: add interruptible execution of loops decorator #272: Directly related—adds wait_for_jobs_to_finish_async to utils/api.py that these notebooks now import and call.
• update: add TE convergence NB #285: Modifies the implementation of wait_for_jobs_to_finish_async in utils/api.py, which this PR's notebooks depend upon.
• update: fix awaiting for queued job #295: Alters wait_for_jobs_to_finish_async behavior for queued/submitted job handling, affecting the async utility these notebooks call.

Suggested reviewers

• timurbazhirov

Poem

🐰 Hops with glee through async dreams,
Where await and coroutines reign supreme,
Six notebooks now dance with grace,
As sync gives way to async's pace! ✨

🚥 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 PR title clearly and specifically describes the main change: migrating job-wait calls from synchronous to asynchronous implementation across multiple example notebooks.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.

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

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Commit unit tests in branch fix/SOF-7805

Warning

There were issues while running some tools. Please review the errors and either fix the tool's configuration or disable the tool if it's a critical failure.

🔧 Ruff (0.15.10)

other/jarvis/run_job_using_material_from_jarvis_db.ipynb

Unexpected end of JSON input

examples/reproducing_publications/band_gaps_for_interface_bilayer_twisted_molybdenum_disulfide.ipynb

Unexpected end of JSON input

examples/job/ml-train-model-predict-properties.ipynb

Unexpected end of JSON input

• 4 others

---

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

• X
• Mastodon
• Reddit
• LinkedIn

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

[coderabbitai]

Actionable comments posted: 1

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (2)

examples/reproducing_publications/band_structure_for_interface_bilayer_twisted_molybdenum_disulfide.ipynb (1)

749-753: ⚠️ Potential issue | 🔴 Critical

Async wait call is unresolved due to stale import.

Line 753 uses wait_for_jobs_to_finish_async, but Line 749 imports wait_for_jobs_to_finish, which does not exist in utils.api. This will raise NameError at runtime.

🔧 Proposed fix

-from utils.api import wait_for_jobs_to_finish
+from utils.api import wait_for_jobs_to_finish_async

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In
`@examples/reproducing_publications/band_structure_for_interface_bilayer_twisted_molybdenum_disulfide.ipynb`
around lines 749 - 753, The code imports wait_for_jobs_to_finish but calls
wait_for_jobs_to_finish_async, causing a NameError; update the import to pull
wait_for_jobs_to_finish_async from utils.api (or alternatively change the call
to wait_for_jobs_to_finish if that is the intended function). Locate the import
statement that currently reads "from utils.api import wait_for_jobs_to_finish"
and replace it with "from utils.api import wait_for_jobs_to_finish_async" (or
change the await call to use wait_for_jobs_to_finish), ensuring the symbol names
in the notebook (wait_for_jobs_to_finish_async and job_ids usage) match the
imported function.

examples/reproducing_publications/band_gaps_for_interface_bilayer_twisted_molybdenum_disulfide.ipynb (1)

762-766: ⚠️ Potential issue | 🔴 Critical

Fix import and function signature for async job completion waiter.

The import statement references a non-existent function wait_for_jobs_to_finish, and the function call passes an unsupported poll_interval parameter. The actual function in utils/api is wait_for_jobs_to_finish_async and accepts only endpoint and job_ids parameters.

🔧 Proposed fix

-from utils.api import wait_for_jobs_to_finish
+from utils.api import wait_for_jobs_to_finish_async

-await wait_for_jobs_to_finish_async(job_endpoints, job_ids, poll_interval=60)
+await wait_for_jobs_to_finish_async(job_endpoints, job_ids)

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In
`@examples/reproducing_publications/band_gaps_for_interface_bilayer_twisted_molybdenum_disulfide.ipynb`
around lines 762 - 766, Replace the incorrect import and call: change the import
to bring in wait_for_jobs_to_finish_async from utils.api (instead of
wait_for_jobs_to_finish) and call it with only the two required parameters
(endpoint and job_ids); in practice update the import to use
wait_for_jobs_to_finish_async and remove the unsupported poll_interval argument
so the call becomes await wait_for_jobs_to_finish_async(job_endpoints, job_ids),
keeping existing job_endpoints and job_ids variable names.

🤖 Prompt for all review comments with AI agents

Verify each finding against the current code and only fix it if needed.

Inline comments:
In `@examples/workflow/qe_scf_calculation.ipynb`:
- Line 340: The notebook calls await wait_for_jobs_to_finish_async(...) but
never imports it, causing a NameError; update the imports (or the call) so the
async function is available: either import wait_for_jobs_to_finish_async from
its module where wait_for_jobs_to_finish is imported, or replace the await call
with the synchronous wait_for_jobs_to_finish(...) call (and remove await).
Locate the reference to wait_for_jobs_to_finish_async and the existing import of
wait_for_jobs_to_finish to make the corresponding change.

---

Outside diff comments:
In
`@examples/reproducing_publications/band_gaps_for_interface_bilayer_twisted_molybdenum_disulfide.ipynb`:
- Around line 762-766: Replace the incorrect import and call: change the import
to bring in wait_for_jobs_to_finish_async from utils.api (instead of
wait_for_jobs_to_finish) and call it with only the two required parameters
(endpoint and job_ids); in practice update the import to use
wait_for_jobs_to_finish_async and remove the unsupported poll_interval argument
so the call becomes await wait_for_jobs_to_finish_async(job_endpoints, job_ids),
keeping existing job_endpoints and job_ids variable names.

In
`@examples/reproducing_publications/band_structure_for_interface_bilayer_twisted_molybdenum_disulfide.ipynb`:
- Around line 749-753: The code imports wait_for_jobs_to_finish but calls
wait_for_jobs_to_finish_async, causing a NameError; update the import to pull
wait_for_jobs_to_finish_async from utils.api (or alternatively change the call
to wait_for_jobs_to_finish if that is the intended function). Locate the import
statement that currently reads "from utils.api import wait_for_jobs_to_finish"
and replace it with "from utils.api import wait_for_jobs_to_finish_async" (or
change the await call to use wait_for_jobs_to_finish), ensuring the symbol names
in the notebook (wait_for_jobs_to_finish_async and job_ids usage) match the
imported function.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro

Run ID: 7efe8225-d7fd-4574-b355-aee4202c6337

📥 Commits

Reviewing files that changed from the base of the PR and between 9cfa8aa and 9f35a3b.

📒 Files selected for processing (7)

• examples/job/get-file-from-job.ipynb
• examples/job/ml-train-model-predict-properties.ipynb
• examples/job/run-simulations-and-extract-properties.ipynb
• examples/reproducing_publications/band_gaps_for_interface_bilayer_twisted_molybdenum_disulfide.ipynb
• examples/reproducing_publications/band_structure_for_interface_bilayer_twisted_molybdenum_disulfide.ipynb
• examples/workflow/qe_scf_calculation.ipynb
• other/jarvis/run_job_using_material_from_jarvis_db.ipynb

[coderabbitai]

⚠️ Potential issue | 🔴 Critical

🧩 Analysis chain

🏁 Script executed:

# First, locate the notebook file
find . -name "qe_scf_calculation.ipynb" -type f

Repository: Exabyte-io/api-examples

Length of output: 111

---

🏁 Script executed:

# Also check if utils/api module exists and examine its structure
find . -path "*/utils/api*" -type f

Repository: Exabyte-io/api-examples

Length of output: 81

---

🏁 Script executed:

# Parse the notebook to find the relevant cells
python3 << 'EOF'
import json

with open('./examples/workflow/qe_scf_calculation.ipynb', 'r') as f:
    notebook = json.load(f)

# Find cells that contain imports (around cell with wait_for_jobs_to_finish)
for i, cell in enumerate(notebook.get('cells', [])):
    if cell.get('cell_type') == 'code':
        source = ''.join(cell.get('source', []))
        if 'wait_for_jobs_to_finish' in source or 'job_endpoints' in source:
            print(f"\n=== Cell {i} ===")
            print(f"Source:\n{source}")
            if 'execution_count' in cell:
                print(f"Execution count: {cell['execution_count']}")
EOF

Repository: Exabyte-io/api-examples

Length of output: 1963

---

🏁 Script executed:

# Check utils/api.py for the function definitions
rg "def wait_for_jobs_to_finish" utils/api.py -A 5

Repository: Exabyte-io/api-examples

Length of output: 329

---

wait_for_jobs_to_finish_async is used but not imported.

The notebook imports wait_for_jobs_to_finish but calls await wait_for_jobs_to_finish_async(). This will fail at runtime with NameError.

🔧 Proposed fix

-from utils.api import wait_for_jobs_to_finish
+from utils.api import wait_for_jobs_to_finish_async

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@examples/workflow/qe_scf_calculation.ipynb` at line 340, The notebook calls
await wait_for_jobs_to_finish_async(...) but never imports it, causing a
NameError; update the imports (or the call) so the async function is available:
either import wait_for_jobs_to_finish_async from its module where
wait_for_jobs_to_finish is imported, or replace the await call with the
synchronous wait_for_jobs_to_finish(...) call (and remove await). Locate the
reference to wait_for_jobs_to_finish_async and the existing import of
wait_for_jobs_to_finish to make the corresponding change.