server.cpp

#include <cstdio>
#include <iostream>
#include <algorithm>
#include <memory>
#include <string>
#include <thread>
#include <array>
#include <queue>
#include <mutex>
#include <vector>
#include <filesystem>
#include <crow.h>
#include <ranges>
#include "threadpool.h"
#include <sstream>
#include <chrono>
#define CROW_MAIN

using std::string;
using std::mutex;
using std::lock_guard;
using std::make_shared;
using std::queue;
using std::vector;
namespace fs = std::filesystem;
namespace rv = std::ranges::views;

constexpr string to_st(const auto& i){ std::stringstream ss; ss << i; return ss.str(); }

static inline auto uid(const std::string& s){
	std::stringstream ss;
	std::hash<string> h;
	const auto t0 = std::chrono::system_clock::now();
	ss << s << '|' << std::chrono::duration_cast<std::chrono::nanoseconds>(t0.time_since_epoch()).count();
	return to_st(h(ss.str()));
}

static inline string exec(const char* cmd) {
	std::array<char, 128> buffer;
	std::unique_ptr<FILE, decltype(&pclose)> pipe(popen(cmd, "r"), pclose);
	string result;
	if (!pipe)
		return "Command failed";
	else {
		while (fgets(buffer.data(), buffer.size(), pipe.get()) != nullptr)
			result += buffer.data();
	}
	return result;
}

constexpr auto reify(const std::ranges::range auto& r){ return vector(r.begin(), r.end()); }

constexpr string strip(const string& s){ return s; }

constexpr string strip(string s, char ch){
	s.erase(std::remove_if(s.begin(), s.end(), [=](char c){ return c == ch; }), s.end());
	return s;
}

constexpr string strip(string s, char ch, auto... chs){
	return strip(strip(s, ch), chs...);
}

constexpr vector<string> splitOn(const string& s, const string& delim){
	vector<string> ret;
	long long int start = 0;
	for(size_t dp = s.find_first_of(delim, 0); start >= 0
	   ; start = (dp == std::string::npos) ? -1 : dp + 1, dp = s.find_first_of(delim, start))
		if(auto n = s.substr(start, dp - start); !n.empty())
			ret.emplace_back(n);
	return ret;
}

template <typename T>
constexpr auto q_to_v(queue<T> qcopy){
	vector<T> v;
	v.reserve(qcopy.size());
	while(!qcopy.empty())
		v.push_back(qcopy.front()), qcopy.pop();
	return v;
}

int main(){
	crow::SimpleApp app;
	typedef std::array<string, 2> guy;
	auto commissions = make_shared<queue<guy>>();
	auto queue_mutex = make_shared<mutex>()
	   , train_mutex = make_shared<mutex>();
	auto pool = make_shared<threadpool<>>(avail_threads() / 2);

	auto run = [=](const string& cmd){
		CROW_LOG_INFO << "running \'" << cmd;
		return exec(cmd.c_str());
	};

	auto poppe = [=](){
		lock_guard<mutex> qlock(*queue_mutex);
		commissions->pop();
		CROW_LOG_INFO << commissions->size() << " left in queue";
	};

	auto training_loop = [=](){
		lock_guard<mutex> lock(*train_mutex);
		while(!commissions->empty()){
			auto& [id, prompt] = commissions->front();
			CROW_LOG_INFO << "Launched training for prompt: " + prompt;
			run(string("sh train.sh \"") + strip(prompt, '\'', '\"') + "\"");
			CROW_LOG_INFO << run(string("sh upload.sh ") + id);
			CROW_LOG_INFO << "Finished training for prompt: " + prompt;
			poppe();
		}
	};

	auto enqueue = [=](const guy& thing){
		lock_guard<mutex> lock(*queue_mutex);
		commissions->push(thing);
		auto& [name, prompt] = thing;
		CROW_LOG_INFO << name << " queued with prompt: " << prompt;
	};

	CROW_ROUTE(app, "/create")
		.methods("GET"_method, "POST"_method)([=](const crow::request& req) -> string {
			if(auto prompt = req.url_params.get("prompt"); prompt == nullptr){
				CROW_LOG_INFO << "No prompt specified";
				return "Error: Can't train a NeRF without a prompt!";
			} else {
				CROW_LOG_INFO << prompt << " commissioned";
				auto id = uid(prompt);
				enqueue({id, prompt});
				pool->enqueue(training_loop);
				CROW_LOG_INFO << "Launched training loop";
				return "Scheduled training for " + id;
			}
		});

	CROW_ROUTE(app, "/list")([&](){
		std::vector<string> fin = splitOn(exec("ls *zip"), "\n")
						  , q = reify( q_to_v(*commissions) 
									 | rv::transform([](const guy& i){ return i[0] + ": '" + i[1] + "'"; }));
		crow::json::wvalue ret;
		ret["pending"] = q;
		return ret;
	});

	app.port(80).run();
}