Added defaults to get_harvester_config API (#15820)

This commit is contained in:
Izumi Hoshino
2023-07-20 13:25:16 -05:00
committed by GitHub
parent bc00bc71d3
commit ec5ef087fe
2 changed files with 37 additions and 15 deletions
+16 -7
View File
@@ -15,6 +15,13 @@ from chia.consensus.constants import ConsensusConstants
from chia.plot_sync.sender import Sender
from chia.plotting.manager import PlotManager
from chia.plotting.util import (
DEFAULT_DECOMPRESSOR_THREAD_COUNT,
DEFAULT_DISABLE_CPU_AFFINITY,
DEFAULT_ENFORCE_GPU_INDEX,
DEFAULT_GPU_INDEX,
DEFAULT_MAX_COMPRESSION_LEVEL_ALLOWED,
DEFAULT_PARALLEL_DECOMPRESSOR_COUNT,
DEFAULT_USE_GPU_HARVESTING,
HarvestingMode,
PlotRefreshEvents,
PlotRefreshResult,
@@ -85,15 +92,17 @@ class Harvester:
self.state_changed_callback: Optional[StateChangedProtocol] = None
self.parallel_read: bool = config.get("parallel_read", True)
context_count = config.get("parallel_decompressor_count", 5)
thread_count = config.get("decompressor_thread_count", 0)
context_count = config.get("parallel_decompressor_count", DEFAULT_PARALLEL_DECOMPRESSOR_COUNT)
thread_count = config.get("decompressor_thread_count", DEFAULT_DECOMPRESSOR_THREAD_COUNT)
if thread_count == 0:
thread_count = multiprocessing.cpu_count() // 2
disable_cpu_affinity = config.get("disable_cpu_affinity", False)
max_compression_level_allowed = config.get("max_compression_level_allowed", 7)
use_gpu_harvesting = config.get("use_gpu_harvesting", False)
gpu_index = config.get("gpu_index", 0)
enforce_gpu_index = config.get("enforce_gpu_index", False)
disable_cpu_affinity = config.get("disable_cpu_affinity", DEFAULT_DISABLE_CPU_AFFINITY)
max_compression_level_allowed = config.get(
"max_compression_level_allowed", DEFAULT_MAX_COMPRESSION_LEVEL_ALLOWED
)
use_gpu_harvesting = config.get("use_gpu_harvesting", DEFAULT_USE_GPU_HARVESTING)
gpu_index = config.get("gpu_index", DEFAULT_GPU_INDEX)
enforce_gpu_index = config.get("enforce_gpu_index", DEFAULT_ENFORCE_GPU_INDEX)
try:
self._mode = self.plot_manager.configure_decompressor(
+21 -8
View File
@@ -17,6 +17,15 @@ from chia.util.streamable import Streamable, streamable
log = logging.getLogger(__name__)
DEFAULT_PARALLEL_DECOMPRESSOR_COUNT = 5
DEFAULT_DECOMPRESSOR_THREAD_COUNT = 0
DEFAULT_DISABLE_CPU_AFFINITY = False
DEFAULT_MAX_COMPRESSION_LEVEL_ALLOWED = 7
DEFAULT_USE_GPU_HARVESTING = False
DEFAULT_GPU_INDEX = 0
DEFAULT_ENFORCE_GPU_INDEX = False
DEFAULT_RECURSIVE_PLOT_SCAN = False
@streamable
@dataclass(frozen=True)
@@ -99,7 +108,7 @@ def get_plot_filenames(root_path: Path) -> Dict[Path, List[Path]]:
# Returns a map from directory to a list of all plots in the directory
all_files: Dict[Path, List[Path]] = {}
config = load_config(root_path, "config.yaml")
recursive_scan: bool = config["harvester"].get("recursive_plot_scan", False)
recursive_scan: bool = config["harvester"].get("recursive_plot_scan", DEFAULT_RECURSIVE_PLOT_SCAN)
for directory_name in get_plot_directories(root_path, config):
try:
directory = Path(directory_name).resolve()
@@ -161,13 +170,17 @@ def get_harvester_config(root_path: Path) -> Dict[str, Any]:
)
return {
"use_gpu_harvesting": config["harvester"].get("use_gpu_harvesting"),
"gpu_index": config["harvester"].get("gpu_index"),
"enforce_gpu_index": config["harvester"].get("enforce_gpu_index"),
"disable_cpu_affinity": config["harvester"].get("disable_cpu_affinity"),
"parallel_decompressor_count": config["harvester"].get("parallel_decompressor_count"),
"decompressor_thread_count": config["harvester"].get("decompressor_thread_count"),
"recursive_plot_scan": config["harvester"].get("recursive_plot_scan"),
"use_gpu_harvesting": config["harvester"].get("use_gpu_harvesting", DEFAULT_USE_GPU_HARVESTING),
"gpu_index": config["harvester"].get("gpu_index", DEFAULT_GPU_INDEX),
"enforce_gpu_index": config["harvester"].get("enforce_gpu_index", DEFAULT_ENFORCE_GPU_INDEX),
"disable_cpu_affinity": config["harvester"].get("disable_cpu_affinity", DEFAULT_DISABLE_CPU_AFFINITY),
"parallel_decompressor_count": config["harvester"].get(
"parallel_decompressor_count", DEFAULT_PARALLEL_DECOMPRESSOR_COUNT
),
"decompressor_thread_count": config["harvester"].get(
"decompressor_thread_count", DEFAULT_DECOMPRESSOR_THREAD_COUNT
),
"recursive_plot_scan": config["harvester"].get("recursive_plot_scan", DEFAULT_RECURSIVE_PLOT_SCAN),
"plots_refresh_parameter": plots_refresh_parameter,
}