Your Name
3 months ago
parent
edc293cb6f
commit
6be667bba4
@ -0,0 +1,393 @@
|
||||
from typing import List, Callable, Union, Optional
|
||||
from loguru import logger
|
||||
from swarms.structs.base_swarm import BaseSwarm
|
||||
from queue import PriorityQueue
|
||||
from concurrent.futures import (
|
||||
ThreadPoolExecutor,
|
||||
as_completed,
|
||||
)
|
||||
import time
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
|
||||
class SwarmRunData(BaseModel):
|
||||
"""
|
||||
Pydantic model to capture metadata about each swarm's execution.
|
||||
"""
|
||||
|
||||
swarm_name: str
|
||||
task: str
|
||||
priority: int
|
||||
start_time: Optional[float] = None
|
||||
end_time: Optional[float] = None
|
||||
duration: Optional[float] = None
|
||||
status: str = "Pending"
|
||||
retries: int = 0
|
||||
result: Optional[str] = None
|
||||
exception: Optional[str] = None
|
||||
|
||||
|
||||
class FederatedSwarmModel(BaseModel):
|
||||
"""
|
||||
Pydantic base model to capture and log data for the FederatedSwarm system.
|
||||
"""
|
||||
|
||||
task: str
|
||||
swarms_data: List[SwarmRunData] = Field(default_factory=list)
|
||||
|
||||
def add_swarm(self, swarm_name: str, task: str, priority: int):
|
||||
swarm_data = SwarmRunData(
|
||||
swarm_name=swarm_name, task=task, priority=priority
|
||||
)
|
||||
self.swarms_data.append(swarm_data)
|
||||
|
||||
def update_swarm_status(
|
||||
self,
|
||||
swarm_name: str,
|
||||
status: str,
|
||||
start_time: float = None,
|
||||
end_time: float = None,
|
||||
retries: int = 0,
|
||||
result: str = None,
|
||||
exception: str = None,
|
||||
):
|
||||
for swarm in self.swarms_data:
|
||||
if swarm.name == swarm_name:
|
||||
swarm.status = status
|
||||
if start_time:
|
||||
swarm.start_time = start_time
|
||||
if end_time:
|
||||
swarm.end_time = end_time
|
||||
swarm.duration = end_time - swarm.start_time
|
||||
swarm.retries = retries
|
||||
swarm.result = result
|
||||
swarm.exception = exception
|
||||
break
|
||||
|
||||
|
||||
class FederatedSwarm:
|
||||
def __init__(
|
||||
self,
|
||||
swarms: List[Union[BaseSwarm, Callable]],
|
||||
max_workers: int = 4,
|
||||
):
|
||||
"""
|
||||
Initializes the FederatedSwarm with a list of swarms or callable objects and
|
||||
sets up a priority queue and thread pool for concurrency.
|
||||
|
||||
Args:
|
||||
swarms (List[Union[BaseSwarm, Callable]]): A list of swarms (BaseSwarm) or callable objects.
|
||||
max_workers (int): The maximum number of concurrent workers (threads) to run swarms in parallel.
|
||||
"""
|
||||
self.swarms = PriorityQueue()
|
||||
self.max_workers = max_workers
|
||||
self.thread_pool = ThreadPoolExecutor(
|
||||
max_workers=self.max_workers
|
||||
)
|
||||
self.task_queue = []
|
||||
self.future_to_swarm = {}
|
||||
self.results = {}
|
||||
self.validate_swarms(swarms)
|
||||
|
||||
def init_metadata(self, task: str):
|
||||
"""
|
||||
Initializes the Pydantic base model to capture metadata about the current task and swarms.
|
||||
"""
|
||||
self.metadata = FederatedSwarmModel(task=task)
|
||||
for priority, swarm in list(self.swarms.queue):
|
||||
swarm_name = (
|
||||
swarm.__class__.__name__
|
||||
if hasattr(swarm, "__class__")
|
||||
else str(swarm)
|
||||
)
|
||||
self.metadata.add_swarm(
|
||||
swarm_name=swarm_name, task=task, priority=priority
|
||||
)
|
||||
logger.info(f"Metadata initialized for task '{task}'.")
|
||||
|
||||
def validate_swarms(
|
||||
self, swarms: List[Union[BaseSwarm, Callable]]
|
||||
):
|
||||
"""
|
||||
Validates and adds swarms to the priority queue, ensuring each swarm has a `run(task)` method.
|
||||
|
||||
Args:
|
||||
swarms (List[Union[BaseSwarm, Callable]]): List of swarms with an optional priority value.
|
||||
"""
|
||||
for swarm, priority in swarms:
|
||||
if not callable(swarm):
|
||||
raise TypeError(f"{swarm} is not callable.")
|
||||
|
||||
if hasattr(swarm, "run"):
|
||||
logger.info(f"{swarm} has a 'run' method.")
|
||||
else:
|
||||
raise AttributeError(
|
||||
f"{swarm} does not have a 'run(task)' method."
|
||||
)
|
||||
|
||||
self.swarms.put((priority, swarm))
|
||||
logger.info(
|
||||
f"Swarm {swarm} added with priority {priority}."
|
||||
)
|
||||
|
||||
def run_parallel(
|
||||
self,
|
||||
task: str,
|
||||
timeout: Optional[float] = None,
|
||||
retries: int = 0,
|
||||
):
|
||||
"""
|
||||
Runs all swarms in parallel with prioritization and optional timeout.
|
||||
|
||||
Args:
|
||||
task (str): The task to be passed to the `run` method of each swarm.
|
||||
timeout (Optional[float]): Maximum time allowed for each swarm to run.
|
||||
retries (int): Number of retries allowed for failed swarms.
|
||||
"""
|
||||
logger.info(
|
||||
f"Running task '{task}' in parallel with timeout: {timeout}, retries: {retries}"
|
||||
)
|
||||
self.init_metadata(task)
|
||||
|
||||
while not self.swarms.empty():
|
||||
priority, swarm = self.swarms.get()
|
||||
swarm_name = (
|
||||
swarm.__class__.__name__
|
||||
if hasattr(swarm, "__class__")
|
||||
else str(swarm)
|
||||
)
|
||||
future = self.thread_pool.submit(
|
||||
self._run_with_retry,
|
||||
swarm,
|
||||
task,
|
||||
retries,
|
||||
timeout,
|
||||
swarm_name,
|
||||
)
|
||||
self.future_to_swarm[future] = swarm
|
||||
|
||||
for future in as_completed(self.future_to_swarm):
|
||||
swarm = self.future_to_swarm[future]
|
||||
try:
|
||||
result = future.result()
|
||||
swarm_name = (
|
||||
swarm.__class__.__name__
|
||||
if hasattr(swarm, "__class__")
|
||||
else str(swarm)
|
||||
)
|
||||
self.metadata.update_swarm_status(
|
||||
swarm_name=swarm_name,
|
||||
status="Completed",
|
||||
result=result,
|
||||
)
|
||||
logger.info(
|
||||
f"Swarm {swarm_name} completed successfully."
|
||||
)
|
||||
except Exception as e:
|
||||
swarm_name = (
|
||||
swarm.__class__.__name__
|
||||
if hasattr(swarm, "__class__")
|
||||
else str(swarm)
|
||||
)
|
||||
self.metadata.update_swarm_status(
|
||||
swarm_name=swarm_name,
|
||||
status="Failed",
|
||||
exception=str(e),
|
||||
)
|
||||
logger.error(f"Swarm {swarm_name} failed: {e}")
|
||||
self.results[swarm] = "Failed"
|
||||
|
||||
def run_sequentially(
|
||||
self,
|
||||
task: str,
|
||||
retries: int = 0,
|
||||
timeout: Optional[float] = None,
|
||||
):
|
||||
"""
|
||||
Runs all swarms sequentially in order of priority.
|
||||
|
||||
Args:
|
||||
task (str): The task to pass to the `run` method of each swarm.
|
||||
retries (int): Number of retries for failed swarms.
|
||||
timeout (Optional[float]): Optional time limit for each swarm.
|
||||
"""
|
||||
logger.info(f"Running task '{task}' sequentially.")
|
||||
|
||||
while not self.swarms.empty():
|
||||
priority, swarm = self.swarms.get()
|
||||
try:
|
||||
logger.info(
|
||||
f"Running swarm {swarm} with priority {priority}."
|
||||
)
|
||||
self._run_with_retry(swarm, task, retries, timeout)
|
||||
logger.info(f"Swarm {swarm} completed successfully.")
|
||||
except Exception as e:
|
||||
logger.error(f"Swarm {swarm} failed with error: {e}")
|
||||
|
||||
def _run_with_retry(
|
||||
self,
|
||||
swarm: Union[BaseSwarm, Callable],
|
||||
task: str,
|
||||
retries: int,
|
||||
timeout: Optional[float],
|
||||
swarm_name: str,
|
||||
):
|
||||
"""
|
||||
Helper function to run a swarm with a retry mechanism and optional timeout.
|
||||
|
||||
Args:
|
||||
swarm (Union[BaseSwarm, Callable]): The swarm to run.
|
||||
task (str): The task to pass to the swarm.
|
||||
retries (int): The number of retries allowed for the swarm in case of failure.
|
||||
timeout (Optional[float]): Maximum time allowed for the swarm to run.
|
||||
swarm_name (str): Name of the swarm (used for metadata).
|
||||
"""
|
||||
attempts = 0
|
||||
start_time = time.time()
|
||||
while attempts <= retries:
|
||||
try:
|
||||
logger.info(
|
||||
f"Running swarm {swarm}. Attempt: {attempts + 1}"
|
||||
)
|
||||
self.metadata.update_swarm_status(
|
||||
swarm_name=swarm_name,
|
||||
status="Running",
|
||||
start_time=start_time,
|
||||
)
|
||||
if hasattr(swarm, "run"):
|
||||
if timeout:
|
||||
start_time = time.time()
|
||||
swarm.run(task)
|
||||
duration = time.time() - start_time
|
||||
if duration > timeout:
|
||||
raise TimeoutError(
|
||||
f"Swarm {swarm} timed out after {duration:.2f}s."
|
||||
)
|
||||
else:
|
||||
swarm.run(task)
|
||||
else:
|
||||
swarm(task)
|
||||
end_time = time.time()
|
||||
self.metadata.update_swarm_status(
|
||||
swarm_name=swarm_name,
|
||||
status="Completed",
|
||||
end_time=end_time,
|
||||
retries=attempts,
|
||||
)
|
||||
return "Success"
|
||||
except Exception as e:
|
||||
logger.error(f"Swarm {swarm} failed: {e}")
|
||||
attempts += 1
|
||||
if attempts > retries:
|
||||
end_time = time.time()
|
||||
self.metadata.update_swarm_status(
|
||||
swarm_name=swarm_name,
|
||||
status="Failed",
|
||||
end_time=end_time,
|
||||
retries=attempts,
|
||||
exception=str(e),
|
||||
)
|
||||
logger.error(f"Swarm {swarm} exhausted retries.")
|
||||
raise
|
||||
|
||||
def add_swarm(
|
||||
self, swarm: Union[BaseSwarm, Callable], priority: int
|
||||
):
|
||||
"""
|
||||
Adds a new swarm to the FederatedSwarm at runtime.
|
||||
|
||||
Args:
|
||||
swarm (Union[BaseSwarm, Callable]): The swarm to add.
|
||||
priority (int): The priority level for the swarm.
|
||||
"""
|
||||
self.swarms.put((priority, swarm))
|
||||
logger.info(
|
||||
f"Swarm {swarm} added dynamically with priority {priority}."
|
||||
)
|
||||
|
||||
def queue_task(self, task: str):
|
||||
"""
|
||||
Adds a task to the internal task queue for batch processing.
|
||||
|
||||
Args:
|
||||
task (str): The task to queue.
|
||||
"""
|
||||
self.task_queue.append(task)
|
||||
logger.info(f"Task '{task}' added to the queue.")
|
||||
|
||||
def process_task_queue(self):
|
||||
"""
|
||||
Processes all tasks in the task queue.
|
||||
"""
|
||||
for task in self.task_queue:
|
||||
logger.info(f"Processing task: {task}")
|
||||
self.run_parallel(task)
|
||||
self.task_queue = []
|
||||
|
||||
def log_swarm_results(self):
|
||||
"""
|
||||
Logs the results of all swarms after execution.
|
||||
"""
|
||||
logger.info("Logging swarm results...")
|
||||
for swarm, result in self.results.items():
|
||||
logger.info(f"Swarm {swarm}: {result}")
|
||||
|
||||
def get_swarm_status(self) -> dict:
|
||||
"""
|
||||
Retrieves the status of each swarm (completed, running, failed).
|
||||
|
||||
Returns:
|
||||
dict: Dictionary containing swarm statuses.
|
||||
"""
|
||||
status = {}
|
||||
for future, swarm in self.future_to_swarm.items():
|
||||
if future.done():
|
||||
status[swarm] = "Completed"
|
||||
elif future.running():
|
||||
status[swarm] = "Running"
|
||||
else:
|
||||
status[swarm] = "Failed"
|
||||
return status
|
||||
|
||||
def cancel_running_swarms(self):
|
||||
"""
|
||||
Cancels all currently running swarms by shutting down the thread pool.
|
||||
"""
|
||||
logger.warning("Cancelling all running swarms...")
|
||||
self.thread_pool.shutdown(wait=False)
|
||||
logger.info("All running swarms cancelled.")
|
||||
|
||||
|
||||
# Example Usage:
|
||||
|
||||
|
||||
# class ExampleSwarm(BaseSwarm):
|
||||
# def run(self, task: str):
|
||||
# logger.info(f"ExampleSwarm is processing task: {task}")
|
||||
|
||||
|
||||
# def example_callable(task: str):
|
||||
# logger.info(f"Callable is processing task: {task}")
|
||||
|
||||
|
||||
# if __name__ == "__main__":
|
||||
# swarms = [(ExampleSwarm(), 1), (example_callable, 2)]
|
||||
# federated_swarm = FederatedSwarm(swarms)
|
||||
|
||||
# # Run in parallel
|
||||
# federated_swarm.run_parallel(
|
||||
# "Process data", timeout=10, retries=3
|
||||
# )
|
||||
|
||||
# # Run sequentially
|
||||
# federated_swarm.run_sequentially("Process data sequentially")
|
||||
|
||||
# # Log results
|
||||
# federated_swarm.log_swarm_results()
|
||||
|
||||
# # Get status of swarms
|
||||
# status = federated_swarm.get_swarm_status()
|
||||
# logger.info(f"Swarm statuses: {status}")
|
||||
|
||||
# # Cancel running swarms (if needed)
|
||||
# # federated_swarm.cancel_running_swarms()
|
||||
Loading…
Reference in new issue