1 /* Copyright 2018 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 #ifndef TENSORFLOW_CORE_DISTRIBUTED_RUNTIME_COLLECTIVE_PARAM_RESOLVER_DISTRIBUTED_H_ 16 #define TENSORFLOW_CORE_DISTRIBUTED_RUNTIME_COLLECTIVE_PARAM_RESOLVER_DISTRIBUTED_H_ 17 18 #include "tensorflow/core/common_runtime/collective_param_resolver_local.h" 19 20 namespace tensorflow { 21 class ConfigProto; 22 class WorkerCacheInterface; 23 class DeviceResolverDistributed; 24 class DeviceMgr; 25 26 class CollectiveParamResolverDistributed : public CollectiveParamResolverLocal { 27 public: 28 CollectiveParamResolverDistributed(const ConfigProto& config, 29 const DeviceMgr* dev_mgr, 30 DeviceResolverDistributed* dev_resolver, 31 WorkerCacheInterface* worker_cache, 32 const string& task_name); 33 34 void CompleteParamsAsync(const string& device, CollectiveParams* cp, 35 CancellationManager* cancel_mgr, 36 const StatusCallback& done) override; 37 38 void CompleteGroupAsync(const CompleteGroupRequest* request, 39 CompleteGroupResponse* response, 40 CancellationManager* cancel_mgr, 41 const StatusCallback& done) override; 42 43 void CompleteInstanceAsync(const CompleteInstanceRequest* request, 44 CompleteInstanceResponse* response, 45 CancellationManager* cancel_mgr, 46 const StatusCallback& done) override; 47 48 protected: 49 // Returns true iff there's an entry for this group_key in the 50 // local group_table_. 51 bool GroupIsCached(int32 group_key) LOCKS_EXCLUDED(group_mu_); 52 53 // Updates group_table_ with contents of resp. 54 Status UpdateGroupCache(const CompleteGroupResponse& resp) 55 LOCKS_EXCLUDED(group_mu_); 56 57 // Finds the GroupRec that corresponds to cp->group_key and also 58 // populates cp->group from that GroupRec. 59 // 60 // Semantics are like those of CompleteGroupLocal but will make a 61 // remote call to the group leader if necessary. 62 void CompleteGroupDistributed(const string& device, CollectiveParams* cp, 63 CancellationManager* cancel_mgr, 64 const GroupRecCallback& done); 65 66 // Returns true iff there's an entry for this instance_key in the 67 // local instance_table_. 68 bool InstanceIsCached(int32 instance_key) LOCKS_EXCLUDED(instance_mu_); 69 70 // Updates instance_table_ with contents of resp. 71 void UpdateInstanceCache(const GroupRec* gr, CollectiveParams* cp, 72 const CompleteInstanceResponse& resp, 73 const StatusCallback& done) 74 LOCKS_EXCLUDED(instance_mu_, gr->mu, group_mu_); 75 76 // Finish populating *cp. Semantics are like those of 77 // CompleteInstanceLocal but will make a remote call to the group 78 // leader if necessary. 79 void CompleteInstanceDistributed(const string& device, const GroupRec* gr, 80 CollectiveParams* cp, 81 CancellationManager* cancel_mgr, 82 const StatusCallback& done) 83 LOCKS_EXCLUDED(instance_mu_, gr->mu, group_mu_); 84 85 WorkerCacheInterface* worker_cache_; // Not owned 86 const string group_leader_; 87 }; 88 89 } // namespace tensorflow 90 #endif // TENSORFLOW_CORE_DISTRIBUTED_RUNTIME_COLLECTIVE_PARAM_RESOLVER_DISTRIBUTED_H_ 91