xref: /third_party/mindspore/mindspore/ccsrc/minddata/dataset/engine/cache/cache_server.h

/**
 * Copyright 2020-2021 Huawei Technologies Co., Ltd

 * 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.
*/

#ifndef MINDSPORE_CCSRC_MINDDATA_DATASET_ENGINE_CACHE_SERVER_H_
#define MINDSPORE_CCSRC_MINDDATA_DATASET_ENGINE_CACHE_SERVER_H_

#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <algorithm>
#include <atomic>
#include <chrono>
#include <iostream>
#include <memory>
#include <mutex>
#include <string>
#include <utility>
#include <vector>
#include <map>
#include <set>
#include <thread>
#include "minddata/dataset/engine/cache/cache_arena.h"
#include "minddata/dataset/engine/cache/cache_hw.h"
#include "minddata/dataset/engine/cache/cache_numa.h"
#include "minddata/dataset/engine/cache/cache_service.h"
#include "minddata/dataset/engine/cache/cache_grpc_server.h"
#include "minddata/dataset/engine/cache/cache_pool.h"
#include "minddata/dataset/core/tensor.h"
#include "minddata/dataset/util/allocator.h"
#include "minddata/dataset/util/arena.h"
#include "minddata/dataset/util/lock.h"
#include "minddata/dataset/util/random.h"
#include "minddata/dataset/util/semaphore.h"
#include "minddata/dataset/util/service.h"
#include "minddata/dataset/util/services.h"
#include "minddata/dataset/util/system_pool.h"
#include "minddata/dataset/util/queue.h"
#include "minddata/dataset/util/task_manager.h"

namespace mindspore {
namespace dataset {
/// \brief A server which provides CacheService services.
class CacheServer : public Service {
 public:
  friend class Services;
  using cache_index = std::map<connection_id_type, std::unique_ptr<CacheService>>;
  // Only allow new service to be created if left memory is more than 15% of our hard memory cap
  constexpr static float kMemoryBottomLineForNewService = 0.15;

  class Builder {
   public:
    Builder();

    ~Builder() = default;

    /// \brief Getter functions
    const std::string &GetTop() const { return top_; }
    int32_t GetNumWorkers() const { return num_workers_; }
    int32_t GetPort() const { return port_; }
    int32_t GetSharedMemorySzInGb() const { return shared_memory_sz_in_gb_; }
    float GetMemoryCapRatio() const { return memory_cap_ratio_; }
    int8_t GetLogLevel() const { return log_level_; }

    Builder &SetRootDirectory(std::string root) {
      top_ = std::move(root);
      return *this;
    }
    Builder &SetNumWorkers(int32_t n) {
      num_workers_ = n;
      return *this;
    }
    Builder &SetPort(int32_t p) {
      port_ = p;
      return *this;
    }
    Builder &SetSharedMemorySizeInGB(int32_t sz) {
      shared_memory_sz_in_gb_ = sz;
      return *this;
    }
    Builder &SetMemoryCapRatio(float ratio) {
      memory_cap_ratio_ = ratio;
      return *this;
    }
    Builder &SetLogLevel(int8_t log_level) {
      log_level_ = log_level;
      return *this;
    }

    Status SanityCheck();

    void Print(std::ostream &out) const {
      out << "Summary of the cache server configuration\n"
          << "Spill directory: " << (GetTop().empty() ? "None" : GetTop()) << "\n"
          << "Number of parallel workers: " << GetNumWorkers() << "\n"
          << "Tcp/ip port: " << GetPort() << "\n"
          << "Shared memory size (in GB): " << GetSharedMemorySzInGb() << "\n"
          << "Memory cap ratio: " << GetMemoryCapRatio() << "\n"
          << "Log level: " << std::to_string(GetLogLevel());
    }

    friend std::ostream &operator<<(std::ostream &out, const Builder &bld) {
      bld.Print(out);
      return out;
    }

