launcher/src/main.rs

#![deny(clippy::wildcard_imports)]
// Opt-in heap profiler (the `dhat` feature): a counting global allocator that
// dumps per-callsite retained-heap stats on graceful exit. Debug builds only
// in practice — the normal release build never pays for it.
#[cfg(feature = "dhat")]
#[global_allocator]
static DHAT_ALLOC: dhat::Alloc = dhat::Alloc;

// mimalloc (add-on image builds): trims back to the OS — see the feature
// doc in Cargo.toml for why the Pi needs it.
#[cfg(feature = "mimalloc")]
#[global_allocator]
static MI_ALLOC: mimalloc::MiMalloc = mimalloc::MiMalloc;

use crate::docker::sync;
use anyhow::Result;
use clap::Parser;
fluent_i18n::i18n!("locales", fallback = "en");

mod after_stop;
mod ansi;
mod before_serve;
mod bench;
mod bench_artifacts;
mod bench_depth;
mod bench_io;
mod bench_remote;
mod bench_report;
mod builtin_scripts;
mod cli;
mod cli_sub;
mod cluster_state;
mod cold;
mod config;
mod detect;
mod docker;
mod doctor;
mod ends;
mod engine;
mod exec;
mod gconfig;
mod handlers;
mod health;
mod home_assistant;
mod host_keys;
mod locale;
mod mikrotik;
mod native;
mod orchestrate;
mod power;
mod pre_serve;
mod promote;
mod publish;
mod publish_render;
mod readiness;
mod recipe_repos;
mod registry;
mod report;
mod runlock;
mod ssh;
mod staging;
mod stamp;
mod tmux;
#[cfg(test)]
mod ui_test_lock;
mod validate;
mod web;
use crate::cli::{Cli, Commands};
use crate::config::{Config, DeployTarget, load_recipe};
use crate::handlers::{handle_exec, handle_health, handle_logs, handle_power};

