global_scheduler.patch

diff --git a/sarathi/benchmark/benchmark_runner.py b/sarathi/benchmark/benchmark_runner.py
index 083dcb3..36bf343 100644
--- a/sarathi/benchmark/benchmark_runner.py
+++ b/sarathi/benchmark/benchmark_runner.py
@@ -2,6 +2,9 @@ import logging
 import os
 import time
 
+from typing import Optional, List
+from queue import Queue
+
 import ray
 import wandb
 from tqdm import tqdm
@@ -11,30 +14,23 @@ from sarathi.benchmark.config import BenchmarkConfig
 from sarathi.benchmark.entities import Request
 from sarathi.benchmark.request_generator import RequestGeneratorRegistry
 from sarathi.benchmark.utils.random import set_seeds
-from sarathi.config import ReplicaConfig
+from sarathi.config import ReplicaConfig, BaseGlobalSchedulerTypeConfig
 from sarathi.metrics.metrics_store import MetricsStore
 from sarathi.types import ReplicaResourceMapping, ResourceMapping
 from sarathi.utils import get_ip
+from sarathi.core.datatypes.sequence import SamplerOutputs, Sequence
+from sarathi.utils import Counter
+from sarathi.engine.multi_replica_llm_engine import MultiReplicaLLMEngine
 
 logger = logging.getLogger(__name__)
 
 
-class BenchmarkRunner:
+class BenchmarkRunnerLauncher:
 
-    def __init__(
-        self,
-        replica_id: int,
-        config: BenchmarkConfig,
-        resource_mapping: ResourceMapping,
-    ) -> None:
-        self.replica_id = replica_id
+    def __init__(self, config: BenchmarkConfig) -> None:
         self.config = config
 
-        replica_config = ReplicaConfig(
-            replica_id,
-            self.config.output_dir,
-            resource_mapping,
-        )
+        replica_config = ReplicaConfig(0, self.config.output_dir)
         os.makedirs(replica_config.output_dir, exist_ok=True)
 
         set_seeds(self.config.seed)
@@ -44,19 +40,12 @@ class BenchmarkRunner:
         )
         self.requests = request_generator.generate()
 
-        # select every nth request for this replica
-        # e.g. if there are 4 replicas, and this is the 2nd replica, then
-        # we will select the 2nd, 6th, 10th, ... requests
-        # round robin scheduling
-        self.requests = self.requests[self.replica_id :: self.config.num_replicas]
+        self.config.metrics_config.wandb_project = None
 
-        if self.config.num_replicas > 1:
-            # disable per-replica wandb logging for multi-replica runs
-            # so that we can aggregate metrics across all replicas
-            self.config.metrics_config.wandb_project = None
+        self.system_config = self.config.create_system_config(replica_config)
+        self.system_config.num_replicas = self.config.num_replicas
 
-        system_config = self.config.create_system_config(replica_config)
-        self.llm_engine = LLMEngine.from_system_config(system_config)
+        self.llm_engine = MultiReplicaLLMEngine(self.system_config)
 
         if wandb.run is not None:
             wandb.config.update(self.config.to_dict())
@@ -95,9 +84,10 @@ class BenchmarkRunner:
 
         num_processed_requests = 0
         num_steps = 0
+
         pbar = tqdm(
             total=len(self.requests),
-            desc=f"Replica {self.replica_id} processed requests",
+            desc=f"Total processed requests",
         )
         start_time = time.monotonic()
 
@@ -119,7 +109,11 @@ class BenchmarkRunner:
         pbar.close()
 
         logger.info(
-            f"Replica {self.replica_id} exiting after processing {len(self.requests)} ({num_steps} iterations), Total time taken: {end_time - start_time:.2f} seconds"
+            f"{num_processed_requests} requests processed and exited before completing all requests"
+        )
+
+        logger.info(
+            f"Exiting after processing {len(self.requests)} ({num_steps} iterations), Total time taken: {end_time - start_time:.2f} seconds"
         )
 
         if self.config.enable_profiling:
@@ -130,150 +124,23 @@ class BenchmarkRunner:
         first_request_time = time.monotonic()
         while index < len(self.requests):
             request = self.requests[index]
+
             self.llm_engine.add_request(
                 **self._get_input_params(request, first_request_time)
             )
             index += 1
 
-    def run(self) -> None:
+    def run_benchmark(self) -> None:
         self.llm_engine.reset_metrics()
         self._add_requests()
+        self.llm_engine.start_engine_execution()
         self._run()
         self.llm_engine.pull_worker_metrics()
         metric_store = self.llm_engine.get_metric_store()
         return metric_store
 
-
-class BenchmarkRunnerLauncher:
-
-    def __init__(self, config: BenchmarkConfig) -> None:
-        self.config = config
-        self.is_multi_replica = self.config.num_replicas > 1
-
-        ray.init(ignore_reinit_error=True)
-
-        self._validate_cluster_resources()
-        self.runners = self._create_runners()
-
-        if self.is_multi_replica:
-            self.aggregate_metric_store = self._create_aggregate_metric_store()
-
-    def _validate_cluster_resources(self):
-        num_replicas = self.config.num_replicas
-        num_gpus_required = num_replicas * self.config.parallel_config.world_size
-
-        available_resources = ray.available_resources()
-
-        assert (
-            available_resources["GPU"] >= num_gpus_required
-        ), f"Insufficient GPUs. Required: {num_gpus_required}, Available: {available_resources['GPU']}"
-
-    def _get_replica_resource_mapping(self) -> ReplicaResourceMapping:
-        if self.config.replica_resource_mapping:
-            assert len(self.config.replica_resource_mapping) == self.config.num_replicas
-            logger.info(
-                f"Replica resource mapping: {self.config.replica_resource_mapping}"
-            )
-            return self.config.replica_resource_mapping
-
-        cluster_resources_keys = list(ray.available_resources().keys())
-        num_gpus = ray.available_resources()["GPU"]
-        ip_addresses = [
-            x
-            for x in cluster_resources_keys
-            if x.startswith("node:") and x != "node:__internal_head__"
-        ]
-
-        runner_ip = f"node:{get_ip()}"
-
-        ip_addresses.remove(runner_ip)
-        ip_addresses.insert(0, runner_ip)
-
-        num_nodes = len(ip_addresses)
-        assert num_nodes > 0, "No nodes found in the cluster"
-        assert num_gpus > 0, "No GPUs found in the cluster"
-        assert (
-            num_gpus % num_nodes == 0
-        ), f"Number of GPUs ({num_gpus}) is not a multiple of number of nodes ({num_nodes})"
-        num_gpus_per_node = int(num_gpus // num_nodes)
-        num_replicas = self.config.num_replicas
-        num_gpus_per_replica = self.config.parallel_config.world_size
-
-        assert (
-            num_gpus >= num_replicas * num_gpus_per_replica
-        ), f"Insufficient GPUs. Required: {num_replicas * num_gpus_per_replica}, Available: {num_gpus}"
-
-        replica_resource_mapping = []
-
-        available_gpus = []
-        for ip_address in ip_addresses:
-            for gpu_id in reversed(range(num_gpus_per_node)):
-                available_gpus.append((ip_address, gpu_id))
-
-        for _ in range(num_replicas):
-            resource_mapping = []
-            for _ in range(num_gpus_per_replica):
-                resource_mapping.append(available_gpus.pop(0))
-            replica_resource_mapping.append(resource_mapping)
-
-        logger.info(f"Replica resource mapping: {replica_resource_mapping}")
-
-        return replica_resource_mapping
-
-    def _create_runners(self):
-        replica_resource_mapping = self._get_replica_resource_mapping()
-
-        if not self.is_multi_replica:
-            return [BenchmarkRunner(0, self.config, replica_resource_mapping[0])]
-
-        runner_class = ray.remote(num_cpus=1)(BenchmarkRunner)
-
-        runners = []
-
-        for replica_id in range(self.config.num_replicas):
-            runners.append(
-                runner_class.options(
-                    resources={
-                        replica_resource_mapping[replica_id][0][0]: 0.01,
-                    },
-                ).remote(replica_id, self.config, replica_resource_mapping[replica_id])
-            )
-
-        return runners
-
-    def _create_aggregate_metric_store(self):
-        replica_config = ReplicaConfig(
-            replica_id=0,  # dummy replica id
-            output_dir=self.config.output_dir,
-        )
-        metrics_store = MetricsStore.get_instance(
-            replica_config,
-            self.config.model_config,
-            self.config.metrics_config,
-        )
-
-        if wandb.run is not None:
-            wandb.config.update(self.config.to_dict())
-
-        metrics_store.mark_initial_memory_profiling_done()
-
-        return metrics_store
-
     def run(self):
-        if self.is_multi_replica:
-            ray.get([runner.warmup.remote() for runner in self.runners])
-
-            runner_metrics = ray.get([runner.run.remote() for runner in self.runners])
-
-            for runner_metric in runner_metrics:
-                self.aggregate_metric_store.merge(runner_metric)
-
-            if wandb.run is not None:
-                wandb.config.update(self.config.__dict__)
-
-            self.aggregate_metric_store.plot()
-        else:
-            metric_store = self.runners[0].run()
-            metric_store.plot()
-
-        wandb.finish()
+        metric_store = self.run_benchmark()
+        metric_store.plot()
+        if wandb.run is not None:
+            wandb.finish()
diff --git a/sarathi/benchmark/config.py b/sarathi/benchmark/config.py
index bbc5d10..4780e24 100644
--- a/sarathi/benchmark/config.py
+++ b/sarathi/benchmark/config.py
@@ -2,7 +2,7 @@ import datetime
 from dataclasses import dataclass, field
 from typing import Optional
 
-from sarathi.config import BaseEndpointConfig
+from sarathi.config import BaseEndpointConfig, BaseGlobalSchedulerTypeConfig
 from sarathi.config.base_poly_config import BasePolyConfig
 from sarathi.config.flat_dataclass import create_flat_dataclass
 from sarathi.logger import init_logger
diff --git a/sarathi/config/config.py b/sarathi/config/config.py
index a36d583..63e3605 100644
--- a/sarathi/config/config.py
+++ b/sarathi/config/config.py
@@ -7,7 +7,7 @@ from sarathi.config.base_poly_config import BasePolyConfig
 from sarathi.config.flat_dataclass import create_flat_dataclass
 from sarathi.logger import init_logger
 from sarathi.transformers_utils.config import get_config
-from sarathi.types import AttentionBackend, ResourceMapping, SchedulerType
+from sarathi.types import AttentionBackend, ResourceMapping, SchedulerType, GlobalSchedulerType
 from sarathi.utils.hf_utils import get_and_verify_dtype, get_and_verify_max_len
 
 logger = init_logger(__name__)
@@ -221,7 +221,6 @@ class SimpleChunkingSchedulerConfig(BaseSchedulerConfig):
 
 @dataclass
 class OrcaSchedulerConfig(BaseSchedulerConfig):
-
     def get_max_num_batched_tokens(self, max_model_len: int):
         return self.max_num_seqs * max_model_len
 
@@ -232,7 +231,6 @@ class OrcaSchedulerConfig(BaseSchedulerConfig):
 
 @dataclass
 class FasterTransformerSchedulerConfig(BaseSchedulerConfig):
-
     def get_max_num_batched_tokens(self, max_model_len: int):
         return self.max_num_seqs * max_model_len
 
@@ -356,6 +354,28 @@ class WorkerConfig:
             )
 
 
