1 /* Copyright 2015 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 #include "tensorflow/core/common_runtime/device_mgr.h"
17
18 #include <memory>
19 #include <vector>
20
21 #include "tensorflow/core/common_runtime/local_device.h"
22 #include "tensorflow/core/framework/device_attributes.pb.h"
23 #include "tensorflow/core/lib/core/errors.h"
24 #include "tensorflow/core/platform/logging.h"
25 #include "tensorflow/core/util/device_name_utils.h"
26
27 namespace tensorflow {
28
~DeviceMgr()29 DeviceMgr::~DeviceMgr() {}
30
StaticDeviceMgr(std::vector<std::unique_ptr<Device>> devices)31 StaticDeviceMgr::StaticDeviceMgr(std::vector<std::unique_ptr<Device>> devices)
32 : devices_(std::move(devices)),
33 name_backing_store_(128),
34 cpu_device_(nullptr) {
35 for (auto& d : devices_) {
36 // Register under the (1) full name and (2) canonical name.
37 for (const string& name :
38 DeviceNameUtils::GetNamesForDeviceMappings(d->parsed_name())) {
39 device_map_[CopyToBackingStore(name)] = d.get();
40 }
41 // Register under the (3) local name and (4) legacy local name.
42 for (const string& name :
43 DeviceNameUtils::GetLocalNamesForDeviceMappings(d->parsed_name())) {
44 device_map_[CopyToBackingStore(name)] = d.get();
45 }
46 const auto& t = d->device_type();
47 device_type_counts_[t]++;
48 device_incarnation_set_.insert(d->attributes().incarnation());
49 if (cpu_device_ == nullptr && t == "CPU" && d->parsed_name().id == 0) {
50 cpu_device_ = d.get();
51 }
52 }
53 }
54
StaticDeviceMgr(std::unique_ptr<Device> device)55 StaticDeviceMgr::StaticDeviceMgr(std::unique_ptr<Device> device)
56 : StaticDeviceMgr([&device] {
57 std::vector<std::unique_ptr<Device>> vector;
58 vector.push_back(std::move(device));
59 return vector;
60 }()) {}
61
~StaticDeviceMgr()62 StaticDeviceMgr::~StaticDeviceMgr() {
63 // Release resources ahead of destroying the device manager as the resource
64 // destructors (e.g. ~IteratorResource) assume devices still exist.
65 for (auto& device : devices_) {
66 device->ClearResourceMgr();
67 }
68 }
69
CopyToBackingStore(StringPiece s)70 StringPiece StaticDeviceMgr::CopyToBackingStore(StringPiece s) {
71 size_t n = s.size();
72 char* space = name_backing_store_.Alloc(n);
73 memcpy(space, s.data(), n);
74 return StringPiece(space, n);
75 }
76
ListDeviceAttributes(std::vector<DeviceAttributes> * devices) const77 void StaticDeviceMgr::ListDeviceAttributes(
78 std::vector<DeviceAttributes>* devices) const {
79 devices->reserve(devices_.size());
80 for (const auto& dev : devices_) {
81 devices->emplace_back(dev->attributes());
82 }
83 }
84
ListDevices() const85 std::vector<Device*> StaticDeviceMgr::ListDevices() const {
86 std::vector<Device*> devices(devices_.size());
87 for (size_t i = 0; i < devices_.size(); ++i) {
88 devices[i] = devices_[i].get();
89 }
90 return devices;
91 }
92
DebugString() const93 string StaticDeviceMgr::DebugString() const {
94 string out;
95 for (const auto& dev : devices_) {
96 strings::StrAppend(&out, dev->name(), "\n");
97 }
98 return out;
99 }
100
DeviceMappingString() const101 string StaticDeviceMgr::DeviceMappingString() const {
102 string out;
103 for (const auto& dev : devices_) {
104 if (!dev->attributes().physical_device_desc().empty()) {
105 strings::StrAppend(&out, dev->name(), " -> ",
106 dev->attributes().physical_device_desc(), "\n");
107 }
108 }
109 return out;
110 }
111
LookupDevice(StringPiece name,Device ** device) const112 Status StaticDeviceMgr::LookupDevice(StringPiece name, Device** device) const {
113 auto iter = device_map_.find(name);
114 if (iter == device_map_.end()) {
115 std::vector<StringPiece> device_names;
116 for (auto&& itr : device_map_) {
117 device_names.push_back(itr.first);
118 }
119 VLOG(1) << "Unknown device: " << name
120 << " all devices: " << absl::StrJoin(device_names, ", ");
121 return errors::InvalidArgument(name, " unknown device.");
122 }
123 *device = iter->second;
124 return Status::OK();
125 }
126
ContainsDevice(int64 device_incarnation) const127 bool StaticDeviceMgr::ContainsDevice(int64 device_incarnation) const {
128 return device_incarnation_set_.contains(device_incarnation);
129 }
130
ClearContainers(gtl::ArraySlice<string> containers) const131 void StaticDeviceMgr::ClearContainers(
132 gtl::ArraySlice<string> containers) const {
133 Status s;
134 for (const auto& dev : devices_) {
135 if (containers.empty()) {
136 s.Update(dev->resource_manager()->Cleanup(
137 dev->resource_manager()->default_container()));
138 } else {
139 for (const string& c : containers) {
140 s.Update(dev->resource_manager()->Cleanup(c));
141 }
142 }
143 if (!s.ok()) {
144 LOG(WARNING) << s;
145 }
146 }
147 }
148
NumDeviceType(const string & type) const149 int StaticDeviceMgr::NumDeviceType(const string& type) const {
150 auto iter = device_type_counts_.find(type);
151 if (iter != device_type_counts_.end()) return iter->second;
152 return 0;
153 }
154
HostCPU() const155 Device* StaticDeviceMgr::HostCPU() const { return cpu_device_; }
156
157 } // namespace tensorflow
158