chore: adjust number of workers

Author: pycorn0729

webgenie/constants.py

@@ -1,6 +1,6 @@
 import bittensor as bt
 import os
-
+import psutil
 # Change this value when updating your code base.
 # Define the version of the webgenie.
 __VERSION__ = "1.1.18" # version
@@ -88,3 +88,11 @@
 NEURON_EPOCH_LENGTH = int(os.getenv("NEURON_EPOCH_LENGTH", 25)) # neuron epoch length
 
 MAX_NUMBER_OF_TASKS_PER_SESSION = 18 # max number of tasks per session
+
+NUMBER_OF_CONCURRENT_WORKERS = max(
+    1,
+    min(
+        os.cpu_count(),
+        (psutil.virtual_memory().total) // (1024 * 1024 * 1024 * 4)
+    )
+)

webgenie/rewards/lighthouse_reward/lighthouse_reward.py

@@ -7,9 +7,9 @@
 import numpy as np
 from typing import List
 
+from webgenie.constants import NUMBER_OF_CONCURRENT_WORKERS
 from webgenie.rewards.reward import Reward
 from webgenie.tasks import Task, Solution
-
 from .get_lighthouse_score import get_lighthouse_score
 
 
@@ -40,9 +40,9 @@ async def reward(self, task: Task, solutions: List[Solution]) -> np.ndarray:
         htmls = [solution.html for solution in solutions]
 
         # Use ProcessPoolExecutor for parallel processing
-        with multiprocessing.Pool(processes=os.cpu_count()) as pool:
+        with multiprocessing.Pool(processes=NUMBER_OF_CONCURRENT_WORKERS) as pool:
             # Convert solutions into chunks for parallel processing
-            chunk_size = max(1, len(htmls) // os.cpu_count()) 
+            chunk_size = max(1, len(htmls) // NUMBER_OF_CONCURRENT_WORKERS) 
 
             html_chunks = [htmls[i:i + chunk_size] for i in range(0, len(htmls), chunk_size)]
 

webgenie/rewards/visual_reward/visual_reward.py

@@ -11,7 +11,7 @@
 from datetime import datetime
 from typing import List
 
-from webgenie.constants import WORK_DIR
+from webgenie.constants import WORK_DIR, NUMBER_OF_CONCURRENT_WORKERS
 from webgenie.rewards.reward import Reward
 from webgenie.rewards.visual_reward.common.browser import start_browser, stop_browser
 from webgenie.rewards.visual_reward.high_level_matching_score import high_level_matching_score
@@ -80,11 +80,11 @@ async def reward(self, task: Task, solutions: List[Solution]) -> np.ndarray:
         current_work_dir = f"{WORK_DIR}/task_{timestamp}_{task.task_id}"
         os.makedirs(current_work_dir, exist_ok=True)
 
-        bt.logging.info(f"The number of cpu cores: {os.cpu_count()}")
+        bt.logging.info(f"The number of concurrent workers: {NUMBER_OF_CONCURRENT_WORKERS}")
         # Use ProcessPoolExecutor for parallel processing
-        with multiprocessing.Pool(processes=os.cpu_count()) as pool:
+        with multiprocessing.Pool(processes=NUMBER_OF_CONCURRENT_WORKERS) as pool:
             # Convert solutions into chunks for parallel processing
-            chunk_size = max(1, len(solutions) // os.cpu_count())
+            chunk_size = max(1, len(solutions) // NUMBER_OF_CONCURRENT_WORKERS)
             solution_chunks = [solutions[i:i + chunk_size] for i in range(0, len(solutions), chunk_size)]
 
             # Create partial tasks for each chunk