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 | #ifndef TENSORFLOW_CORE_UTIL_DEVICE_NAME_UTILS_H_ |
17 | #define TENSORFLOW_CORE_UTIL_DEVICE_NAME_UTILS_H_ |
18 | |
19 | #include <string> |
20 | |
21 | #include "tensorflow/core/lib/core/status.h" |
22 | #include "tensorflow/core/lib/core/stringpiece.h" |
23 | |
24 | namespace tensorflow { |
25 | |
26 | // In TensorFlow a device name is a string of the following form: |
27 | // /job:<name>/replica:<replica>/task:<task>/device:<type>:<device_num> |
28 | // |
29 | // <name> is a short identifier conforming to the regexp |
30 | // [a-zA-Z][_a-zA-Z]* |
31 | // <type> is a supported device type (e.g. 'cpu' or 'gpu') |
32 | // <replica>, <task>, <device_num> are small non-negative integers and are |
33 | // densely allocated (except in tests). |
34 | // |
35 | // For some purposes, we also allow device patterns, which can specify |
36 | // some or none of the specific fields above, with missing components, |
37 | // or "<component>:*" indicating "any value allowed for that component. |
38 | // |
39 | // For example: |
40 | // "/job:param_server" - Consider any devices in the "param_server" job |
41 | // "/device:cpu:*" - Consider any cpu devices in any job/task/replica |
42 | // "/job:*/replica:*/task:*/device:cpu:*" - Consider any cpu devices in any |
43 | // job/task/replica |
44 | // "/job:w/replica:0/task:0/device:gpu:*" - Consider any gpu devices in |
45 | // replica 0, task 0, of job "w" |
46 | class DeviceNameUtils { |
47 | public: |
48 | // Returns a fully qualified device name given the parameters. |
49 | static std::string FullName(const std::string& job, int replica, int task, |
50 | const std::string& type, int id); |
51 | |
52 | struct ParsedName { |
53 | void Clear() { |
54 | has_job = false; |
55 | has_replica = false; |
56 | has_task = false; |
57 | has_type = false; |
58 | has_id = false; |
59 | job.clear(); |
60 | replica = 0; |
61 | task = 0; |
62 | type.clear(); |
63 | id = 0; |
64 | } |
65 | |
66 | bool operator==(const ParsedName& other) const { |
67 | return (has_job ? (other.has_job && job == other.job) : !other.has_job) && |
68 | (has_replica ? (other.has_replica && replica == other.replica) |
69 | : !other.has_replica) && |
70 | (has_task ? (other.has_task && task == other.task) |
71 | : !other.has_task) && |
72 | (has_type ? (other.has_type && type == other.type) |
73 | : !other.has_type) && |
74 | (has_id ? (other.has_id && id == other.id) : !other.has_id); |
75 | } |
76 | |
77 | bool operator!=(const ParsedName& other) const { |
78 | return !operator==(other); |
79 | } |
80 | |
81 | bool has_job = false; |
82 | std::string job; |
83 | bool has_replica = false; |
84 | int replica = 0; |
85 | bool has_task = false; |
86 | int task = 0; |
87 | bool has_type = false; |
88 | std::string type; |
89 | bool has_id = false; |
90 | int id = 0; |
91 | }; |
92 | |
93 | // Parses the device name, first as a full name, then, if it fails, as a |
94 | // global one. Returns `false` if both attempts fail. |
95 | static bool ParseFullOrLocalName(StringPiece fullname, ParsedName* parsed); |
96 | |
97 | // Parses "fullname" into "*parsed". Returns true iff succeeds. |
98 | // Legacy names like "/cpu:0" that don't contain "device", |
99 | // are parsed to mean their current counterparts "/device:CPU:0". More |
100 | // specifically, the lower case "cpu" and "gpu" is capitalized and "device" |
101 | // is added. "/tpu:0" is not treated the same way - it has use the current |
102 | // full syntax. |
103 | // Also, note that lower case "cpu" and "gpu" device types in current syntax |
104 | // are not capitalized. For example, "/device:CPU:0" is different from |
105 | // "/device:cpu:0" |
106 | static bool ParseFullName(StringPiece fullname, ParsedName* parsed); |
107 | |
108 | // Canonicalizes "fullname" into "*canonical_name". Uses a fully specified |
109 | // basename to fill in fields that are missing. Accepts both legacy, newer |
110 | // and local versions of the device spec. Returns the newer version of the |
111 | // device spec. If we were unable to interpret / parse "fullname" returns |
112 | // an error and *canonical_name is set to "". |
113 | static Status CanonicalizeDeviceName(StringPiece fullname, |
114 | StringPiece basename, |
115 | std::string* canonical_name); |
116 | |
117 | // Returns true if "name" specifies any non-trivial constraint on the device. |
118 | static bool HasSomeDetails(const ParsedName& name) { |
119 | return name.has_job || name.has_replica || name.has_task || name.has_type || |
120 | name.has_id; |
121 | } |
122 | |
123 | // Returns true if more_specific is a specification of |
124 | // less_specific, i.e. everywhere that less-specific has a |
125 | // non-wildcard component value, more_specific has the same value |
126 | // for that component. |
127 | static bool IsSpecification(const ParsedName& less_specific, |
128 | const ParsedName& more_specific); |
129 | |
130 | // Makes minimal changes to more_specific so that it becomes a |
131 | // specification of less_specific. |
132 | static void EnsureSpecification(ParsedName* more_specific, |
133 | const ParsedName& less_specific); |
134 | |
135 | // Like IsSpecification, but the second argument "name" must have a |
136 | // non-wildcard value for all of its components. |
137 | static bool IsCompleteSpecification(const ParsedName& pattern, |
138 | const ParsedName& name); |
139 | |
140 | // True iff there exists any possible device name that is a specification of |
141 | // both "a" and "b". |
142 | static bool AreCompatibleDevNames(const ParsedName& a, const ParsedName& b); |
143 | |
144 | // Merges the device specifications in "*target" and "other", and |
145 | // stores the result in "*target". Returns OK if "*target" and |
146 | // "other" are compatible, otherwise returns an error. |
147 | static Status MergeDevNames(ParsedName* target, const ParsedName& other) { |
148 | return MergeDevNames(target, other, false); |
149 | } |
150 | static Status MergeDevNames(ParsedName* target, const ParsedName& other, |
151 | bool allow_soft_placement); |
152 | // Same as MergeDevNames with allow_soft_placement=true, but instead of |
153 | // clearing conflicting fields, overrides them with `other`'s values. |
154 | static Status MergeOverrideDevNames(ParsedName* target, |
155 | const ParsedName& other); |
156 | |
157 | // Merges the device specifications in "*target" and "other", and |
158 | // stores the result in "*target" by setting all unset values in target with |
159 | // corresponding set ones in other. |
160 | static void MergeUnsetDevNames(ParsedName* target, const ParsedName& other); |
161 | |
162 | // Returns true iff devices identified by 'src' and 'dst' are in the |
163 | // same address space. |
164 | static bool IsSameAddressSpace(StringPiece src, StringPiece dst); |
165 | static bool IsSameAddressSpace(const ParsedName& src, const ParsedName& dst); |
166 | |
167 | // Returns true iff devices identified by 'a' and 'b' are in different |
168 | // address space. |
169 | static bool IsDifferentAddressSpace(const ParsedName& a, const ParsedName& b); |
170 | |
171 | // Returns the an address space specification containing only the |
172 | // job/replica/task of the given name. |
173 | static const ParsedName AddressSpace(const ParsedName& name); |
174 | |
175 | // Returns the local device given its "type" and "id". |
176 | static std::string LocalName(StringPiece type, int id); |
177 | |
178 | // Returns a short local device name (cpu:0, gpu:1, etc) based on |
179 | // the given fullname. |
180 | static std::string LocalName(StringPiece fullname); |
181 | |
182 | // If "name" is a valid local device name (cpu:0, gpu:1, etc.), |
183 | // fills in parsed.type and parsed.id accordingly. Returns true iff |
184 | // succeeds. |
185 | static bool ParseLocalName(StringPiece name, ParsedName* parsed); |
186 | |
187 | // Splits a fully-qualified device name into a task identifier and a |
188 | // relative device identifier. It first parses "name" using |
189 | // ParseFullName(), then assigns *task with everything except for |
190 | // the local device component, and assigns the relative device |
191 | // component into *device. This function will still return true if |
192 | // the task component is empty, but it requires the relative device |
193 | // component to be fully specified. |
194 | static bool SplitDeviceName(StringPiece name, std::string* task, |
195 | std::string* device); |
196 | |
197 | // Get the task name from ParsedName. Return false if the task component is |
198 | // not fully specified. |
199 | static bool GetTaskName(const ParsedName& pn, std::string* task); |
200 | |
201 | static std::string ParsedNameToString(const ParsedName& pn); |
202 | |
203 | // Returns canonical and legacy full names for the given parsed |
204 | // device name 'pn'. The returned string names are often useful to |
205 | // look up devices from a mapping. |
206 | static std::vector<string> GetNamesForDeviceMappings(const ParsedName& pn); |
207 | |
208 | // Returns canonical and legacy local names for the given parsed device name |
209 | // 'pn'. The returned string names are often useful to look up devices from a |
210 | // mapping. |
211 | static std::vector<string> GetLocalNamesForDeviceMappings( |
212 | const ParsedName& pn); |
213 | |
214 | // Returns name of the CPU:0 device on the same host as the device |
215 | // `device_name`. |
216 | static Status DeviceNameToCpuDeviceName(const std::string& device_name, |
217 | std::string* host_device_name); |
218 | }; |
219 | |
220 | std::ostream& operator<<(std::ostream& os, |
221 | const DeviceNameUtils::ParsedName& x); |
222 | |
223 | } // namespace tensorflow |
224 | |
225 | #endif // TENSORFLOW_CORE_UTIL_DEVICE_NAME_UTILS_H_ |
226 | |