Fixing some missed changes in last merge with main.

Author: aman2930

jetstream/core/orchestrator.py

@@ -1611,109 +1611,6 @@ def list_adapters_from_tensorstore(self):
     return listed_adapters
 
 
-  def load_adapter_to_tensorstore(
-          self,
-          adapter_id: str,
-          adapter_path: str):
-    """Load the adapter to adapter_tensorstore for each engine."""
-    logger.info("Loading adapter_id=%s from %s.",
-        adapter_id, adapter_path)
-
-    for idx, tensorstore in enumerate(self._prefill_adapterstore):
-      try:
-        engine = self._prefill_engines[idx]
-        adapter_params, adapter_config = engine.load_single_adapter(
-            adapter_path)
-
-        if not adapter_params or not adapter_config:
-          raise ValueError(
-              f"Failed to load adapter={adapter_id} from {adapter_path}.")
-
-        tensorstore.register_adapter(
-            adapter_id,
-            adapter_path,
-            adapter_config)
-
-        asyncio.run(tensorstore.load_adapter(adapter_id, adapter_params, True))
-
-        logger.info("Successfully loaded '%s' in engine_%d.",
-            adapter_id, idx)
-        engine.print_stats(f"After loading '{adapter_id}' in engine_{idx}")
-
-      except Exception as e:
-        logger.info("Adapter loading failed with error: %s", str(e))
-        raise e
-
-    for idx, tensorstore in enumerate(self._generate_adapterstore):
-      try:
-        engine = self._generate_engines[idx]
-        adapter_params, adapter_config = engine.load_single_adapter(
-            adapter_path)
-
-        if not adapter_params or not adapter_config:
-          raise ValueError(
-              f"Failed to load adapter={adapter_id} from {adapter_path}.")
-
-        tensorstore.register_adapter(
-            adapter_id,
-            adapter_path,
-            adapter_config)
-
-        asyncio.run(tensorstore.load_adapter(adapter_id, adapter_params, True))
-
-        logger.info("Successfully loaded '%s' in engine_%d.",
-            adapter_id, idx)
-        engine.print_stats(f"After loading '{adapter_id}' in engine_{idx}")
-
-      except Exception as e:
-        logger.info("Adapter loading failed with error: %s", str(e))
-        raise e
-
-
-  def unload_adapter_from_tensorstore(
-          self,
-          adapter_id: str):
-    """Unload the adapter from adapter_tensorstore of each engine."""
-    logger.info("Unloading adapter_id=%s", adapter_id)
-
-    for idx, tensorstore in enumerate(self._prefill_adapterstore):
-      try:
-        engine = self._prefill_engines[idx]
-        asyncio.run(tensorstore.unload_adapter(adapter_id))
-
-        logger.info("Successfully unloaded '%s' in engine_%d.",
-            adapter_id, idx)
-        engine.print_stats(f"After unloading '{adapter_id}' in engine_{idx}")
-
-      except Exception as e:
-        logger.info("Adapter unloading failed with error: %s", str(e))
-        raise e
-
-    for idx, tensorstore in enumerate(self._generate_adapterstore):
-      try:
-        engine = self._generate_engines[idx]
-        asyncio.run(tensorstore.unload_adapter(adapter_id))
-
-        logger.info("Successfully unloaded '%s' in engine_%d.",
-            adapter_id, idx)
-        engine.print_stats(f"After unloading '{adapter_id}' in engine_{idx}")
-
-      except Exception as e:
-        logger.info("Adapter unloading failed with error: %s", str(e))
-        raise e
-
-
-  def list_adapters_from_tensorstore(self):
-    """List all the adapters from the adapter_tensorstore of each engine."""
-    logger.info("Listing loaded adapters.")
-
-    listed_adapters = {}
-    for tensorstore in self._generate_adapterstore:
-      listed_adapters.update(tensorstore.adapter_registry)
-
-    return listed_adapters
-
-
 class LLMOrchestrator(jetstream_pb2_grpc.OrchestratorServicer):
   """Coordinates a set of prefill and generate slices for LLM decoding."""
 

jetstream/core/proto/jetstream_pb2_grpc.py

@@ -19,8 +19,13 @@
 
 
 class OrchestratorStub(object):
-    """TODO: Merge this with main JetStream core once we settle on an API.
+  """TODO: Merge this with main JetStream core once we settle on an API."""
 
+  def __init__(self, channel):
+    """Constructor.
+
+    Args:
+        channel: A grpc.Channel.
     """
     self.Decode = channel.unary_stream(
         "/jetstream_proto.Orchestrator/Decode",
@@ -35,23 +40,19 @@ class OrchestratorStub(object):
 
 
 class OrchestratorServicer(object):
-    """TODO: Merge this with main JetStream core once we settle on an API.
-
-    """
+  """TODO: Merge this with main JetStream core once we settle on an API."""
 
-    def Decode(self, request, context):
-        """Query LLM to generate text or tokens.
-        """
-        context.set_code(grpc.StatusCode.UNIMPLEMENTED)
-        context.set_details('Method not implemented!')
-        raise NotImplementedError('Method not implemented!')
+  def Decode(self, request, context):
+    """Query LLM to generate text or tokens."""
+    context.set_code(grpc.StatusCode.UNIMPLEMENTED)
+    context.set_details("Method not implemented!")
+    raise NotImplementedError("Method not implemented!")
 
-    def HealthCheck(self, request, context):
-        """Checks if the model server is live.
-        """
-        context.set_code(grpc.StatusCode.UNIMPLEMENTED)
-        context.set_details('Method not implemented!')
-        raise NotImplementedError('Method not implemented!')
+  def HealthCheck(self, request, context):
+    """Checks if the model server is live."""
+    context.set_code(grpc.StatusCode.UNIMPLEMENTED)
+    context.set_details("Method not implemented!")
+    raise NotImplementedError("Method not implemented!")
 
 
 def add_OrchestratorServicer_to_server(servicer, server):
@@ -73,9 +74,9 @@ def add_OrchestratorServicer_to_server(servicer, server):
   server.add_generic_rpc_handlers((generic_handler,))
 
 
- # This class is part of an EXPERIMENTAL API.
+# This class is part of an EXPERIMENTAL API.
 class Orchestrator(object):
-    """TODO: Merge this with main JetStream core once we settle on an API.
+  """TODO: Merge this with main JetStream core once we settle on an API."""
 
   @staticmethod
   def Decode(