From 9205bef5560e5bae485ae276206331db3d3f8b54 Mon Sep 17 00:00:00 2001
From: xla authors <google-ml-automation@google.com>
Date: Wed, 5 Feb 2025 05:44:03 -0800
Subject: [PATCH] Parse XLA_FLAGS environment variable every time,
 conditionally on xla_flags_reset flag.

PiperOrigin-RevId: 723476384
---
 xla/debug_options_flags.cc  |  8 ++++++++
 xla/parse_flags_from_env.cc | 13 +++++++++----
 xla/parse_flags_from_env.h  |  6 ++++--
 xla/xla.proto               |  3 ++-
 4 files changed, 23 insertions(+), 7 deletions(-)

diff --git a/xla/debug_options_flags.cc b/xla/debug_options_flags.cc
index fc4488573b3c0..e710c1796bbf7 100644
--- 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
index 331cccb536be8..b955f69ff32fc 100644
--- 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
index e135a664d93b4..e4f749df5478e 100644
--- a/xla/parse_flags_from_env.h
+++ b/xla/parse_flags_from_env.h
@@ -62,7 +62,8 @@ 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
@@ -70,7 +71,8 @@ void ParseFlagsFromEnvAndDieIfUnknown(absl::string_view envvar,
 //
 // 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
index e67e86c553409..dd9dbd7ded053 100644
--- 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.