#[tokio::main]
async fn main() -> Result<()> {
    // dhat heap profiler (the `dhat` feature): alive for the whole run, so
    // the dump at process exit covers every allocation — the retained-heap
    // view points at the leak. Writes dhat-heap.json to cwd on exit.
    #[cfg(feature = "dhat")]
    let _dhat = dhat::Profiler::new_heap();
    // Locale: check LANG env var first (macOS sys-locale ignores it).
    // Normalize "ru_RU.UTF-8" → "ru-RU" for unic-langid.
    let locale = std::env::var("LANG")
        .ok()
        .map(|l| l.split('.').next().unwrap_or(&l).replace('_', "-"));
    let _ = fluent_i18n::set_locale(locale.as_deref());
    locale::set_global_locale(locale.clone()); // tokio workers miss the thread-local; spawn_with_locale re-applies
    tracing_subscriber::fmt()
        .with_env_filter(
            tracing_subscriber::EnvFilter::try_from_default_env()
                .unwrap_or_else(|_| tracing_subscriber::EnvFilter::new("info")),
        )
        .init();
    let cli = Cli::parse();

    let (config_path, raw_config) = crate::gconfig::read_config(&cli.config).await?;
    let mut config: Config = toml::from_str(&raw_config)?;
    crate::config::merge_cluster_sources(&mut config).await?;
    let dry_run = cli.dry_run;
    // Recipe repos: clone/pull the configured git repos into the local
    // cache and append their roots to the recipe + image-build search
    // paths (first repo wins on a recipe-name collision). Best-effort —
    // an unreachable repo degrades to a warning and never blocks a serve
    // of a recipe from a reachable repo. Skipped in dry-run (no git I/O).
    if !config.recipe_repos.is_empty() && !dry_run {
        let roots = crate::recipe_repos::sync_recipe_repos(&config).await;
        for r in &roots {
            // recipe search: the repo's `recipes/` dir (each `<name>/recipe.toml`);
            // image builds: the repo ROOT (its `docker-images/<dir>` build contexts)
            config
                .recipe_sources
                .push(r.join("recipes").display().to_string());
            config.docker_image_sources.push(r.display().to_string());
        }
    }
    // `config` subcommands are pure path/file utilities: dispatch them
    // before cluster_sources merging so they work even with broken sources.
    if let Commands::Config(cmd) = &cli.command {
        return crate::gconfig::run(cmd, &cli.config, dry_run);
    }

    match cli.command {
        Commands::Serve {
            recipe,
            cluster,
            node,
            keep,
            skip_before_serve,
            no_follow,
            no_wait,
            auto_power_on,
            sync_image,
            sync_image_source,
            build_node,
            force_build,
            force,
            solo,
            mirror,
            hf_cache,
            stage,
            fit,
            cold,
            clear,
            no_push,
            image_source,
        } => {
            handle_serve(
                &config,
                &recipe,
                &cluster,
                node,
                docker::ServeOpts {
                    keep,
                    skip_before_serve,
                    no_wait_ready: no_wait,
                    dry_run,
                    auto_power_on,
                    sync_image,
                    sync_image_source,
                    no_push,
                    image_source,
                    build_node,
                    force_build,
                    force,
                    solo,
                    mirror,
                    hf_cache,
                    stage,
                    fit,
                    cold,
                    clear,
                },
                !no_follow,
                dry_run,
            )
            .await?;
        }
        Commands::SyncImage(a) => {
            sync::handle_sync_image(&config, &a.recipe, &a.cluster, a.source.as_deref(), dry_run)
                .await?;
        }
        Commands::Build {
            recipe,
            build_node,
            cluster,
            force_build,
            no_push,
            image_source,
        } => {
            docker::build::handle_build_cmd(
                &config,
                &docker::build::BuildCmdArgs {
                    recipe,
                    build_node,
                    cluster,
                    force_build,
                    no_push,
                    image_source,
                },
                dry_run,
            )
            .await?;
        }
        Commands::Logs {
            recipe: recipe_name,
            cluster,
            follow,
            node,
        } => {
            handle_logs(&config, &recipe_name, &cluster, follow, node, dry_run).await?;
        }
        Commands::Stop {
            recipe: recipe_name,
            cluster,
            node,
            power_off,
            no_wait,
        } => {
            // Native (DGPP) recipes stop via the engine's own cluster manager,
            // not docker labels — route by the recipe's deploy target.
            let recipe_obj =
                crate::config::load_recipe(&recipe_name, &config.recipe_sources).await?;
            if recipe_obj.deploy == crate::config::DeployTarget::Native {
                let cname = if !cluster.is_empty() && cluster != "auto" {
                    cluster.clone()
                } else {
                    recipe_obj.cluster.clone().ok_or_else(|| {
                        anyhow::anyhow!("native stop needs --cluster (or recipe cluster)")
                    })?
                };
                let cl = crate::config::select_cluster_by_name(&config, &cname)?;
                native::native_stop(&config, &recipe_obj, &cl, dry_run).await?;
                if !dry_run {
                    crate::runlock::mark_stop(&recipe_name, &config.recipe_sources)?;
                }
                if power_off {
                    let names: Vec<String> = cl.nodes.iter().map(|n| n.name.clone()).collect();
                    crate::power::power_off(&config, &cl, &names, dry_run).await?;
                }
            } else {
                crate::docker::stop::handle(
                    &config,
                    &recipe_name,
                    &cluster,
                    &node,
                    !no_wait,
                    power_off,
                    dry_run,
                )
                .await?;
            }
        }
        Commands::Bench {
            recipe,
            cluster,
            node,
            label,
            concurrency,
            seed,
            parallel,
            perf_only,
            depth,
            tokenizer,
            ssh_host,
            no_think,
            timeout,
        } => {
            let args = crate::bench::Args {
                recipe,
                cluster,
                node,
                label,
                concurrency,
                seed,
                parallel,
                perf_only,
                depth,
                tokenizer,
                ssh_host,
                no_think,
                timeout,
            };
            crate::bench::run(&config, &args, dry_run).await?;
        }
        Commands::Health {
            cluster: cluster_name,
            fingerprint,
            format,
        } => {
            handle_health(&config, cluster_name, fingerprint, format).await?;
        }
        Commands::Doctor {
            recipe,
            cluster,
            repair,
        } => {
            doctor::doctor(&config, &recipe, &cluster, repair, dry_run).await?;
        }
        Commands::Publish(a) => {
            publish::publish(
                &config,
                &a.recipe,
                a.cluster.as_deref(),
                a.out,
                a.gist,
                &publish::RecipeRegistryOpts {
                    url: a.registry.clone(),
                    token: a.token.clone(),
                },
                publish::PushOpts {
                    push: a.push,
                    node: a.node.clone(),
                    force: a.force,
                },
            )
            .await?;
        }
        Commands::Registry(cmd) => registry::handle(cmd).await?,
        Commands::Tmux(tmux_cmd) => {
            crate::tmux::handle_tmux(&config, tmux_cmd, dry_run).await?;
        }
        Commands::InitHostKeys { cluster } => {
            host_keys::init_host_keys(
                &config,
                cluster.as_deref(),
                config_path.to_str().unwrap_or("config.toml"),
            )
            .await?;
        }
        Commands::Config(_) => unreachable!("dispatched above, before merging"),
        Commands::Web {
            port,
            bind,
            token,
            open,
            insecure,
            repo_sync_interval,
        } => {
            crate::web::run(
                config,
                port,
                bind,
                token,
                open,
                insecure,
                repo_sync_interval,
            )
            .await?;
        }
        Commands::Exec {
            cluster,
            recipe,
            node,
            alias,
            command,
        } => {
            handle_exec(&config, cluster, recipe, node, alias, command, dry_run).await?;
        }
        Commands::Power { action } => {
            handle_power(&config, action, dry_run).await?;
        }
        Commands::Report {
            serve,
            port,
            open,
            archive,
            unarchive,
            archived,
        } => {
            crate::report::run(
                &config,
                serve || open,
                port,
                open,
                archive,
                unarchive,
                archived,
            )
            .await?;
        }
        Commands::Runlock(cmd) => crate::runlock::run(&config, &cmd)?,
        Commands::Repos(a) => {
            let cmd = a.command.clone().unwrap_or(crate::cli::ReposCommands::List);
            crate::recipe_repos::run(&config, &cmd, dry_run).await?
        }
        Commands::Archive { url } => crate::recipe_repos::archive(&url, false, dry_run)?,
        Commands::Unarchive { url } => crate::recipe_repos::archive(&url, true, dry_run)?,
        Commands::Promote {
            recipe,
            from,
            to,
            apply,
        } => crate::promote::run(&config, &recipe, &from, &to, apply, dry_run)?,
        Commands::Validate => {
            let v = crate::validate::validate_all(&config);
            let errs = v
                .iter()
                .filter(|x| x.level == crate::validate::Level::Error)
                .count();
            if v.is_empty() {
                println!("validate: all recipes pass");
            } else {
                for x in &v {
                    let tag = match x.level {
                        crate::validate::Level::Error => "ERROR",
                        crate::validate::Level::Warning => "WARN ",
                    };
                    println!("{tag} {} [{}]: {}", x.recipe, x.rule, x.message);
                }
                eprintln!("validate: {errs} error(s), {} warning(s)", v.len() - errs);
                if errs > 0 {
                    std::process::exit(1);
                }
            }
        }
    }

    Ok(())
}

