/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "tensorflow/core/runtime_fallback/runtime/fallback_batch_kernel.h"

#include <cstdint>
#include <cstdlib>
#include <optional>
#include <string>

#include "absl/log/log.h"
#include "absl/status/status.h"
#include "absl/strings/numbers.h"
#include "absl/strings/string_view.h"
#include "tensorflow/core/framework/op_kernel.h"
#include "tensorflow/core/kernels/batching_util/bounded_executor.h"
#include "tensorflow/core/lib/monitoring/gauge.h"
#include "tensorflow/core/platform/types.h"
#include "tensorflow/core/runtime_fallback/kernel/kernel_fallback_compat_request_state.h"
#include "tensorflow/core/tfrt/fallback/op_kernel_runner.h"
#include "tensorflow/core/tfrt/utils/error_util.h"
#include "tensorflow/core/tfrt/utils/fallback_tensor.h"
#include "tfrt/host_context/async_value_ref.h"  // from @tf_runtime
#include "tfrt/host_context/execution_context.h"  // from @tf_runtime
#include "tfrt/host_context/host_context.h"  // from @tf_runtime

namespace tensorflow {
namespace tfrt_stub {
namespace {

// Op attributes.
constexpr char kEnableAdaptiveSchedulerAttr[] = "_enable_adaptive_scheduler";
constexpr char kMinInflightBatchesAttr[] = "_min_inflight_batches";
constexpr char kInitialInflightBatchesAttr[] = "_initial_inflight_batches";
constexpr char kMaxInflightBatchesAttr[] = "_max_inflight_batches";
constexpr char kBatchesToAverageOverAttr[] = "_batches_to_average_over";
// Default thread count in the per-process batching thread pool.
// The value is the same as the TF batch kernel BatchKernel.

}  // namespace

int32_t
BatchFunctionFallbackKernelBase::NumBatchThreadsFromEnvironmentWithDefault(
    int default_num_batch_threads) {
  int32_t num;
  const char* val = std::getenv("TF_NUM_BATCH_THREADS");

  return (val && absl::SimpleAtoi(val, &num)) ? num : default_num_batch_threads;
}

thread::ThreadPool*
BatchFunctionFallbackKernelBase::GetOrCreateBatchThreadsPool() {
  static thread::ThreadPool* shared_thread_pool = [&]() -> thread::ThreadPool* {
    serving::BoundedExecutor::Options options;

    options.num_threads =
        NumBatchThreadsFromEnvironmentWithDefault(kBatchThreadPoolSize);

    options.thread_name = "adaptive_batch_threads";

    auto status_or_executor = serving::BoundedExecutor::Create(options);
    if (!status_or_executor.ok()) {
      LOG(WARNING) << "Failed to create a batch threads pool with error "
                   << status_or_executor.status();
      return nullptr;
    }
    static serving::BoundedExecutor* executor =
        status_or_executor.value().release();
    return new thread::ThreadPool(executor);
  }();
  return shared_thread_pool;
}

BatchFunctionFallbackKernelBase::BatchFunctionFallbackKernelBase(
    OpKernelConstruction* c)
    : AsyncOpKernel(c) {
  OP_REQUIRES_OK(c, c->GetAttr("container", &container_));
  OP_REQUIRES_OK(c, c->GetAttr("shared_name", &shared_name_));
  OP_REQUIRES_OK(c, c->GetAttr("batching_queue", &batcher_queue_));
  OP_REQUIRES_OK(c, c->GetAttr("num_batch_threads", &num_batch_threads_));
  OP_REQUIRES_OK(c, c->GetAttr("max_batch_size", &max_batch_size_));
  OP_REQUIRES_OK(c, c->GetAttr("batch_timeout_micros", &batch_timeout_micros_));
  OP_REQUIRES_OK(c, c->GetAttr("max_enqueued_batches", &max_enqueued_batches_));
  OP_REQUIRES_OK(c, c->GetAttr("allowed_batch_sizes", &allowed_batch_sizes_));

  OP_REQUIRES_OK(c, c->GetAttr("low_priority_max_batch_size",
                               &low_priority_max_batch_size_));
  OP_REQUIRES_OK(c, c->GetAttr("low_priority_batch_timeout_micros",
                               &low_priority_batch_timeout_micros_));
  OP_REQUIRES_OK(c, c->GetAttr("low_priority_allowed_batch_sizes",
                               &low_priority_allowed_batch_sizes_));
  OP_REQUIRES_OK(c, c->GetAttr("low_priority_max_enqueued_batches",
                               &low_priority_max_enqueued_batches_));
  OP_REQUIRES_OK(c,
                 c->GetAttr("mixed_priority_policy", &mixed_priority_policy_));
  OP_REQUIRES_OK(c, c->GetAttr("batch_padding_policy", &batch_padding_policy_));

  if (shared_name_.empty()) {
    // If shared_name is not supplied, use name instead (prevent collisions by
    // default).
    shared_name_ = name();
  }

  VLOG(1) << "BatchFunctionFallbackKernel(" << this
          << ") container attribute: \"" << container_
          << "\", shared_name attribute: \"" << shared_name_
          << "\", batching_queue attribute: \"" << batcher_queue_ << "\"";

  if (c->HasAttr("enable_large_batch_splitting")) {
    OP_REQUIRES_OK(c, c->GetAttr("enable_large_batch_splitting",
                                 &enable_large_batch_splitting_));
    has_attribute_enable_large_batch_splitting_ = true;
  } else {
    enable_large_batch_splitting_ = false;
    has_attribute_enable_large_batch_splitting_ = false;
  }

  if (c->HasAttr("disable_padding")) {
    OP_REQUIRES_OK(c, c->GetAttr("disable_padding", &disable_padding_));
  } else {
    disable_padding_ = false;
  }

  if (c->HasAttr("enable_priority_aware_batch_scheduler")) {
    OP_REQUIRES_OK(c, c->GetAttr("enable_priority_aware_batch_scheduler",
                                 &enable_priority_aware_batch_scheduler_));
  }

  // Helper function `SetAdaptiveBatchSchedulerOptions` calls
  // `OP_REQUIRES_OK`, which exits the current function upon error.
  // So validate status of `op-kernel-construction`.
  SetAdaptiveBatchSchedulerOptions(c, num_batch_threads_);
  if (!c->status().ok()) {
    return;
  }

  if (enable_adaptive_batch_threads_) {
    // One scheduler instance contains a couple of queue instances,
    // `batcher_queue_` is the key to find queue for this batch-op in the
    // graph.
    // Use `shared_name_` and name() as prefix for `batcher_queue_`.
    // Note name() is unique per session (from session metadata).
    batcher_queue_ = name() + "/" + shared_name_ + batcher_queue_;
  }

  OP_REQUIRES_OK(c, ValidateAllowedBatchSizes());
}

absl::Status BatchFunctionFallbackKernelBase::ValidateAllowedBatchSizes()
    const {
  if (allowed_batch_sizes_.empty()) {
    return absl::OkStatus();
  }
  int32_t last_size = 0;
  for (size_t i = 0; i < allowed_batch_sizes_.size(); ++i) {
    const int32_t size = allowed_batch_sizes_.at(i);
    if (i > 0 && size <= last_size) {
      return errors::InvalidArgument(
          "allowed_batch_sizes entries must be monotonically increasing");
    }

    if ((!enable_large_batch_splitting_) &&
        (i == allowed_batch_sizes_.size() - 1) && (size != max_batch_size_)) {
      return errors::InvalidArgument(
          "final entry in allowed_batch_sizes must equal max_batch_size when "
          "enable_large_batch_splitting is False");
    }

    last_size = size;
  }
  return absl::OkStatus();
}

void BatchFunctionFallbackKernelBase::SetAdaptiveBatchSchedulerOptions(
    OpKernelConstruction* c, int32_t num_batch_threads) {
  if (c->HasAttr(kEnableAdaptiveSchedulerAttr)) {
    OP_REQUIRES_OK(c, c->GetAttr(kEnableAdaptiveSchedulerAttr,
                                 &enable_adaptive_batch_threads_));
  }

  if (num_batch_threads <= 0) {
    enable_adaptive_batch_threads_ = true;
  }

  if (!enable_adaptive_batch_threads_) {
    // adaptive_batch_scheduler_options_ is nullopt.
    return;
  }

  // adaptive_batch_scheduler_options_ is not nullopt
  AdaptiveBatchSchedulerOptions options;

  if (c->HasAttr(kBatchesToAverageOverAttr)) {
    OP_REQUIRES_OK(c, c->GetAttr(kBatchesToAverageOverAttr,
                                 &options.batches_to_average_over));
  }

  if (c->HasAttr(kMinInflightBatchesAttr)) {
    OP_REQUIRES_OK(c, c->GetAttr(kMinInflightBatchesAttr,
                                 &options.min_in_flight_batches_limit));
  }

  if (c->HasAttr(kInitialInflightBatchesAttr)) {
    OP_REQUIRES_OK(c, c->GetAttr(kInitialInflightBatchesAttr,
                                 &options.initial_in_flight_batches_limit));
  }

  if (c->HasAttr(kMaxInflightBatchesAttr)) {
    OP_REQUIRES_OK(c, c->GetAttr(kMaxInflightBatchesAttr,
                                 &options.max_in_flight_batches_limit));
  }

  // At this point, the batch kernel is configured to use adaptive scheduling.
  // To validate or return error at kernel construction time, invokes
  // `GetOrCreateBatchThreadsPool` and validates returned `thread_pool` is
  // valid.
  // Note`GetOrCreateBatchThreadsPool` creates the thread pool once and
  // re-uses the thread-pool instance afterwards.
  thread::ThreadPool* thread_pool = GetOrCreateBatchThreadsPool();
  OP_REQUIRES(
      c, thread_pool != nullptr,
      errors::FailedPrecondition("Failed to create batch threads pool"));

  adaptive_batch_scheduler_options_ = options;
}

}  // namespace tfrt_stub
}  // namespace tensorflow
