flodl 0.7.0

//! Worker construction: `channels()` factory + `new()` thread-local builder.

use std::sync::{Arc, Mutex, mpsc};

use crate::autograd::{NoGradGuard, Variable};
use crate::data::BatchDataSet;
use crate::tensor::cuda_event::{CudaEvent, CudaEventFlags};
use crate::tensor::cuda_stream::{CudaStream, StreamGuard};
use crate::distributed::nccl::NcclRankComm;
use crate::nn::{Module, Optimizer, Parameter};
use crate::tensor::{Device, Result, Tensor, TensorError};

use super::super::{
    ApplyPolicy, CheckpointFn, ControlMsg, EvalFn,
    MetricsMsg, ParamSnapshot, TimingMsg, WorkerConfig,
};
use super::{GpuWorker, WorkerChannels, WorkerEndpoints};

impl<M: Module> GpuWorker<M> {
    /// Create the channel pairs for one worker.
    ///
    /// Returns (worker-side senders/receiver, coordinator-side receivers/sender).
    /// Call this on the main thread, then pass the worker-side halves into
    /// [`GpuWorker::new`] inside the spawned thread.
    pub(crate) fn channels() -> (WorkerEndpoints, WorkerChannels) {
        let (timing_tx, timing_rx) = mpsc::channel();
        let (metrics_tx, metrics_rx) = mpsc::channel();
        let (param_tx, param_rx) = mpsc::channel();
        let (final_param_tx, final_param_rx) = mpsc::channel();
        let (control_tx, control_rx) = mpsc::channel();
        (
            (timing_tx, metrics_tx, param_tx, final_param_tx, control_rx),
            WorkerChannels { timing_rx, metrics_rx, param_rx, final_param_rx, control_tx },
        )
    }

