|
| 1 | +#include "tcp_server.h" |
| 2 | +#include "llama.h" |
| 3 | +#include "utils.h" |
| 4 | + |
| 5 | +#include <iostream> |
| 6 | + |
| 7 | +#include <stdarg.h> |
| 8 | +#include <stdio.h> |
| 9 | +#include <stdlib.h> |
| 10 | +#include <stdbool.h> |
| 11 | +#include <string.h> |
| 12 | +#include <errno.h> |
| 13 | + |
| 14 | +#include <signal.h> |
| 15 | +#include <unistd.h> |
| 16 | +#include <sys/wait.h> |
| 17 | + |
| 18 | +#include <sys/types.h> |
| 19 | +#include <sys/socket.h> |
| 20 | +#include <arpa/inet.h> |
| 21 | +#include <netdb.h> |
| 22 | + |
| 23 | +class PosixStream : public std::istream { |
| 24 | + public: |
| 25 | + PosixStream(int fd) : std::istream(&buf), buf(fd) {} |
| 26 | + ~PosixStream() { close(buf.get_fd()); } |
| 27 | + |
| 28 | + private: |
| 29 | + class PosixStreamBuf : public std::streambuf { |
| 30 | + public: |
| 31 | + PosixStreamBuf(int fd) : fd(fd) {} |
| 32 | + int get_fd() const { return fd; } |
| 33 | + |
| 34 | + protected: |
| 35 | + virtual int_type underflow() { |
| 36 | + if (gptr() < egptr()) { |
| 37 | + return traits_type::to_int_type(*gptr()); |
| 38 | + } |
| 39 | + |
| 40 | + ssize_t num_read = ::read(fd, buffer, BUFFER_SIZE); |
| 41 | + if (num_read <= 0) { |
| 42 | + return traits_type::eof(); |
| 43 | + } |
| 44 | + |
| 45 | + setg(buffer, buffer, buffer + num_read); |
| 46 | + return traits_type::to_int_type(*gptr()); |
| 47 | + } |
| 48 | + |
| 49 | + private: |
| 50 | + static const int BUFFER_SIZE = 1024; |
| 51 | + int fd; |
| 52 | + char buffer[BUFFER_SIZE]; |
| 53 | + }; |
| 54 | + |
| 55 | + PosixStreamBuf buf; |
| 56 | +}; |
| 57 | + |
| 58 | +void die(const char *msg, ...) |
| 59 | +{ |
| 60 | + va_list ap; |
| 61 | + |
| 62 | + va_start(ap, msg); |
| 63 | + vfprintf(stderr, msg, ap); |
| 64 | + va_end(ap); |
| 65 | + fputc('\n', stderr); |
| 66 | + exit(1); |
| 67 | +} |
| 68 | + |
| 69 | +static char *read_argument(uint8_t **param_buf, size_t *param_buf_size, FILE *instream) { |
| 70 | + bool done = false; |
| 71 | + uint8_t *buf = *param_buf; |
| 72 | + size_t bufsize = *param_buf_size; |
| 73 | + size_t bufpos = 0; |
| 74 | + while (!done) { |
| 75 | + if (bufpos == bufsize) { |
| 76 | + bufsize += 1024; |
| 77 | + buf = (uint8_t *)realloc(buf, bufsize); |
| 78 | + if (!buf) { |
| 79 | + die("failed to allocate memory"); |
| 80 | + } |
| 81 | + } |
| 82 | + |
| 83 | + int c = fgetc(instream); |
| 84 | + if (c == EOF) { |
| 85 | + die("unexpected EOF client socket"); |
| 86 | + } |
| 87 | + buf[bufpos++] = (uint8_t)c; |
| 88 | + if (c == 0) { |
| 89 | + // done reading argument |
| 90 | + break; |
| 91 | + } |
| 92 | + } |
| 93 | + *param_buf = buf; |
| 94 | + *param_buf_size = bufsize; |
| 95 | + return strdup((char *)buf); |
| 96 | +} |
| 97 | + |
| 98 | +static int read_arguments(int argc, char **argv, FILE *instream) { |
| 99 | + int i = 1; |
| 100 | + size_t param_buf_size = 0; |
| 101 | + uint8_t *param_buf = nullptr; |
| 102 | + |
| 103 | + for (i = 1; i < argc; i++) { |
| 104 | + argv[i] = read_argument(¶m_buf, ¶m_buf_size, instream); |
| 105 | + } |
| 106 | + |
| 107 | + free(param_buf); |
| 108 | + return i; |
| 109 | +} |
| 110 | + |
| 111 | +static int serve_model(llama_context * ctx, |
| 112 | + gpt_params params, |
| 113 | + int sock_fd) |
| 114 | +{ |
| 115 | + int argc; |
| 116 | + char **argv; |
| 117 | + FILE *instream = fdopen(sock_fd, "r"); |
| 118 | + FILE *outstream = fdopen(sock_fd, "w"); |
| 119 | + setvbuf(instream, NULL, _IONBF, 0); |
| 120 | + |
| 121 | + // start by reading the parameter count |
| 122 | + if (fscanf(instream, "%d\n", &argc) != 1) { |
| 123 | + fprintf(outstream, "Error: First line must be character count\n"); |
| 124 | + fflush(outstream); |
| 125 | + return 1; |
| 126 | + } |
| 127 | + |
| 128 | + argc += 1; // add one extra argument to emulate the program command line |
| 129 | + argv = (char **)malloc(argc * sizeof *argv); |
| 130 | + argv[0] = nullptr; |
| 131 | + if (read_arguments(argc, argv, instream) != argc) { |
| 132 | + fprintf(outstream, "Error: Failed to read arguments\n"); |
| 133 | + fflush(outstream); |
| 134 | + } |
| 135 | + |
| 136 | + if (gpt_params_parse(argc, argv, params) == false) { |
| 137 | + fprintf(outstream, "Error: Failed to parse parameters\n"); |
| 138 | + fflush(outstream); |
| 139 | + return 1; |
| 140 | + } |
| 141 | + |
| 142 | + for (int i = 1; i < argc; i++) { |
| 143 | + free(argv[i]); |
| 144 | + } |
| 145 | + free(argv); |
| 146 | + |
| 147 | + PosixStream tcp_instream(sock_fd); |
| 148 | + |
| 149 | + return run(ctx, params, tcp_instream, outstream, outstream); |
| 150 | +} |
| 151 | + |
| 152 | +int listen_tcp(llama_context * ctx, gpt_params params) { |
| 153 | + int listen_fd; |
| 154 | + int status; |
| 155 | + pid_t child; |
| 156 | + struct addrinfo hints; |
| 157 | + struct addrinfo *servinfo, *p; |
| 158 | + int yes = 1; |
| 159 | + |
| 160 | + memset(&hints, 0, sizeof hints); |
| 161 | + hints.ai_family = AF_INET; |
| 162 | + hints.ai_socktype = SOCK_STREAM; |
| 163 | + hints.ai_flags = AI_PASSIVE; |
| 164 | + |
| 165 | + // This should only ever listen on a loopback address. Access from outside |
| 166 | + // should be proxied via socat or similar software |
| 167 | + status = getaddrinfo("127.0.0.1", params.listen_port.c_str(), &hints, &servinfo); |
| 168 | + if (status) { |
| 169 | + die("getaddrinfo error: %s", gai_strerror(status)); |
| 170 | + } |
| 171 | + |
| 172 | + // bind to the first addrinfo we can from the getaddrinfo results |
| 173 | + for (p = servinfo; p != NULL; p = p->ai_next) { |
| 174 | + listen_fd = socket(p->ai_family, p->ai_socktype, p->ai_protocol); |
| 175 | + if (listen_fd == -1) { |
| 176 | + perror("server: socket"); |
| 177 | + continue; |
| 178 | + } |
| 179 | + |
| 180 | + if (setsockopt(listen_fd, SOL_SOCKET, SO_REUSEADDR | SO_REUSEPORT, &yes, sizeof yes)) { |
| 181 | + die("setsockopt error: %s", params.listen_port.c_str(), strerror(errno)); |
| 182 | + } |
| 183 | + |
| 184 | + if (bind(listen_fd, p->ai_addr, p->ai_addrlen) == 0) { |
| 185 | + struct sockaddr_in addr_in; |
| 186 | + socklen_t addr_in_len = sizeof(addr_in); |
| 187 | + memset(&addr_in, 0, addr_in_len); |
| 188 | + getsockname(listen_fd, (struct sockaddr*)&addr_in, &addr_in_len); |
| 189 | + |
| 190 | + printf("Listening on %s:%d\n", inet_ntoa(addr_in.sin_addr), ntohs(addr_in.sin_port)); |
| 191 | + break; |
| 192 | + } |
| 193 | + |
| 194 | + close(listen_fd); |
| 195 | + perror("server: bind"); |
| 196 | + } |
| 197 | + |
| 198 | + freeaddrinfo(servinfo); |
| 199 | + |
| 200 | + if (p == NULL) { |
| 201 | + die("failed to bind: %s", strerror(errno)); |
| 202 | + } |
| 203 | + |
| 204 | + if (listen(listen_fd, 20)) { |
| 205 | + die("listen error: %s", strerror(errno)); |
| 206 | + } |
| 207 | + // Don't track child processes, so ignore SIGCHLD to prevent zombies |
| 208 | + signal(SIGCHLD, SIG_IGN); |
| 209 | + |
| 210 | + for (;;) { |
| 211 | + struct sockaddr_in client_addr; |
| 212 | + socklen_t client_addr_len = 0; |
| 213 | + memset(&client_addr, 0, sizeof(client_addr)); |
| 214 | + |
| 215 | + int sock_fd = accept(listen_fd, |
| 216 | + (struct sockaddr *)&client_addr, |
| 217 | + &client_addr_len); |
| 218 | + if (sock_fd < 0) { |
| 219 | + fprintf(stderr, "accept error: %s\n", strerror(errno)); |
| 220 | + break; |
| 221 | + } |
| 222 | + |
| 223 | + child = fork(); |
| 224 | + if (child == 0) { |
| 225 | + // close the listen_fd since we won't use it in the child |
| 226 | + close(listen_fd); |
| 227 | + int ret = serve_model(ctx, params, sock_fd); |
| 228 | + close(sock_fd); |
| 229 | + return ret; |
| 230 | + } else { |
| 231 | + // close the client since we won't use it in the server |
| 232 | + close(sock_fd); |
| 233 | + sock_fd = 0; |
| 234 | + } |
| 235 | + } |
| 236 | + close(listen_fd); |
| 237 | + |
| 238 | + // ignore SIGTERM since we'll send it to the group |
| 239 | + signal(SIGTERM, SIG_IGN); |
| 240 | + // tell children to exit |
| 241 | + kill(0, SIGTERM); |
| 242 | + // wait for children to terminate |
| 243 | + wait(&status); |
| 244 | + return 0; |
| 245 | +} |
0 commit comments