• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
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 {
ClearParsedName53     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.
HasSomeDetails(const ParsedName & name)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.
MergeDevNames(ParsedName * target,const ParsedName & other)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   // Returns true iff devices identified by 'src' and 'dst' are in the
158   // same address space.
159   static bool IsSameAddressSpace(StringPiece src, StringPiece dst);
160   static bool IsSameAddressSpace(const ParsedName& src, const ParsedName& dst);
161 
162   // Returns true iff devices identified by 'a' and 'b' are in different
163   // address space.
164   static bool IsDifferentAddressSpace(const ParsedName& a, const ParsedName& b);
165 
166   // Returns the an address space specification containing only the
167   // job/replica/task of the given name.
168   static const ParsedName AddressSpace(const ParsedName& name);
169 
170   // Returns the local device given its "type" and "id".
171   static std::string LocalName(StringPiece type, int id);
172 
173   // Returns a short local device name (cpu:0, gpu:1, etc) based on
174   // the given fullname.
175   static std::string LocalName(StringPiece fullname);
176 
177   // If "name" is a valid local device name (cpu:0, gpu:1, etc.),
178   // fills in parsed.type and parsed.id accordingly. Returns true iff
179   // succeeds.
180   static bool ParseLocalName(StringPiece name, ParsedName* parsed);
181 
182   // Splits a fully-qualified device name into a task identifier and a
183   // relative device identifier. It first parses "name" using
184   // ParseFullName(), then assigns *task with everything except for
185   // the local device component, and assigns the relative device
186   // component into *device.  This function will still return true if
187   // the task component is empty, but it requires the relative device
188   // component to be fully specified.
189   static bool SplitDeviceName(StringPiece name, std::string* task,
190                               std::string* device);
191 
192   // Get the task name from ParsedName. Return false if the task component is
193   // not fully specified.
194   static bool GetTaskName(const ParsedName& pn, std::string* task);
195 
196   static std::string ParsedNameToString(const ParsedName& pn);
197 
198   // Returns canonical and legacy full names for the given parsed
199   // device name 'pn'. The returned string names are often useful to
200   // look up devices from a mapping.
201   static std::vector<string> GetNamesForDeviceMappings(const ParsedName& pn);
202 
203   // Returns canonical and legacy local names for the given parsed device name
204   // 'pn'. The returned string names are often useful to look up devices from a
205   // mapping.
206   static std::vector<string> GetLocalNamesForDeviceMappings(
207       const ParsedName& pn);
208 
209   // Returns name of the CPU:0 device on the same host as the device
210   // `device_name`.
211   static Status DeviceNameToCpuDeviceName(const std::string& device_name,
212                                           std::string* host_device_name);
213 };
214 
215 std::ostream& operator<<(std::ostream& os,
216                          const DeviceNameUtils::ParsedName& x);
217 
218 }  // namespace tensorflow
219 
220 #endif  // TENSORFLOW_CORE_UTIL_DEVICE_NAME_UTILS_H_
221