Chia Plotter

plot-o-matic.sh

#!/usr/bin/env bash

######################## REQUIRED ENVIRONMENT VARIABLES ########################
# FARMER_PUBKEY
#   The farmer public key used to generate the plots.
# POOL_PUBKEY
#   The pool public key used to generate the plots.
# SCRATCH_DIR
#   An existing directory where plotters put their temp files.
#   WRANING: Everything under this directory will be deleted.
#   For safety, this cannot be the system's root directory.
# OUTPUT_DIR
#   Where the finished plots should be placed.
#   For safety, this cannot be a subdirectory of SCRATCH_DIR.


######################## OPTIONAL ENVIRONMENT VARIABLES ########################
# FARMER_CA_DIR
# FARMER_HOST
# FARMER_PORT           (default 8447)
#   If supplied, these are used to set up a remote harvester.
#   This will allow you to harvest your plots until they are downloaded
#   or moved to another location.
# CHIA_HOME             (default $HOME/chia-blockchain)
#   The directory where the Chia blockchain tools are (or will be) installed.
# LOG_DIR               (default $HOME/plot-o-matic-logs)
#   A directory to store this script's output in addition to plotter logs.

[[ -z "${FARMER_CA_DIR}" ]] && FARMER_CA_DIR=""
[[ -z "${FARMER_HOST}"   ]] && FARMER_HOST=""
[[ -z "${FARMER_PORT}"   ]] && FARMER_PORT=8447
[[ -z "${CHIA_HOME}"     ]] && CHIA_HOME="${HOME}/chia-blockchain"
[[ -z "${LOG_DIR}"       ]] && LOG_DIR="${HOME}/plot-o-matic-logs"


######################### OPTIONAL PLOTTING VARIABLES ##########################
# PLOTTER_LIMIT         (default 0)
#   The upper bound limit on the number of parallel plotters to run,
#   regardless how much memory, processors, or scratch space is available.
#   If negative or zero, only limit based on available resources.
# PLOTTER_BUCKETS       (default 128)
#   The number of buckets to use during plotting.
# PLOTTER_THREADS       (default 2)
#   The number of threads to use during plotting.
# PLOTTER_MEMORY_MIB    (default 3584 or 3.5 GiB)
#   The amount of memory in MiB (powers of 1024) to use during plotting.
#   This is used to calculate the maximum number of parallel plotters
#   supported by the system's memory.
# PLOTTER_SCRATCH_MIB   (default 245760 or 245 GiB)
#   The amount of scratch space in MiB required for a single plotter.
#   This is used to calculate the maximum number of parallel plotters
#   supported by the SCRATCH_DIR disk capacity.
# MIN_FREE_MEMORY_MIB   (default 512)
#   Amount of memory to reserve for system use.
# EXTRA_PLOTTER_OPTS
#   A string of additional options to pass to the plotter.
# DELAY_SECONDS         (default 600)
#   Wait until the previous plotter has been running this long
#   before spawning another plotter.  If DELAY_UNTIL_TEXT is set as well,
#   the next plotter will run when either of the conditions are met.
# DELAY_UNTIL_TEXT
#   Wait until this text appears in the previous plotter's logs
#   before spawning another plotter.  If DELAY_SECONDS is set as well,
#   the next plotter will run when either of the conditions are met.

[[ -z "${PLOTTER_LIMIT}"       ]] && PLOTTER_LIMIT=0
[[ -z "${PLOTTER_BUCKETS}"     ]] && PLOTTER_BUCKETS=128
[[ -z "${PLOTTER_THREADS}"     ]] && PLOTTER_THREADS=2
[[ -z "${PLOTTER_MEMORY_MIB}"  ]] && PLOTTER_MEMORY_MIB=$(( 3 * 1024 + 512 ))
[[ -z "${PLOTTER_SCRATCH_MIB}" ]] && PLOTTER_SCRATCH_MIB=$(( 245 * 1024 ))
[[ -z "${MIN_FREE_MEMORY_MIB}" ]] && MIN_FREE_MEMORY_MIB=$(( 512 ))
[[ -z "${EXTRA_PLOTTER_OPTS}"  ]] && EXTRA_PLOTTER_OPTS=""
[[ -z "${DELAY_SECONDS}"       ]] && DELAY_SECONDS=$(( 10 * 60 ))
[[ -z "${DELAY_UNTIL_TEXT}"    ]] && DELAY_UNTIL_TEXT=""


