Pipeline: Add Runpod docker configuration (#222)

* Begin work on Runpod docker config

* Reduce docker image size

* Create .dockerignore
This commit is contained in:
2026-05-12 00:44:09 -04:00
committed by GitHub
parent 26635ace84
commit b8ebe03921
17 changed files with 425 additions and 90 deletions

View File

@@ -50,6 +50,8 @@ std::optional<ApplicationOptions> ParseArguments(const int argc, char** argv) {
opt("prompt-dir", prog_opts::value<std::string>()->default_value(""),
"Directory containing named prompt files (e.g. BREWERY_GENERATION.md)."
" Required when not using --mocked.");
opt("n-gpu-layers", prog_opts::value<int>()->default_value(0),
"Number of layers to offload to GPU");
};
add_sampling_options();
@@ -85,6 +87,7 @@ std::optional<ApplicationOptions> ParseArguments(const int argc, char** argv) {
const bool use_mocked = var_map["mocked"].as<bool>();
const std::string model_path = var_map["model"].as<std::string>();
const int n_gpu_layers = var_map["n-gpu-layers"].as<int>();
// Enforce mutual exclusivity before any further configuration is applied.
if (use_mocked && !model_path.empty()) {
@@ -110,6 +113,7 @@ std::optional<ApplicationOptions> ParseArguments(const int argc, char** argv) {
options.generator.use_mocked = use_mocked;
options.generator.model_path = model_path;
options.generator.n_gpu_layers = n_gpu_layers;
// Only populate sampling config when the user explicitly overrides at
// least one value. Leaving it as std::nullopt lets LlamaGenerator fall

View File

@@ -89,6 +89,7 @@ LlamaGenerator::LlamaGenerator(
}
n_ctx_ = sampling.n_ctx;
n_gpu_layers_ = options.generator.n_gpu_layers;
this->Load(model_path);
}

View File

@@ -12,6 +12,7 @@
#include <utility>
#include "data_generation/llama_generator.h"
#include "ggml-backend.h"
#include "llama.h"
// Maximum batch size for decode operations. Capping the batch prevents
@@ -22,7 +23,12 @@ void LlamaGenerator::Load(const std::string& model_path) {
context_.reset();
model_.reset();
const llama_model_params model_params = llama_model_default_params();
// Specifically load dynamic ggml backends (like CUDA) that are provided
// externally before attempting to load a model.
ggml_backend_load_all();
llama_model_params model_params = llama_model_default_params();
model_params.n_gpu_layers = n_gpu_layers_;
LlamaGenerator::ModelHandle loaded_model(
llama_model_load_from_file(model_path.c_str(), model_params));
if (!loaded_model) {