Fixing some failures due to merge conflicts.

Author: aman2930

jetstream/core/lora/adapter_tensorstore.py

@@ -87,11 +87,12 @@ class AdapterTensorStore:
   model to server multiple different LoRA adapters in a single batch.
 
   Args:
+    engine: Engine corresponding to the adapter tensorstore
+    adapters_dir_path: GCS path storing all the adapters
     hbm_memory_budget (int): The maximum amount of HBM (in bytes) to use for
         storing LoRA adapter weights.
     cpu_memory_budget (int): The maximum amount of CPU RAM (in bytes) to use
         for storing LoRA adapter weights.
-    total_slots: Number of generate slots. This is also equals to max_concurrent_decodes.
   """
 
   def __init__(self,
@@ -111,7 +112,7 @@ def __init__(self,
     self.current_cpu_usage: int = 0
     self.running_requests: int = 0    # Number of async tasks which are in "loading" state
     self.decoding_adapters_cache: Dict[str, Any] = {}
-    self.total_slots = total_slots
+    self.total_slots = engine.max_concurrent_decodes
     self.lock = asyncio.Lock()        # Use an asyncio Lock for thread safety
 
 

jetstream/core/orchestrator.py

@@ -304,15 +304,6 @@ def __init__(
     self._metrics_collector = metrics_collector
     self._multi_sampling = multi_sampling
 
-    total_slots = 0
-    for engine in self._generate_engines:
-      total_slots += engine.max_concurrent_decodes
-
-    self._adapter_tensorstore = adapter_tensorstore.AdapterTensorStore(
-        hbm_memory_budget=(20 * (1024 ** 3)),       # 20 GB HBM
-        cpu_memory_budget=(100 * (1024 ** 3)),      # 100 GB RAM
-        total_slots=total_slots)
-
     # Stages 1-4 represent the life cycle of a request.
     # Stage 1
     # At first, a request is placed here in order to get prefilled.
@@ -562,7 +553,8 @@ def _export_lora_request_info(self):
     if self._metrics_collector:
       for idx, engine in enumerate(self._generate_engines):
         max_loras += engine.max_concurrent_decodes
-        if self._generate_adapterstore and idx < len(self._generate_adapterstore):
+        if (self._generate_adapterstore and
+            idx < len(self._generate_adapterstore)):
           adapters_list_str += asyncio.run(
               self._generate_adapterstore[idx].get_hbm_loaded_adapters())
 
@@ -908,6 +900,10 @@ def _insert_if_possible(
     # Check if there are any free my_slots. We don't want to block here since
     # we can still generate if we can't insert. We do this in a while loop to
     # insert as many sequences as possible.
+    adapter_tensorstore = None
+    if self._generate_adapterstore:
+      adapter_tensorstore = self._generate_adapterstore[idx]
+
     while True:
       my_slots_size = my_slots.qsize()
 
@@ -979,7 +975,9 @@ def _insert_if_possible(
           #request_id=new_request.request_id,
       )
 
-      self._adapter_tensorstore.insert_adapter_in_cache(new_request.adapter_id, slot)
+      if adapter_tensorstore:
+        adapter_tensorstore.insert_adapter_in_cache(
+            new_request.adapter_id, slot)
 
       ThreadDebugLog(
           thread_name,
@@ -1120,6 +1118,10 @@ def _generate_thread(self, idx: int):
     my_generate_backlog = self._generate_backlogs[idx]
     my_detokenize_backlog = self._detokenize_backlogs[idx]
 
+    adapter_tensorstore = None
+    if self._generate_adapterstore and idx < len(self._generate_adapterstore):
+      adapter_tensorstore = self._generate_adapterstore[idx]
+
     # Keep track of what step tokens were generated at.
     generate_timestep = 0
     # State to store things like running kv cache in.
@@ -1185,9 +1187,14 @@ def _generate_thread(self, idx: int):
           my_slots.qsize() < max_concurrent_decodes
       ), "At this point we must have some requests inserted into the slots."
 
+      decoding_adapters_cache = None
+
+      if adapter_tensorstore:
+        decoding_adapters_cache = adapter_tensorstore.decoding_adapters_cache
+
       # Now we actually take a generate step on requests in the slots.
       decode_state, sampled_tokens = generate_engine.generate(
-          generate_params, decode_state, self._adapter_tensorstore.decoding_adapters_cache,
+          generate_params, decode_state, decoding_adapters_cache
       )
       sampled_tokens.copy_to_host_async()
       # Respond to detokenization backpressure.