1/* Copyright 2015 The TensorFlow Authors. All Rights Reserved.
2
3Licensed under the Apache License, Version 2.0 (the "License");
4you may not use this file except in compliance with the License.
5You may obtain a copy of the License at
6
7 http://www.apache.org/licenses/LICENSE-2.0
8
9Unless required by applicable law or agreed to in writing, software
10distributed under the License is distributed on an "AS IS" BASIS,
11WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12See the License for the specific language governing permissions and
13limitations 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
27namespace tensorflow {
28
29DeviceMgr::~DeviceMgr() {}
30
31StaticDeviceMgr::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
55StaticDeviceMgr::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
62StaticDeviceMgr::~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
70StringPiece 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
77void 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
85std::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
93string 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
101string 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
112Status 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 OkStatus();
125}
126
127bool StaticDeviceMgr::ContainsDevice(int64_t device_incarnation) const {
128 return device_incarnation_set_.contains(device_incarnation);
129}
130
131void 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
149int 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
155Device* StaticDeviceMgr::HostCPU() const { return cpu_device_; }
156
157} // namespace tensorflow
158