caption-flow 0.1.0__py3-none-any.whl → 0.2.0__py3-none-any.whl

caption_flow-0.2.0.dist-info/METADATA

@@ -0,0 +1,369 @@
+Metadata-Version: 2.4
+Name: caption-flow
+Version: 0.2.0
+Summary: Self-contained distributed community captioning system
+Author-email: bghira <bghira@users.github.com>
+License: MIT
+Keywords: captioning,distributed,vllm,dataset,community
+Classifier: Development Status :: 4 - Beta
+Classifier: Intended Audience :: Developers
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Requires-Python: <3.13,>=3.10
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: websockets>=12.0
+Requires-Dist: pyarrow>=14.0.0
+Requires-Dist: click>=8.1.0
+Requires-Dist: pydantic>=2.0.0
+Requires-Dist: aiofiles>=23.0.0
+Requires-Dist: rich>=13.0.0
+Requires-Dist: cryptography>=41.0.0
+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: 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
+Requires-Dist: datasets<5.0.0,>=4.0.0
+Requires-Dist: boto3<2.0.0,>=1.40.11
+Provides-Extra: dev
+Requires-Dist: pytest>=7.4.0; extra == "dev"
+Requires-Dist: pytest-asyncio>=0.21.0; extra == "dev"
+Requires-Dist: pytest-cov>=4.1.0; extra == "dev"
+Requires-Dist: black>=23.0.0; extra == "dev"
+Requires-Dist: ruff>=0.1.0; extra == "dev"
+Requires-Dist: mypy>=1.5.0; extra == "dev"
+Dynamic: license-file
+
+# CaptionFlow
+
+scalable, fault-tolerant **vLLM-powered image captioning**. this "first round" focuses on a fast websocket orchestrator plus lightweight gpu workers that batch requests through vLLM.
+
+* **orchestrator**: hands out work in chunked shards, collects captions, checkpoints progress, and keeps simple stats.
+* **workers (vLLM)**: connect to the orchestrator, stream in image samples, batch them, and generate 1..N captions per image using prompts supplied by the orchestrator.
+* **config-driven**: all components read YAML config; flags can override.
+* **tui monitor (optional)**: a monitor client is wired into the CLI; ship a `monitor` module to enable it.
+
+> no conda. just `venv` + `pip`.
+
+---
+
+## install
+
+```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
+```
+
+## quickstart (single box)
+
+1. copy + edit the sample configs
+
+```bash
+cp orchestrator.yaml my-orchestrator.yaml
+cp worker.yaml my-worker.yaml
+cp monitor.yaml my-monitor.yaml   # optional; requires a monitor module
+```
+
+set a unique shared token in both `my-orchestrator.yaml` and `my-worker.yaml` (see `auth.worker_tokens` in the orchestrator config and `worker.token` in the worker config). if you use private hugging face datasets/models, export `HUGGINGFACE_HUB_TOKEN` before starting workers.
+
+2. start the orchestrator
+
+```bash
+caption-flow orchestrator --config my-orchestrator.yaml
+```
+
+3. start one or more vLLM workers
+
+```bash
+# gpu 0 on the same host
+caption-flow worker --config my-worker.yaml --gpu-id 0
+
+# your second GPU
+caption-flow worker --config my-worker.yaml --gpu-id 1
+```
+
+4. (optional) start the monitor
+
+```bash
+caption-flow monitor --config my-monitor.yaml
+```
+
+5. (optional) scan/fix chunks on disk if you had crashes
+
+```bash
+caption-flow scan_chunks --data-dir ./caption_data --checkpoint-dir ./checkpoints --fix
+```
+
+---
+
+## how it’s wired
+
+### orchestrator
+
+* **websocket server** (default `0.0.0.0:8765`) with three client roles: workers, data-feeders, and admin.
+* **dataset control**: the orchestrator centrally defines the dataset (`huggingface` or `local`) and version/name. it chunk-slices shards and assigns work.
+* **vLLM config broadcast**: model, tp size, dtype, max seq len, memory targets, batching, sampling params, and **inference prompts** are all pushed to workers; workers can apply many changes without a model reload.
+* **storage + checkpoints**: captions buffer to disk with periodic checkpoints. chunk state is tracked so restarts don’t double-work.
+* **auth**: token lists for `worker`, `monitor`, and `admin` roles.
+
+start flags you’ll likely use:
+
+```text
+--config PATH                # yaml config for the orchestrator
+--port INT, --host STR       # bind controls
+--data-dir PATH              # overrides storage.data_dir
+--cert PATH, --key PATH      # enable TLS (or use --no-ssl for ws:// in dev)
+--vllm                       # use the vLLM-style orchestrator (webdataset/hf)
+```
+
+### vLLM worker
+
+* **one process per gpu**. select the device with `--gpu-id` (or `worker.gpu_id` in YAML).
+* **gets its marching orders** from the orchestrator: dataset info, model, prompts, batch size, and sampling.
+* **resilient**: detects disconnects, abandons the current chunk cleanly, clears queues, reconnects, and resumes.
+* **batched generate()**: images are resized down for consistent batching; each image can get multiple captions (one per prompt).
+
+start flags you’ll likely use:
+
+```text
+--config PATH                 # yaml for the worker
+--server URL                  # ws(s)://host:port
+--token STR                   # must match an allowed worker token on the orchestrator
+--name STR                    # display name
+--batch-size INT              # override vLLM batch size
+--vllm                        # use the vLLM worker implementation
+--gpu-id INT                  # which gpu to use
+--precision STR, --model STR  # optional overrides for dtype/model
+--no-verify-ssl               # accept self-signed certs in dev
+```
+
+### (optional) monitor
+
+* a CLI entry exists for a TUI monitor; wire in a `monitor` module to enable it. config lives in `monitor.yaml` or inside `orchestrator.yaml` under `monitor:`.
+
+---
+
+## configuration
+
+### config discovery order
+
+for any component, the CLI looks for config in this order (first match wins):
+
+1. `--config /path/to/file.yaml`
+2. `./<component>.yaml` (current directory)
+3. `~/.caption-flow/<component>.yaml`
+4. `$XDG_CONFIG_HOME/caption-flow/<component>.yaml`
+5. `/etc/caption-flow/<component>.yaml`
+6. any `$XDG_CONFIG_DIRS` entries under `caption-flow/`
+7. `./examples/<component>.yaml` (fallback)
+
+### orchestrator.yaml (highlights)
+
+```yaml
+orchestrator:
+  host: 0.0.0.0
+  port: 8765
+  # ssl:
+  #   cert: /path/fullchain.pem
+  #   key:  /path/privkey.pem
+
+  dataset:
+    type: huggingface   # or "local"
+    path: <hf-dataset-or-local-path>
+    name: <logical-name>
+    version: "1.0"
+
+  vllm:
+    model: Qwen/Qwen2.5-VL-3B-Instruct
+    tensor_parallel_size: 1
+    max_model_len: 16384
+    dtype: float16
+    gpu_memory_utilization: 0.92
+    enforce_eager: true
+    disable_mm_preprocessor_cache: true
+    limit_mm_per_prompt: { image: 1 }
+
+    batch_size: 8
+
+    sampling:
+      temperature: 0.7
+      top_p: 0.95
+      max_tokens: 256
+      repetition_penalty: 1.05
+      skip_special_tokens: true
+      stop: ["<|end|>", "<|endoftext|>", "<|im_end|>"]
+
+    inference_prompts:
+      - "describe this image in detail"
+      - "provide a comprehensive description of the visual content"
+      - "what are the key elements in this image?"
+
+  storage:
+    data_dir: ./caption_data
+    checkpoint_dir: ./checkpoints
+    caption_buffer_size: 100
+    checkpoint_interval: 1000
+
+  # chunking/queueing
+  chunk_size: 1000
+  chunks_per_request: 2
+  chunk_buffer_multiplier: 3
+  min_chunk_buffer: 10
+
+  auth:
+    worker_tokens:
+      - { token: "example-worker-token", name: "Example Worker" }
+    monitor_tokens:
+      - { token: "letmein", name: "Default monitor" }
+    admin_tokens:
+      - { token: "admin-secret-2024", name: "Admin" }
+```
+
+### worker.yaml (highlights)
+
+```yaml
+worker:
+  server: ws://localhost:8765   # use wss:// in prod
+  token: example-worker-token
+  name: local-gpu
+  gpu_id: 0
+  vllm: true
+
+  # local queues
+  readahead_size: 256
+  inference_queue_size: 128
+```
+
+### monitor.yaml (optional)
+
+```yaml
+monitor:
+  server: ws://localhost:8765
+  token: letmein
+  refresh_rate: 1.0
+  show_contributors: true
+  show_quality_metrics: true
+  max_activity_items: 20
+  show_chunk_progress: true
+  show_worker_queues: true
+  show_throughput_graph: true
+```
+
+---
+
+## tls / certificates
+
+use the built-in helpers during development:
+
+```bash
+# self-signed certs for quick local testing
+caption-flow generate_cert --self-signed --domain localhost --output-dir ./certs
+
+# inspect any certificate file
+caption-flow inspect_cert ./certs/fullchain.pem
+```
+
+then point the orchestrator at the resulting cert/key (or run `--no-ssl` for dev-only ws\://).
+
+---
+
+## tips & notes
+
+* **multi-gpu**: start one worker process per gpu (set `--gpu-id` or `worker.gpu_id`).
+* **throughput**: tune `vllm.batch_size` in the orchestrator config (or override with `--batch-size` at worker start). higher isn’t always better; watch VRAM.
+* **prompts**: add more strings under `vllm.inference_prompts` to get multiple captions per image; the worker returns only non-empty generations.
+* **private HF**: if your dataset/model needs auth, export `HUGGINGFACE_HUB_TOKEN` before `caption-flow worker ...`.
+* **self-signed ssl**: pass `--no-verify-ssl` to workers/monitors in dev.
+* **recovery**: if you hard-crash mid-run, `caption-flow scan_chunks --fix` can reset abandoned chunks so the orchestrator can reissue them cleanly.
+
+---
+
+## roadmap
+
+* hot config reload via the admin websocket path.
+* dedicated data-feeder clients (separate from gpu workers) that push samples into the orchestrator.
+* richer monitor TUI.
+
+PRs welcome. keep it simple and fast.
+
+## architecture
+
+```
+┌─────────────┐     WebSocket      ┌─────────────┐
+│   Worker    │◄──────────────────►│             │
+└─────────────┘                    │             │     ┌──────────────┐
+                                   │ Orchestrator│────►│Arrow/Parquet │
+┌─────────────┐                    │             │     │   Storage    │
+│   Worker    │◄──────────────────►│             │     └──────────────┘
+└─────────────┘                    └─────────────┘
+                                           ▲
+┌─────────────┐                           │
+│   Monitor   │◄──────────────────────────┘
+└─────────────┘
+```
+
+## Storage Schema
+
+### captions.parquet
+- `job_id`: Unique job identifier
+- `dataset`: Dataset name
+- `shard`: Shard identifier
+- `item_key`: Item within shard
+- `caption`: Generated caption text
+- `contributor_id`: Worker who generated it
+- `timestamp`: Generation time
+- `quality_score`: Optional quality metric
+
+### jobs.parquet
+- `job_id`: Unique identifier
+- `dataset`: Dataset name
+- `shard`: Shard identifier
+- `status`: pending/processing/completed/failed
+- `assigned_to`: Worker ID
+- `timestamp`: Status change time
+
+### contributors.parquet
+- `contributor_id`: Unique identifier
+- `name`: Display name
+- `total_captions`: Lifetime count
+- `trust_level`: Quality tier (0-5)
+
+## Development
+
+```bash
+# Install with dev dependencies
+pip install -e ".[dev]"
+
+# Run tests
+pytest
+
+# Format code
+black src/
+ruff --fix src/
+
+# Type checking
+mypy src/
+```
+
+## Community Contribution
+
+To contribute compute:
+
+1. Install caption-flow: `pip install caption-flow`
+2. Get a worker token from the project maintainer
+3. Run: `caption-flow worker --server wss://project.domain.com:8765 --token YOUR_TOKEN`
+
+Your contributions will be tracked and attributed in the final dataset!
+
+## License
+
+MIT

caption_flow-0.2.0.dist-info/RECORD

@@ -0,0 +1,29 @@
+caption_flow/__init__.py,sha256=NLPJ25lRN7xHqncXweINDNwbt0q8lgjZ30G21zlPdRs,303
+caption_flow/cli.py,sha256=DVVN4e4uL0jL0gRTaIC5BL0DBU2IU_2yUOi4lg6-lEw,28639
+caption_flow/models.py,sha256=qo6lQiO10UISbaBVr6Cs-fSW_pmjwE6kmiTmmU_l3Wk,2140
+caption_flow/monitor.py,sha256=MltOwBqcFwni1XEPWu5dIO-os5NKDbH_LInOBXUWHAY,7870
+caption_flow/orchestrator.py,sha256=vLW_w5KuRn9Asy_343DxZDRxiUs0xYgbfuuNGgqIf7k,76403
+caption_flow/storage.py,sha256=hC6ZHT_PHFoUVjqD5JUwy3_79oAD1e1H30neA_xsz7s,40748
+caption_flow/utils/__init__.py,sha256=F1BChVoCsj9zn1GJRBOLHET1kLW6xrAmsbzcR7hHy6Y,202
+caption_flow/utils/auth.py,sha256=UrxX2n8OEEcfMD1Ey27TxGfrJFmUCpC59x-SCrQJoVE,2253
+caption_flow/utils/caption_utils.py,sha256=esUMAdcCkNjRroZ0Bhxv0_yKlLtMf0XeDCTt-5k6bik,5309
+caption_flow/utils/certificates.py,sha256=eu4blQZEkL9NRaY1ynQWg1asvDorRYhGRZea7STonJE,4635
+caption_flow/utils/checkpoint_tracker.py,sha256=8tsTFF-HcygitK92YcS-QWzeg-qRm9AuCpQoQRfC8M0,3335
+caption_flow/utils/chunk_tracker.py,sha256=hKn8CN6ubErc9kuCWZMj12ZCZKxVlqXqAEocbzjfa-k,17296
+caption_flow/utils/dataset_loader.py,sha256=qjoRuPnCv_2nGPfrdqf45AgBXlthw1HwqZ1IqwIXzH4,20792
+caption_flow/utils/image_processor.py,sha256=Zl8TAv9gYPdAYat3UiTuuNdIb2fXNfZ35AxsxuovJTs,5650
+caption_flow/utils/job_queue.py,sha256=itdfXcrkvGjmXn4qtpgMF63k1ufRBaejDe4V6WcxzgU,1104
+caption_flow/utils/json_utils.py,sha256=IiZYn8uCM-3pYmyIbX2fmaOIyutArn67SqAyp0ggNpU,5396
+caption_flow/utils/prompt_template.py,sha256=AKp0diSZqNBMwZkpiTNjw8-bbQwHStr7QZTOJ7o1dC4,4345
+caption_flow/utils/shard_processor.py,sha256=CRda6M4xh4U0vwvYlzq9nJEzz4d_4yzUBosYAeBcPEA,10854
+caption_flow/utils/shard_tracker.py,sha256=Wt2oE-O85F2FxSnqIocJiaYeFn00OVVjIiklZIZRGL8,3233
+caption_flow/utils/vllm_config.py,sha256=TC7Rmjk0zRKbBXbWUXrFL4Z58hzax_-4L0pXZn09hdM,6019
+caption_flow/workers/base.py,sha256=jPm_Xw4Lxd0cnrPs-biBqKRQKkTOJLvHLolmp0Gb1CI,7530
+caption_flow/workers/caption.py,sha256=NZ9kTjk2uOoNwyyNSkB_arYk213vLr5mowHN-OjiFkk,54631
+caption_flow/workers/data.py,sha256=0Tg8NE0wdONeMlivYQ4nvbcfWdLuU51O7vR8_YSnJgo,14813
+caption_flow-0.2.0.dist-info/licenses/LICENSE,sha256=hIahDEOTzuHCU5J2nd07LWwkLW7Hko4UFO__ffsvB-8,34523
+caption_flow-0.2.0.dist-info/METADATA,sha256=6qwt05U0S23Omjz1yR6VzLq_wRHbRx_xl3YzhwHyDLc,11900
+caption_flow-0.2.0.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+caption_flow-0.2.0.dist-info/entry_points.txt,sha256=KnVlyrGKZj6p2zNyuEnCx4Y6jvJ4V-mcfN0lddPKTlQ,55
+caption_flow-0.2.0.dist-info/top_level.txt,sha256=_bXpKRutqded0FQ80dCChIz26ETV7tL4d4e2E_Y1FXs,13
+caption_flow-0.2.0.dist-info/RECORD,,

caption_flow/worker.py

@@ -1,300 +0,0 @@
-"""Worker node for distributed captioning."""
-
-import asyncio
-import json
-import logging
-import ssl
-from typing import Dict, Any, Optional
-from pathlib import Path
-
-import websockets
-import websockets.exceptions
-from websockets.client import WebSocketClientProtocol
-
-from .models import Job, JobStatus
-from .utils.image_processor import ImageProcessor
-
-logger = logging.getLogger(__name__)
-
-
-class Worker:
-    """Worker node that processes captioning jobs."""
-
-    def __init__(self, config: Dict[str, Any]):
-        self.config = config
-        self.server_url = config["server"]
-        self.token = config["token"]
-        self.name = config.get("name", "worker")
-        self.batch_size = config.get("batch_size", 32)
-
-        # Dataset configuration will be received from orchestrator
-        self.dataset_config = None
-        self.dataset_type = None
-        self.dataset_path = None
-
-        # SSL configuration
-        self.ssl_context = self._setup_ssl()
-
-        # Components
-        self.image_processor = ImageProcessor()
-
-        # State
-        self.worker_id: Optional[str] = None
-        self.websocket: Optional[WebSocketClientProtocol] = None
-        self.running = False
-        self.current_job: Optional[Job] = None
-
-        # Metrics
-        self.processed_count = 0
-        self.error_count = 0
-
-    def _setup_ssl(self) -> Optional[ssl.SSLContext]:
-        """Configure SSL context."""
-        # Check if URL is WSS (requires SSL)
-        if self.server_url.startswith("ws://"):
-            logger.warning(
-                "Using insecure WebSocket connection (ws://). Consider using wss:// for production."
-            )
-            return None  # No SSL for ws://
-
-        if not self.config.get("verify_ssl", True):
-            # Disable SSL verification for development
-            context = ssl.create_default_context()
-            context.check_hostname = False
-            context.verify_mode = ssl.CERT_NONE
-            return context
-
-        return ssl.create_default_context()
-
-    async def start(self):
-        """Start the worker and connect to orchestrator."""
-        self.running = True
-
-        while self.running:
-            try:
-                await self._connect_and_run()
-            except websockets.exceptions.ConnectionClosed as e:
-                logger.warning(f"Connection closed: {e}")
-                if self.running:
-                    logger.info("Reconnecting in 5 seconds...")
-                    await asyncio.sleep(5)
-            except Exception as e:
-                logger.error(f"Connection error: {e}")
-                if self.running:
-                    logger.info("Reconnecting in 5 seconds...")
-                    await asyncio.sleep(5)
-
-    async def _connect_and_run(self):
-        """Connect to orchestrator and process jobs."""
-        logger.info(f"Connecting to {self.server_url}")
-
-        try:
-            async with websockets.connect(self.server_url, ssl=self.ssl_context) as websocket:
-                self.websocket = websocket
-
-                # Authenticate
-                await websocket.send(json.dumps({"token": self.token, "name": self.name}))
-
-                # Wait for welcome message with dataset configuration
-                welcome = await websocket.recv()
-                welcome_data = json.loads(welcome)
-
-                if "error" in welcome_data:
-                    logger.error(f"Authentication failed: {welcome_data['error']}")
-                    self.running = False
-                    return
-
-                self.worker_id = welcome_data.get("worker_id")
-
-                # Extract and store dataset configuration from orchestrator
-                if "dataset_config" in welcome_data:
-                    self.dataset_config = welcome_data["dataset_config"]
-                    self.dataset_type = self.dataset_config.get("dataset_type")
-                    self.dataset_path = self.dataset_config.get("dataset_path")
-                    logger.info(
-                        f"Received dataset configuration from orchestrator: "
-                        f"type={self.dataset_type}, path={self.dataset_path}"
-                    )
-                else:
-                    logger.warning("No dataset configuration received from orchestrator")
-
-                logger.info(f"Connected as {self.worker_id}")
-
-                # Create tasks for concurrent operations
-                tasks = [
-                    asyncio.create_task(self._heartbeat_loop()),
-                    asyncio.create_task(self._job_processing_loop()),
-                    asyncio.create_task(self._message_handler()),
-                ]
-
-                try:
-                    # Wait for any task to complete (usually due to connection close)
-                    done, pending = await asyncio.wait(tasks, return_when=asyncio.FIRST_COMPLETED)
-
-                    # Cancel remaining tasks
-                    for task in pending:
-                        task.cancel()
-                        try:
-                            await task
-                        except asyncio.CancelledError:
-                            pass
-
-                    # Check if we had an error in completed tasks
-                    for task in done:
-                        try:
-                            task.result()
-                        except websockets.exceptions.ConnectionClosed:
-                            logger.info("WebSocket connection closed")
-                        except Exception as e:
-                            logger.error(f"Task error: {e}")
-
-                except websockets.exceptions.ConnectionClosed:
-                    logger.info("Connection closed by orchestrator")
-
-        except websockets.exceptions.ConnectionClosed as e:
-            logger.info(f"Failed to connect: {e}")
-            raise
-        except Exception as e:
-            logger.error(f"Unexpected error in connection: {e}")
-            raise
-        finally:
-            self.websocket = None
-            self.current_job = None
-
-    async def _job_processing_loop(self):
-        """Main loop for requesting and processing jobs."""
-        while self.running and self.websocket:
-            try:
-                # Request a job
-                await self.websocket.send(json.dumps({"type": "request_job"}))
-
-                # Wait a bit for response
-                await asyncio.sleep(1)
-
-                if self.current_job:
-                    await self._process_job(self.current_job)
-                    self.current_job = None
-                else:
-                    # No job available, wait before requesting again
-                    await asyncio.sleep(5)
-
-            except websockets.exceptions.ConnectionClosed:
-                logger.info("Connection closed during job processing")
-                break
-            except Exception as e:
-                logger.error(f"Job processing error: {e}")
-                self.error_count += 1
-                await asyncio.sleep(1)
-
-    async def _message_handler(self):
-        """Handle incoming messages from orchestrator."""
-        try:
-            async for message in self.websocket:
-                try:
-                    data = json.loads(message)
-                    msg_type = data.get("type")
-
-                    if msg_type == "job":
-                        job_data = data["job"]
-                        self.current_job = Job(**job_data)
-                        logger.info(f"Received job {self.current_job.job_id}")
-
-                    elif msg_type == "no_jobs":
-                        logger.debug("No jobs available")
-
-                    elif msg_type == "ack":
-                        logger.debug(f"Job {data['job_id']} acknowledged")
-                        self.processed_count += 1
-
-                except json.JSONDecodeError as e:
-                    logger.error(f"Invalid message format: {e}")
-                except Exception as e:
-                    logger.error(f"Error handling message: {e}")
-
-        except websockets.exceptions.ConnectionClosed:
-            logger.info("Connection closed while waiting for messages")
-        except Exception as e:
-            logger.error(f"Message handler error: {e}")
-
-    async def _process_job(self, job: Job):
-        """Process a single captioning job."""
-        if not self.websocket:
-            logger.warning(f"No websocket connection, skipping job {job.job_id}")
-            return
-
-        logger.info(f"Processing job {job.job_id}")
-
-        try:
-            # Load and preprocess images
-            images = await self._load_images(job)
-
-            # TODO: Here you would integrate your captioning model
-            # For now, using placeholder
-            caption = f"[Generated caption for {job.item_key}]"
-
-            # Submit result
-            await self.websocket.send(
-                json.dumps(
-                    {
-                        "type": "submit_caption",
-                        "job_id": job.job_id,
-                        "dataset": job.dataset,
-                        "shard": job.shard,
-                        "item_key": job.item_key,
-                        "caption": caption,
-                    }
-                )
-            )
-
-            logger.info(f"Completed job {job.job_id}")
-
-        except websockets.exceptions.ConnectionClosed:
-            logger.warning(f"Connection lost while processing job {job.job_id}")
-            raise  # Re-raise to trigger reconnection
-        except Exception as e:
-            logger.error(f"Failed to process job {job.job_id}: {e}")
-
-            # Report failure if still connected
-            if self.websocket:
-                try:
-                    await self.websocket.send(
-                        json.dumps({"type": "job_failed", "job_id": job.job_id, "error": str(e)})
-                    )
-                except:
-                    pass  # Connection might be closed
-
-    async def _load_images(self, job: Job):
-        """Load and preprocess images for a job."""
-        # This would load actual images from the dataset
-        # Now can use self.dataset_type and self.dataset_path received from orchestrator
-        # For now, returning placeholder
-        return []
-
-    async def _heartbeat_loop(self):
-        """Send periodic heartbeats to orchestrator."""
-        while self.running and self.websocket:
-            try:
-                await self.websocket.send(
-                    json.dumps(
-                        {
-                            "type": "heartbeat",
-                            "processed": self.processed_count,
-                            "errors": self.error_count,
-                        }
-                    )
-                )
-                await asyncio.sleep(30)
-            except websockets.exceptions.ConnectionClosed:
-                logger.info("Connection closed during heartbeat")
-                break
-            except Exception as e:
-                logger.error(f"Heartbeat error: {e}")
-                break
-
-    async def shutdown(self):
-        """Graceful shutdown."""
-        logger.info("Shutting down worker...")
-        self.running = False
-
-        if self.websocket:
-            await self.websocket.close()