Skip to content

UN-2609 Added graceful shutdown support for tool containers with signal handling and configurable timeout#194

Merged
kirtimanmishrazipstack merged 6 commits intomainfrom
UN-2609-Graceful-shutdown-for-tool-script
Jul 15, 2025
Merged

UN-2609 Added graceful shutdown support for tool containers with signal handling and configurable timeout#194
kirtimanmishrazipstack merged 6 commits intomainfrom
UN-2609-Graceful-shutdown-for-tool-script

Conversation

@kirtimanmishrazipstack
Copy link
Contributor

@kirtimanmishrazipstack kirtimanmishrazipstack commented Jul 14, 2025

What

  • Added graceful shutdown support for structure tool container with signal handling and configurable timeout

Why

  • Tool container should end gracefully

How

  • Changes in tool/sctructure to accept graceful shutdown

Can 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)

  • No. It is is related to tool which accepts SIGTERM

Relevant Docs

Related Issues or PRs

Dependencies Versions / Env Variables

Notes on Testing

# Structure Tool SIGTERM Testing Documentation

## Overview
This document outlines the changes made to test SIGTERM handling in the structure tool using local unstract-sdk modifications. The goal was to integrate local SDK changes that include SIGTERM signal handling and verify the tool responds gracefully to termination signals.

## Directory Structure & Workflow

### Build vs Run Locations
- **Build Location**: `/home/kirtimanmishra/Zipstack/Repositories/` (repositories root)
- **Run Location**: `/home/kirtimanmishra/Zipstack/Repositories/unstract/tools/structure/` (structure tool directory)

### Why Build from Repositories Root?
Building from the repositories root allows access to both:
- `unstract-sdk/` directory (contains local SDK changes)
- `unstract/tools/structure/` directory (contains the structure tool)

This enables copying the local SDK into the Docker build context.

## Changes Made

### 1. Dockerfile Changes
**File**: `unstract/tools/structure/Dockerfile`

Added before the requirements installation:
# Copy the unstract-sdk directory first
COPY unstract-sdk ${APP_HOME}/unstract-sdk/

**Purpose**: Copies the local unstract-sdk (with SIGTERM changes) into the Docker image build context.

### 2. Requirements.txt Changes
**File**: `unstract/tools/structure/requirements.txt`

**Before**:
unstract-sdk[aws]~=0.74.0

**After**:
# unstract-sdk[aws]~=0.74.0
-e ./unstract-sdk
s3fs

**Purpose**: 
- Commented out the PyPI version
- Added local editable install of unstract-sdk
- Added s3fs dependency

### 3. Main.py Changes
**File**: `unstract/tools/structure/src/main.py`

**Added import**:
import time

**Added in run method** (after initial setup):
# Add delay for SIGTERM testing
self.stream_log("Starting 30-second delay for SIGTERM testing...")
time.sleep(30)
self.stream_log("Delay completed, continuing with tool execution...")

**Purpose**: Provides a 30-second window during tool execution to test SIGTERM signal handling.

## Build Instructions

### Prerequisites
1. Local unstract-sdk with SIGTERM changes
2. Structure tool source code
3. Docker installed

### Step 1: Copy SDK to Build Context
# From repositories root
cd /home/kirtimanmishra/Zipstack/Repositories/
cp -r unstract-sdk unstract/tools/structure/

### Step 2: Build Docker Image
# From repositories root
cd /home/kirtimanmishra/Zipstack/Repositories/
docker build -f unstract/tools/structure/Dockerfile -t structure-tool-test .

**Alternative** (from structure tool directory):
# From structure tool directory (after copying SDK)
cd /home/kirtimanmishra/Zipstack/Repositories/unstract/tools/structure/
docker build -t structure-tool-test .

## Run Instructions

### Navigate to Structure Tool Directory
cd /home/kirtimanmishra/Zipstack/Repositories/unstract/tools/structure/

### Run Container (Detached for SIGTERM Testing)
docker run -d --name sigterm-test \
  --network unstract-network \
  --env-file .env \
  -v "$(pwd)/data_dir:/app/data_dir" \
  structure-tool-test:latest \
  --command RUN \
  --settings '{"prompt_registry_id": "479a3cbc-6872-4945-b839-366f6a04751d", "enable_single_pass_extraction": true, "summarize_as_source": true}' \
  --log-level DEBUG

## SIGTERM Testing Methods

### Method 1: docker stop 
# Get container ID
docker ps

# Send SIGTERM with 10-second grace period
docker stop <container-id>

**What happens**:
1. Docker sends SIGTERM to container
2. Waits 10 seconds for graceful shutdown
3. Sends SIGKILL if container hasn't stopped

### Method 2: Ctrl+C (Interactive Mode)
# Run in interactive mode
docker run --rm -it ... structure-tool-test ...

# Press Ctrl+C during the 30-second delay
# This sends SIGINT, which should also be handled

### Method 3: Direct SIGTERM Signal
# Send SIGTERM directly
docker kill --signal=SIGTERM <container-id>

# Or by name
docker kill --signal=SIGTERM sigterm-test


Screenshots

...

Checklist

I have read and understood the Contribution Guidelines.

@coderabbitai
Copy link
Contributor

coderabbitai bot commented Jul 14, 2025

Summary by CodeRabbit

  • New Features

    • Improved handling of system signals to enhance stability when receiving termination or interrupt requests.
  • Chores

    • Updated version number to v0.76.0.