+@dataclass
+class BaseGlobalSchedulerTypeConfig(BasePolyConfig):
+    scheduler_type: str = field(
+        default="pull",
+        metadata={"help": "Replica level scheduler type either pull or RR"},
+    )
+
+
+@dataclass
+class PullGlobalSchedulerConfig(BaseGlobalSchedulerTypeConfig):
+    @staticmethod
+    def get_type():
+        return GlobalSchedulerType.PULL
+    
+
+@dataclass
+class RoundRobinGlobalSchedulerConfig(BaseGlobalSchedulerTypeConfig):
+    @staticmethod
+    def get_type():
+        return GlobalSchedulerType.ROUND_ROBIN
+
+
 @dataclass
 class SystemConfig:
     replica_config: ReplicaConfig = field(default_factory=ReplicaConfig)
@@ -367,6 +387,10 @@ class SystemConfig:
         default_factory=SarathiSchedulerConfig
     )
     metrics_config: MetricsConfig = field(default_factory=MetricsConfig)
+    num_replicas: int = field(default=1, metadata={"help": "Number of replicas."})
+    global_scheduler_config : BaseGlobalSchedulerTypeConfig = field(
+        default_factory=BaseGlobalSchedulerTypeConfig
+    )
 
 
 @dataclass
diff --git a/sarathi/core/datatypes/sequence.py b/sarathi/core/datatypes/sequence.py
index 07dbf1b..7b61e90 100644
--- a/sarathi/core/datatypes/sequence.py
+++ b/sarathi/core/datatypes/sequence.py
@@ -1,6 +1,8 @@
 """Sequence and its related classes."""
 
-from typing import List, Optional
+from typing import List, Optional,Any
+from dataclasses import dataclass, field
+import random
 
 from sarathi.core.datatypes.block import LogicalTokenBlock
 from sarathi.core.datatypes.sampling_params import SamplingParams
@@ -8,6 +10,11 @@ from sarathi.core.datatypes.sequence_state import SequenceState
 from sarathi.core.datatypes.sequence_status import SequenceStatus
 
 
+@dataclass(order=True)
+class SequenceWithPriority:
+    priority : float
+    seq : Any=field(compare=False)
+
 class Sequence:
     """Stores the data, status, and block information of a sequence.
 
@@ -216,6 +223,9 @@ class Sequence:
             f"prompt_stage_processing_finished={self.prompt_stage_processing_finished})"
         )
 
+    @property
+    def arrived_at(self) -> float:
+        return self.arrival_time
 
 class SequenceScheduleMetadata:
     """Metadata generated by the scheduler for sequence that has been scheduled.
diff --git a/sarathi/core/scheduler/base_scheduler.py b/sarathi/core/scheduler/base_scheduler.py
index 6c73d59..d69dc9f 100644
--- a/sarathi/core/scheduler/base_scheduler.py
+++ b/sarathi/core/scheduler/base_scheduler.py
@@ -1,15 +1,19 @@
 from abc import ABC, abstractmethod
 from typing import List
+from queue import PriorityQueue
 
 from sarathi.config import BaseSchedulerConfig, CacheConfig, ModelConfig, ParallelConfig
 from sarathi.core.block_space_manager.block_space_manager_registry import (
     BlockSpaceManagerRegistry,
 )
 from sarathi.core.datatypes.scheduler_output import SchedulerOutputs
-from sarathi.core.datatypes.sequence import Sequence, SequenceStatus
+from sarathi.core.datatypes.sequence import Sequence, SequenceStatus, SequenceWithPriority
+from sarathi.core.sequence_manager.engine_sequence_manager import EngineSequenceManager
 from sarathi.core.policy import PolicyFactory
 from sarathi.logger import init_logger
 from sarathi.metrics.metrics_store import MetricsStore
+from sarathi.utils.threading_utils import synchronized
+
 
 logger = init_logger(__name__)
 
@@ -22,6 +26,9 @@ class BaseScheduler(ABC):
         scheduler_config: BaseSchedulerConfig,
         cache_config: CacheConfig,
         parallel_config: ParallelConfig,
+        waiting_queue : PriorityQueue,
+        replica_seq_manager : EngineSequenceManager,
+        metric_store : MetricsStore,
     ) -> None:
         self.metrics_store = MetricsStore.get_instance()
         self.model_config = model_config
@@ -42,13 +49,18 @@ class BaseScheduler(ABC):
             model_config.max_model_len,
         )
         self.prompt_limit = model_config.max_model_len
+        self.replica_seq_manager = replica_seq_manager
+        self.new_seqs: List[Sequence] = []
+        self.metrics_store = metric_store
+        self.seq_seen = set()
 
         # number of running batches should be less than or equal to the number of pipeline stages
         self.num_running_batches = 0
 
         # TODO(zhuohan): Use deque instead of list for better performance.
         # Sequence groups in the WAITING state.
