modelq 0.1.0__tar.gz

modelq-0.1.0/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2024 ModelsLab
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

modelq-0.1.0/PKG-INFO

@@ -0,0 +1,17 @@
+Metadata-Version: 2.1
+Name: modelq
+Version: 0.1.0
+Summary: Celery-like task queue for ML inference.
+Author: Tanmaypatil123
+Author-email: tanmay@modelslab.com
+Requires-Python: >=3.9,<4.0
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.9
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Requires-Dist: click (>=8.0.0,<9.0.0)
+Requires-Dist: redis (>=4.0.0,<5.0.0)
+Description-Content-Type: text/markdown
+
+# celery-plus-plus

modelq-0.1.0/README.md

@@ -0,0 +1 @@
+# celery-plus-plus

modelq-0.1.0/modelq/__init__.py

@@ -0,0 +1,6 @@
+from .app import ModelQ
+
+
+__all__ = [
+    "ModelQ"
+]

modelq-0.1.0/modelq/app/__init__.py

@@ -0,0 +1,5 @@
+from .base import ModelQ
+
+__all__ = [
+    "ModelQ"
+]

modelq-0.1.0/modelq/app/base.py

@@ -0,0 +1,263 @@
+from typing import Optional, Dict, Any, Generator
+import redis
+import json
+import functools
+import threading
+import time
+import sqlite3
+import uuid
+import logging
+from modelq.app.tasks import Task
+from modelq.exceptions import TaskProcessingError, TaskTimeoutError
+from modelq.app.cache import Cache
+from modelq.app.middleware import Middleware
+
+# Set up logging configuration
+logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
+logger = logging.getLogger(__name__)
+
+class ModelQ:
+    """ModelQ class for managing machine learning tasks with Redis queueing and streaming."""
+
+    def __init__(
+            self,
+            host: str = "localhost",
+            server_id: Optional[str] = None,
+            username: str = None,
+            port: str = 6379,
+            db: int = 0,
+            password: str = None,
+            ssl: bool = False,
+            ssl_cert_reqs: any = None,
+            cache_db_path: str = "cache.db",
+            **kwargs
+    ):
+        self.redis_client = self._connect_to_redis(
+            host=host,
+            port=port,
+            db=db,
+            password=password,
+            username=username,
+            ssl=ssl,
+            ssl_cert_reqs=ssl_cert_reqs,
+            **kwargs
+        )
+        self.server_id = server_id or str(uuid.uuid4())
+        self.allowed_tasks = set()
+        self.cache_db_path = cache_db_path
+        self.cache = Cache(db_path=cache_db_path)
+        self.task_configurations: Dict[str, Dict[str, Any]] = {}
+        self.middleware: Middleware = None
+        self.register_server()
+        self.worker_threads = []
+
+    def _connect_to_redis(
+            self, host: str, port: str, db: int, password: str, ssl: bool, ssl_cert_reqs: any, username: str
+    ) -> redis.Redis:
+        if host == "localhost":
+            connection = redis.Redis(host="localhost", db=3)
+        else:
+            connection = redis.Redis(
+                host=host,
+                port=port,
+                password=password,
+                username=username,
+                ssl=ssl,
+                ssl_cert_reqs=ssl
+            )
+
+        return connection
+
+    def register_server(self):
+        """Registers the server in Redis with its capabilities."""
+        self.redis_client.hset("servers", self.server_id, json.dumps({"allowed_tasks": list(self.allowed_tasks), "status": "idle"}))
+
+    def update_server_status(self, status: str):
+        """Updates the server status in Redis."""
+        server_data = json.loads(self.redis_client.hget("servers", self.server_id))
+        server_data["status"] = status
+        self.redis_client.hset("servers", self.server_id, json.dumps(server_data))
+
+    def enqueue_task(self, task_name: str, payload: dict):
+        task = {
+            **task_name,
+            "status": "queued"
+        }
+
+        self.redis_client.rpush("ml_tasks", json.dumps(task))
+
+    def task(self, task_class=Task, timeout: Optional[int] = None, stream: bool = False, retries: int = 0):
+        """Decorator to create a task. Allows specifying a custom task class, timeout, streaming support, and retries."""
+        def decorator(func):
+            @functools.wraps(func)
+            def wrapper(*args, **kwargs):
+                task_name = func.__name__
+                payload = {
+                    "args": args,
+                    "kwargs": kwargs,
+                    "timeout": timeout,
+                    "stream": stream,
+                    "retries": retries
+                }
+                task = task_class(task_name=task_name, payload=payload)
+                if stream:
+                    task.stream = True
+                self.enqueue_task(task.to_dict(), payload=payload)
+                # Store task in cache
+                task.store_in_cache(self.cache_db_path)
+                return task
+            # Attach the function to the instance so it can be called by process_task
+            setattr(self, func.__name__, func)
+            self.allowed_tasks.add(func.__name__)
+            self.register_server()
+            return wrapper
+        return decorator
+
+    def start_workers(self, no_of_workers: int = 1):
+        if any(thread.is_alive() for thread in self.worker_threads):
+            return  # Workers are already running
+
+        self.check_middleware("before_worker_boot")
+        
+        def worker_loop(worker_id):
+            while True:
+                try:
+                    # Update server status to idle while waiting for tasks
+                    self.update_server_status(f"worker_{worker_id}: idle")
+                    task_data = self.redis_client.blpop("ml_tasks")
+                    if task_data:
+                        # Update server status to busy when a task is picked up
+                        self.update_server_status(f"worker_{worker_id}: busy")
+                        _, task_json = task_data
+                        task_dict = json.loads(task_json)
+                        task = Task.from_dict(task_dict)
+                        if task.task_name in self.allowed_tasks:
+                            try:
+                                logger.info(f"Worker {worker_id} started processing task: {task.task_name}")
+                                start_time = time.time()
+                                self.process_task(task)
+                                end_time = time.time()
+                                logger.info(f"Worker {worker_id} finished task: {task.task_name} in {end_time - start_time:.2f} seconds")
+                            except TaskProcessingError as e:
+                                logger.error(f"Worker {worker_id} encountered a TaskProcessingError while processing task '{task.task_name}': {e}")
+                                if task.payload.get("retries", 0) > 0:
+                                    task.payload["retries"] -= 1
+                                    self.enqueue_task(task.to_dict(), payload=task.payload)
+                            except Exception as e:
+                                logger.error(f"Worker {worker_id} encountered an unexpected error while processing task '{task.task_name}': {e}")
+                                if task.payload.get("retries", 0) > 0:
+                                    task.payload["retries"] -= 1
+                                    self.enqueue_task(task.to_dict(), payload=task.payload)
+                        else:
+                            # Requeue the task if this server cannot process it
+                            self.redis_client.rpush("ml_tasks", task_json)
+                except Exception as e:
+                    logger.error(f"Worker {worker_id} crashed with error: {e}. Restarting worker...")
+        
+        for i in range(no_of_workers):
+            worker_thread = threading.Thread(target=worker_loop, args=(i,))
+            worker_thread.daemon = True
+            worker_thread.start()
+            self.worker_threads.append(worker_thread)
+
+        # Log after all workers have started
+        task_names = ', '.join(self.allowed_tasks) if self.allowed_tasks else 'No tasks registered'
+        logger.info(f"ModelQ worker started successfully with {no_of_workers} worker(s). Connected to Redis at {self.redis_client.connection_pool.connection_kwargs['host']}:{self.redis_client.connection_pool.connection_kwargs['port']}. Registered tasks: {task_names}")
+
+    def check_middleware(self, middleware_event: str):
+        logger.info(f"Middleware event triggered: {middleware_event}")
+        if self.middleware:
+            self.middleware.execute(event=middleware_event)
+
+    def process_task(self, task: Task) -> None:
+        """Processes a given task."""
+        if task.task_name in self.allowed_tasks:
+            task_function = getattr(self, task.task_name, None)
+            if task_function:
+                try:
+                    logger.info(f"Processing task: {task.task_name} with args: {task.payload.get('args', [])} and kwargs: {task.payload.get('kwargs', {})}")
+                    start_time = time.time()
+                    timeout = task.payload.get("timeout", None)
+                    if task.payload.get("stream", False):
+                        for result in task_function(*task.payload.get("args", []), **task.payload.get("kwargs", {})):
+                            task.status = "in_progress"
+                            self.redis_client.xadd(f"task_stream:{task.task_id}", {"result": json.dumps(result)})
+                        # Mark the task as completed when streaming ends
+                        task.status = "completed"
+                        self.redis_client.set(f"task_result:{task.task_id}", json.dumps(task.to_dict()), ex=3600)
+                    else:
+                        if timeout:
+                            result = self._run_with_timeout(task_function, timeout, *task.payload.get("args", []), **task.payload.get("kwargs", {}))
+                        else:
+                            result = task_function(*task.payload.get("args", []), **task.payload.get("kwargs", {}))
+                        result_str = task._convert_to_string(result)
+                        task.result = result_str
+                        task.status = "completed"
+                        self.redis_client.set(f"task_result:{task.task_id}", json.dumps(task.to_dict()), ex=3600)
+                    end_time = time.time()
+                    logger.info(f"Task {task.task_name} completed successfully in {end_time - start_time:.2f} seconds")
+                    # Store updated task status in cache
+                    task.store_in_cache(self.cache_db_path)
+                except Exception as e:
+                    task.status = "failed"
+                    task.result = str(e)
+                    self.redis_client.set(f"task_result:{task.task_id}", json.dumps(task.to_dict()), ex=3600)
+                    # Store failed task status in cache
+                    task.store_in_cache(self.cache_db_path)
+                    logger.error(f"Task {task.task_name} failed with error: {e}")
+                    raise TaskProcessingError(task.task_name, str(e))
+            else:
+                task.status = "failed"
+                task.result = "Task function not found"
+                self.redis_client.set(f"task_result:{task.task_id}", json.dumps(task.to_dict()), ex=3600)
+                # Store failed task status in cache
+                task.store_in_cache(self.cache_db_path)
+                logger.error(f"Task {task.task_name} failed because the task function was not found")
+                raise TaskProcessingError(task.task_name, "Task function not found")
+        else:
+            task.status = "failed"
+            task.result = "Task not allowed"
+            self.redis_client.set(f"task_result:{task.task_id}", json.dumps(task.to_dict()), ex=3600)
+            # Store failed task status in cache
+            task.store_in_cache(self.cache_db_path)
+            logger.error(f"Task {task.task_name} is not allowed")
+            raise TaskProcessingError(task.task_name, "Task not allowed")
+
+    def _run_with_timeout(self, func, timeout, *args, **kwargs):
+        """Runs a function with a timeout."""
+        result = [None]
+        exception = [None]
+
+        def target():
+            try:
+                result[0] = func(*args, **kwargs)
+            except Exception as e:
+                exception[0] = e
+
+        thread = threading.Thread(target=target)
+        thread.start()
+        thread.join(timeout)
+        if thread.is_alive():
+            logger.error(f"Task exceeded timeout of {timeout} seconds")
+            raise TaskTimeoutError(f"Task exceeded timeout of {timeout} seconds")
+        if exception[0]:
+            raise exception[0]
+        return result[0]
+
+    def get_all_queued_tasks(self) -> list:
+        """Retrieves all tasks with status 'queued' from the SQLite cache database."""
+        with sqlite3.connect(self.cache_db_path) as conn:
+            cursor = conn.cursor()
+            cursor.execute('SELECT task_id, task_name, status FROM tasks WHERE status = ?', ("queued",))
+            rows = cursor.fetchall()
+            return [{"task_id": row[0], "task_name": row[1], "status": row[2]} for row in rows]
+
+    def get_task_status(self, task_id: str) -> Optional[str]:
+        """Retrieves the status of a particular task by task ID from the SQLite cache database."""
+        with sqlite3.connect(self.cache_db_path) as conn:
+            cursor = conn.cursor()
+            cursor.execute('SELECT status FROM tasks WHERE task_id = ?', (task_id,))
+            row = cursor.fetchone()
+            if row:
+                return row[0]
+            return None