    /// Build a GpuWorker inside a spawned thread.
    ///
    /// `model_factory` creates the model on `config.device` (thread-local, Rc-based).
    /// `optim_factory` creates the optimizer for the model's parameters.
    /// `initial_params`/`initial_buffers` from `WorkerConfig` are copied into the
    /// model's Variables to synchronize all workers to the same starting state.
    #[allow(clippy::too_many_arguments)]
    pub(crate) fn new<F, G, O>(
        config: &WorkerConfig,
        model_factory: F,
        optim_factory: G,
        dataset: Arc<dyn BatchDataSet>,
        nccl_comm: Option<NcclRankComm>,
        checkpoint_fn: Option<CheckpointFn<M>>,
        eval_fn: Option<EvalFn<M>>,
        eval_dataset: Option<Arc<dyn BatchDataSet>>,
        timing_tx: mpsc::Sender<TimingMsg>,
        metrics_tx: mpsc::Sender<MetricsMsg>,
        param_tx: mpsc::Sender<ParamSnapshot>,
        final_param_tx: mpsc::Sender<ParamSnapshot>,
        control_rx: mpsc::Receiver<ControlMsg>,
        outer_optimizer: Option<Box<dyn crate::distributed::OuterOptimizer>>,
    ) -> Result<Self>
    where
        F: FnOnce(Device) -> Result<M>,
        G: FnOnce(&[Parameter]) -> O,
        O: Optimizer + 'static,
    {
        // Set the per-thread log prefix so every flodl log line from this
        // worker carries its identity. Thread-based DDP (`Ddp::wrap`: one
        // process hosts every rank) shows `[rN]`. Process-per-rank children
        // are spawned by the cluster launcher, which ALREADY line-prefixes
        // their stdout/stderr with `[host:dev:rN]` (see
        // `launcher::forward_lines`) -- and that wrap also tags raw lines
        // (libtorch warnings, bench `final eval=` / `done:` prints) the
        // in-process logger never sees. Setting the same prefix in-process
        // would double it on flodl-log-macro lines, so skip it there,
        // detected by the launcher's per-rank env marker.
        let local_dev = match config.device {
            Device::CUDA(d) => d,
            _ => 0,
        };
        if std::env::var(crate::distributed::cluster::ENV_LOCAL_RANK).is_err() {
            crate::log::set_thread_device(local_dev, Some(config.rank));
        }

        // Create CUDA streams first (before model construction) so model
        // parameters are allocated on the same stream used by subsequent
        // forward/backward passes. Without this, AccumulateGrad nodes end
        // up on the default stream while gradients arrive on compute_stream,
        // triggering libtorch's "stream does not match" warning and breaking
        // CUDA graph capture.
        let (compute_stream, comm_stream, copy_done) = if config.device.is_cuda() {
            let cs = CudaStream::new(config.device, false)?;
            let ms = CudaStream::new(config.device, false)?;
            let ev = CudaEvent::new(CudaEventFlags::DisableTiming)?;
            // Record initial event so first wait_event is a no-op
            ev.record_on(&ms)?;
            (Some(cs), Some(ms), Some(ev))
        } else {
            (None, None, None)
        };

        // Build the model under compute_stream so every leaf tensor
        // (parameters, buffers) and the AccumulateGrad nodes created at
        // first backward belong to the training stream.
        let model = {
            let _guard = compute_stream.as_ref().map(StreamGuard::new);
            model_factory(config.device)?
        };
        let params = model.parameters();
        let buffers = model.buffers();
        crate::distributed::ddp_run::ensure_trainable_params(
            params.len(),
            &format!("GpuWorker rank {}", config.rank),
        )?;

        // Copy initial params into model variables on compute_stream
        // (no_grad: leaf tensors with requires_grad).
        if params.len() != config.initial_params.len() {
            return Err(TensorError::new(&format!(
                "GpuWorker rank {}: model has {} params but config has {}",
                config.rank, params.len(), config.initial_params.len()
            )));
        }
        {
            let _guard = compute_stream.as_ref().map(StreamGuard::new);
            let _no_grad = NoGradGuard::new();
            for (p, src) in params.iter().zip(&config.initial_params) {
                p.variable.data().copy_(src, false)?;
            }
        }

        // Copy initial buffers into model buffers on compute_stream.
        if buffers.len() != config.initial_buffers.len() {
            return Err(TensorError::new(&format!(
                "GpuWorker rank {}: model has {} buffers but config has {}",
                config.rank, buffers.len(), config.initial_buffers.len()
            )));
        }
        {
            let _guard = compute_stream.as_ref().map(StreamGuard::new);
            for (b, src) in buffers.iter().zip(&config.initial_buffers) {
                b.get().copy_(src, false)?;
            }
        }

        // Eagerly materialize each parameter's AccumulateGrad node under
        // compute_stream and hold a strong reference so it survives
        // between iterations. The node captures the current CUDA stream
        // at construction time into its input_metadata. If the node is
        // GCed and re-created on the autograd engine's worker thread
        // (whose current stream is the device default), libtorch fires
        // the "AccumulateGrad stream does not match" warning on every
        // DDP run that uses a non-default training stream.
        let grad_accumulators: Vec<crate::tensor::GradAccumulatorHandle> = {
            let _guard = compute_stream.as_ref().map(StreamGuard::new);
            let mut handles = Vec::with_capacity(params.len());
            for p in &params {
                if let Some(h) = p.variable.ensure_grad_accumulator()? {
                    handles.push(h);
                }
            }
            handles
        };

        // Create optimizer for this replica's parameters on compute_stream
        // so optimizer state tensors (momentum, Adam moments, ...) are
        // allocated on the same stream as the gradients that will update them.
        let optimizer = {
            let _guard = compute_stream.as_ref().map(StreamGuard::new);
            optim_factory(&params)
        };

        // Extract variable handles (for snapshot/load)
        let param_vars: Vec<Variable> = params.iter().map(|p| p.variable.clone()).collect();
        let buffer_list = buffers;

        // Create prefetch worker for async H2D (VRAM gauge).
        // Cap depth at 512 to avoid huge channel allocations when
        // batch_bytes is tiny (e.g. toy test datasets).
        // Depth 0 = skip prefetch entirely (sync fallback for tight VRAM).
        // Skip entirely when the dataset fits in a single batch (nothing to
        // prefetch ahead of).
        //
        // Note: activation_reserve=0 here because we haven't measured the
        // training activation peak yet. The first run_epoch_plan() will
        // force depth=0 (sync) to calibrate, then adjust on subsequent chunks.
        //
        // Reservation staging: wrap the dataset so the live prefetch path
        // and the background stager share one sample-keyed tier. Dormant
        // (budget 0, pass-through) until the coordinator's first
        // StageAdvisory arrives, so non-progressive runs and tests never
        // pay for it. `FLODL_STAGER=off` disables the whole layer (the
        // dataset stays unwrapped): the A/B discriminator for staging
        // benefit measurements and the operational escape hatch.
        let stager_off = std::env::var("FLODL_STAGER")
            .map(|v| v.eq_ignore_ascii_case("off") || v == "0")
            .unwrap_or(false);
        let (dataset, stager) = if stager_off {
            crate::verbose!("  ddp-worker: rank {} stager disabled (FLODL_STAGER=off)", config.rank);
            if config.disk_stage_gb > 0 {
                crate::verbose!(
                    "  ddp-worker: rank {} disk_stage ignored (the disk tier \
                     lives under the stager's cache, and the stager is off)",
                    config.rank
                );
            }
            (dataset, None)
        } else {
            let stage_cache =
                Arc::new(crate::data::sample_cache::SampleCache::new(dataset.len()));
            // Local-disk overflow tier under the stager's cache — the
            // same RAM → disk → source cascade the solo loader builds
            // (`DataLoaderBuilder::disk_stage`): samples the RAM budget
            // declines spill to an ephemeral per-rank pack file (the
            // pack name is pid-unique, so co-hosted ranks sharing a
            // directory never collide).
            if config.disk_stage_gb > 0 {
                let dir = config
                    .disk_stage_dir
                    .clone()
                    .unwrap_or_else(std::env::temp_dir);
                stage_cache.attach_disk(crate::data::sample_cache::DiskStage::create(
                    &dir,
                    config.disk_stage_gb.saturating_mul(1 << 30),
                    dataset.len(),
                )?);
            }
            let stream_pool = Arc::new(Mutex::new(super::stager::StreamPool::new()));
            let dataset: Arc<dyn BatchDataSet> =
                Arc::new(super::stager::StagedBatchDataSet::new(
                    dataset,
                    Arc::clone(&stage_cache),
                    Arc::clone(&stream_pool),
                ));
            let stager = super::stager::spawn_stager(
                Arc::clone(&dataset),
                stage_cache,
                stream_pool,
                config.seed,
                config.rank,
                config.world_size,
                config.augment,
                config.ram_max_usage,
                config.sample_cache,
            );
            (dataset, Some(stager))
        };

        // The epoch's work is picks (samples × augment views).
        let total_batches =
            dataset.len() * config.augment.max(1) / config.batch_size.max(1);
        // Device sample pool switch: builder/config knob AND the
        // `FLODL_VRAM_POOL=off` runtime kill-switch (A/B runs) — one
        // shared parse with the solo loader (audit D7).
        let pool_off = crate::data::vram_pool::vram_pool_env_off();
        let vram_pool_enabled = config.vram_pool && !pool_off;
        let (prefetch, per_sample_bytes) = if config.device.is_cuda() && total_batches > 1 {
            let sample = dataset.get_batch(&[0])?;
            let psb: usize = sample.iter().map(|t| t.nbytes()).sum();
            drop(sample);
            let depth = crate::data::prefetch_depth_from_vram(
                psb, config.batch_size, config.device, config.vram_max_usage, 0,
            ).min(512);
            crate::debug!(
                "  ddp-worker: rank {} constructor prefetch sizing: psb={} depth={} (used, total)={:?}",
                config.rank, psb, depth,
                crate::tensor::cuda_memory_info_idx(config.device.index() as i32)
            );
            // Reset peak stats so first run_epoch_plan gets a clean baseline.
            crate::tensor::cuda_reset_peak_stats_idx(config.device.index() as i32);
            if depth > 0 {
                // Device sample pool: coordinator-paced epochs have no
                // governor, so the worker signals the pool's budget
                // moment itself at the first post-calibration plan
                // boundary (activation peak measured = the honest
                // probe).
                if config.vram_pool && pool_off {
                    crate::verbose!(
                        "  ddp-worker: rank {} vram pool disabled (FLODL_VRAM_POOL=off)",
                        config.rank
                    );
                }
                crate::debug!(
                    "  ddp-worker: rank {} prefetch depth={} vram_pool={}",
                    config.rank, depth, vram_pool_enabled
                );
                let pw = crate::data::prefetch::PrefetchWorker::new(
                    Arc::clone(&dataset), config.device, depth,
                    vram_pool_enabled,
                    config.augment,
                );
                (Some(pw), psb)
            } else {
                (None, psb)
            }
        } else {
            (None, 0)
        };

        // Allocate scratch buffers for weight-space divergence
        // measurement AND for cluster-mode NCCL abort recovery (the
        // retry path in `sync_now_nccl` restores params from this
        // scratch after a peer-death abort). Allocated whenever an
        // NCCL comm is attached — the divergence value is near-zero in
        // Sync mode but the recovery path needs the buffer regardless,
        // so paying the alloc once is simpler than threading a
        // `cluster_mode` flag through the constructor.
        // Alloc failure must fail construction loudly: a swallowed None
        // here is indistinguishable from "no comm attached" and only
        // surfaces later as a misleading abort-retry error (params) or a
        // silent restore skip (buffers).
        let pre_sync_scratch = if nccl_comm.is_some() {
            let scratch: Vec<Tensor> = param_vars.iter()
                .map(|v| Tensor::zeros_like(&v.data()))
                .collect::<Result<_>>()
                .map_err(|e| TensorError::new(&format!(
                    "GpuWorker r{}: pre-sync param scratch alloc failed: {e}",
                    config.rank,
                )))?;
            Some(scratch)
        } else {
            None
        };
        // Companion scratch for the f32 buffers riding the NCCL sync
        // (mover-averaged running stats) — same abort-recovery rationale,
        // same gating. Tiny next to the param scratch (per-channel stat
        // vectors vs full weight matrices). Empty (no f32 buffers) is the
        // one legitimate None once a comm is attached.
        let pre_sync_buffer_scratch = if nccl_comm.is_some() {
            let scratch: Vec<Tensor> = buffer_list.iter()
                .map(|b| b.get())
                .filter(|t| t.dtype() == crate::tensor::DType::Float32)
                .map(|t| Tensor::zeros_like(&t))
                .collect::<Result<_>>()
                .map_err(|e| TensorError::new(&format!(
                    "GpuWorker r{}: pre-sync buffer scratch alloc failed: {e}",
                    config.rank,
                )))?;
            (!scratch.is_empty()).then_some(scratch)
        } else {
            None
        };

        // Adopt the model's shared aggregated-metrics slot (Graph
        // exposes one; other Modules default to None and get a
        // private slot). Captured BEFORE the `model` move into
        // `GpuWorker` below so the slot lookup still has a reference
        // to read from.
        let aggregated_slot = model
            .aggregated_metrics_slot()
            .unwrap_or_else(|| Arc::new(Mutex::new(None)));
        Ok(GpuWorker {
            model,
            optimizer: Box::new(optimizer),
            param_vars,
            buffer_list,
            rank: config.rank,
            world_size: config.world_size,
            device: config.device,
            epoch_callback_role: None,
            compute_stream,
            comm_stream,
            copy_done,
            pending_param_h2d: false,
            last_h2d_wait_ms: 0.0,
            last_update_at: None,
            h2d_wait_ms_total: 0.0,
            prof_enabled: crate::log::enabled(crate::log::Verbosity::Debug),
            snapshot_ns_total: 0,
            snapshot_count: 0,
            snapshot_pinned_params: Vec::new(),
            snapshot_pinned_buffers: Vec::new(),
            pinned_fallback_logged: false,
            compute_ms_run_total: 0.0,
            data_ms_run_total: 0.0,
            ctrl_msgs_handled: 0,
            nccl_abort_handle: nccl_comm.as_ref().map(|c| c.abort_handle()),
            nccl_abort_slot: None,
            nccl_comm,
            nccl_session_mailbox: None,
            local_dead_ranks: None,
            timing_tx,
            metrics_tx,
            param_tx,
            final_param_tx,
            control_rx,
            dataset,
            partition: Vec::new(), // filled by first StartEpoch from coordinator
            batch_size: config.batch_size,
            base_seed: config.seed,
            augment: config.augment.max(1),
            transform: config.transform.clone(),
            local_step: 0,
            nccl_sync_seq: 0,
            steps_since_avg: 0,
            steps_at_snapshot: 0,
            gamma: config.gamma,
            bf16_wire: config.bf16_wire,
            current_version: 0,
            current_epoch: 0,
            pending_plan: None,
            global_step: 0,
            scheduler: None,
            lr_scale: 1.0,
            aggregated_metrics: aggregated_slot,
            metrics_stream_tx: None,
            eval_stream_tx: None,
            checkpoint_fn,
            eval_fn,
            eval_dataset,
            save_path: config.save_path.clone(),
            prefetch,
            stager,
            per_sample_bytes,
            vram_max_usage: config.vram_max_usage,
            ram_max_usage: config.ram_max_usage,
            activation_peak_bytes: 0,
            // Nothing to signal when the pool is off: pre-latch so the
            // install boundary (and its flow-reserve depth collapse)
            // never runs for a disabled pool.
            vram_pool_budget_sent: !vram_pool_enabled,
            max_grad_norm: config.max_grad_norm,
            // EASGD elastic blending is an Async-only concept: Sync and
            // Cadence MUST full-overwrite to the consensus each window. Gate
            // structurally on the worker's policy here -- the single point
            // every worker (threaded, single-host, cluster) is built through
            // -- so a stray `easgd_alpha` from ANY upstream config path can
            // never blend a non-async worker. The config value alone is too
            // weak a guard for this invariant, and the only honored value is
            // already mode-defaulted to `Some` for CpuAsync.
            easgd_alpha: if matches!(config.policy, ApplyPolicy::Async) {
                config.easgd_alpha
            } else {
                None
            },
            timeline: config.timeline.clone(),
            pre_sync_scratch,
            pre_sync_buffer_scratch,
            outer_optimizer,
            outer_prev_global: None,
            _grad_accumulators: grad_accumulators,
        })
    }
}