chia-blockchain 2.5.1rc1__py3-none-any.whl

chia/plotters/plotters.py

@@ -0,0 +1,577 @@
+from __future__ import annotations
+
+import argparse
+import binascii
+import os
+from enum import Enum
+from pathlib import Path
+from typing import Any, Optional
+
+from chia.plotters.bladebit import get_bladebit_install_info, plot_bladebit
+from chia.plotters.chiapos import get_chiapos_install_info, plot_chia
+from chia.plotters.madmax import get_madmax_install_info, plot_madmax
+
+
+class Options(Enum):
+    TMP_DIR = 1
+    TMP_DIR2 = 2
+    FINAL_DIR = 3
+    K = 4
+    MEMO = 5
+    ID = 6
+    BUFF = 7
+    NUM_BUCKETS = 8
+    STRIPE_SIZE = 9
+    NUM_THREADS = 10
+    NOBITFIELD = 11
+    PLOT_COUNT = 12
+    MADMAX_NUM_BUCKETS_PHRASE3 = 13
+    MADMAX_WAITFORCOPY = 14
+    POOLKEY = 15
+    FARMERKEY = 16
+    MADMAX_TMPTOGGLE = 17
+    POOLCONTRACT = 18
+    MADMAX_RMULTI2 = 19
+    BLADEBIT_WARMSTART = 20
+    BLADEBIT_NONUMA = 21
+    VERBOSE = 22
+    OVERRIDE_K = 23
+    ALT_FINGERPRINT = 24
+    EXCLUDE_FINAL_DIR = 25
+    CONNECT_TO_DAEMON = 26
+    BLADEBIT_NO_CPU_AFFINITY = 27
+    BLADEBIT_CACHE = 28
+    BLADEBIT_F1_THREAD = 29
+    BLADEBIT_FP_THREAD = 30
+    BLADEBIT_C_THREAD = 31
+    BLADEBIT_P2_THREAD = 32
+    BLADEBIT_P3_THREAD = 33
+    BLADEBIT_ALTERNATE = 34
+    BLADEBIT_NO_T1_DIRECT = 35
+    BLADEBIT_NO_T2_DIRECT = 36
+    COMPRESSION = 37
+    BLADEBIT_DEVICE_INDEX = 38
+    CUDA_TMP_DIR = 40
+    BLADEBIT_HYBRID_128_MODE = 41
+    BLADEBIT_HYBRID_16_MODE = 42
+
+
+chia_plotter_options = [
+    Options.TMP_DIR,
+    Options.TMP_DIR2,
+    Options.K,
+    Options.MEMO,
+    Options.ID,
+    Options.BUFF,
+    Options.NUM_BUCKETS,
+    Options.STRIPE_SIZE,
+    Options.NUM_THREADS,
+    Options.NOBITFIELD,
+    Options.OVERRIDE_K,
+    Options.ALT_FINGERPRINT,
+    Options.POOLCONTRACT,
+    Options.FARMERKEY,
+    Options.POOLKEY,
+    Options.PLOT_COUNT,
+    Options.EXCLUDE_FINAL_DIR,
+    Options.CONNECT_TO_DAEMON,
+    Options.FINAL_DIR,
+    Options.COMPRESSION,
+]
+
+madmax_plotter_options = [
+    Options.K,
+    Options.PLOT_COUNT,
+    Options.NUM_THREADS,
+    Options.NUM_BUCKETS,
+    Options.MADMAX_NUM_BUCKETS_PHRASE3,
+    Options.TMP_DIR,
+    Options.TMP_DIR2,
+    Options.MADMAX_WAITFORCOPY,
+    Options.POOLKEY,
+    Options.FARMERKEY,
+    Options.POOLCONTRACT,
+    Options.MADMAX_TMPTOGGLE,
+    Options.MADMAX_RMULTI2,
+    Options.CONNECT_TO_DAEMON,
+    Options.FINAL_DIR,
+]
+
+bladebit_cuda_plotter_options = [
+    Options.NUM_THREADS,
+    Options.PLOT_COUNT,
+    Options.FARMERKEY,
+    Options.POOLKEY,
+    Options.POOLCONTRACT,
+    Options.CUDA_TMP_DIR,
+    Options.TMP_DIR2,
+    Options.ID,
+    Options.BLADEBIT_WARMSTART,
+    Options.BLADEBIT_NONUMA,
+    Options.BLADEBIT_NO_CPU_AFFINITY,
+    Options.VERBOSE,
+    Options.CONNECT_TO_DAEMON,
+    Options.FINAL_DIR,
+    Options.COMPRESSION,
+    Options.BLADEBIT_DEVICE_INDEX,
+    Options.BLADEBIT_HYBRID_128_MODE,
+    Options.BLADEBIT_HYBRID_16_MODE,
+]
+
+bladebit_ram_plotter_options = [
+    Options.NUM_THREADS,
+    Options.PLOT_COUNT,
+    Options.FARMERKEY,
+    Options.POOLKEY,
+    Options.POOLCONTRACT,
+    Options.ID,
+    Options.BLADEBIT_WARMSTART,
+    Options.BLADEBIT_NONUMA,
+    Options.BLADEBIT_NO_CPU_AFFINITY,
+    Options.VERBOSE,
+    Options.CONNECT_TO_DAEMON,
+    Options.FINAL_DIR,
+    Options.COMPRESSION,
+]
+
+bladebit_disk_plotter_options = [
+    Options.NUM_THREADS,
+    Options.PLOT_COUNT,
+    Options.FARMERKEY,
+    Options.POOLKEY,
+    Options.POOLCONTRACT,
+    Options.ID,
+    Options.BLADEBIT_WARMSTART,
+    Options.BLADEBIT_NONUMA,
+    Options.VERBOSE,
+    Options.CONNECT_TO_DAEMON,
+    Options.FINAL_DIR,
+    Options.BLADEBIT_NO_CPU_AFFINITY,
+    Options.BLADEBIT_CACHE,
+    Options.BLADEBIT_F1_THREAD,
+    Options.BLADEBIT_FP_THREAD,
+    Options.BLADEBIT_C_THREAD,
+    Options.BLADEBIT_P2_THREAD,
+    Options.BLADEBIT_P3_THREAD,
+    Options.BLADEBIT_ALTERNATE,
+    Options.TMP_DIR,
+    Options.TMP_DIR2,
+    Options.NUM_BUCKETS,
+    Options.MEMO,
+    Options.BLADEBIT_NO_T1_DIRECT,
+    Options.BLADEBIT_NO_T2_DIRECT,
+    Options.COMPRESSION,
+]
+
+
+def get_plotters_root_path(root_path: Path) -> Path:
+    return root_path / "plotters"
+
+
+def build_parser(subparsers, root_path, option_list, name, plotter_desc):
+    parser = subparsers.add_parser(name, help=plotter_desc)
+    for option in option_list:
+        if option is Options.K:
+            parser.add_argument(
+                "-k",
+                "--size",
+                type=int,
+                help="K value.",
+                default=32,
+            )
+        if option is Options.NUM_BUCKETS:
+            u_default = 0 if name == "chiapos" else 256
+            parser.add_argument(
+                "-u",
+                "--buckets",
+                type=int,
+                help="Number of buckets.",
+                default=u_default,
+            )
+        if option is Options.STRIPE_SIZE:
+            parser.add_argument(
+                "-s",
+                "--stripes",
+                type=int,
+                help="Stripe size.",
+                default=0,
+            )
+        if option is Options.CUDA_TMP_DIR:
+            parser.add_argument(
+                "-t",
+                "--tmp_dir",
+                type=str,
+                dest="tmpdir",
+                help="Temporary directory 1.",
+                required=False,  # Unlike `Options.TMP_DIR`, this is not required
+            )
+        if option is Options.TMP_DIR:
+            parser.add_argument(
+                "-t",
+                "--tmp_dir",
+                type=str,
+                dest="tmpdir",
+                help="Temporary directory 1.",
+                required=True,
+            )
+        if option is Options.TMP_DIR2:
+            parser.add_argument(
+                "-2",
+                "--tmp_dir2",
+                type=str,
+                dest="tmpdir2",
+                help="Temporary directory 2.",
+                default="",
+            )
+        if option is Options.FINAL_DIR:
+            parser.add_argument(
+                "-d",
+                "--final_dir",
+                type=str,
+                dest="finaldir",
+                help="Final directory.",
+                required=True,
+            )
+        if option is Options.BUFF:
+            parser.add_argument(
+                "-b",
+                "--buffer",
+                type=int,
+                help="Size of the buffer, in MB.",
+                default=0,
+            )
+        if option is Options.NUM_THREADS:
+            r_default = 4 if name == "madmax" else 0
+            parser.add_argument(
+                "-r",
+                "--threads",
+                type=int,
+                help="Num threads.",
+                default=r_default,
+            )
+        if option is Options.NOBITFIELD:
+            parser.add_argument(
+                "-e",
+                "--nobitfield",
+                action="store_true",
+                help="Disable bitfield.",
+                default=False,
+            )
+        if option is Options.MEMO:
+            parser.add_argument(
+                "-m",
+                "--memo",
+                type=binascii.unhexlify,
+                help="Memo variable.",
+            )
+        if option is Options.ID:
+            parser.add_argument(
+                "-i",
+                "--id",
+                type=binascii.unhexlify,
+                help="Plot id",
+            )
+        if option is Options.PLOT_COUNT:
+            parser.add_argument(
+                "-n",
+                "--count",
+                type=int,
+                help="Number of plots to create (default = 1)",
+                default=1,
+            )
+        if option is Options.MADMAX_NUM_BUCKETS_PHRASE3:
+            parser.add_argument(
+                "-v",
+                "--buckets3",
+                type=int,
+                help="Number of buckets for phase 3+4 (default = 256)",
+                default=256,
+            )
+        if option is Options.MADMAX_WAITFORCOPY:
+            parser.add_argument(
+                "-w",
+                "--waitforcopy",
+                action="store_true",
+                help="Wait for copy to start next plot",
+                default=False,
+            )
+        if option is Options.MADMAX_TMPTOGGLE:
+            parser.add_argument(
+                "-G",
+                "--tmptoggle",
+                action="store_true",
+                help="Alternate tmpdir/tmpdir2 (default = false)",
+                default=False,
+            )
+        if option is Options.POOLCONTRACT:
+            parser.add_argument(
+                "-c",
+                "--contract",
+                type=str,
+                help="Pool Contract Address (64 chars)",
+                default="",
+            )
+        if option is Options.MADMAX_RMULTI2:
+            parser.add_argument(
+                "-K",
+                "--rmulti2",
+                type=int,
+                help="Thread multiplier for P2 (default = 1)",
+                default=1,
+            )
+        if option is Options.POOLKEY:
+            parser.add_argument(
+                "-p",
+                "--pool-key",
+                type=binascii.unhexlify,
+                help="Pool Public Key (48 bytes)",
+                default="",
+            )
+        if option is Options.FARMERKEY:
+            parser.add_argument(
+                "-f",
+                "--farmerkey",
+                type=binascii.unhexlify,
+                help="Farmer Public Key (48 bytes)",
+                default="",
+            )
+        if option is Options.BLADEBIT_WARMSTART:
+            parser.add_argument(
+                "-w",
+                "--warmstart",
+                action="store_true",
+                help="Warm start",
+                default=False,
+            )
+        if option is Options.BLADEBIT_NONUMA:
+            parser.add_argument(
+                "--nonuma",
+                action="store_true",
+                help="Disable numa",
+                default=False,
+            )
+        if option is Options.VERBOSE:
+            parser.add_argument(
+                "-v",
+                "--verbose",
+                action="store_true",
+                help="Set verbose",
+                default=False,
+            )
+        if option is Options.OVERRIDE_K:
+            parser.add_argument(
+                "--override-k",
+                dest="override",
+                action="store_true",
+                help="Force size smaller than 32",
+                default=False,
+            )
+        if option is Options.ALT_FINGERPRINT:
+            parser.add_argument(
+                "-a",
+                "--alt_fingerprint",
+                type=int,
+                default=None,
+                help="Enter the alternative fingerprint of the key you want to use",
+            )
+        if option is Options.EXCLUDE_FINAL_DIR:
+            parser.add_argument(
+                "-x",
+                "--exclude_final_dir",
+                action="store_true",
+                help="Skips adding [final dir] to harvester for farming",
+                default=False,
+            )
+        if option is Options.CONNECT_TO_DAEMON:
+            parser.add_argument(
+                "-D",
+                "--connect-to-daemon",
+                action="store_true",
+                help=argparse.SUPPRESS,
+                default=False,
+            )
+        if option is Options.BLADEBIT_NO_CPU_AFFINITY:
+            parser.add_argument(
+                "--no-cpu-affinity",
+                action="store_true",
+                help="Disable assigning automatic thread affinity",
+                default=False,
+            )
+        if option is Options.BLADEBIT_CACHE:
+            parser.add_argument(
+                "--cache",
+                type=str,
+                help="Size of cache to reserve for I/O",
+            )
+        if option is Options.BLADEBIT_F1_THREAD:
+            parser.add_argument(
+                "--f1-threads",
+                type=int,
+                help="Override the thread count for F1 generation",
+            )
+        if option is Options.BLADEBIT_FP_THREAD:
+            parser.add_argument(
+                "--fp-threads",
+                type=int,
+                help="Override the thread count for forward propagation",
+            )
+        if option is Options.BLADEBIT_C_THREAD:
+            parser.add_argument(
+                "--c-threads",
+                type=int,
+                help="Override the thread count for C table processing",
+            )
+        if option is Options.BLADEBIT_P2_THREAD:
+            parser.add_argument(
+                "--p2-threads",
+                type=int,
+                help="Override the thread count for Phase 2",
+            )
+        if option is Options.BLADEBIT_P3_THREAD:
+            parser.add_argument(
+                "--p3-threads",
+                type=int,
+                help="Override the thread count for Phase 3",
+            )
+        if option is Options.BLADEBIT_ALTERNATE:
+            parser.add_argument(
+                "--alternate",
+                action="store_true",
+                help="Halves the temp2 cache size requirements by alternating bucket writing methods between tables",
+                default=False,
+            )
+        if option is Options.BLADEBIT_NO_T1_DIRECT:
+            parser.add_argument(
+                "--no-t1-direct",
+                action="store_true",
+                help="Disable direct I/O on the temp 1 directory",
+                default=False,
+            )
+        if option is Options.BLADEBIT_NO_T2_DIRECT:
+            parser.add_argument(
+                "--no-t2-direct",
+                action="store_true",
+                help="Disable direct I/O on the temp 2 directory",
+                default=False,
+            )
+        if option is Options.COMPRESSION:
+            parser.add_argument(
+                "--compress",
+                type=int,
+                help="Compression level",
+                default=1,
+            )
+        if option is Options.BLADEBIT_DEVICE_INDEX:
+            parser.add_argument(
+                "--device",
+                type=int,
+                help="The CUDA device index",
+                default=0,
+            )
+        if option is Options.BLADEBIT_HYBRID_128_MODE:
+            parser.add_argument(
+                "--disk-128",
+                action="store_true",
+                help="Enable hybrid disk plotting for 128G system RAM",
+                default=False,
+            )
+        if option is Options.BLADEBIT_HYBRID_16_MODE:
+            parser.add_argument(
+                "--disk-16",
+                action="store_true",
+                help="Enable hybrid disk plotting for 16G system RAM",
+                default=False,
+            )
+
+
+def call_plotters(root_path: Path, args):
+    # Add `plotters` section in CHIA_ROOT.
+    chia_root_path = root_path
+    root_path = get_plotters_root_path(root_path)
+    if not root_path.is_dir():
+        if os.path.exists(root_path):
+            try:
+                os.remove(root_path)
+            except Exception as e:
+                print(f"Exception deleting old root path: {type(e)} {e}.")
+
+    if not os.path.exists(root_path):
+        print(f"Creating plotters folder within CHIA_ROOT: {root_path}")
+        try:
+            os.mkdir(root_path)
+        except Exception as e:
+            print(f"Cannot create plotters root path {root_path} {type(e)} {e}.")
+
+    plotters = argparse.ArgumentParser("chia plotters", description="Available options.")
+    subparsers = plotters.add_subparsers(help="Available options", dest="plotter")
+
+    build_parser(subparsers, root_path, chia_plotter_options, "chiapos", "Create a plot with the default chia plotter")
+    build_parser(subparsers, root_path, madmax_plotter_options, "madmax", "Create a plot with madMAx")
+
+    bladebit_parser = subparsers.add_parser("bladebit", help="Create a plot with bladebit")
+    subparsers_bb = bladebit_parser.add_subparsers(dest="plot_type", required=True)
+    build_parser(subparsers_bb, root_path, bladebit_cuda_plotter_options, "cudaplot", "Creat a plot using CUDA")
+    build_parser(subparsers_bb, root_path, bladebit_ram_plotter_options, "ramplot", "Create a plot using RAM")
+    build_parser(subparsers_bb, root_path, bladebit_disk_plotter_options, "diskplot", "Create a plot using disk")
+
+    subparsers.add_parser("version", help="Show plotter versions")
+
+    deprecation_warning = (
+        "[DEPRECATED] 'chia plotters install' is no longer available. Use install-plotter.sh/ps1 instead."
+    )
+    subparsers.add_parser("install", help=deprecation_warning, add_help=False)
+
+    deprecation_warning_bb2 = "[DEPRECATED] 'chia plotters bladebit2' was integrated to 'chia plotters bladebit'"
+    subparsers.add_parser("bladebit2", help=deprecation_warning_bb2, add_help=False)
+
+    known_args = plotters.parse_known_args(args)
+    maybe_plotter = vars(known_args[0]).get("plotter")
+    if maybe_plotter == "install":
+        print(deprecation_warning)
+        return
+    elif maybe_plotter == "bladebit2":
+        print(deprecation_warning_bb2)
+        return
+
+    args = plotters.parse_args(args)
+
+    if args.plotter is None:
+        plotters.print_help()
+    elif args.plotter == "chiapos":
+        plot_chia(args, chia_root_path)
+    elif args.plotter == "madmax":
+        plot_madmax(args, chia_root_path, root_path)
+    elif args.plotter == "bladebit":
+        plot_bladebit(args, chia_root_path, root_path)
+    elif args.plotter == "version":
+        show_plotters_version(chia_root_path)
+
+
+def get_available_plotters(root_path) -> dict[str, Any]:
+    plotters_root_path: Path = get_plotters_root_path(root_path)
+    plotters: dict[str, Any] = {}
+    chiapos: Optional[dict[str, Any]] = get_chiapos_install_info()
+    bladebit: Optional[dict[str, Any]] = get_bladebit_install_info(plotters_root_path)
+    madmax: Optional[dict[str, Any]] = get_madmax_install_info(plotters_root_path)
+
+    if chiapos is not None:
+        plotters["chiapos"] = chiapos
+    if bladebit is not None:
+        plotters["bladebit"] = bladebit
+    if madmax is not None:
+        plotters["madmax"] = madmax
+
+    return plotters
+
+
+def show_plotters_version(root_path: Path):
+    info = get_available_plotters(root_path)
+    if "chiapos" in info and "version" in info["chiapos"]:
+        print(f"chiapos: {info['chiapos']['version']}")
+    if "bladebit" in info and "version" in info["bladebit"]:
+        if info["bladebit"]["cuda_support"]:
+            print(f"bladebit: {info['bladebit']['version']} (CUDA ready)")
+        else:
+            print(f"bladebit: {info['bladebit']['version']}")
+    if "madmax" in info and "version" in info["madmax"]:
+        print(f"madmax: {info['madmax']['version']}")

