MemoryOS 2.0.3__py3-none-any.whl

memos/mem_scheduler/task_schedule_modules/local_queue.py

@@ -0,0 +1,247 @@
+"""
+Local Queue implementation for SchedulerMessageItem objects.
+This module provides a local-based queue implementation that can replace
+the local memos_message_queue functionality in BaseScheduler.
+"""
+
+from typing import TYPE_CHECKING
+
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+from memos.log import get_logger
+from memos.mem_scheduler.general_modules.misc import AutoDroppingQueue as Queue
+from memos.mem_scheduler.schemas.message_schemas import ScheduleMessageItem
+from memos.mem_scheduler.schemas.task_schemas import DEFAULT_STREAM_KEY_PREFIX
+from memos.mem_scheduler.task_schedule_modules.orchestrator import SchedulerOrchestrator
+from memos.mem_scheduler.utils.status_tracker import TaskStatusTracker
+from memos.mem_scheduler.webservice_modules.redis_service import RedisSchedulerModule
+
+
+logger = get_logger(__name__)
+
+
+class SchedulerLocalQueue(RedisSchedulerModule):
+    def __init__(
+        self,
+        maxsize: int = 0,
+        stream_key_prefix: str = DEFAULT_STREAM_KEY_PREFIX,
+        orchestrator: SchedulerOrchestrator | None = None,
+        status_tracker: TaskStatusTracker | None = None,
+    ):
+        """
+        Initialize the SchedulerLocalQueue with a maximum queue size limit.
+        Arguments match SchedulerRedisQueue for compatibility.
+
+        Args:
+            maxsize (int): Maximum number of messages allowed in each individual queue.
+            stream_key_prefix (str): Prefix for stream keys (simulated).
+            orchestrator: SchedulerOrchestrator instance (ignored).
+            status_tracker: TaskStatusTracker instance (ignored).
+        """
+        super().__init__()
+
+        self.stream_key_prefix = stream_key_prefix or "local_queue"
+
+        self.max_internal_message_queue_size = maxsize
+
+        # Dictionary to hold per-stream queues: key = stream_key, value = Queue[ScheduleMessageItem]
+        self.queue_streams: dict[str, Queue[ScheduleMessageItem]] = {}
+
+        self.orchestrator = orchestrator
+        self.status_tracker = status_tracker
+
+        self._is_listening = False
+        self._message_handler: Callable[[ScheduleMessageItem], None] | None = None
+
+        logger.info(
+            f"SchedulerLocalQueue initialized with max_internal_message_queue_size={self.max_internal_message_queue_size}"
+        )
+
+    def get_stream_key(self, user_id: str, mem_cube_id: str, task_label: str) -> str:
+        stream_key = f"{self.stream_key_prefix}:{user_id}:{mem_cube_id}:{task_label}"
+        return stream_key
+
+    def put(
+        self, message: ScheduleMessageItem, block: bool = True, timeout: float | None = None
+    ) -> None:
+        """
+        Put a message into the appropriate internal queue based on user_id and mem_cube_id.
+
+        If the corresponding queue does not exist, it is created automatically.
+        This method uses a local in-memory queue (not Redis) for buffering messages.
+
+        Args:
+            message (ScheduleMessageItem): The message to enqueue.
+            block (bool): If True, block if the queue is full; if False, raise Full immediately.
+            timeout (float | None): Maximum time to wait for the queue to become available.
+                                   If None, block indefinitely. Ignored if block=False.
+
+        Raises:
+            queue.Full: If the queue is full and block=False or timeout expires.
+            Exception: Any underlying error during queue.put() operation.
+        """
+        stream_key = self.get_stream_key(
+            user_id=message.user_id, mem_cube_id=message.mem_cube_id, task_label=message.label
+        )
+
+        message.stream_key = stream_key
+
+        # Create the queue if it doesn't exist yet
+        if stream_key not in self.queue_streams:
+            logger.info(f"Creating new internal queue for stream: {stream_key}")
+            self.queue_streams[stream_key] = Queue(maxsize=self.max_internal_message_queue_size)
+
+        try:
+            self.queue_streams[stream_key].put(item=message, block=block, timeout=timeout)
+            logger.info(
+                f"Message successfully put into queue '{stream_key}'. Current size: {self.queue_streams[stream_key].qsize()}"
+            )
+        except Exception as e:
+            logger.error(f"Failed to put message into queue '{stream_key}': {e}", exc_info=True)
+            raise  # Re-raise to maintain caller expectations
+
+    def get(
+        self,
+        stream_key: str,
+        block: bool = True,
+        timeout: float | None = None,
+        batch_size: int | None = 1,
+    ) -> list[ScheduleMessageItem]:
+        if batch_size is not None and batch_size <= 0:
+            logger.warning(
+                f"get() called with invalid batch_size: {batch_size}. Returning empty list."
+            )
+            return []
+
+        # Return empty list if queue does not exist
+        if stream_key not in self.queue_streams:
+            logger.error(f"Stream {stream_key} does not exist when trying to get messages.")
+            return []
+
+        # Ensure we always request a batch so we get a list back
+        effective_batch_size = batch_size if batch_size is not None else 1
+
+        # Note: Assumes custom Queue implementation supports batch_size parameter
+        res = self.queue_streams[stream_key].get(
+            block=block, timeout=timeout, batch_size=effective_batch_size
+        )
+        logger.debug(
+            f"Retrieved {len(res)} messages from queue '{stream_key}'. Current size: {self.queue_streams[stream_key].qsize()}"
+        )
+        return res
+
+    def get_nowait(self, stream_key: str, batch_size: int | None = 1) -> list[ScheduleMessageItem]:
+        """
+        Non-blocking version of get(). Equivalent to get(stream_key, block=False, batch_size=batch_size).
+
+        Returns immediately with available messages or an empty list if queue is empty.
+
+        Args:
+            stream_key (str): The stream/queue identifier.
+            batch_size (int | None): Number of messages to retrieve in a batch.
+                                   If None, retrieves one message.
+
+        Returns:
+            List[ScheduleMessageItem]: Retrieved messages or empty list if queue is empty.
+        """
+        logger.debug(f"get_nowait() called for {stream_key} with batch_size: {batch_size}")
+        return self.get(stream_key=stream_key, block=False, batch_size=batch_size)
+
+    def get_messages(self, batch_size: int) -> list[ScheduleMessageItem]:
+        """
+        Get messages from all streams in round-robin or sequential fashion.
+        Equivalent to SchedulerRedisQueue.get_messages.
+        """
+        messages = []
+        # Snapshot keys to avoid runtime modification issues
+        stream_keys = list(self.queue_streams.keys())
+
+        # Simple strategy: try to get up to batch_size messages across all streams
+        # We can just iterate and collect.
+
+        # Calculate how many to get per stream to be fair?
+        # Or just greedy? Redis implementation uses a complex logic.
+        # For local, let's keep it simple: just iterate and take what's available (non-blocking)
+
+        for stream_key in stream_keys:
+            if len(messages) >= batch_size:
+                break
+
+            needed = batch_size - len(messages)
+            # Use get_nowait to avoid blocking
+            fetched = self.get_nowait(stream_key=stream_key, batch_size=needed)
+            messages.extend(fetched)
+
+        return messages
+
+    def qsize(self) -> dict:
+        """
+        Return the current size of all internal queues as a dictionary.
+
+        Each key is the stream name, and each value is the number of messages in that queue.
+        Also includes 'total_size'.
+
+        Returns:
+            Dict[str, int]: Mapping from stream name to current queue size.
+        """
+        sizes = {stream: queue.qsize() for stream, queue in self.queue_streams.items()}
+        total_size = sum(sizes.values())
+        sizes["total_size"] = total_size
+        logger.debug(f"Current queue sizes: {sizes}")
+        return sizes
+
+    def clear(self, stream_key: str | None = None) -> None:
+        if stream_key:
+            if stream_key in self.queue_streams:
+                self.queue_streams[stream_key].clear()
+        else:
+            for queue in self.queue_streams.values():
+                queue.clear()
+
+    @property
+    def unfinished_tasks(self) -> int:
+        """
+        Calculate the total number of unprocessed messages across all queues.
+
+        This is a convenience property for monitoring overall system load.
+
+        Returns:
+            int: Sum of all message counts in all internal queues.
+        """
+        # qsize() now includes "total_size", so we need to be careful not to double count if we use qsize() values
+        # But qsize() implementation above sums values from queue_streams, then adds total_size.
+        # So sum(self.queue_streams.values().qsize()) is safer.
+        total = sum(queue.qsize() for queue in self.queue_streams.values())
+        logger.debug(f"Total unfinished tasks across all queues: {total}")
+        return total
+
+    def get_stream_keys(self, stream_key_prefix: str | None = None) -> list[str]:
+        """
+        Return list of active stream keys.
+        """
+        prefix = stream_key_prefix or self.stream_key_prefix
+        return [k for k in self.queue_streams if k.startswith(prefix)]
+
+    def size(self) -> int:
+        """
+        Total size of all queues.
+        """
+        return sum(q.qsize() for q in self.queue_streams.values())
+
+    def empty(self) -> bool:
+        """
+        Check if all queues are empty.
+        """
+        return self.size() == 0
+
+    def full(self) -> bool:
+        """
+        Check if any queue is full (approximate).
+        """
+        if self.max_internal_message_queue_size <= 0:
+            return False
+        return any(
+            q.qsize() >= self.max_internal_message_queue_size for q in self.queue_streams.values()
+        )