modelq-0.1.0/modelq/app/cache/__init__.py

@@ -0,0 +1,5 @@
+from .base import Cache
+
+__all__ = [
+    "Cache"
+]

modelq-0.1.0/modelq/app/cache/base.py

@@ -0,0 +1,76 @@
+import sqlite3
+import os
+import json
+from typing import Optional
+from modelq.app.tasks import Task
+
+class Cache:
+
+    def __init__(self, db_path: str = "cache.db") -> None:
+        self.db_path = db_path
+        self._initialize_db()
+    
+    def _initialize_db(self):
+        """Initializes the SQLite database if it doesn't exist."""
+        if not os.path.exists(self.db_path):
+            with sqlite3.connect(self.db_path) as conn:
+                cursor = conn.cursor()
+                cursor.execute(
+                    '''
+                    CREATE TABLE tasks (
+                        task_id TEXT PRIMARY KEY,
+                        task_name TEXT,
+                        payload TEXT,
+                        status TEXT,
+                        result TEXT,
+                        timestamp REAL
+                    )
+                    '''
+                )
+                conn.commit()
+    
+    def _convert_to_string(self, data) -> str:
+        """Converts any data type to a string representation."""
+        try:
+            if isinstance(data, (dict, list, int, float, bool)):
+                return json.dumps(data)
+            return str(data)
+        except TypeError:
+            return str(data)
+
+    def store_task(self, task: Task) -> None:
+        """Stores a new task or updates an existing one in the SQLite database."""
+        with sqlite3.connect(self.db_path) as conn:
+            cursor = conn.cursor()
+            task_id = self._convert_to_string(task.task_id)
+            task_name = self._convert_to_string(task.task_name)
+            payload = self._convert_to_string(task.payload)
+            status = self._convert_to_string(task.status)
+            result = self._convert_to_string(task.result)
+            timestamp = task.timestamp if isinstance(task.timestamp, (int, float)) else None
+
+            cursor.execute(
+                '''
+                INSERT OR REPLACE INTO tasks (task_id, task_name, payload, status, result, timestamp)
+                VALUES (?, ?, ?, ?, ?, ?)
+                ''',
+                (task_id, task_name, payload, status, result, timestamp)
+            )
+            conn.commit()
+
+    def get_task(self, task_id: str) -> Optional[Task]:
+        """Retrieves a task from the SQLite database by its task ID."""
+        with sqlite3.connect(self.db_path) as conn:
+            cursor = conn.cursor()
+            cursor.execute('SELECT * FROM tasks WHERE task_id = ?', (task_id,))
+            row = cursor.fetchone()
+            if row:
+                return Task.from_dict({
+                    "task_id": row[0],
+                    "task_name": row[1],
+                    "payload": json.loads(row[2]),
+                    "status": row[3],
+                    "result": json.loads(row[4]) if row[4] and row[4].startswith(('{', '[')) else row[4],
+                    "timestamp": row[5]
+                })
+            return None

