Parse XLA_FLAGS environment variable every time, conditionally on xla…

…_flags_reset flag. PiperOrigin-RevId: 723476384
openxla · Feb 5, 2025 · 9205bef · 9205bef
1 parent 8cd6948
commit 9205bef
Show file tree

Hide file tree

Showing 4 changed files with 23 additions and 7 deletions.
diff --git a/xla/debug_options_flags.cc b/xla/debug_options_flags.cc
@@ -1031,6 +1031,10 @@ void MakeDebugOptionsFlags(std::vector<tsl::Flag>* flag_list,
       "and \"test_undeclared_outputs_dir\" have a special meaning: They cause "
       "us to dump into the directory specified by the environment variable "
       "TEST_UNDECLARED_OUTPUTS_DIR."));
+  flag_list->push_back(tsl::Flag(
+      "xla_flags_reset", bool_setter_for(&DebugOptions::set_xla_flags_reset),
+      debug_options->xla_flags_reset(),
+      "Whether to reset XLA_FLAGS next time to parse."));
   flag_list->push_back(tsl::Flag(
       "xla_gpu_unsupported_annotate_with_emitter_loc",
       bool_setter_for(
@@ -2199,6 +2203,10 @@ void AppendDebugOptionsFlags(std::vector<tsl::Flag>* flag_list,
 
 xla::DebugOptions GetDebugOptionsFromFlags() {
   absl::call_once(flags_init, &AllocateFlags, nullptr);
+  if (flag_values->xla_flags_reset()) {
+    ParseFlagsFromEnvAndDieIfUnknown("XLA_FLAGS", *flag_objects,
+                                     /*reset_envvar=*/true);
+  }
   return *flag_values;
 }
 

diff --git a/xla/parse_flags_from_env.cc b/xla/parse_flags_from_env.cc
@@ -193,14 +193,19 @@ static absl::Mutex env_argv_mu(absl::kConstInit);
 static void DieIfEnvHasUnknownFlagsLeft(absl::string_view envvar);
 
 void ParseFlagsFromEnvAndDieIfUnknown(absl::string_view envvar,
-                                      const std::vector<tsl::Flag>& flag_list) {
-  ParseFlagsFromEnvAndIgnoreUnknown(envvar, flag_list);
+                                      const std::vector<tsl::Flag>& flag_list,
+                                      const bool reset_envvar) {
+  ParseFlagsFromEnvAndIgnoreUnknown(envvar, flag_list, reset_envvar);
   DieIfEnvHasUnknownFlagsLeft(envvar);
 }
 
-void ParseFlagsFromEnvAndIgnoreUnknown(
-    absl::string_view envvar, const std::vector<tsl::Flag>& flag_list) {
+void ParseFlagsFromEnvAndIgnoreUnknown(absl::string_view envvar,
+                                       const std::vector<tsl::Flag>& flag_list,
+                                       const bool reset_envvar) {
   absl::MutexLock lock(&env_argv_mu);
+  if (reset_envvar) {
+    EnvArgvs().erase(envvar);
+  }
   auto* env_argv = &EnvArgvs()[envvar];
   SetArgvFromEnv(envvar, env_argv);  // a no-op if already initialized
 

diff --git a/xla/parse_flags_from_env.h b/xla/parse_flags_from_env.h
@@ -62,15 +62,17 @@ namespace xla {
 // Raises a fatal error if any flags in `envvar` were not recognized, or if flag
 // parsing failed.
 void ParseFlagsFromEnvAndDieIfUnknown(absl::string_view envvar,
-                                      const std::vector<tsl::Flag>& flag_list);
+                                      const std::vector<tsl::Flag>& flag_list,
+                                      bool reset_envvar = false);
 
 // Calls tsl::Flags::Parse(argc, argv, flag_list) against any as yet
 // unrecognized flags passed in the environment variable `envvar`, and returns
 // its return value.
 //
 // Ignores unknown flags, raises a fatal if flag parsing failed.
 void ParseFlagsFromEnvAndIgnoreUnknown(absl::string_view envvar,
-                                       const std::vector<tsl::Flag>& flag_list);
+                                       const std::vector<tsl::Flag>& flag_list,
+                                       bool reset_envvar = false);
 
 // Used only for testing.  Not to be used by clients.
 void ResetFlagsFromEnvForTesting(absl::string_view envvar, int** pargc,

diff --git a/xla/xla.proto b/xla/xla.proto
@@ -906,6 +906,7 @@ message DebugOptions {
 
   // Directory to dump into.
   string xla_dump_to = 109;
+  bool xla_flags_reset = 364;
 
   // If specified, will only dump modules which match this regexp.
   string xla_dump_hlo_module_re = 110;
@@ -1127,7 +1128,7 @@ message DebugOptions {
 
   // Note: when adding a new flag, please add it to one of the hardware-specific
   // or hardware-agnostic sections at the top of this proto message.
-  // Next id: 364
+  // Next id: 365
 
   // Extra options to pass to the compilation backend (e.g. LLVM); specific
   // interpretation of these values is left to the backend.