server: cancel prompt processing & non-streamed requests when connection closed

common/arg.cpp

@@ -1879,6 +1879,13 @@ gpt_params_context gpt_params_parser_init(gpt_params & params, llama_example ex,
             params.slot_prompt_similarity = std::stof(value);
         }
     ).set_examples({LLAMA_EXAMPLE_SERVER}));
+    add_opt(llama_arg(
+        {"--testing-sampler-delay-millis"}, "N",
+        format("for tests: delay in milliseconds to add to each sampling (default: %d)", params.testing_sampler_delay_millis),
+        [](gpt_params & params, int value) {
+            params.testing_sampler_delay_millis = value;
+        }
+    ).set_examples({LLAMA_EXAMPLE_SERVER}));
     add_opt(llama_arg(
         {"--lora-init-without-apply"},
         format("load LoRA adapters without applying them (apply later via POST /lora-adapters) (default: %s)", params.lora_init_without_apply ? "enabled" : "disabled"),

common/common.h

@@ -299,6 +299,8 @@ struct gpt_params {
 
     float slot_prompt_similarity = 0.5f;
 
+    int testing_sampler_delay_millis = 0;
+
     // batched-bench params
     bool is_pp_shared = false;
 

examples/server/httplib.h

@@ -457,6 +457,7 @@ class DataSink {
 
   std::function<bool(const char *data, size_t data_len)> write;
   std::function<bool()> is_writable;
+  std::function<bool()> is_alive;
   std::function<void()> done;
   std::function<void(const Headers &trailer)> done_with_trailer;
   std::ostream os;
@@ -639,6 +640,7 @@ class Stream {
 
   virtual bool is_readable() const = 0;
   virtual bool is_writable() const = 0;
+  virtual bool is_alive() const = 0;
 
   virtual ssize_t read(char *ptr, size_t size) = 0;
   virtual ssize_t write(const char *ptr, size_t size) = 0;
@@ -2135,6 +2137,7 @@ class BufferStream final : public Stream {
 
   bool is_readable() const override;
   bool is_writable() const override;
+  bool is_alive() const override;
   ssize_t read(char *ptr, size_t size) override;
   ssize_t write(const char *ptr, size_t size) override;
   void get_remote_ip_and_port(std::string &ip, int &port) const override;
@@ -2945,6 +2948,7 @@ class SocketStream final : public Stream {
 
   bool is_readable() const override;
   bool is_writable() const override;
+  bool is_alive() const override;
   ssize_t read(char *ptr, size_t size) override;
   ssize_t write(const char *ptr, size_t size) override;
   void get_remote_ip_and_port(std::string &ip, int &port) const override;
@@ -2975,6 +2979,7 @@ class SSLSocketStream final : public Stream {
 
   bool is_readable() const override;
   bool is_writable() const override;
+  bool is_alive() const override;
   ssize_t read(char *ptr, size_t size) override;
   ssize_t write(const char *ptr, size_t size) override;
   void get_remote_ip_and_port(std::string &ip, int &port) const override;
@@ -4088,6 +4093,7 @@ inline bool write_content(Stream &strm, const ContentProvider &content_provider,
   };
 
   data_sink.is_writable = [&]() -> bool { return strm.is_writable(); };
+  data_sink.is_alive = [&]() -> bool { return strm.is_alive(); };
 
   while (offset < end_offset && !is_shutting_down()) {
     if (!strm.is_writable()) {
@@ -4134,6 +4140,7 @@ write_content_without_length(Stream &strm,
   };
 
   data_sink.is_writable = [&]() -> bool { return strm.is_writable(); };
+  data_sink.is_alive = [&]() -> bool { return strm.is_alive(); };
 
   data_sink.done = [&](void) { data_available = false; };
 
@@ -4186,6 +4193,7 @@ write_content_chunked(Stream &strm, const ContentProvider &content_provider,
   };
 
   data_sink.is_writable = [&]() -> bool { return strm.is_writable(); };
+  data_sink.is_alive = [&]() -> bool { return strm.is_alive(); };
 
   auto done_with_trailer = [&](const Headers *trailer) {
     if (!ok) { return; }
@@ -5484,6 +5492,10 @@ inline bool SocketStream::is_writable() const {
          is_socket_alive(sock_);
 }
 
+inline bool SocketStream::is_alive() const {
+  return is_socket_alive(sock_);
+}
+
 inline ssize_t SocketStream::read(char *ptr, size_t size) {
 #ifdef _WIN32
   size =
@@ -5558,6 +5570,8 @@ inline bool BufferStream::is_readable() const { return true; }
 
 inline bool BufferStream::is_writable() const { return true; }
 
+inline bool BufferStream::is_alive() const { return true; }
+
 inline ssize_t BufferStream::read(char *ptr, size_t size) {
 #if defined(_MSC_VER) && _MSC_VER < 1910
   auto len_read = buffer._Copy_s(ptr, size, size, position);
@@ -8348,6 +8362,10 @@ inline bool SSLSocketStream::is_writable() const {
          is_socket_alive(sock_);
 }
 
+inline bool SSLSocketStream::is_alive() const {
+  return is_socket_alive(sock_);
+}
+
 inline ssize_t SSLSocketStream::read(char *ptr, size_t size) {
   if (SSL_pending(ssl_) > 0) {
     return SSL_read(ssl_, ptr, static_cast<int>(size));