Add pipeline run stop support for Vertex AI orchestrator#4604
Add pipeline run stop support for Vertex AI orchestrator#4604martinpanacek789 wants to merge 1 commit intozenml-io:developfrom
Conversation
|
|
Implement _stop_run() on VertexOrchestrator, enabling users to stop running pipeline runs via CLI or dashboard. Handles both static (PipelineJob) and dynamic (CustomJob) pipelines by calling the appropriate Google Cloud cancel API. Add tests for graceful stop, forceful stop (static and dynamic), and missing run ID error case. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
martinpanacek789
force-pushed
the
feat/vertex-stop-run
branch
from
d5c4c6e to
f8fb789
Compare
| if graceful: | ||
| logger.info( | ||
| "Graceful stop requested - letting the Vertex AI pipeline " | ||
| "finish naturally." | ||
| ) | ||
| return |
There was a problem hiding this comment.
I think this should work, but only as long as step heartbeat is enabled. So we might need to check for that here, and fail if it's not as otherwise I don't think we can gracefully shut down on vertex. This is something that @bcdurak can confirm maybe.
| if METADATA_ORCHESTRATOR_RUN_ID in run.run_metadata: | ||
| run_id = run.run_metadata[METADATA_ORCHESTRATOR_RUN_ID] | ||
| elif run.orchestrator_run_id is not None: | ||
| run_id = run.orchestrator_run_id |
There was a problem hiding this comment.
This should usually be fine (as in, they should be the same value), but it runs the risk that some users (on purpose or accidentally) modify the orchestrator run ID in the metadata.
For that reason, I would suggest we swap these conditions around and use the run.orchestrator_run_id as the primary source of truth if it exists.
schustmi
added
the
release-notes
|
|
|
🔒 This PR has been automatically closed due to 4 weeks of inactivity. |
3 participants
Summary
_stop_run()onVertexOrchestrator, enabling users to stop running Vertex AI pipeline runs via CLI (zenml pipeline runs stop) or the dashboardPipelineJob.cancel()) and dynamic pipelines (CustomJob.cancel())Test plan
mypypasses on the changed filebash scripts/format.shproduces no changes🤖 Generated with Claude Code