Added cuda and plot compression support for bladebit 3 (#15774)

* Added cuda and plot compression support for bladebit 3

* Update chia/plotters/bladebit.py

Co-authored-by: Arvid Norberg <arvid@libtorrent.org>

* Update chia/plotters/bladebit.py

Co-authored-by: Arvid Norberg <arvid@libtorrent.org>

* Fixed issues

* Fixed lint error

---------

Co-authored-by: Arvid Norberg <arvid@libtorrent.org>

chia/daemon/server.py

@@ -831,28 +831,52 @@ class WebSocketServer:
 
     def _bladebit_plotting_command_args(self, request: Any, ignoreCount: bool) -> List[str]:
         plot_type = request["plot_type"]
-        assert plot_type == "ramplot" or plot_type == "diskplot"
+        if plot_type not in ["ramplot", "diskplot", "cudaplot"]:
+            raise ValueError(f"Unknown plot_type: {plot_type}")
 
         command_args: List[str] = []
 
-        if plot_type == "ramplot":
-            w = request.get("w", False)  # Warm start
-            m = request.get("m", False)  # Disable NUMA
-            no_cpu_affinity = request.get("no_cpu_affinity", False)
-
-            if w is True:
-                command_args.append("--warmstart")
-            if m is True:
-                command_args.append("--nonuma")
-            if no_cpu_affinity is True:
-                command_args.append("--no-cpu-affinity")
-
-            return command_args
-
-        # if plot_type == "diskplot"
+        # Common options among diskplot, ramplot, cudaplot
         w = request.get("w", False)  # Warm start
         m = request.get("m", False)  # Disable NUMA
         no_cpu_affinity = request.get("no_cpu_affinity", False)
+        compress = request.get("compress", None)  # Compression level
+
+        if w is True:
+            command_args.append("--warmstart")
+        if m is True:
+            command_args.append("--nonuma")
+        if no_cpu_affinity is True:
+            command_args.append("--no-cpu-affinity")
+        if compress is not None and str(compress).isdigit():
+            command_args.append("--compress")
+            command_args.append(str(compress))
+
+        # ramplot don't accept any more options
+        if plot_type == "ramplot":
+            return command_args
+
+        # Options only applicable for cudaplot
+        if plot_type == "cudaplot":
+            device_index = request.get("device", None)
+            no_direct_downloads = request.get("no_direct_downloads", False)
+            t1 = request.get("t", None)  # Temp directory
+            t2 = request.get("t2", None)  # Temp2 directory
+
+            if device_index is not None and str(device_index).isdigit():
+                command_args.append("--device")
+                command_args.append(str(device_index))
+            if no_direct_downloads:
+                command_args.append("--no-direct-downloads")
+            if t1 is not None:
+                command_args.append("-t")
+                command_args.append(t1)
+            if t2 is not None:
+                command_args.append("-2")
+                command_args.append(t2)
+            return command_args
+
+        # if plot_type == "diskplot"
         # memo = request["memo"]
         t1 = request["t"]  # Temp directory
         t2 = request.get("t2")  # Temp2 directory
@@ -867,44 +891,37 @@ class WebSocketServer:
         no_t1_direct = request.get("no_t1_direct", False)
         no_t2_direct = request.get("no_t2_direct", False)
 
-        if w is True:
-            command_args.append("--warmstart")
-        if m is True:
-            command_args.append("--nonuma")
-        if no_cpu_affinity is True:
-            command_args.append("--no-cpu-affinity")
-
         command_args.append("-t")
         command_args.append(t1)
-        if t2:
+        if t2 is not None:
             command_args.append("-2")
             command_args.append(t2)
-        if u:
+        if u is not None:
             command_args.append("-u")
             command_args.append(str(u))
-        if cache:
+        if cache is not None:
             command_args.append("--cache")
             command_args.append(str(cache))
-        if f1_threads:
+        if f1_threads is not None:
             command_args.append("--f1-threads")
             command_args.append(str(f1_threads))
-        if fp_threads:
+        if fp_threads is not None:
             command_args.append("--fp-threads")
             command_args.append(str(fp_threads))
-        if c_threads:
+        if c_threads is not None:
             command_args.append("--c-threads")
             command_args.append(str(c_threads))
-        if p2_threads:
+        if p2_threads is not None:
             command_args.append("--p2-threads")
             command_args.append(str(p2_threads))
-        if p3_threads:
+        if p3_threads is not None:
             command_args.append("--p3-threads")
             command_args.append(str(p3_threads))
-        if alternate:
+        if alternate is not None:
             command_args.append("--alternate")
-        if no_t1_direct:
+        if no_t1_direct is not None:
             command_args.append("--no-t1-direct")
-        if no_t2_direct:
+        if no_t2_direct is not None:
             command_args.append("--no-t2-direct")
 
         return command_args
@@ -943,7 +960,7 @@ class WebSocketServer:
             # plotter command must be either
             # 'chia plotters bladebit ramplot' or 'chia plotters bladebit diskplot'
             plot_type = request["plot_type"]
-            assert plot_type == "diskplot" or plot_type == "ramplot"
+            assert plot_type == "diskplot" or plot_type == "ramplot" or plot_type == "cudaplot"
             command_args.append(plot_type)
 
         command_args.extend(self._common_plotting_command_args(request, ignoreCount))

chia/plotters/bladebit.py

@@ -7,7 +7,7 @@ import os
 import sys
 import traceback
 from pathlib import Path
-from typing import Any, Dict, Optional, Tuple
+from typing import Any, Dict, List, Literal, Optional, Tuple, Union
 
 from chia.plotters.plotters_util import get_venv_bin, reset_loop_policy_for_windows, run_command, run_plotter
 from chia.plotting.create_plots import resolve_plot_keys
@@ -52,6 +52,24 @@ def meets_memory_requirement(plotters_root_path: Path) -> Tuple[bool, Optional[s
     return have_enough_memory, warning_string
 
 
+def is_cudaplot_available(plotters_root_path: Path) -> bool:
+    bladebit_executable_path = get_bladebit_executable_path(plotters_root_path)
+    if not bladebit_executable_path.exists():
+        return False
+    try:
+        proc = run_command(
+            [os.fspath(bladebit_executable_path), "cudacheck"],
+            "Failed to call bladebit with cudacheck command",
+            capture_output=True,
+            text=True,
+            check=False,
+        )
+        return proc.returncode == 0
+    except Exception as e:
+        print(f"Failed to determine whether bladebit supports cuda: {e}")
+        return False
+
+
 def get_bladebit_src_path(plotters_root_path: Path) -> Path:
     return plotters_root_path / BLADEBIT_PLOTTER_DIR
 
@@ -60,43 +78,60 @@ def get_bladebit_package_path() -> Path:
     return Path(os.path.dirname(sys.executable)) / "bladebit"
 
 
-def get_bladebit_exec_venv_path() -> Optional[Path]:
+def get_bladebit_exec_path(with_cuda: bool = False) -> str:
+    if with_cuda:
+        return "bladebit_cuda.exe" if sys.platform in ["win32", "cygwin"] else "bladebit_cuda"
+    return "bladebit.exe" if sys.platform in ["win32", "cygwin"] else "bladebit"
+
+
+def get_bladebit_exec_venv_path(with_cuda: bool = False) -> Optional[Path]:
     venv_bin_path = get_venv_bin()
     if not venv_bin_path:
         return None
-    if sys.platform in ["win32", "cygwin"]:
-        return venv_bin_path / "bladebit.exe"
-    else:
-        return venv_bin_path / "bladebit"
+    bladebit_exec = get_bladebit_exec_path(with_cuda)
+    return venv_bin_path / bladebit_exec
 
 
-def get_bladebit_exec_src_path(plotters_root_path: Path) -> Path:
+def get_bladebit_exec_src_path(plotters_root_path: Path, with_cuda: bool = False) -> Path:
     bladebit_src_dir = get_bladebit_src_path(plotters_root_path)
     build_dir = "build/Release" if sys.platform in ["win32", "cygwin"] else "build"
-    bladebit_exec = "bladebit.exe" if sys.platform in ["win32", "cygwin"] else "bladebit"
+    bladebit_exec = get_bladebit_exec_path(with_cuda)
     return bladebit_src_dir / build_dir / bladebit_exec
 
 
-def get_bladebit_exec_package_path() -> Path:
+def get_bladebit_exec_package_path(with_cuda: bool = False) -> Path:
     bladebit_package_dir = get_bladebit_package_path()
-    bladebit_exec = "bladebit.exe" if sys.platform in ["win32", "cygwin"] else "bladebit"
+    bladebit_exec = get_bladebit_exec_path(with_cuda)
     return bladebit_package_dir / bladebit_exec
 
 
 def get_bladebit_executable_path(plotters_root_path: Path) -> Path:
-    bladebit_exec_venv_path = get_bladebit_exec_venv_path()
+    # Search for bladebit executable which supports CUDA at the first priority
+    bladebit_exec_venv_path = get_bladebit_exec_venv_path(with_cuda=True)
     if bladebit_exec_venv_path is not None and bladebit_exec_venv_path.exists():
         return bladebit_exec_venv_path
-    bladebit_exec_src_path = get_bladebit_exec_src_path(plotters_root_path)
+    bladebit_exec_src_path = get_bladebit_exec_src_path(plotters_root_path, with_cuda=True)
     if bladebit_exec_src_path.exists():
         return bladebit_exec_src_path
-    return get_bladebit_exec_package_path()
+    bladebit_exec_package_path = get_bladebit_exec_package_path(with_cuda=True)
+    if bladebit_exec_package_path.exists():
+        return bladebit_exec_package_path
+
+    bladebit_exec_venv_path = get_bladebit_exec_venv_path(with_cuda=False)
+    if bladebit_exec_venv_path is not None and bladebit_exec_venv_path.exists():
+        return bladebit_exec_venv_path
+    bladebit_exec_src_path = get_bladebit_exec_src_path(plotters_root_path, with_cuda=False)
+    if bladebit_exec_src_path.exists():
+        return bladebit_exec_src_path
+    return get_bladebit_exec_package_path(with_cuda=False)
 
 
-def get_bladebit_version(plotters_root_path: Path):
+def get_bladebit_version(
+    plotters_root_path: Path,
+) -> Union[Tuple[Literal[False], str], Tuple[None, str], Tuple[Literal[True], List[str]]]:
     bladebit_executable_path = get_bladebit_executable_path(plotters_root_path)
     if not bladebit_executable_path.exists():
-        # (NotFound, "")
+        # (found=False, "")
         return False, ""
 
     try:
@@ -108,20 +143,22 @@ def get_bladebit_version(plotters_root_path: Path):
             check=False,
         )
         if proc.returncode != 0:
+            # (found=unknown, errMsg)
             return None, proc.stderr.strip()
 
-        # (Found, versionStr)
+        # (found=True, versionStr)
         version_str: str = proc.stdout.strip()
         return True, version_str.split(".")
     except Exception as e:
-        # (Unknown, Exception)
-        return None, e
+        # (found=unknown, errMsg)
+        return None, str(e)
 
 
 def get_bladebit_install_info(plotters_root_path: Path) -> Optional[Dict[str, Any]]:
     info: Dict[str, Any] = {"display_name": "BladeBit Plotter"}
     installed: bool = False
     supported: bool = is_bladebit_supported()
+    cuda_available: bool = is_cudaplot_available(plotters_root_path)
 
     bladebit_executable_path = get_bladebit_executable_path(plotters_root_path)
     if bladebit_executable_path.exists():
@@ -147,9 +184,33 @@ def get_bladebit_install_info(plotters_root_path: Path) -> Optional[Dict[str, An
         if memory_warning is not None:
             info["bladebit_memory_warning"] = memory_warning
 
+    info["cuda_support"] = cuda_available
+
     return info
 
 
+# @TODO Set valid progress logs
+progress_bladebit_cuda = {
+    "Finished F1 sort": 0.01,
+    "Finished forward propagating table 2": 0.06,
+    "Finished forward propagating table 3": 0.12,
+    "Finished forward propagating table 4": 0.2,
+    "Finished forward propagating table 5": 0.28,
+    "Finished forward propagating table 6": 0.36,
+    "Finished forward propagating table 7": 0.42,
+    "Finished prunning table 6": 0.43,
+    "Finished prunning table 5": 0.48,
+    "Finished prunning table 4": 0.51,
+    "Finished prunning table 3": 0.55,
+    "Finished prunning table 2": 0.58,
+    "Finished compressing tables 1 and 2": 0.66,
+    "Finished compressing tables 2 and 3": 0.73,
+    "Finished compressing tables 3 and 4": 0.79,
+    "Finished compressing tables 4 and 5": 0.85,
+    "Finished compressing tables 5 and 6": 0.92,
+    "Finished compressing tables 6 and 7": 0.98,
+}
+
 progress_bladebit_ram = {
     "Finished F1 sort": 0.01,
     "Finished forward propagating table 2": 0.06,
@@ -171,7 +232,6 @@ progress_bladebit_ram = {
     "Finished compressing tables 6 and 7": 0.98,
 }
 
-
 progress_bladebit_disk = {
     # "Running Phase 1": 0.01,
     "Finished f1 generation in ": 0.01,
@@ -231,7 +291,12 @@ def plot_bladebit(args, chia_root_path, root_path):
             args.connect_to_daemon,
         )
     )
-    plot_type = "ramplot" if args.plot_type == "ramplot" else "diskplot"
+    if args.plot_type == "ramplot" or args.plot_type == "diskplot" or args.plot_type == "cudaplot":
+        plot_type = args.plot_type
+    else:
+        plot_type = "diskplot"
+        print("plot_type is automatically set to diskplot")
+
     call_args = [
         os.fspath(bladebit_executable_path),
         "--threads",
@@ -261,6 +326,14 @@ def plot_bladebit(args, chia_root_path, root_path):
         call_args.append("--no-cpu-affinity")
     if args.verbose:
         call_args.append("--verbose")
+    if (
+        "compress" in args
+        and args.compress is not None
+        and str(args.compress).isdigit()
+        and int(version_or_exception[0]) >= 3
+    ):
+        call_args.append("--compress")
+        call_args.append(str(args.compress))
 
     call_args.append(plot_type)
 
@@ -297,11 +370,21 @@ def plot_bladebit(args, chia_root_path, root_path):
         call_args.append("--no-t1-direct")
     if "no_t2_direct" in args and args.no_t2_direct:
         call_args.append("--no-t2-direct")
+    if "device" in args and str(args.device).isdigit():
+        call_args.append("--device")
+        call_args.append(args.device)
+    if "no_direct_downloads" in args and args.no_direct_downloads is not None:
+        call_args.append("--no-direct-downloads")
 
     call_args.append(args.finaldir)
 
     try:
-        progress = progress_bladebit_ram if plot_type == "ramplot" else progress_bladebit_disk
+        if plot_type == "cudaplot":
+            progress = progress_bladebit_cuda
+        elif plot_type == "ramplot":
+            progress = progress_bladebit_ram
+        else:
+            progress = progress_bladebit_disk
         asyncio.run(run_plotter(chia_root_path, args.plotter, call_args, progress))
     except Exception as e:
         print(f"Exception while plotting: {e} {type(e)}")

chia/plotters/plotters.py

@@ -49,6 +49,9 @@ class Options(Enum):
     BLADEBIT_ALTERNATE = 34
     BLADEBIT_NO_T1_DIRECT = 35
     BLADEBIT_NO_T2_DIRECT = 36
+    COMPRESSION = 37
+    BLADEBIT_DEVICE_INDEX = 38
+    BLADEBIT_NO_DIRECT_DOWNLOADS = 39
 
 
 chia_plotter_options = [
@@ -71,6 +74,7 @@ chia_plotter_options = [
     Options.EXCLUDE_FINAL_DIR,
     Options.CONNECT_TO_DAEMON,
     Options.FINAL_DIR,
+    Options.COMPRESSION,
 ]
 
 madmax_plotter_options = [
@@ -91,6 +95,24 @@ madmax_plotter_options = [
     Options.FINAL_DIR,
 ]
 
+bladebit_cuda_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,
+    Options.BLADEBIT_DEVICE_INDEX,
+    Options.BLADEBIT_NO_DIRECT_DOWNLOADS,
+]
+
 bladebit_ram_plotter_options = [
     Options.NUM_THREADS,
     Options.PLOT_COUNT,
@@ -104,6 +126,7 @@ bladebit_ram_plotter_options = [
     Options.VERBOSE,
     Options.CONNECT_TO_DAEMON,
     Options.FINAL_DIR,
+    Options.COMPRESSION,
 ]
 
 bladebit_disk_plotter_options = [
@@ -132,6 +155,7 @@ bladebit_disk_plotter_options = [
     Options.MEMO,
     Options.BLADEBIT_NO_T1_DIRECT,
     Options.BLADEBIT_NO_T2_DIRECT,
+    Options.COMPRESSION,
 ]
 
 
@@ -416,6 +440,27 @@ def build_parser(subparsers, root_path, option_list, name, plotter_desc):
                 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_NO_DIRECT_DOWNLOADS:
+            parser.add_argument(
+                "--no-direct-downloads",
+                action="store_true",
+                help="Don't allocate host tables using pinned buffers",
+                default=False,
+            )
 
 
 def call_plotters(root_path: Path, args):
@@ -444,6 +489,7 @@ def call_plotters(root_path: Path, args):
 
     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")
 
@@ -489,12 +535,8 @@ def get_available_plotters(root_path) -> Dict[str, Any]:
 
     if chiapos is not None:
         plotters["chiapos"] = chiapos
-    if bladebit and bladebit.get("version") is not None:
-        bladebit_major_version = bladebit["version"].split(".")[0]
-        if bladebit_major_version == "2":
-            plotters["bladebit2"] = bladebit
-        else:
-            plotters["bladebit"] = bladebit
+    if bladebit is not None:
+        plotters["bladebit"] = bladebit
     if madmax is not None:
         plotters["madmax"] = madmax
 
@@ -506,8 +548,9 @@ def show_plotters_version(root_path: 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"]:
-        print(f"bladebit: {info['bladebit']['version']}")
-    if "bladebit2" in info and "version" in info["bladebit2"]:
-        print(f"bladebit: {info['bladebit2']['version']}")
+        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']}")

tests/core/daemon/test_daemon.py

@@ -1222,6 +1222,12 @@ async def test_bad_json(daemon_connection_and_temp_keychain: Tuple[aiohttp.Clien
         response={
             "success": True,
             "plotters": {
+                "bladebit": {
+                    "can_install": True,
+                    "cuda_support": False,
+                    "display_name": "BladeBit Plotter",
+                    "installed": False,
+                },
                 "chiapos": {"display_name": "Chia Proof of Space", "installed": True, "version": chiapos_version},
                 "madmax": {"can_install": True, "display_name": "madMAx Plotter", "installed": False},
             },