Refactoring and fixing lint errors.

Author: aman2930

jetstream/core/lora/adapter_tensorstore.py

@@ -104,8 +104,8 @@ def __init__(self,
 
   def register_adapter(self,
       adapter_id: str,
-      adapter_path: str = None,
-      adapter_config: Dict[str, Any] = None):
+      adapter_path: str | None = None,
+      adapter_config: Dict[str, Any] | None = None):
     """Registers a new LoRA adatper."""
     """
     Registers a LoRA adapter with the TensorStore.  This also loads the adapter; 
@@ -284,6 +284,7 @@ async def load_adapter(
 
     try:
       if adapter_weights is None:
+        adapter_path = f"{self.adapters_dir_path}/{adapter_id}"
         adapter_weights, adapter_config = self.engine.load_single_adapter(adapter_path)
 
       if adapter_weights is None:

jetstream/core/lora/multi_lora_inference_api.py

@@ -40,7 +40,7 @@ def models(
     """ListAdapters all loaded LoRA adapters."""
 
     try:
-      adapters = self._driver.listAdaptersFromTensorstore()
+      adapters = self._driver.list_adapters_from_tensorstore()
 
       adapter_infos = []
       for adapter_id, adapter_data in adapters.items():
@@ -77,7 +77,7 @@ def load_lora_adapter(
     """Load a LoRA adapter as mentioned in the request."""
 
     try:
-      self._driver.loadAdapterToTensorstore(request.adapter_id, request.adapter_path)
+      self._driver.load_adapter_to_tensorstore(request.adapter_id, request.adapter_path)
 
       return multi_lora_decoding_pb2.LoadAdapterResponse(success=True)
     except Exception as e:
@@ -93,7 +93,7 @@ def unload_lora_adapter(
     """Unload a LoRA adapter as mentioned in the request."""
 
     try:
-      self._driver.unloadAdapterFromTensorstore(request.adapter_id)
+      self._driver.unload_adapter_from_tensorstore(request.adapter_id)
       return multi_lora_decoding_pb2.UnloadAdapterResponse(success=True)
     except Exception as e:
       logging.info(f"Loading of adapter_id={request.adapter_id} failed with error: {str(e)}")

jetstream/core/metrics/prometheus.py

@@ -254,14 +254,14 @@ def __init__(self, model_name: Optional[str] = None):
 
   _kv_cache_utilization = Gauge(
       name="kv_cache_utilization_perc",
-      documentation="Percentage of kv-cache utilized by the requests under processing.",
+      documentation="kv-cache utilization % by the requests under processing.",
       labelnames=["id"],
       multiprocess_mode="sum",
   )
 
   _lora_request_info = Gauge(
       name="lora_request_info",
-      documentation="Information about LoRA adapters loaded into TPU Memory for serving current requests.",
+      documentation="LoRA adapters loaded into HBM for processing requests.",
       labelnames=[
         "id",
         "max_lora",
@@ -322,4 +322,5 @@ def get_kv_cache_utilization_metric(self):
     return self._kv_cache_utilization.labels(**self.universal_labels)
 
   def get_lora_request_info_metric(self, max_lora: int, loaded_adapters: str):
-    return self._lora_request_info.labels(**self.universal_labels, max_lora=max_lora, running_lora_adapters=loaded_adapters)
+    return self._lora_request_info.labels(**self.universal_labels,
+        max_lora=max_lora, running_lora_adapters=loaded_adapters)