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