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 // 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