# Model Evaluation Pipeline README

This document provides a technical guide for developers on how to evaluate a model checkpoint using the provided shell script. The pipeline is designed to be executed in an environment with access to an S3-compatible object store and a local scratch space for temporary files.

## Overview

The core of the evaluation process is a two-step pipeline:

1.  **Model Conversion**: A specified model checkpoint is downloaded from an S3 store. The pipeline supports two checkpoint types:

      * **Legacy Checkpoint**: A single file in Composer `.pt` or Flower `.npz` format.
      * **Structured Checkpoint**: A directory containing multiple state files, identified by a `run_uuid`, `desired_steps`, and `n_experts`. This format is used for models with Mixture-of-Experts (MoE) layers.
        The checkpoint is then converted into the standard Hugging Face `safetensors` format, which is required by the evaluation framework.

2.  **Evaluation Execution**: The converted `safetensors` model is passed to the `lighteval` framework via `torchrun` for distributed execution. This step runs a suite of evaluation tasks and saves the results.
3.  The configuration files for which tasks lighteval executes are available in ``mt_dao/conf/lighteval/``.

The entire process is encapsulated in a single bash script for ease of use and reproducibility.

## Prerequisites

Before running the script, ensure the following are configured in your environment:

  * **Python Environment Manager**: The script uses `uv` to manage the Python environment and dependencies.
  * **PyTorch Distributed**: `torchrun` is required to launch the evaluation script.
  * **S3 Access**: The environment must have network access to the S3 endpoint.
  * **Local Scratch Space**: A fast, local storage directory (e.g., `/local/scratch`) is needed to temporarily store the large `safetensors` model files during the run.

## Evaluation Script

Below is the complete, self-contained script for running the evaluation pipeline. It handles both legacy and structured checkpoint types.

```bash
#!/bin/bash
set -e # Exit immediately if a command exits with a non-zero status.

# --- 1. Configuration ---
# Set the project path to your repository's root
PROJECT_PATH="$HOME/projects/mt_dao"

# Set the endpoint for your S3-compatible object store
export S3_ENDPOINT_URL=http://anon.anon:9000

# Define the S3 bucket where checkpoints are stored
CHECKPOINTS_BUCKET="checkpoints"

# Define the model and tokenizer types
# These should match the configurations used during training.
TOKENIZER="SMOLLM_SHARED"
MODEL_TYPE="DEPTOE_4E_MAX"

# --- CHECKPOINT SELECTION ---
# Choose the type of checkpoint to evaluate: "legacy" or "structured"
CHECKPOINT_TYPE="structured" # Or "legacy"

# --- Configuration for a 'structured' checkpoint ---
# Used for MoE models stored as multiple files in a directory structure.
RUN_UUID="f-amd-bs1024-bnxt_re-fsdp-c128-1918k-20250615_16"
DESIRED_STEPS="64"
N_EXPERTS="4" # Number of experts used in the MoE model.

# --- Configuration for a 'legacy' checkpoint ---
# A single .pt or .npz file. Uncomment and set if CHECKPOINT_TYPE="legacy"
# LEGACY_CHECKPOINT_PATH="tune-deptoe_4e_max-20250721_010855/ep0-ba1024-rank0.pt"
# LEGACY_CHECKPOINT_PATH="f-amd-bs1024-bnxt_re-fsdp-c128-1918k-20250615_16/server/64/current_server_parameters.npz"


# Define a temporary local directory to store the converted model
# This should be on a fast local disk (like /local/scratch on an HPC node).
LOCAL_TMPDIR="/local/scratch"

# --- Do not modify below this line ---
# Generate a unique identifier and temporary directory for the evaluation run
if [[ "$CHECKPOINT_TYPE" == "structured" ]]; then
    if [ -z "$RUN_UUID" ] || [ -z "$DESIRED_STEPS" ]; then
        echo "Error: For 'structured' checkpoint type, RUN_UUID and DESIRED_STEPS must be set."
        exit 1
    fi
    # The 'checkpoint_id' is a logical path used for logging and identification.
    CHECKPOINT_ID="${RUN_UUID}/server/${DESIRED_STEPS}"
    SAVE_TENSORS_DIR="${LOCAL_TMPDIR}/SAVE_${RUN_UUID}_${DESIRED_STEPS}_safetensors"
elif [[ "$CHECKPOINT_TYPE" == "legacy" ]]; then
    if [ -z "$LEGACY_CHECKPOINT_PATH" ]; then
        echo "Error: For 'legacy' checkpoint type, LEGACY_CHECKPOINT_PATH must be set."
        exit 1
    fi
    CHECKPOINT_ID="${LEGACY_CHECKPOINT_PATH}"
    SAVE_TENSORS_DIR="${LOCAL_TMPDIR}/SAVE_$(basename ${LEGACY_CHECKPOINT_PATH})_safetensors"
else
    echo "Error: Invalid CHECKPOINT_TYPE. Must be 'legacy' or 'structured'."
    exit 1
fi

# --- 2. Convert the Model Checkpoint to SafeTensors ---
echo "Step 2: Converting checkpoint ${CHECKPOINT_ID} to safetensors..."
CONVERT_CMD="TMPDIR=$LOCAL_TMPDIR uv run python -m mt_dao.eval_light.download_mosaic_model \
    --safe_tensors_dir \"${SAVE_TENSORS_DIR}\" \
    --tokenizer \"${TOKENIZER}\" \
    --checkpoint_dir \"${CHECKPOINTS_BUCKET}\" \
    --model_type \"${MODEL_TYPE}\""

if [[ "$CHECKPOINT_TYPE" == "structured" ]]; then
    CONVERT_CMD+=" --run_uuid \"${RUN_UUID}\" --desired_steps \"${DESIRED_STEPS}\" --n_experts \"${N_EXPERTS}\""
else # legacy
    CONVERT_CMD+=" --checkpoint \"${CHECKPOINT_ID}\""
fi

eval $CONVERT_CMD

# Check if the conversion was successful
if [ $? -ne 0 ]; then
    echo "Error: Model conversion failed."
    exit 1
fi
echo "Conversion complete. SafeTensors saved to ${SAVE_TENSORS_DIR}"

# --- 3. Run the Evaluation ---
echo "Step 3: Starting evaluation..."
# The 'checkpoint' argument here is used for logging/tracking purposes by the eval script.
TMPDIR=$LOCAL_TMPDIR uv run torchrun -m mt_dao.eval_light.eval_mosaic \
    --checkpoint "${CHECKPOINT_ID}" \
    --safe_tensors_dir "${SAVE_TENSORS_DIR}" \
    --model_args "revision=main,dtype=bfloat16,vllm,gpu_memory_utilisation=0.8,max_model_length=2048" \
    --custom_tasks "${PROJECT_PATH}/mt_dao/eval_light/tasks.py" \
    --tasks "${PROJECT_PATH}/mt_dao/conf/lighteval/smollm2_base.txt" \
    --output_dir "./evals_new" \
    --save_details

if [ $? -ne 0 ]; then
    echo "Error: Evaluation run failed."
    # Still attempt to clean up before exiting
else
    echo "Evaluation complete."
fi

# --- 4. Cleanup ---
echo "Step 4: Cleaning up temporary files..."
rm -rf "${SAVE_TENSORS_DIR}"
echo "Cleanup finished."

```

