feat(auditor): add blocking submit option to auditor SDK - #1025
feat(auditor): add blocking submit option to auditor SDK#1025parkanzky wants to merge 2 commits into
Conversation
📝 WalkthroughWalkthroughThe SDK now returns synchronous and asynchronous audit job resource handles. The handles support status polling, log retrieval, completion waiting, and secure artifact downloads. Tests and documentation use the resource interface. ChangesAuditor job resources
Sequence Diagram(s)sequenceDiagram
participant Client
participant AuditorPluginResource
participant AuditorJobResource
participant JobsAPI
Client->>AuditorPluginResource: submit audit job
AuditorPluginResource->>JobsAPI: submit job
JobsAPI-->>AuditorPluginResource: job name
AuditorPluginResource-->>Client: AuditorJobResource handle
Client->>AuditorJobResource: wait_until_done()
AuditorJobResource->>JobsAPI: poll status and retrieve logs
JobsAPI-->>AuditorJobResource: status and logs
Client->>AuditorJobResource: download_artifacts()
AuditorJobResource->>JobsAPI: download completed artifacts
JobsAPI-->>AuditorJobResource: archive
Possibly related PRs
Suggested labels: Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches 💡 1📝 Generate docstrings 💡
🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 4
🧹 Nitpick comments (7)
docs/auditor/sdk-resources.mdx (2)
226-226: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winUse the canonical Fern URL.
Replace
#auditorjobresourcewith the canonical Fern navigation URL for this section.As per coding guidelines, use canonical Fern navigation URLs for internal links.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@docs/auditor/sdk-resources.mdx` at line 226, Update the internal link in the AsyncAuditorPluginResource.submit() documentation to replace the local `#auditorjobresource` anchor with the canonical Fern navigation URL for the AuditorJobResource section.Source: Coding guidelines
65-116: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winSeparate reference and workflow content.
Lines 65-90 are REFERENCE content. Lines 92-116 are HOW-TO content. Move the workflow to a separate HOW-TO page and cross-link it.
As per coding guidelines, each documentation page must fit one Diataxis quadrant.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@docs/auditor/sdk-resources.mdx` around lines 65 - 116, Separate the “Submit and wait for an audit job” workflow, including both code examples, from the reference page into a dedicated HOW-TO page. Keep the submit argument and AuditorJobResource method tables in sdk-resources.mdx as reference content, then add a cross-link from that page to the new workflow page.Source: Coding guidelines
plugins/nemo-auditor/README.md (1)
84-95: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winPut remote and local workflows in tab sets.
These are parallel alternatives. Separate them into supported tab sets.
As per coding guidelines, use tab sets for parallel alternatives.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@plugins/nemo-auditor/README.md` around lines 84 - 95, Update the README examples around the remote job workflow and the local audit workflow to present them as separate tab sets, with each alternative in its own tab and the existing commands preserved within the appropriate tab.Source: Coding guidelines
plugins/nemo-auditor/src/nemo_auditor/sdk_resources/job_resources.py (2)
331-333: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueDrop the lambda and reuse
_extract_tarin the sync path.
asyncio.to_threadforwards positional arguments. The sync twin at Lines 219-220 duplicates the same extraction block.♻️ Proposed change
- await asyncio.to_thread( - lambda: _extract_tar(resp.content, output_path), - ) + await asyncio.to_thread(_extract_tar, resp.content, output_path)Then replace Lines 219-220 in the sync method:
_extract_tar(resp.content, output_path)🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@plugins/nemo-auditor/src/nemo_auditor/sdk_resources/job_resources.py` around lines 331 - 333, Update the asynchronous extraction call around _extract_tar to pass _extract_tar directly to asyncio.to_thread with resp.content and output_path as positional arguments, removing the lambda; also replace the duplicated extraction block in the synchronous method with a direct _extract_tar(resp.content, output_path) call.
174-193: 🚀 Performance & Scalability | 🔵 Trivial | 🏗️ Heavy liftFull log re-pagination every second scales poorly.
wait_until_donecallsget_logseach second, andget_logsre-pages the complete history. Transferred bytes grow quadratically with job duration, andaccept_logsthen discards everything except the tail.Persist the last
page_cursoracross polls, or apply backoff toWAIT_INTERVAL_SECONDS.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@plugins/nemo-auditor/src/nemo_auditor/sdk_resources/job_resources.py` around lines 174 - 193, The polling flow around get_logs and wait_until_done re-fetches the complete log history on every poll, causing quadratic transfer growth. Persist and reuse the latest page_cursor between polls so each wait_until_done iteration fetches only newly available pages, while preserving the existing log parsing and accept_logs behavior.plugins/nemo-auditor/tests/test_sdk_resources.py (2)
676-689: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAssert the archive member was extracted.
The test checks the return path and the URL only. Extraction could silently produce nothing and the test would still pass.
💚 Proposed addition
assert result == tmp_path + assert (tmp_path / "report.jsonl").read_bytes() == b"report data" platform._client.get.assert_called_with(🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@plugins/nemo-auditor/tests/test_sdk_resources.py` around lines 676 - 689, Update test_download_artifacts_extracts_tarball to assert that the expected archive member created by _make_tar_bytes exists under tmp_path after resource.download_artifacts completes, while preserving the existing return-path and request-URL assertions.
633-646: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd coverage for log parsing and level routing.
All wait tests return empty log pages.
get_logspagination and_WaitLogCollector.accept_logslevel handling stay untested, and that is where the KeyError defect onjob_resources.pyLines 58-71 lives. Add a case with a multi-page response and entries that carryname,levelname, andmessage, plus one entry missinglevelname.The async class also lacks the terminal-failure and
_poll_safethreshold tests that the sync class has.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@plugins/nemo-auditor/tests/test_sdk_resources.py` around lines 633 - 646, Extend wait-test coverage around test_wait_until_done_polls_until_completed with multi-page log responses containing name, levelname, and message fields plus an entry without levelname, verifying get_logs pagination and _WaitLogCollector.accept_logs routing without KeyError. Add matching async tests for terminal failure and _poll_safe threshold behavior, mirroring the existing synchronous-class coverage.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@docs/auditor/sdk-resources.mdx`:
- Around line 106-108: Correct the download_artifacts example and its preceding
comment so the output path consistently reflects extraction into ./my-reports,
removing the incorrect <job-name> subdirectory reference.
- Around line 88-90: Align the documented wait_until_done() contract with its
async implementation: update the async resource’s terminal-failure handling so
it raises RuntimeError instead of only logging and returning, preserving normal
completion behavior. Use the wait_until_done() implementation and its
terminal-status failure branch as the change points, or remove the RuntimeError
guarantee from the documentation if raising cannot be supported.
In `@plugins/nemo-auditor/src/nemo_auditor/sdk_resources/job_resources.py`:
- Around line 159-172: Update both wait_until_done loops in
plugins/nemo-auditor/src/nemo_auditor/sdk_resources/job_resources.py:159-172 and
:271-284, including their synchronous and asynchronous variants, to accept and
enforce a timeout. Reuse _status_is_complete and the known-status definitions so
each loop exits on None, any unknown status, or a known terminal status, while
preserving log collection and final-status reporting.
- Around line 58-71: Update accept_logs to safely handle log entries missing
name or levelname by using defaults or skipping incomplete entries, and ensure
_try_parse_log_message validates all three required keys—message, name, and
levelname—so malformed platform logs cannot propagate KeyError through
wait_until_done.
---
Nitpick comments:
In `@docs/auditor/sdk-resources.mdx`:
- Line 226: Update the internal link in the AsyncAuditorPluginResource.submit()
documentation to replace the local `#auditorjobresource` anchor with the canonical
Fern navigation URL for the AuditorJobResource section.
- Around line 65-116: Separate the “Submit and wait for an audit job” workflow,
including both code examples, from the reference page into a dedicated HOW-TO
page. Keep the submit argument and AuditorJobResource method tables in
sdk-resources.mdx as reference content, then add a cross-link from that page to
the new workflow page.
In `@plugins/nemo-auditor/README.md`:
- Around line 84-95: Update the README examples around the remote job workflow
and the local audit workflow to present them as separate tab sets, with each
alternative in its own tab and the existing commands preserved within the
appropriate tab.
In `@plugins/nemo-auditor/src/nemo_auditor/sdk_resources/job_resources.py`:
- Around line 331-333: Update the asynchronous extraction call around
_extract_tar to pass _extract_tar directly to asyncio.to_thread with
resp.content and output_path as positional arguments, removing the lambda; also
replace the duplicated extraction block in the synchronous method with a direct
_extract_tar(resp.content, output_path) call.
- Around line 174-193: The polling flow around get_logs and wait_until_done
re-fetches the complete log history on every poll, causing quadratic transfer
growth. Persist and reuse the latest page_cursor between polls so each
wait_until_done iteration fetches only newly available pages, while preserving
the existing log parsing and accept_logs behavior.
In `@plugins/nemo-auditor/tests/test_sdk_resources.py`:
- Around line 676-689: Update test_download_artifacts_extracts_tarball to assert
that the expected archive member created by _make_tar_bytes exists under
tmp_path after resource.download_artifacts completes, while preserving the
existing return-path and request-URL assertions.
- Around line 633-646: Extend wait-test coverage around
test_wait_until_done_polls_until_completed with multi-page log responses
containing name, levelname, and message fields plus an entry without levelname,
verifying get_logs pagination and _WaitLogCollector.accept_logs routing without
KeyError. Add matching async tests for terminal failure and _poll_safe threshold
behavior, mirroring the existing synchronous-class coverage.
🪄 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: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 906753d2-d7f7-4c10-8c28-8a291a4dcb6e
📒 Files selected for processing (6)
docs/auditor/sdk-resources.mdxe2e/auditor/test_audit_job.pyplugins/nemo-auditor/README.mdplugins/nemo-auditor/src/nemo_auditor/sdk.pyplugins/nemo-auditor/src/nemo_auditor/sdk_resources/job_resources.pyplugins/nemo-auditor/tests/test_sdk_resources.py
| def accept_logs(self, current_logs: list[dict[str, str]]) -> None: | ||
| for log in current_logs[len(self.seen_logs) :]: | ||
| self.seen_logs.append(log) | ||
| if not log["name"].startswith("nemo_auditor"): | ||
| continue | ||
| level = log["levelname"].lower() | ||
| if level == "info": | ||
| logger.info(log["message"]) | ||
| elif level in {"warning", "warn"}: | ||
| logger.warning(log["message"]) | ||
| self.warning_occurred = True | ||
| elif level == "error": | ||
| logger.error(log["message"]) | ||
| self.error_occurred = True |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
accept_logs raises KeyError on partial log payloads.
_try_parse_log_message only validates the message key (Line 119). Any payload without name or levelname makes Line 61 or Line 63 raise KeyError. The KeyError escapes wait_until_done, so a single malformed platform log entry aborts the blocking submit.
Use defaults, or require all three keys in _try_parse_log_message.
🐛 Proposed fix
def accept_logs(self, current_logs: list[dict[str, str]]) -> None:
for log in current_logs[len(self.seen_logs) :]:
self.seen_logs.append(log)
- if not log["name"].startswith("nemo_auditor"):
+ if not log.get("name", "").startswith("nemo_auditor"):
continue
- level = log["levelname"].lower()
+ level = log.get("levelname", "").lower()📝 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.
| def accept_logs(self, current_logs: list[dict[str, str]]) -> None: | |
| for log in current_logs[len(self.seen_logs) :]: | |
| self.seen_logs.append(log) | |
| if not log["name"].startswith("nemo_auditor"): | |
| continue | |
| level = log["levelname"].lower() | |
| if level == "info": | |
| logger.info(log["message"]) | |
| elif level in {"warning", "warn"}: | |
| logger.warning(log["message"]) | |
| self.warning_occurred = True | |
| elif level == "error": | |
| logger.error(log["message"]) | |
| self.error_occurred = True | |
| def accept_logs(self, current_logs: list[dict[str, str]]) -> None: | |
| for log in current_logs[len(self.seen_logs) :]: | |
| self.seen_logs.append(log) | |
| if not log.get("name", "").startswith("nemo_auditor"): | |
| continue | |
| level = log.get("levelname", "").lower() | |
| if level == "info": | |
| logger.info(log["message"]) | |
| elif level in {"warning", "warn"}: | |
| logger.warning(log["message"]) | |
| self.warning_occurred = True | |
| elif level == "error": | |
| logger.error(log["message"]) | |
| self.error_occurred = True |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@plugins/nemo-auditor/src/nemo_auditor/sdk_resources/job_resources.py` around
lines 58 - 71, Update accept_logs to safely handle log entries missing name or
levelname by using defaults or skipping incomplete entries, and ensure
_try_parse_log_message validates all three required keys—message, name, and
levelname—so malformed platform logs cannot propagate KeyError through
wait_until_done.
| def wait_until_done(self) -> None: | ||
| """Block until the job reaches a terminal status, streaming logs along the way.""" | ||
| log_collector = _WaitLogCollector.create() | ||
| job_status = self.get_job_status() | ||
| while job_status != "completed": | ||
| _pause(WAIT_INTERVAL_SECONDS) | ||
| current_logs = self._poll_safe(self.get_logs, log_collector.seen_logs) | ||
| log_collector.accept_logs(current_logs) | ||
| if job_status in TERMINAL_INCOMPLETE_STATUSES: | ||
| log_collector.error_occurred = True | ||
| logger.error(f"Audit job terminated with status {job_status!r}.") | ||
| break | ||
| job_status = self._poll_safe(self.get_job_status, job_status) | ||
| log_collector.log_final_status() |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | 🏗️ Heavy lift
Both wait_until_done loops can spin forever. Each loop exits only on "completed" or on a member of TERMINAL_INCOMPLETE_STATUSES. get_job_status returns resp.json().get("status"), so None is reachable, and any platform status outside that set never terminates. _status_is_complete already handles unknown statuses explicitly, so the wait loops are inconsistent with it.
plugins/nemo-auditor/src/nemo_auditor/sdk_resources/job_resources.py#L159-L172: add a timeout parameter, and break onNoneor on any status outside the known set.plugins/nemo-auditor/src/nemo_auditor/sdk_resources/job_resources.py#L271-L284: apply the same timeout and terminal-status handling in the async twin.
📍 Affects 1 file
plugins/nemo-auditor/src/nemo_auditor/sdk_resources/job_resources.py#L159-L172(this comment)plugins/nemo-auditor/src/nemo_auditor/sdk_resources/job_resources.py#L271-L284
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@plugins/nemo-auditor/src/nemo_auditor/sdk_resources/job_resources.py` around
lines 159 - 172, Update both wait_until_done loops in
plugins/nemo-auditor/src/nemo_auditor/sdk_resources/job_resources.py:159-172 and
:271-284, including their synchronous and asynchronous variants, to accept and
enforce a timeout. Reuse _status_is_complete and the known-status definitions so
each loop exits on None, any unknown status, or a known terminal status, while
preserving log collection and final-status reporting.
|
643bff8 to
5993975
Compare
Signed-off-by: Paul A. Parkanzky <parkanzky@users.noreply.github.com>
Signed-off-by: Paul A. Parkanzky <parkanzky@users.noreply.github.com>
5993975 to
01fe4e5
Compare
There was a problem hiding this comment.
🧹 Nitpick comments (4)
plugins/nemo-auditor/tests/test_sdk_resources.py (4)
568-568: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueRemove the unused
_STATUS_PAYLOADconstant.No test in the added blocks references it. The status tests build inline dicts instead.
♻️ Proposed cleanup
-_STATUS_PAYLOAD = {"name": "audit-job-abc123", "status": "active", "workspace": "default"} - -🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@plugins/nemo-auditor/tests/test_sdk_resources.py` at line 568, Remove the unused _STATUS_PAYLOAD constant from the test module; keep the status tests’ existing inline payload dictionaries unchanged.
699-708: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDerive the seed count from
MAX_CONSECUTIVE_POLL_ERRORS.The literal
4encodes an assumption that the constant equals 5. If the constant changes, this test passes for the wrong reason:_poll_safereturns the fallback and thepytest.raisesblock fails with an unclear message.♻️ Proposed refactor
def test_poll_safe_raises_after_max_consecutive_errors(self) -> None: _, resource = self._make_resource() - resource._consecutive_poll_errors = 4 + resource._consecutive_poll_errors = MAX_CONSECUTIVE_POLL_ERRORS - 1Import it alongside the resource classes:
from nemo_auditor.sdk_resources.job_resources import ( MAX_CONSECUTIVE_POLL_ERRORS, AsyncAuditorJobResource, AuditorJobResource, )🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@plugins/nemo-auditor/tests/test_sdk_resources.py` around lines 699 - 708, Update test_poll_safe_raises_after_max_consecutive_errors to import MAX_CONSECUTIVE_POLL_ERRORS from the job_resources module and initialize resource._consecutive_poll_errors to MAX_CONSECUTIVE_POLL_ERRORS - 1 instead of the literal 4, keeping the existing exception and reset assertions unchanged.
674-687: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAssert the extracted artifact, not only the return path.
Both download tests confirm the URL and returned directory. Neither confirms that extraction produced
report.jsonl.download_artifactsroutes throughsafe_extract_tar/_extract_tar, so extraction is the behavior worth pinning.💚 Proposed assertion
assert result == tmp_path + assert (tmp_path / "report.jsonl").read_bytes() == b"report data" platform._client.get.assert_called_with(🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@plugins/nemo-auditor/tests/test_sdk_resources.py` around lines 674 - 687, Update the download artifact tests around test_download_artifacts_extracts_tarball to assert that extraction creates report.jsonl under tmp_path, in addition to the existing return-path and URL checks. Cover the safe_extract_tar/_extract_tar extraction outcome without changing the download setup.
711-780: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd async
_poll_safecoverage.The sync class tests the transient-fallback path and the raise-after-max path.
AsyncAuditorJobResource._poll_safehas separate await logic and no equivalent tests. Mirror both cases with an async failing callable.Also consider asserting the terminal-failure path (
status: "error") for the asyncwait_until_done, matchingtest_wait_until_done_exits_on_terminal_failure.Want me to draft these tests?
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@plugins/nemo-auditor/tests/test_sdk_resources.py` around lines 711 - 780, Extend TestAsyncAuditorJobResource with async _poll_safe tests using a failing async callable: verify transient failures return the fallback value and repeated failures raise after the configured maximum attempts, mirroring the synchronous coverage. Also add async wait_until_done coverage for a terminal {"status": "error"} response, matching the synchronous terminal-failure test.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Nitpick comments:
In `@plugins/nemo-auditor/tests/test_sdk_resources.py`:
- Line 568: Remove the unused _STATUS_PAYLOAD constant from the test module;
keep the status tests’ existing inline payload dictionaries unchanged.
- Around line 699-708: Update test_poll_safe_raises_after_max_consecutive_errors
to import MAX_CONSECUTIVE_POLL_ERRORS from the job_resources module and
initialize resource._consecutive_poll_errors to MAX_CONSECUTIVE_POLL_ERRORS - 1
instead of the literal 4, keeping the existing exception and reset assertions
unchanged.
- Around line 674-687: Update the download artifact tests around
test_download_artifacts_extracts_tarball to assert that extraction creates
report.jsonl under tmp_path, in addition to the existing return-path and URL
checks. Cover the safe_extract_tar/_extract_tar extraction outcome without
changing the download setup.
- Around line 711-780: Extend TestAsyncAuditorJobResource with async _poll_safe
tests using a failing async callable: verify transient failures return the
fallback value and repeated failures raise after the configured maximum
attempts, mirroring the synchronous coverage. Also add async wait_until_done
coverage for a terminal {"status": "error"} response, matching the synchronous
terminal-failure test.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 118b6ca3-fda7-430c-b915-812987104a0a
📒 Files selected for processing (6)
docs/auditor/sdk-resources.mdxe2e/auditor/test_audit_job.pyplugins/nemo-auditor/README.mdplugins/nemo-auditor/src/nemo_auditor/sdk.pyplugins/nemo-auditor/src/nemo_auditor/sdk_resources/job_resources.pyplugins/nemo-auditor/tests/test_sdk_resources.py
🚧 Files skipped from review as they are similar to previous changes (5)
- e2e/auditor/test_audit_job.py
- plugins/nemo-auditor/README.md
- plugins/nemo-auditor/src/nemo_auditor/sdk.py
- docs/auditor/sdk-resources.mdx
- plugins/nemo-auditor/src/nemo_auditor/sdk_resources/job_resources.py
1 participant
Add ability to block on auditor job completion within the SDK. This encapsulates the polling that a user would have to do to wait for an Auditor job to complete.
Summary by CodeRabbit
New Features
Documentation