add support for creating v2 plots in block_tools

Author: arvidn

chia/_tests/conftest.py

@@ -1288,6 +1288,7 @@ async def farmer_harvester_2_simulators_zero_bits_plot_filter(
                 num_og_plots=0,
                 num_pool_plots=0,
                 num_non_keychain_plots=0,
+                num_v2_plots=0,
                 config_overrides=config_overrides,
             )
             for _ in range(2)

chia/_tests/core/custom_types/test_proof_of_space.py

@@ -112,25 +112,23 @@ def b32(key: str) -> bytes32:
         id="v2 plot size 0",
         pos_challenge=bytes32(b"1" * 32),
         plot_size=PlotSize.make_v2(0),
+        pool_contract_puzzle_hash=bytes32(b"1" * 32),
         plot_public_key=G1Element(),
-        pool_public_key=G1Element(),
         expected_error="Plot size (0) is lower than the minimum (28)",
     ),
     ProofOfSpaceCase(
         id="v2 plot size 34",
         pos_challenge=bytes32(b"1" * 32),
         plot_size=PlotSize.make_v2(34),
+        pool_contract_puzzle_hash=bytes32(b"1" * 32),
         plot_public_key=G1Element(),
-        pool_public_key=G1Element(),
         expected_error="Plot size (34) is higher than the maximum (32)",
     ),
     ProofOfSpaceCase(
         id="Not passing the plot filter v2",
-        pos_challenge=b32("3d29ea79d19b3f7e99ebf764ae53697cbe143603909873946af6ab1ece606861"),
+        pos_challenge=b32("4cfaacbd2782db64d07cf490ca938534adb07dfbd2f92b0e479e2e5b196178db"),
         plot_size=PlotSize.make_v2(32),
-        pool_public_key=g1(
-            "b6449c2c68df97c19e884427e42ee7350982d4020571ead08732615ff39bd216bfd630b6460784982bec98b49fea79d0"
-        ),
+        pool_contract_puzzle_hash=bytes32(b"1" * 32),
         plot_public_key=g1(
             "879526b4e7b616cfd64984d8ad140d0798b048392a6f11e2faf09054ef467ea44dc0dab5e5edb2afdfa850c5c8b629cc"
         ),
@@ -160,16 +158,15 @@ def test_verify_and_get_quality_string(caplog: pytest.LogCaptureFixture, case: P
 @datacases(
     ProofOfSpaceCase(
         id="v2 plot are not implemented",
-        plot_size=PlotSize.make_v2(30),
-        pos_challenge=b32("47deb938e145d25d7b3b3c85ca9e3972b76c01aeeb78a02fe5d3b040d282317e"),
+        plot_size=PlotSize.make_v2(28),
+        pos_challenge=b32("6a85c4c5f21e5728c6668b01c0758e33bb7f3ae20d38703d4863ad151f983d9c"),
         plot_public_key=g1(
             "afa3aaf09c03885154be49216ee7fb2e4581b9c4a4d7e9cc402e27280bf0cfdbdf1b9ba674e301fd1d1450234b3b1868"
         ),
-        pool_public_key=g1(
-            "b6449c2c68df97c19e884427e42ee7350982d4020571ead08732615ff39bd216bfd630b6460784982bec98b49fea79d0"
-        ),
-        expected_error="NotImplementedError",
+        pool_contract_puzzle_hash=bytes32(b"1" * 32),
+        expected_error="Did not pass the plot filter",
     ),
+    # TODO: todo_v2_plots add test case that passes the plot filter
 )
 def test_verify_and_get_quality_string_v2(caplog: pytest.LogCaptureFixture, case: ProofOfSpaceCase) -> None:
     pos = make_pos(

chia/_tests/core/test_farmer_harvester_rpc.py

@@ -87,7 +87,7 @@ async def non_zero_plots() -> bool:
         res = await harvester_rpc_client.get_plots()
         nonlocal harvester_plots
         harvester_plots = res["plots"]
-        return len(harvester_plots) > 0
+        return len(harvester_plots) == 340
 
     await time_out_assert(10, non_zero_plots)
 
@@ -293,7 +293,7 @@ async def wait_for_plot_sync() -> bool:
     await farmer_api.farmer.update_pool_state()
 
     pool_plot_count: int = (await farmer_rpc_client.get_pool_state())["pool_state"][0]["plot_count"]
-    assert pool_plot_count == 5
+    assert pool_plot_count == 325
 
     # TODO: Maybe improve this to not remove from Receiver directly but instead from the harvester and then wait for
     #       plot sync event.
@@ -333,8 +333,8 @@ def test_plot_matches_filter(filter_item: FilterItem, match: bool) -> None:
 @pytest.mark.parametrize(
     "endpoint, filtering, sort_key, reverse, expected_plot_count",
     [
-        (FarmerRpcClient.get_harvester_plots_valid, [], "filename", False, 20),
-        (FarmerRpcClient.get_harvester_plots_valid, [], "size", True, 20),
+        (FarmerRpcClient.get_harvester_plots_valid, [], "filename", False, 340),
+        (FarmerRpcClient.get_harvester_plots_valid, [], "size", True, 340),
         (
             FarmerRpcClient.get_harvester_plots_valid,
             [FilterItem("pool_contract_puzzle_hash", None)],
@@ -349,6 +349,20 @@ def test_plot_matches_filter(filter_item: FilterItem, match: bool) -> None:
             False,
             4,
         ),
+        (
+            FarmerRpcClient.get_harvester_plots_valid,
+            [FilterItem("strength", "2")],
+            "size",
+            True,
+            320,
+        ),
+        (
+            FarmerRpcClient.get_harvester_plots_valid,
+            [FilterItem("strength", "0")],
+            "size",
+            True,
+            20,
+        ),
         (FarmerRpcClient.get_harvester_plots_invalid, [], None, True, 13),
         (FarmerRpcClient.get_harvester_plots_invalid, ["invalid_0"], None, False, 6),
         (FarmerRpcClient.get_harvester_plots_invalid, ["inval", "lid_1"], None, False, 2),

chia/_tests/farmer_harvester/test_filter_prefix_bits.py

@@ -84,7 +84,7 @@ async def have_connections() -> bool:
     await harvester_rpc_cl.await_closed()
 
 
-@pytest.mark.parametrize(argnames=["peak_height", "eligible_plots"], argvalues=[(5495999, 0), (5496000, 1)])
+@pytest.mark.parametrize(argnames=["peak_height", "eligible_plots"], argvalues=[(5495999, 9), (5496000, 10)])
 @pytest.mark.anyio
 async def test_filter_prefix_bits_with_farmer_harvester(
     farmer_harvester_with_filter_size_9: tuple[HarvesterService, FarmerAPI],

chia/plotting/create_plots.py

@@ -5,15 +5,17 @@
 from pathlib import Path
 from typing import Optional
 
-from chia_rs import AugSchemeMPL, G1Element, PrivateKey
+from chia_rs import AugSchemeMPL, G1Element, PrivateKey, create_v2_plot
 from chia_rs.sized_bytes import bytes32
+from chia_rs.sized_ints import uint8
 from chiapos import DiskPlotter
 
 from chia.daemon.keychain_proxy import KeychainProxy, connect_to_keychain_and_validate, wrap_local_keychain
 from chia.plotting.util import Params, stream_plot_info_ph, stream_plot_info_pk
 from chia.types.blockchain_format.proof_of_space import (
     calculate_plot_id_ph,
     calculate_plot_id_pk,
+    calculate_plot_id_v2,
     generate_plot_public_key,
 )
 from chia.util.bech32m import decode_puzzle_hash
@@ -275,3 +277,63 @@ async def create_plots(
         log.info(created_path.name)
 
     return created_plots, existing_plots
+
+
+async def create_v2_plots(
+    final_dir: Path,
+    *,
+    pool_ph: bytes32,
+    farmer_pk: G1Element,
+    size: int = 28,
+    strength: int = 2,
+    num: int = 1,
+    use_datetime: bool = True,
+    test_private_keys: Optional[list[PrivateKey]] = None,
+) -> tuple[dict[bytes32, Path], dict[bytes32, Path]]:
+    log.info(
+        f"Creating {num} plots of size {size}, pool contract address:  "
+        f"{pool_ph} farmer public key: {bytes(farmer_pk).hex()}"
+    )
+
+    final_dir.mkdir(parents=True, exist_ok=True)
+
+    created_plots: dict[bytes32, Path] = {}
+    existing_plots: dict[bytes32, Path] = {}
+    for i in range(num):
+        # Generate a random master secret key
+        if test_private_keys is not None:
+            assert len(test_private_keys) == num
+            sk: PrivateKey = test_private_keys[i]
+        else:
+            sk = AugSchemeMPL.key_gen(bytes32.secret())
+
+        # The plot public key is the combination of the harvester and farmer keys
+        # New plots will also include a taproot of the keys, for extensibility
+        plot_public_key = generate_plot_public_key(master_sk_to_local_sk(sk).get_g1(), farmer_pk, include_taproot=True)
+
+        # The plot id is based on the harvester, farmer, pool contract puzzle
+        # hash and strength
+        plot_id = calculate_plot_id_v2(pool_ph, plot_public_key, uint8(strength))
+        plot_memo = stream_plot_info_ph(pool_ph, farmer_pk, sk)
+
+        dt_string = datetime.now().strftime("%Y-%m-%d-%H-%M")
+
+        if use_datetime:
+            filename: str = f"plot-k{size}-{dt_string}-{plot_id}.plot2"
+        else:
+            filename = f"plot-k{size}-{plot_id}.plot2"
+        full_path: Path = final_dir / filename
+
+        if not full_path.exists():
+            log.info(f"Starting plot {i + 1}/{num}")
+            create_v2_plot(str(full_path.absolute()), size, strength, plot_id, plot_memo)
+            created_plots[plot_id] = full_path
+        else:
+            log.info(f"Plot {filename} already exists")
+            existing_plots[plot_id] = full_path
+
+    log.info(f"Created a total of {len(created_plots)} new plots")
+    for created_path in created_plots.values():
+        log.info(created_path.name)
+
+    return created_plots, existing_plots

chia/plotting/manager.py

@@ -278,8 +278,8 @@ def _refresh_task(self, sleep_interval_ms: int):
                     f"total_result.removed {len(total_result.removed)}, "
                     f"total_duration {total_result.duration:.2f} seconds"
                 )
-            except Exception as e:
-                log.error(f"_refresh_callback raised: {e} with the traceback: {traceback.format_exc()}")
+            except Exception:
+                log.exception("_refresh_callback raised")
                 self.reset()
 
     def refresh_batch(self, plot_paths: list[Path], plot_directories: set[Path]) -> PlotRefreshResult:

chia/simulator/block_tools.py

@@ -63,7 +63,7 @@
 from chia.consensus.vdf_info_computation import get_signage_point_vdf_info
 from chia.daemon.keychain_proxy import KeychainProxy, connect_to_keychain_and_validate, wrap_local_keychain
 from chia.full_node.bundle_tools import simple_solution_generator, simple_solution_generator_backrefs
-from chia.plotting.create_plots import PlotKeys, create_plots
+from chia.plotting.create_plots import PlotKeys, create_plots, create_v2_plots
 from chia.plotting.manager import PlotManager
 from chia.plotting.prover import PlotVersion, QualityProtocol, V1Prover, V2Prover, V2Quality
 from chia.plotting.util import (
@@ -319,6 +319,7 @@ def __init__(
         self.temp_dir.mkdir(parents=True, exist_ok=True)
         self.expected_plots: dict[bytes32, Path] = {}
         self.created_plots: int = 0
+        self.created_plots2: int = 0
         self.total_result = PlotRefreshResult()
 
         def test_callback(event: PlotRefreshEvents, update_result: PlotRefreshResult) -> None:
@@ -338,6 +339,29 @@ def test_callback(event: PlotRefreshEvents, update_result: PlotRefreshResult) ->
                 assert self.total_result.processed == update_result.processed
                 assert self.total_result.duration == update_result.duration
                 assert update_result.remaining == 0
+
+                expected_plots: set[str] = set()
+                found_plots: set[str] = set()
+                if len(self.plot_manager.plots) != len(self.expected_plots):
+                    for pid, filename in self.expected_plots.items():
+                        expected_plots.add(filename.name)
+                    for filename, _ in self.plot_manager.plots.items():
+                        found_plots.add(filename.name)
+                    print(f"directory: {self.plot_dir}")
+                    print(f"expected: {len(expected_plots)}")
+                    for f in expected_plots:
+                        print(f)
+                    print(f"plot manager: {len(found_plots)}")
+                    for f in found_plots:
+                        print(f)
+                    diff = found_plots.difference(expected_plots)
+                    print(f"found unexpected: {len(diff)}")
+                    for f in diff:
+                        print(f)
+                    diff = expected_plots.difference(found_plots)
+                    print(f"not found: {len(diff)}")
+                    for f in diff:
+                        print(f)
                 assert len(self.plot_manager.plots) == len(self.expected_plots)
 
         self.plot_manager: PlotManager = PlotManager(
@@ -498,11 +522,25 @@ async def setup_plots(
         num_og_plots: int = 15,
         num_pool_plots: int = 5,
         num_non_keychain_plots: int = 3,
+        num_v2_plots: int = 320,
         plot_size: int = 20,
         bitfield: bool = True,
     ) -> bool:
+        # we have 20 v1 plots, each about 16.6 MB
+        # in order to balance that out with v2 plots, we need approximately the
+        # same size on disk. A v2 k-18 plot is expected to be about 1 MB, so we
+        # need about 16 times more v2 plots, i.e. 320.
+        # Additionally, the current reference plots are quite a bit larger than
+        # the expected final plots. So until we have the optimized plot format
+        # completed, the actual disk space for the v2 plots will be larger.
+        print(
+            f"setup_plots({num_og_plots}, {num_pool_plots}, "
+            f"{num_non_keychain_plots}, {num_v2_plots}) "
+            f"plot-dir: {self.plot_dir}"
+        )
         self.add_plot_directory(self.plot_dir)
         assert self.created_plots == 0
+        assert self.created_plots2 == 0
         existing_plots: bool = True
         # OG Plots
         for i in range(num_og_plots):
@@ -525,6 +563,11 @@ async def setup_plots(
             )
             if plot.new_plot:
                 existing_plots = False
+        # v2 plots
+        for i in range(num_v2_plots):
+            plot = await self.new_plot2(plot_size=18)
+            if plot.new_plot:
+                existing_plots = False
         await self.refresh_plots()
         assert len(self.plot_manager.plots) == len(self.expected_plots)
         return existing_plots
@@ -604,6 +647,48 @@ async def new_plot(
             shutil.rmtree(self.temp_dir, ignore_errors=True)
             sys.exit(1)
 
+    async def new_plot2(
+        self,
+        path: Optional[Path] = None,
+        exclude_plots: bool = False,
+        plot_size: int = 18,
+    ) -> BlockToolsNewPlotResult:
+        final_dir = self.plot_dir
+        if path is not None:
+            final_dir = path
+            final_dir.mkdir(parents=True, exist_ok=True)
+
+        # No datetime in the filename, to get deterministic filenames and not re-plot
+        created, existed = await create_v2_plots(
+            final_dir=Path(final_dir),
+            size=plot_size,
+            pool_ph=self.pool_ph,
+            farmer_pk=self.farmer_pk,
+            use_datetime=False,
+            test_private_keys=[AugSchemeMPL.key_gen(std_hash(self.created_plots2.to_bytes(2, "big")))],
+        )
+        self.created_plots2 += 1
+
+        plot_id_new: Optional[bytes32] = None
+        path_new: Optional[Path] = None
+        new_plot: bool = True
+
+        if len(created):
+            assert len(existed) == 0
+            plot_id_new, path_new = next(iter(created.items()))
+
+        if len(existed):
+            assert len(created) == 0
+            plot_id_new, path_new = next(iter(existed.items()))
+            new_plot = False
+        assert plot_id_new is not None
+        assert path_new is not None
+
+        if not exclude_plots:
+            self.expected_plots[plot_id_new] = path_new
+
+        return BlockToolsNewPlotResult(plot_id_new, new_plot)
+
     async def refresh_plots(self) -> None:
         self.plot_manager.refresh_parameter = replace(
             self.plot_manager.refresh_parameter, batch_size=uint32(4 if len(self.expected_plots) % 3 == 0 else 3)
@@ -2080,6 +2165,7 @@ async def create_block_tools_async(
     num_og_plots: int = 15,
     num_pool_plots: int = 5,
     num_non_keychain_plots: int = 3,
+    num_v2_plots: int = 320,
 ) -> BlockTools:
     global create_block_tools_async_count
     create_block_tools_async_count += 1
@@ -2090,6 +2176,7 @@ async def create_block_tools_async(
         num_og_plots=num_og_plots,
         num_pool_plots=num_pool_plots,
         num_non_keychain_plots=num_non_keychain_plots,
+        num_v2_plots=num_v2_plots,
     )
 
     return bt

chia/simulator/start_simulator.py

@@ -117,7 +117,9 @@ async def async_main(
         plot_dir=plot_dir,
     )
     await bt.setup_keys(fingerprint=fingerprint, reward_ph=farming_puzzle_hash)
-    await bt.setup_plots(num_og_plots=PLOTS, num_pool_plots=0, num_non_keychain_plots=0, plot_size=PLOT_SIZE)
+    await bt.setup_plots(
+        num_og_plots=PLOTS, num_pool_plots=0, num_non_keychain_plots=0, plot_size=PLOT_SIZE, num_v2_plots=PLOTS
+    )
     # Everything after this is not simulator specific, excluding the if test_mode.
     initialize_logging(
         service_name=SERVICE_NAME,

chia/types/blockchain_format/proof_of_space.py

@@ -59,6 +59,14 @@ def make_pos(
 
 def get_plot_id(pos: ProofOfSpace) -> bytes32:
     assert pos.pool_public_key is None or pos.pool_contract_puzzle_hash is None
+
+    plot_size = pos.size()
+    if plot_size.size_v2 is not None:
+        assert pos.pool_contract_puzzle_hash
+        # v2 plots have a fixed k-size, so we store the *strength* in this field
+        # instead
+        return calculate_plot_id_v2(pos.pool_contract_puzzle_hash, pos.plot_public_key, uint8(plot_size.size_v2))
+
     if pos.pool_public_key is None:
         assert pos.pool_contract_puzzle_hash is not None
         return calculate_plot_id_ph(pos.pool_contract_puzzle_hash, pos.plot_public_key)
@@ -203,6 +211,10 @@ def calculate_plot_id_ph(
     return std_hash(bytes(pool_contract_puzzle_hash) + bytes(plot_public_key))
 
 
+def calculate_plot_id_v2(pool_contract_puzzle_hash: bytes32, plot_public_key: G1Element, strength: uint8) -> bytes32:
+    return std_hash(bytes(pool_contract_puzzle_hash) + bytes(plot_public_key) + strength.stream_to_bytes())
+
+
 def generate_taproot_sk(local_pk: G1Element, farmer_pk: G1Element) -> PrivateKey:
     taproot_message: bytes = bytes(local_pk + farmer_pk) + bytes(local_pk) + bytes(farmer_pk)
     taproot_hash: bytes32 = std_hash(taproot_message)