chia/plotters/plotters_util.py

@@ -0,0 +1,133 @@
+from __future__ import annotations
+
+import asyncio
+import contextlib
+import json
+import os
+import signal
+import subprocess
+import sys
+from collections.abc import Iterator
+from datetime import datetime
+from pathlib import Path
+from typing import Optional, TextIO
+
+from chia.util.chia_version import chia_short_version
+from chia.util.config import lock_and_load_config
+from chia.util.task_referencer import create_referenced_task
+
+
+@contextlib.contextmanager
+def get_optional_beta_plot_log_file(root_path: Path, plotter: str) -> Iterator[Optional[TextIO]]:
+    beta_log_path: Optional[Path] = None
+    with lock_and_load_config(root_path, "config.yaml") as config:
+        if config.get("beta", {}).get("enabled", False):
+            file_name = f"{plotter}_{datetime.now().strftime('%m_%d_%Y__%H_%M_%S')}.log"
+            beta_log_path = Path(config["beta"]["path"]) / chia_short_version() / "plotting" / file_name
+            beta_log_path.parent.mkdir(parents=True, exist_ok=True)
+    if beta_log_path is not None:
+        with open(beta_log_path, "w") as file:
+            yield file
+    else:
+        yield None
+
+
+# https://kevinmccarthy.org/2016/07/25/streaming-subprocess-stdin-and-stdout-with-asyncio-in-python/
+async def _read_stream(stream, callback):
+    while True:
+        line = await stream.readline()
+        if line:
+            callback(line)
+        else:
+            break
+
+
+def parse_stdout(out, progress):
+    out = out.rstrip()
+    print(out, flush=True)
+    for k, v in progress.items():
+        if k in out:
+            print(f"Progress update: {v}", flush=True)
+
+
+async def run_plotter(root_path, plotter, args, progress_dict):
+    orig_sigint_handler = signal.getsignal(signal.SIGINT)
+    installed_sigint_handler = False
+    process = await asyncio.create_subprocess_exec(
+        *args, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE
+    )
+
+    def sigint_handler(signum, frame):
+        process.terminate()
+
+    # For Windows, we'll install a SIGINT handler to catch Ctrl-C (KeyboardInterrupt isn't raised)
+    if sys.platform in {"win32", "cygwin"}:
+        signal.signal(signal.SIGINT, sigint_handler)
+        installed_sigint_handler = True
+
+    with get_optional_beta_plot_log_file(root_path, plotter) as log_file:
+        if log_file is not None:
+            log_file.write(json.dumps(args) + "\n")
+
+        def process_stdout_line(line_bytes: bytes) -> None:
+            line_str = line_bytes.decode("UTF8")
+            parse_stdout(line_str, progress_dict)
+            if log_file is not None:
+                log_file.write(line_str)
+
+        def process_stderr_line(line_bytes: bytes) -> None:
+            err_str = f"STDERR: {line_bytes.decode('UTF8')}"
+            print(err_str)
+            if log_file is not None:
+                log_file.write(err_str)
+
+        try:
+            await asyncio.wait(
+                [
+                    create_referenced_task(
+                        _read_stream(
+                            process.stdout,
+                            process_stdout_line,
+                        )
+                    ),
+                    create_referenced_task(
+                        _read_stream(
+                            process.stderr,
+                            process_stderr_line,
+                        )
+                    ),
+                ]
+            )
+
+            await process.wait()
+        except Exception as e:
+            print(f"Caught exception while invoking plotter: {e}")
+        finally:
+            # Restore the original SIGINT handler
+            if installed_sigint_handler:
+                signal.signal(signal.SIGINT, orig_sigint_handler)
+
+
+def run_command(args, exc_description, *, check=True, **kwargs) -> subprocess.CompletedProcess:
+    try:
+        proc = subprocess.run(args, check=check, **kwargs)
+    except Exception as e:
+        raise RuntimeError(f"{exc_description} {e}")
+    return proc
+
+
+def reset_loop_policy_for_windows():
+    asyncio.set_event_loop_policy(asyncio.WindowsProactorEventLoopPolicy())
+
+
+def get_venv_bin():
+    venv_dir = os.environ.get("VIRTUAL_ENV", None)
+    if not venv_dir:
+        return None
+
+    venv_path = Path(venv_dir)
+
+    if sys.platform == "win32":
+        return venv_path / "Scripts"
+    else:
+        return venv_path / "bin"