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 string FullName(const string& job, int replica, int task, 50 const 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 has_job = false; 78 string job; 79 bool has_replica = false; 80 int replica = 0; 81 bool has_task = false; 82 int task = 0; 83 bool has_type = false; 84 string type; 85 bool has_id = false; 86 int id = 0; 87 }; 88 // Parses "fullname" into "*parsed". Returns true iff succeeds. 89 // Legacy names like "/cpu:0" that don't contain "device", 90 // are parsed to mean their current counterparts "/device:CPU:0". More 91 // specifically, the lower case "cpu" and "gpu" is capitalized and "device" 92 // is added. "/tpu:0" is not treated the same way - it has use the current 93 // full syntax. 94 // Also, note that lower case "cpu" and "gpu" device types in current syntax 95 // are not capitalized. For example, "/device:CPU:0" is different from 96 // "/device:cpu:0" 97 static bool ParseFullName(StringPiece fullname, ParsedName* parsed); 98 99 // Canonicalizes "fullname" into "*canonical_name". Uses a fully specified 100 // basename to fill in fields that are missing. Accepts both legacy, newer 101 // and local versions of the device spec. Returns the newer version of the 102 // device spec. If we were unable to interpret / parse "fullname" returns 103 // an error and *canonical_name is set to "". 104 static Status CanonicalizeDeviceName(StringPiece fullname, 105 StringPiece basename, 106 string* canonical_name); 107 108 // Returns true if "name" specifies any non-trivial constraint on the device. HasSomeDetails(const ParsedName & name)109 static bool HasSomeDetails(const ParsedName& name) { 110 return name.has_job || name.has_replica || name.has_task || name.has_type || 111 name.has_id; 112 } 113 114 // Returns true if more_specific is a specification of 115 // less_specific, i.e. everywhere that less-specific has a 116 // non-wildcard component value, more_specific has the same value 117 // for that component. 118 static bool IsSpecification(const ParsedName& less_specific, 119 const ParsedName& more_specific); 120 121 // Makes minimal changes to more_specific so that it becomes a 122 // specification of less_specific. 123 static void EnsureSpecification(ParsedName* more_specific, 124 const ParsedName& less_specific); 125 126 // Like IsSpecification, but the second argument "name" must have a 127 // non-wildcard value for all of its components. 128 static bool IsCompleteSpecification(const ParsedName& pattern, 129 const ParsedName& name); 130 131 // True iff there exists any possible device name that is a specification of 132 // both "a" and "b". 133 static bool AreCompatibleDevNames(const ParsedName& a, const ParsedName& b); 134 135 // Merges the device specifications in "*target" and "other", and 136 // stores the result in "*target". Returns OK if "*target" and 137 // "other" are compatible, otherwise returns an error. MergeDevNames(ParsedName * target,const ParsedName & other)138 static Status MergeDevNames(ParsedName* target, const ParsedName& other) { 139 return MergeDevNames(target, other, false); 140 } 141 static Status MergeDevNames(ParsedName* target, const ParsedName& other, 142 bool allow_soft_placement); 143 // Same as MergeDevNames with allow_soft_placement=true, but instead of 144 // clearing conflicting fields, overrides them with `other`'s values. 145 static Status MergeOverrideDevNames(ParsedName* target, 146 const ParsedName& other); 147 148 // Returns true iff devices identified by 'src' and 'dst' are in the 149 // same address space. 150 static bool IsSameAddressSpace(StringPiece src, StringPiece dst); 151 static bool IsSameAddressSpace(const ParsedName& src, const ParsedName& dst); 152 153 // Returns true iff devices identified by 'a' and 'b' are in different 154 // address space. 155 static bool IsDifferentAddressSpace(const ParsedName& a, const ParsedName& b); 156 157 // Returns the an address space specification containing only the 158 // job/replica/task of the given name. 159 static const ParsedName AddressSpace(const ParsedName& name); 160 161 // Returns the local device given its "type" and "id". 162 static string LocalName(StringPiece type, int id); 163 164 // Returns a short local device name (cpu:0, gpu:1, etc) based on 165 // the given fullname. 166 static string LocalName(StringPiece fullname); 167 168 // If "name" is a valid local device name (cpu:0, gpu:1, etc.), 169 // fills in parsed.type and parsed.id accordingly. Returns true iff 170 // succeeds. 171 static bool ParseLocalName(StringPiece name, ParsedName* parsed); 172 173 // Splits a fully-qualified device name into a task identifier and a 174 // relative device identifier. It first parses "name" using 175 // ParseFullName(), then assigns *task with everything except for 176 // the local device component, and assigns the relative device 177 // component into *device. This function will still return true if 178 // the task component is empty, but it requires the relative device 179 // component to be fully specified. 180 static bool SplitDeviceName(StringPiece name, string* task, string* device); 181 182 // Get the task name from ParsedName. Return false if the task component is 183 // not fully specified. 184 static bool GetTaskName(const ParsedName& pn, string* task); 185 186 static string ParsedNameToString(const ParsedName& pn); 187 188 // Returns canonical and legacy full names for the given parsed 189 // device name 'pn'. The returned string names are often useful to 190 // look up devices from a mapping. 191 static std::vector<string> GetNamesForDeviceMappings(const ParsedName& pn); 192 193 // Returns canonical and legacy local names for the given parsed device name 194 // 'pn'. The returned string names are often useful to look up devices from a 195 // mapping. 196 static std::vector<string> GetLocalNamesForDeviceMappings( 197 const ParsedName& pn); 198 199 // Returns name of the CPU:0 device on the same host as the device 200 // `device_name`. 201 static Status DeviceNameToCpuDeviceName(const string& device_name, 202 string* host_device_name); 203 }; 204 205 std::ostream& operator<<(std::ostream& os, 206 const DeviceNameUtils::ParsedName& x); 207 208 } // namespace tensorflow 209 210 #endif // TENSORFLOW_CORE_UTIL_DEVICE_NAME_UTILS_H_ 211