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_UTIL_DEVICE_NAME_UTILS_H_ 17 #define TENSORFLOW_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 static bool ParseFullName(StringPiece fullname, ParsedName* parsed); 90 91 // Canonicalizes "fullname". Accepts both legacy, newer and local versions of 92 // the device spec. Returns the newer version of the device spec. If we were 93 // unable to interpret / parse "fullname" returns "". 94 static string CanonicalizeDeviceName(StringPiece fullname); 95 96 // Returns true if "name" specifies any non-trivial constraint on the device. HasSomeDetails(const ParsedName & name)97 static bool HasSomeDetails(const ParsedName& name) { 98 return name.has_job || name.has_replica || name.has_task || name.has_type || 99 name.has_id; 100 } 101 102 // Returns true if more_specific is a specification of 103 // less_specific, i.e. everywhere that less-specific has a 104 // non-wildcard component value, more_specific has the same value 105 // for that component. 106 static bool IsSpecification(const ParsedName& less_specific, 107 const ParsedName& more_specific); 108 109 // Like IsSpecification, but the second argument "name" must have a 110 // non-wildcard value for all of its components. 111 static bool IsCompleteSpecification(const ParsedName& pattern, 112 const ParsedName& name); 113 114 // True iff there exists any possible complete device name that is 115 // a specification of both "a" and "b". AreCompatibleDevNames(const ParsedName & a,const ParsedName & b)116 static inline bool AreCompatibleDevNames(const ParsedName& a, 117 const ParsedName& b) { 118 return IsSpecification(a, b) || IsSpecification(b, a); 119 } 120 121 // Merges the device specifications in "*target" and "other", and 122 // stores the result in "*target". Returns OK if "*target" and 123 // "other" are compatible, otherwise returns an error. MergeDevNames(ParsedName * target,const ParsedName & other)124 static Status MergeDevNames(ParsedName* target, const ParsedName& other) { 125 return MergeDevNames(target, other, false); 126 } 127 static Status MergeDevNames(ParsedName* target, const ParsedName& other, 128 bool allow_soft_placement); 129 130 // Returns true iff devices identified by 'src' and 'dst' are in the 131 // same address space. 132 static bool IsSameAddressSpace(StringPiece src, StringPiece dst); 133 static bool IsSameAddressSpace(const ParsedName& src, const ParsedName& dst); 134 135 // Returns the local device given its "type" and "id". 136 static string LocalName(StringPiece type, int id); 137 138 // Returns a short local device name (cpu:0, gpu:1, etc) based on 139 // the given fullname. 140 static string LocalName(StringPiece fullname); 141 142 // If "name" is a valid local device name (cpu:0, gpu:1, etc.), 143 // fills in parsed.type and parsed.id accordingly. Returns true iff 144 // succeeds. 145 static bool ParseLocalName(StringPiece name, ParsedName* parsed); 146 147 // Splits a fully-qualified device name into a task identifier and a 148 // relative device identifier. It first parses "name" using 149 // ParseFullName(), then assigns *task with everything except for 150 // the local device component, and assigns the relative device 151 // component into *device. This function will still return true if 152 // the task component is empty, but it requires the relative device 153 // component to be fully specified. 154 static bool SplitDeviceName(StringPiece name, string* task, string* device); 155 156 static string ParsedNameToString(const ParsedName& pn); 157 158 // Returns canonical and legacy full names for the given parsed 159 // device name 'pn'. The returned string names are often useful to 160 // look up devices from a mapping. 161 static std::vector<string> GetNamesForDeviceMappings(const ParsedName& pn); 162 163 // Returns canonical and legacy local names for the given parsed device name 164 // 'pn'. The returned string names are often useful to look up devices from a 165 // mapping. 166 static std::vector<string> GetLocalNamesForDeviceMappings( 167 const ParsedName& pn); 168 }; 169 170 } // namespace tensorflow 171 172 #endif // TENSORFLOW_UTIL_DEVICE_NAME_UTILS_H_ 173