-        self.waiting: List[Sequence] = []
+        # self.waiting : PriorityQueue = PriorityQueue()
+        self.waiting : PriorityQueue = waiting_queue
         # Sequence groups in the RUNNING state.
         self.running: List[Sequence] = []
 
@@ -56,19 +68,37 @@ class BaseScheduler(ABC):
         self._iteration_id = -1
 
     def add_seq(self, seq: Sequence) -> None:
-        # Add sequence groups to the waiting queue.
-        self.waiting.append(seq)
+        # Add sequence groups to the waiting queue. 
+        wrapped_seq = SequenceWithPriority(seq.arrived_at, seq)
+
+        self.waiting.put(wrapped_seq)
 
     def has_unfinished_seqs(self) -> bool:
-        return self.waiting or self.running
+        return self.waiting.qsize() > 0 or self.running
 
     def get_num_unfinished_seqs(self) -> int:
-        return len(self.waiting) + len(self.running)
+        return self.waiting.qsize() + len(self.running)
 
     @abstractmethod
     def _schedule(self) -> SchedulerOutputs:
         pass
 
+    @synchronized
+    def add_to_new_seqs(self, seq: Sequence) -> None:
+        self.new_seqs.append(seq)
+
+    @synchronized
+    def get_new_seqs(
+        self,
+    ) -> List[Sequence]:
+        new_seqs = self.new_seqs
+        self.new_seqs = []
+        return new_seqs
+    
+    @synchronized
+    def add_seq_to_seq_manager(self, seq: Sequence) -> None:
+        self.replica_seq_manager.add_seq(seq)
+
     def schedule(self) -> SchedulerOutputs:
         # Schedule sequence groups.
         # This function call changes the internal states of the scheduler
@@ -82,7 +112,7 @@ class BaseScheduler(ABC):
                 preempted_seq_ids=[],
                 scheduled_seq_metadata_list=[],
             )
-
+        
         scheduler_outputs = self._schedule()
 
         if not scheduler_outputs.is_empty():
@@ -119,7 +149,10 @@ class BaseScheduler(ABC):
     ) -> None:
         assert seq.is_executing()
         self._free_seq(seq)
-        self.waiting.insert(0, seq)
+
+        wrapped_seq = SequenceWithPriority(seq.arrived_at, seq)
+        
+        self.waiting.put(wrapped_seq)
 
     def _check_request_prompt_length(self, seq: Sequence) -> bool:
         if seq.get_len() > self.prompt_limit:
@@ -128,7 +161,7 @@ class BaseScheduler(ABC):
                 f" and exceeds limit of {self.prompt_limit}"
             )
             seq.set_status(SequenceStatus.FINISHED_IGNORED)
-            self.waiting.pop(0)
+            self.waiting.get(block=False)
             return False
 
-        return True
+        return True
\ No newline at end of file
diff --git a/sarathi/core/scheduler/faster_transformer_scheduler.py b/sarathi/core/scheduler/faster_transformer_scheduler.py
index b103630..8ae05d3 100644
--- a/sarathi/core/scheduler/faster_transformer_scheduler.py
+++ b/sarathi/core/scheduler/faster_transformer_scheduler.py
@@ -1,5 +1,6 @@
 import time
 from typing import List
+from queue import PriorityQueue
 
 from sarathi.config import (
     CacheConfig,
@@ -12,8 +13,11 @@ from sarathi.core.block_space_manager.faster_transformer_block_space_manager imp
 )
 from sarathi.core.datatypes.scheduler_output import SchedulerOutputs
 from sarathi.core.datatypes.sequence import SequenceScheduleMetadata
+from sarathi.core.sequence_manager.engine_sequence_manager import EngineSequenceManager
 from sarathi.core.scheduler.base_scheduler import BaseScheduler
 from sarathi.logger import init_logger
+from sarathi.metrics.metrics_store import MetricsStore
+
 
 logger = init_logger(__name__)
 
@@ -26,8 +30,11 @@ class FasterTransformerScheduler(BaseScheduler):
         scheduler_config: FasterTransformerSchedulerConfig,
         cache_config: CacheConfig,
         parallel_config: ParallelConfig,
+        waiting_queue : PriorityQueue,
+        replica_seq_manager : EngineSequenceManager,
+        metric_store : MetricsStore,
     ) -> None:
-        super().__init__(model_config, scheduler_config, cache_config, parallel_config)
+        super().__init__(model_config, scheduler_config, cache_config, parallel_config, waiting_queue, replica_seq_manager, metric_store)
 
     def get_block_space_manager_class(self):
         return FasterTransformerBlockSpaceManager
@@ -58,9 +65,9 @@ class FasterTransformerScheduler(BaseScheduler):
         # Optimization: We do not sort the waiting queue since the preempted
         # sequences are added to the front and the new sequences
         # are added to the back.
-        while self.waiting:
-            seq = self.waiting[0]
-
+        while self.waiting.qsize() > 0:
+            seq_wrapped = self.waiting.queue[0]
+            seq = seq_wrapped.seq
             # This is required to handle benchmarking where
             # we set request arrival time ahead of time
             if seq.arrival_time > now:
@@ -77,7 +84,8 @@ class FasterTransformerScheduler(BaseScheduler):
             if len(self.running) + 1 > self.scheduler_config.max_num_seqs:
                 break
 
