PR #25166: Add slow operation alarm for argument initialization in the HLO runner.

Imported from GitHub PR #25166

Copybara import of the project:

--
faa4fd1 by Ilia Sergachev <isergachev@nvidia.com>:

Add slow operation alarm for argument initialization in the HLO runner.

Merging this change closes #25166

FUTURE_COPYBARA_INTEGRATE_REVIEW=#25166 from openxla:add_slow_arg_init_alarm faa4fd1
PiperOrigin-RevId: 750111996

Author: sergachev

xla/tools/multihost_hlo_runner/BUILD

@@ -168,6 +168,7 @@ cc_library(
         "//xla/service:hlo_module_config",
         "//xla/service:hlo_module_util",
         "//xla/service:hlo_proto_cc",
+        "//xla/service:slow_operation_alarm",
         "//xla/tests:test_utils",
         "//xla/tools:hlo_control_flow_flattening",
         "//xla/tsl/platform:env",

xla/tools/multihost_hlo_runner/functional_hlo_runner.cc

@@ -65,6 +65,7 @@ limitations under the License.
 #include "xla/service/hlo.pb.h"
 #include "xla/service/hlo_module_config.h"
 #include "xla/service/hlo_module_util.h"
+#include "xla/service/slow_operation_alarm.h"
 #include "xla/shape_layout.h"
 #include "xla/shape_util.h"
 #include "xla/status_macros.h"
@@ -1161,6 +1162,11 @@ FunctionalHloRunner::CreateArgumentsOnDevice(
         client, executable, running_options, flatten_arguments);
   }
 
+  SlowOperationAlarm alarm(
+      absl::Seconds(5),
+      absl::StrFormat("Argument initialization is slow. Consider changing "
+                      "--hlo_argument_mode."));
+
   absl::Span<PjRtDevice* const> addressable_devices =
       executable->addressable_devices();
   size_t num_addressable_devices = addressable_devices.size();
@@ -1185,7 +1191,7 @@ FunctionalHloRunner::CreateArgumentsOnDevice(
           ModuleArgumentMode::kUseZerosAsInput;
 
   for (int i = 0; i < num_addressable_devices; ++i) {
-    VLOG(3) << "Creating fake argument for device " << i;
+    VLOG(3) << "Creating fake arguments for device " << i;
     LiteralVec& argument_literals =
         per_device_argument_literals[addressable_devices[i]->id()];
     int executable_idx = hlo_modules.size() == 1