broccoli-workers 0.1.2__tar.gz → 0.1.3__tar.gz

{broccoli_workers-0.1.2 → broccoli_workers-0.1.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: broccoli-workers
-Version: 0.1.2
+Version: 0.1.3
 Summary: Distributed task queue with priorities, dependencies, and chains, backed by Redis
 Author-email: Success Oguntuyi <oguntuyisuccess@gmail.com>
 License: MIT License
@@ -37,6 +37,7 @@ Requires-Python: >=3.11
 Description-Content-Type: text/markdown
 License-File: LICENSE
 Requires-Dist: redis>=4.0.0
+Requires-Dist: torch>=3.0.0
 Provides-Extra: dev
 Requires-Dist: pytest>=7.0; extra == "dev"
 Requires-Dist: pytest-cov; extra == "dev"

{broccoli_workers-0.1.2 → broccoli_workers-0.1.3}/broccoli/core/chain/chain.py

@@ -43,15 +43,26 @@ class Chain:
 
     @classmethod
     def from_dict(cls, data: dict) -> "Chain":
-        """Create a Chain object from a dictionary."""
+        """Create a Chain object from a dictionary, safely handling result JSON."""
         data = cls._normalize_redis_mapping(data)
+
+        # Safely deserialize 'result' field
+        raw_result = data.get("result")
+        if raw_result and isinstance(raw_result, str) and raw_result.strip():
+            try:
+                result = json.loads(raw_result)
+            except json.JSONDecodeError:
+                result = None
+        else:
+            result = None
+
         return cls(
             chain_id=data.get("chain_id"),
             completion_task=data.get("completion_task") or None,
             status=data.get("status", "pending"),
             completed_tasks=int(data.get("completed_tasks", 0)),
             current_task=int(data.get("current_task", 0)),
-            result=json.loads(data.get("result")) if data.get("result") else None,
+            result=result,
             failed=data.get("failed", "False") == "True",
             total_tasks=int(data.get("total_tasks", 0))
             if data.get("total_tasks")

{broccoli_workers-0.1.2 → broccoli_workers-0.1.3}/broccoli/core/chain/task_chain.py

@@ -5,6 +5,7 @@ import uuid
 from typing import Any, Dict, List, Optional, Union
 
 from broccoli.core.chain.chain import Chain
+from broccoli.core.chain.chain_queue import ChainQueue
 from broccoli.core.redis_controller import RedisController
 from broccoli.core.task.task import Task
 from broccoli.core.task.task_queue import TaskQueue
@@ -24,7 +25,7 @@ class TaskChain:
         redis_config: dict = None,
     ):
         redis_config = redis_config or {}
-        self.queue = TaskQueue(
+        self.queue = ChainQueue(
             redis_url=redis_url,
             queue_name="chain_tasks:queue",
             task_prefix="chain",
@@ -167,3 +168,14 @@ class TaskChain:
                 return {"status": "Not Found"}
 
         return data
+
+    def get_all_chain_tasks(self, chain_id: str) -> List[Dict[str, Any]]:
+        """Retrieve all tasks associated with a chain."""
+        tasks_json = self._redis.get(f"chain:{chain_id}:tasks")
+        if not tasks_json:
+            return []
+        try:
+            return json.loads(tasks_json)
+        except json.JSONDecodeError:
+            logger.error(f"Failed to decode tasks for chain {chain_id}")
+            return []

{broccoli_workers-0.1.2 → broccoli_workers-0.1.3}/broccoli/core/task/task.py

@@ -86,3 +86,15 @@ class Task:
             created_at=data.get("created_at") or datetime.now().isoformat(),
             updated_at=data.get("updated_at") or None,
         )
+
+    def __post_init__(self):
+        """Inject the task_id into the payload automatically."""
+        # Ensure payload is a dict (it is by default)
+        if self.payload is None:
+            self.payload = {}
+        # Always set the task_id in the payload to match this task's ID
+        if "task_id" not in self.payload:
+            self.payload["task_id"] = self.task_id
+        # (Optional) Ensure depends_on is a list
+        if self.depends_on is None:
+            self.depends_on = []

{broccoli_workers-0.1.2 → broccoli_workers-0.1.3}/broccoli/core/task/task_queue.py

@@ -222,6 +222,8 @@ class TaskQueue:
         """
         # Always persist task metadata first so workers can read it after pop().
         self._redis.hset(f"{self.task_prefix}:{task.task_id}", mapping=task.to_dict())
+        if len(task.depends_on) == 0:
+            task.depends_on = None
 
         if task.depends_on:
             # Build the key list for the Lua script:

{broccoli_workers-0.1.2 → broccoli_workers-0.1.3}/broccoli/workers/base_worker.py

@@ -225,9 +225,7 @@ class BaseWorker(ABC):
                     exc_info=True,
                 )
         if task.payload.get("__chain_id"):
-            logger.info(
-                f"Task {task.task_id} is part of chain {task.chain_id}; skipping result storage"
-            )
+            self.result.store_task(task)
             if task.payload.get("__is_last_task"):
                 logger.info(
                     f"Task {task.task_id} is the last task in chain {task.chain_id}; running chain completion handlers"

{broccoli_workers-0.1.2 → broccoli_workers-0.1.3}/broccoli_workers.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: broccoli-workers
-Version: 0.1.2
+Version: 0.1.3
 Summary: Distributed task queue with priorities, dependencies, and chains, backed by Redis
 Author-email: Success Oguntuyi <oguntuyisuccess@gmail.com>
 License: MIT License
@@ -37,6 +37,7 @@ Requires-Python: >=3.11
 Description-Content-Type: text/markdown
 License-File: LICENSE
 Requires-Dist: redis>=4.0.0
+Requires-Dist: torch>=3.0.0
 Provides-Extra: dev
 Requires-Dist: pytest>=7.0; extra == "dev"
 Requires-Dist: pytest-cov; extra == "dev"

{broccoli_workers-0.1.2 → broccoli_workers-0.1.3}/broccoli_workers.egg-info/requires.txt

@@ -1,4 +1,5 @@
 redis>=4.0.0
+torch>=3.0.0
 
 [dev]
 pytest>=7.0

{broccoli_workers-0.1.2 → broccoli_workers-0.1.3}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "broccoli-workers"
-version = "0.1.2"
+version = "0.1.3"
 description = "Distributed task queue with priorities, dependencies, and chains, backed by Redis"
 readme = "README.md"
 authors = [
@@ -24,6 +24,7 @@ classifiers = [
 requires-python = ">=3.11"   # ← changed from >=3.9
 dependencies = [
     "redis>=4.0.0",
+    "torch>=3.0.0",
 ]
 
 [project.optional-dependencies]