feat(inference): enhance CLI with detailed options and improve distributed setup

- Reworked CLI description and examples for better usability and clarity.
- Introduced comprehensive option groups for general, sampling, input, and distributed configurations.
- Standardized parameter validation, including required flags and value checks.
- Improved distributed inference support by refining transport initialization and error handling.

Author: Harikeshav-R

src/inference/main.cpp

@@ -178,164 +178,174 @@ void chat(Transformer *transformer, Tokenizer *tokenizer, Sampler *sampler, cons
 }
 
 int main(int argc, char *argv[]) {
-    CLI::App app{"LEAP Inference Engine\nA high-performance, distributed LLM inference runner.\n"};
+    CLI::App app{"LEAP Inference Engine - High-performance distributed LLM inference"};
+    app.description("A specialized runner for LEAP models supporting single-node and distributed pipelined inference.\n"
+                    "Example (Single Node):\n"
+                    "  ./inference model.bin -p \"Hello, how are you?\"\n\n"
+                    "Example (Distributed Master):\n"
+                    "  ./inference model.bin --role master --next-host 192.168.1.10 --next-port 9999 --split 16\n\n"
+                    "Example (Distributed Worker):\n"
+                    "  ./inference model.bin --role worker --port 9999 --next-host 192.168.1.11 --next-port 9999 --split 16 --end 32");
+    
     argv = app.ensure_utf8(argv);
 
-    std::string checkpoint_path;
+    // Core arguments
+    std::string model_path;
     std::string tokenizer_path = "tokenizer.bin";
+    app.add_option("model", model_path, "Path to the model checkpoint file (.bin)")
+        ->required()
+        ->check(CLI::ExistingFile);
+
+    // General Configuration
+    auto *gen_group = app.add_option_group("General", "Basic inference settings");
+    std::string mode = "generate";
+    bool chat_mode = false;
+    gen_group->add_flag("-c,--chat", chat_mode, "Run in interactive chat mode");
+    gen_group->add_option("-t,--tokenizer", tokenizer_path, "Path to the tokenizer file")
+        ->check(CLI::ExistingFile)
+        ->capture_default_str();
+    int n_predict = 4096;
+    gen_group->add_option("-n,--n-predict", n_predict, "Maximum number of tokens to generate (0 = model max)")
+        ->capture_default_str();
+    unsigned long long rng_seed = 0;
+    gen_group->add_option("--seed", rng_seed, "Random seed for reproducibility (0 = use current time)")
+        ->capture_default_str();
+
+    // Sampling Configuration
+    auto *sample_group = app.add_option_group("Sampling", "Tokens selection parameters");
     float temperature = 1.0f;
+    sample_group->add_option("--temp", temperature, "Temperature for sampling (higher = more creative, 0.0 = greedy)")
+        ->check(CLI::NonNegativeNumber)
+        ->capture_default_str();
     float topp = 0.9f;
-    int steps = 4096;
+    sample_group->add_option("--top-p", topp, "Top-P (nucleus) sampling threshold")
+        ->check(CLI::Range(0.0, 1.0))
+        ->capture_default_str();
+
+    // Input
+    auto *input_group = app.add_option_group("Input", "Prompt configuration");
     std::string prompt;
-    unsigned long long rng_seed = 0;
-    std::string mode = "generate";
+    input_group->add_option("-p,--prompt", prompt, "Initial input prompt for the model");
     std::string system_prompt;
-
-    // Distributed args
-    std::string dist_mode_str = "single";
-    std::string ip = "0.0.0.0";
-    std::string master_ip = "";
+    input_group->add_option("--system", system_prompt, "System prompt for chat mode");
+
+    // Distributed Inference
+    auto *dist_group = app.add_option_group("Distributed", "Configuration for multi-node pipeline parallelism");
+    std::string role = "single";
+    dist_group->add_option("--role", role, "Node role in the cluster")
+        ->check(CLI::IsMember({"single", "master", "worker"}))
+        ->capture_default_str();
+    std::string transport_type = "tcp";
+    dist_group->add_option("--transport", transport_type, "Network transport protocol")
+        ->check(CLI::IsMember({"tcp", "udp", "kernel"}))
+        ->capture_default_str();
+    
+    std::string host = "0.0.0.0";
+    dist_group->add_option("--host", host, "Local IP to bind for incoming connections (Workers)")
+        ->capture_default_str();
     int port = 9999;
-    int split_layer = 0;
-    std::string next_ip = "";
+    dist_group->add_option("--port", port, "Local port to bind for incoming connections (Workers)")
+        ->check(CLI::Range(1, 65535))
+        ->capture_default_str();
+    
+    std::string next_host = "";
+    dist_group->add_option("--next-host", next_host, "IP address of the next node in the pipeline");
     int next_port = 0;
+    dist_group->add_option("--next-port", next_port, "Port of the next node in the pipeline")
+        ->check(CLI::Range(0, 65535));
+    
+    int split_layer = 0;
+    dist_group->add_option("--split", split_layer, "Layer index to start processing on this node")
+        ->capture_default_str();
     int end_layer = 0;
+    dist_group->add_option("--end", end_layer, "Layer index to stop processing (exclusive, 0 = until end)")
+        ->capture_default_str();
 
-    // --- Options ---
-    app.add_option("checkpoint", checkpoint_path, "Path to the model checkpoint file (e.g., model.bin)")
-        ->required()
-        ->check(CLI::ExistingFile);
-
-    // Group: General
-    auto *gen_opts = app.add_option_group("General Configuration");
-    gen_opts->add_option("-z,--tokenizer", tokenizer_path, "Path to the tokenizer file")
-            ->check(CLI::ExistingFile);
-    gen_opts->add_option("-m,--mode", mode, "Inference mode: 'generate' for completion, 'chat' for dialog")
-            ->check(CLI::IsMember({"generate", "chat"}));
-    gen_opts->add_option("-n,--steps", steps, "Maximum number of steps to run (0 = max_seq_len)");
-    gen_opts->add_option("-s,--seed", rng_seed, "Random seed (0 = use time)");
-
-    // Group: Sampling
-    auto *sample_opts = app.add_option_group("Sampling Parameters");
-    sample_opts->add_option("-t,--temperature", temperature, "Temperature for sampling [0.0, inf)")
-               ->check(CLI::NonNegativeNumber);
-    sample_opts->add_option("-p,--top-p", topp, "Top-P (Nucleus) sampling probability [0.0, 1.0]")
-               ->check(CLI::Range(0.0, 1.0));
-
-    // Group: Input
-    auto *input_opts = app.add_option_group("Input");
-    input_opts->add_option("-i,--prompt", prompt, "Initial user prompt");
-    input_opts->add_option("-y,--system-prompt", system_prompt, "System prompt (only used in chat mode)");
-
-    // Group: Distributed
-    auto *dist_opts = app.add_option_group("Distributed Inference", "Configuration for multi-node inference");
-    dist_opts->add_option("--dist", dist_mode_str, "Distributed mode")
-             ->check(CLI::IsMember({"single", "master", "worker", "master-udp", "worker-udp", "worker-kernel"}));
-    
-    dist_opts->add_option("--ip", ip, "Bind IP address (Worker only). Default: 0.0.0.0");
-    dist_opts->add_option("--port", port, "Bind Port (Worker only). Default: 9999")
-             ->check(CLI::Range(1, 65535));
-    dist_opts->add_option("--split", split_layer, "Layer index to split at (start layer for worker)");
-    dist_opts->add_option("--end-layer", end_layer, "Layer index to stop at (exclusive, for worker). Default: n_layers");
-    
-    dist_opts->add_option("--master-ip", master_ip, "Master IP address (required for worker-kernel mode)");
-    
-    dist_opts->add_option("--next-ip", next_ip, "Target IP address for outgoing connection (Master -> Worker 1, Worker -> Next)");
-    dist_opts->add_option("--next-port", next_port, "Target Port for outgoing connection")
-             ->check(CLI::Range(0, 65535)); // 0 means unset
+    std::string master_host = "";
+    dist_group->add_option("--master-host", master_host, "Master IP address (required for kernel transport)");
 
     CLI11_PARSE(app, argc, argv);
 
-    // --- Post-Processing / Defaults ---
-    if (rng_seed <= 0) rng_seed = static_cast<unsigned int>(std::time(nullptr));
-    if (steps < 0) steps = 0;
+    // Post-processing
+    if (chat_mode) mode = "chat";
+    if (rng_seed == 0) rng_seed = static_cast<unsigned int>(std::time(nullptr));
+    if (n_predict < 0) n_predict = 0;
 
     try {
-        auto transformer = Transformer::create(checkpoint_path);
+        auto transformer = Transformer::create(model_path);
 
-        std::cout << "Model loaded successfully." << std::endl;
-        std::cout << "Config: ["
+        std::cout << "Model loaded: " << model_path << std::endl;
+        std::cout << "Architecture: "
                 << transformer->config.n_layers << " layers, "
                 << transformer->config.dim << " dim, "
                 << transformer->config.n_heads << " heads, "
                 << transformer->config.vocab_size << " vocab, "
-                << transformer->config.seq_len << " seq_len]" << std::endl;
+                << transformer->config.seq_len << " context" << std::endl;
 
-        if (steps == 0 || steps > transformer->config.seq_len) {
-            steps = transformer->config.seq_len;
+        if (n_predict == 0 || n_predict > transformer->config.seq_len) {
+            n_predict = transformer->config.seq_len;
         }
 
         if (end_layer == 0) end_layer = transformer->config.n_layers;
 
         // Setup Distributed Mode
-        DistributedMode dist_mode = DistributedMode::Single;
+        DistributedMode dist_role = DistributedMode::Single;
+        if (role == "master") dist_role = DistributedMode::Master;
+        else if (role == "worker") dist_role = DistributedMode::Worker;
+
         std::unique_ptr<Transport> transport = nullptr;
 
-        if (dist_mode_str == "master") {
-            dist_mode = DistributedMode::Master;
+        if (dist_role == DistributedMode::Master) {
             if (split_layer <= 0 || split_layer >= transformer->config.n_layers) {
-                std::cerr << "Error: Invalid split layer for master mode. Must be > 0 and < n_layers." << std::endl;
-                return 1;
+                throw std::runtime_error("Invalid --split layer for master. Must be > 0 and < n_layers.");
             }
-            if (next_ip.empty() || next_port == 0) {
-                std::cerr << "Error: Master mode requires --next-ip and --next-port to connect to the first worker." << std::endl;
-                return 1;
+            if (next_host.empty() || next_port == 0) {
+                throw std::runtime_error("Master role requires --next-host and --next-port to connect to workers.");
             }
 
-            transport = std::make_unique<TcpTransport>(next_ip, next_port, false, "", 0);
-            transport->initialize();
-        } else if (dist_mode_str == "worker") {
-            dist_mode = DistributedMode::Worker;
-            if (split_layer <= 0 || split_layer >= transformer->config.n_layers) {
-                std::cerr << "Error: Invalid split layer for worker mode." << std::endl;
-                return 1;
+            if (transport_type == "tcp") {
+                transport = std::make_unique<TcpTransport>(next_host, next_port, false, "", 0);
+            } else if (transport_type == "udp") {
+                transport = std::make_unique<UdpTransport>(host, port, false, next_host, next_port);
+            } else {
+                throw std::runtime_error("Master role does not support 'kernel' transport (Workers only).");
             }
-            transport = std::make_unique<TcpTransport>(ip, port, true, next_ip, next_port);
-            transport->initialize();
-        } else if (dist_mode_str == "udp") {
-             std::cerr << "Error: Use specific modes: master-udp, worker-udp, worker-kernel" << std::endl;
-             return 1;
-        } else if (dist_mode_str == "master-udp") {
-            dist_mode = DistributedMode::Master;
-            if (next_ip.empty() || next_port == 0) {
-                std::cerr << "Error: Master UDP mode requires --next-ip and --next-port." << std::endl;
-                return 1;
+        } else if (dist_role == DistributedMode::Worker) {
+            if (split_layer <= 0 || split_layer >= transformer->config.n_layers) {
+                throw std::runtime_error("Invalid --split layer for worker.");
             }
-            // Bind to local ip:port, send to next_ip:next_port
-            transport = std::make_unique<UdpTransport>(ip, port, false, next_ip, next_port);
-            transport->initialize();
-        } else if (dist_mode_str == "worker-udp") {
-            dist_mode = DistributedMode::Worker;
-            transport = std::make_unique<UdpTransport>(ip, port, true, next_ip, next_port);
-            transport->initialize();
-        } else if (dist_mode_str == "worker-kernel") {
+            
+            if (transport_type == "tcp") {
+                transport = std::make_unique<TcpTransport>(host, port, true, next_host, next_port);
+            } else if (transport_type == "udp") {
+                transport = std::make_unique<UdpTransport>(host, port, true, next_host, next_port);
+            } else if (transport_type == "kernel") {
 #ifndef __linux__
-            std::cerr << "Error: --dist worker-kernel is only supported on Linux." << std::endl;
-            return 1;
+                throw std::runtime_error("Kernel transport is only supported on Linux.");
 #else
-            dist_mode = DistributedMode::Worker;
-            std::string target_ip = master_ip.empty() ? ip : master_ip;
-            transport = std::make_unique<KernelTransport>(target_ip, port, next_ip, next_port);
-            transport->initialize();
+                std::string target = master_host.empty() ? host : master_host;
+                transport = std::make_unique<KernelTransport>(target, port, next_host, next_port);
 #endif
-        } else if (dist_mode_str != "single") {
-             // CLI11 validation covers this, but safe fallback
-             std::cerr << "Error: Unknown distributed mode: " << dist_mode_str << std::endl;
-             return 1;
+            }
+        }
+
+        if (transport) {
+            transport->initialize();
         }
 
         DistributedConfig dist_config;
-        dist_config.mode = dist_mode;
+        dist_config.mode = dist_role;
         dist_config.split_layer = split_layer;
         dist_config.end_layer = end_layer;
         dist_config.transport = transport.get();
-        dist_config.next_ip = next_ip;
+        dist_config.next_ip = next_host;
         dist_config.next_port = next_port;
-        dist_config.is_tail = (dist_mode == DistributedMode::Worker && next_ip.empty());
+        dist_config.is_tail = (dist_role == DistributedMode::Worker && next_host.empty());
 
         transformer->set_distributed_config(dist_config);
 
-        if (dist_mode == DistributedMode::Worker) {
+        if (dist_role == DistributedMode::Worker) {
+            std::cout << "Node started as Worker [Layers " << split_layer << " to " << end_layer << "]" << std::endl;
             transformer->worker_loop();
             return 0;
         }
@@ -345,12 +355,12 @@ int main(int argc, char *argv[]) {
         Sampler sampler(transformer->config.vocab_size, temperature, topp, rng_seed);
 
         if (mode == "generate") {
-            generate(transformer.get(), &tokenizer, &sampler, prompt, steps);
+            generate(transformer.get(), &tokenizer, &sampler, prompt, n_predict);
         } else if (mode == "chat") {
-            chat(transformer.get(), &tokenizer, &sampler, prompt, system_prompt, steps);
+            chat(transformer.get(), &tokenizer, &sampler, prompt, system_prompt, n_predict);
         }
     } catch (const std::exception &e) {
-        std::cerr << "Runtime Error: " << e.what() << std::endl;
+        std::cerr << "\n[Error] " << e.what() << std::endl;
         return 1;
     }