modelq-0.1.0/modelq/app/middleware/__init__.py

@@ -0,0 +1,5 @@
+from modelq.app.middleware.base import Middleware
+
+__all__ = [
+    "Middleware"
+]

modelq-0.1.0/modelq/app/middleware/base.py

@@ -0,0 +1,18 @@
+from abc import ABC , abstractmethod
+
+
+class Middleware(ABC) :
+    def __init__(self) -> None:
+        pass
+    
+    def execute(self,event):
+        if event == "before_worker_boot":
+            self.before_worker_boot()
+
+    @abstractmethod
+    def before_worker_boot(self):
+        "Called before the worker process starts up."
+        pass
+
+    
+

modelq-0.1.0/modelq/app/tasks/__init__.py

@@ -0,0 +1,5 @@
+from .base import Task
+
+__all__ = [
+    "Task"
+]

modelq-0.1.0/modelq/app/tasks/base.py

@@ -0,0 +1,127 @@
+import uuid
+import time
+import json
+import redis
+import sqlite3
+import base64
+from typing import Any, Optional, Generator
+from modelq.exceptions import TaskTimeoutError
+from PIL import Image, PngImagePlugin
+import io
+
+
+class Task:
+
+    def __init__(self, task_name: str, payload: dict, timeout: int = 15):
+        self.task_id = str(uuid.uuid4())
+        self.task_name = task_name
+        self.payload = payload
+        self.status = "queued"
+        self.result = None
+        self.timestamp = time.time()
+        self.timeout = timeout
+        self.stream = False
+        self.combined_result = ""
+
+    def to_dict(self):
+        return {
+            "task_id": self.task_id,
+            "task_name": self.task_name,
+            "payload": self.payload,
+            "status": self.status,
+            "result": self.result,
+            "timestamp": self.timestamp,
+            "stream": self.stream
+        }
+
+    @staticmethod
+    def from_dict(data: dict) -> 'Task':
+        task = Task(task_name=data["task_name"], payload=data["payload"])
+        task.task_id = data["task_id"]
+        task.status = data["status"]
+        task.result = data.get("result")
+        task.timestamp = data["timestamp"]
+        task.stream = data.get("stream", False)
+        return task
+
+    def _convert_to_string(self, data: Any) -> str:
+        """Converts any data type to a string representation, including PIL images."""
+        # print(type(data))
+        try:
+            if isinstance(data, (dict, list, int, float, bool)):
+                return json.dumps(data)
+            elif isinstance(data, (Image.Image, PngImagePlugin.PngImageFile)):
+                print("here")
+                buffered = io.BytesIO()
+                data.save(buffered, format="PNG")
+                return "data:image/png;base64," + base64.b64encode(buffered.getvalue()).decode("utf-8")
+            return str(data)
+        except TypeError:
+            return str(data)
+
+    def store_in_cache(self, db_path: str = "cache.db") -> None:
+        """Stores the task in the SQLite cache database."""
+        with sqlite3.connect(db_path) as conn:
+            cursor = conn.cursor()
+            task_id = self._convert_to_string(self.task_id)
+            task_name = self._convert_to_string(self.task_name)
+            payload = self._convert_to_string(self.payload)
+            status = self._convert_to_string(self.status)
+            result = self._convert_to_string(self.result)
+            timestamp = self.timestamp if isinstance(self.timestamp, (int, float)) else None
+
+            cursor.execute(
+                '''
+                INSERT OR REPLACE INTO tasks (task_id, task_name, payload, status, result, timestamp)
+                VALUES (?, ?, ?, ?, ?, ?)
+                ''',
+                (task_id, task_name, payload, status, result, timestamp)
+            )
+            conn.commit()
+
+    def get_result(self, redis_client: redis.Redis, timeout: int = None) -> Any:
+        """Waits for the result of the task until the timeout."""
+
+        if not timeout:
+            timeout = self.timeout
+
+        start_time = time.time()
+        while time.time() - start_time < timeout:
+            task_json = redis_client.get(f"task_result:{self.task_id}")
+            if task_json:
+                task_data = json.loads(task_json)
+                self.result = task_data.get("result")
+                self.status = task_data.get("status")
+                # Store the updated task status in cache
+                self.store_in_cache()
+                return self.result
+            time.sleep(1)
+        raise TaskTimeoutError(self.task_id)
+
+    def get_stream(self, redis_client: redis.Redis) -> Generator[Any, None, None]:
+        """Generator to yield results from a streaming task."""
+        stream_key = f"task_stream:{self.task_id}"
+        last_id = "0"
+        completed = False
+
+        while not completed:
+            results = redis_client.xread({stream_key: last_id}, block=1000, count=10)
+            if results:
+                for _, messages in results:
+                    for message_id, message_data in messages:
+                        result = json.loads(message_data[b"result"].decode("utf-8"))
+                        yield result
+                        last_id = message_id
+                        # Concatenate result for storing combined response
+                        self.combined_result += result
+            # Check if the task is marked as completed after yielding current messages
+            task_json = redis_client.get(f"task_result:{self.task_id}")
+            if task_json:
+                task_data = json.loads(task_json)
+                if task_data.get("status") == "completed":
+                    completed = True
+                    # Store the completed task status and combined result in cache
+                    self.status = "completed"
+                    self.result = self.combined_result
+                    self.store_in_cache()
+        return