/// Handle `serve` command: launch inference containers on a cluster.
///
/// Dispatches to `docker::serve` or `tmux::serve` based on the recipe's
/// `deploy` field. Docker-only CLI opts (sync_image, build_node,
/// force_build, keep, auto_power_on) are ignored with a warning when
/// `deploy = "tmux"`.
pub(crate) async fn handle_serve(
    config: &Config,
    recipe: &str,
    cluster: &str,
    node: Option<String>,
    opts: docker::ServeOpts,
    follow: bool,
    dry_run: bool,
) -> Result<()> {
    let recipe_obj = load_recipe(recipe, &config.recipe_sources).await?;
    match recipe_obj.deploy {
        DeployTarget::Docker => {
            let cluster = docker::serve(config, recipe, cluster, node.as_deref(), opts).await?;
            if follow && !dry_run {
                tracing::info!("{}", fluent_i18n::t!("container-log-header"));
                docker::logs::follow_with_liveness(
                    config,
                    &cluster,
                    recipe,
                    recipe_obj.port,
                    node.as_deref(),
                    dry_run,
                )
                .await?;
            }
        }
        DeployTarget::Native => {
            // Head-coordinated native launch (DGPP): the engine's cluster
            // manager runs on rank 0 and brings up the peers itself.
            native::native_serve(config, recipe, cluster, node.as_deref(), &opts).await?;
            if follow && !dry_run {
                // DGPP logs live in its own log dir on rank 0 (no docker
                // label to tail); point the operator at the log file.
                tracing::info!("DGPP serving; engine log on rank 0: ~/dgpp/log/sm121-{recipe}.log");
            }
        }
        DeployTarget::Tmux => {
            // Warn about docker-only CLI opts that are silently ignored.
            let ignored: Vec<&str> = [
                (opts.sync_image, "--sync-image"),
                (opts.sync_image_source.is_some(), "--sync-image-source"),
                (opts.build_node.is_some(), "--build-node"),
                (opts.force_build, "--force-build"),
                (opts.keep, "--keep"),
                (opts.solo.is_some(), "--solo/--no-solo"),
                (opts.mirror.is_some(), "--mirror"),
                (opts.fit.is_some(), "--fit"),
            ]
            .iter()
            .filter_map(|(set, flag)| if *set { Some(*flag) } else { None })
            .collect();
            if !ignored.is_empty() {
                tracing::warn!(
                    "{}",
                    fluent_i18n::t!("tmux-deploy-ignored-opts", { "opts" => &ignored.join(", ") })
                );
            }
            crate::tmux::serve::tmux_serve(
                config,
                recipe,
                cluster,
                opts.skip_before_serve,
                dry_run,
                opts.auto_power_on,
                opts.force,
            )
            .await?;
            if !dry_run {
                // no readiness probe for tmux; cluster stays None (unresolved here)
                crate::runlock::mark_start(recipe, &config.recipe_sources, None)?;
            }
            if follow && !dry_run {
                tracing::info!(
                    "{}",
                    fluent_i18n::t!("tmux-following-logs", { "name" => recipe.to_string() })
                );
                crate::tmux::ops::tmux_logs(config, recipe, cluster, true, dry_run).await?;
            }
        }
    }
    Ok(())
}