-            seq = self.waiting.pop(0)
+            seq_wrapped = self.waiting.get()
+            seq = seq_wrapped.seq
             self._allocate(seq)
             self.running.append(seq)
             scheduled_seq_metadata_list.append(
diff --git a/sarathi/core/scheduler/global_scheduler.py b/sarathi/core/scheduler/global_scheduler.py
new file mode 100644
index 0000000..dde1064
--- /dev/null
+++ b/sarathi/core/scheduler/global_scheduler.py
@@ -0,0 +1,172 @@
+import random
+import time
+
+from typing import List, Optional, Dict
+from threading import Thread
+from queue import PriorityQueue
+
+from sarathi.core.datatypes.sampling_params import SamplingParams
+from sarathi.logger import init_logger
+from sarathi.core.sequence_manager.engine_sequence_manager import EngineSequenceManager
+from sarathi.core.datatypes.sequence import SamplerOutputs, Sequence, SequenceWithPriority
+from sarathi.utils.threading_utils import synchronized
+
+logger = init_logger(__name__)
+
+
+class GlobalScheduler:
+    def __init__(self, config, num_replicas, sequence_counter, ):
+        logger.info(
+            f"GlobalScheduler initialized with {num_replicas} replicas"
+        )
+
+        self.config = config
+        self.num_replicas = num_replicas
+        self.replica_llm_engine_mapping = {}
+        self.seq_counter = sequence_counter
+        self.seq_map = None
+        self.new_seq_list = None        
+
+    def init_queue(self):
+        pass
+    
+    def get_replica_queue(self, replica_id):
+        pass
+
+    def get_replica_queue_mapping(self):
+        pass
+
+    def get_seq_map(self):
+        return self.seq_map
+    
+    def get_new_seq_list(self):
+        return self.new_seq_list
+
+    def set_replica_llm_engine(self, replica_id, replica_llm_engine):
+        self.replica_llm_engine_mapping[replica_id] = replica_llm_engine
+    
+    def _assign_queue(self, seq: Sequence, replica_id : int):
+        pass
+    
+    @synchronized
+    def assign_seq_replica(self, seq : Sequence) -> None:
+        pass
+
+    def assign_replica(
+        self,
+        prompt: Optional[str],
+        sampling_params: SamplingParams,
+        prompt_token_ids: Optional[List[int]] = None,
+        arrival_time: Optional[float] = None,
+        seq_id: Optional[str] = None,
+    ):
+        pass
+
+    def get_num_unfinished_requests(self):
+        pass
+
+    def has_unfinished_requests(self):
+        pass
+
+    def get_replica_id(self):
+        pass
+
+
+class PullScheduler(GlobalScheduler):
+    """
+    PullScheduler is a global scheduler that assigns requests to a global request queue and the replicas pull requests from the queue.
+    """
+
+    def __init__(self, config, num_replicas, sequence_counter):
+        super().__init__(config, num_replicas, sequence_counter)
+        logger.info(f"PullScheduler initialized with {num_replicas} replicas")
+
+    def init_queue(self):
+        self.replica_queue_mapping = {"global": PriorityQueue()}
+        self.seq_map: Dict[str, Sequence] = None
+        self.new_seq_list = None
+
+    def get_replica_queue_mapping(self):
+        return self.replica_queue_mapping
+
+    def get_replica_queue(self, replica_id):
+        return self.replica_queue_mapping["global"]
+
+    def assign_replica(
+        self,
+        prompt: Optional[str],
+        sampling_params: SamplingParams,
+        prompt_token_ids: Optional[List[int]] = None,
+        arrival_time: Optional[float] = None,
+        seq_id: Optional[str] = None,
+    ):
+        pass
+    
+    def _assign_queue(self, seq, replica_id):
+        wrapped_seq = SequenceWithPriority(seq.arrived_at, seq)
+        self.replica_queue_mapping["global"].put(wrapped_seq)
+    
+    @synchronized
+    def assign_seq_replica(self, seq : Sequence) -> None:
+        self._assign_queue(seq, None)
+        
+    def get_num_unfinished_requests(self):
+        return self.replica_queue_mapping["global"].qsize()
+
+    def has_unfinished_requests(self):
+        return not self.replica_queue_mapping["global"].empty()
+
+
+class RoundRobinScheduler(GlobalScheduler):
+    """
+    RoundRobinScheduler is a global scheduler that assigns requests to replicas in a round-robin manner.
+    """
+
+    def __init__(self, config, num_replicas, sequence_counter):
+        super().__init__(config, num_replicas, sequence_counter)
+        self.current_replica_id = 0
+        logger.info(f"RoundRobinScheduler initialized with {num_replicas} replicas")
+
+    def init_queue(self):
+        self.replica_queue_mapping = {
+            replica_id: PriorityQueue() for replica_id in range(self.num_replicas)
+        }
+
+    def get_replica_queue_mapping(self):
+        return self.replica_queue_mapping
+
+    def get_replica_queue(self, replica_id):
+        return self.replica_queue_mapping[replica_id]
+    
+    def _assign_queue(self, seq, replica_id):
+        logger.info(f"[ROUND ROBIN SCHED] Assigning request to replica {replica_id}")
+        wrapped_seq = SequenceWithPriority(seq.arrived_at, seq)
+        self.replica_queue_mapping[replica_id].put(wrapped_seq)
+
+    @synchronized
+    def assign_seq_replica(self, seq: Sequence) -> None:
+        replica_id = self.current_replica_id
+        self._assign_queue(seq, replica_id)
+        self.current_replica_id = (self.current_replica_id + 1) % self.num_replicas
+
+    def assign_replica(
+        self,
+        prompt: Optional[str],
+        sampling_params: SamplingParams,
+        prompt_token_ids: Optional[List[int]] = None,
+        arrival_time: Optional[float] = None,
+        seq_id: Optional[str] = None,
+    ):
+        pass
+
+    def get_num_unfinished_requests(self):
+        num_unfinished_requests = 0
+        for replica_id in range(self.num_replicas):
+            num_unfinished_requests += self.replica_queue_mapping[replica_id].qsize
+        return num_unfinished_requests
+
+    def has_unfinished_requests(self):
+        for replica_id in range(self.num_replicas):
+            if not self.replica_queue_mapping[replica_id].empty():
+                return True
+        return False
diff --git a/sarathi/core/scheduler/orca_scheduler.py b/sarathi/core/scheduler/orca_scheduler.py
index c80ad04..daae28d 100644
--- a/sarathi/core/scheduler/orca_scheduler.py
+++ b/sarathi/core/scheduler/orca_scheduler.py
@@ -1,5 +1,6 @@
 import time
 from typing import List
+from queue import PriorityQueue
 
 from sarathi.config import CacheConfig, ModelConfig, OrcaSchedulerConfig, ParallelConfig
 from sarathi.core.block_space_manager.orca_block_space_manager import (
@@ -7,8 +8,10 @@ from sarathi.core.block_space_manager.orca_block_space_manager import (
 )
 from sarathi.core.datatypes.scheduler_output import SchedulerOutputs
 from sarathi.core.datatypes.sequence import SequenceScheduleMetadata
+from sarathi.core.sequence_manager.engine_sequence_manager import EngineSequenceManager
 from sarathi.core.scheduler.base_scheduler import BaseScheduler
 from sarathi.logger import init_logger
+from sarathi.metrics.metrics_store import MetricsStore
 
 logger = init_logger(__name__)
 
@@ -21,8 +24,11 @@ class OrcaScheduler(BaseScheduler):
         scheduler_config: OrcaSchedulerConfig,
         cache_config: CacheConfig,
         parallel_config: ParallelConfig,
+        waiting_queue : PriorityQueue,
+        replica_seq_manager : EngineSequenceManager,
+        metric_store : MetricsStore,
     ) -> None:
-        super().__init__(model_config, scheduler_config, cache_config, parallel_config)
+        super().__init__(model_config, scheduler_config, cache_config, parallel_config,waiting_queue, replica_seq_manager, metric_store)
 
     def get_block_space_manager_class(self):
         return OrcaBlockSpaceManager
@@ -48,8 +54,9 @@ class OrcaScheduler(BaseScheduler):
         # Optimization: We do not sort the waiting queue since the preempted
         # sequences are added to the front and the new sequences
         # are added to the back.
-        while self.waiting:
-            seq = self.waiting[0]
+        while self.waiting.qsize() > 0:
+            seq_wrapped = self.waiting.queue[0]
+            seq = seq_wrapped.seq
 
             # This is required to handle benchmarking where we set request arrival time ahead of time
             if seq.arrival_time > now:
@@ -66,7 +73,8 @@ class OrcaScheduler(BaseScheduler):
             if len(self.running) + 1 > self.scheduler_config.max_num_seqs:
                 break
 
-            seq = self.waiting.pop(0)
+            seq_wrapped = self.waiting.get()
+            seq = seq_wrapped.seq
             self._allocate(seq)
             self.running.append(seq)
             scheduled_seq_metadata_list.append(
diff --git a/sarathi/core/scheduler/sarathi_scheduler.py b/sarathi/core/scheduler/sarathi_scheduler.py
index bd47563..49d33b2 100644
--- a/sarathi/core/scheduler/sarathi_scheduler.py
+++ b/sarathi/core/scheduler/sarathi_scheduler.py
@@ -1,5 +1,7 @@
+import copy
 import time
 from typing import List
+from queue import PriorityQueue
 
 import numpy as np
 
@@ -13,9 +15,11 @@ from sarathi.core.block_space_manager.sarathi_block_space_manager import (
     SarathiBlockSpaceManager,
 )
 from sarathi.core.datatypes.scheduler_output import SchedulerOutputs
-from sarathi.core.datatypes.sequence import Sequence, SequenceScheduleMetadata
+from sarathi.core.datatypes.sequence import Sequence, SequenceScheduleMetadata, SequenceWithPriority
+from sarathi.core.sequence_manager.engine_sequence_manager import EngineSequenceManager
 from sarathi.core.scheduler.base_scheduler import BaseScheduler
 from sarathi.logger import init_logger
+from sarathi.metrics.metrics_store import MetricsStore
 
 logger = init_logger(__name__)
 
@@ -28,8 +32,11 @@ class SarathiScheduler(BaseScheduler):
         scheduler_config: SarathiSchedulerConfig,
         cache_config: CacheConfig,
         parallel_config: ParallelConfig,
+        waiting_queue : PriorityQueue,
+        replica_seq_manager : EngineSequenceManager,
+        metric_store : MetricsStore,
     ) -> None:
-        super().__init__(model_config, scheduler_config, cache_config, parallel_config)
+        super().__init__(model_config, scheduler_config, cache_config, parallel_config, waiting_queue, replica_seq_manager, metric_store)
 
         self.chunk_size = self.scheduler_config.chunk_size
         self.enable_dynamic_chunking_schedule = (
@@ -193,8 +200,9 @@ class SarathiScheduler(BaseScheduler):
         # Optimization: We do not sort the waiting queue since the preempted
         # sequence groups are added to the front and the new sequence groups
         # are added to the back.
-        while self.waiting:
-            seq = self.waiting[0]
+        while self.waiting.qsize() > 0:
+            seq_wrapped = self.waiting.queue[0]
+            seq = seq_wrapped.seq
 
             # This is required to handle benchmarking where we set request arrival time ahead of time
             if seq.arrival_time > now:
@@ -223,8 +231,9 @@ class SarathiScheduler(BaseScheduler):
 
             if next_num_prefill_tokens == 0:
                 break
-
-            seq = self.waiting.pop(0)
+            
+            seq_wrapped = self.waiting.get()
+            seq = seq_wrapped.seq
             self._allocate(seq)
             num_batched_tokens += next_num_prefill_tokens
             scheduled_seq_metadata_list.append(
@@ -232,6 +241,13 @@ class SarathiScheduler(BaseScheduler):
                     seq, prompt_chunk_len=next_num_prefill_tokens
                 )
             )
+
+            if seq.seq_id not in self.seq_seen:
+                self.add_seq_to_seq_manager(seq)
+                self.add_to_new_seqs(copy.deepcopy(seq))
+                self.seq_seen.add(seq.seq_id)
+            self.metrics_store.on_request_arrival(seq)
+
             running.append(seq)
 
         # make sure that prefills are at the start of the batch, so that we don't violate assumptions
@@ -243,4 +259,4 @@ class SarathiScheduler(BaseScheduler):
             ignored_seq_ids=ignored_seq_ids,
             preempted_seq_ids=preempted_seq_ids,
             scheduled_seq_metadata_list=scheduled_seq_metadata_list,
-        )
+        )
\ No newline at end of file
diff --git a/sarathi/core/scheduler/simple_chunking_scheduler.py b/sarathi/core/scheduler/simple_chunking_scheduler.py
index 95b4329..7d6e43c 100644
--- a/sarathi/core/scheduler/simple_chunking_scheduler.py
+++ b/sarathi/core/scheduler/simple_chunking_scheduler.py
@@ -1,6 +1,7 @@
 import time
 from enum import Enum, auto
 from typing import List
