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