UN-2609 Added graceful shutdown support for tool containers with signal handling and configurable timeout#194
Conversation
Summary by CodeRabbit
Summary by CodeRabbit
WalkthroughSignal handling for SIGTERM and SIGINT has been added to the Changes
Sequence Diagram(s)sequenceDiagram
participant User
participant ToolEntrypoint
participant OS
User->>ToolEntrypoint: launch()
ToolEntrypoint->>ToolEntrypoint: Register _signal_handler for SIGTERM/SIGINT
ToolEntrypoint->>ToolEntrypoint: Parse arguments and execute tool
OS-->>ToolEntrypoint: Send SIGTERM/SIGINT
ToolEntrypoint->>ToolEntrypoint: _signal_handler logs signal (does not exit)
📜 Recent review detailsConfiguration used: CodeRabbit UI 📒 Files selected for processing (1)
✅ Files skipped from review due to trivial changes (1)
✨ Finishing Touches
🪧 TipsChatThere are 3 ways to chat with CodeRabbit:
SupportNeed help? Create a ticket on our support page for assistance with any issues or questions. Note: Be mindful of the bot's finite context window. It's strongly recommended to break down tasks such as reading entire modules into smaller chunks. For a focused discussion, use review comments to chat about specific files and their changes, instead of using the PR comments. CodeRabbit Commands (Invoked using PR comments)
Other keywords and placeholders
CodeRabbit Configuration File (
|
There was a problem hiding this comment.
Actionable comments posted: 1
🔭 Outside diff range comments (1)
src/unstract/sdk/tool/entrypoint.py (1)
12-37: Missing configurable timeout functionality mentioned in PR objectives.The PR objectives explicitly mention "graceful shutdown support for tool containers with signal handling and configurable timeout", but the current implementation lacks any timeout mechanism. This is a significant gap between the stated requirements and the actual implementation.
The implementation needs:
- A configurable timeout parameter
- A mechanism to enforce the timeout during graceful shutdown
- Integration with the tool execution to allow proper cleanup
Do you want me to generate a complete implementation that includes the configurable timeout mechanism mentioned in the PR objectives?
🧹 Nitpick comments (1)
src/unstract/sdk/tool/entrypoint.py (1)
31-34: Good signal registration placement, but fix whitespace issue.The signal registration is correctly placed before tool execution begins. However, there's a whitespace issue on line 34.
Fix the whitespace issue:
# Register signal handlers for graceful shutdown signal.signal(signal.SIGTERM, ToolEntrypoint._signal_handler) signal.signal(signal.SIGINT, ToolEntrypoint._signal_handler) - +Once you implement the actual graceful shutdown mechanism, you'll also need to check the shutdown flag during tool execution and handle cleanup appropriately.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
Cache: Disabled due to Reviews > Disable Cache setting
Knowledge Base: Disabled due to Reviews > Disable Knowledge Base setting
📒 Files selected for processing (1)
src/unstract/sdk/tool/entrypoint.py(3 hunks)
🧰 Additional context used
🪛 Ruff (0.11.9)
src/unstract/sdk/tool/entrypoint.py
13-13: Missing return type annotation for staticmethod _signal_handler
Add return type annotation: None
(ANN205)
13-13: Missing type annotation for function argument signum
(ANN001)
13-13: Missing type annotation for function argument frame
(ANN001)
34-34: Blank line contains whitespace
Remove whitespace from blank line
(W293)
🔇 Additional comments (1)
src/unstract/sdk/tool/entrypoint.py (1)
2-2: LGTM!The signal import is correctly added and necessary for the signal handling functionality.
There was a problem hiding this comment.
Actionable comments posted: 0
♻️ Duplicate comments (2)
src/unstract/sdk/tool/entrypoint.py (2)
13-19: Incomplete graceful shutdown implementation.The current signal handler only logs the signal without implementing actual graceful shutdown logic or configurable timeout as mentioned in the PR objectives. This doesn't provide the "graceful shutdown support" described in the PR summary.
The implementation needs:
- Actual graceful shutdown mechanism with shutdown flags
- Configurable timeout functionality
- Proper coordination with the tool execution lifecycle
Consider implementing a more complete solution as suggested in the previous review comments.
1-2: Missing import for type annotation.The
Anytype is used in the signal handler but not imported.Add the missing import:
import logging import signal +from typing import Any
🧹 Nitpick comments (1)
src/unstract/sdk/tool/entrypoint.py (1)
34-34: Remove whitespace from blank line.Static analysis flagged unnecessary whitespace on this blank line.
- +
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
Cache: Disabled due to Reviews > Disable Cache setting
Knowledge Base: Disabled due to Reviews > Disable Knowledge Base setting
📒 Files selected for processing (1)
src/unstract/sdk/tool/entrypoint.py(2 hunks)
🧰 Additional context used
🧬 Code Graph Analysis (1)
src/unstract/sdk/tool/entrypoint.py (3)
src/unstract/sdk/tool/base.py (1)
BaseTool(24-313)src/unstract/sdk/tool/executor.py (1)
ToolExecutor(16-85)src/unstract/sdk/tool/parser.py (1)
ToolArgsParser(7-56)
🪛 Ruff (0.11.9)
src/unstract/sdk/tool/entrypoint.py
14-14: Undefined name Any
(F821)
34-34: Blank line contains whitespace
Remove whitespace from blank line
(W293)
🔇 Additional comments (1)
src/unstract/sdk/tool/entrypoint.py (1)
31-33: Signal registration looks correct.The signal handler registration is properly implemented and follows standard practices for handling SIGTERM and SIGINT signals.
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (1)
src/unstract/sdk/tool/entrypoint.py (1)
32-35: Good placement of signal handler registration.Registering the signal handlers early in the launch method ensures they're active throughout the tool's execution lifecycle. The placement before argument parsing is appropriate.
Minor formatting issue to address:
signal.signal(signal.SIGTERM, ToolEntrypoint._signal_handler) signal.signal(signal.SIGINT, ToolEntrypoint._signal_handler) - + parsed_args = ToolArgsParser.parse_args(args)
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
Cache: Disabled due to Reviews > Disable Cache setting
Knowledge Base: Disabled due to Reviews > Disable Knowledge Base setting
📒 Files selected for processing (1)
src/unstract/sdk/tool/entrypoint.py(2 hunks)
🧰 Additional context used
🧬 Code Graph Analysis (1)
src/unstract/sdk/tool/entrypoint.py (3)
src/unstract/sdk/tool/base.py (1)
BaseTool(24-313)src/unstract/sdk/tool/executor.py (1)
ToolExecutor(16-85)src/unstract/sdk/tool/parser.py (1)
ToolArgsParser(7-56)
🪛 Ruff (0.11.9)
src/unstract/sdk/tool/entrypoint.py
15-15: Dynamically typed expressions (typing.Any) are disallowed in frame
(ANN401)
35-35: Blank line contains whitespace
Remove whitespace from blank line
(W293)
🔇 Additional comments (2)
src/unstract/sdk/tool/entrypoint.py (2)
1-9: LGTM! Well-structured imports and logger setup.The imports are appropriate for signal handling functionality, and the module-level logger setup follows best practices for use in the signal handler.
14-20: Excellent implementation of signal handling!The signal handler implementation is well-designed:
- Uses proper type annotations
- Implements the improved signal name retrieval using
signal.Signals(signum)as suggested in past reviews- Correctly logs the signal without exiting, which aligns with K8s graceful shutdown patterns
- Allows the container to complete within the K8s grace period before receiving SIGKILL
This approach is optimal for containerized environments where the orchestrator manages timeout enforcement.
chandrasekharan-zipstack
left a comment
There was a problem hiding this comment.
Make sure to bump the SDK version also
gaya3-zipstack
left a comment
There was a problem hiding this comment.
Increment the SDK version in init.py
gaya3-zipstack
left a comment
There was a problem hiding this comment.
Increment version in init.py
…m:Zipstack/unstract-sdk into UN-2609-Graceful-shutdown-for-tool-script
What
Why
How
tool/sctructureto accept graceful shutdownCan this PR break any existing features. If yes, please list possible items. If no, please explain why. (PS: Admins do not merge the PR without this section filled)
Relevant Docs
Related Issues or PRs
Dependencies Versions / Env Variables
Notes on Testing
Screenshots
...
Checklist
I have read and understood the Contribution Guidelines.