pythontensorflow-serving

TensorFlow Serving: Update model_config (add additional models) at runtime


I'm busy configuring a TensorFlow Serving client that asks a TensorFlow Serving server to produce predictions on a given input image, for a given model.

If the model being requested has not yet been served, it is downloaded from a remote URL to a folder where the server's models are located. (The client does this). At this point I need to update the model_config and trigger the server to reload it.

This functionality appears to exist (based on https://github.com/tensorflow/serving/pull/885 and https://github.com/tensorflow/serving/blob/master/tensorflow_serving/apis/model_service.proto#L22), but I can't find any documentation on how to actually use it.

I am essentially looking for a python script with which I can trigger the reload from client side (or otherwise to configure the server to listen for changes and trigger the reload itself).


Solution

  • So it took me ages of trawling through pull requests to finally find a code example for this. For the next person who has the same question as me, here is an example of how to do this. (You'll need the tensorflow_serving package for this; pip install tensorflow-serving-api).

    Based on this pull request (which at the time of writing hadn't been accepted and was closed since it needed review): https://github.com/tensorflow/serving/pull/1065

    from tensorflow_serving.apis import model_service_pb2_grpc
    from tensorflow_serving.apis import model_management_pb2
    from tensorflow_serving.config import model_server_config_pb2
    
    import grpc
    
    def add_model_config(host, name, base_path, model_platform):
      channel = grpc.insecure_channel(host) 
      stub = model_service_pb2_grpc.ModelServiceStub(channel)
      request = model_management_pb2.ReloadConfigRequest() 
      model_server_config = model_server_config_pb2.ModelServerConfig()
    
      #Create a config to add to the list of served models
      config_list = model_server_config_pb2.ModelConfigList()       
      one_config = config_list.config.add()
      one_config.name= name
      one_config.base_path=base_path
      one_config.model_platform=model_platform
    
      model_server_config.model_config_list.CopyFrom(config_list)
    
      request.config.CopyFrom(model_server_config)
    
      print(request.IsInitialized())
      print(request.ListFields())
    
      response = stub.HandleReloadConfigRequest(request,10)
      if response.status.error_code == 0:
          print("Reload sucessfully")
      else:
          print("Reload failed!")
          print(response.status.error_code)
          print(response.status.error_message)
    
    
    add_model_config(host="localhost:8500", 
                        name="my_model", 
                        base_path="/models/my_model", 
                        model_platform="tensorflow")