caption-flow 0.4.2__tar.gz → 0.5.0__tar.gz

{caption_flow-0.4.2/src/caption_flow.egg-info → caption_flow-0.5.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: caption-flow
-Version: 0.4.2
+Version: 0.5.0
 Summary: Self-contained distributed community captioning system
 Author-email: bghira <bghira@users.github.com>
 License: MIT
@@ -11,7 +11,8 @@ Classifier: License :: OSI Approved :: MIT License
 Classifier: Programming Language :: Python :: 3
 Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
-Requires-Python: <3.13,>=3.11
+Classifier: Programming Language :: Python :: 3.13
+Requires-Python: <3.14,>=3.11
 Description-Content-Type: text/markdown
 License-File: LICENSE
 Requires-Dist: websockets>=12.0
@@ -25,7 +26,8 @@ Requires-Dist: pyyaml>=6.0
 Requires-Dist: certbot>=2.0.0
 Requires-Dist: numpy>=1.24.0
 Requires-Dist: pillow>=10.0.0
-Requires-Dist: vllm<0.11.0,>=0.10.0
+Requires-Dist: vllm<0.20.0,>=0.19.0
+Requires-Dist: transformers<6.0.0,>=5.0.0
 Requires-Dist: webdataset<2.0.0,>=1.0.2
 Requires-Dist: pandas<3.0.0,>=2.3.1
 Requires-Dist: arrow<2.0.0,>=1.3.0
@@ -64,13 +66,12 @@ a fast websocket-based orchestrator paired with lightweight gpu workers achieves
 
 ---
 
-## install
+## install from pypi
 
 ```bash
 python -m venv .venv
 source .venv/bin/activate  # windows: .venv\Scripts\activate
-pip install --upgrade pip
-pip install -e .  # installs the `caption-flow` command
+pip install caption-flow
 ```
 
 ## quickstart (single box)

{caption_flow-0.4.2 → caption_flow-0.5.0}/README.md

@@ -16,13 +16,12 @@ a fast websocket-based orchestrator paired with lightweight gpu workers achieves
 
 ---
 
-## install
+## install from pypi
 
 ```bash
 python -m venv .venv
 source .venv/bin/activate  # windows: .venv\Scripts\activate
-pip install --upgrade pip
-pip install -e .  # installs the `caption-flow` command
+pip install caption-flow
 ```
 
 ## quickstart (single box)

{caption_flow-0.4.2 → caption_flow-0.5.0}/pyproject.toml

@@ -1,9 +1,9 @@
 [project]
 name = "caption-flow"
-version = "0.4.2"
+version = "0.5.0"
 description = "Self-contained distributed community captioning system"
 readme = "README.md"
-requires-python = ">=3.11,<3.13"
+requires-python = ">=3.11,<3.14"
 license = { text = "MIT" }
 authors = [{ name = "bghira", email = "bghira@users.github.com" }]
 keywords = ["captioning", "distributed", "vllm", "dataset", "community"]
@@ -14,6 +14,7 @@ classifiers = [
     "Programming Language :: Python :: 3",
     "Programming Language :: Python :: 3.11",
     "Programming Language :: Python :: 3.12",
+    "Programming Language :: Python :: 3.13",
 ]
 
 dependencies = [
@@ -28,7 +29,8 @@ dependencies = [
     "certbot>=2.0.0",
     "numpy>=1.24.0",
     "pillow>=10.0.0",
-    "vllm (>=0.10.0,<0.11.0)",
+    "vllm (>=0.19.0,<0.20.0)",
+    "transformers (>=5.0.0,<6.0.0)",
     "webdataset (>=1.0.2,<2.0.0)",
     "pandas (>=2.3.1,<3.0.0)",
     "arrow (>=1.3.0,<2.0.0)",
@@ -64,7 +66,7 @@ where = ["src"]
 
 [tool.black]
 line-length = 100
-target-version = ['py310']
+target-version = ['py311', 'py312', 'py313']
 
 [tool.ruff]
 line-length = 100
@@ -85,7 +87,7 @@ ignore = [
 ]
 
 [tool.mypy]
-python_version = "3.11"
+python_version = "3.13"
 warn_return_any = true
 warn_unused_configs = true
 disallow_untyped_defs = true

{caption_flow-0.4.2 → caption_flow-0.5.0}/src/caption_flow/cli.py

@@ -214,6 +214,7 @@ def validate_orchestrator_auth_config(
 
     Raises:
         ValueError: If no auth configuration is found
+
     """
     # Check if auth is already in the orchestrator section
     if "auth" in config_data:
@@ -277,7 +278,7 @@ def validate_orchestrator_auth_config(
             "[yellow]Warning: No admin tokens configured - admin operations will be unavailable[/yellow]"
         )
 
-    console.print(f"[green]✓ Auth validation passed[/green]")
+    console.print("[green]✓ Auth validation passed[/green]")
     return config_data
 
 
@@ -363,6 +364,13 @@ def orchestrator(ctx, config: Optional[str], **kwargs):
 @click.option("--gpu-id", type=int, help="GPU device ID (for vLLM)")
 @click.option("--precision", help="Model precision (for vLLM)")
 @click.option("--model", help="Model name (for vLLM)")
+@click.option(
+    "--when_finished",
+    type=click.Choice(["stay_connected", "shutdown", "post_exec_hook"]),
+    default="stay_connected",
+    help="Action when all captions are complete (default: stay_connected)",
+)
+@click.option("--post_exec_hook", help="Path to executable for post_exec_hook action")
 @click.pass_context
 def worker(ctx, config: Optional[str], **kwargs):
     """Start a worker node."""
@@ -376,7 +384,18 @@ def worker(ctx, config: Optional[str], **kwargs):
         config_data = base_config
 
     # Apply CLI overrides (only non-None values)
-    for key in ["server", "token", "name", "batch_size", "gpu_id", "precision", "model"]:
+    cli_overrides = [
+        "server",
+        "token",
+        "name",
+        "batch_size",
+        "gpu_id",
+        "precision",
+        "model",
+        "when_finished",
+        "post_exec_hook",
+    ]
+    for key in cli_overrides:
         if kwargs.get(key) is not None:
             config_data[key] = kwargs[key]
 
@@ -391,6 +410,14 @@ def worker(ctx, config: Optional[str], **kwargs):
         console.print("[red]Error: --token required (or set in config)[/red]")
         sys.exit(1)
 
+    # Validate when_finished logic
+    when_finished = config_data.get("when_finished", "stay_connected")
+    if when_finished == "post_exec_hook" and not config_data.get("post_exec_hook"):
+        console.print(
+            "[red]Error: --post_exec_hook required when --when_finished=post_exec_hook[/red]"
+        )
+        sys.exit(1)
+
     # Choose worker type
     if kwargs.get("vllm") or config_data.get("vllm"):
         from .workers.caption import CaptionWorker
@@ -853,7 +880,7 @@ def add(ctx, role: str, name: str, token_value: str, no_reload: bool):
             console.print("[yellow]No admin token specified, skipping orchestrator reload[/yellow]")
             console.print("[dim]Use --token to reload orchestrator config[/dim]")
         else:
-            console.print(f"[cyan]Reloading orchestrator config...[/cyan]")
+            console.print("[cyan]Reloading orchestrator config...[/cyan]")
             success = asyncio.run(
                 _reload_orchestrator_config(server, admin_token, config_data, no_verify_ssl)
             )
@@ -923,7 +950,7 @@ def remove(ctx, role: str, identifier: str, no_reload: bool):
         elif not admin_token:
             console.print("[yellow]No admin token specified, skipping orchestrator reload[/yellow]")
         else:
-            console.print(f"[cyan]Reloading orchestrator config...[/cyan]")
+            console.print("[cyan]Reloading orchestrator config...[/cyan]")
             success = asyncio.run(
                 _reload_orchestrator_config(server, admin_token, config_data, no_verify_ssl)
             )
@@ -1008,7 +1035,7 @@ def reload_config(
         sys.exit(1)
 
     # Validate and normalize auth configuration for reload
-    console.print(f"[cyan]Validating configuration...[/cyan]")
+    console.print("[cyan]Validating configuration...[/cyan]")
     if "orchestrator" in new_cfg:
         orchestrator_config = new_cfg["orchestrator"]
         try:

{caption_flow-0.4.2 → caption_flow-0.5.0}/src/caption_flow/models.py

@@ -253,6 +253,20 @@ class StorageContents:
                     logger.warning(f"Row missing columns: {missing_cols}")
 
 
+class WhenFinished(Enum):
+    """Actions to take when caption worker finishes all work."""
+
+    STAY_CONNECTED = "stay_connected"
+    SHUTDOWN = "shutdown"
+    POST_EXEC_HOOK = "post_exec_hook"
+
+    def __str__(self):
+        return self.value
+
+    def to_json(self):
+        return self.value
+
+
 class ExportError(Exception):
     """Base exception for export-related errors."""
 

{caption_flow-0.4.2 → caption_flow-0.5.0}/src/caption_flow/workers/caption.py

@@ -18,7 +18,7 @@ import websockets
 from huggingface_hub import get_token
 from PIL import Image
 
-from ..models import ProcessingStage, StageResult
+from ..models import ProcessingStage, StageResult, WhenFinished
 from ..processors import (
     HuggingFaceDatasetWorkerProcessor,
     LocalFilesystemWorkerProcessor,
@@ -93,9 +93,11 @@ class MultiStageVLLMManager:
 
         # Load tokenizer and processor
         self.tokenizers[model_name] = AutoTokenizer.from_pretrained(
-            model_name, trust_remote_code=True, use_fast=True
+            model_name, trust_remote_code=True
+        )
+        self.processors[model_name] = AutoProcessor.from_pretrained(
+            model_name, trust_remote_code=True
         )
-        self.processors[model_name] = AutoProcessor.from_pretrained(model_name)
 
         # Initialize LLM
         vllm_params = {
@@ -217,6 +219,15 @@ class CaptionWorker(BaseWorker):
         # Processing control
         self.should_stop_processing = Event()
 
+        # When finished behavior
+        when_finished_str = config.get("when_finished", "stay_connected")
+        self.when_finished = WhenFinished(when_finished_str)
+        self.post_exec_hook = config.get("post_exec_hook")
+
+        # Track consecutive no work responses to detect completion
+        self.consecutive_no_work = 0
+        self.no_work_threshold = 3  # Consider work complete after 3 consecutive no_work responses
+
     def _init_metrics(self):
         """Initialize worker metrics."""
         self.items_processed = 0
@@ -347,13 +358,105 @@ class CaptionWorker(BaseWorker):
                     self.assigned_units.append(unit)
             logger.info(f"Received {len(assignment.units)} work units")
 
+            # Reset no work counter since we got work
+            self.consecutive_no_work = 0
+
         elif msg_type == "no_work":
             logger.info("No work available")
+            self.consecutive_no_work += 1
+
+            # Check if all work appears to be complete
+            if self.consecutive_no_work >= self.no_work_threshold:
+                logger.info(
+                    f"Received {self.consecutive_no_work} consecutive 'no work' responses, "
+                    "assuming all captions are complete"
+                )
+                await self._handle_work_completion()
+                return
+
             await asyncio.sleep(10)
 
             if self.websocket and self.connected.is_set():
                 await self.websocket.send(json.dumps({"type": "get_work_units", "count": 2}))
 
+    async def _handle_work_completion(self):
+        """Handle when all captions are complete based on when_finished setting."""
+        logger.info(f"Handling work completion with action: {self.when_finished}")
+
+        if self.when_finished == WhenFinished.STAY_CONNECTED:
+            logger.info("Staying connected and continuing to poll for work")
+            return
+
+        elif self.when_finished == WhenFinished.SHUTDOWN:
+            logger.info("All captions complete, shutting down worker")
+            await self._shutdown_worker()
+
+        elif self.when_finished == WhenFinished.POST_EXEC_HOOK:
+            if not self.post_exec_hook:
+                logger.error("post_exec_hook action specified but no hook path provided")
+                return
+
+            logger.info(f"All captions complete, executing hook: {self.post_exec_hook}")
+            await self._execute_post_hook()
+
+    async def _shutdown_worker(self):
+        """Clean up and shut down the worker."""
+        logger.info("Shutting down worker...")
+
+        # Stop processing
+        self.should_stop_processing.set()
+
+        # Set running to False to exit the main loop
+        self.running = False
+
+        # Disconnect gracefully
+        if self.websocket:
+            try:
+                await self.websocket.close()
+            except Exception as e:
+                logger.warning(f"Error closing websocket: {e}")
+
+    async def _execute_post_hook(self):
+        """Execute the post-execution hook."""
+        import os
+        from pathlib import Path
+
+        hook_path = Path(self.post_exec_hook)
+
+        if not hook_path.exists():
+            logger.error(f"Post-exec hook not found: {hook_path}")
+            return
+
+        if not os.access(hook_path, os.X_OK):
+            logger.error(f"Post-exec hook is not executable: {hook_path}")
+            return
+
+        try:
+            logger.info(f"Executing post-exec hook: {hook_path}")
+
+            # Execute the hook
+            process = await asyncio.create_subprocess_exec(
+                str(hook_path), stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE
+            )
+
+            stdout, stderr = await process.communicate()
+
+            if process.returncode == 0:
+                logger.info("Post-exec hook completed successfully")
+                if stdout:
+                    logger.info(f"Hook output: {stdout.decode()}")
+            else:
+                logger.error(f"Post-exec hook failed with return code {process.returncode}")
+                if stderr:
+                    logger.error(f"Hook error: {stderr.decode()}")
+
+        except Exception as e:
+            logger.error(f"Error executing post-exec hook: {e}")
+
+        finally:
+            # Shut down after executing hook
+            await self._shutdown_worker()
+
     def _parse_stages_config(self, vllm_config: Dict[str, Any]) -> List[ProcessingStage]:
         """Parse stages configuration from vLLM config."""
         stages_config = vllm_config.get("stages", [])

{caption_flow-0.4.2 → caption_flow-0.5.0/src/caption_flow.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: caption-flow
-Version: 0.4.2
+Version: 0.5.0
 Summary: Self-contained distributed community captioning system
 Author-email: bghira <bghira@users.github.com>
 License: MIT
@@ -11,7 +11,8 @@ Classifier: License :: OSI Approved :: MIT License
 Classifier: Programming Language :: Python :: 3
 Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
-Requires-Python: <3.13,>=3.11
+Classifier: Programming Language :: Python :: 3.13
+Requires-Python: <3.14,>=3.11
 Description-Content-Type: text/markdown
 License-File: LICENSE
 Requires-Dist: websockets>=12.0
@@ -25,7 +26,8 @@ Requires-Dist: pyyaml>=6.0
 Requires-Dist: certbot>=2.0.0
 Requires-Dist: numpy>=1.24.0
 Requires-Dist: pillow>=10.0.0
-Requires-Dist: vllm<0.11.0,>=0.10.0
+Requires-Dist: vllm<0.20.0,>=0.19.0
+Requires-Dist: transformers<6.0.0,>=5.0.0
 Requires-Dist: webdataset<2.0.0,>=1.0.2
 Requires-Dist: pandas<3.0.0,>=2.3.1
 Requires-Dist: arrow<2.0.0,>=1.3.0
@@ -64,13 +66,12 @@ a fast websocket-based orchestrator paired with lightweight gpu workers achieves
 
 ---
 
-## install
+## install from pypi
 
 ```bash
 python -m venv .venv
 source .venv/bin/activate  # windows: .venv\Scripts\activate
-pip install --upgrade pip
-pip install -e .  # installs the `caption-flow` command
+pip install caption-flow
 ```
 
 ## quickstart (single box)

{caption_flow-0.4.2 → caption_flow-0.5.0}/src/caption_flow.egg-info/requires.txt

@@ -9,7 +9,8 @@ pyyaml>=6.0
 certbot>=2.0.0
 numpy>=1.24.0
 pillow>=10.0.0
-vllm<0.11.0,>=0.10.0
+vllm<0.20.0,>=0.19.0
+transformers<6.0.0,>=5.0.0
 webdataset<2.0.0,>=1.0.2
 pandas<3.0.0,>=2.3.1
 arrow<2.0.0,>=1.3.0

{caption_flow-0.4.2 → caption_flow-0.5.0}/tests/test_cli.py

@@ -857,5 +857,126 @@ class TestReloadConfigCommand:
         assert "✓ Auth validation passed" in result.output
 
 
+class TestWhenFinishedCLI:
+    """Test CLI validation for when_finished parameters."""
+
+    def test_worker_command_when_finished_validation_success(self, runner):
+        """Test successful worker command with when_finished parameters."""
+        with patch("caption_flow.workers.caption.CaptionWorker") as mock_worker:
+            mock_worker_instance = Mock()
+            mock_worker.return_value = mock_worker_instance
+            mock_worker_instance.start = AsyncMock()
+
+            # Test valid combinations
+            result = runner.invoke(
+                main,
+                [
+                    "worker",
+                    "--server",
+                    "ws://localhost:8765",
+                    "--token",
+                    "test-token",
+                    "--vllm",
+                    "--when_finished",
+                    "stay_connected",
+                ],
+            )
+            # This will fail due to async mocking, but should pass validation
+            assert "--post_exec_hook required" not in result.output
+
+    def test_worker_command_when_finished_shutdown(self, runner):
+        """Test worker command with shutdown option."""
+        with patch("caption_flow.workers.caption.CaptionWorker") as mock_worker:
+            mock_worker_instance = Mock()
+            mock_worker.return_value = mock_worker_instance
+            mock_worker_instance.start = AsyncMock()
+
+            result = runner.invoke(
+                main,
+                [
+                    "worker",
+                    "--server",
+                    "ws://localhost:8765",
+                    "--token",
+                    "test-token",
+                    "--vllm",
+                    "--when_finished",
+                    "shutdown",
+                ],
+            )
+            assert "--post_exec_hook required" not in result.output
+
+    def test_worker_command_when_finished_post_hook_missing(self, runner):
+        """Test worker command fails when post_exec_hook is missing."""
+        result = runner.invoke(
+            main,
+            [
+                "worker",
+                "--server",
+                "ws://localhost:8765",
+                "--token",
+                "test-token",
+                "--vllm",
+                "--when_finished",
+                "post_exec_hook",
+                # Missing --post_exec_hook
+            ],
+        )
+
+        assert result.exit_code == 1
+        assert "--post_exec_hook required when --when_finished=post_exec_hook" in result.output
+
+    def test_worker_command_when_finished_post_hook_valid(self, runner, tmp_path):
+        """Test worker command with valid post_exec_hook."""
+        hook_script = tmp_path / "test_hook.sh"
+        hook_script.write_text("#!/bin/bash\necho 'test'")
+        hook_script.chmod(0o755)
+
+        with patch("caption_flow.workers.caption.CaptionWorker") as mock_worker:
+            mock_worker_instance = Mock()
+            mock_worker.return_value = mock_worker_instance
+            mock_worker_instance.start = AsyncMock()
+
+            result = runner.invoke(
+                main,
+                [
+                    "worker",
+                    "--server",
+                    "ws://localhost:8765",
+                    "--token",
+                    "test-token",
+                    "--vllm",
+                    "--when_finished",
+                    "post_exec_hook",
+                    "--post_exec_hook",
+                    str(hook_script),
+                ],
+            )
+            # Should pass validation
+            assert "--post_exec_hook required" not in result.output
+
+    def test_worker_command_when_finished_default(self, runner):
+        """Test worker command uses default when_finished value."""
+        with patch("caption_flow.workers.caption.CaptionWorker") as mock_worker:
+            mock_worker_instance = Mock()
+            mock_worker.return_value = mock_worker_instance
+            mock_worker_instance.start = AsyncMock()
+
+            result = runner.invoke(
+                main,
+                [
+                    "worker",
+                    "--server",
+                    "ws://localhost:8765",
+                    "--token",
+                    "test-token",
+                    "--vllm",
+                    # No when_finished specified, should default to stay_connected
+                ],
+            )
+            # Should work without errors
+            assert "--post_exec_hook required" not in result.output
+
+
 if __name__ == "__main__":
     pytest.main([__file__])

{caption_flow-0.4.2 → caption_flow-0.5.0}/tests/test_worker_caption.py

@@ -12,7 +12,7 @@ from PIL import Image
 # Import pytest-asyncio
 pytest_plugins = ("pytest_asyncio",)
 import pytest_asyncio
-from caption_flow.models import Caption, JobId, ProcessingStage
+from caption_flow.models import Caption, JobId, ProcessingStage, WhenFinished
 from caption_flow.processors import WorkAssignment, WorkUnit
 from caption_flow.storage import StorageManager
 
@@ -859,9 +859,9 @@ class TestCaptionWorker:
                 work_failed_call = sent_data
                 break
 
-        assert (
-            work_failed_call is not None
-        ), "work_failed message should have been sent for incomplete unit"
+        assert work_failed_call is not None, (
+            "work_failed message should have been sent for incomplete unit"
+        )
         assert work_failed_call["unit_id"] == "unit1"
         assert "Processing incomplete" in work_failed_call["error"]
         assert "1/3 items processed" in work_failed_call["error"]
@@ -1119,5 +1119,216 @@ class TestCaptionWorkerConfigReload:
         assert worker.items_failed == 1
 
 
+class TestWhenFinishedFunctionality:
+    """Test suite for when_finished functionality."""
+
+    @pytest.fixture
+    def worker_config_stay_connected(self):
+        """Create test worker config with stay_connected behavior."""
+        return {
+            "name": "test_worker",
+            "token": "test_token",
+            "server": "ws://localhost:8765",
+            "server_url": "ws://localhost:8765",
+            "gpu_id": 0,
+            "batch_image_processing": True,
+            "when_finished": "stay_connected",
+        }
+
+    @pytest.fixture
+    def worker_config_shutdown(self):
+        """Create test worker config with shutdown behavior."""
+        return {
+            "name": "test_worker",
+            "token": "test_token",
+            "server": "ws://localhost:8765",
+            "server_url": "ws://localhost:8765",
+            "gpu_id": 0,
+            "batch_image_processing": True,
+            "when_finished": "shutdown",
+        }
+
+    @pytest.fixture
+    def worker_config_post_hook(self, tmp_path):
+        """Create test worker config with post_exec_hook behavior."""
+        hook_script = tmp_path / "test_hook.sh"
+        hook_script.write_text("#!/bin/bash\necho 'Hook executed'\nexit 0")
+        hook_script.chmod(0o755)
+
+        return {
+            "name": "test_worker",
+            "token": "test_token",
+            "server": "ws://localhost:8765",
+            "server_url": "ws://localhost:8765",
+            "gpu_id": 0,
+            "batch_image_processing": True,
+            "when_finished": "post_exec_hook",
+            "post_exec_hook": str(hook_script),
+        }
+
+    def test_when_finished_initialization_stay_connected(self, worker_config_stay_connected):
+        """Test worker initialization with stay_connected setting."""
+        worker = create_fast_caption_worker(worker_config_stay_connected)
+
+        assert worker.when_finished == WhenFinished.STAY_CONNECTED
+        assert worker.post_exec_hook is None
+        assert worker.consecutive_no_work == 0
+        assert worker.no_work_threshold == 3
+
+    def test_when_finished_initialization_shutdown(self, worker_config_shutdown):
+        """Test worker initialization with shutdown setting."""
+        worker = create_fast_caption_worker(worker_config_shutdown)
+
+        assert worker.when_finished == WhenFinished.SHUTDOWN
+        assert worker.post_exec_hook is None
+
+    def test_when_finished_initialization_post_hook(self, worker_config_post_hook):
+        """Test worker initialization with post_exec_hook setting."""
+        worker = create_fast_caption_worker(worker_config_post_hook)
+
+        assert worker.when_finished == WhenFinished.POST_EXEC_HOOK
+        assert worker.post_exec_hook is not None
+        assert "test_hook.sh" in worker.post_exec_hook
+
+    @pytest.mark.asyncio
+    async def test_consecutive_no_work_counter(self, worker_config_stay_connected):
+        """Test consecutive no_work counter logic."""
+        worker = create_fast_caption_worker(worker_config_stay_connected)
+
+        # Mock websocket
+        mock_websocket = AsyncMock()
+        worker.websocket = mock_websocket
+        worker.connected.set()
+
+        # Simulate receiving work assignment - should reset counter
+        assignment_data = {
+            "type": "work_assignment",
+            "assignment": {
+                "assignment_id": "test_assignment",
+                "worker_id": "test_worker",
+                "units": [],
+                "assigned_at": datetime.now().isoformat(),
+            },
+        }
+
+        worker.consecutive_no_work = 2
+        await worker._handle_message(assignment_data)
+        assert worker.consecutive_no_work == 0
+
+        # Simulate receiving no_work messages
+        no_work_data = {"type": "no_work"}
+
+        # First no_work - should increment counter but not trigger completion
+        await worker._handle_message(no_work_data)
+        assert worker.consecutive_no_work == 1
+
+        # Second no_work - should increment counter but not trigger completion
+        await worker._handle_message(no_work_data)
+        assert worker.consecutive_no_work == 2
+
+        # Mock the completion handler to avoid actual shutdown
+        with patch.object(worker, "_handle_work_completion") as mock_completion:
+            # Third no_work - should trigger completion
+            await worker._handle_message(no_work_data)
+            assert worker.consecutive_no_work == 3
+            mock_completion.assert_called_once()
+
+    @pytest.mark.asyncio
+    async def test_handle_work_completion_stay_connected(self, worker_config_stay_connected):
+        """Test work completion handling with stay_connected action."""
+        worker = create_fast_caption_worker(worker_config_stay_connected)
+
+        # Should just log and return
+        await worker._handle_work_completion()
+        # No assertions needed - just verify it doesn't crash
+
+    @pytest.mark.asyncio
+    async def test_handle_work_completion_shutdown(self, worker_config_shutdown):
+        """Test work completion handling with shutdown action."""
+        worker = create_fast_caption_worker(worker_config_shutdown)
+
+        # Mock websocket
+        mock_websocket = AsyncMock()
+        worker.websocket = mock_websocket
+
+        await worker._handle_work_completion()
+
+        # Should set running to False and stop processing
+        assert worker.running is False
+        assert worker.should_stop_processing.is_set()
+
+    @pytest.mark.asyncio
+    async def test_handle_work_completion_post_hook_success(self, worker_config_post_hook):
+        """Test work completion handling with successful post_exec_hook."""
+        worker = create_fast_caption_worker(worker_config_post_hook)
+
+        # Mock websocket
+        mock_websocket = AsyncMock()
+        worker.websocket = mock_websocket
+
+        await worker._handle_work_completion()
+
+        # Should execute hook and then shutdown
+        assert worker.running is False
+        assert worker.should_stop_processing.is_set()
+
+    @pytest.mark.asyncio
+    async def test_handle_work_completion_post_hook_missing_path(self, worker_config_shutdown):
+        """Test work completion handling when post_exec_hook path is missing."""
+        worker = create_fast_caption_worker(worker_config_shutdown)
+        worker.when_finished = WhenFinished.POST_EXEC_HOOK
+        worker.post_exec_hook = None  # Missing hook path
+
+        # Store initial state
+        initial_running_state = worker.running
+
+        # Should log error and return without crashing
+        await worker._handle_work_completion()
+
+        # Should not change running state when hook path is missing
+        assert worker.running == initial_running_state
+        assert worker.should_stop_processing.is_set() is False
+
+    @pytest.mark.asyncio
+    async def test_execute_post_hook_nonexistent_file(self, tmp_path):
+        """Test post-exec hook with nonexistent file."""
+        config = {
+            "name": "test_worker",
+            "token": "test_token",
+            "server": "ws://localhost:8765",
+            "server_url": "ws://localhost:8765",
+            "gpu_id": 0,
+            "when_finished": "post_exec_hook",
+            "post_exec_hook": str(tmp_path / "nonexistent.sh"),
+        }
+
+        worker = create_fast_caption_worker(config)
+
+        # Should handle missing file gracefully
+        await worker._execute_post_hook()
+
+    @pytest.mark.asyncio
+    async def test_execute_post_hook_not_executable(self, tmp_path):
+        """Test post-exec hook with non-executable file."""
+        hook_script = tmp_path / "non_executable.sh"
+        hook_script.write_text("#!/bin/bash\necho 'test'")
+        # Don't make it executable
+
+        config = {
+            "name": "test_worker",
+            "token": "test_token",
+            "server": "ws://localhost:8765",
+            "server_url": "ws://localhost:8765",
+            "gpu_id": 0,
+            "when_finished": "post_exec_hook",
+            "post_exec_hook": str(hook_script),
+        }
+
+        worker = create_fast_caption_worker(config)
+
+        # Should handle non-executable file gracefully
+        await worker._execute_post_hook()
+
+
 if __name__ == "__main__":
     pytest.main([__file__, "-v", "-s"])