- Created separate adapter_tensorstore for each engine.

- Implemented unapply lora from base_params
- Fixed some comments from the PR

Author: aman2930

jetstream/core/lora/adapter_tensorstore.py

@@ -26,6 +26,7 @@
 import functools
 from typing import Dict, Optional, Any
 import numpy as np
+from jetstream.engine import engine_api
 
 
 def _get_size_of_pytree(params):
@@ -82,8 +83,14 @@ class AdapterTensorStore:
   """
 
 
-  def __init__(self, hbm_memory_budget: int, cpu_memory_budget: int):
+  def __init__(self,
+      engine: engine_api.Engine,
+      adapters_dir_path: str,
+      hbm_memory_budget: int,
+      cpu_memory_budget: int):
     """Initializes the AdapterTensorStore."""
+    self.engine = engine                                      # Possibly MaxEngine object
+    self.adapters_dir_path = adapters_dir_path.rstrip("/")    # All Adapters path without trailing `/`
     self.hbm_memory_budget = hbm_memory_budget
     self.cpu_memory_budget = cpu_memory_budget
     self.adapter_registry: Dict[str, AdapterMetadata] = {}    # All known adapters
@@ -95,26 +102,49 @@ def __init__(self, hbm_memory_budget: int, cpu_memory_budget: int):
     self.lock = asyncio.Lock()        # Use an asyncio Lock for thread safety
 
 
-  def register_adapter(self, adapter_id: str, adapter_path: str, config: Dict[str, Any]):
+  def register_adapter(self,
+      adapter_id: str,
+      adapter_path: str = None,
+      adapter_config: Dict[str, Any] = None):
     """Registers a new LoRA adatper."""
     """
-    Registers a LoRA adapter with the TensorStore.  This does *not* load
-    the adapter; it simply adds metadata about the adapter to the registry.
+    Registers a LoRA adapter with the TensorStore.  This also loads the adapter; 
+    IF called without adapter_config. Because in this case, it needs
+    to get adapter_config from the engine's load_single_adapter() call, which
+    also provides the adapter_params. So in that case it is beneficial to load
+    the adapter to HBM. This call path is expected only from the direct inference
+    request.
+    OTHERWISE, it simply adds metadata about the adapter to the registry.
 
     Args:
       adapter_id (str): A unique identifier for the adapter.
       adapter_path (str): The path to the adapter weights (file or directory).