## Usage Guide

1.  **Configure the Script**:

      * Open the script and review the variables in the **Configuration** section.
      * Set `CHECKPOINT_TYPE` to either `"legacy"` or `"structured"`.
      * If using `"structured"`, set the `RUN_UUID`, `DESIRED_STEPS`, and `N_EXPERTS` variables.
      * If using `"legacy"`, set the `LEGACY_CHECKPOINT_PATH` variable to the full S3 path of the checkpoint file relative to the bucket.

2.  **Save the Script**: Save the content as a shell script (e.g., `run_eval.sh`).

3.  **Make it Executable**: `chmod +x run_eval.sh`

4.  **Execute**: `./run_eval.sh`

## Workflow Breakdown

1.  **Configuration**: The script begins by setting environment variables and parameters. Based on the selected `CHECKPOINT_TYPE`, it constructs a unique `CHECKPOINT_ID` for logging and a `SAVE_TENSORS_DIR` for the temporary model files. `TMPDIR` is set to the high-performance local disk to ensure that intermediate files from the `uv run` are also stored efficiently.

2.  **Model Conversion**: The `mt_dao.eval_light.download_mosaic_model` Python module is invoked with parameters corresponding to the chosen checkpoint type.

      * For a **structured** checkpoint, the script is called with `--run_uuid`, `--desired_steps`, and `--n_experts`. It retrieves multiple state files from the appropriate S3 directory (`s3://<bucket>/<run_uuid>/server/<steps>`), merges the expert layers if necessary, and reconstructs the full model state.
      * For a **legacy** checkpoint, the script is called with `--checkpoint`. It handles the logic for downloading the specified single checkpoint file (`.pt` or `.npz`).
        In both cases, it initializes a Composer model, loads the state dictionary, and saves the final model and tokenizer to the `safetensors` format in the directory defined by `SAVE_TENSORS_DIR`.

3.  **Evaluation**: `torchrun` launches the `mt_dao.eval_light.eval_mosaic` module. This script loads the `safetensors` model from the local temporary directory and runs the specified evaluation tasks. The original `CHECKPOINT_ID` is passed for metadata and logging purposes.

4.  **Cleanup**: Upon completion or failure of the evaluation step, the script removes the temporary `SAVE_TENSORS_DIR` and all its contents to free up space on the local scratch disk.