status() { echo "[$(date)]" "$@"; }
warning() { status "$@" 1>&2; }
fail() { warning "$@"; exit 1; }

. "/etc/profile"
set -u
set -o pipefail

# Include a date in all local logs so they don't get overwritten between runs.
log_date=$(date +"%Y%m%d-%H%M")
mkdir -p "${LOG_DIR}" || LOG_DIR="/tmp"
exec &> >(tee "${LOG_DIR}/${log_date}-runner.log")


###################### Sanity Checks #######################

export SCRIPT_LOCK="/var/lock/plot-o-matic.lock"
if command -v flock &>/dev/null; then
    exec 9> "/var/lock/plot-o-matic.lock"
    if ! flock --exclusive --nonblocking 9; then
        if (( ${IGNORE_MUTEX:-} )); then
            warning "Another instance of this script is already running"
        else
            fail "Another instance of this script is already running"
        fi
    fi
fi

if [[ ! -d "${SCRATCH_DIR}" ]]; then
    fail "SCRATCH_DIR '${SCRATCH_DIR}' doesn't exist"
fi

if [[ ! -d "${OUTPUT_DIR}" ]]; then
    fail "OUTPUT_DIR '${OUTPUT_DIR}' doesn't exist"
fi

SCRATCH_DIR=$(readlink -f "${SCRATCH_DIR}")
OUTPUT_DIR=$(readlink -f "${OUTPUT_DIR}")
if [[ "${OUTPUT_DIR}" == "${SCRATCH_DIR}" ]]; then
    fail "OUTPUT_DIR cannot be the same as SCRATCH_DIR"

elif [[ "${OUTPUT_DIR}" == "${SCRATCH_DIR}/"* ]]; then
    fail "OUTPUT_DIR cannot be a subdirectory of SCRATCH_DIR"
fi

if [[ "${SCRATCH_DIR}" == "/" ]]; then
    fail "For safety, SCRATCH_DIR cannot be the root directory"
fi

if [[ -z "${FARMER_PUBKEY}" || -z "${POOL_PUBKEY}" ]]; then
    fail "FARMER_PUBKEY and POOL_PUBKEY must be set"
fi


######################## Chia Setup ########################

chia_updated=0
status "Downloading and installing Chia"
if [[ ! -d "${CHIA_HOME}" ]]; then
    clone_log="${LOG_DIR}/${log_date}-download.log"
    if ! git clone "https://github.com/Chia-Network/chia-blockchain.git" \
           -b "latest" \
           --recurse-submodules \
           "${CHIA_HOME}" &>"${clone_log}"
    then
        fail "Failed to clone Chia into ${CHIA_HOME}, see ${clone_log} for details"
    fi
    chia_updated=1
fi

# Update the Chia installation every three days.
# That's `-mtime +2` because mtime looks at `> int(days)`, not `>=`.
if [[ -n "$(find "${CHIA_HOME}" -maxdepth 0 -mtime +2)" ]]; then
    update_log="${LOG_DIR}/${log_date}-update.log"
    if ! (
        cd "${CHIA_HOME}" \
        && git fetch \
        && git checkout -f "latest" \
        && git pull --recurse-submodules \
        && git submodule update --init --recursive
    ) &>"${update-Log}"; then
        fail "Failed to update Chia repo in ${CHIA_HOME}, see ${update_log} for details"
    fi
    touch "${CHIA_HOME}"
    chia_updated=1
fi

if [[ ! -e "${CHIA_HOME}/activate" ]] || (( chia_updated )); then
    build_log="${LOG_DIR}/${log_date}-build.log"
    if ! ( cd "${CHIA_HOME}" && chmod +x "./install.sh" && ./install.sh ) &>"${build_log}"; then
        fail "Failed to build Chia repo in ${CHIA_HOME}, see ${build_log} for details"
    fi
fi

if [[ ! -e "${CHIA_HOME}/activate" ]]; then
    fail "Can't find Chia environment's 'activate' script in ${CHIA_HOME}"
fi

status "Activating Chia environment"
. "${CHIA_HOME}/activate"
chia init
chia plots add -d "${OUTPUT_DIR}"

if [[ -n "${FARMER_CA_DIR}" && -d "${FARMER_CA_DIR}" ]]; then
    status "Adding farmer CA certificates"
    chia init -c "${FARMER_CA_DIR}"
fi

if [[ -n "${FARMER_HOST}" && -n "${FARMER_PORT}" ]]; then
    FARMER_PEER="${FARMER_HOST}:${FARMER_PORT}"
    status "Setting farmer peer to ${FARMER_PEER}"
    chia configure --set-farmer-peer "${FARMER_PEER}"
