Home | History | Annotate | Download | only in distributed_runtime
      1 /* Copyright 2016 The TensorFlow Authors. All Rights Reserved.
      2 
      3 Licensed under the Apache License, Version 2.0 (the "License");
      4 you may not use this file except in compliance with the License.
      5 You may obtain a copy of the License at
      6 
      7     http://www.apache.org/licenses/LICENSE-2.0
      8 
      9 Unless required by applicable law or agreed to in writing, software
     10 distributed under the License is distributed on an "AS IS" BASIS,
     11 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
     12 See the License for the specific language governing permissions and
     13 limitations under the License.
     14 ==============================================================================*/
     15 
     16 #ifndef TENSORFLOW_CORE_DISTRIBUTED_RUNTIME_SERVER_LIB_H_
     17 #define TENSORFLOW_CORE_DISTRIBUTED_RUNTIME_SERVER_LIB_H_
     18 
     19 #include <memory>
     20 
     21 #include "tensorflow/core/lib/core/status.h"
     22 #include "tensorflow/core/platform/macros.h"
     23 #include "tensorflow/core/protobuf/tensorflow_server.pb.h"
     24 
     25 namespace tensorflow {
     26 
     27 // This library supports a registration/factory-based mechanism for
     28 // creating TensorFlow server objects. Each server implementation must
     29 // have an accompanying implementation of ServerFactory, and create a
     30 // static "registrar" object that calls `ServerFactory::Register()`
     31 // with an instance of the factory class. See "rpc/grpc_server_lib.cc"
     32 // for an example.
     33 
     34 // Represents a single TensorFlow server that exports Master and Worker
     35 // services.
     36 class ServerInterface {
     37  public:
     38   ServerInterface() {}
     39   virtual ~ServerInterface() {}
     40 
     41   // Starts the server running asynchronously. Returns OK on success, otherwise
     42   // returns an error.
     43   virtual Status Start() = 0;
     44 
     45   // Stops the server asynchronously. Returns OK on success, otherwise returns
     46   // an error.
     47   //
     48   // After calling `Stop()`, the caller may call `Join()` to block until the
     49   // server has stopped.
     50   virtual Status Stop() = 0;
     51 
     52   // Blocks until the server has stopped. Returns OK on success, otherwise
     53   // returns an error.
     54   virtual Status Join() = 0;
     55 
     56   // Returns a target string that can be used to connect to this server using
     57   // `tensorflow::NewSession()`.
     58   virtual const string target() const = 0;
     59 
     60  private:
     61   TF_DISALLOW_COPY_AND_ASSIGN(ServerInterface);
     62 };
     63 
     64 class ServerFactory {
     65  public:
     66   // Creates a new server based on the given `server_def`, and stores
     67   // it in `*out_server`. Returns OK on success, otherwise returns an
     68   // error.
     69   virtual Status NewServer(const ServerDef& server_def,
     70                            std::unique_ptr<ServerInterface>* out_server) = 0;
     71 
     72   // Returns true if and only if this factory can create a server
     73   // based on the given `server_def`.
     74   virtual bool AcceptsOptions(const ServerDef& server_def) = 0;
     75 
     76   virtual ~ServerFactory() {}
     77 
     78   // For each `ServerFactory` subclass, an instance of that class must
     79   // be registered by calling this method.
     80   //
     81   // The `server_type` must be unique to the server factory.
     82   static void Register(const string& server_type, ServerFactory* factory);
     83 
     84   // Looks up a factory that can create a server based on the given
     85   // `server_def`, and stores it in `*out_factory`. Returns OK on
     86   // success, otherwise returns an error.
     87   static Status GetFactory(const ServerDef& server_def,
     88                            ServerFactory** out_factory);
     89 };
     90 
     91 // Creates a server based on the given `server_def`, and stores it in
     92 // `*out_server`. Returns OK on success, otherwise returns an error.
     93 Status NewServer(const ServerDef& server_def,
     94                  std::unique_ptr<ServerInterface>* out_server);
     95 
     96 }  // namespace tensorflow
     97 
     98 #endif  // TENSORFLOW_CORE_DISTRIBUTED_RUNTIME_SERVER_LIB_H_
     99