diff --git a/jellybench_py/compatibility.py b/jellybench_py/compatibility.py new file mode 100644 index 0000000..e085e4e --- /dev/null +++ b/jellybench_py/compatibility.py @@ -0,0 +1,189 @@ +class CompTool: + """Compatibility Tool for checking system compatibility.""" + + def __init__( + self, + logger, + system_info: dict, + ffmpeg_binary: str | None = None, + skip_prompts: bool | None = False, + ) -> None: + """ + Initializes the compatibility tool. + + :param logger: Logger instance. + :param skip_prompts: If True, skips user prompts. + :param system_info: System information dictionary. + """ + self.logger = logger + self.skip_prompts = skip_prompts + self.system_info = system_info + self.ffmpeg_binary = ffmpeg_binary + self.devices = { + "cpu": False, + "gpu": {}, + } + + def set_gpu(self, idx: int) -> None: + """Adds a GPU device to the compatibility check.""" + if idx not in self.devices["gpu"]: + self.devices["gpu"][idx] = { + "compatible": None, + "limit": None, + } + self.logger.info(f"GPU {idx} added to the compatibility check list.") + else: + self.logger.warning(f"GPU {idx} already exists in the compatibility check.") + + def set_cpu(self, state: bool) -> None: + """Adds a CPU device to the compatibility check.""" + self.devices["cpu"] = state + self.logger.info("CPU added to the compatibility check list.") + + def set_ffmpeg_binary(self, ffmpeg_binary: str) -> None: + """Sets the ffmpeg binary for the compatibility check.""" + self.ffmpeg_binary = ffmpeg_binary + self.logger.info(f"FFmpeg binary set to {ffmpeg_binary}.") + + def run(self): + """Runs the compatibility checks.""" + print("Running compatibility checks...") + if self.devices["cpu"]: + print("Running CPU compatibility check...") + if self.devices["gpu"]: + print("Running GPU compatibility check...") + for idx, device in self.devices["gpu"].items(): + print(f"Checking GPU {idx}...") + gpu_data = self.system_info["gpu"][idx] + if gpu_data["vendor"] == "nvidia": + nvidia_handler() + elif gpu_data["vendor"] == "amd": + amd_handler() + elif gpu_data["vendor"] == "intel": + intel_handler() + else: + print(f"Unknown GPU vendor: {gpu_data['vendor']}") + pass # Placeholder for actual implementation + + +def nvidia_handler(): + """Handles NVIDIA-specific compatibility checks.""" + print("NVIDIA GPU detected, checking driver limits...") + pass # Placeholder for actual implementation + + +def amd_handler(): + """Handles AMD-specific compatibility checks.""" + print("AMD GPU detected, checking compatibility...") + pass # Placeholder for actual implementation + + +def intel_handler(): + """Handles Intel-specific compatibility checks.""" + print("Intel GPU detected, checking compatibility...") + pass # Placeholder for actual implementation + + +# To be implemented +""" +def check_driver_limit_demo(device: dict, ffmpeg_binary: str, gpu_idx: int): + def build_test_cmd(worker_ammount: int, ffmpeg_binary: str, gpu_arg) -> str: + # Build an ffmpeg command to test {n}-concurrent NvEnc Streams + if hwi.platform.system().lower() == "windows": + base_cmd = Constants.NVENC_TEST_WINDOWS.BASE_CMD.format( + ffmpeg=ffmpeg_binary, gpu=gpu_arg + ) + worker_base = Constants.NVENC_TEST_WINDOWS.WORKER_CMD + else: + base_cmd = Constants.NVENC_TEST_LINUX.BASE_CMD.format( + ffmpeg=ffmpeg_binary, gpu=gpu_arg + ) + worker_base = Constants.NVENC_TEST_LINUX.WORKER_CMD + worker_commands = [] + for i in range(1, worker_ammount + 1): + bitrate = f"{i}M" + worker_command = worker_base.format(bitrate=bitrate) + worker_commands.append(worker_command) + full_command = base_cmd + " ".join(worker_commands) + return full_command + + def parse_driver(driver_raw: str) -> int: + # parse windows specific driver output into NVIDIA Version (e.g. 32.0.15.6603 --> 566.03) + split_driver = driver_raw.split(".") + if len(split_driver) != 4: + return -1 + + major_version = list(split_driver[-2]) + if 0 < len(major_version): + major_version = major_version[-1] + else: + return -1 + minor_version = split_driver[-1] + + if not (major_version.isdigit() and minor_version.isdigit()): + return -1 + + driver_version = float(major_version + minor_version) / 100 + limit = get_nvenc_session_limit(driver_version) + print( + f"| Your driver ({driver_version}) should only allow {limit} NvEnc sessions" + ) + return limit + + limit = 8 # default driver limit + + if "configuration" in device and "driver" in device["configuration"]: + driver_raw = device["configuration"]["driver"] + driver_limit = parse_driver(driver_raw) + if 0 < driver_limit: + limit = driver_limit + + gpu_arg = format_gpu_arg(hwi.platform.system(), device, gpu_idx) + worker_ammount = limit + 1 + print(f"| Testing with {worker_ammount} workers...", end="") + command = build_test_cmd(worker_ammount, ffmpeg_binary, gpu_arg) + + skip_device = False + limited_driver = 0 + successful_count, failure_reason = worker.test_command(command, ffmpeg_log) + if successful_count == worker_ammount: + print(" success!") + + elif 0 < successful_count < worker_ammount: + print(styled(" limited!", [Style.BG_RED])) + print( + f"| > Your GPU driver does only allow {successful_count} concurrent NvEnc sessions!" + ) + skip_device = confirm( + message="| > Do you want to skip GPU tests?", + default=False, + automate=skip_prompts, + ) + limited_driver = successful_count + else: + print(" Error!") + print("| > Your GPU is not capable of running NvEnc Streams!") + print(f"| > FFmpeg: {failure_reason}") + print("| > Please run the tool again and disable GPU tests") + exit() + print(styled("Done", [Style.GREEN])) + print() + return limited_driver, skip_device + +def test_command(ffmpeg_cmd: str, passed_logger: Logger): + passed_logger.info(f"> > Running 1x ffmpeg {ffmpeg_cmd}") + ffmpeg_cmd_list = shlex.split(ffmpeg_cmd) + successful_stream_count = 0 + raw_worker_data = run_ffmpeg(1, ffmpeg_cmd_list, passed_logger) + + failure_reason = raw_worker_data[1] + process_output = raw_worker_data[0] + nvenc_limit_reasons = ["incompatible client key", "out of memory"] + if failure_reason in nvenc_limit_reasons or failure_reason is None: + success_pattern = r"Output #\d+, null, to 'pipe:'" + successful_streams = re.findall(success_pattern, process_output) + successful_stream_count = len(successful_streams) + passed_logger.info("< < Finished ffmpeg") + + return successful_stream_count, failure_reason +""" diff --git a/jellybench_py/core.py b/jellybench_py/core.py index d361e4f..68656ff 100644 --- a/jellybench_py/core.py +++ b/jellybench_py/core.py @@ -30,6 +30,7 @@ from jellybench_py import hwi, worker from jellybench_py.api import ApiClient, ApiError +from jellybench_py.compatibility import CompTool from jellybench_py.constant import Constants, Style from jellybench_py.util import ( confirm, @@ -727,6 +728,10 @@ def cli() -> None: print("| Obtaining System Information...", end="") system_info = hwi.get_system_info() + + # initialize compatibility tool + comp_tool = CompTool(system_info=system_info, logger=main_log) + print(" success!") print("| Detected System Config:") print(f"| OS: {system_info['os']['pretty_name']}") @@ -775,6 +780,7 @@ def cli() -> None: if not args.disable_cpu: main_log.info("Using all available CPUs") supported_types.append("cpu") + comp_tool.set_cpu(True) # GPU Logic gpus = system_info["gpu"] @@ -815,6 +821,7 @@ def cli() -> None: gpu = gpus[gpu_idx] main_log.info(f"Using GPU {args.gpu_input}") supported_types.append(gpu["vendor"]) + comp_tool.set_gpu(gpu_idx) # Error if all hardware disabled if args.gpu_input == 0 and args.disable_cpu: @@ -884,6 +891,7 @@ def cli() -> None: ffmpeg_binary = os.path.abspath(ffmpeg_binary) ffmpeg_binary = ffmpeg_binary.replace("\\", "\\\\") + comp_tool.set_ffmpeg(ffmpeg_binary) print(styled("Done", [Style.GREEN])) print() @@ -906,6 +914,7 @@ def cli() -> None: print() # Test for NvEnc Limits + comp_tool.run() if gpu and gpu["vendor"] == "nvidia": print(styled("Testing for driver limits: ", [Style.BOLD]) + "(NVIDIA)") limited_driver, skip_device = check_driver_limit(gpu, ffmpeg_binary, gpu_idx)