modelq-0.1.0/modelq/app/utils/__init__.py

@@ -0,0 +1,5 @@
+from .base64 import base64_to_image
+
+__all__ = [
+    "base64_to_image"
+]

modelq-0.1.0/modelq/app/utils/base64.py

@@ -0,0 +1,14 @@
+from PIL import Image
+import io
+import base64
+
+def base64_to_image(base64_string):
+    split = base64_string.split(",")
+    if len(split) > 1 :
+        base64_data = split[1]
+    else :
+
+        base64_data = base64_string
+    img = Image.open(io.BytesIO(base64.b64decode(base64_data)))
+
+    return img

modelq-0.1.0/modelq/exceptions.py

@@ -0,0 +1,13 @@
+class TaskTimeoutError(Exception):
+    """Custom exception to indicate task timeout."""
+    def __init__(self, task_id : str) -> None:
+        super().__init__(f"Task {task_id} timed out waiting for result.")
+        self.task_id = task_id
+
+
+class TaskProcessingError(Exception):
+    """Custom exception to indicate an error occurred during task processing."""
+    def __init__(self, task_name: str, message: str):
+        super().__init__(f"Error processing task {task_name}: {message}")
+        self.task_name = task_name
+        self.message = message

modelq-0.1.0/pyproject.toml

@@ -0,0 +1,18 @@
+[tool.poetry]
+name = "modelq"
+version = "0.1.0"
+description = "Celery-like task queue for ML inference."
+authors = ["Tanmaypatil123 <tanmay@modelslab.com>"]
+readme = "README.md"
+
+[tool.poetry.scripts]
+modelq = "modelq.app:run_worker"
+
+[tool.poetry.dependencies]
+python = "^3.9"
+click = "^8.0.0"
+redis = "^4.0.0"
+
+[build-system]
+requires = ["poetry-core"]
+build-backend = "poetry.core.masonry.api"