    Status Build() {
      // Get information of numa architecture and adjust num_workers_ based on numa count
      hw_info_ = std::make_shared<CacheServerHW>();
      RETURN_IF_NOT_OK(hw_info_->GetNumaNodeInfo());
      std::string warning_string;
      if (num_workers_ == -1) {
        // if the user did not provide a value for num_workers, set it to half of num_cpu as default and adjust it if
        // the default is not the optimal.
        int32_t dft_num_workers = std::thread::hardware_concurrency() > 2 ? std::thread::hardware_concurrency() / 2 : 1;
        num_workers_ = AdjustNumWorkers(dft_num_workers);
      } else {
        // if the users have given their own value, adjust it and provide a warning if it got changed.
        int32_t num_workers_new = AdjustNumWorkers(num_workers_);
        if (num_workers_ != num_workers_new) {
          warning_string =
            "The configuration of workers on the cache server is dependent on the NUMA architecture of the server. "
            "The current setting is not the optimal for the NUMA architecture. Re-adjusting the number of workers "
            "to optimal setting of " +
            std::to_string(num_workers_new) + ".\n";
          MS_LOG(INFO) << warning_string;
        }
        num_workers_ = num_workers_new;
      }
      RETURN_IF_NOT_OK(SanityCheck());
      // We need to bring up the Task Manager by bringing up the Services singleton.
      RETURN_IF_NOT_OK(Services::CreateInstance());
      RETURN_IF_NOT_OK(CacheServer::CreateInstance(top_, num_workers_, port_, shared_memory_sz_in_gb_,
                                                   memory_cap_ratio_, log_level_, std::move(hw_info_)));
      return Status(StatusCode::kSuccess, warning_string);
    }

   private:
    std::string top_;
    int32_t num_workers_;
    int32_t port_;
    int32_t shared_memory_sz_in_gb_;
    float memory_cap_ratio_;
    int8_t log_level_;
    std::shared_ptr<CacheServerHW> hw_info_;

    /// \brief Sanity checks on the shared memory.
    /// \return Status object
    Status IpcResourceCleanup();

    /// \brief Adjust the value of num_workers if it's not the optimal to NUMA architecture.
    int32_t AdjustNumWorkers(int32_t num_workers);
  };

  CacheServer(const CacheServer &) = delete;
  CacheServer &operator=(const CacheServer &) = delete;
  CacheServer(CacheServer &&) = delete;
  CacheServer &operator=(CacheServer &) = delete;
  Status DoServiceStart() override;
  Status DoServiceStop() override;
  ~CacheServer() override { (void)ServiceStop(); }

  static Status CreateInstance(const std::string &spill_path, int32_t num_workers, int32_t port,
                               int32_t shared_memory_sz, float memory_cap_ratio, int8_t log_level,
                               std::shared_ptr<CacheServerHW> hw_info) {
    std::call_once(init_instance_flag_, [&]() -> Status {
      auto &SvcManager = Services::GetInstance();
      RETURN_IF_NOT_OK(SvcManager.AddHook(&instance_, spill_path, num_workers, port, shared_memory_sz, memory_cap_ratio,
                                          log_level, hw_info));
      return Status::OK();
    });
    return Status::OK();
  }

  static CacheServer &GetInstance() { return *instance_; }

  /// \brief For the current demonstration, a cache client contacts cache server using a Queue.
  /// \param rq
  /// \return Status object
  Status PushRequest(int32_t queue_id, CacheServerRequest *rq) {
    RETURN_UNEXPECTED_IF_NULL(rq);
    RETURN_IF_NOT_OK(cache_q_->operator[](queue_id)->Add(rq));
    return Status::OK();
  }

  /// \\brief Kick off server threads. Never return unless error out.
  Status Run(SharedMessage::queue_id_t msg_qid);

  /// \brief Get a free tag
  /// \param q[in] pointer to a pointer to a CacheServerRequest
  /// \return Status object
  static Status GetFreeRequestTag(CacheServerRequest **q);

  /// \brief Return a tag to the free list
  /// \param p[in] pointer to already finished CacheServerRequest tag
  /// \return Status object
  static Status ReturnRequestTag(CacheServerRequest *p);

  /// Return an instance of the numa control
  std::shared_ptr<CacheServerHW> GetHWControl() { return hw_info_; }

  /// \brief Set CPU affinity
  Status SetAffinity(const Task &tk, numa_id_t numa_node) { return hw_info_->SetAffinity(tk, numa_node); }

  /// \brief return number of workers
  auto GetNumWorkers() const { return num_workers_; }

  /// \brief return number of grpc workers
  auto GetNumGrpcWorkers() const { return num_grpc_workers_; }

  /// \brief return number of numa nodes
  auto GetNumaNodeCount() const { return hw_info_->GetNumaNodeCount(); }

  /// \brief Assign a worker by a numa id
  /// \return worker id
  worker_id_t GetWorkerByNumaId(numa_id_t node_id) const;

  /// \brief Randomly pick a worker
  /// \return worker id
  worker_id_t GetRandomWorker() const;

  /// \brief Check if we bind threads to numa cores
  bool IsNumaAffinityOn() const { return numa_affinity_; }

  /// \brief Return the memory cap ratio
  float GetMemoryCapRatio() const { return memory_cap_ratio_; }