+from queue import PriorityQueue
 
 from sarathi.config import (
     CacheConfig,
@@ -14,8 +15,10 @@ from sarathi.core.block_space_manager.vllm_block_space_manager import (
 from sarathi.core.datatypes.scheduler_output import SchedulerOutputs
 from sarathi.core.datatypes.sequence import Sequence, SequenceScheduleMetadata
 from sarathi.core.datatypes.sequence_status import SequenceStatus
+from sarathi.core.sequence_manager.engine_sequence_manager import EngineSequenceManager
 from sarathi.core.scheduler.base_scheduler import BaseScheduler
 from sarathi.logger import init_logger
+from sarathi.metrics.metrics_store import MetricsStore
 
 logger = init_logger(__name__)
 
@@ -33,8 +36,11 @@ class SimpleChunkingScheduler(BaseScheduler):
         scheduler_config: SimpleChunkingSchedulerConfig,
         cache_config: CacheConfig,
         parallel_config: ParallelConfig,
+        waiting_queue : PriorityQueue,
+        replica_seq_manager : EngineSequenceManager,
+        metric_store : MetricsStore,
     ) -> None:
-        super().__init__(model_config, scheduler_config, cache_config, parallel_config)
+        super().__init__(model_config, scheduler_config, cache_config, parallel_config, waiting_queue, replica_seq_manager, metric_store)
 
         self.chunk_size = self.scheduler_config.chunk_size
         self.whose_turn = Turn.PREFILL
@@ -115,8 +121,9 @@ class SimpleChunkingScheduler(BaseScheduler):
                 scheduled_seq_metadata_list=scheduled_seq_metadata_list,
             )
 
-        while self.waiting and self.whose_turn == Turn.PREFILL:
-            seq = self.waiting[0]
+        while self.waiting.qsize() > 0 and self.whose_turn == Turn.PREFILL:
+            seq_wrapped = self.waiting.queue[0]
+            seq = seq_wrapped.seq
             # This is required to handle benchmarking where
             # we set request arrival time ahead of time
             if seq.arrival_time > now:
@@ -141,7 +148,8 @@ class SimpleChunkingScheduler(BaseScheduler):
                 # not enough space to allocate the sequence
                 break
 
-            self.waiting.pop(0)
+            seq_wrapped = self.waiting.get()
+            seq = seq_wrapped.seq
             self._allocate(seq)
             self.running.append(seq)
             num_batched_tokens += next_num_prefill_tokens
diff --git a/sarathi/core/scheduler/vllm_scheduler.py b/sarathi/core/scheduler/vllm_scheduler.py
index 84cb3b1..e4bcfca 100644
--- a/sarathi/core/scheduler/vllm_scheduler.py
+++ b/sarathi/core/scheduler/vllm_scheduler.py
@@ -1,14 +1,17 @@
 import time
 from typing import List
+from queue import PriorityQueue
 
 from sarathi.config import CacheConfig, ModelConfig, ParallelConfig, VllmSchedulerConfig
 from sarathi.core.block_space_manager.vllm_block_space_manager import (
     VLLMBlockSpaceManager,
 )
 from sarathi.core.datatypes.scheduler_output import SchedulerOutputs
-from sarathi.core.datatypes.sequence import Sequence, SequenceScheduleMetadata
+from sarathi.core.datatypes.sequence import Sequence, SequenceScheduleMetadata, SequenceWithPriority
+from sarathi.core.sequence_manager.engine_sequence_manager import EngineSequenceManager
 from sarathi.core.scheduler.base_scheduler import BaseScheduler
 from sarathi.logger import init_logger
+from sarathi.metrics.metrics_store import MetricsStore
 
 logger = init_logger(__name__)
 
@@ -21,8 +24,12 @@ class VLLMScheduler(BaseScheduler):
         scheduler_config: VllmSchedulerConfig,
         cache_config: CacheConfig,
         parallel_config: ParallelConfig,
+        waiting_queue : PriorityQueue,
+        replica_seq_manager : EngineSequenceManager,
+        metric_store : MetricsStore,
+
     ) -> None:
