#!/bin/bash

# Dataloader profiling script with torch.profiler enabled
# This will generate detailed Chrome traces viewable in TensorBoard

CUDA_VISIBLE_DEVICES=0 /home/victor/anaconda3/envs/suno_env_fa2_fa3_pt29/bin/python -u profile_dataloader.py \
    --data_dir="/app2/suno/data/auk_v0" \
    --train_metas_filename="metas_v8_tr_mini.jsonl" \
    --batch_store_size=1 \
    \
    --n_layer=2 \
    --n_head=32 \
    --d_head=128 \
    \
    --block_size=8000 \
    --batch_size=1 \
    \
    --allow_skip=False \
    --num_batches_to_profile=10 \
    --num_workers_oracle=6 \
    \
    --use_torch_profiler=True \
    --torch_profiler_wait=2 \
    --torch_profiler_warmup=2 \
    --torch_profiler_active=5 \
    --torch_profiler_repeat=1 \
    --torch_profiler_output_dir="./profiler_traces"

echo ""
echo "=========================================="
echo "Profiling complete!"
echo "=========================================="

# Upload trace to S3
TRACE_FILE=$(find ./profiler_traces/traces -name "rank0_trace.json.gz" | head -1)
if [ -f "$TRACE_FILE" ]; then
    echo "Uploading trace to S3..."
    aws s3 cp "$TRACE_FILE" s3://suno-data/traces/dataloader/trace.json.gz
    echo "Trace uploaded to: s3://suno-data/traces/dataloader/trace.json.gz"
else
    echo "Warning: Trace file not found"
fi

echo ""
echo "View traces with:"
echo "  1. Download from S3: aws s3 cp s3://suno-data/traces/dataloader/trace.json.gz ."
echo "  2. Open chrome://tracing in Chrome/Edge"
echo "  3. Load the .json.gz file"
echo "=========================================="