fi


###################### Plotter Setup #######################

# Calculating disk limits
scratch_size_mib=$(df -m "${SCRATCH_DIR}" --output=size | tail -n 1)
max_scratch_plotters=$(( scratch_size_mib / PLOTTER_SCRATCH_MIB ))
if (( max_scratch_plotters < 1 )); then
    explanation="${scratch_size_mib} MiB total < ${PLOTTER_SCRATCH_MIB} MiB/plotter"
    fail "Scratch volume at '${SCRATCH_DIR}' too small to support plotting (${explanation})"
fi
status "Scratch volume at '${SCRATCH_DIR}' supports up to ${max_scratch_plotters} plotters"

# Calculating processor limits
processor_count=$(grep -c "^processor" "/proc/cpuinfo")
max_cpu_plotters=$(( processor_count ))
if (( max_cpu_plotters < 1 )); then
    fail "This computer doesn't have a CPU?"
fi
status "Processor supports up to ${max_cpu_plotters} plotters"

# Calculating memory limits
mem_size_kib=$(grep "^MemTotal" "/proc/meminfo" | awk '{ print $2; }')
mem_size_mib=$(( mem_size_kib / 1024 ))
usable_mem_size_mib=$(( mem_size_mib - MIN_FREE_MEMORY_MIB ))
max_mem_plotters=$(( usable_mem_size_mib / PLOTTER_MEMORY_MIB ))
if (( max_mem_plotters < 1 )); then
    explanation="${mem_size_mib} MiB total - ${MIN_FREE_MEMORY_MIB} MiB reserved < ${PLOTTER_MEMORY_MIB} MiB/plotter"
    fail "Memory too small to support plotting (${explanation})"
fi
status "Memory supports up to ${max_mem_plotters} at ${PLOTTER_MEMORY_MIB} MiB/plotter"

# Determining the limiting factor
plotter_count=$(printf "%s\n" "${max_scratch_plotters}" "${max_cpu_plotters}" "${max_mem_plotters}" | sort -g | head -n 1)
if (( PLOTTER_LIMIT > 0 && plotter_count > PLOTTER_LIMIT )); then
    plotter_count="${PLOTTER_LIMIT}"
    status "Limiting to ${plotter_count} due to PLOTTER_LIMIT value"

elif (( max_scratch_plotters == plotter_count )); then
    status "Limiting to ${plotter_count} plotters due to scratch size"

elif (( max_mem_plotters == plotter_count )); then
    status "Limiting to ${plotter_count} plotters due to memory"

elif (( max_cpu_plotters == plotter_count )); then
    status "Limiting to ${plotter_count} plotters due to processor count"
fi


##################### Plotter Manager ######################

job_count() { jobs | wc -l; }

memory_avail_mib() {
    local memory_avail_kib=$(grep "^MemAvailable" "/proc/meminfo" | awk '{ print $2; }')
    echo $(( memory_avail_kib / 1024 ))
}

scratch_avail_mib() { df -m "${SCRATCH_DIR}" --output=avail | tail -n 1; }

progress() {
    # Similar to `status`, but suppresses duplicate messages.
    if [[ "$*" != "${last_progress:-}" ]]; then
        status "$@"
        last_progress="$*"
    fi
}

graceful_shutdown() {
    status "Shutdown request received (Ctrl-C again to force stop)"

    while (( $(job_count) > 0 )); do
        progress "Waiting on $(job_count) plotters to complete"
        sleep 10 || break
    done

    # If we broke out of the last loop early, the user requested a force stop.
    if (( $(job_count) > 0 )); then
        status "Killing $(job_count) plotters:"
        jobs -l
        kill $(jobs -p)
    fi

    status "Plotting complete"
    exit
}

maybe_ts() {
    if command -v ts &>/dev/null; then
        ts "$@"
    else
        cat
    fi
}

maybe_flock() {
    if command -v flock &>/dev/null && [[ -n "${LOCKFILE}" ]]; then
        flock --verbose --exclusive "${LOCKFILE}" "$@"
    else
        "$@"
    fi
}

get_lock_for() {
    local target_path="${1}"
    local target_device_id=$(stat --format "%d" "${target_path}")
    echo "/var/lock/plot-o-matic.d${target_device_id}.mv.lock"
}