  /// \brief Function to handle a row request
  /// \param[in] cache_req A row request to handle
  /// \param[out] internal_request Indicator if the request is an internal request
  /// \return Status object
  Status ProcessRowRequest(CacheServerRequest *cache_req, bool *internal_request);

  /// \brief Function to handle an admin request
  /// \param[in] cache_req An admin request to handle
  /// \return Status object
  Status ProcessAdminRequest(CacheServerRequest *cache_req);

  /// \brief Function to handle a session request
  /// \param[in] cache_req A session request to handle
  /// \return Status object
  Status ProcessSessionRequest(CacheServerRequest *cache_req);

  /// \brief How a request is handled.
  /// \note that it can be process immediately by a grpc thread or routed to a server thread
  /// which is pinned to some numa node core.
  Status ProcessRequest(CacheServerRequest *cache_req);

  void GlobalShutdown();

  /// \brief This returns where we attach to the shared memory.
  /// Some gRPC requests will ask for a shared memory block, and
  /// we can't return the absolute address as this makes no sense
  /// in the client. So instead we will return an address relative
  /// to the base address of the shared memory where we attach to.
  /// \return Base address of the shared memory.
  const void *SharedMemoryBaseAddr() const { return shm_->SharedMemoryBaseAddr(); }

  /// \brief Return the public key of the shared memory.
  int32_t GetKey() const { return shm_->GetKey(); }

  Status AllocateSharedMemory(int32_t client_id, size_t sz, void **p);

  void DeallocateSharedMemory(int32_t client_id, void *p);

 private:
  static std::once_flag init_instance_flag_;
  static CacheServer *instance_;
  mutable RWLock rwLock_;
  mutable RWLock sessions_lock_;
  std::string top_;
  cache_index all_caches_;
  std::set<session_id_type> active_sessions_;
  std::shared_ptr<QueueList<CacheServerRequest *>> cache_q_;
  std::shared_ptr<CacheServerGreeterImpl> comm_layer_;
  TaskGroup vg_;
  int32_t num_workers_;
  int32_t num_grpc_workers_;
  int32_t port_;
  int32_t shared_memory_sz_in_gb_;
  int8_t log_level_;  // log_level is saved here for informational purpose only. It's not a functional field.
  std::atomic<bool> global_shutdown_;
  float memory_cap_ratio_;
  std::shared_ptr<CacheServerHW> hw_info_;
  std::map<worker_id_t, Task *> numa_tasks_;
  bool numa_affinity_;
  std::vector<int32_t> shutdown_qIDs_;
  std::unique_ptr<CachedSharedMemory> shm_;

  /// \brief Constructor
  /// \param spill_path Top directory for spilling buffers to.
  /// \param num_workers Number of threads for handling requests.
  explicit CacheServer(const std::string &spill_path, int32_t num_workers, int32_t port, int32_t share_memory_sz_in_gb,
                       float memory_cap_ratio, int8_t log_level, std::shared_ptr<CacheServerHW> hw_info);

  /// \brief Locate a cache service from connection id.
  /// \return Pointer to cache service. Null if not found
  CacheService *GetService(connection_id_type id) const;

  /// \brief Going over existing cache service and calculate how much we have consumed so far, a new cache service
  /// can only be created if there is still enough avail memory left
  /// \param cache_mem_sz Requested memory for a new cache service
  /// \return Status object
  Status GlobalMemoryCheck(uint64_t cache_mem_sz);

  /// \brief Create a cache service. We allow multiple clients to create the same cache service.
  /// Subsequent duplicate requests are ignored. The first cache client to create the service will be given
  /// a special unique cookie.
  /// \return Status object
  Status CreateService(CacheRequest *rq, CacheReply *reply);

  /// \brief Destroy a cache service
  /// \param rq
  /// \return Status object
  Status DestroyCache(CacheRequest *rq);

  /// \brief Entry point for all internal server threads.
  Status ServerRequest(worker_id_t worker_id);

  /// \brief Entry point for all grpc threads.
  /// \return
  Status RpcRequest(worker_id_t worker_id);

  Status DestroySession(CacheRequest *rq);

  /// \brief Create a connection id from a session id and a crc
  /// \param session_id
  /// \param crc
  /// \return connection id
  connection_id_type GetConnectionID(session_id_type session_id, uint32_t crc) const;

  /// \brief Extract the session id from a connection id
  /// \param connection_id
  /// \return session id
  session_id_type GetSessionID(connection_id_type connection_id) const;

  /// \brief Generate a session ID for the client
  /// \return Session ID
  session_id_type GenerateSessionID();

  /// \brief Handle kAllocateSharedBlock request
  /// \param rq CacheRequest
  /// \param reply CacheReply
  /// \return Status object
  Status AllocateSharedMemory(CacheRequest *rq, CacheReply *reply);