-      config (dict): Config of the loRA adapter.
+      adapter_config (dict): Config of the loRA adapter.
 
     Raises:
       ValueError: If an adapter with the same ID is already registered.
     """
     if adapter_id in self.adapter_registry:
-      raise ValueError(f"Adapter with ID '{adapter_id}' already registered.")
+      logging.warning(f"Adapter with ID '{adapter_id}' already registered.")
+      return
+
+    if adapter_path is None:
+      adapter_path = f"{self.adapters_dir_path}/{adapter_id}"
+
+    adapter_params = None
+    if adapter_config is None:
+      adapter_params, adapter_config = self.engine.load_single_adapter(adapter_path)
+
+    if adapter_config is None:
+      raise ValueError(f"Failed to read adapter_config from {adapter_path}")
+
     self.adapter_registry[adapter_id] = AdapterMetadata(
         adapter_id=adapter_id,
         adapter_path=adapter_path,
-        config=config)
+        config=adapter_config)
+
+    if adapter_params is not None:
+      asyncio.run(self.load_adapter(adapter_id, adapter_params, True))
 
 
   async def _transfer_to_hbm(self, adapter_id: str):
@@ -254,7 +284,10 @@ async def load_adapter(
 
     try:
       if adapter_weights is None:
-        raise ValueError("Adapter weights for adapter_id={adapter_id} is None.")
+        adapter_weights, adapter_config = self.engine.load_single_adapter(adapter_path)
+
+      if adapter_weights is None:
+        raise ValueError("Failed to load adapter_weights from {adapter_path}.")
 
       async with self.lock:       # Critical section for memory management
         adapter_weights_as_jnp_array = _as_jnp_array(adapter_weights)
@@ -303,21 +336,36 @@ async def load_adapter(
         self.running_requests -= 1
 
 
-  def get_lora_config(self, adapter_id):
+  def get_lora_config(self, adapter_id: str, load_if_not_loaded: bool = False):
     """Getter for the LoRA adapter config."""
     metadata = self.adapter_registry.get(adapter_id)
+
+    if load_if_not_loaded and metadata is None:
+      self.register_adapter(adapter_id)
+      metadata = self.adapter_registry.get(adapter_id)
+
+    if metadata is None:
+      raise ValueError(f"LoRA adapter with id={adapter_id} is not loaded.")
+
     return metadata.config
 
 
-  def get_lora_weights(self, adapter_id, to_hbm: bool = True):
+  def get_lora_weights(self,
+      adapter_id,
+      to_hbm: bool = True,
+      load_if_not_loaded: bool = False):
     """Retrieves the unified LoRA parameters for the given adapter IDs.
        Handles HBM/CPU placement.
     """
 
     metadata = self.adapter_registry.get(adapter_id)
 
+    if load_if_not_loaded and metadata is None:
+      self.register_adapter(adapter_id)
+      metadata = self.adapter_registry.get(adapter_id)
+
     if metadata is None:
-      raise ValueError(f"Adapter with ID '{adapter_id}' not registered.")
+      raise ValueError(f"LoRA adapter with id={adapter_id} is not loaded.")
 
     if metadata.status != "loaded_hbm" and metadata.status != "loaded_cpu":
       asyncio.run(self.load_adapter(adapter_id, None, to_hbm))    # Start loading (async)

jetstream/core/lora/multi_lora_inference_api.py

@@ -17,17 +17,11 @@
 
 import logging
 import grpc
-import time
-import uuid
 
-from typing import Any, AsyncIterator, Optional, Tuple, cast
+from typing import Optional
 from jetstream.core import orchestrator
-from jetstream.core.lora import adapter_tensorstore
 from jetstream.core.proto import multi_lora_decoding_pb2_grpc
 from jetstream.core.proto import multi_lora_decoding_pb2
-from jetstream.core.utils import async_multifuture
-from jetstream.core.utils.return_sample import ReturnSample
-from jetstream.engine import engine_api, tokenizer_api, token_utils
 
 
 class MultiLoraManager(multi_lora_decoding_pb2_grpc.v1Servicer):
@@ -105,155 +99,3 @@ def unload_lora_adapter(
       logging.info(f"Loading of adapter_id={request.adapter_id} failed with error: {str(e)}")
       return multi_lora_decoding_pb2.UnloadAdapterResponse(success=False, error_message=str(e))
 
-
-  def _get_prefill_content(
-      self, request: multi_lora_decoding_pb2.CompletionRequest
-  ) -> Tuple[str | list[int], bool]:
-    which_content = request.WhichOneof("content")
-    content = getattr(request, which_content)
-    if which_content == "text_content":
-      return cast(multi_lora_decoding_pb2.CompletionRequest.TextContent, content).text, False
-    else:
-      return (
-          list(
-              cast(multi_lora_decoding_pb2.CompletionRequest.TokenContent, content).token_ids
-          ),
-          True,
-      )
-
-  def process_client_side_tokenization_response(self, response: Any):
-    samples = []
-    for sample in response:
-      samples.append(
-          multi_lora_decoding_pb2.CompletionResponse.StreamContent.Sample(
-              token_ids=sample.token_ids,
-          )
-      )
-    return multi_lora_decoding_pb2.CompletionResponse(
-        stream_content=multi_lora_decoding_pb2.CompletionResponse.StreamContent(
-            samples=samples
-        )
-    )
-
-  def should_buffer_response(self, response: Any) -> bool:
-    for item in response:
-      if item.text and token_utils.is_byte_token(item.text[-1]):
-        # If any sample ends in bytes, this means we might still need to
-        # decode more bytes to compose the string.
-        return True
-
-  def process_server_side_tokenization_response(
-      self, response: Any, buffered_response_list
-  ):
-    # Flush the buffered responses to each sample of current response.
-    current_response_with_flushed_buffer = list(
-        zip(*buffered_response_list, response)
-    )
-    # Empty buffer: [[s0_cur], [s1_cur], ...]
-    # Has buffer:
-    # [[s0_b0, s0_b1, ..., s0_cur], [s1_b0, s1_b1, ..., s1_cur], ...]
-    current_response_with_flushed_buffer = cast(
-        list[list[ReturnSample]], current_response_with_flushed_buffer
-    )
-    # Form correct sample(s) and return as StreamContent for this iteration.
-    samples = []
-    for sample in current_response_with_flushed_buffer:
-      text = []
-      token_ids = []
-      for resp in sample:
-        text.extend(resp.text)
-        token_ids.extend(resp.token_ids)
-      samples.append(
-          multi_lora_decoding_pb2.CompletionResponse.StreamContent.Sample(
-              text=token_utils.text_tokens_to_str(text),
-              token_ids=token_ids,
-          )
-      )
-    return multi_lora_decoding_pb2.CompletionResponse(
-        stream_content=multi_lora_decoding_pb2.CompletionResponse.StreamContent(
-            samples=samples
-        )
-    )
-
-  async def completions(  # pylint: disable=invalid-overridden-method
-      self,
-      request: multi_lora_decoding_pb2.CompletionRequest,
-      context: Optional[grpc.aio.ServicerContext] = None,
-  ) -> AsyncIterator[multi_lora_decoding_pb2.CompletionResponse]:
-
-    """Decode."""
-    if context is None:
-      logging.warning(
-          "LLM orchestrator is being used in offline test mode, and will not"
-          " respond to gRPC queries - only direct function calls."
-      )
-    is_client_side_tokenization = False
-    return_channel = async_multifuture.AsyncMultifuture()
-    if context:
-      context.add_done_callback(return_channel.cancel)
-
-    prefill_content, is_client_side_tokenization = self._get_prefill_content(
-        request
-    )
-
-    # Wrap request as an ActiveRequest.
-    active_request = orchestrator.ActiveRequest(
-        request_id=uuid.uuid4(),
-        max_tokens=request.max_tokens,
-        prefill_content=prefill_content,
-        is_client_side_tokenization=is_client_side_tokenization,
-        return_channel=return_channel,
-        adapter_id=request.adapter_id,
-        metadata=orchestrator.ActiveRequestMetadata(
-            start_time=request.metadata.start_time,
-            prefill_enqueue_time=time.perf_counter(),
-        ),
-    )
-    # The first stage is being prefilled, all other stages are handled
-    # inside the driver (transfer, generate*N, detokenize).
-    try:
-      self._driver.place_request_on_prefill_queue(active_request)
-    except queue.Full:
-      # Safely abort the gRPC server thread with a retriable error.
-      await _abort_or_raise(
-          context=context,
-          code=grpc.StatusCode.RESOURCE_EXHAUSTED,
-          details=(
-              "The driver prefill queue is full and more requests cannot be"
-              " handled. You may retry this request."
-          ),
-      )
-    logging.info(
-        "Placed request on the prefill queue.",
-    )
-    # When an active request is created a queue is instantiated. New tokens
-    # are placed there during the decoding loop, we pop from that queue by
-    # using the .next method on the active request.
-    # Yielding allows for the response to be a streaming grpc call - which
-    # can be called via iterating over a for loop on the client side.
-    # The DecodeResponse stream should consume all generated tokens in
-    # return_channel when complete signal is received (AsyncMultifuture
-    # promises this).
-    buffered_response_list = []
-    async for response in active_request.return_channel:
-      response = cast(list[ReturnSample], response)
-      if is_client_side_tokenization:
-        # If is_client_side_tokenization, the client should request with token
-        # ids, and the JetStream server will return token ids as response.
-        # The client should take care of tokenization and detokenization.
-        yield self.process_client_side_tokenization_response(response)
-      else:
-        # Buffer response mechanism is used to handle streaming
-        # detokenization with special character (For some edge cases with
-        # SentencePiece tokenizer, it requires to decode a complete sequence
-        # instead of a single token).
-        if self.should_buffer_response(response):
-          buffered_response_list.append(response)
-          continue
-        yield self.process_server_side_tokenization_response(
-            response, buffered_response_list
-        )
-        # Reset buffer after flushed.
-        buffered_response_list = []
-
-