content: import asyncio
  import concurrent.futures
  import os
  import time
  from pathlib import Path
  from typing import List, Optional
  
  import torch
  from grpc import aio
  from grpc_reflection.v1alpha import reflection
  from loguru import logger
  from tqdm import tqdm
  
  from lorax_server.cache import Cache
  from lorax_server.interceptor import ExceptionInterceptor
  from lorax_server.models import Model, get_model
  from lorax_server.pb import generate_pb2, generate_pb2_grpc
  from lorax_server.tracing import UDSOpenTelemetryAioServerInterceptor
  from lorax_server.utils import PBASE, S3, map_pbase_model_id_to_s3
  from lorax_server.utils.adapter import (
      adapter_source_enum_to_string,
      download_adapter,
      enum_string_to_adapter_source,
      is_base_model,
  )
  from lorax_server.utils.punica import LORAX_PUNICA_TRITON_DISABLED, has_sgmv
  from lorax_server.utils.state import set_max_prefill_tokens, set_speculative_tokens
  
  
  class LoraxService(generate_pb2_grpc.LoraxServiceServicer):
      """
      Implementation of the LoraxService gRPC service.
  
      Args:
          model (Model): The model used for inference.
          cache (Cache): The cache used for storing and retrieving batches.
          server_urls (List[str]): List of server URLs for service discovery.
      """
  
      def __init__(self, model: Model, cache: Cache, server_urls: List[str]):
          self.cache = cache
          self.model = model
          self.server_urls = server_urls
          # For some reason, inference_mode does not work well with GLOO which we use on CPU
          if model.device.type == "cuda":
              # Force inference mode for the lifetime of LoraxService
              self._inference_mode_raii_guard = torch._C._InferenceMode(True)
  
      async def Info(self, request, context):
          return self.model.info
  
      async def Health(self, request, context):
          if self.model.device.type == "cuda":
              torch.zeros((2, 2)).cuda()
          return generate_pb2.HealthResponse()
  
      async def ServiceDiscovery(self, request, context):
          return generate_pb2.ServiceDiscoveryResponse(urls=self.server_urls)
  
      async def ClearCache(self, request, context):
          if request.HasField("id"):
              self.cache.delete(request.id)
          else:
              self.cache.clear()
  
          return generate_pb2.ClearCacheResponse()
  
      async def FilterBatch(self, request, context):
          batch = self.cache.pop(request.batch_id)
          if batch is None:
              raise ValueError(f"Batch ID {request.batch_id} not found in cache.")
          filtered_batch = batch.filter(request.request_ids)
          self.cache.set(filtered_batch)
  
          return generate_pb2.FilterBatchResponse(batch=filtered_batch.to_pb())
  
      async def Warmup(self, request: generate_pb2.WarmupRequest, context):
          set_max_prefill_tokens(request.max_prefill_tokens)
  
          batch = self.model.batch_type.from_pb(
              request.batch,
              self.model.tokenizer,
              self.model.tokenizers,
              self.model.processor,
              self.model.model.config,
              self.model.dtype,
              self.model.device,
          )
          max_supported_total_tokens = self.model.warmup(batch, request.max_new_tokens)
  
          return generate_pb2.WarmupResponse(max_supported_total_tokens=max_supported_total_tokens)
  
      async def Prefill(self, request: generate_pb2.PrefillRequest, context):
          batch = self.model.batch_type.from_pb(
              request.batch,
              self.model.tokenizer,
              self.model.tokenizers,
              self.model.processor,
              self.model.model.config,
              self.model.dtype,
              self.model.device,
          )
  
          if self.model.supports_chunking:
              if request.HasField("cached_batch"):
                  cached_batch = self.cache.pop(request.cached_batch.id)
                  if cached_batch is None:
                      raise ValueError(f"Batch ID {request.cached_batch.id} not found in cache.")
                  batch = self.model.batch_type.concatenate([cached_batch, batch])
  
          generations, next_batch = self.model.generate_token(batch)
          self.cache.set(next_batch)
  
          if self.model.profiler:
              self.model.profiler_steps += 1
              if self.model.profiler_steps == 10:
                  self.model.profiler.stop()
                  print(self.model.profiler.key_averages())
  
          return generate_pb2.PrefillResponse(
              generations=[generation.to_pb() for generation in generations],
              batch=next_batch.to_pb() if next_batch else None,
          )
  
      async def Classify(self, request: generate_pb2.ClassifyRequest, context):
          if not self.model.supports_classification:
              raise ValueError("Model does not support classification")
  
          batch = self.model.batch_type.from_pb(
              request.batch,
              self.model.tokenizer,
              self.model.tokenizers,
              self.model.processor,
              self.model.model.config,
              self.model.dtype,
              self.model.device,
          )
          predicated_token_class, confidence_scores = self.model.classify(batch)
          ner_results = self.model.batch_type.to_pb_classify(batch, predicated_token_class, confidence_scores)
          return ner_results
  
      async def Embed(self, request: generate_pb2.EmbedRequest, context):
          if not self.model.supports_embeddings:
              raise ValueError("Model does not support embeddings")
  
          batch = self.model.batch_type.from_pb(
              request.batch,
              self.model.tokenizer,
              self.model.tokenizers,
              self.model.processor,
              self.model.model.config,
              self.model.dtype,
              self.model.device,
          )
          embeddings = self.model.embed(batch)
          embeddings_pb = self.model.batch_type.to_pb_embed(batch, embeddings)
          return embeddings_pb
  
      async def Decode(self, request: generate_pb2.DecodeRequest, context):
          if len(request.batches) == 0:
              raise ValueError("Must provide at least one batch")
  
          batches = []
          for batch_pb in request.batches:
              batch = self.cache.pop(batch_pb.id)
              if batch is None:
                  raise ValueError(f"Batch ID {batch_pb.id} not found in cache.")
              batches.append(batch)
  
          if len(batches) == 0:
              raise ValueError("All batches are empty")
  
          if len(batches) > 1:
              batch = self.model.batch_type.concatenate(batches)
          else:
              batch = batches[0]
  
          generations, next_batch = self.model.generate_token(batch)
          self.cache.set(next_batch)
  
          return generate_pb2.DecodeResponse(
              generations=[generation.to_pb() for generation in generations],
              batch=next_batch.to_pb() if next_batch else None,
          )
  
      async def DownloadAdapter(self, request: generate_pb2.DownloadAdapterRequest, context):
          if (
              len(request.adapter_parameters.adapter_ids) == 1
              and request.adapter_parameters.adapter_ids[0] in self.model.preloaded_adapter_memory_fractions
          ):
              logger.info("Adapter is already preloaded. Skipping.")
              return generate_pb2.DownloadAdapterResponse(
                  downloaded=True,
                  memory_fraction=self.model.preloaded_adapter_memory_fractions[
                      request.adapter_parameters.adapter_ids[0]
                  ],
              )
  
          return download_adapter(request, self.model)
  
      async def LoadAdapter(self, request: generate_pb2.LoadAdapterRequest, context):
          adapter_parameters = request.adapter_parameters
          if is_base_model(adapter_parameters):
              logger.info("No adapter to load for base model. Skipping.")
              return generate_pb2.LoadAdapterResponse(loaded=False)
  
          if request.adapter_index in self.model.loaded_adapters:
              logger.info(f"Adapter {request.adapter_index} is already loaded. Skipping.")
              return generate_pb2.LoadAdapterResponse(loaded=True)
  
          try:
              adapter_source = adapter_source_enum_to_string(request.adapter_source)
              adapter_index = request.adapter_index
              api_token = request.api_token
  
              if adapter_source == PBASE:
                  for i in range(len(adapter_parameters.adapter_ids)):
                      adapter_id = adapter_parameters.adapter_ids[i]
                      adapter_id = map_pbase_model_id_to_s3(adapter_id, api_token)
                      adapter_parameters.adapter_ids[i] = adapter_id
                  adapter_source = S3
  
              self.model.load_adapter(adapter_parameters, adapter_source, adapter_index, api_token)
  
              return generate_pb2.LoadAdapterResponse(loaded=True)
          except Exception:
              logger.exception("Error when loading adapter")
              raise
  
      async def OffloadAdapter(self, request: generate_pb2.OffloadAdapterRequest, context):
          adapter_parameters = request.adapter_parameters
          if is_base_model(adapter_parameters):
              logger.info("No adapter to offload for base model. Skipping.")
              return generate_pb2.OffloadAdapterResponse(offloaded=False)
  
          try:
              adapter_idx = request.adapter_index
              adapter_source = adapter_source_enum_to_string(request.adapter_source)
              adapter_index = request.adapter_index
  
              offloaded = self.model.offload_adapter(adapter_idx, adapter_source, adapter_index)
              if offloaded:
                  # Ensure there is enough memory for the next adapter
                  torch.cuda.empty_cache()
                  torch.cuda.synchronize(self.model.device)
  
              return generate_pb2.OffloadAdapterResponse(offloaded=offloaded)
          except Exception:
              logger.exception("Error when offloading adapter")
              raise
  
  
  def serve(
      model_id: str,
      adapter_id: str,
      revision: Optional[str],
      sharded: bool,
      quantize: Optional[str],
      compile: bool,
      dtype: Optional[str],
      trust_remote_code: bool,
      uds_path: Path,
      source: str,
      adapter_source: str,
      speculative_tokens: int,
      preloaded_adapter_ids: List[str],
      merge_adapter_weights: bool,
      preloaded_adapter_source: str,
      embedding_dim: Optional[int] = None,
  ):
      async def serve_inner(
          model_id: str,
          adapter_id: str,
          adapter_source: str,
          revision: Optional[str],
          sharded: bool,
          quantize: Optional[str],
          compile: bool,
          dtype: Optional[str],
          trust_remote_code: bool,
          speculative_tokens: int,
          preloaded_adapter_ids: List[str],
          merge_adapter_weights: bool,
          preloaded_adapter_source: str,
          embedding_dim: Optional[int] = None,
      ):
          unix_socket_template = "unix://{}-{}"
          if sharded:
              server_urls = [unix_socket_template.format(uds_path, rank) for rank in range(int(os.environ["WORLD_SIZE"]))]
              local_url = server_urls[int(os.environ["RANK"])]
          else:
              local_url = unix_socket_template.format(uds_path, 0)
              server_urls = [local_url]
  
          if adapter_source == PBASE and adapter_id != "":
              logger.info("Got a PBASE adapter source, mapping model ID to S3")
              api_token = os.getenv("PREDIBASE_API_TOKEN")
              adapter_id = map_pbase_model_id_to_s3(adapter_id, api_token)
              adapter_source = S3
  
          try:
              model = get_model(
                  model_id,
                  adapter_id,
                  revision,
                  sharded,
                  quantize,
                  compile,
                  dtype,
                  trust_remote_code,
                  source,
                  adapter_source,
                  merge_adapter_weights,
                  embedding_dim,
              )
          except Exception:
              logger.exception("Error when initializing model")
              raise
  
          if quantize == "gptq":
              try:
                  # When using GPTQ, Exllama kernels need some global kernels
                  # For which we have the finale shapes only after the model has loaded
                  # This will allocate those buffers.
                  from lorax_server.utils.gptq.exllamav2 import (
                      create_exllama_buffers,
                      set_device,
                  )
  
                  set_device(model.device)
                  create_exllama_buffers()
              except ImportError:
                  pass
  
          # set speculative decoding tokens
          speculative_tokens = max(model.max_speculative_tokens, speculative_tokens)
          if speculative_tokens > 0:
              # Only use ngram speculation if the model does not support speculative tokens itself
              use_ngram = model.max_speculative_tokens == 0
              set_speculative_tokens(speculative_tokens, use_ngram=use_ngram)
  
          if preloaded_adapter_ids:
              logger.info(f"Preloading {len(preloaded_adapter_ids)} adapters")
  
              _adapter_source = enum_string_to_adapter_source(preloaded_adapter_source)
              adapter_preload_api_token = None
              if _adapter_source == generate_pb2.AdapterSource.PBASE:
                  # Derive the predibase token from an env variable if we are using predibase adapters.
                  adapter_preload_api_token = os.getenv("PREDIBASE_API_TOKEN")
              elif _adapter_source == generate_pb2.AdapterSource.HUB:
                  # Use global token during init
                  adapter_preload_api_token = os.getenv("HUGGING_FACE_HUB_TOKEN")
  
              preloaded_adapters = [
                  generate_pb2.PreloadedAdapter(
                      adapter_parameters=generate_pb2.AdapterParameters(adapter_ids=[adapter_id]),
                      adapter_source=_adapter_source,
                      adapter_index=i + 1,
                  )
                  for i, adapter_id in enumerate(preloaded_adapter_ids)
              ]
  
              download_requests = [
                  generate_pb2.DownloadAdapterRequest(
                      adapter_parameters=adapter_info.adapter_parameters,
                      adapter_source=adapter_info.adapter_source,
                      api_token=adapter_preload_api_token,
                  )
                  for adapter_info in preloaded_adapters
              ]
              models = [model] * len(download_requests)
              are_preloaded = [True] * len(download_requests)
  
              # Download adapters
              t0 = time.time()
              with concurrent.futures.ThreadPoolExecutor(max_workers=20) as executor:
                  download_responses = list(
                      tqdm(
                          executor.map(
                              download_adapter, download_requests, models, are_preloaded
                          ),
                          total=len(download_requests),
                      )
                  )
              logger.info(f"Downloaded {len(download_requests)} adapters in {time.time() - t0:.2f}s")
  
              if not all(download_responses):
                  raise RuntimeError("Failed to download all adapters")
  
              def load_adapter(adapter_info: generate_pb2.PreloadedAdapter) -> bool:
                  _adapter_source = adapter_source_enum_to_string(adapter_info.adapter_source)
                  _adapter_id = adapter_info.adapter_parameters.adapter_ids[0]
                  if _adapter_source == PBASE:
                      _adapter_id = map_pbase_model_id_to_s3(_adapter_id, api_token=adapter_preload_api_token)
                      _adapter_source = S3
  
                  model.load_adapter(
                      generate_pb2.AdapterParameters(adapter_ids=[_adapter_id]),
                      _adapter_source,
                      adapter_index=adapter_info.adapter_index,
                      api_token=None,
                      dynamic=True,
                  )
                  return True
  
              # Load adapters
              t0 = time.time()
              with concurrent.futures.ThreadPoolExecutor(max_workers=20) as executor:
                  responses = list(tqdm(executor.map(load_adapter, preloaded_adapters), total=len(preloaded_adapters)))
  
              if not all(responses):
                  raise RuntimeError("Failed to preload all adapters")
  
              logger.info(f"Preloaded {len(preloaded_adapters)} adapters in {time.time() - t0:.2f}s")
  
              adapter_memory_fractions = [r.memory_fraction for r in download_responses]
              model.register_preloaded_adapters(preloaded_adapters, adapter_memory_fractions)
  
          server = aio.server(
              interceptors=[
                  ExceptionInterceptor(),
                  UDSOpenTelemetryAioServerInterceptor(),
              ],
              options=[
                  # Set the maximum possible message length: i32::MAX
                  ("grpc.max_receive_message_length", (1 << 31) - 1)
              ],
          )
          generate_pb2_grpc.add_LoraxServiceServicer_to_server(LoraxService(model, Cache(), server_urls), server)
          SERVICE_NAMES = (
              generate_pb2.DESCRIPTOR.services_by_name["LoraxService"].full_name,
              reflection.SERVICE_NAME,
          )
          reflection.enable_server_reflection(SERVICE_NAMES, server)
          server.add_insecure_port(local_url)
  
          await server.start()
  
          # Log SGMV kernel status
          if not LORAX_PUNICA_TRITON_DISABLED and not model.dynamic_adapter_loading_enabled:
              logger.info("Trion kernel is enabled, multi-LoRA inference will be fast!")
          if has_sgmv():
              logger.info("SGMV kernel is enabled, multi-LoRA inference will be fast!")
          else:
              logger.info("Punica kernels are disabled, multi-LoRA inference may be slow")
  
          logger.info("Server started at {}".format(local_url))
  
          try:
              await server.wait_for_termination()
          except KeyboardInterrupt:
              logger.info("Signal received. Shutting down")
              await server.stop(0)
  
      asyncio.run(
          serve_inner(
              model_id,
              adapter_id,
              adapter_source,
              revision,
              sharded,
              quantize,
              compile,
              dtype,
              trust_remote_code,
              speculative_tokens,
              preloaded_adapter_ids,
              merge_adapter_weights,
              preloaded_adapter_source,
              embedding_dim,
          )
      )
  