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_WORKER_H_
     17 #define TENSORFLOW_CORE_DISTRIBUTED_RUNTIME_WORKER_H_
     18 
     19 #include <unordered_map>
     20 
     21 #include "tensorflow/core/distributed_runtime/graph_mgr.h"
     22 #include "tensorflow/core/distributed_runtime/partial_run_mgr.h"
     23 #include "tensorflow/core/distributed_runtime/session_mgr.h"
     24 #include "tensorflow/core/distributed_runtime/worker_interface.h"
     25 
     26 namespace tensorflow {
     27 
     28 class CancellationManager;
     29 class Device;
     30 struct WorkerEnv;
     31 struct WorkerSession;
     32 
     33 // A TensorFlow Worker runs registered graphs and supports worker-to-worker
     34 // Tensor transfer.
     35 //
     36 // See `../protobuf/worker_service.proto` for more details about each method.
     37 //
     38 // This class may be subclassed to provide specialized implementations of
     39 // particular methods for different transport mechanism. For example,
     40 // `GrpcWorker` specializes the `RecvTensorAsync()` method to support a more
     41 // efficient gRPC data structure for handling large binary data.
     42 class Worker : public WorkerInterface {
     43  public:
     44   Worker(WorkerEnv* env);
     45   virtual ~Worker() {}
     46 
     47   void GetStatusAsync(const GetStatusRequest* request,
     48                       GetStatusResponse* response,
     49                       StatusCallback done) override;
     50 
     51   void CreateWorkerSessionAsync(const CreateWorkerSessionRequest* request,
     52                                 CreateWorkerSessionResponse* response,
     53                                 StatusCallback done) override;
     54 
     55   void DeleteWorkerSessionAsync(const DeleteWorkerSessionRequest* request,
     56                                 DeleteWorkerSessionResponse* response,
     57                                 StatusCallback done) override;
     58 
     59   void RegisterGraphAsync(const RegisterGraphRequest* request,
     60                           RegisterGraphResponse* response,
     61                           StatusCallback done) override;
     62 
     63   void DeregisterGraphAsync(const DeregisterGraphRequest* request,
     64                             DeregisterGraphResponse* response,
     65                             StatusCallback done) override;
     66 
     67   void RunGraphAsync(CallOptions* opts, RunGraphRequestWrapper* request,
     68                      MutableRunGraphResponseWrapper* response,
     69                      StatusCallback done) override;
     70 
     71   MutableRunGraphRequestWrapper* CreateRunGraphRequest() override;
     72 
     73   MutableRunGraphResponseWrapper* CreateRunGraphResponse() override;
     74 
     75   void CleanupGraphAsync(const CleanupGraphRequest* request,
     76                          CleanupGraphResponse* response,
     77                          StatusCallback done) override;
     78 
     79   void CleanupAllAsync(const CleanupAllRequest* request,
     80                        CleanupAllResponse* response,
     81                        StatusCallback done) override;
     82 
     83   void RecvTensorAsync(CallOptions* opts, const RecvTensorRequest* request,
     84                        TensorResponse* response, StatusCallback done) override;
     85 
     86   void LoggingAsync(const LoggingRequest* request, LoggingResponse* response,
     87                     StatusCallback done) override;
     88 
     89   void TracingAsync(const TracingRequest* request, TracingResponse* response,
     90                     StatusCallback done) override;
     91 
     92  protected:
     93   WorkerEnv* const env_;  // Not owned.
     94 
     95   Status PrepareRecvTensor(const Rendezvous::ParsedKey& parsed,
     96                            Device** src_dev);
     97 
     98   void AbortStep(int64);
     99 
    100  private:
    101   PartialRunMgr partial_run_mgr_;
    102 
    103   mutex mu_;
    104   CancellationManager* cancellation_manager_ GUARDED_BY(mu_);
    105 
    106   Status PrepareRunGraph(RunGraphRequestWrapper* req,
    107                          GraphMgr::NamedTensors* in,
    108                          GraphMgr::NamedTensors* out);
    109 
    110   void DoRunGraph(CallOptions* opts, RunGraphRequestWrapper* request,
    111                   MutableRunGraphResponseWrapper* response,
    112                   StatusCallback done);
    113 
    114   void DoPartialRunGraph(CallOptions* opts, RunGraphRequestWrapper* request,
    115                          MutableRunGraphResponseWrapper* response,
    116                          StatusCallback done);
    117 
    118   TF_DISALLOW_COPY_AND_ASSIGN(Worker);
    119 };
    120 
    121 }  // namespace tensorflow
    122 
    123 #endif  // TENSORFLOW_CORE_DISTRIBUTED_RUNTIME_WORKER_H_
    124