run_plotter() {
    local plotter_temp=$(mktemp -d -p "${SCRATCH_DIR}")
    {
        status "Starting plotter PID $$"
        if chia plots create \
          --size 32 \
          --farmer_public_key "${FARMER_PUBKEY}" \
          --pool_public_key "${POOL_PUBKEY}" \
          --buckets "${PLOTTER_BUCKETS}" \
          --num_threads "${PLOTTER_THREADS}" \
          --buffer "${PLOTTER_MEMORY_MIB}" \
          --tmp_dir "${plotter_temp}" \
          --final_dir "${plotter_temp}" \
          ${EXTRA_PLOTTER_OPTS} 2>&1
        then
            move_plots "${plotter_temp}/"*".plot"
        else
            warning "Plotter exited with $? status"
        fi
        status "Plotter complete"
    } | maybe_ts -s "[%H:%M:%S]" | tee "${plotter_temp}/plotter.log"
    rm -rf "${plotter_temp}"
}

move_plots() {
    # The default plot moving logic can be a bit dumb and usually ends up doing a file
    # copy instead of a file move even when the directories are on the same mountpoint.
    # Furthermore, parallel copies can result in all copies failing when the destination
    # is nearing full capacity.
    # It's better to get a single full plot copied than no plots at all.
    local output_dir_lock=$(get_lock_for "${OUTPUT_DIR}")
    for plot_path in "$@"; do
        # In case we actually are crossing a disk boundary, move/copy the file to a temp name first.
        # This prevents the plot from being seen by the harvester until the file is fully copied.
        local plot_name=$(basename "${plot_path}")
        local temp_path="${OUTPUT_DIR}/.${plot_name}.tmp"
        local final_path="${OUTPUT_DIR}/${plot_name}"

        status "Moving ${plot_name}"
        if LOCKFILE="${output_dir_lock}" maybe_flock mv -f "${plot_path}" "${temp_path}"; then
            # Rename the plot back to its original name.
            mv -f "${temp_path}" "${final_path}"

        else
            if [[ -e "${final_path}" ]]; then
                rm -f "${final_path}"
            fi
            local delay=$(( 120 + RANDOM % 300 ))
            warning "Failed to move plot, will retry in ${delay} seconds"
            sleep "${delay}"
        fi
    done
}

# If there are any old plot files that didn't finish copying, move them before the purge.
status "Checking for unmoved plot files"
IFS=$'\n' old_plot_files=( $(find "${SCRATCH_DIR}" -type f -name '*.plot') )
if (( ${#old_plot_files[@]} > 0 )); then
    move_plots "${old_plot_files[@]}"
fi

warning "Purging scratch directory '${SCRATCH_DIR}'"
sleep 5 && rm -rf "${SCRATCH_DIR:?}/"*

# If the user requests a stop, try to shutdown gracefully.
# Otherwise, they'll have to manually kill off plotters.
trap "graceful_shutdown" INT

n=1
last_log=""
last_start_time=0
while :; do
    while (( $(job_count) >= plotter_count )); do
        progress "Waiting on any previous plotter to complete..."
        sleep 60
    done

    # We really really really don't want the plotter to OOM or use swap.
    while (( $(memory_avail_mib) < PLOTTER_MEMORY_MIB )); do
        progress "Waiting for enough available memory..."
        sleep 60
    done

    # While the Chia plotter goes into a retry loop on failed writes,
    # it's better to have one plotter successfully run to comletion
    # than have all of the plotter fail.
    # Note that this check isn't foolproof.  It assumes that all running
    # plotters are currently using their peak amount of disk space,
    # which likely isn't the case.
    while (( $(scratch_avail_mib) < PLOTTER_SCRATCH_MIB )); do
        progress "Waiting for enough available scratch space..."
        sleep 60
    done

    while [[ -n "${last_log}" ]]; do
        if (( DELAY_SECONDS <= 0 || SECONDS - last_start_time >= DELAY_SECONDS )); then
            break
        elif [[ -n "${DELAY_UNTIL_TEXT}" ]] && grep -q "${DELAY_UNTIL_TEXT}" "${last_log}"; then
            break
        fi
        progress "Waiting on the previous plotter to make progress..."
        sleep 60
    done

    progress "Spawning plotter ${n}"
    log_date=$(date +"%Y%m%d-%H%M")
    plotter_log="${LOG_DIR}/${log_date}-plotter-${n}.log"
    run_plotter &>"${plotter_log}" &

    (( n += 1 ))
    last_log="${plotter_log}"
    last_start_time="${SECONDS}"
done