add cpu restriction feature

compress.py

@@ -2,6 +2,8 @@ import os
 import subprocess
 import logging
 from pathlib import Path
+from concurrent.futures import ThreadPoolExecutor, as_completed
+import psutil
 
 # Set up logging
 logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
@@ -12,6 +14,19 @@ HANDBRAKE_PATH = r"HandBrakeCLI.exe"
 PRESET_FILE = "xiaowan.json"
 PRESET_NAME = "1080P x264 xiaowan"
 
+# New constant for CPU cores per process
+CORES_PER_PROCESS = 4  # Adjust this value as needed
+
+def set_cpu_affinity():
+    # Get the current process
+    process = psutil.Process()
+    # Get the number of CPU cores
+    num_cores = psutil.cpu_count()
+    # Calculate which cores to use
+    cores_to_use = list(range(CORES_PER_PROCESS))
+    # Set the CPU affinity
+    process.cpu_affinity(cores_to_use)
+
 def compress_video(input_path, output_path):
     try:
         command = [
@@ -21,27 +36,47 @@ def compress_video(input_path, output_path):
             "-i", str(input_path),
             "-o", str(output_path)
         ]
-        subprocess.run(command, check=True)
+        # Run the subprocess with CPU affinity set
+        process = subprocess.Popen(command)
+        # Set CPU affinity for the subprocess
+        psutil.Process(process.pid).cpu_affinity(list(range(CORES_PER_PROCESS)))
+        # Wait for the process to complete
+        process.wait()
+        if process.returncode != 0:
+            raise subprocess.CalledProcessError(process.returncode, command)
         logging.info(f"Successfully compressed: {input_path}")
     except subprocess.CalledProcessError as e:
         logging.error(f"Error compressing {input_path}: {e}")
 
 def process_folder(input_folder, output_folder):
-    for root, _, files in os.walk(input_folder):
-        for file in files:
-            input_path = Path(root) / file
-            relative_path = input_path.relative_to(input_folder)
-            output_path = Path(output_folder) / relative_path
+    input_folder = Path(input_folder).resolve()
+    output_folder = Path(output_folder).resolve()
+    
+    # Calculate the number of workers based on available cores and cores per process
+    max_workers = max(1, psutil.cpu_count() // CORES_PER_PROCESS)
+    
+    with ThreadPoolExecutor(max_workers=max_workers) as executor:
+        futures = []
+        for root, _, files in os.walk(input_folder):
+            for file in files:
+                input_path = Path(root) / file
+                relative_path = input_path.relative_to(input_folder)
+                output_path = output_folder / relative_path
 
-            # Check if it's a video file (you may want to add more extensions)
-            if input_path.suffix.lower() in ['.mp4', '.avi', '.mkv', '.mov']:
-                if input_path.stat().st_size > SIZE_THRESHOLD:
-                    output_path.parent.mkdir(parents=True, exist_ok=True)
-                    compress_video(input_path, output_path)
+                if input_path.suffix.lower() in ['.mp4', '.avi', '.mkv', '.mov']:
+                    if input_path.stat().st_size > SIZE_THRESHOLD:
+                        output_path.parent.mkdir(parents=True, exist_ok=True)
+                        futures.append(executor.submit(compress_video, input_path, output_path))
+                    else:
+                        logging.info(f"Skipping {input_path} (size < 5GB)")
                 else:
-                    logging.info(f"Skipping {input_path} (size < 5GB)")
-            else:
-                logging.info(f"Skipping non-video file: {input_path}")
+                    logging.info(f"Skipping non-video file: {input_path}")
+        
+        for future in as_completed(futures):
+            try:
+                future.result()
+            except Exception as e:
+                logging.error(f"An error occurred during compression: {e}")
 
 def main():
     input_folder = input("Enter the input folder path: ")
@@ -59,6 +94,8 @@ def main():
 
 if __name__ == "__main__":
     try:
+        # Set CPU affinity for the main process
+        set_cpu_affinity()
         main()
     except Exception as e:
         logging.exception(f"An unexpected error occurred: {e}")