vllm_shim.hpp

Go to the documentation of this file.

#pragma once
#include <continuum/backend/backend.hpp>
 
#include <atomic>
#include <cstddef>
#include <cstdint>
#include <map>
#include <mutex>
#include <optional>
#include <string>
#include <vector>
 
namespace continuum::backend {
 
class VllmShimBackend : public Backend {
 public:
  struct PrefixState {
    std::string prefix_text;
    std::string model;
  };
 
  std::vector<std::uint8_t> export_state(const BackendState& state) const override;
  std::optional<BackendState> import_state(const std::vector<std::uint8_t>& bytes) override;
  std::size_t state_count() const;
  std::size_t rewarm_count() const { return rewarm_count_.load(); }
 
  static std::string ExtractJsonString(const std::string& body, const std::string& key);
  static std::int32_t ExtractJsonInt(const std::string& body, const std::string& key);
 
  BackendCapabilities capabilities() const override;
  BackendRunResult run_with_cache(
      const ir::Node& node,
      const std::vector<continuum::Value>& inputs,
      const std::optional<BackendState>& prefix_state,
      std::int32_t remaining_tokens) override;
 
 private:
  static constexpr std::size_t kMaxStates = 4096;
  BackendState remember(std::string prefix_text, std::string model);
  std::optional<PrefixState> lookup(const BackendState& state) const;
 
  mutable std::mutex mu_;
  std::map<std::uint64_t, PrefixState> states_;  // ordered: begin() is the oldest
  std::uint64_t next_state_id_ = 1;
  std::atomic<std::size_t> rewarm_count_{0};
};
 
}  // namespace continuum::backend