  /// \brief Handle kFreeSharedBlock request
  /// \param rq
  /// \return Status object
  Status FreeSharedMemory(CacheRequest *rq);

  /// \brief Handle CacheRow request
  /// \note There are two different implementation depends if shared memory is used for transportation.
  /// \return Status object
  Status FastCacheRow(CacheRequest *rq, CacheReply *reply);
  Status CacheRow(CacheRequest *rq, CacheReply *reply);

  /// \brief Internal function to get statistics
  /// \param rq
  /// \param reply
  /// \return Status object
  Status GetStat(CacheRequest *rq, CacheReply *reply);

  /// \brief Internal function to get cache state
  /// \param rq
  /// \param reply
  /// \return Status object
  Status GetCacheState(CacheRequest *rq, CacheReply *reply);

  /// \brief Cache a schema request
  /// \param rq
  /// \return Status object
  Status CacheSchema(CacheRequest *rq);

  /// \brief Fetch a schema request
  /// \param rq
  /// \param reply
  /// \return Status object
  Status FetchSchema(CacheRequest *rq, CacheReply *reply);

  /// \brief Mark Build phase done (for non-mappable case)
  /// \param rq
  /// \return Status object
  Status BuildPhaseDone(CacheRequest *rq);

  /// \brief A proper shutdown of the server
  /// \return Status object
  Status AcknowledgeShutdown(CacheServerRequest *cache_req);

  /// \brief Find keys that will be cache miss
  /// \return Status object
  Status GetCacheMissKeys(CacheRequest *rq, CacheReply *reply);

  /// \brief Toggle write mode for a service
  Status ToggleWriteMode(CacheRequest *rq);

  /// \brief List the sessions and their caches
  /// \param reply
  /// \return Status object
  Status ListSessions(CacheReply *reply);

  /// \brief Connect request by a pipeline
  Status ConnectReset(CacheRequest *rq);

  /// \brief This is an internal structure used by Batch processing.
  /// This is how it works internally. For batch fetch/cache, the grpc thread
  /// will intercept the request and breaks it down into multiple internal requests
  /// and spread over all the server workers. But each internal request consumes
  /// one free tag and we may run out of free tags if they don't return promptly.
  /// So we will let the server thread return the free tag immediately but the put
  /// the return code in this following structure. GRPC thread must wait until all
  /// the rc come back.
  class BatchWait : public std::enable_shared_from_this<BatchWait> {
   public:
    explicit BatchWait(int n) : expected_(n), num_back_(0) {
      expected_ = n;
      rc_lists_.reserve(expected_);
    }
    ~BatchWait() = default;

    std::weak_ptr<BatchWait> GetBatchWait() { return weak_from_this(); }

    Status Set(Status rc) {
      CHECK_FAIL_RETURN_UNEXPECTED(expected_ > num_back_, "Programming error");
      std::unique_lock<std::mutex> lck(mux_);
      rc_lists_.push_back(std::move(rc));
      ++num_back_;
      if (num_back_ == expected_) {
        wp_.Set();
      }
      return Status::OK();
    }

    Status Wait() { return wp_.Wait(); }

    Status GetRc() {
      Status rc;
      for (auto &cache_rc : rc_lists_) {
        if (cache_rc.IsError() && cache_rc != StatusCode::kMDInterrupted && rc.IsOk()) {
          rc = cache_rc;
        }
      }
      return rc;
    }

   private:
    std::mutex mux_;
    WaitPost wp_;
    int64_t expected_;
    int64_t num_back_;
    std::vector<Status> rc_lists_;
  };

  /// \brief Internal function to do row batch fetch
  /// \param rq Request
  /// \param reply Reply
  /// \return Status object
  Status BatchFetchRows(CacheRequest *rq, CacheReply *reply);

  /// \brief Main function to fetch rows in batch. The output is a contiguous memory which will be decoded
  /// by the CacheClient. Cache miss is not an error, and will be coded in the output to mark an empty row.
  /// \param[in] v A vector of row id.
  /// \param[out] out A contiguous memory buffer that holds the requested rows.
  /// \return Status object
  Status BatchFetch(const std::shared_ptr<flatbuffers::FlatBufferBuilder> &fbb, WritableSlice *out);
  Status BatchCacheRows(CacheRequest *rq);

  Status InternalFetchRow(CacheRequest *rq);
  Status InternalCacheRow(CacheRequest *rq, CacheReply *reply);
};
}  // namespace dataset
}  // namespace mindspore
#endif  // MINDSPORE_CCSRC_MINDDATA_DATASET_CORE_CACHE_TENSOR_H_