Chia plot copying script -- To copy from the "final" plot location to farm disks, safely filling farm disks

xfer_plots.py

#!/usr/bin/env python3
import os
from pprint import pprint as pp
import shutil
import socket
import sys
import typing as th


# Largest plot size seen is 108,935,751,977, and smallest is 108,232,347,648. For
# file copying purposes we'll add some buffer on the expected space required for
# a final plot.
MAX_PLOT_SIZE_bytes = 120_000_000_000

# The plot dirs to check for finished .plot files to move to farm dirs...
finishedPlotDirs = [
    "/home/russ/plotstash",
    "/home/russ/chia/plots/d1/final",
    "/home/russ/chia/plots/d2/final",
]

# The list of farm dir destinations. Script will always copy to the emptiest one.
farmDirs = [
    "/home/russ/chia/farms/f1",
    "/home/russ/chia/farms/f2",
    "/home/russ/chia/farms/f3",
    "/home/russ/chia/farms/f4",
    "/home/russ/chia/farms/f5",
    "/home/russ/chia/farms/f6",
    "/home/russ/chia/farms/f7",
    "/home/russ/chia/farms/f8",
    "/home/russ/chia/farms/f9",
    "/home/russ/chia/farms/f10",
    "/home/russ/chia/farms/f11",
    "/home/russ/chia/farms/f12",
    "/home/russ/chia/farms/f13",
    "/home/russ/chia/farms/f14",
    "/home/russ/chia/farms/f15",
    "/home/russ/chia/farms/f16",
    "/home/russ/chia/farms/f17",
    "/home/russ/chia/farms/f18",
    "/home/russ/chia/farms/f19",
    "/home/russ/chia/farms/f20",
    "/home/russ/chia/farms/f21",
    "/home/russ/chia/farms/f22",
    "/home/russ/chia/farms/f23",
]


def ensureSingletonProcess() -> None:
    # Here we'll bind to an abstract socket for a nice and clean singleton implementation
    # without any pidfile for socketfile clutter, no semaphore pain, etc...
    #  - for more on abstract sockets (linux only) see here: https://goo.gl/mnjQfQ
    lockSock = socket.socket(socket.AF_UNIX, socket.SOCK_DGRAM)
    try:
        lockSock.bind(b"\0chia_xfer_plots")
    except socket.error:
        print("Another copier is already running")
        sys.exit(1)
    ensureSingletonProcess._lockSock = lockSock # to prevent garbage collection


def getFreeDiskSpace_bytes(driveDir: str) -> int:
    try:
        return shutil.disk_usage(driveDir).free
    except Exception:
        return 0


def moveFile(src: str, dst: str) -> None:
    print(f"Starting copy of {src} --> {os.path.dirname(dst)}")
    # copy with a large buffer size to help with large file copy
    with open(src, 'rb') as fsrc:
        with open(dst, 'wb') as fdst:
            shutil.copyfileobj(fsrc, fdst, 10485760)  # <-- long process!!
    print(f"Removing: {src}")
    os.remove(src)
    print(f"Move complete!")


def getSoloPlotPaths(farmDir: str) -> th.List[str]:
    soloPlotDir = os.path.join(farmDir, "plots")
    fnames = os.listdir(soloPlotDir)
    fpaths = [os.path.join(farmDir, "plots", fname) for fname in fnames
              if fname.endswith(".plot")]
    return fpaths


def getSoloPlotCount(farmDir: str) -> int:
    try:
        soloPlotPaths = getSoloPlotPaths(farmDir)
        return len(soloPlotPaths)
    except:
        return 0


def smartMovePlot(plotPath: str) -> None:
    # Strategy here is to fill drives up one by one, not spread it out. This is to help with
    # portability, and avoid using up more USB ports than needed.
    farms_free = [(farmDir, getFreeDiskSpace_bytes(farmDir + "/plots")) for farmDir in farmDirs]
    sorted_farms_free = sorted(farms_free, key = lambda t: t[1])  # ascending
    pp(sorted_farms_free)
    for (farmDir, free) in sorted_farms_free: # smallest free to largest free
        if free > MAX_PLOT_SIZE_bytes:
            print(f"Fullest farm that can take a plot is: {farmDir}")
            dstDir = os.path.join(farmDir, "pplots")
            dstPath = os.path.join(dstDir, os.path.basename(plotPath))
            moveFile(plotPath, dstPath)
            return
    print("NO ROOM FOR NEW POOLING PLOTS! Time to swap out solo plots...")

    # no free disk space left! Start replacing solo plots with pooled plots...
    #   - do this a drive at a time, with smallest non-zero solo plot count first
    solo_plots = [(farmDir, getSoloPlotCount(farmDir)) for farmDir in farmDirs]
    sorted_solo_plots = sorted(solo_plots, key = lambda t: t[1])  # ascending

    print("SOLO PLOT COUNTS:")
    pp(sorted_solo_plots)

    for (farmDir, numSolos) in sorted_solo_plots:
        if numSolos > 0:
            print(f"Selected farm (with the least number of solo plots) is: {farmDir}")
            doomedSoloPlotPath = getSoloPlotPaths(farmDir)[0]  # first one (arbitrary)
            dstDir = os.path.join(farmDir, "pplots")
            dstPath = os.path.join(dstDir, os.path.basename(plotPath))
            try:
                os.remove(doomedSoloPlotPath)  # clear space by removing the solo plot
                moveFile(plotPath, dstPath)    # move the new plot into the newly available space
                return  # done!
            except Exception as e:
                print(f"Error caught removing {doomedSoloPlotPath} and/or moving {plotPath}: {e}")
                sys.exit(3)  # if we removed a solo plot, next run will fill the hole the old way

    print("ERROR: NO FREE SPACE, AND NO SOLO PLOTS TO REMOVE!")
    sys.exit(2)


def main() -> None:
    ensureSingletonProcess()
    copyCount = 0
    for finishedPlotDir in finishedPlotDirs:
        for fname in os.listdir(finishedPlotDir):
            if fname.endswith(".plot"):
                smartMovePlot(os.path.join(finishedPlotDir, fname))
                copyCount += 1
        if copyCount == 0:
            print(f"No finished plots in finishedPlotDir {finishedPlotDir}")

    print("DONE!")


if __name__ == "__main__":
    main()