memos/mem_scheduler/task_schedule_modules/orchestrator.py

@@ -0,0 +1,74 @@
+"""
+Scheduler Orchestrator for Redis-backed task queues.
+
+This module provides an orchestrator class that works with `SchedulerRedisQueue` to:
+- Broker tasks from Redis streams according to per-user priority weights.
+- Maintain a cache of fetched messages and assemble balanced batches across
+  `(user_id, mem_cube_id, task_label)` groups.
+
+Stream format:
+- Keys follow: `{prefix}:{user_id}:{mem_cube_id}:{task_label}`
+
+Default behavior:
+- All users have priority 1, so fetch sizes are equal per user.
+"""
+
+from __future__ import annotations
+
+from memos.log import get_logger
+from memos.mem_scheduler.schemas.task_schemas import (
+    ADD_TASK_LABEL,
+    ANSWER_TASK_LABEL,
+    DEFAULT_PENDING_CLAIM_MIN_IDLE_MS,
+    PREF_ADD_TASK_LABEL,
+    QUERY_TASK_LABEL,
+    TaskPriorityLevel,
+)
+from memos.mem_scheduler.webservice_modules.redis_service import RedisSchedulerModule
+
+
+logger = get_logger(__name__)
+
+
+class SchedulerOrchestrator(RedisSchedulerModule):
+    def __init__(self):
+        """
+        Args:
+            queue: An instance of `SchedulerRedisQueue`.
+        """
+        # Cache of fetched messages grouped by (user_id, mem_cube_id, task_label)
+        self._cache = None
+        self.tasks_priorities = {
+            ADD_TASK_LABEL: TaskPriorityLevel.LEVEL_1,
+            QUERY_TASK_LABEL: TaskPriorityLevel.LEVEL_1,
+            ANSWER_TASK_LABEL: TaskPriorityLevel.LEVEL_1,
+        }
+
+        # Per-task minimum idle time (ms) before claiming pending messages
+        # Default fallback handled in `get_task_idle_min`.
+        self.tasks_min_idle_ms = {
+            # Preferential add tasks: allow claiming pending sooner (10 minute)
+            PREF_ADD_TASK_LABEL: 600_000,
+        }
+
+    def get_stream_priorities(self) -> None | dict:
+        return None
+
+    def get_task_priority(self, task_label: str):
+        return self.tasks_priorities.get(task_label, TaskPriorityLevel.LEVEL_3)
+
+    def get_task_idle_min(self, task_label: str) -> int:
+        idle_min = self.tasks_min_idle_ms.get(task_label, DEFAULT_PENDING_CLAIM_MIN_IDLE_MS)
+        return idle_min
+
+    def get_stream_quotas(self, stream_keys, consume_batch_size) -> dict:
+        stream_priorities = self.get_stream_priorities()
+        stream_quotas = {}
+        for stream_key in stream_keys:
+            if stream_priorities is None:
+                # Distribute per-stream evenly
+                stream_quotas[stream_key] = consume_batch_size
+            else:
+                # TODO: not implemented yet
+                stream_quotas[stream_key] = consume_batch_size
+        return stream_quotas