-        super().__init__(model_config, scheduler_config, cache_config, parallel_config)
+        super().__init__(model_config, scheduler_config, cache_config, parallel_config, waiting_queue, replica_seq_manager, metric_store)
 
         self.max_num_batched_tokens = self.scheduler_config.get_max_num_batched_tokens(
             self.model_config.max_model_len
@@ -46,8 +53,10 @@ class VLLMScheduler(BaseScheduler):
         # Optimization: We do not sort the waiting queue since the preempted
         # sequence groups are added to the front and the new sequence groups
         # are added to the back.
-        while self.waiting:
-            seq = self.waiting[0]
+        while self.waiting.qsize() > 0:
+            seq_wrapped = self.waiting.queue[0]
+            seq = seq_wrapped.seq
+
             # This is required to handle benchmarking where
             # we set request arrival time ahead of time
             if seq.arrival_time > now:
@@ -68,8 +77,9 @@ class VLLMScheduler(BaseScheduler):
 
             if len(self.running) + 1 > self.scheduler_config.max_num_seqs:
                 break
-
-            seq = self.waiting.pop(0)
+            
+            seq_wrapped = self.waiting.get()
+            seq = seq_wrapped.seq
             self._allocate(seq)
             num_batched_tokens += num_prompt_tokens
             scheduled_seq_metadata_list.append(
diff --git a/sarathi/core/sequence_manager/base_sequence_manager.py b/sarathi/core/sequence_manager/base_sequence_manager.py
index 87e7887..5c1495d 100644
--- a/sarathi/core/sequence_manager/base_sequence_manager.py
+++ b/sarathi/core/sequence_manager/base_sequence_manager.py
@@ -20,6 +20,9 @@ class BaseSequenceManager(ABC):
     def __init__(self, config: SystemConfig):
         self.seq_map: Dict[str, Sequence] = {}
 
+    def set_seq_map(self, seq_map : Dict[str, Sequence]):
+        self.seq_map = seq_map
+
     @synchronized
     def add_seq(self, seq: Sequence) -> None:
         assert seq.seq_id not in self.seq_map
diff --git a/sarathi/engine/base_llm_engine.py b/sarathi/engine/base_llm_engine.py
index c5b046a..7d9eab7 100644
--- a/sarathi/engine/base_llm_engine.py
+++ b/sarathi/engine/base_llm_engine.py
@@ -3,6 +3,8 @@ import math
 import time
 from functools import partial
 from typing import Any, Dict, List, Optional, Tuple
+from threading import Thread, Event
+from queue import Queue, PriorityQueue
 
 import zmq
 
@@ -22,11 +24,12 @@ from sarathi.metrics.cpu_timer import CpuTimer
 from sarathi.metrics.metrics_store import MetricsStore
 from sarathi.transformers_utils.tokenizer import get_tokenizer
 from sarathi.utils import Counter, get_ip, unset_cuda_visible_devices
-from sarathi.utils.threading_utils import synchronized
+from sarathi.utils.threading_utils import synchronized, exit_on_error
 
 logger = init_logger(__name__)
 
 _MAX_WORKER_CONCURRENCY = 1
+SCHEDULER_LOOP_DELAY = 0.01
 
 ModelParallelRank = Tuple[int, int]
 
@@ -48,6 +51,9 @@ class BaseLLMEngine:
     def __init__(
         self,
         config: SystemConfig,
+        seq_counter: Optional[Counter] = None,
+        seq_waiting_queue: Optional[PriorityQueue] = None,
+        global_output_queue: Optional[Queue] = None,
     ) -> None:
         logger.info(
             "Initializing an LLM engine with config: "
@@ -59,6 +65,7 @@ class BaseLLMEngine:
         )
         # TODO(woosuk): Print more configs in debug mode.
 
+        self.has_started_execution_loops = False
         self.config = config
         self._verify_args()
 
@@ -69,7 +76,11 @@ class BaseLLMEngine:
         )
 
         self.seq_manager = EngineSequenceManager(self.tokenizer, config)
-        self.seq_counter = Counter()
+
+        if seq_counter is None:
+            self.seq_counter = Counter()
+        else:
+            self.seq_counter = seq_counter
 
         self.metrics_store = MetricsStore.get_or_create_instance(
             config.replica_config,
@@ -96,6 +107,8 @@ class BaseLLMEngine:
 
         self.mark_initial_memory_profiling_done()
 
+        self.scheduler_queue = PriorityQueue()
+
         # Create the scheduler.
         self.scheduler = SchedulerRegistry.get(
             config.scheduler_config.get_type(),
@@ -103,14 +116,34 @@ class BaseLLMEngine:
             config.scheduler_config,
             config.cache_config,
             config.parallel_config,
+            seq_waiting_queue,
+            self.seq_manager,
+            self.metrics_store,
         )
 
+        self.waiting_queue = seq_waiting_queue
+
         self._scheduler_timer = CpuTimer(CpuOperationMetrics.SCHEDULE)
         self._process_model_outputs_timer = CpuTimer(
             CpuOperationMetrics.PROCESS_MODEL_OUTPUTS
         )
 
-        self.new_seqs: List[Sequence] = []
+
+        if global_output_queue is not None:
+            # We're dealing with a multi-replica setup here
+            # We need to share outputs between the replica-level LLM engine and the MultiReplicaLLMEngine class
+            # So we use a shared global output queue to share outputs between the two classes
+            self.global_output_queue: Queue = global_output_queue
+            # We will need to start a output thread similar to the one in pipeline_parallel_llm_engine.py
+            # to collect outputs from the global output queue and feed into the master output queue
+            self.schedule_event = Event()
+            self.pull_event = Event()
+            self.scheduler_timer_thread = Thread(target=self._scheduler_timer_loop, daemon=True)
+            self.schedule_thread = Thread(target=self._schedule_loop, daemon=True)
+        else:
+            # Just a dummy queue to satisfy the requirements
+            self.global_output_queue = Queue()
+
 
         self._run_workers("wait_till_ready")
 
@@ -280,6 +313,7 @@ class BaseLLMEngine:
         all_request_outputs = self.seq_manager.generate_request_outputs(
             ignored_seqs, seq_metadata_list
         )
+        self.global_output_queue.put(all_request_outputs)
         return all_request_outputs
 
     def get_model_config(self) -> ModelConfig:
@@ -341,17 +375,22 @@ class BaseLLMEngine:
         self.scheduler.add_seq(seq)
         self.metrics_store.on_request_arrival(seq)
 
+    def add_seq(self, seq : Sequence):
+        self.seq_manager.add_seq(seq)
+
+        self._append_new_seq(copy.deepcopy(seq))
+        self.scheduler.add_seq(seq)
+        self.metrics_store.on_request_arrival(seq)
+
     @synchronized
     def _append_new_seq(self, seq: Sequence) -> None:
-        self.new_seqs.append(seq)
+        self.scheduler.add_to_new_seqs(seq)
 
     @synchronized
     def _get_new_seqs(
         self,
     ) -> List[Sequence]:
-        new_seqs = self.new_seqs
-        self.new_seqs = []
-        return new_seqs
+        return self.scheduler.get_new_seqs()
 
     def get_num_unfinished_requests(self) -> int:
         """Gets the number of unfinished requests."""
@@ -360,7 +399,7 @@ class BaseLLMEngine:
     def has_unfinished_requests(self) -> bool:
         """Returns True if there are unfinished requests."""
         return self.scheduler.has_unfinished_seqs()
-
+ 
     def step(self) -> List[RequestOutput]:
         """Performs one decoding iteration and returns newly generated results.
 
@@ -432,6 +471,29 @@ class BaseLLMEngine:
         for other_output in all_outputs[1:]:
             assert output == other_output
         return output
+    
+    def start_execution_loops(self) -> None:
+        if self.has_started_execution_loops:
+            return
+        self.has_started_execution_loops = True
+        self.schedule_event.set()
+        self.schedule_thread.start()
+        self.scheduler_timer_thread.start()
+    
+    @exit_on_error
+    def _scheduler_timer_loop(self) -> None:
+        while True:
+            time.sleep(SCHEDULER_LOOP_DELAY)
+            self.schedule_event.set()
+
+    @exit_on_error
+    def _schedule_loop(self) -> None:
+        while True:
+            self.schedule_event.wait()
+            self.schedule_event.clear()
+
+            if self.has_unfinished_requests():
+                self.step()
 
     def _run_worker(
         self,
diff --git a/sarathi/engine/llm_engine.py b/sarathi/engine/llm_engine.py
index 5687c1b..fed077a 100644
--- a/sarathi/engine/llm_engine.py
+++ b/sarathi/engine/llm_engine.py
@@ -1,17 +1,34 @@
+from typing import Optional, Any, List, Dict
+from queue import Queue, PriorityQueue
+from collections import deque
+
+from sarathi.core.datatypes.sequence import Sequence
 from sarathi.config import SystemConfig
 from sarathi.engine.base_llm_engine import BaseLLMEngine
 from sarathi.engine.pipeline_parallel_llm_engine import PipelineParallelLLMEngine
 
+from sarathi.utils import Counter
+
 
 class LLMEngine:
 
     @classmethod
-    def from_system_config(cls, config: SystemConfig) -> "LLMEngine":
+    def from_system_config(
+        cls,
+        config: SystemConfig,
+        sequence_counter: Optional[Counter] = None,
+        sequence_waiting_list: Optional[PriorityQueue] = None,
+        global_output_queue: Optional[Queue] = None,
+        seq_map : Optional[Dict[str, Sequence]] = None,
+        new_seq_global: Optional[List[Sequence]] = None,
+    ) -> "LLMEngine":
         """Creates an LLM engine from the engine arguments."""
         # Create the engine configs.
         if config.parallel_config.pipeline_parallel_size > 1:
             engine = PipelineParallelLLMEngine(config)
         else:
-            engine = BaseLLMEngine(config)
+            engine = BaseLLMEngine(
+                config, sequence_counter, sequence_waiting_list, global_output_queue
+            )
 
         return engine
diff --git a/sarathi/engine/multi_replica_llm_engine.py b/sarathi/engine/multi_replica_llm_engine.py
new file mode 100644
index 0000000..7f40c5e
--- /dev/null
+++ b/sarathi/engine/multi_replica_llm_engine.py
@@ -0,0 +1,387 @@
+import copy
+import math
+import time
+from functools import partial
+from typing import Any, Dict, List, Optional, Tuple
+from threading import Thread, Event
+from queue import Queue, PriorityQueue
+
+import ray
+import zmq
+
+import os
+from sarathi import LLMEngine
+from sarathi.config import ModelConfig, SystemConfig
+from sarathi.core.datatypes.comm_info import CommInfo
+from sarathi.core.datatypes.request_output import RequestOutput
+from sarathi.core.datatypes.sampling_params import SamplingParams
+from sarathi.core.datatypes.scheduler_output import SchedulerOutputs
+from sarathi.core.scheduler.global_scheduler import (
+    PullScheduler,
+    RoundRobinScheduler,
+)
+from sarathi.core.datatypes.sequence import SamplerOutputs, Sequence, SequenceMetadata
+from sarathi.core.datatypes.step_inputs import StepInputs
+from sarathi.core.scheduler.scheduler_registry import SchedulerRegistry
+from sarathi.core.sequence_manager.engine_sequence_manager import EngineSequenceManager
+from sarathi.engine.ray_utils import RayWorker, initialize_cluster, ray
+from sarathi.logger import init_logger
+from sarathi.config import ReplicaConfig
+from sarathi.types import ReplicaResourceMapping, ResourceMapping
+from sarathi.metrics.constants import CpuOperationMetrics
+from sarathi.metrics.cpu_timer import CpuTimer
+from sarathi.metrics.metrics_store import MetricsStore
+from sarathi.transformers_utils.tokenizer import get_tokenizer
+from sarathi.utils import Counter, get_ip, unset_cuda_visible_devices
+from sarathi.utils.threading_utils import synchronized, exit_on_error
+
+logger = init_logger(__name__)
+
+SCHEDULER_LOOP_DELAY = 0.01
+
+
+class MultiReplicaLLMEngine:
+    """A Multi-replica LLM engine that receives requests and generates texts across multiple replicas of the same model.
+
+    Args:
+        config; System Config: The system configuration for the engine.
+        replica_resource_mapping: The resource mapping for the replicas. (Basically the ip and gpu id for each replica)
+    """
+
+    def __init__(
+        self,
+        config: SystemConfig,
+        replica_resource_mapping: Optional[ReplicaResourceMapping] = None,
+    ) -> None:
+        logger.info(
+            "Initializing a Multi Replica LLM engine with config: "
+            f"model={config.model_config.model!r}, "
+            f"dtype={config.model_config.dtype}, "
+            f"tensor_parallel_size={config.parallel_config.tensor_parallel_size}, "
+            f"pipeline_parallel_size={config.parallel_config.pipeline_parallel_size}, "
+            f"seed={config.model_config.seed})"
+        )
+
+        logger.info(f"Replica resource mapping: {replica_resource_mapping}")
+
+        self.config = config
+        self._verify_args()
+
+        ray.init(ignore_reinit_error=True)
+
+        if replica_resource_mapping is None:
+            self._validate_cluster_resources()
+            self.replica_resource_mapping = self._get_replica_resource_mapping()
+        else:
+            self.replica_resource_mapping = replica_resource_mapping
+
+        self.sequence_counter = Counter()
+
+        self.global_scheduler_type = self.config.global_scheduler_config.scheduler_type
+
+        self.aggregate_metric_store = self._create_aggregate_metric_store()
+
+        self.global_output_queue: Queue = Queue()
+
+        self.request_processing_queue = Queue()
+
+        self.global_scheduler = self._get_global_scheduler()
+        self.global_scheduler.init_queue()
+
+        self._init_replica_engines()
+
+        # start daemon thread for the assign_thread function
+        self._setup_tokenizer_threads(num_tokenizers=self.config.num_replicas*2)
+
+    def _setup_tokenizer_threads(self, num_tokenizers):
+        self.tokenizer_pool = []
+        for _ in range(num_tokenizers):
+            tokenizer = get_tokenizer(
+                self.config.model_config.model,
+                trust_remote_code=self.config.model_config.trust_remote_code,
+                revision=self.config.model_config.revision,
+            )
+            self.tokenizer_pool.append(tokenizer)
+        
+        # Spawn the daemon threads that keep pulling from the self.request_processing_queue
+        self.tokenizer_queue_threads = []
+        for i in range(num_tokenizers):
+            tokenizer_queue_thread = Thread(target=self._tokenizer_queue_worker, args=(self.tokenizer_pool[i],), daemon=True)
+            tokenizer_queue_thread.start()
+    
+    @exit_on_error
+    def _tokenizer_queue_worker(self, tokenizer):
+        while True:
+            prompt,sampling_params,prompt_token_ids,arrival_time,seq_id = self.request_processing_queue.get()
+            self._tokenize_and_generate_sequence(tokenizer, prompt,sampling_params,prompt_token_ids,arrival_time,seq_id)
+
+
+    def _tokenize_and_generate_sequence(
+        self,
+        tokenizer,
+        prompt: Optional[str],
+        sampling_params: SamplingParams,
+        prompt_token_ids: Optional[List[int]] = None,
+        arrival_time: Optional[float] = None,
+        seq_id: Optional[str] = None,
+    ):
+        try:
+            if arrival_time is None:
+                arrival_time = time.monotonic()
+
+            if prompt_token_ids is None:
+                prompt_token_ids = tokenizer.encode(prompt)
+
+            if not seq_id:
+                seq_id = str(next(self.sequence_counter))
+
+            block_size = self.config.cache_config.block_size
+            eos_token_id = tokenizer.eos_token_id
+
+            seq = Sequence(
+                seq_id,
+                prompt,
+                prompt_token_ids,
+                block_size,
+                eos_token_id,
+                arrival_time,
+                sampling_params,
+            )
+
+            self.global_scheduler.assign_seq_replica(seq)
+        except Exception as e:
+            logger.error(f"Error in tokenizing and generating sequence: {e}")
+
+
+    def _init_replica_engines(self):
+        self.replica_llm_engine_mapping = {}
+        for replica_id in range(self.config.num_replicas):
+
+            replica_config = ReplicaConfig(
+                replica_id=replica_id,
+                output_dir=self.config.replica_config.output_dir,
+                resource_mapping=self.replica_resource_mapping[replica_id],
+            )
+            os.makedirs(replica_config.output_dir, exist_ok=True)
+
+            system_config = copy.deepcopy(self.config)
+            system_config.replica_config = replica_config
+            # Spawn a replica LLM engine for each replica in their own thread
+            # self._spawn_replica_engine(replica_id, system_config, self.sequence_counter, self.global_scheduler.get_replica_queue(replica_id), self.global_output_queue, self.global_scheduler.get_seq_map())
+            replica_thread = Thread(
+                target=self._spawn_replica_engine,
+                args=(
+                    replica_id,
+                    system_config,
+                    self.sequence_counter,
+                    self.global_scheduler.get_replica_queue(replica_id),
+                    self.global_output_queue,
+                    self.global_scheduler.get_seq_map(),
+                    self.global_scheduler.get_new_seq_list(),
+                ),
+            )
+            replica_thread.start()
+
+            # Wait for the replica thread to finish
+            replica_thread.join()
+
+    def _get_global_scheduler(self):
+        if self.global_scheduler_type == "pull":
+            return PullScheduler(
+                self.config,
+                self.config.num_replicas,
+                self.sequence_counter            
+            )
+        elif self.global_scheduler_type == "round_robin":
+            return RoundRobinScheduler(
+                self.config,
+                self.config.num_replicas,
+                self.sequence_counter
+            )
+        else:
+            raise ValueError(
+                f"Unknown global scheduler type: {self.global_scheduler_type}"
+            )
+
+    def _spawn_replica_engine(
+        self,
+        replica_id,
+        system_config,
+        sequence_counter,
+        replica_llm_engine_queue,
+        global_output_queue,
+        seq_map,
+        new_seq_global,
+    ):
+        replica_llm_engine = LLMEngine.from_system_config(
+            system_config,
+            sequence_counter,
+            replica_llm_engine_queue,
+            global_output_queue,
+        )
+
+        self.global_scheduler.set_replica_llm_engine(replica_id, replica_llm_engine)
+
+        self.replica_llm_engine_mapping[replica_id] = replica_llm_engine
+
+    def _validate_cluster_resources(self):
+
+        num_replicas = self.config.num_replicas
+        num_gpus_required = num_replicas * self.config.parallel_config.world_size
+
+        available_resources = ray.available_resources()
+
+        assert (
+            available_resources["GPU"] >= num_gpus_required
+        ), f"Insufficient GPUs. Required: {num_gpus_required}, Available: {available_resources['GPU']}"
+
+    def _get_replica_resource_mapping(self) -> ReplicaResourceMapping:
+
+        cluster_resources_keys = list(ray.available_resources().keys())
+        num_gpus = ray.available_resources()["GPU"]
+        ip_addresses = [
+            x
+            for x in cluster_resources_keys
+            if x.startswith("node:") and x != "node:__internal_head__"
+        ]
+
+        runner_ip = f"node:{get_ip()}"
+
+        ip_addresses.remove(runner_ip)
+        ip_addresses.insert(0, runner_ip)
+
+        num_nodes = len(ip_addresses)
+        assert num_nodes > 0, "No nodes found in the cluster"
+        assert num_gpus > 0, "No GPUs found in the cluster"
+        assert (
+            num_gpus % num_nodes == 0
+        ), f"Number of GPUs ({num_gpus}) is not a multiple of number of nodes ({num_nodes})"
+        num_gpus_per_node = int(num_gpus // num_nodes)
+        num_replicas = self.config.num_replicas
+        num_gpus_per_replica = self.config.parallel_config.world_size
+
+        assert (
+            num_gpus >= num_replicas * num_gpus_per_replica
+        ), f"Insufficient GPUs. Required: {num_replicas * num_gpus_per_replica}, Available: {num_gpus}"
+
+        replica_resource_mapping = []
+
+        available_gpus = []
+        for ip_address in ip_addresses:
+            for gpu_id in reversed(range(num_gpus_per_node)):
+                available_gpus.append((ip_address, gpu_id))
+
+        for _ in range(num_replicas):
+            resource_mapping = []
+            for _ in range(num_gpus_per_replica):
+                resource_mapping.append(available_gpus.pop(0))
+            replica_resource_mapping.append(resource_mapping)
+
+        logger.info(f"Replica resource mapping: {replica_resource_mapping}")
+
+        return replica_resource_mapping
+
+    def _validate_parallel_config(self) -> None:
+        assert self.config.parallel_config.pipeline_parallel_size == 1
+
+    def _verify_args(self) -> None:
+        self._validate_parallel_config()
+        self.config.model_config.verify_with_parallel_config(
+            self.config.parallel_config
+        )
+
+    def add_request(
+        self,
+        prompt: Optional[str],
+        sampling_params: SamplingParams,
+        prompt_token_ids: Optional[List[int]] = None,
+        arrival_time: Optional[float] = None,
+        seq_id: Optional[str] = None,
+    ) -> None:
+        """
+        Add a new request to a replica LLM engine
+        The entire Request allocation is abstracted out to the Global Scheduler
+        """
+        self.request_processing_queue.put(
+            (
+                prompt,
+                sampling_params,
+                prompt_token_ids,
+                arrival_time,
+                seq_id,
+            )
+        )
+
+    def get_num_unfinished_requests(self) -> int:
+        """Gets the number of unfinished requests."""
+        return self.global_scheduler.get_num_unfinished_requests()
+
+    def has_unfinished_requests(self) -> bool:
+        """Returns True if there are unfinished requests."""
+        return self.global_scheduler.get_num_unfinished_requests()
+
+    def get_metrics_store(self, replica_id: int) -> MetricsStore:
+        return self.replica_llm_engine_mapping[replica_id].get_metrics_store()
+
+    def pull_worker_metrics(self, replica_id: int) -> None:
+        self.replica_llm_engine_mapping[replica_id].pull_worker_metrics()
+
+    def _reset_replice_metrics(self, replica_id: int) -> None:
+        self.replica_llm_engine_mapping[replica_id].reset_metrics()
+
+    def reset_metrics(self) -> None:
+        for replica_id in range(self.config.num_replicas):
+            self._reset_replice_metrics(replica_id)
+
+    def start_profiling(self):
+        for replica_id in range(self.config.num_replicas):
+            self.replica_llm_engine_mapping[replica_id].start_profiling()
+
+    def stop_profiling(self):
+        for replica_id in range(self.config.num_replicas):
+            self.replica_llm_engine_mapping[replica_id].stop_profiling()
+
+    def _pull_worker_metrics(self, replica_id: int) -> None:
+        self.replica_llm_engine_mapping[replica_id].pull_worker_metrics()
+
+    def pull_worker_metrics(self) -> None:
+        for replica_id in range(self.config.num_replicas):
+            self._pull_worker_metrics(replica_id)
+
+    def _get_metric_store(self, replica_id: int) -> MetricsStore:
+        return self.replica_llm_engine_mapping[replica_id].get_metric_store()
+
+    def get_metric_store(self) -> MetricsStore:
+        for replica_id in range(self.config.num_replicas):
+            self.aggregate_metric_store.merge(self._get_metric_store(replica_id))
+        return self.aggregate_metric_store
+
+    def _create_aggregate_metric_store(self):
+        replica_config = ReplicaConfig(
+            replica_id=0,  # dummy replica id
+            output_dir=self.config.replica_config.output_dir,
+        )
+        metrics_store = MetricsStore.get_or_create_instance(
+            replica_config,
+            self.config.model_config,
+            self.config.metrics_config,
+        )
+
+        metrics_store.mark_initial_memory_profiling_done()
+
+        return metrics_store
+    
+    def start_engine_execution(self):
+        for replica_id in range(self.config.num_replicas):
+            self.replica_llm_engine_mapping[replica_id].start_execution_loops()
+
+    def step(self) -> List[RequestOutput]:
+        """
+        The step function drains the global output queue and returns the list of request outputs.
+        The replica LLM engines push the request outputs to the global output queue.
+        The replica level LLM engines continously keep calling the step function
+        """
+        output_list = []
+        while not self.global_output_queue.empty():
+            output_list.extend(self.global_output_queue.get())
+        return output_list
diff --git a/sarathi/types.py b/sarathi/types.py
index fbf2cbb..41630aa 100644
--- a/sarathi/types.py
+++ b/sarathi/types.py
@@ -37,3 +37,7 @@ class RequestLengthGeneratorType(Enum):
 class AttentionBackend(Enum):
     FLASHINFER = "FLASHINFER"
     NO_OP = "NO_OP"
+
+class GlobalSchedulerType(Enum):
+    PULL = "PULL"
+    ROUND_ROBIN = "ROUND_ROBIN"
\ No newline at end of file
diff --git a/sarathi/utils/__init__.py b/sarathi/utils/__init__.py
index 83074c6..1f67ffe 100644
--- a/sarathi/utils/__init__.py
+++ b/sarathi/utils/__init__.py
@@ -9,6 +9,8 @@ from typing import AsyncIterator, List, Tuple, TypeVar, Union
 import psutil
 import torch
 
+from sarathi.utils.threading_utils import synchronized
+
 T = TypeVar("T")
 
 
@@ -71,7 +73,6 @@ def is_port_in_use(port: int) -> bool:
     with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
         return s.connect_ex(("localhost", port)) == 0
 
-
 def get_random_port() -> int:
     port = None
     while not port or is_port_in_use(port):