• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1 /* Copyright 2017 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 // This module exports ParseFlagsFromEnvAndDieIfUnknown(), which allows other
17 // modules to parse flags from an environtment variable, or a file named by the
18 // environment variable.
19 
20 #include <stdio.h>
21 #include <stdlib.h>
22 #include <string.h>
23 #include <memory>
24 #include <unordered_map>
25 #include <vector>
26 
27 #include "absl/strings/ascii.h"
28 #include "absl/strings/str_format.h"
29 #include "absl/strings/str_join.h"
30 #include "absl/types/span.h"
31 #include "tensorflow/compiler/xla/parse_flags_from_env.h"
32 #include "tensorflow/compiler/xla/types.h"
33 #include "tensorflow/core/platform/logging.h"
34 #include "tensorflow/core/platform/macros.h"
35 #include "tensorflow/core/platform/mutex.h"
36 #include "tensorflow/core/platform/types.h"
37 #include "tensorflow/core/util/command_line_flags.h"
38 
39 namespace xla {
40 
41 static const char kWS[] = " \t\r\n";  // whitespace
42 
43 // The following struct represents an argv[]-style array, parsed
44 // from data gleaned from the environment.
45 //
46 // As usual, an anonymous namespace is advisable to avoid
47 // constructor/destructor collisions with other "private" types
48 // in the same named namespace.
49 namespace {
50 
51 // Functor which deletes objects by calling `free`.  Necessary to free strdup'ed
52 // strings created by AppendToEnvArgv.
53 struct FreeDeleter {
operator ()xla::__anonfa92b6720111::FreeDeleter54   void operator()(char* ptr) { free(ptr); }
55 };
56 
57 struct EnvArgv {
EnvArgvxla::__anonfa92b6720111::EnvArgv58   EnvArgv() : initialized(false), argc(0) {}
59   bool initialized;         // whether the other fields have been set.
60   int argc;                 // elements used in argv[]
61   std::vector<char*> argv;  // flag arguments parsed from environment string.
62   // saved values from argv[] to avoid leaks
63   std::vector<std::unique_ptr<char, FreeDeleter>> argv_save;
64 };
65 }  // anonymous namespace
66 
67 // Append the string s0[0, .., s0len-1] concatenated with s1[0, .., s1len-1] as
68 // a newly allocated nul-terminated string to the array *a.  If s0==nullptr, a
69 // nullptr is appended without increasing a->argc.
AppendToEnvArgv(const char * s0,size_t s0len,const char * s1,size_t s1len,EnvArgv * a)70 static void AppendToEnvArgv(const char* s0, size_t s0len, const char* s1,
71                             size_t s1len, EnvArgv* a) {
72   if (s0 == nullptr) {
73     a->argv.push_back(nullptr);
74     a->argv_save.push_back(nullptr);
75   } else {
76     string s = string(s0, s0len) + string(s1, s1len);
77     char* str = strdup(s.c_str());
78     a->argv.push_back(str);
79     a->argv_save.emplace_back(str);
80     a->argc++;
81   }
82 }
83 
84 // Like s.find_first_of(x, pos), but return s.size() when find_first_of() would
85 // return string::npos.  This avoids if-statements elsewhere.
FindFirstOf(const string & s,const char * x,size_t pos)86 static size_t FindFirstOf(const string& s, const char* x, size_t pos) {
87   size_t result = s.find_first_of(x, pos);
88   return result == string::npos ? s.size() : result;
89 }
90 
91 // Like s.find_first_not_of(x, pos), but return s.size() when
92 // find_first_not_of() would return string::npos.  This avoids if-statements
93 // elsewhere.
FindFirstNotOf(const string & s,const char * x,size_t pos)94 static size_t FindFirstNotOf(const string& s, const char* x, size_t pos) {
95   size_t result = s.find_first_not_of(x, pos);
96   return result == string::npos ? s.size() : result;
97 }
98 
99 // Given a string containing flags, parse them into the XLA command line flags.
100 // The parse is best effort, and gives up on the first syntax error.
ParseArgvFromString(const string & flag_str,EnvArgv * a)101 static void ParseArgvFromString(const string& flag_str, EnvArgv* a) {
102   size_t b = FindFirstNotOf(flag_str, kWS, 0);
103   while (b != flag_str.size() && flag_str[b] == '-') {
104     // b is the index of the start of a flag.
105     // Set e to the index just past the end of the flag.
106     size_t e = b;
107     while (e != flag_str.size() && isascii(flag_str[e]) &&
108            (strchr("-_", flag_str[e]) != nullptr ||
109             absl::ascii_isalnum(flag_str[e]))) {
110       e++;
111     }
112     if (e != flag_str.size() && flag_str[e] == '=' &&
113         e + 1 != flag_str.size() && strchr("'\"", flag_str[e + 1]) != nullptr) {
114       // A flag of the form  --flag="something in double or single quotes"
115       int c;
116       e++;  // point just past '='
117       size_t eflag = e;
118       char quote = flag_str[e];
119       e++;  // point just past quote
120       // Put in value the string with quotes removed.
121       string value;
122       for (; e != flag_str.size() && (c = flag_str[e]) != quote; e++) {
123         if (quote == '"' && c == '\\' && e + 1 != flag_str.size()) {
124           // Handle backslash in double quoted strings.  They are literal in
125           // single-quoted strings.
126           e++;
127           c = flag_str[e];
128         }
129         value += c;
130       }
131       if (e != flag_str.size()) {  // skip final " or '
132         e++;
133       }
134       AppendToEnvArgv(flag_str.data() + b, eflag - b, value.data(),
135                       value.size(), a);
136     } else {  // A flag without a quoted value.
137       e = FindFirstOf(flag_str, kWS, e);
138       AppendToEnvArgv(flag_str.data() + b, e - b, "", 0, a);
139     }
140     b = FindFirstNotOf(flag_str, kWS, e);
141   }
142 }
143 
144 // Call ParseArgvFromString(..., a) on a string derived from the setting of the
145 // environment variable `envvar`, or a file it points to.
SetArgvFromEnv(absl::string_view envvar,EnvArgv * a)146 static void SetArgvFromEnv(absl::string_view envvar, EnvArgv* a) {
147   if (!a->initialized) {
148     static const char kDummyArgv[] = "<argv[0]>";
149     AppendToEnvArgv(kDummyArgv, strlen(kDummyArgv), nullptr, 0,
150                     a);  // dummy argv[0]
151     const char* env = getenv(string(envvar).c_str());
152     if (env == nullptr || env[0] == '\0') {
153       // nothing
154     } else if (env[strspn(env, kWS)] == '-') {  // flags in env var value
155       ParseArgvFromString(env, a);
156     } else {  // assume it's a file name
157       FILE* fp = fopen(env, "r");
158       if (fp != nullptr) {
159         string str;
160         char buf[512];
161         int n;
162         while ((n = fread(buf, 1, sizeof(buf), fp)) > 0) {
163           str.append(buf, n);
164         }
165         fclose(fp);
166         ParseArgvFromString(str, a);
167       }
168     }
169     AppendToEnvArgv(nullptr, 0, nullptr, 0, a);  // add trailing nullptr to *a.
170     a->initialized = true;
171   }
172 }
173 
174 // The simulated argv[] parsed from the environment, one for each different
175 // environment variable we've seen.
EnvArgvs()176 static std::unordered_map<string, EnvArgv>& EnvArgvs() {
177   static auto* env_argvs = new std::unordered_map<string, EnvArgv>();
178   return *env_argvs;
179 }
180 
181 // Used to protect accesses to env_argvs.
182 static tensorflow::mutex env_argv_mu(tensorflow::LINKER_INITIALIZED);
183 
ParseFlagsFromEnvAndDieIfUnknown(absl::string_view envvar,const std::vector<tensorflow::Flag> & flag_list)184 bool ParseFlagsFromEnvAndDieIfUnknown(
185     absl::string_view envvar, const std::vector<tensorflow::Flag>& flag_list) {
186   tensorflow::mutex_lock lock(env_argv_mu);
187   auto* env_argv = &EnvArgvs()[string(envvar)];
188   SetArgvFromEnv(envvar, env_argv);  // a no-op if already initialized
189 
190   if (VLOG_IS_ON(1)) {
191     VLOG(1) << "For env var " << envvar << " found arguments:";
192     for (int i = 0; i < env_argv->argc; i++) {
193       VLOG(1) << "  argv[" << i << "] = " << env_argv->argv[i];
194     }
195   }
196 
197   bool result =
198       tensorflow::Flags::Parse(&env_argv->argc, &env_argv->argv[0], flag_list);
199 
200   // There's always at least one unparsed argc, namely the fake argv[0].
201   if (result && env_argv->argc != 1) {
202     // Skip the first argv, which is the fake argv[0].
203     auto unknown_flags = absl::MakeSpan(env_argv->argv);
204     unknown_flags.remove_prefix(1);
205 
206     // Some flags are set on XLA_FLAGS, others on TF_XLA_FLAGS.  If we find an
207     // unrecognized flag, suggest the alternative.
208     string alternate_envvar;
209     if (envvar == "TF_XLA_FLAGS") {
210       alternate_envvar = "XLA_FLAGS";
211     } else if (envvar == "XLA_FLAGS") {
212       alternate_envvar = "TF_XLA_FLAGS";
213     }
214     string did_you_mean;
215     if (!alternate_envvar.empty()) {
216       did_you_mean = absl::StrFormat(
217           "\nPerhaps you meant to specify these on the %s envvar?",
218           alternate_envvar);
219     }
220 
221     LOG(FATAL) << "Unknown flag" << (unknown_flags.size() > 1 ? "s" : "")
222                << " in " << envvar << ": " << absl::StrJoin(unknown_flags, " ")
223                << did_you_mean;
224     return false;
225   }
226   return result;
227 }
228 
229 // Testing only.
230 //
231 // Resets the env_argv struct so that subsequent calls to
232 // ParseFlagsFromEnvAndDieIfUnknown() will parse the environment variable (or
233 // the file it points to) anew, and set *pargc, and *pargv to point to the
234 // internal locations of the argc and argv constructed from the environment.
ResetFlagsFromEnvForTesting(absl::string_view envvar,int ** pargc,std::vector<char * > ** pargv)235 void ResetFlagsFromEnvForTesting(absl::string_view envvar, int** pargc,
236                                  std::vector<char*>** pargv) {
237   tensorflow::mutex_lock lock(env_argv_mu);
238   EnvArgvs().erase(string(envvar));
239   auto& env_argv = EnvArgvs()[string(envvar)];
240   *pargc = &env_argv.argc;
241   *pargv = &env_argv.argv;
242 }
243 
244 }  // namespace xla
245