from __future__ import annotations import logging import threading import time import traceback from collections.abc import Callable from concurrent.futures.thread import ThreadPoolExecutor from pathlib import Path from typing import Any from chia_rs import ConsensusConstants, G1Element from chiapos import decompressor_context_queue from chia.consensus.pos_quality import UI_ACTUAL_SPACE_CONSTANT_FACTOR, _expected_plot_size from chia.plotting.cache import Cache, CacheEntry from chia.plotting.prover import get_prover_from_file from chia.plotting.util import ( HarvestingMode, PlotInfo, PlotRefreshEvents, PlotRefreshResult, PlotsRefreshParameter, get_plot_filenames, ) from chia.util.batches import to_batches log = logging.getLogger(__name__) class PlotManager: plots: dict[Path, PlotInfo] plot_filename_paths: dict[str, tuple[str, set[str]]] plot_filename_paths_lock: threading.Lock failed_to_open_filenames: dict[Path, int] no_key_filenames: set[Path] farmer_public_keys: list[G1Element] pool_public_keys: list[G1Element] cache: Cache match_str: str | None open_no_key_filenames: bool last_refresh_time: float refresh_parameter: PlotsRefreshParameter log: Any _lock: threading.Lock _refresh_thread: threading.Thread | None _refreshing_enabled: bool _refresh_callback: Callable _initial: bool max_compression_level_allowed: int context_count: int constants: ConsensusConstants def __init__( self, root_path: Path, refresh_callback: Callable, constants: ConsensusConstants, match_str: str | None = None, open_no_key_filenames: bool = False, refresh_parameter: PlotsRefreshParameter = PlotsRefreshParameter(), ): self.root_path = root_path self.plots = {} self.plot_filename_paths = {} self.plot_filename_paths_lock = threading.Lock() self.failed_to_open_filenames = {} self.no_key_filenames = set() self.farmer_public_keys = [] self.pool_public_keys = [] # Since `compression_level` property was added to Cache structure, # previous cache file formats needs to be reset # When user downgrades harvester, it looks 'plot_manager.dat` while # latest harvester reads/writes 'plot_manager_v2.dat` self.cache = Cache(self.root_path.resolve() / "cache" / "plot_manager_v2.dat") self.match_str = match_str self.open_no_key_filenames = open_no_key_filenames self.last_refresh_time = 0 self.refresh_parameter = refresh_parameter self.log = logging.getLogger(__name__) self._lock = threading.Lock() self._refresh_thread = None self._refreshing_enabled = False self._refresh_callback = refresh_callback self._initial = True self.max_compression_level_allowed = 0 self.context_count = 0 self.constants = constants def __enter__(self): self._lock.acquire() def __exit__(self, exc_type, exc_value, exc_traceback): self._lock.release() def configure_decompressor( self, context_count: int, thread_count: int, disable_cpu_affinity: bool, max_compression_level_allowed: int, use_gpu_harvesting: bool, gpu_index: int, enforce_gpu_index: bool, decompressor_timeout: int, ) -> HarvestingMode: if max_compression_level_allowed > 7: log.error( "Currently only compression levels up to 7 are allowed, " f"while {max_compression_level_allowed} was specified." "Setting max_compression_level_allowed to 7." ) max_compression_level_allowed = 7 is_using_gpu = decompressor_context_queue.init( context_count, thread_count, disable_cpu_affinity, max_compression_level_allowed, use_gpu_harvesting, gpu_index, enforce_gpu_index, decompressor_timeout, ) if not is_using_gpu and use_gpu_harvesting: log.error( "GPU harvesting failed initialization. " f"Falling back to CPU harvesting: {context_count} decompressors count, {thread_count} threads." ) self.max_compression_level_allowed = max_compression_level_allowed self.context_count = context_count return HarvestingMode.GPU if is_using_gpu else HarvestingMode.CPU def reset(self) -> None: with self: self.last_refresh_time = time.time() self.plots.clear() self.plot_filename_paths.clear() self.failed_to_open_filenames.clear() self.no_key_filenames.clear() self._initial = True def set_refresh_callback(self, callback: Callable): self._refresh_callback = callback def set_public_keys(self, farmer_public_keys: list[G1Element], pool_public_keys: list[G1Element]): self.farmer_public_keys = farmer_public_keys self.pool_public_keys = pool_public_keys def initial_refresh(self) -> bool: return self._initial def public_keys_available(self) -> bool: return len(self.farmer_public_keys) > 0 and len(self.pool_public_keys) > 0 def plot_count(self) -> int: with self: return len(self.plots) def get_duplicates(self) -> list[Path]: result = [] for plot_filename, paths_entry in self.plot_filename_paths.items(): _, duplicated_paths = paths_entry for path in duplicated_paths: result.append(Path(path) / plot_filename) return result def needs_refresh(self) -> bool: return time.time() - self.last_refresh_time > float(self.refresh_parameter.interval_seconds) def start_refreshing(self, sleep_interval_ms: int = 1000): self._refreshing_enabled = True if self._refresh_thread is None or not self._refresh_thread.is_alive(): self.cache.load() self._refresh_thread = threading.Thread(target=self._refresh_task, args=(sleep_interval_ms,)) self._refresh_thread.start() def stop_refreshing(self) -> None: self._refreshing_enabled = False if self._refresh_thread is not None and self._refresh_thread.is_alive(): self._refresh_thread.join() self._refresh_thread = None def trigger_refresh(self) -> None: log.debug("trigger_refresh") self.last_refresh_time = 0 def _refresh_task(self, sleep_interval_ms: int): while self._refreshing_enabled: try: while not self.needs_refresh() and self._refreshing_enabled: time.sleep(sleep_interval_ms / 1000.0) if not self._refreshing_enabled: return plot_filenames: dict[Path, list[Path]] = get_plot_filenames(self.root_path) plot_directories: set[Path] = set(plot_filenames.keys()) plot_paths: set[Path] = set() for paths in plot_filenames.values(): plot_paths.update(paths) total_result: PlotRefreshResult = PlotRefreshResult() total_size = len(plot_paths) self._refresh_callback(PlotRefreshEvents.started, PlotRefreshResult(remaining=total_size)) # First drop all plots we have in plot_filename_paths but not longer in the filesystem or set in config for path in list(self.failed_to_open_filenames.keys()): if path not in plot_paths: del self.failed_to_open_filenames[path] for path in self.no_key_filenames.copy(): if path not in plot_paths: self.no_key_filenames.remove(path) filenames_to_remove: list[str] = [] for plot_filename, paths_entry in self.plot_filename_paths.items(): loaded_path, duplicated_paths = paths_entry loaded_plot = Path(loaded_path) / Path(plot_filename) if loaded_plot not in plot_paths: filenames_to_remove.append(plot_filename) with self: if loaded_plot in self.plots: del self.plots[loaded_plot] total_result.removed.append(loaded_plot) # No need to check the duplicates here since we drop the whole entry continue paths_to_remove: list[str] = [] for path_str in duplicated_paths: loaded_plot = Path(path_str) / Path(plot_filename) if loaded_plot not in plot_paths: paths_to_remove.append(path_str) for path_str in paths_to_remove: duplicated_paths.remove(path_str) for filename in filenames_to_remove: del self.plot_filename_paths[filename] for batch in to_batches(sorted(list(plot_paths)), self.refresh_parameter.batch_size): batch_result: PlotRefreshResult = self.refresh_batch(batch.entries, plot_directories) if not self._refreshing_enabled: self.log.debug("refresh_plots: Aborted") break # Set the remaining files since `refresh_batch()` doesn't know them but we want to report it batch_result.remaining = batch.remaining total_result.loaded += batch_result.loaded total_result.processed += batch_result.processed total_result.duration += batch_result.duration self._refresh_callback(PlotRefreshEvents.batch_processed, batch_result) if batch.remaining == 0: break batch_sleep = self.refresh_parameter.batch_sleep_milliseconds self.log.debug(f"refresh_plots: Sleep {batch_sleep} milliseconds") time.sleep(float(batch_sleep) / 1000.0) if self._refreshing_enabled: self._refresh_callback(PlotRefreshEvents.done, total_result) # Reset the initial refresh indication self._initial = False # Cleanup unused cache self.log.debug(f"_refresh_task: cached entries before cleanup: {len(self.cache)}") remove_paths: list[Path] = [] for path, cache_entry in self.cache.items(): if cache_entry.expired(Cache.expiry_seconds) and path not in self.plots: remove_paths.append(path) elif path in self.plots: cache_entry.bump_last_use() self.cache.remove(remove_paths) self.log.debug(f"_refresh_task: cached entries removed: {len(remove_paths)}") if self.cache.changed(): self.cache.save() self.last_refresh_time = time.time() self.log.debug( f"_refresh_task: total_result.loaded {len(total_result.loaded)}, " f"total_result.removed {len(total_result.removed)}, " f"total_duration {total_result.duration:.2f} seconds" ) except Exception: log.exception("_refresh_callback raised") self.reset() def refresh_batch(self, plot_paths: list[Path], plot_directories: set[Path]) -> PlotRefreshResult: start_time: float = time.time() result: PlotRefreshResult = PlotRefreshResult(processed=len(plot_paths)) counter_lock = threading.Lock() log.debug(f"refresh_batch: {len(plot_paths)} files in directories {plot_directories}") if self.match_str is not None: log.info(f'Only loading plots that contain "{self.match_str}" in the file or directory name') def process_file(file_path: Path) -> PlotInfo | None: if not self._refreshing_enabled: return None filename_str = str(file_path) if self.match_str is not None and self.match_str not in filename_str: return None if ( file_path in self.failed_to_open_filenames and (time.time() - self.failed_to_open_filenames[file_path]) < self.refresh_parameter.retry_invalid_seconds ): # Try once every `refresh_parameter.retry_invalid_seconds` seconds to open the file return None if file_path in self.plots: return self.plots[file_path] entry: tuple[str, set[str]] | None = self.plot_filename_paths.get(file_path.name) if entry is not None: _loaded_parent, duplicates = entry if str(file_path.parent) in duplicates: log.debug(f"Skip duplicated plot {file_path!s}") return None try: if not file_path.exists(): return None stat_info = file_path.stat() cache_entry = self.cache.get(file_path) cache_hit = cache_entry is not None if not cache_hit: prover = get_prover_from_file(str(file_path)) log.debug(f"process_file {file_path!s}") expected_size = ( _expected_plot_size(prover.get_param(), self.constants) * UI_ACTUAL_SPACE_CONSTANT_FACTOR ) # TODO: consider checking if the file was just written to (which would mean that the file is still # being copied). A segfault might happen in this edge case. param = prover.get_param() level = prover.get_compression_level() if ( level == 0 and stat_info.st_size < 0.98 * expected_size and param.size_v1 is not None and param.size_v1 >= 30 ): log.warning( f"Not farming plot {file_path}. " f"Size is {stat_info.st_size / (1024**3)} GiB, " f"but expected at least: {expected_size / (1024**3)} GiB. " "We assume the file is being copied." ) return None cache_entry = CacheEntry.from_prover(prover) self.cache.update(file_path, cache_entry) assert cache_entry is not None level = cache_entry.prover.get_compression_level() if level > self.max_compression_level_allowed: log.warning( f"Not farming plot {file_path}. Plot compression level: {level}, " f"max compression level allowed: {self.max_compression_level_allowed}." ) return None if level > 0 and self.context_count == 0: log.warning( f"Not farming compressed plot {file_path}. Plot compression level: {level}, " f"because parallel_decompressor_count is set to 0 in config.yaml. Use a non-zero value" " to start harvesting compressed plots." ) return None # Only use plots that correct keys associated with them if cache_entry.farmer_public_key not in self.farmer_public_keys: log.warning(f"Plot {file_path} has a farmer public key that is not in the farmer's pk list.") self.no_key_filenames.add(file_path) if not self.open_no_key_filenames: return None if cache_entry.pool_public_key is not None and cache_entry.pool_public_key not in self.pool_public_keys: log.warning(f"Plot {file_path} has a pool public key that is not in the farmer's pool pk list.") self.no_key_filenames.add(file_path) if not self.open_no_key_filenames: return None # If a plot is in `no_key_filenames` the keys were missing in earlier refresh cycles. We can remove # the current plot from that list if its in there since we passed the key checks above. if file_path in self.no_key_filenames: self.no_key_filenames.remove(file_path) with self.plot_filename_paths_lock: paths: tuple[str, set[str]] | None = self.plot_filename_paths.get(file_path.name) if paths is None: paths = (str(Path(cache_entry.prover.get_filename()).parent), set()) self.plot_filename_paths[file_path.name] = paths else: paths[1].add(str(Path(cache_entry.prover.get_filename()).parent)) log.warning(f"Have multiple copies of the plot {file_path.name} in {[paths[0], *paths[1]]}.") return None new_plot_info: PlotInfo = PlotInfo( cache_entry.prover, cache_entry.pool_public_key, cache_entry.pool_contract_puzzle_hash, cache_entry.plot_public_key, stat_info.st_size, stat_info.st_mtime, ) cache_entry.bump_last_use() with counter_lock: result.loaded.append(new_plot_info) if file_path in self.failed_to_open_filenames: del self.failed_to_open_filenames[file_path] except Exception as e: tb = traceback.format_exc() log.error(f"Failed to open file {file_path}. {e} {tb}") self.failed_to_open_filenames[file_path] = int(time.time()) return None log.debug(f"Found plot {file_path} of size {new_plot_info.prover.get_param()}, cache_hit: {cache_hit}") return new_plot_info with self, ThreadPoolExecutor(thread_name_prefix="plot-process-") as executor: plots_refreshed: dict[Path, PlotInfo] = {} for new_plot in executor.map(process_file, plot_paths): if new_plot is not None: plots_refreshed[Path(new_plot.prover.get_filename())] = new_plot self.plots.update(plots_refreshed) result.duration = time.time() - start_time self.log.debug( f"refresh_batch: loaded {len(result.loaded)}, " f"removed {len(result.removed)}, processed {result.processed}, " f"remaining {result.remaining}, batch_size {self.refresh_parameter.batch_size}, " f"duration: {result.duration:.2f} seconds" ) return result