File size: 4,121 Bytes
9c774f1 1e9f296 4305cb6 1e9f296 9c774f1 1e9f296 4305cb6 1e9f296 9c774f1 ffb9b72 1e9f296 9c774f1 1e9f296 76cdce6 561ef2b 1e9f296 9c774f1 469ab8d 9c774f1 1e9f296 9c774f1 469ab8d 76cdce6 4305cb6 9c774f1 4305cb6 9c774f1 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 |
#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 \"") + 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, strip(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();
}
|