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 #include "tensorflow/core/common_runtime/local_device.h"
21 #include "tensorflow/core/framework/device_attributes.pb.h"
22 #include "tensorflow/core/lib/core/errors.h"
23 #include "tensorflow/core/platform/logging.h"
24 #include "tensorflow/core/util/device_name_utils.h"
25
26 namespace tensorflow {
27
DeviceMgr(std::vector<std::unique_ptr<Device>> devices)28 DeviceMgr::DeviceMgr(std::vector<std::unique_ptr<Device>> devices)
29 : devices_(std::move(devices)), name_backing_store_(128) {
30 for (auto& d : devices_) {
31 CHECK(d->device_mgr_ == nullptr);
32 d->device_mgr_ = this;
33
34 // Register under the (1) full name and (2) canonical name.
35 for (const string& name :
36 DeviceNameUtils::GetNamesForDeviceMappings(d->parsed_name())) {
37 device_map_[CopyToBackingStore(name)] = d.get();
38 }
39 // Register under the (3) local name and (4) legacy local name.
40 for (const string& name :
41 DeviceNameUtils::GetLocalNamesForDeviceMappings(d->parsed_name())) {
42 device_map_[CopyToBackingStore(name)] = d.get();
43 }
44 device_type_counts_[d->device_type()]++;
45 }
46 }
47
DeviceMgr(std::unique_ptr<Device> device)48 DeviceMgr::DeviceMgr(std::unique_ptr<Device> device)
49 : DeviceMgr([&device] {
50 std::vector<std::unique_ptr<Device>> vector;
51 vector.push_back(std::move(device));
52 return vector;
53 }()) {}
54
CopyToBackingStore(StringPiece s)55 StringPiece DeviceMgr::CopyToBackingStore(StringPiece s) {
56 size_t n = s.size();
57 char* space = name_backing_store_.Alloc(n);
58 memcpy(space, s.data(), n);
59 return StringPiece(space, n);
60 }
61
ListDeviceAttributes(std::vector<DeviceAttributes> * devices) const62 void DeviceMgr::ListDeviceAttributes(
63 std::vector<DeviceAttributes>* devices) const {
64 devices->reserve(devices_.size());
65 for (const auto& dev : devices_) {
66 devices->emplace_back(dev->attributes());
67 }
68 }
69
ListDevices() const70 std::vector<Device*> DeviceMgr::ListDevices() const {
71 std::vector<Device*> devices(devices_.size());
72 for (size_t i = 0; i < devices_.size(); ++i) {
73 devices[i] = devices_[i].get();
74 }
75 return devices;
76 }
77
DebugString() const78 string DeviceMgr::DebugString() const {
79 string out;
80 for (const auto& dev : devices_) {
81 strings::StrAppend(&out, dev->name(), "\n");
82 }
83 return out;
84 }
85
DeviceMappingString() const86 string DeviceMgr::DeviceMappingString() const {
87 string out;
88 for (const auto& dev : devices_) {
89 if (!dev->attributes().physical_device_desc().empty()) {
90 strings::StrAppend(&out, dev->name(), " -> ",
91 dev->attributes().physical_device_desc(), "\n");
92 }
93 }
94 return out;
95 }
96
LookupDevice(StringPiece name,Device ** device) const97 Status DeviceMgr::LookupDevice(StringPiece name, Device** device) const {
98 auto iter = device_map_.find(name);
99 if (iter == device_map_.end()) {
100 std::vector<StringPiece> device_names;
101 for (auto&& itr : device_map_) {
102 device_names.push_back(itr.first);
103 }
104 VLOG(1) << "Unknown device: " << name
105 << " all devices: " << str_util::Join(device_names, ", ");
106 return errors::InvalidArgument(name, " unknown device.");
107 }
108 *device = iter->second;
109 return Status::OK();
110 }
111
ClearContainers(gtl::ArraySlice<string> containers) const112 void DeviceMgr::ClearContainers(gtl::ArraySlice<string> containers) const {
113 Status s;
114 for (const auto& dev : devices_) {
115 if (containers.empty()) {
116 s.Update(dev->resource_manager()->Cleanup(
117 dev->resource_manager()->default_container()));
118 } else {
119 for (const string& c : containers) {
120 s.Update(dev->resource_manager()->Cleanup(c));
121 }
122 }
123 if (!s.ok()) {
124 LOG(WARNING) << s;
125 }
126 }
127 }
128
NumDeviceType(const string & type) const129 int DeviceMgr::NumDeviceType(const string& type) const {
130 auto iter = device_type_counts_.find(type);
131 if (iter != device_type_counts_.end()) return iter->second;
132 return 0;
133 }
134
135 } // namespace tensorflow
136