88 lines
2.2 KiB
C++
88 lines
2.2 KiB
C++
#pragma once
|
|
|
|
#include <rapidjson/document.h>
|
|
|
|
#include <random>
|
|
|
|
#include "TaskExecutor.hpp"
|
|
|
|
namespace rj = rapidjson;
|
|
|
|
namespace daggy::executors::task {
|
|
|
|
namespace daggy_runner {
|
|
struct Capacity
|
|
{
|
|
ssize_t cores;
|
|
ssize_t memoryMB;
|
|
void operator==(const Capacity &other)
|
|
{
|
|
cores = other.cores;
|
|
memoryMB = other.memoryMB;
|
|
}
|
|
};
|
|
|
|
std::string capacityToJSON(const Capacity &cap);
|
|
Capacity capacityFromJSON(const rj::Value &spec);
|
|
Capacity capacityFromTask(const Task &task);
|
|
|
|
void validateTaskParameters(const ConfigValues &job);
|
|
} // namespace daggy_runner
|
|
|
|
class DaggyRunnerTaskExecutor : public TaskExecutor
|
|
{
|
|
public:
|
|
using Command = std::vector<std::string>;
|
|
|
|
DaggyRunnerTaskExecutor();
|
|
~DaggyRunnerTaskExecutor() override;
|
|
|
|
// Validates the job to ensure that all required values are set and are of
|
|
// the right type,
|
|
bool validateTaskParameters(const ConfigValues &job) override;
|
|
|
|
std::vector<ConfigValues> expandTaskParameters(
|
|
const ConfigValues &job, const ConfigValues &expansionValues) override;
|
|
|
|
// Runs the task
|
|
std::future<AttemptRecord> execute(DAGRunID runID,
|
|
const std::string &taskName,
|
|
const Task &task) override;
|
|
|
|
bool stop(DAGRunID runID, const std::string &taskName) override;
|
|
|
|
std::string description() const;
|
|
|
|
void addRunner(const std::string &url);
|
|
|
|
private:
|
|
void monitor();
|
|
|
|
struct RunningTask
|
|
{
|
|
std::promise<AttemptRecord> prom;
|
|
DAGRunID runID;
|
|
std::string taskName;
|
|
std::string runnerURL;
|
|
uint32_t retries;
|
|
daggy_runner::Capacity resources;
|
|
};
|
|
|
|
// Resolves jobs through polling
|
|
std::atomic<bool> running_;
|
|
std::thread monitorWorker_;
|
|
|
|
struct RunnerCapacity
|
|
{
|
|
daggy_runner::Capacity current;
|
|
daggy_runner::Capacity total;
|
|
};
|
|
std::mutex runnersGuard_;
|
|
std::unordered_map<std::string, RunnerCapacity> runners_;
|
|
|
|
std::mutex rtGuard_;
|
|
std::unordered_map<std::pair<DAGRunID, std::string>, RunningTask>
|
|
runningTasks_;
|
|
};
|
|
} // namespace daggy::executors::task
|