PR #33414: [ROCm] Add missing keepalive timeout for rbe builds in rocm

Imported from GitHub PR openxla/xla#33414

📝 Summary of Changes
Set keepalive time to prevent connection drop from bazel client during rbe builds

🎯 Justification
Set keepalive time to prevent connection drop from bazel client during rbe builds

🚀 Kind of Contribution
Please remove what does not apply: 🐛 Bug Fix

📊 Benchmark (for Performance Improvements)
Not relevant

🧪 Unit Tests:
Not relevant

🧪 Execution Tests:
Not relevant

Copybara import of the project:

--
76a26b318dac967b17a64abfbbbb00ca26980b5a by Alexandros Theodoridis <[email protected]>:

Add missing keepalive timeout for rbe builds in rocm

Merging this change closes #33414

FUTURE_COPYBARA_INTEGRATE_REVIEW=openxla/xla#33414 from ROCm:ci_add_missing_keepalive_timeout_for_rbe_builds 76a26b318dac967b17a64abfbbbb00ca26980b5a
PiperOrigin-RevId: 827818530

Author: alekstheod

tensorflow/compiler/tf2xla/kernels/all_reduce_op.cc

@@ -94,9 +94,9 @@ class CollectiveReduceV2Op : public XlaOpKernel {
 
  private:
   DataType dtype_ = DT_INVALID;
-  string merge_op_name_;
-  string final_op_name_;
-  string communication_hint_;
+  std::string merge_op_name_;
+  std::string final_op_name_;
+  std::string communication_hint_;
 
   CollectiveReduceV2Op(const CollectiveReduceV2Op&) = delete;
   void operator=(const CollectiveReduceV2Op&) = delete;

tensorflow/compiler/tf2xla/kernels/batch_norm_op.cc

@@ -48,7 +48,7 @@ class FusedBatchNormOp : public XlaOpKernel {
     OP_REQUIRES_OK(ctx, ctx->GetAttr("is_training", &is_training_));
     OP_REQUIRES_OK(
         ctx, ctx->GetAttr("exponential_avg_factor", &exponential_avg_factor_));
-    string data_format_str;
+    std::string data_format_str;
     OP_REQUIRES_OK(ctx, ctx->GetAttr("data_format", &data_format_str));
     OP_REQUIRES(
         ctx, FormatFromString(data_format_str, &data_format_),
@@ -61,7 +61,7 @@ class FusedBatchNormOp : public XlaOpKernel {
                   errors::InvalidArgument(
                       "FusedBatchNormEx supports at most 1 side input."));
       add_side_input_ = (num_side_inputs == 1);
-      string activation_mode;
+      std::string activation_mode;
       OP_REQUIRES_OK(ctx, ctx->GetAttr("activation_mode", &activation_mode));
       OP_REQUIRES(ctx,
                   activation_mode == "Identity" || activation_mode == "Relu",
@@ -249,7 +249,7 @@ class FusedBatchNormGradOp : public XlaOpKernel {
   explicit FusedBatchNormGradOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {
     OP_REQUIRES_OK(ctx, ctx->GetAttr("epsilon", &epsilon_));
     OP_REQUIRES_OK(ctx, ctx->GetAttr("is_training", &is_training_));
-    string data_format_str;
+    std::string data_format_str;
     OP_REQUIRES_OK(ctx, ctx->GetAttr("data_format", &data_format_str));
     OP_REQUIRES(
         ctx, FormatFromString(data_format_str, &data_format_),

tensorflow/compiler/tf2xla/kernels/bcast_ops.cc

@@ -66,9 +66,11 @@ class BCastArgsOp : public XlaOpKernel {
     Tensor output(val_type, TensorShape({len}));
     for (int64_t i = 0; i < len; ++i) {
       if (val_type == DT_INT32) {
-        output.flat<int32>()(i) = static_cast<int32>(bcast.output_shape()[i]);
+        output.flat<int32_t>()(i) =
+            static_cast<int32_t>(bcast.output_shape()[i]);
       } else {
-        output.flat<int64>()(i) = static_cast<int64>(bcast.output_shape()[i]);
+        output.flat<int64_t>()(i) =
+            static_cast<int64_t>(bcast.output_shape()[i]);
       }
     }
     ctx->SetConstantOutput(0, output);
@@ -129,9 +131,9 @@ class BCastGradArgsOp : public XlaOpKernel {
     Tensor constant(val_type, TensorShape({len}));
     for (int64_t i = 0; i < len; ++i) {
       if (val_type == DT_INT32) {
-        constant.flat<int32>()(i) = static_cast<int32>(v[i]);
+        constant.flat<int32_t>()(i) = static_cast<int32_t>(v[i]);
       } else {
-        constant.flat<int64>()(i) = static_cast<int64>(v[i]);
+        constant.flat<int64_t>()(i) = static_cast<int64_t>(v[i]);
       }
     }
     ctx->SetConstantOutput(idx, constant);

tensorflow/compiler/tf2xla/kernels/bias_ops.cc

@@ -28,7 +28,7 @@ namespace {
 class BiasOp : public XlaOpKernel {
  public:
   explicit BiasOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {
-    string data_format;
+    std::string data_format;
     if (ctx->GetAttr("data_format", &data_format).ok()) {
       OP_REQUIRES(ctx, FormatFromString(data_format, &data_format_),
                   errors::InvalidArgument("Invalid data format"));

tensorflow/compiler/tf2xla/kernels/bucketize_op.cc

@@ -55,7 +55,7 @@ class BucketizeOp : public XlaOpKernel {
                                         /*broadcast_dimensions=*/{0}),
                                 xla::S32);
     xla::XlaOp buckets = xla::Reduce(
-        comparison, /*init_value=*/xla::ConstantR0<int32>(builder, 0),
+        comparison, /*init_value=*/xla::ConstantR0<int32_t>(builder, 0),
         /*computation=*/xla::CreateScalarAddComputation(xla::S32, builder),
         /*dimensions_to_reduce=*/{0});
     context->SetOutput(0, buckets);

tensorflow/compiler/tf2xla/kernels/case_op.cc

@@ -66,7 +66,7 @@ XlaCaseOp::GetPrunedBranchesAndIndex(XlaOpKernelContext* ctx) {
     return {unpruned_branches_, ctx->Input(0)};
   }
 
-  int32_t branch_index = branch_index_literal.Get<int32>({});
+  int32_t branch_index = branch_index_literal.Get<int32_t>({});
   if (branch_index < 0 || branch_index >= unpruned_branches_.size()) {
     branch_index = unpruned_branches_.size() - 1;
   }
@@ -187,7 +187,8 @@ void XlaCaseOp::Compile(XlaOpKernelContext* ctx) {
 
       // Add any TensorArray gradients touched by the then/else computation to
       // the enclosing graph.
-      for (const string& grad_source : update.tensor_array_gradients_accessed) {
+      for (const std::string& grad_source :
+           update.tensor_array_gradients_accessed) {
         VLOG(5) << "TensorArray " << resource->name() << " accessed gradient "
                 << grad_source;
         XlaResource* gradient;
@@ -289,7 +290,7 @@ void XlaCaseOp::Compile(XlaOpKernelContext* ctx) {
       // Set token input for this "case" op.
       std::vector<xla::XlaOp> token_inputs;
       token_inputs.reserve(token_input_nodes_.size());
-      for (const string& node_name : token_input_nodes_) {
+      for (const std::string& node_name : token_input_nodes_) {
         auto token_or = compiler->GetNodeToken(node_name);
         OP_REQUIRES_OK(ctx, token_or.status());
         token_inputs.push_back(token_or.value());

tensorflow/compiler/tf2xla/kernels/case_op.h

@@ -65,8 +65,8 @@ class XlaCaseOp : public XlaOpKernel {
   DataTypeVector input_types_;
   DataTypeVector output_types_;
   bool has_token_input_output_;
-  std::vector<string> token_input_nodes_;
-  string original_node_name_;
+  std::vector<std::string> token_input_nodes_;
+  std::string original_node_name_;
   // Whether to propagate compile time consts into the cond branches.
   // This is not supported by default now since it may cause HBM memory
   // overheads.

tensorflow/compiler/tf2xla/kernels/categorical_op.cc

@@ -185,7 +185,7 @@ class StatelessCategoricalOp : public CategoricalOp {
 
  private:
   DataType dtype_;
-  string device_type_string_;
+  std::string device_type_string_;
 
   StatelessCategoricalOp(const StatelessCategoricalOp&) = delete;
   void operator=(const StatelessCategoricalOp&) = delete;

tensorflow/compiler/tf2xla/kernels/const_op.cc

@@ -38,7 +38,7 @@ template <typename DstT,
                                   std::is_same<DstT, bfloat16>::value>::type* =
               nullptr>
 DstT CastTo(int32_t src) {
-  return absl::bit_cast<DstT>(static_cast<uint16>(src));
+  return absl::bit_cast<DstT>(static_cast<uint16_t>(src));
 }
 
 // Returns scalar constant with the value in the tensor, if the given proto has

tensorflow/compiler/tf2xla/kernels/conv_op_helpers.cc

@@ -163,8 +163,8 @@ absl::Status CheckConvAttrs(const ConvOpAttrs& attrs) {
 absl::Status ConvBackpropComputeDimensionsV2XlaShapes(
     absl::string_view label, int num_spatial_dims,
     const xla::Shape& input_shape, const xla::Shape& filter_shape,
-    const xla::Shape& out_backprop_shape, absl::Span<const int32> dilations,
-    const std::vector<int32>& strides, Padding padding,
+    const xla::Shape& out_backprop_shape, absl::Span<const int32_t> dilations,
+    const std::vector<int32_t>& strides, Padding padding,
     TensorFormat data_format, ConvBackpropDimensions* dims,
     absl::Span<const int64_t> explicit_paddings) {
   TensorShape input_tensor_shape, filter_tensor_shape,
@@ -203,7 +203,7 @@ absl::StatusOr<ConvOpAttrs> ConvOpAttrs::Create(int num_spatial_dims,
         ctx->GetAttr("explicit_paddings", &attrs.explicit_paddings));
   }
 
-  string data_format;
+  std::string data_format;
   TF_RETURN_IF_ERROR(ctx->GetAttr("data_format", &data_format));
   if (!FormatFromString(data_format, &attrs.data_format)) {
     return errors::InvalidArgument("Invalid data format: ", data_format);
@@ -231,7 +231,7 @@ absl::StatusOr<ConvNDOpAttrs> ConvNDOpAttrs::Create(OpKernelConstruction* ctx) {
         ctx->GetAttr("explicit_paddings", &attrs.explicit_paddings));
   }
 
-  string data_format_str;
+  std::string data_format_str;
   TF_RETURN_IF_ERROR(ctx->GetAttr("data_format", &data_format_str));
   if (!(data_format_str == "CHANNELS_LAST" ||
         data_format_str == "CHANNELS_FIRST")) {