Summary by CodeRabbit

  • New Features
    • Improved handling of termination signals (SIGTERM and SIGINT) to ensure proper logging and process management when running tools.

Walkthrough

Signal handling for SIGTERM and SIGINT has been added to the ToolEntrypoint class. A static method _signal_handler logs when these signals are received without terminating the process. The launch method now registers this handler for both signals before proceeding with argument parsing and tool execution. Additionally, the SDK version was updated from v0.75.0 to v0.76.0.

Changes

File Change Summary
src/unstract/sdk/tool/entrypoint.py Added static _signal_handler for SIGTERM/SIGINT; updated launch to register this handler.
src/unstract/sdk/init.py Updated SDK version string from "v0.75.0" to "v0.76.0".

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)
Loading

📜 Recent 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

📥 Commits

Reviewing files that changed from the base of the PR and between 88c1d17 and 68ce169.

📒 Files selected for processing (1)
  • src/unstract/sdk/__init__.py (1 hunks)
✅ Files skipped from review due to trivial changes (1)
  • src/unstract/sdk/init.py
✨ Finishing Touches
  • 📝 Generate Docstrings

🪧 Tips

Chat

There are 3 ways to chat with CodeRabbit:

  • Review comments: Directly reply to a review comment made by CodeRabbit. Example:
    • I pushed a fix in commit <commit_id>, please review it.
    • Explain this complex logic.
    • Open a follow-up GitHub issue for this discussion.
  • Files and specific lines of code (under the "Files changed" tab): Tag @coderabbitai in a new review comment at the desired location with your query. Examples:
    • @coderabbitai explain this code block.
    • @coderabbitai modularize this function.
  • PR comments: Tag @coderabbitai in a new PR comment to ask questions about the PR branch. For the best results, please provide a very specific query, as very limited context is provided in this mode. Examples:
    • @coderabbitai gather interesting stats about this repository and render them as a table. Additionally, render a pie chart showing the language distribution in the codebase.
    • @coderabbitai read src/utils.ts and explain its main purpose.
    • @coderabbitai read the files in the src/scheduler package and generate a class diagram using mermaid and a README in the markdown format.
    • @coderabbitai help me debug CodeRabbit configuration file.

Support

Need 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)

  • @coderabbitai pause to pause the reviews on a PR.
  • @coderabbitai resume to resume the paused reviews.
  • @coderabbitai review to trigger an incremental review. This is useful when automatic reviews are disabled for the repository.
  • @coderabbitai full review to do a full review from scratch and review all the files again.
  • @coderabbitai summary to regenerate the summary of the PR.
  • @coderabbitai generate docstrings to generate docstrings for this PR.
  • @coderabbitai generate sequence diagram to generate a sequence diagram of the changes in this PR.
  • @coderabbitai resolve resolve all the CodeRabbit review comments.
  • @coderabbitai configuration to show the current CodeRabbit configuration for the repository.
  • @coderabbitai help to get help.

Other keywords and placeholders

  • Add @coderabbitai ignore anywhere in the PR description to prevent this PR from being reviewed.
  • Add @coderabbitai summary to generate the high-level summary at a specific location in the PR description.
  • Add @coderabbitai anywhere in the PR title to generate the title automatically.

CodeRabbit Configuration File (.coderabbit.yaml)

  • You can programmatically configure CodeRabbit by adding a .coderabbit.yaml file to the root of your repository.
  • Please see the configuration documentation for more information.
  • If your editor has YAML language server enabled, you can add the path at the top of this file to enable auto-completion and validation: # yaml-language-server: $schema=https://coderabbit.ai/integrations/schema.v2.json

Documentation and Community

  • Visit our Documentation for detailed information on how to use CodeRabbit.
  • Join our Discord Community to get help, request features, and share feedback.
  • Follow us on X/Twitter for updates and announcements.

Copy link
Contributor

@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: 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:

  1. A configurable timeout parameter
  2. A mechanism to enforce the timeout during graceful shutdown
  3. 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

📥 Commits

Reviewing files that changed from the base of the PR and between bbbbaaf and 523a09f.

📒 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.

Copy link
Contributor

@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: 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:

  1. Actual graceful shutdown mechanism with shutdown flags
  2. Configurable timeout functionality
  3. 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 Any type 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

📥 Commits

Reviewing files that changed from the base of the PR and between 523a09f and d3ca118.

📒 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.

Copy link
Contributor

@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: 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

📥 Commits

Reviewing files that changed from the base of the PR and between d3ca118 and 88c1d17.

📒 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.

Copy link
Contributor

@chandrasekharan-zipstack chandrasekharan-zipstack left a comment

Choose a reason for hiding this comment

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

Make sure to bump the SDK version also

Copy link
Contributor

@gaya3-zipstack gaya3-zipstack left a comment

Choose a reason for hiding this comment

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

Increment the SDK version in init.py

Copy link
Contributor

@gaya3-zipstack gaya3-zipstack left a comment

Choose a reason for hiding this comment

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

Increment version in init.py

…m:Zipstack/unstract-sdk into UN-2609-Graceful-shutdown-for-tool-script
Copy link
Contributor

@gaya3-zipstack gaya3-zipstack left a comment

Choose a reason for hiding this comment

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

Looks good.

@kirtimanmishrazipstack kirtimanmishrazipstack merged commit de5cafe into main Jul 15, 2025
2 checks passed
@kirtimanmishrazipstack kirtimanmishrazipstack deleted the UN-2609-Graceful-shutdown-for-tool-script branch July 15, 2025 07:44
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.

4 participants