diff --git a/config_sample.yml b/config_sample.yml index 3070642..c92f673 100644 --- a/config_sample.yml +++ b/config_sample.yml @@ -67,6 +67,11 @@ developer: # NOTE: It's recommended to enable this, but if something breaks, turn this off. #uvloop: False + # Set process to use a higher priority + # For realtime process priority, run as administrator or sudo + # Otherwise, the priority will be set to high + #realtime_process_priority: False + # Options for model overrides and loading # Please read the comments to understand how arguments are handled between initial and API loads model: diff --git a/main.py b/main.py index 2ec98db..c62a381 100644 --- a/main.py +++ b/main.py @@ -9,6 +9,8 @@ import signal from loguru import logger from typing import Optional +import psutil + from common import config, gen_logging, sampling, model from common.args import convert_args_to_dict, init_argparser from common.auth import load_auth_keys @@ -143,6 +145,19 @@ def entrypoint(arguments: Optional[dict] = None): logger.warning("EXPERIMENTAL: Running program with Uvloop/Winloop.") + # Set the process priority + if unwrap(developer_config.get("realtime_process_priority"), False): + current_process = psutil.Process(os.getpid()) + if platform.system() == "Windows": + current_process.nice(psutil.REALTIME_PRIORITY_CLASS) + else: + current_process.nice(psutil.IOPRIO_CLASS_RT) + + logger.warning( + "EXPERIMENTAL: Process priority set to Realtime. \n" + "If you're not running on administrator/sudo, the priority is set to high." + ) + # Enter into the async event loop asyncio.run(entrypoint_async()) diff --git a/pyproject.toml b/pyproject.toml index 031e228..f89846f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -28,6 +28,9 @@ dependencies = [ "tokenizers", "lm-format-enforcer >= 0.9.6", "aiofiles", + "aiohttp", + "huggingface_hub", + "psutil", # Improved asyncio loops "uvloop ; platform_system == 'Linux' and platform_machine == 'x86_64'", @@ -35,10 +38,6 @@ dependencies = [ # TEMP: Remove once 2.x is fixed in upstream "numpy < 2.0.0", - - # TODO: Maybe move these to a downloader feature? - "aiohttp", - "huggingface_hub", ] [project.urls]