rebase

slekkala1 · slekkala1 · commit 33f3a1077be4 · 2025-10-29T11:58:31.000-07:00
diff --git a/src/llama_stack/providers/utils/job_scheduler/__init__.py b/src/llama_stack/providers/utils/job_scheduler/__init__.py
@@ -0,0 +1,47 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+
+from .api import JobStatus, Scheduler
+from .config import CelerySchedulerConfig, InlineSchedulerConfig, SchedulerConfig
+
+
+async def scheduler_impl(config: SchedulerConfig) -> Scheduler:
+    """
+    Factory function to instantiate scheduler implementations.
+
+    Args:
+        config: Scheduler configuration (InlineSchedulerConfig or CelerySchedulerConfig)
+
+    Returns:
+        Scheduler: An initialized scheduler instance
+
+    Raises:
+        ValueError: If the config type is unknown
+    """
+    impl: Scheduler
+    if isinstance(config, InlineSchedulerConfig):
+        from .inline import InlineSchedulerImpl
+
+        impl = InlineSchedulerImpl(config)
+    elif isinstance(config, CelerySchedulerConfig):
+        from .celery import CelerySchedulerImpl
+
+        impl = CelerySchedulerImpl(config)
+    else:
+        raise ValueError(f"Unknown scheduler config type: {type(config)}")
+
+    await impl.initialize()
+    return impl
+
+
+__all__ = [
+    "JobStatus",
+    "Scheduler",
+    "SchedulerConfig",
+    "InlineSchedulerConfig",
+    "CelerySchedulerConfig",
+    "scheduler_impl",
+]
diff --git a/src/llama_stack/providers/utils/job_scheduler/api.py b/src/llama_stack/providers/utils/job_scheduler/api.py
@@ -0,0 +1,135 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+
+from collections.abc import Awaitable, Callable
+from enum import StrEnum
+from typing import Protocol
+
+
+class JobStatus(StrEnum):
+    PENDING = "pending"
+    RUNNING = "running"
+    COMPLETED = "completed"
+    FAILED = "failed"
+    CANCELLED = "cancelled"
+
+
+class Scheduler(Protocol):
+    """
+    Abstract scheduler protocol for managing async jobs.
+
+    This provides a pluggable backend for job scheduling and execution.
+    """
+
+    async def schedule_job(
+        self,
+        job_type: str,
+        job_data: dict,
+        metadata: dict | None = None,
+    ) -> str:
+        """
+        Schedule a new job for execution.
+
+        Args:
+            job_type: Type of job (e.g., "batch_processing", "log_aggregation", "metrics_collection")
+            job_data: Job-specific data and parameters
+            metadata: Additional metadata for tracking
+
+        Returns:
+            job_id: UUID for tracking the job
+        """
+        ...
+
+    async def get_job_info(self, job_id: str) -> dict:
+        """
+        Get complete information about a job.
+
+        Returns:
+            {
+                "job_id": str,
+                "job_type": str,
+                "status": JobStatus,
+                "created_at": datetime,
+                "started_at": datetime | None,
+                "completed_at": datetime | None,
+                "progress": float,  # 0.0 to 1.0
+                "metadata": dict,
+                "error": str | None,  # Error message if status == FAILED
+                "result": dict | None,  # Job result if status == COMPLETED
+            }
+        """
+        ...
+
+    async def cancel_job(self, job_id: str) -> bool:
+        """
+        Cancel a pending or running job.
+
+        Args:
+            job_id: Job to cancel
+
+        Returns:
+            True if job was cancelled, False if not found or already completed
+        """
+        ...
+
+    async def delete_job(self, job_id: str) -> bool:
+        """
+        Delete a completed or cancelled job.
+
+        Args:
+            job_id: Job to delete
+
+        Returns:
+            True if job was deleted, False if not found
+        """
+        ...
+
+    async def list_jobs(
+        self,
+        job_type: str | None = None,
+        status: JobStatus | None = None,
+        limit: int = 100,
+        offset: int = 0,
+    ) -> list[dict]:
+        """
+        List jobs with optional filtering.
+
+        Args:
+            job_type: Filter by job type (e.g., "batch_processing", "log_aggregation", "metrics_collection")
+            status: Filter by status
+            limit: Maximum number of jobs to return
+            offset: Offset for pagination
+
+        Returns:
+            List of job info dictionaries
+        """
+        ...
+
+    def register_job_executor(
+        self,
+        job_type: str,
+        executor: Callable[[dict], Awaitable[dict]],
+    ) -> None:
+        """
+        Register a job executor function for a specific job type.
+
+        This allows components to register custom job execution logic with the scheduler.
+        When a job of the registered type is executed, the scheduler will call the
+        registered executor function.
+
+        Args:
+            job_type: The type of job (e.g., "vector_store_file_batch")
+            executor: Async function that takes job_data dict and returns result dict
+        """
+        ...
+
+    async def initialize(self) -> None:
+        """Initialize the scheduler (connect to backend, etc.)"""
+        ...
+
+    async def shutdown(self) -> None:
+        """Gracefully shutdown the scheduler"""
+        ...
diff --git a/src/llama_stack/providers/utils/job_scheduler/celery/__init__.py b/src/llama_stack/providers/utils/job_scheduler/celery/__init__.py
@@ -0,0 +1,9 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+
+from .scheduler import CelerySchedulerImpl
+
+__all__ = ["CelerySchedulerImpl"]
diff --git a/src/llama_stack/providers/utils/job_scheduler/celery/scheduler.py b/src/llama_stack/providers/utils/job_scheduler/celery/scheduler.py
@@ -0,0 +1,81 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+
+from collections.abc import Awaitable, Callable
+
+from ..api import JobStatus, Scheduler
+from ..config import CelerySchedulerConfig
+
+
+class CelerySchedulerImpl(Scheduler):
+    """
+    Celery-based scheduler implementation for distributed multi-worker deployments.
+
+    This scheduler uses Celery to distribute jobs across multiple worker processes
+    or machines. It provides:
+    - Persistent job queue (via Redis/RabbitMQ broker)
+    - Multi-worker coordination
+    - Crash recovery (jobs survive server restarts)
+    - Distributed task execution
+
+    This is suitable for:
+    - Production deployments
+    - Multi-worker scenarios
+    - High availability requirements
+    """
+
+    def __init__(self, config: CelerySchedulerConfig):
+        self.config = config
+        self._job_executors: dict[str, Callable[[dict], Awaitable[dict]]] = {}
+        # TODO: Initialize Celery app with broker and result backend
+        raise NotImplementedError("Celery scheduler not yet implemented")
+
+    def register_job_executor(
+        self,
+        job_type: str,
+        executor: Callable[[dict], Awaitable[dict]],
+    ) -> None:
+        """Register a job executor function for a specific job type."""
+        self._job_executors[job_type] = executor
+
+    async def initialize(self) -> None:
+        """Initialize the Celery scheduler."""
+        raise NotImplementedError("Celery scheduler not yet implemented")
+
+    async def shutdown(self) -> None:
+        """Gracefully shutdown the Celery scheduler."""
+        raise NotImplementedError("Celery scheduler not yet implemented")
+
+    async def schedule_job(
+        self,
+        job_type: str,
+        job_data: dict,
+        metadata: dict | None = None,
+    ) -> str:
+        """Schedule a new job for execution via Celery."""
+        raise NotImplementedError("Celery scheduler not yet implemented")
+
+    async def get_job_info(self, job_id: str) -> dict:
+        """Get complete information about a job from Celery result backend."""
+        raise NotImplementedError("Celery scheduler not yet implemented")
+
+    async def cancel_job(self, job_id: str) -> bool:
+        """Cancel a pending or running Celery job."""
+        raise NotImplementedError("Celery scheduler not yet implemented")
+
+    async def delete_job(self, job_id: str) -> bool:
+        """Delete a completed or cancelled job from Celery result backend."""
+        raise NotImplementedError("Celery scheduler not yet implemented")
+
+    async def list_jobs(
+        self,
+        job_type: str | None = None,
+        status: JobStatus | None = None,
+        limit: int = 100,
+        offset: int = 0,
+    ) -> list[dict]:
+        """List jobs from Celery result backend with optional filtering."""
+        raise NotImplementedError("Celery scheduler not yet implemented")
diff --git a/src/llama_stack/providers/utils/job_scheduler/config.py b/src/llama_stack/providers/utils/job_scheduler/config.py
@@ -0,0 +1,58 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+
+from typing import Annotated, Literal
+
+from pydantic import BaseModel, Field
+
+from llama_stack.core.storage.datatypes import KVStoreReference
+
+
+class SchedulerConfig(BaseModel):
+    """Base class for scheduler configurations."""
+
+    type: str
+
+
+class InlineSchedulerConfig(SchedulerConfig):
+    """
+    Configuration for inline (asyncio-based) scheduler.
+
+    This scheduler runs jobs in the same process using asyncio tasks.
+    Suitable for development and single-worker deployments.
+    """
+
+    type: Literal["inline"] = "inline"
+    kvstore: KVStoreReference = Field(
+        description="KVStore reference for persisting job state",
+    )
+    max_concurrent_jobs: int = Field(
+        default=10,
+        description="Maximum number of jobs that can run concurrently",
+    )
+
+
+class CelerySchedulerConfig(SchedulerConfig):
+    """
+    Configuration for Celery-based distributed scheduler.
+
+    This scheduler distributes jobs across multiple worker processes/machines.
+    Suitable for production and multi-worker deployments.
+    """
+
+    type: Literal["celery"] = "celery"
+    broker_url: str = Field(
+        description="Celery broker URL (e.g., 'redis://localhost:6379/0')",
+    )
+    result_backend: str = Field(
+        description="Celery result backend URL (e.g., 'redis://localhost:6379/1')",
+    )
+
+
+SchedulerConfigUnion = Annotated[
+    InlineSchedulerConfig | CelerySchedulerConfig,
+    Field(discriminator="type"),
+]
diff --git a/src/llama_stack/providers/utils/job_scheduler/inline/__init__.py b/src/llama_stack/providers/utils/job_scheduler/inline/__init__.py
@@ -0,0 +1,9 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+
+from .scheduler import InlineSchedulerImpl
+
+__all__ = ["InlineSchedulerImpl"]
diff --git a/src/llama_stack/providers/utils/job_scheduler/inline/scheduler.py b/src/llama_stack/providers/utils/job_scheduler/inline/scheduler.py
