diff --git a/pyproject.toml b/pyproject.toml index 16f662d6ba..648aad3ae1 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -25,8 +25,30 @@ classifiers = [ "Topic :: Scientific/Engineering :: Artificial Intelligence", ] dependencies = [ + "unsloth_zoo>=2026.6.4", + "wheel>=0.42.0", + "packaging", + "torch>=2.4.0,<2.11.0", + "torchvision", + "numpy", + "tqdm", + "psutil", + "tyro", + "protobuf", + "xformers>=0.0.27.post2 ; ('linux' in sys_platform or sys_platform == 'win32') and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", + "bitsandbytes>=0.45.5,!=0.46.0,!=0.48.0", + "triton>=3.0.0 ; ('linux' in sys_platform)", + "triton-windows ; (sys_platform == 'win32') and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", + "sentencepiece>=0.2.0", + "datasets>=3.4.1,!=4.0.*,!=4.1.0,<4.4.0", + "accelerate>=0.34.1", + "peft>=0.18.0,!=0.11.0", + "huggingface_hub>=0.34.0", + "hf_transfer", + "diffusers", + "transformers>=4.51.3,!=4.52.0,!=4.52.1,!=4.52.2,!=4.52.3,!=4.53.0,!=4.54.0,!=4.55.0,!=4.55.1,!=4.57.0,!=4.57.4,!=4.57.5,!=5.0.0,!=5.1.0,<=5.5.0", + "trl>=0.18.2,!=0.19.0,<=0.24.0", "typer", - "rich", "pydantic", "pyyaml", "nest-asyncio", @@ -55,7 +77,6 @@ studio = [ "frontend/.git*", "backend/requirements/**/*", "backend/plugins/**/*", - "backend/assets/**/*.jinja", "backend/core/data_recipe/oxc-validator/*.json", "backend/core/data_recipe/oxc-validator/*.mjs", ] @@ -69,11 +90,11 @@ triton = [ "triton>=3.0.0 ; ('linux' in sys_platform)", "triton-windows ; (sys_platform == 'win32') and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", ] - -huggingfacenotorch = [ +huggingface = [ "unsloth_zoo>=2026.6.4", "wheel>=0.42.0", "packaging", + "torchvision", "numpy", "tqdm", "psutil", @@ -89,12 +110,10 @@ huggingfacenotorch = [ "transformers>=4.51.3,!=4.52.0,!=4.52.1,!=4.52.2,!=4.52.3,!=4.53.0,!=4.54.0,!=4.55.0,!=4.55.1,!=4.57.0,!=4.57.4,!=4.57.5,!=5.0.0,!=5.1.0,<=5.5.0", "trl>=0.18.2,!=0.19.0,<=0.24.0", "sentence-transformers", -] -huggingface = [ - "unsloth[huggingfacenotorch]", - "unsloth_zoo>=2026.6.4", - "torchvision", - "unsloth[triton]", + "typer", + "pydantic", + "pyyaml", + "nest-asyncio", ] windows = [ "unsloth[huggingface]", @@ -105,240 +124,126 @@ base = [ "unsloth[huggingface]", ] cu118only = [ - "xformers @ https://download.pytorch.org/whl/cu118/xformers-0.0.22.post7%2Bcu118-cp39-cp39-manylinux2014_x86_64.whl ; python_version=='3.9' and ('linux' in sys_platform)", - "xformers @ https://download.pytorch.org/whl/cu118/xformers-0.0.22.post7%2Bcu118-cp310-cp310-manylinux2014_x86_64.whl ; python_version=='3.10' and ('linux' in sys_platform)", - "xformers @ https://download.pytorch.org/whl/cu118/xformers-0.0.22.post7%2Bcu118-cp311-cp311-manylinux2014_x86_64.whl ; python_version=='3.11' and ('linux' in sys_platform)", + "xformers==0.0.22.post7 ; ('linux' in sys_platform or sys_platform == 'win32') and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", ] cu121only = [ - "xformers @ https://download.pytorch.org/whl/cu121/xformers-0.0.22.post7-cp39-cp39-manylinux2014_x86_64.whl ; python_version=='3.9' and ('linux' in sys_platform)", - "xformers @ https://download.pytorch.org/whl/cu121/xformers-0.0.22.post7-cp310-cp310-manylinux2014_x86_64.whl ; python_version=='3.10' and ('linux' in sys_platform)", - "xformers @ https://download.pytorch.org/whl/cu121/xformers-0.0.22.post7-cp311-cp311-manylinux2014_x86_64.whl ; python_version=='3.11' and ('linux' in sys_platform)", + "xformers==0.0.22.post7 ; ('linux' in sys_platform or sys_platform == 'win32') and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", ] cu118onlytorch211 = [ - "xformers @ https://download.pytorch.org/whl/cu118/xformers-0.0.23%2Bcu118-cp39-cp39-manylinux2014_x86_64.whl ; python_version=='3.9' and ('linux' in sys_platform)", - "xformers @ https://download.pytorch.org/whl/cu118/xformers-0.0.23%2Bcu118-cp310-cp310-manylinux2014_x86_64.whl ; python_version=='3.10' and ('linux' in sys_platform)", - "xformers @ https://download.pytorch.org/whl/cu118/xformers-0.0.23%2Bcu118-cp311-cp311-manylinux2014_x86_64.whl ; python_version=='3.11' and ('linux' in sys_platform)", + "xformers==0.0.23 ; ('linux' in sys_platform or sys_platform == 'win32') and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", ] cu121onlytorch211 = [ - "xformers @ https://download.pytorch.org/whl/cu121/xformers-0.0.23-cp39-cp39-manylinux2014_x86_64.whl ; python_version=='3.9' and ('linux' in sys_platform)", - "xformers @ https://download.pytorch.org/whl/cu121/xformers-0.0.23-cp310-cp310-manylinux2014_x86_64.whl ; python_version=='3.10' and ('linux' in sys_platform)", - "xformers @ https://download.pytorch.org/whl/cu121/xformers-0.0.23-cp311-cp311-manylinux2014_x86_64.whl ; python_version=='3.11' and ('linux' in sys_platform)", + "xformers==0.0.23 ; ('linux' in sys_platform or sys_platform == 'win32') and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", ] cu118onlytorch212 = [ - "xformers @ https://download.pytorch.org/whl/cu118/xformers-0.0.23.post1%2Bcu118-cp39-cp39-manylinux2014_x86_64.whl ; python_version=='3.9' and ('linux' in sys_platform)", - "xformers @ https://download.pytorch.org/whl/cu118/xformers-0.0.23.post1%2Bcu118-cp310-cp310-manylinux2014_x86_64.whl ; python_version=='3.10' and ('linux' in sys_platform)", - "xformers @ https://download.pytorch.org/whl/cu118/xformers-0.0.23.post1%2Bcu118-cp311-cp311-manylinux2014_x86_64.whl ; python_version=='3.11' and ('linux' in sys_platform)", + "xformers==0.0.23.post1 ; ('linux' in sys_platform or sys_platform == 'win32') and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", ] cu121onlytorch212 = [ - "xformers @ https://download.pytorch.org/whl/cu121/xformers-0.0.23.post1-cp39-cp39-manylinux2014_x86_64.whl ; python_version=='3.9' and ('linux' in sys_platform)", - "xformers @ https://download.pytorch.org/whl/cu121/xformers-0.0.23.post1-cp310-cp310-manylinux2014_x86_64.whl ; python_version=='3.10' and ('linux' in sys_platform)", - "xformers @ https://download.pytorch.org/whl/cu121/xformers-0.0.23.post1-cp311-cp311-manylinux2014_x86_64.whl ; python_version=='3.11' and ('linux' in sys_platform)", + "xformers==0.0.23.post1 ; ('linux' in sys_platform or sys_platform == 'win32') and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", ] cu118onlytorch220 = [ - "xformers @ https://download.pytorch.org/whl/cu118/xformers-0.0.24%2Bcu118-cp39-cp39-manylinux2014_x86_64.whl ; python_version=='3.9' and ('linux' in sys_platform)", - "xformers @ https://download.pytorch.org/whl/cu118/xformers-0.0.24%2Bcu118-cp310-cp310-manylinux2014_x86_64.whl ; python_version=='3.10' and ('linux' in sys_platform)", - "xformers @ https://download.pytorch.org/whl/cu118/xformers-0.0.24%2Bcu118-cp311-cp311-manylinux2014_x86_64.whl ; python_version=='3.11' and ('linux' in sys_platform)", + "xformers==0.0.24 ; ('linux' in sys_platform or sys_platform == 'win32') and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", ] cu121onlytorch220 = [ - "xformers @ https://download.pytorch.org/whl/cu121/xformers-0.0.24-cp39-cp39-manylinux2014_x86_64.whl ; python_version=='3.9' and ('linux' in sys_platform)", - "xformers @ https://download.pytorch.org/whl/cu121/xformers-0.0.24-cp310-cp310-manylinux2014_x86_64.whl ; python_version=='3.10' and ('linux' in sys_platform)", - "xformers @ https://download.pytorch.org/whl/cu121/xformers-0.0.24-cp311-cp311-manylinux2014_x86_64.whl ; python_version=='3.11' and ('linux' in sys_platform)", + "xformers==0.0.24 ; ('linux' in sys_platform or sys_platform == 'win32') and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", ] cu118onlytorch230 = [ - "xformers @ https://download.pytorch.org/whl/cu118/xformers-0.0.27%2Bcu118-cp39-cp39-manylinux2014_x86_64.whl ; python_version=='3.9' and ('linux' in sys_platform)", - "xformers @ https://download.pytorch.org/whl/cu118/xformers-0.0.27%2Bcu118-cp310-cp310-manylinux2014_x86_64.whl ; python_version=='3.10' and ('linux' in sys_platform)", - "xformers @ https://download.pytorch.org/whl/cu118/xformers-0.0.27%2Bcu118-cp311-cp311-manylinux2014_x86_64.whl ; python_version=='3.11' and ('linux' in sys_platform)", - "xformers @ https://download.pytorch.org/whl/cu118/xformers-0.0.27%2Bcu118-cp312-cp312-manylinux2014_x86_64.whl ; python_version=='3.12' and ('linux' in sys_platform)", + "xformers==0.0.27 ; ('linux' in sys_platform or sys_platform == 'win32') and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", ] cu121onlytorch230 = [ - "xformers @ https://download.pytorch.org/whl/cu121/xformers-0.0.27-cp39-cp39-manylinux2014_x86_64.whl ; python_version=='3.9' and ('linux' in sys_platform)", - "xformers @ https://download.pytorch.org/whl/cu121/xformers-0.0.27-cp310-cp310-manylinux2014_x86_64.whl ; python_version=='3.10' and ('linux' in sys_platform)", - "xformers @ https://download.pytorch.org/whl/cu121/xformers-0.0.27-cp311-cp311-manylinux2014_x86_64.whl ; python_version=='3.11' and ('linux' in sys_platform)", - "xformers @ https://download.pytorch.org/whl/cu121/xformers-0.0.27-cp312-cp312-manylinux2014_x86_64.whl ; python_version=='3.12' and ('linux' in sys_platform)", + "xformers==0.0.27 ; ('linux' in sys_platform or sys_platform == 'win32') and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", ] cu118onlytorch240 = [ - "xformers @ https://download.pytorch.org/whl/cu118/xformers-0.0.27.post2%2Bcu118-cp39-cp39-manylinux2014_x86_64.whl ; python_version=='3.9' and ('linux' in sys_platform)", - "xformers @ https://download.pytorch.org/whl/cu118/xformers-0.0.27.post2%2Bcu118-cp310-cp310-manylinux2014_x86_64.whl ; python_version=='3.10' and ('linux' in sys_platform)", - "xformers @ https://download.pytorch.org/whl/cu118/xformers-0.0.27.post2%2Bcu118-cp311-cp311-manylinux2014_x86_64.whl ; python_version=='3.11' and ('linux' in sys_platform)", - "xformers @ https://download.pytorch.org/whl/cu118/xformers-0.0.27.post2%2Bcu118-cp312-cp312-manylinux2014_x86_64.whl ; python_version=='3.12' and ('linux' in sys_platform)", + "xformers==0.0.27.post2 ; ('linux' in sys_platform or sys_platform == 'win32') and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", ] cu121onlytorch240 = [ - "xformers @ https://download.pytorch.org/whl/cu121/xformers-0.0.28.post1-cp39-cp39-manylinux_2_28_x86_64.whl ; python_version=='3.9' and ('linux' in sys_platform)", - "xformers @ https://download.pytorch.org/whl/cu121/xformers-0.0.28.post1-cp310-cp310-manylinux_2_28_x86_64.whl ; python_version=='3.10' and ('linux' in sys_platform)", - "xformers @ https://download.pytorch.org/whl/cu121/xformers-0.0.28.post1-cp311-cp311-manylinux_2_28_x86_64.whl ; python_version=='3.11' and ('linux' in sys_platform)", - "xformers @ https://download.pytorch.org/whl/cu121/xformers-0.0.28.post1-cp312-cp312-manylinux_2_28_x86_64.whl ; python_version=='3.12' and ('linux' in sys_platform)", + "xformers==0.0.27.post2 ; ('linux' in sys_platform or sys_platform == 'win32') and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", ] cu124onlytorch240 = [ - "xformers @ https://download.pytorch.org/whl/cu124/xformers-0.0.28.post1-cp39-cp39-manylinux_2_28_x86_64.whl ; python_version=='3.9' and ('linux' in sys_platform)", - "xformers @ https://download.pytorch.org/whl/cu124/xformers-0.0.28.post1-cp310-cp310-manylinux_2_28_x86_64.whl ; python_version=='3.10' and ('linux' in sys_platform)", - "xformers @ https://download.pytorch.org/whl/cu124/xformers-0.0.28.post1-cp311-cp311-manylinux_2_28_x86_64.whl ; python_version=='3.11' and ('linux' in sys_platform)", - "xformers @ https://download.pytorch.org/whl/cu124/xformers-0.0.28.post1-cp312-cp312-manylinux_2_28_x86_64.whl ; python_version=='3.12' and ('linux' in sys_platform)", - "xformers @ https://download.pytorch.org/whl/cu124/xformers-0.0.28.post1-cp39-cp39-win_amd64.whl ; python_version=='3.9' and (sys_platform == 'win32')", - "xformers @ https://download.pytorch.org/whl/cu124/xformers-0.0.28.post1-cp310-cp310-win_amd64.whl ; python_version=='3.10' and (sys_platform == 'win32')", - "xformers @ https://download.pytorch.org/whl/cu124/xformers-0.0.28.post1-cp311-cp311-win_amd64.whl ; python_version=='3.11' and (sys_platform == 'win32')", - "xformers @ https://download.pytorch.org/whl/cu124/xformers-0.0.28.post1-cp312-cp312-win_amd64.whl ; python_version=='3.12' and (sys_platform == 'win32')", + "xformers==0.0.28.post1 ; ('linux' in sys_platform or sys_platform == 'win32') and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", ] cu118onlytorch250 = [ - "xformers @ https://download.pytorch.org/whl/cu118/xformers-0.0.28.post2-cp39-cp39-manylinux_2_28_x86_64.whl ; python_version=='3.9' and ('linux' in sys_platform)", - "xformers @ https://download.pytorch.org/whl/cu118/xformers-0.0.28.post2-cp310-cp310-manylinux_2_28_x86_64.whl ; python_version=='3.10' and ('linux' in sys_platform)", - "xformers @ https://download.pytorch.org/whl/cu118/xformers-0.0.28.post2-cp311-cp311-manylinux_2_28_x86_64.whl ; python_version=='3.11' and ('linux' in sys_platform)", - "xformers @ https://download.pytorch.org/whl/cu118/xformers-0.0.28.post2-cp312-cp312-manylinux_2_28_x86_64.whl ; python_version=='3.12' and ('linux' in sys_platform)", + "xformers==0.0.28.post2 ; ('linux' in sys_platform or sys_platform == 'win32') and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", ] cu121onlytorch250 = [ - "xformers @ https://download.pytorch.org/whl/cu121/xformers-0.0.28.post2-cp39-cp39-manylinux_2_28_x86_64.whl ; python_version=='3.9' and ('linux' in sys_platform)", - "xformers @ https://download.pytorch.org/whl/cu121/xformers-0.0.28.post2-cp310-cp310-manylinux_2_28_x86_64.whl ; python_version=='3.10' and ('linux' in sys_platform)", - "xformers @ https://download.pytorch.org/whl/cu121/xformers-0.0.28.post2-cp311-cp311-manylinux_2_28_x86_64.whl ; python_version=='3.11' and ('linux' in sys_platform)", - "xformers @ https://download.pytorch.org/whl/cu121/xformers-0.0.28.post2-cp312-cp312-manylinux_2_28_x86_64.whl ; python_version=='3.12' and ('linux' in sys_platform)", + "xformers==0.0.28.post2 ; ('linux' in sys_platform or sys_platform == 'win32') and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", ] cu124onlytorch250 = [ - "xformers @ https://download.pytorch.org/whl/cu124/xformers-0.0.28.post2-cp39-cp39-manylinux_2_28_x86_64.whl ; python_version=='3.9' and ('linux' in sys_platform)", - "xformers @ https://download.pytorch.org/whl/cu124/xformers-0.0.28.post2-cp310-cp310-manylinux_2_28_x86_64.whl ; python_version=='3.10' and ('linux' in sys_platform)", - "xformers @ https://download.pytorch.org/whl/cu124/xformers-0.0.28.post2-cp311-cp311-manylinux_2_28_x86_64.whl ; python_version=='3.11' and ('linux' in sys_platform)", - "xformers @ https://download.pytorch.org/whl/cu124/xformers-0.0.28.post2-cp312-cp312-manylinux_2_28_x86_64.whl ; python_version=='3.12' and ('linux' in sys_platform)", - "xformers @ https://download.pytorch.org/whl/cu124/xformers-0.0.28.post2-cp39-cp39-win_amd64.whl ; python_version=='3.9' and (sys_platform == 'win32')", - "xformers @ https://download.pytorch.org/whl/cu124/xformers-0.0.28.post2-cp310-cp310-win_amd64.whl ; python_version=='3.10' and (sys_platform == 'win32')", - "xformers @ https://download.pytorch.org/whl/cu124/xformers-0.0.28.post2-cp311-cp311-win_amd64.whl ; python_version=='3.11' and (sys_platform == 'win32')", - "xformers @ https://download.pytorch.org/whl/cu124/xformers-0.0.28.post2-cp312-cp312-win_amd64.whl ; python_version=='3.12' and (sys_platform == 'win32')", + "xformers==0.0.28.post2 ; ('linux' in sys_platform or sys_platform == 'win32') and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", ] cu118onlytorch251 = [ - "xformers @ https://download.pytorch.org/whl/cu118/xformers-0.0.29.post1-cp39-cp39-manylinux_2_28_x86_64.whl ; python_version=='3.9' and ('linux' in sys_platform)", - "xformers @ https://download.pytorch.org/whl/cu118/xformers-0.0.29.post1-cp310-cp310-manylinux_2_28_x86_64.whl ; python_version=='3.10' and ('linux' in sys_platform)", - "xformers @ https://download.pytorch.org/whl/cu118/xformers-0.0.29.post1-cp311-cp311-manylinux_2_28_x86_64.whl ; python_version=='3.11' and ('linux' in sys_platform)", - "xformers @ https://download.pytorch.org/whl/cu118/xformers-0.0.29.post1-cp312-cp312-manylinux_2_28_x86_64.whl ; python_version=='3.12' and ('linux' in sys_platform)", + "xformers==0.0.29.post1 ; ('linux' in sys_platform or sys_platform == 'win32') and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", ] cu121onlytorch251 = [ - "xformers @ https://download.pytorch.org/whl/cu121/xformers-0.0.29.post1-cp39-cp39-manylinux_2_28_x86_64.whl ; python_version=='3.9' and ('linux' in sys_platform)", - "xformers @ https://download.pytorch.org/whl/cu121/xformers-0.0.29.post1-cp310-cp310-manylinux_2_28_x86_64.whl ; python_version=='3.10' and ('linux' in sys_platform)", - "xformers @ https://download.pytorch.org/whl/cu121/xformers-0.0.29.post1-cp311-cp311-manylinux_2_28_x86_64.whl ; python_version=='3.11' and ('linux' in sys_platform)", - "xformers @ https://download.pytorch.org/whl/cu121/xformers-0.0.29.post1-cp312-cp312-manylinux_2_28_x86_64.whl ; python_version=='3.12' and ('linux' in sys_platform)", + "xformers==0.0.29.post1 ; ('linux' in sys_platform or sys_platform == 'win32') and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", ] cu124onlytorch251 = [ - "xformers @ https://download.pytorch.org/whl/cu124/xformers-0.0.29.post1-cp39-cp39-manylinux_2_28_x86_64.whl ; python_version=='3.9' and ('linux' in sys_platform)", - "xformers @ https://download.pytorch.org/whl/cu124/xformers-0.0.29.post1-cp310-cp310-manylinux_2_28_x86_64.whl ; python_version=='3.10' and ('linux' in sys_platform)", - "xformers @ https://download.pytorch.org/whl/cu124/xformers-0.0.29.post1-cp311-cp311-manylinux_2_28_x86_64.whl ; python_version=='3.11' and ('linux' in sys_platform)", - "xformers @ https://download.pytorch.org/whl/cu124/xformers-0.0.29.post1-cp312-cp312-manylinux_2_28_x86_64.whl ; python_version=='3.12' and ('linux' in sys_platform)", - "xformers @ https://download.pytorch.org/whl/cu124/xformers-0.0.29.post1-cp39-cp39-win_amd64.whl ; python_version=='3.9' and (sys_platform == 'win32')", - "xformers @ https://download.pytorch.org/whl/cu124/xformers-0.0.29.post1-cp310-cp310-win_amd64.whl ; python_version=='3.10' and (sys_platform == 'win32')", - "xformers @ https://download.pytorch.org/whl/cu124/xformers-0.0.29.post1-cp311-cp311-win_amd64.whl ; python_version=='3.11' and (sys_platform == 'win32')", - "xformers @ https://download.pytorch.org/whl/cu124/xformers-0.0.29.post1-cp312-cp312-win_amd64.whl ; python_version=='3.12' and (sys_platform == 'win32')", + "xformers==0.0.29.post1 ; ('linux' in sys_platform or sys_platform == 'win32') and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", ] cu118onlytorch260 = [ - "xformers @ https://download.pytorch.org/whl/cu118/xformers-0.0.29.post3-cp39-cp39-manylinux_2_28_x86_64.whl ; python_version=='3.9' and ('linux' in sys_platform)", - "xformers @ https://download.pytorch.org/whl/cu118/xformers-0.0.29.post3-cp310-cp310-manylinux_2_28_x86_64.whl ; python_version=='3.10' and ('linux' in sys_platform)", - "xformers @ https://download.pytorch.org/whl/cu118/xformers-0.0.29.post3-cp311-cp311-manylinux_2_28_x86_64.whl ; python_version=='3.11' and ('linux' in sys_platform)", - "xformers @ https://download.pytorch.org/whl/cu118/xformers-0.0.29.post3-cp312-cp312-manylinux_2_28_x86_64.whl ; python_version=='3.12' and ('linux' in sys_platform)", + "xformers==0.0.29.post3 ; ('linux' in sys_platform or sys_platform == 'win32') and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", ] cu124onlytorch260 = [ - "xformers @ https://download.pytorch.org/whl/cu124/xformers-0.0.29.post3-cp39-cp39-manylinux_2_28_x86_64.whl ; python_version=='3.9' and ('linux' in sys_platform)", - "xformers @ https://download.pytorch.org/whl/cu124/xformers-0.0.29.post3-cp310-cp310-manylinux_2_28_x86_64.whl ; python_version=='3.10' and ('linux' in sys_platform)", - "xformers @ https://download.pytorch.org/whl/cu124/xformers-0.0.29.post3-cp311-cp311-manylinux_2_28_x86_64.whl ; python_version=='3.11' and ('linux' in sys_platform)", - "xformers @ https://download.pytorch.org/whl/cu124/xformers-0.0.29.post3-cp312-cp312-manylinux_2_28_x86_64.whl ; python_version=='3.12' and ('linux' in sys_platform)", - "xformers @ https://download.pytorch.org/whl/cu124/xformers-0.0.29.post3-cp39-cp39-win_amd64.whl ; python_version=='3.9' and (sys_platform == 'win32')", - "xformers @ https://download.pytorch.org/whl/cu124/xformers-0.0.29.post3-cp310-cp310-win_amd64.whl ; python_version=='3.10' and (sys_platform == 'win32')", - "xformers @ https://download.pytorch.org/whl/cu124/xformers-0.0.29.post3-cp311-cp311-win_amd64.whl ; python_version=='3.11' and (sys_platform == 'win32')", - "xformers @ https://download.pytorch.org/whl/cu124/xformers-0.0.29.post3-cp312-cp312-win_amd64.whl ; python_version=='3.12' and (sys_platform == 'win32')", + "xformers==0.0.29.post3 ; ('linux' in sys_platform or sys_platform == 'win32') and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", ] cu126onlytorch260 = [ - "xformers @ https://download.pytorch.org/whl/cu126/xformers-0.0.29.post3-cp39-cp39-manylinux_2_28_x86_64.whl ; python_version=='3.9' and ('linux' in sys_platform)", - "xformers @ https://download.pytorch.org/whl/cu126/xformers-0.0.29.post3-cp310-cp310-manylinux_2_28_x86_64.whl ; python_version=='3.10' and ('linux' in sys_platform)", - "xformers @ https://download.pytorch.org/whl/cu126/xformers-0.0.29.post3-cp311-cp311-manylinux_2_28_x86_64.whl ; python_version=='3.11' and ('linux' in sys_platform)", - "xformers @ https://download.pytorch.org/whl/cu126/xformers-0.0.29.post3-cp312-cp312-manylinux_2_28_x86_64.whl ; python_version=='3.12' and ('linux' in sys_platform)", - "xformers @ https://download.pytorch.org/whl/cu126/xformers-0.0.29.post3-cp39-cp39-win_amd64.whl ; python_version=='3.9' and (sys_platform == 'win32')", - "xformers @ https://download.pytorch.org/whl/cu126/xformers-0.0.29.post3-cp310-cp310-win_amd64.whl ; python_version=='3.10' and (sys_platform == 'win32')", - "xformers @ https://download.pytorch.org/whl/cu126/xformers-0.0.29.post3-cp311-cp311-win_amd64.whl ; python_version=='3.11' and (sys_platform == 'win32')", - "xformers @ https://download.pytorch.org/whl/cu126/xformers-0.0.29.post3-cp312-cp312-win_amd64.whl ; python_version=='3.12' and (sys_platform == 'win32')", + "xformers==0.0.29.post3 ; ('linux' in sys_platform or sys_platform == 'win32') and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", ] cu118onlytorch270 = [ - "xformers @ https://download.pytorch.org/whl/cu118/xformers-0.0.30-cp39-cp39-manylinux_2_28_x86_64.whl ; python_version=='3.9' and ('linux' in sys_platform)", - "xformers @ https://download.pytorch.org/whl/cu118/xformers-0.0.30-cp310-cp310-manylinux_2_28_x86_64.whl ; python_version=='3.10' and ('linux' in sys_platform)", - "xformers @ https://download.pytorch.org/whl/cu118/xformers-0.0.30-cp311-cp311-manylinux_2_28_x86_64.whl ; python_version=='3.11' and ('linux' in sys_platform)", - "xformers @ https://download.pytorch.org/whl/cu118/xformers-0.0.30-cp312-cp312-manylinux_2_28_x86_64.whl ; python_version=='3.12' and ('linux' in sys_platform)", - "xformers @ https://download.pytorch.org/whl/cu118/xformers-0.0.30-cp39-cp39-win_amd64.whl ; python_version=='3.9' and (sys_platform == 'win32')", - "xformers @ https://download.pytorch.org/whl/cu118/xformers-0.0.30-cp310-cp310-win_amd64.whl ; python_version=='3.10' and (sys_platform == 'win32')", - "xformers @ https://download.pytorch.org/whl/cu118/xformers-0.0.30-cp311-cp311-win_amd64.whl ; python_version=='3.11' and (sys_platform == 'win32')", - "xformers @ https://download.pytorch.org/whl/cu118/xformers-0.0.30-cp312-cp312-win_amd64.whl ; python_version=='3.12' and (sys_platform == 'win32')", + "xformers==0.0.30 ; ('linux' in sys_platform or sys_platform == 'win32') and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", ] cu126onlytorch270 = [ - "xformers @ https://download.pytorch.org/whl/cu126/xformers-0.0.30-cp39-cp39-manylinux_2_28_x86_64.whl ; python_version=='3.9' and ('linux' in sys_platform)", - "xformers @ https://download.pytorch.org/whl/cu126/xformers-0.0.30-cp310-cp310-manylinux_2_28_x86_64.whl ; python_version=='3.10' and ('linux' in sys_platform)", - "xformers @ https://download.pytorch.org/whl/cu126/xformers-0.0.30-cp311-cp311-manylinux_2_28_x86_64.whl ; python_version=='3.11' and ('linux' in sys_platform)", - "xformers @ https://download.pytorch.org/whl/cu126/xformers-0.0.30-cp312-cp312-manylinux_2_28_x86_64.whl ; python_version=='3.12' and ('linux' in sys_platform)", - "xformers @ https://download.pytorch.org/whl/cu126/xformers-0.0.30-cp39-cp39-win_amd64.whl ; python_version=='3.9' and (sys_platform == 'win32')", - "xformers @ https://download.pytorch.org/whl/cu126/xformers-0.0.30-cp310-cp310-win_amd64.whl ; python_version=='3.10' and (sys_platform == 'win32')", - "xformers @ https://download.pytorch.org/whl/cu126/xformers-0.0.30-cp311-cp311-win_amd64.whl ; python_version=='3.11' and (sys_platform == 'win32')", - "xformers @ https://download.pytorch.org/whl/cu126/xformers-0.0.30-cp312-cp312-win_amd64.whl ; python_version=='3.12' and (sys_platform == 'win32')", + "xformers==0.0.30 ; ('linux' in sys_platform or sys_platform == 'win32') and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", ] cu128onlytorch270 = [ - "xformers @ https://download.pytorch.org/whl/cu128/xformers-0.0.30-cp39-cp39-manylinux_2_28_x86_64.whl ; python_version=='3.9' and ('linux' in sys_platform)", - "xformers @ https://download.pytorch.org/whl/cu128/xformers-0.0.30-cp310-cp310-manylinux_2_28_x86_64.whl ; python_version=='3.10' and ('linux' in sys_platform)", - "xformers @ https://download.pytorch.org/whl/cu128/xformers-0.0.30-cp311-cp311-manylinux_2_28_x86_64.whl ; python_version=='3.11' and ('linux' in sys_platform)", - "xformers @ https://download.pytorch.org/whl/cu128/xformers-0.0.30-cp312-cp312-manylinux_2_28_x86_64.whl ; python_version=='3.12' and ('linux' in sys_platform)", - "xformers @ https://download.pytorch.org/whl/cu128/xformers-0.0.30-cp39-cp39-win_amd64.whl ; python_version=='3.9' and (sys_platform == 'win32')", - "xformers @ https://download.pytorch.org/whl/cu128/xformers-0.0.30-cp310-cp310-win_amd64.whl ; python_version=='3.10' and (sys_platform == 'win32')", - "xformers @ https://download.pytorch.org/whl/cu128/xformers-0.0.30-cp311-cp311-win_amd64.whl ; python_version=='3.11' and (sys_platform == 'win32')", - "xformers @ https://download.pytorch.org/whl/cu128/xformers-0.0.30-cp312-cp312-win_amd64.whl ; python_version=='3.12' and (sys_platform == 'win32')", + "xformers==0.0.30 ; ('linux' in sys_platform or sys_platform == 'win32') and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", ] cu118onlytorch271 = [ - "xformers @ https://download.pytorch.org/whl/cu118/xformers-0.0.31.post1-cp39-abi3-manylinux_2_28_x86_64.whl ; ('linux' in sys_platform)", - "xformers @ https://download.pytorch.org/whl/cu118/xformers-0.0.31.post1-cp39-abi3-win_amd64.whl ; (sys_platform == 'win32')", + "xformers==0.0.31.post1 ; ('linux' in sys_platform or sys_platform == 'win32') and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", ] cu126onlytorch271 = [ - "xformers @ https://download.pytorch.org/whl/cu126/xformers-0.0.31.post1-cp39-abi3-manylinux_2_28_x86_64.whl ; ('linux' in sys_platform)", - "xformers @ https://download.pytorch.org/whl/cu126/xformers-0.0.31.post1-cp39-abi3-win_amd64.whl ; (sys_platform == 'win32')", + "xformers==0.0.31.post1 ; ('linux' in sys_platform or sys_platform == 'win32') and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", ] cu128onlytorch271 = [ - "xformers @ https://download.pytorch.org/whl/cu128/xformers-0.0.31.post1-cp39-abi3-manylinux_2_28_x86_64.whl ; ('linux' in sys_platform)", - "xformers @ https://download.pytorch.org/whl/cu128/xformers-0.0.31.post1-cp39-abi3-win_amd64.whl ; (sys_platform == 'win32')", + "xformers==0.0.31.post1 ; ('linux' in sys_platform or sys_platform == 'win32') and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", ] cu118onlytorch280 = [ - "xformers @ https://download.pytorch.org/whl/cu126/xformers-0.0.32.post2-cp39-abi3-manylinux_2_28_x86_64.whl ; ('linux' in sys_platform)", - "xformers @ https://download.pytorch.org/whl/cu126/xformers-0.0.32.post2-cp39-abi3-win_amd64.whl ; (sys_platform == 'win32')", + "xformers==0.0.32.post2 ; ('linux' in sys_platform or sys_platform == 'win32') and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", ] cu126onlytorch280 = [ - "xformers @ https://download.pytorch.org/whl/cu128/xformers-0.0.32.post2-cp39-abi3-manylinux_2_28_x86_64.whl ; ('linux' in sys_platform)", - "xformers @ https://download.pytorch.org/whl/cu128/xformers-0.0.32.post2-cp39-abi3-win_amd64.whl ; (sys_platform == 'win32')", + "xformers==0.0.32.post2 ; ('linux' in sys_platform or sys_platform == 'win32') and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", ] cu128onlytorch280 = [ - "xformers @ https://download.pytorch.org/whl/cu129/xformers-0.0.32.post2-cp39-abi3-manylinux_2_28_x86_64.whl ; ('linux' in sys_platform)", - "xformers @ https://download.pytorch.org/whl/cu129/xformers-0.0.32.post2-cp39-abi3-win_amd64.whl ; (sys_platform == 'win32')", + "xformers==0.0.32.post2 ; ('linux' in sys_platform or sys_platform == 'win32') and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", ] cu130onlytorch280 = [ ] cu126onlytorch290 = [ - "xformers @ https://download.pytorch.org/whl/cu126/xformers-0.0.33.post1-cp39-abi3-manylinux_2_28_x86_64.whl ; ('linux' in sys_platform)", - "xformers @ https://download.pytorch.org/whl/cu126/xformers-0.0.33.post1-cp39-abi3-win_amd64.whl ; (sys_platform == 'win32')", + "xformers==0.0.33.post1 ; ('linux' in sys_platform or sys_platform == 'win32') and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", ] cu128onlytorch290 = [ - "xformers @ https://download.pytorch.org/whl/cu128/xformers-0.0.33.post1-cp39-abi3-manylinux_2_28_x86_64.whl ; ('linux' in sys_platform)", - "xformers @ https://download.pytorch.org/whl/cu128/xformers-0.0.33.post1-cp39-abi3-win_amd64.whl ; (sys_platform == 'win32')", + "xformers==0.0.33.post1 ; ('linux' in sys_platform or sys_platform == 'win32') and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", ] cu130onlytorch290 = [ - "xformers @ https://download.pytorch.org/whl/cu130/xformers-0.0.33.post1-cp39-abi3-manylinux_2_28_x86_64.whl ; ('linux' in sys_platform)", - "xformers @ https://download.pytorch.org/whl/cu130/xformers-0.0.33.post1-cp39-abi3-win_amd64.whl ; (sys_platform == 'win32')", + "xformers==0.0.33.post1 ; ('linux' in sys_platform or sys_platform == 'win32') and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", ] cu126onlytorch291 = [ - "xformers @ https://download.pytorch.org/whl/cu126/xformers-0.0.33.post2-cp39-abi3-manylinux_2_28_x86_64.whl ; ('linux' in sys_platform)", - "xformers @ https://download.pytorch.org/whl/cu126/xformers-0.0.33.post2-cp39-abi3-win_amd64.whl ; (sys_platform == 'win32')", + "xformers==0.0.33.post2 ; ('linux' in sys_platform or sys_platform == 'win32') and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", ] cu128onlytorch291 = [ - "xformers @ https://download.pytorch.org/whl/cu128/xformers-0.0.33.post2-cp39-abi3-manylinux_2_28_x86_64.whl ; ('linux' in sys_platform)", - "xformers @ https://download.pytorch.org/whl/cu128/xformers-0.0.33.post2-cp39-abi3-win_amd64.whl ; (sys_platform == 'win32')", + "xformers==0.0.33.post2 ; ('linux' in sys_platform or sys_platform == 'win32') and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", ] cu130onlytorch291 = [ - "xformers @ https://download.pytorch.org/whl/cu130/xformers-0.0.33.post2-cp39-abi3-manylinux_2_28_x86_64.whl ; ('linux' in sys_platform)", - "xformers @ https://download.pytorch.org/whl/cu130/xformers-0.0.33.post2-cp39-abi3-win_amd64.whl ; (sys_platform == 'win32')", + "xformers==0.0.33.post2 ; ('linux' in sys_platform or sys_platform == 'win32') and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", ] cu126onlytorch2100 = [ - "xformers @ https://download.pytorch.org/whl/cu126/xformers-0.0.34-cp39-abi3-manylinux_2_28_x86_64.whl ; ('linux' in sys_platform)", - "xformers @ https://download.pytorch.org/whl/cu126/xformers-0.0.34-cp39-abi3-win_amd64.whl ; (sys_platform == 'win32')", + "xformers==0.0.34 ; ('linux' in sys_platform or sys_platform == 'win32') and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", ] cu128onlytorch2100 = [ - "xformers @ https://download.pytorch.org/whl/cu128/xformers-0.0.34-cp39-abi3-manylinux_2_28_x86_64.whl ; ('linux' in sys_platform)", - "xformers @ https://download.pytorch.org/whl/cu128/xformers-0.0.34-cp39-abi3-win_amd64.whl ; (sys_platform == 'win32')", + "xformers==0.0.34 ; ('linux' in sys_platform or sys_platform == 'win32') and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", ] cu130onlytorch2100 = [ - "xformers @ https://download.pytorch.org/whl/cu130/xformers-0.0.34-cp39-abi3-manylinux_2_28_x86_64.whl ; ('linux' in sys_platform)", - "xformers @ https://download.pytorch.org/whl/cu130/xformers-0.0.34-cp39-abi3-win_amd64.whl ; (sys_platform == 'win32')", + "xformers==0.0.34 ; ('linux' in sys_platform or sys_platform == 'win32') and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", ] cu118 = [ "unsloth[huggingface]", @@ -352,22 +257,22 @@ cu121 = [ ] cu118-torch211 = [ "unsloth[huggingface]", - "bitsandbytes==0.45.5", + "bitsandbytes>=0.45.5,!=0.46.0,!=0.48.0", "unsloth[cu118onlytorch211]", ] cu121-torch211 = [ "unsloth[huggingface]", - "bitsandbytes==0.45.5", + "bitsandbytes>=0.45.5,!=0.46.0,!=0.48.0", "unsloth[cu121onlytorch211]", ] cu118-torch212 = [ "unsloth[huggingface]", - "bitsandbytes==0.45.5", + "bitsandbytes>=0.45.5,!=0.46.0,!=0.48.0", "unsloth[cu118onlytorch212]", ] cu121-torch212 = [ "unsloth[huggingface]", - "bitsandbytes==0.45.5", + "bitsandbytes>=0.45.5,!=0.46.0,!=0.48.0", "unsloth[cu121onlytorch212]", ] cu118-torch220 = [ @@ -437,17 +342,17 @@ cu124-torch251 = [ ] cu118-torch260 = [ "unsloth[huggingface]", - "bitsandbytes>=0.45.5,!=0.46.0,!=0.48.0", + "bitsandbytes>=0.45.1", "unsloth[cu118onlytorch260]", ] cu124-torch260 = [ "unsloth[huggingface]", - "bitsandbytes>=0.45.5,!=0.46.0,!=0.48.0", + "bitsandbytes>=0.45.1", "unsloth[cu124onlytorch260]", ] cu126-torch260 = [ "unsloth[huggingface]", - "bitsandbytes>=0.45.5,!=0.46.0,!=0.48.0", + "bitsandbytes>=0.45.1", "unsloth[cu126onlytorch260]", ] cu118-torch270 = [ @@ -557,16 +462,19 @@ conda = [ ] colab-torch211 = [ "unsloth[huggingface]", - "bitsandbytes==0.45.5", + "bitsandbytes>=0.45.5,!=0.46.0,!=0.48.0", "unsloth[cu121onlytorch211]", ] +flashattention = [ + "packaging ; ('linux' in sys_platform)", + "ninja ; ('linux' in sys_platform)", + "flash-attn>=2.6.3 ; ('linux' in sys_platform)", +] colab-ampere-torch211 = [ "unsloth[huggingface]", - "bitsandbytes==0.45.5", + "bitsandbytes>=0.45.5,!=0.46.0,!=0.48.0", "unsloth[cu121onlytorch211]", - "packaging", - "ninja", - "flash-attn>=2.6.3 ; ('linux' in sys_platform)", + "unsloth[flashattention]", ] colab-torch220 = [ "unsloth[huggingface]", @@ -577,9 +485,7 @@ colab-ampere-torch220 = [ "unsloth[huggingface]", "bitsandbytes>=0.45.5,!=0.46.0,!=0.48.0", "unsloth[cu121onlytorch220]", - "packaging", - "ninja", - "flash-attn>=2.6.3 ; ('linux' in sys_platform)", + "unsloth[flashattention]", ] colab-new = [ "unsloth_zoo>=2026.6.4", @@ -598,6 +504,10 @@ colab-new = [ "bitsandbytes>=0.45.5,!=0.46.0,!=0.48.0", "unsloth[triton]", "sentence-transformers", + "typer", + "pydantic", + "pyyaml", + "nest-asyncio", ] colab-no-deps = [ "accelerate>=0.34.1", @@ -610,11 +520,6 @@ colab-no-deps = [ colab = [ "unsloth[cu121]", ] -flashattention = [ - "packaging ; ('linux' in sys_platform)", - "ninja ; ('linux' in sys_platform)", - "flash-attn>=2.6.3 ; ('linux' in sys_platform)", -] colab-ampere = [ "unsloth[colab-ampere-torch220]", "unsloth[flashattention]", @@ -633,13 +538,13 @@ cu121-ampere = [ ] cu118-ampere-torch211 = [ "unsloth[huggingface]", - "bitsandbytes==0.45.5", + "bitsandbytes>=0.45.5,!=0.46.0,!=0.48.0", "unsloth[cu118onlytorch211]", "unsloth[flashattention]", ] cu121-ampere-torch211 = [ "unsloth[huggingface]", - "bitsandbytes==0.45.5", + "bitsandbytes>=0.45.5,!=0.46.0,!=0.48.0", "unsloth[cu121onlytorch211]", "unsloth[flashattention]", ] @@ -723,19 +628,19 @@ cu124-ampere-torch251 = [ ] cu118-ampere-torch260 = [ "unsloth[huggingface]", - "bitsandbytes>=0.45.5,!=0.46.0,!=0.48.0", + "bitsandbytes>=0.45.1", "unsloth[cu118onlytorch260]", "unsloth[flashattention]", ] cu124-ampere-torch260 = [ "unsloth[huggingface]", - "bitsandbytes>=0.45.5,!=0.46.0,!=0.48.0", + "bitsandbytes>=0.45.1", "unsloth[cu124onlytorch260]", "unsloth[flashattention]", ] cu126-ampere-torch260 = [ "unsloth[huggingface]", - "bitsandbytes>=0.45.5,!=0.46.0,!=0.48.0", + "bitsandbytes>=0.45.1", "unsloth[cu126onlytorch260]", "unsloth[flashattention]", ] @@ -797,7 +702,6 @@ cu130-ampere-torch280 = [ "unsloth[huggingface]", "bitsandbytes>=0.45.5,!=0.46.0,!=0.48.0", "unsloth[cu130onlytorch280]", - "unsloth[flashattention]", ] cu126-ampere-torch290 = [ "unsloth[huggingface]", @@ -844,464 +748,6 @@ cu130-ampere-torch2100 = [ "bitsandbytes>=0.45.5,!=0.46.0,!=0.48.0", "unsloth[cu130onlytorch2100]", ] -flashattentiontorch260abiFALSEcu12x = [ - "flash-attn @ https://github.com/Dao-AILab/flash-attention/releases/download/v2.7.4.post1/flash_attn-2.7.4.post1+cu12torch2.6cxx11abiFALSE-cp39-cp39-linux_x86_64.whl ; ('linux' in sys_platform) and python_version == '3.9'", - "flash-attn @ https://github.com/Dao-AILab/flash-attention/releases/download/v2.7.4.post1/flash_attn-2.7.4.post1+cu12torch2.6cxx11abiFALSE-cp310-cp310-linux_x86_64.whl ; ('linux' in sys_platform) and python_version == '3.10'", - "flash-attn @ https://github.com/Dao-AILab/flash-attention/releases/download/v2.7.4.post1/flash_attn-2.7.4.post1+cu12torch2.6cxx11abiFALSE-cp311-cp311-linux_x86_64.whl ; ('linux' in sys_platform) and python_version == '3.11'", - "flash-attn @ https://github.com/Dao-AILab/flash-attention/releases/download/v2.7.4.post1/flash_attn-2.7.4.post1+cu12torch2.6cxx11abiFALSE-cp312-cp312-linux_x86_64.whl ; ('linux' in sys_platform) and python_version == '3.12'", - "flash-attn @ https://github.com/Dao-AILab/flash-attention/releases/download/v2.7.4.post1/flash_attn-2.7.4.post1+cu12torch2.6cxx11abiFALSE-cp313-cp313-linux_x86_64.whl ; ('linux' in sys_platform) and python_version == '3.13'", -] -flashattentiontorch260abiTRUEcu12x = [ - "flash-attn @ https://github.com/Dao-AILab/flash-attention/releases/download/v2.7.4.post1/flash_attn-2.7.4.post1+cu12torch2.6cxx11abiTRUE-cp39-cp39-linux_x86_64.whl ; ('linux' in sys_platform) and python_version == '3.9'", - "flash-attn @ https://github.com/Dao-AILab/flash-attention/releases/download/v2.7.4.post1/flash_attn-2.7.4.post1+cu12torch2.6cxx11abiTRUE-cp310-cp310-linux_x86_64.whl ; ('linux' in sys_platform) and python_version == '3.10'", - "flash-attn @ https://github.com/Dao-AILab/flash-attention/releases/download/v2.7.4.post1/flash_attn-2.7.4.post1+cu12torch2.6cxx11abiTRUE-cp311-cp311-linux_x86_64.whl ; ('linux' in sys_platform) and python_version == '3.11'", - "flash-attn @ https://github.com/Dao-AILab/flash-attention/releases/download/v2.7.4.post1/flash_attn-2.7.4.post1+cu12torch2.6cxx11abiTRUE-cp312-cp312-linux_x86_64.whl ; ('linux' in sys_platform) and python_version == '3.12'", - "flash-attn @ https://github.com/Dao-AILab/flash-attention/releases/download/v2.7.4.post1/flash_attn-2.7.4.post1+cu12torch2.6cxx11abiTRUE-cp313-cp313-linux_x86_64.whl ; ('linux' in sys_platform) and python_version == '3.13'", -] -flashattentiontorch250abiFALSEcu12x = [ - "flash-attn @ https://github.com/Dao-AILab/flash-attention/releases/download/v2.7.4.post1/flash_attn-2.7.4.post1+cu12torch2.5cxx11abiFALSE-cp39-cp39-linux_x86_64.whl ; ('linux' in sys_platform) and python_version == '3.9'", - "flash-attn @ https://github.com/Dao-AILab/flash-attention/releases/download/v2.7.4.post1/flash_attn-2.7.4.post1+cu12torch2.5cxx11abiFALSE-cp310-cp310-linux_x86_64.whl ; ('linux' in sys_platform) and python_version == '3.10'", - "flash-attn @ https://github.com/Dao-AILab/flash-attention/releases/download/v2.7.4.post1/flash_attn-2.7.4.post1+cu12torch2.5cxx11abiFALSE-cp311-cp311-linux_x86_64.whl ; ('linux' in sys_platform) and python_version == '3.11'", - "flash-attn @ https://github.com/Dao-AILab/flash-attention/releases/download/v2.7.4.post1/flash_attn-2.7.4.post1+cu12torch2.5cxx11abiFALSE-cp312-cp312-linux_x86_64.whl ; ('linux' in sys_platform) and python_version == '3.12'", - "flash-attn @ https://github.com/Dao-AILab/flash-attention/releases/download/v2.7.4.post1/flash_attn-2.7.4.post1+cu12torch2.5cxx11abiFALSE-cp313-cp313-linux_x86_64.whl ; ('linux' in sys_platform) and python_version == '3.13'", -] -flashattentiontorch250abiTRUEcu12x = [ - "flash-attn @ https://github.com/Dao-AILab/flash-attention/releases/download/v2.7.4.post1/flash_attn-2.7.4.post1+cu12torch2.5cxx11abiTRUE-cp39-cp39-linux_x86_64.whl ; ('linux' in sys_platform) and python_version == '3.9'", - "flash-attn @ https://github.com/Dao-AILab/flash-attention/releases/download/v2.7.4.post1/flash_attn-2.7.4.post1+cu12torch2.5cxx11abiTRUE-cp310-cp310-linux_x86_64.whl ; ('linux' in sys_platform) and python_version == '3.10'", - "flash-attn @ https://github.com/Dao-AILab/flash-attention/releases/download/v2.7.4.post1/flash_attn-2.7.4.post1+cu12torch2.5cxx11abiTRUE-cp311-cp311-linux_x86_64.whl ; ('linux' in sys_platform) and python_version == '3.11'", - "flash-attn @ https://github.com/Dao-AILab/flash-attention/releases/download/v2.7.4.post1/flash_attn-2.7.4.post1+cu12torch2.5cxx11abiTRUE-cp312-cp312-linux_x86_64.whl ; ('linux' in sys_platform) and python_version == '3.12'", - "flash-attn @ https://github.com/Dao-AILab/flash-attention/releases/download/v2.7.4.post1/flash_attn-2.7.4.post1+cu12torch2.5cxx11abiTRUE-cp313-cp313-linux_x86_64.whl ; ('linux' in sys_platform) and python_version == '3.13'", -] -flashattentiontorch240abiFALSEcu12x = [ - "flash-attn @ https://github.com/Dao-AILab/flash-attention/releases/download/v2.7.4.post1/flash_attn-2.7.4.post1+cu12torch2.4cxx11abiFALSE-cp39-cp39-linux_x86_64.whl ; ('linux' in sys_platform) and python_version == '3.9'", - "flash-attn @ https://github.com/Dao-AILab/flash-attention/releases/download/v2.7.4.post1/flash_attn-2.7.4.post1+cu12torch2.4cxx11abiFALSE-cp310-cp310-linux_x86_64.whl ; ('linux' in sys_platform) and python_version == '3.10'", - "flash-attn @ https://github.com/Dao-AILab/flash-attention/releases/download/v2.7.4.post1/flash_attn-2.7.4.post1+cu12torch2.4cxx11abiFALSE-cp311-cp311-linux_x86_64.whl ; ('linux' in sys_platform) and python_version == '3.11'", - "flash-attn @ https://github.com/Dao-AILab/flash-attention/releases/download/v2.7.4.post1/flash_attn-2.7.4.post1+cu12torch2.4cxx11abiFALSE-cp312-cp312-linux_x86_64.whl ; ('linux' in sys_platform) and python_version == '3.12'", - "flash-attn @ https://github.com/Dao-AILab/flash-attention/releases/download/v2.7.4.post1/flash_attn-2.7.4.post1+cu12torch2.4cxx11abiFALSE-cp313-cp313-linux_x86_64.whl ; ('linux' in sys_platform) and python_version == '3.13'", -] -flashattentiontorch240abiTRUEcu12x = [ - "flash-attn @ https://github.com/Dao-AILab/flash-attention/releases/download/v2.7.4.post1/flash_attn-2.7.4.post1+cu12torch2.4cxx11abiTRUE-cp39-cp39-linux_x86_64.whl ; ('linux' in sys_platform) and python_version == '3.9'", - "flash-attn @ https://github.com/Dao-AILab/flash-attention/releases/download/v2.7.4.post1/flash_attn-2.7.4.post1+cu12torch2.4cxx11abiTRUE-cp310-cp310-linux_x86_64.whl ; ('linux' in sys_platform) and python_version == '3.10'", - "flash-attn @ https://github.com/Dao-AILab/flash-attention/releases/download/v2.7.4.post1/flash_attn-2.7.4.post1+cu12torch2.4cxx11abiTRUE-cp311-cp311-linux_x86_64.whl ; ('linux' in sys_platform) and python_version == '3.11'", - "flash-attn @ https://github.com/Dao-AILab/flash-attention/releases/download/v2.7.4.post1/flash_attn-2.7.4.post1+cu12torch2.4cxx11abiTRUE-cp312-cp312-linux_x86_64.whl ; ('linux' in sys_platform) and python_version == '3.12'", - "flash-attn @ https://github.com/Dao-AILab/flash-attention/releases/download/v2.7.4.post1/flash_attn-2.7.4.post1+cu12torch2.4cxx11abiTRUE-cp313-cp313-linux_x86_64.whl ; ('linux' in sys_platform) and python_version == '3.13'", -] -intelgputorch260 = [ - "unsloth_zoo[intelgpu]", - "unsloth[huggingfacenotorch]", - - "pytorch_triton_xpu @ https://download.pytorch.org/whl/pytorch_triton_xpu-3.2.0-cp39-cp39-linux_x86_64.whl#sha256=147607f190a7d7aa24ba454def5977fbbfec792fdae18e4ed278cfec29b69271 ; ('linux' in sys_platform) and python_version == '3.9' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", - "pytorch_triton_xpu @ https://download.pytorch.org/whl/pytorch_triton_xpu-3.2.0-cp310-cp310-linux_x86_64.whl#sha256=23aa423fa1542afc34f67eb3ba8ef20060f6d1b3a4697eaeab22b11c92b30f2b ; ('linux' in sys_platform) and python_version == '3.10' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", - "pytorch_triton_xpu @ https://download.pytorch.org/whl/pytorch_triton_xpu-3.2.0-cp311-cp311-linux_x86_64.whl#sha256=bcfa995229bbfd9ffd8d6c8d9f6428d393e876fa6e23ee3c20e3c0d73ca75ca5 ; ('linux' in sys_platform) and python_version == '3.11' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", - "pytorch_triton_xpu @ https://download.pytorch.org/whl/pytorch_triton_xpu-3.2.0-cp312-cp312-linux_x86_64.whl#sha256=bd340903d03470708df3442438acb8b7e08087ab9e61fbe349b2872bf9257ab0 ; ('linux' in sys_platform) and python_version == '3.12' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", - "pytorch_triton_xpu @ https://download.pytorch.org/whl/pytorch_triton_xpu-3.2.0-cp313-cp313-linux_x86_64.whl#sha256=814dccc8a07159e6eca74bed70091bc8fea2d9dd87b0d91845f9f38cde62f01c ; ('linux' in sys_platform) and python_version == '3.13' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", - - "bitsandbytes @ https://github.com/bitsandbytes-foundation/bitsandbytes/releases/download/continuous-release_main/bitsandbytes-1.33.7.preview-py3-none-manylinux_2_24_x86_64.whl ; ('linux' in sys_platform) and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", - "bitsandbytes @ https://github.com/bitsandbytes-foundation/bitsandbytes/releases/download/continuous-release_main/bitsandbytes-1.33.7.preview-py3-none-win_amd64.whl ; (sys_platform == 'win32') and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", - - "torch @ https://download.pytorch.org/whl/xpu/torch-2.6.0%2Bxpu-cp39-cp39-linux_x86_64.whl#sha256=6a8adf6dc4c089406e8b3a7e58ab57a463bddf9b07130d2576e76eced43e92af ; ('linux' in sys_platform) and python_version == '3.9' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", - "torch @ https://download.pytorch.org/whl/xpu/torch-2.6.0%2Bxpu-cp310-cp310-linux_x86_64.whl#sha256=ff4561cbf07c83bbccaa0f6e9bb0e6dcf721bacd53c9c43c4eb0e7331b4792f9 ; ('linux' in sys_platform) and python_version == '3.10' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", - "torch @ https://download.pytorch.org/whl/xpu/torch-2.6.0%2Bxpu-cp311-cp311-linux_x86_64.whl#sha256=12005f66b810ddd3ab93f86c4522bcfdd412cbd27fc9d189b661ff7509bc5e8a ; ('linux' in sys_platform) and python_version == '3.11' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", - "torch @ https://download.pytorch.org/whl/xpu/torch-2.6.0%2Bxpu-cp312-cp312-linux_x86_64.whl#sha256=c4c5c67625cdacf35765c2b94e61fe166e3c3f4a14521b1212a59ad1b3eb0f2e ; ('linux' in sys_platform) and python_version == '3.12' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", - "torch @ https://download.pytorch.org/whl/xpu/torch-2.6.0%2Bxpu-cp313-cp313-linux_x86_64.whl#sha256=e6864f7a60a5ecc43d5d38f59a16e5dd132384f73dfd3a697f74944026038f7b ; ('linux' in sys_platform) and python_version == '3.13' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", -] -intel-gpu-torch260 = [ - "unsloth[intelgputorch260]" -] -intelgputorch270 = [ - "unsloth_zoo[intelgpu]", - "unsloth[huggingfacenotorch]", - - "pytorch_triton_xpu @ https://download.pytorch.org/whl/pytorch_triton_xpu-3.3.0-cp39-cp39-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl#sha256=749a7098492c6a27b356c97149a4a62973b953eae60bc1b6259260974f344913 ; ('linux' in sys_platform) and python_version == '3.9' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", - "pytorch_triton_xpu @ https://download.pytorch.org/whl/pytorch_triton_xpu-3.3.0-cp310-cp310-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl#sha256=44362e80abd752471a08341093321955b066daa2cfb4810e73b8e3b240850f93 ; ('linux' in sys_platform) and python_version == '3.10' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", - "pytorch_triton_xpu @ https://download.pytorch.org/whl/pytorch_triton_xpu-3.3.0-cp311-cp311-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl#sha256=faa6b8c945a837a080f641bc8ccc77a98fa66980dcd7e62e715fd853737343fd ; ('linux' in sys_platform) and python_version == '3.11' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", - "pytorch_triton_xpu @ https://download.pytorch.org/whl/pytorch_triton_xpu-3.3.0-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl#sha256=40f6fb65b345dc9a61813abe7ac9a585f2c9808f414d140cc2a5f11f53ee063c ; ('linux' in sys_platform) and python_version == '3.12' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", - "pytorch_triton_xpu @ https://download.pytorch.org/whl/pytorch_triton_xpu-3.3.0-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl#sha256=b22b4c02ec71b4bfc862ae3cdfd2871dc0b05d2b1802f5db2196e0f897d581e9 ; ('linux' in sys_platform) and python_version == '3.13' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", - "pytorch_triton_xpu @ https://download.pytorch.org/whl/pytorch_triton_xpu-3.3.0-cp39-cp39-win_amd64.whl#sha256=d4b738d7fa5100c1bd766f91614962828a4810eb57b4df92cd5214a83505a752 ; sys_platform == 'win32' and python_version == '3.9' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", - "pytorch_triton_xpu @ https://download.pytorch.org/whl/pytorch_triton_xpu-3.3.0-cp310-cp310-win_amd64.whl#sha256=143fe8a64d807bcdb7d81bbc062816add325570aa160448454ab6ded4a0a17a1 ; sys_platform == 'win32' and python_version == '3.10' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", - "pytorch_triton_xpu @ https://download.pytorch.org/whl/pytorch_triton_xpu-3.3.0-cp311-cp311-win_amd64.whl#sha256=a8025459ff325d6e3532eb5cf72519db1b178155e7d60aff6c56beb5968fc758 ; sys_platform == 'win32' and python_version == '3.11' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", - "pytorch_triton_xpu @ https://download.pytorch.org/whl/pytorch_triton_xpu-3.3.0-cp312-cp312-win_amd64.whl#sha256=0dd07e6d5b872e42e48f5ee140e609d4554ca3cc509d5bf509ac232267cf358e ; sys_platform == 'win32' and python_version == '3.12' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", - "pytorch_triton_xpu @ https://download.pytorch.org/whl/pytorch_triton_xpu-3.3.0-cp313-cp313-win_amd64.whl#sha256=a936a18182d8e065a9933afc9a3ebbffadd38604969f87c493831214539fc027 ; sys_platform == 'win32' and python_version == '3.13' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", - - "bitsandbytes @ https://github.com/bitsandbytes-foundation/bitsandbytes/releases/download/continuous-release_main/bitsandbytes-1.33.7.preview-py3-none-manylinux_2_24_x86_64.whl ; ('linux' in sys_platform) and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", - "bitsandbytes @ https://github.com/bitsandbytes-foundation/bitsandbytes/releases/download/continuous-release_main/bitsandbytes-1.33.7.preview-py3-none-win_amd64.whl ; (sys_platform == 'win32') and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", - - "torch @ https://download.pytorch.org/whl/xpu/torch-2.7.0%2Bxpu-cp39-cp39-linux_x86_64.whl#sha256=f8ee75e50fcbb37ed5b498299ca2264da99ab278a93fae2358e921e4a6e28273 ; ('linux' in sys_platform) and python_version == '3.9' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", - "torch @ https://download.pytorch.org/whl/xpu/torch-2.7.0%2Bxpu-cp310-cp310-linux_x86_64.whl#sha256=d6fdc342961d98fdcd9d03dfd491a3208bb5f7fbb435841f8f72ce9fdcd2d026 ; ('linux' in sys_platform) and python_version == '3.10' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", - "torch @ https://download.pytorch.org/whl/xpu/torch-2.7.0%2Bxpu-cp311-cp311-linux_x86_64.whl#sha256=74d07f9357df5cf2bf223ad3c84de16346bfaa0504f988fdd5590d3e177e5e86 ; ('linux' in sys_platform) and python_version == '3.11' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", - "torch @ https://download.pytorch.org/whl/xpu/torch-2.7.0%2Bxpu-cp312-cp312-linux_x86_64.whl#sha256=c806d44aa2ca5d225629f6fbc6c994d5deaac2d2cde449195bc8e3522ddd219a ; ('linux' in sys_platform) and python_version == '3.12' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", - "torch @ https://download.pytorch.org/whl/xpu/torch-2.7.0%2Bxpu-cp313-cp313-linux_x86_64.whl#sha256=25d8277b7f01d42e2e014ccbab57a2692b6ec4eff8dcf894eda1b297407cf97a ; ('linux' in sys_platform) and python_version == '3.13' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", - "torch @ https://download.pytorch.org/whl/xpu/torch-2.7.0%2Bxpu-cp39-cp39-win_amd64.whl#sha256=046e85125266ae69c1a0d083e6c092f947ab4b6b41532c16bafe40dbced845df ; sys_platform == 'win32' and python_version == '3.9' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", - "torch @ https://download.pytorch.org/whl/xpu/torch-2.7.0%2Bxpu-cp310-cp310-win_amd64.whl#sha256=9ebaeffb82b0b3e39b6030927d3ebe0eb62a0e9045a3b2d7b0a9e7b15222c0db ; sys_platform == 'win32' and python_version == '3.10' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", - "torch @ https://download.pytorch.org/whl/xpu/torch-2.7.0%2Bxpu-cp311-cp311-win_amd64.whl#sha256=356ba66cee127e7e2c942880bd50e03768306a4ea08d358a0f29c6eebfc4bc81 ; sys_platform == 'win32' and python_version == '3.11' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", - "torch @ https://download.pytorch.org/whl/xpu/torch-2.7.0%2Bxpu-cp312-cp312-win_amd64.whl#sha256=94739e665d9b4d5cd7af5f517cb6103f6f9fb421c095184609653a24524040f5 ; sys_platform == 'win32' and python_version == '3.12' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", - "torch @ https://download.pytorch.org/whl/xpu/torch-2.7.0%2Bxpu-cp313-cp313-win_amd64.whl#sha256=31df3cb674918e89bc8c532baa331dc84f4430e1f9c0ec379232db44cba78355 ; sys_platform == 'win32' and python_version == '3.13' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", -] -intel-gpu-torch270 = [ - "unsloth[intelgputorch270]" -] -intelgputorch280 = [ - "unsloth_zoo[intelgpu]", - "unsloth[huggingfacenotorch]", - - "pytorch_triton_xpu @ https://download.pytorch.org/whl/pytorch_triton_xpu-3.4.0-cp39-cp39-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl#sha256=ac4d8e33986b1c3c5e48151640539272b2187e83016985853111b46fb82c3c94 ; 'linux' in sys_platform and python_version == '3.9' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", - "pytorch_triton_xpu @ https://download.pytorch.org/whl/pytorch_triton_xpu-3.4.0-cp310-cp310-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl#sha256=999fef4c1f711092b9d3086525920545df490de476ecebe899ffc777019ae17f ; 'linux' in sys_platform and python_version == '3.10' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", - "pytorch_triton_xpu @ https://download.pytorch.org/whl/pytorch_triton_xpu-3.4.0-cp311-cp311-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl#sha256=57b09c8c492985ff6a27cd3a22b08e8f7b96b407bd8030967b6efbb9f63b80cf ; 'linux' in sys_platform and python_version == '3.11' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", - "pytorch_triton_xpu @ https://download.pytorch.org/whl/pytorch_triton_xpu-3.4.0-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl#sha256=df4bb3282bac9a3b90231700077110d8680b338416de03c2b7c6133c9b602649 ; 'linux' in sys_platform and python_version == '3.12' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", - "pytorch_triton_xpu @ https://download.pytorch.org/whl/pytorch_triton_xpu-3.4.0-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl#sha256=60da63c99ca827bdcb0df28e0298bf7d066dc607454c6d6176783cb4e79d838b ; 'linux' in sys_platform and python_version == '3.13' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", - "pytorch_triton_xpu @ https://download.pytorch.org/whl/pytorch_triton_xpu-3.4.0-cp39-cp39-win_amd64.whl#sha256=64aea8de349f3e2e0ebf4c24b011a8122531fdffda5776edaef45829cc241cf8 ; sys_platform == 'win32' and python_version == '3.9' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", - "pytorch_triton_xpu @ https://download.pytorch.org/whl/pytorch_triton_xpu-3.4.0-cp310-cp310-win_amd64.whl#sha256=ae573d255b257fdbed319a3440dc9d0a721e31160ab7f6eba1b2226e6a409a1d ; sys_platform == 'win32' and python_version == '3.10' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", - "pytorch_triton_xpu @ https://download.pytorch.org/whl/pytorch_triton_xpu-3.4.0-cp311-cp311-win_amd64.whl#sha256=8e0ea4558e5776d8ddab0264310be9b26aee5641bcac0da023537556d4317b86 ; sys_platform == 'win32' and python_version == '3.11' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", - "pytorch_triton_xpu @ https://download.pytorch.org/whl/pytorch_triton_xpu-3.4.0-cp312-cp312-win_amd64.whl#sha256=4090dde07a4fffc34aaf855701a9db28e9fccb57b368ade520f1a0f8e811c878 ; sys_platform == 'win32' and python_version == '3.12' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", - "pytorch_triton_xpu @ https://download.pytorch.org/whl/pytorch_triton_xpu-3.4.0-cp313-cp313-win_amd64.whl#sha256=a33d0888f3c8df028a2d028842715837d0049524d6c06b9bb11869890a13601a ; sys_platform == 'win32' and python_version == '3.13' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", - - "torch @ https://download.pytorch.org/whl/xpu/torch-2.8.0%2Bxpu-cp39-cp39-linux_x86_64.whl ; 'linux' in sys_platform and python_version == '3.9' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", - "torch @ https://download.pytorch.org/whl/xpu/torch-2.8.0%2Bxpu-cp310-cp310-linux_x86_64.whl ; 'linux' in sys_platform and python_version == '3.10' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", - "torch @ https://download.pytorch.org/whl/xpu/torch-2.8.0%2Bxpu-cp311-cp311-linux_x86_64.whl ; 'linux' in sys_platform and python_version == '3.11' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", - "torch @ https://download.pytorch.org/whl/xpu/torch-2.8.0%2Bxpu-cp312-cp312-linux_x86_64.whl ; 'linux' in sys_platform and python_version == '3.12' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", - "torch @ https://download.pytorch.org/whl/xpu/torch-2.8.0%2Bxpu-cp313-cp313-linux_x86_64.whl ; 'linux' in sys_platform and python_version == '3.13' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", - "torch @ https://download.pytorch.org/whl/xpu/torch-2.8.0%2Bxpu-cp39-cp39-win_amd64.whl#sha256=f2f401276892428e4875cf1d8717c5cbab704b16fc594ccf23795e7b16549a99 ; sys_platform == 'win32' and python_version == '3.9' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", - "torch @ https://download.pytorch.org/whl/xpu/torch-2.8.0%2Bxpu-cp310-cp310-win_amd64.whl#sha256=125c60cd59d51b39581a7e9afcd4679bc3a6b8c1f9440b1bb502a23fdd60571e ; sys_platform == 'win32' and python_version == '3.10' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", - "torch @ https://download.pytorch.org/whl/xpu/torch-2.8.0%2Bxpu-cp311-cp311-win_amd64.whl#sha256=47f1a57258cd460e80b38b2ed6744e31587ab77a96b4215bf59546cb4bab5cc0 ; sys_platform == 'win32' and python_version == '3.11' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", - "torch @ https://download.pytorch.org/whl/xpu/torch-2.8.0%2Bxpu-cp312-cp312-win_amd64.whl#sha256=0937d8943c145a83d9bafc6f80ef28971167817f9eda26066d33f72caf8a6646 ; sys_platform == 'win32' and python_version == '3.12' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", - "torch @ https://download.pytorch.org/whl/xpu/torch-2.8.0%2Bxpu-cp313-cp313-win_amd64.whl#sha256=e034aab1d71760dc80a731531be43673ffe15e99033b82d24e40d2e6d41bd8bf ; sys_platform == 'win32' and python_version == '3.13' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", - - "bitsandbytes @ https://github.com/bitsandbytes-foundation/bitsandbytes/releases/download/continuous-release_main/bitsandbytes-1.33.7.preview-py3-none-manylinux_2_24_x86_64.whl ; ('linux' in sys_platform) and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", - "bitsandbytes @ https://github.com/bitsandbytes-foundation/bitsandbytes/releases/download/continuous-release_main/bitsandbytes-1.33.7.preview-py3-none-win_amd64.whl ; (sys_platform == 'win32') and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", - - "torchvision @ https://download.pytorch.org/whl/xpu/torchvision-0.23.0%2Bxpu-cp39-cp39-manylinux_2_28_x86_64.whl#sha256=6e981c192045fc249c008441179ff237bb00174d818b875b0475730b63f0eaca ; 'linux' in sys_platform and python_version == '3.9' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", - "torchvision @ https://download.pytorch.org/whl/xpu/torchvision-0.23.0%2Bxpu-cp310-cp310-manylinux_2_28_x86_64.whl#sha256=e5ba4805969277175ebfd59cc717093528cc6e3ada89ac2725fc7a3c1fee6169 ; 'linux' in sys_platform and python_version == '3.10' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", - "torchvision @ https://download.pytorch.org/whl/xpu/torchvision-0.23.0%2Bxpu-cp311-cp311-manylinux_2_28_x86_64.whl#sha256=74c39c144104416bc4c5ad8c26ab0c169dc5cc6be58059e01bc3665dd0ef676f ; 'linux' in sys_platform and python_version == '3.11' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", - "torchvision @ https://download.pytorch.org/whl/xpu/torchvision-0.23.0%2Bxpu-cp312-cp312-manylinux_2_28_x86_64.whl#sha256=0acec355b80c3899841184084f365df336c508602812e34a44007b8b60d53af4 ; 'linux' in sys_platform and python_version == '3.12' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", - "torchvision @ https://download.pytorch.org/whl/xpu/torchvision-0.23.0%2Bxpu-cp313-cp313-manylinux_2_28_x86_64.whl#sha256=e2109ae773dad27b98ca17681044b4f876563c37f2382b75de3a371399edcff8 ; 'linux' in sys_platform and python_version == '3.13' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", - "torchvision @ https://download.pytorch.org/whl/xpu/torchvision-0.23.0%2Bxpu-cp39-cp39-win_amd64.whl#sha256=5f7904e7048d414379bc8c1167260f1e84204f105db2d0a2f9c89e87ce1cf205 ; sys_platform == 'win32' and python_version == '3.9' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", - "torchvision @ https://download.pytorch.org/whl/xpu/torchvision-0.23.0%2Bxpu-cp310-cp310-win_amd64.whl#sha256=005fca5e658ca8e37adb63c1a021c84f5e56dfa6cf0d601d89cfe40b9473f79f ; sys_platform == 'win32' and python_version == '3.10' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", - "torchvision @ https://download.pytorch.org/whl/xpu/torchvision-0.23.0%2Bxpu-cp311-cp311-win_amd64.whl#sha256=c6d030f5361461550c0ff1339b5bca8585fc1e84fda2e64b6184e65a581e4f98 ; sys_platform == 'win32' and python_version == '3.11' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", - "torchvision @ https://download.pytorch.org/whl/xpu/torchvision-0.23.0%2Bxpu-cp312-cp312-win_amd64.whl#sha256=91aafd61864cdce27461cbec13ddbf28c1bc6494265a1e4b80131c64a3b7d18f ; sys_platform == 'win32' and python_version == '3.12' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", - "torchvision @ https://download.pytorch.org/whl/xpu/torchvision-0.23.0%2Bxpu-cp313-cp313-win_amd64.whl#sha256=71dc4a6421742ed1e7f585b04a100ad53615c341fbccfbc255aefb38ea9091da ; sys_platform == 'win32' and python_version == '3.13' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", -] -intel-gpu-torch280 = [ - "unsloth[intelgputorch280]" -] -intelgputorch290 = [ - "unsloth_zoo[intelgpu]", - "unsloth[huggingfacenotorch]", - - "pytorch_triton_xpu @ https://download.pytorch.org/whl/pytorch_triton_xpu-3.5.0-cp310-cp310-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl#sha256=c169a1de14c19673b17c751290d467fa282fc90fa5da4314b2e5cdab1f553146 ; platform_system == 'Linux' and python_version == '3.10' and platform_machine == 'x86_64'", - "pytorch_triton_xpu @ https://download.pytorch.org/whl/pytorch_triton_xpu-3.5.0-cp311-cp311-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl#sha256=013d9dd5d6479bd22983161f462e61c8dbe1d82e6730624a7a8d5945507eaa61 ; platform_system == 'Linux' and python_version == '3.11' and platform_machine == 'x86_64'", - "pytorch_triton_xpu @ https://download.pytorch.org/whl/pytorch_triton_xpu-3.5.0-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl#sha256=afc8cabfbf7ed51fd278d1e0f88d6afc157b0201bad4b99d681e4d542f9e66d4 ; platform_system == 'Linux' and python_version == '3.12' and platform_machine == 'x86_64'", - "pytorch_triton_xpu @ https://download.pytorch.org/whl/pytorch_triton_xpu-3.5.0-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl#sha256=0d24c1716088f2764d0d24c64227732195b6a42706c3c5fc89eeb4904bfa0818 ; platform_system == 'Linux' and python_version == '3.13' and platform_machine == 'x86_64'", - "pytorch_triton_xpu @ https://download.pytorch.org/whl/pytorch_triton_xpu-3.5.0-cp310-cp310-win_amd64.whl#sha256=c83ab007311d9cfb6e809ee5a4587d99a9eef4be720b90da4f1aaa68b45139a0 ; sys_platform == 'win32' and python_version == '3.10' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", - "pytorch_triton_xpu @ https://download.pytorch.org/whl/pytorch_triton_xpu-3.5.0-cp311-cp311-win_amd64.whl#sha256=debf75348da8e8c7166b4d4a9b91d1508bb8d6581e339f79f7604b2e6746bacd ; sys_platform == 'win32' and python_version == '3.11' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", - "pytorch_triton_xpu @ https://download.pytorch.org/whl/pytorch_triton_xpu-3.5.0-cp312-cp312-win_amd64.whl#sha256=97337a47425f1963a723475bd61037460e84ba01db4f87a1d662c3718ff6c47e ; sys_platform == 'win32' and python_version == '3.12' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", - "pytorch_triton_xpu @ https://download.pytorch.org/whl/pytorch_triton_xpu-3.5.0-cp313-cp313-win_amd64.whl#sha256=2caf8138695f6abb023ecd02031a2611ba1bf8fff2f19802567cb2fadefe9e87 ; sys_platform == 'win32' and python_version == '3.13' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", - - "torch @ https://download.pytorch.org/whl/xpu/torch-2.9.0%2Bxpu-cp310-cp310-linux_x86_64.whl#sha256=5afbe860ce991825a36b75706a523601087e414b77598ef0d9d3d565741c277d ; platform_system == 'Linux' and python_version == '3.10' and platform_machine == 'x86_64'", - "torch @ https://download.pytorch.org/whl/xpu/torch-2.9.0%2Bxpu-cp311-cp311-linux_x86_64.whl#sha256=607fe419c32d6e8e0556f745742e7cff1d0babce51f54be890e0c1422359c442 ; platform_system == 'Linux' and python_version == '3.11' and platform_machine == 'x86_64'", - "torch @ https://download.pytorch.org/whl/xpu/torch-2.9.0%2Bxpu-cp312-cp312-linux_x86_64.whl#sha256=376bae584d89980b8e59934d248c38d5fa3b7d4687a4df1a19f4bc1d23dcc8c1 ; platform_system == 'Linux' and python_version == '3.12' and platform_machine == 'x86_64'", - "torch @ https://download.pytorch.org/whl/xpu/torch-2.9.0%2Bxpu-cp313-cp313-linux_x86_64.whl#sha256=98d6a06dd7fb185874367b18bd609f05f16fdce4142a5980ca94461949965cd2 ; platform_system == 'Linux' and python_version == '3.13' and platform_machine == 'x86_64'", - "torch @ https://download.pytorch.org/whl/xpu/torch-2.9.0%2Bxpu-cp310-cp310-win_amd64.whl#sha256=47cc68f631f65bd9c84924d052cd04dec7531023caa85e80345e9c94611c887d ; sys_platform == 'win32' and python_version == '3.10' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", - "torch @ https://download.pytorch.org/whl/xpu/torch-2.9.0%2Bxpu-cp311-cp311-win_amd64.whl#sha256=d56c44ab4818aba57e5c7b628f422d014e0d507427170a771c5be85e308b0bc6 ; sys_platform == 'win32' and python_version == '3.11' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", - "torch @ https://download.pytorch.org/whl/xpu/torch-2.9.0%2Bxpu-cp312-cp312-win_amd64.whl#sha256=18cad93aaff76a01ce73aef6935ece7cfc03344b905592ec731446c44d44592b ; sys_platform == 'win32' and python_version == '3.12' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", - "torch @ https://download.pytorch.org/whl/xpu/torch-2.9.0%2Bxpu-cp313-cp313-win_amd64.whl#sha256=579929cdc10a76800ead41289cac191ea36d1b16f5f501d3fc25607d4375cd83 ; sys_platform == 'win32' and python_version == '3.13' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", - - "bitsandbytes @ https://github.com/bitsandbytes-foundation/bitsandbytes/releases/download/continuous-release_main/bitsandbytes-1.33.7.preview-py3-none-manylinux_2_24_x86_64.whl ; ('linux' in sys_platform) and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", - "bitsandbytes @ https://github.com/bitsandbytes-foundation/bitsandbytes/releases/download/continuous-release_main/bitsandbytes-1.33.7.preview-py3-none-win_amd64.whl ; (sys_platform == 'win32') and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", - - "torchvision @ https://download.pytorch.org/whl/xpu/torchvision-0.24.0%2Bxpu-cp310-cp310-manylinux_2_28_x86_64.whl#sha256=cbfae2b79b7549fd368c2462fc8e94f8f26cc450782ee72138e908077c09a519 ; platform_system == 'Linux' and python_version == '3.10' and platform_machine == 'x86_64'", - "torchvision @ https://download.pytorch.org/whl/xpu/torchvision-0.24.0%2Bxpu-cp311-cp311-manylinux_2_28_x86_64.whl#sha256=044fa36ef4b6b43edcd490b75c853fa4b3eb033c2bded29f8fbcf27734713c67 ; platform_system == 'Linux' and python_version == '3.11' and platform_machine == 'x86_64'", - "torchvision @ https://download.pytorch.org/whl/xpu/torchvision-0.24.0%2Bxpu-cp312-cp312-manylinux_2_28_x86_64.whl#sha256=4b91e4bec1d740a6211f02578a79888550b73f3a4e1383035f8f6d72f587212c ; platform_system == 'Linux' and python_version == '3.12' and platform_machine == 'x86_64'", - "torchvision @ https://download.pytorch.org/whl/xpu/torchvision-0.24.0%2Bxpu-cp313-cp313-manylinux_2_28_x86_64.whl#sha256=88239e73ca37254bec84f29cd5887e10ff712de7edbbda3fbb3609cd6190d99e ; platform_system == 'Linux' and python_version == '3.13' and platform_machine == 'x86_64'", - "torchvision @ https://download.pytorch.org/whl/xpu/torchvision-0.24.0%2Bxpu-cp310-cp310-win_amd64.whl#sha256=19c7da8ca767d593e13a88a12bb08d06e34a673f6f26c2f9c191d60e81c02953 ; sys_platform == 'win32' and python_version == '3.10' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", - "torchvision @ https://download.pytorch.org/whl/xpu/torchvision-0.24.0%2Bxpu-cp311-cp311-win_amd64.whl#sha256=9bb0d1421c544ac8e2eca5b47daacaf54706dc9139c003aa5e77ee5f355c5931 ; sys_platform == 'win32' and python_version == '3.11' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", - "torchvision @ https://download.pytorch.org/whl/xpu/torchvision-0.24.0%2Bxpu-cp312-cp312-win_amd64.whl#sha256=6a5194bc736089606342d48a3f6822829b167617e9495d91d753dd1bd46fda18 ; sys_platform == 'win32' and python_version == '3.12' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", - "torchvision @ https://download.pytorch.org/whl/xpu/torchvision-0.24.0%2Bxpu-cp313-cp313-win_amd64.whl#sha256=da47a3ce2bb7f0301a31124668b5908f9b9e92d6241443de15a310ef9632fd83 ; sys_platform == 'win32' and python_version == '3.13' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", -] -intel-gpu-torch290 = [ - "unsloth[intelgputorch290]" -] -intelgputorch271 = [ - "unsloth_zoo[intelgpu]", - "unsloth[huggingfacenotorch]", - - "pytorch_triton_xpu @ https://download.pytorch.org/whl/pytorch_triton_xpu-3.3.1-cp310-cp310-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl#sha256=663ce21364096b268c6687f26f22862cb1001cae0c4ec9f98a0998415f99e2b0 ; platform_system == 'Linux' and python_version == '3.10' and platform_machine == 'x86_64'", - "pytorch_triton_xpu @ https://download.pytorch.org/whl/pytorch_triton_xpu-3.3.1-cp311-cp311-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl#sha256=dd92cc17000bad19f213b6a877d7f10cd71341b703cd188513ce9fff8d42e3dd ; platform_system == 'Linux' and python_version == '3.11' and platform_machine == 'x86_64'", - "pytorch_triton_xpu @ https://download.pytorch.org/whl/pytorch_triton_xpu-3.3.1-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl#sha256=aa5c3ec21a89e967d1dfe61e3d5b1c1ae9620c871ed804771d3378d6a44066f2 ; platform_system == 'Linux' and python_version == '3.12' and platform_machine == 'x86_64'", - "pytorch_triton_xpu @ https://download.pytorch.org/whl/pytorch_triton_xpu-3.3.1-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl#sha256=d1c6f522e11112a311b1a61ba7b40b43ad8305675fa29153017ccb1ad0b6816d ; platform_system == 'Linux' and python_version == '3.13' and platform_machine == 'x86_64'", - "pytorch_triton_xpu @ https://download.pytorch.org/whl/pytorch_triton_xpu-3.3.1-cp310-cp310-win_amd64.whl#sha256=a5c16dcf449a9cb62bc3788f7ec45782bb3ead6edc2637a12b60ef0f8f45dc55 ; sys_platform == 'win32' and python_version == '3.10' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", - "pytorch_triton_xpu @ https://download.pytorch.org/whl/pytorch_triton_xpu-3.3.1-cp311-cp311-win_amd64.whl#sha256=bc2d76ffa4ceed5b38ae34b52dbff643442e1a44d52ca72d7cb520ca1950e9ae ; sys_platform == 'win32' and python_version == '3.11' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", - "pytorch_triton_xpu @ https://download.pytorch.org/whl/pytorch_triton_xpu-3.3.1-cp312-cp312-win_amd64.whl#sha256=b09ca59ce52d6d27b1510df783cde222b703a71857a6fa953f1f155f9f50811a ; sys_platform == 'win32' and python_version == '3.12' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", - "pytorch_triton_xpu @ https://download.pytorch.org/whl/pytorch_triton_xpu-3.3.1-cp313-cp313-win_amd64.whl#sha256=1260c4a4bad426b6cd3c8f3e1a21835381c6f217bf434bcb55fedec08a206dea ; sys_platform == 'win32' and python_version == '3.13' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", - - "torch @ https://download.pytorch.org/whl/xpu/torch-2.7.1%2Bxpu-cp310-cp310-linux_x86_64.whl#sha256=231c3fbd88a75d94de5ccbbb7f4f9a96cb3c58b3d891c2a1b469d38df95f9be6 ; platform_system == 'Linux' and python_version == '3.10' and platform_machine == 'x86_64'", - "torch @ https://download.pytorch.org/whl/xpu/torch-2.7.1%2Bxpu-cp311-cp311-linux_x86_64.whl#sha256=78edcc27709dd819fc820f5eb9421bd10d3f3dcb14adb25ee60766c76f0e67f3 ; platform_system == 'Linux' and python_version == '3.11' and platform_machine == 'x86_64'", - "torch @ https://download.pytorch.org/whl/xpu/torch-2.7.1%2Bxpu-cp312-cp312-linux_x86_64.whl#sha256=b443df40bc9cb7d648a9f8f9ed1d5c3a1203e561ebd0a61dd55fb8a58833d5ec ; platform_system == 'Linux' and python_version == '3.12' and platform_machine == 'x86_64'", - "torch @ https://download.pytorch.org/whl/xpu/torch-2.7.1%2Bxpu-cp313-cp313-linux_x86_64.whl#sha256=412b58ffcceebea399c9a1bcdb22896aa10385c2650a8c4f8a677fb11c49b448 ; platform_system == 'Linux' and python_version == '3.13' and platform_machine == 'x86_64'", - "torch @ https://download.pytorch.org/whl/xpu/torch-2.7.1%2Bxpu-cp310-cp310-win_amd64.whl#sha256=2591228dc2cb73c78daf24277c4449ba9474f94cd31938147249269fe89d05d6 ; sys_platform == 'win32' and python_version == '3.10' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", - "torch @ https://download.pytorch.org/whl/xpu/torch-2.7.1%2Bxpu-cp311-cp311-win_amd64.whl#sha256=1aacb86e9a9684ffc8bde3db14b251d00df7019a9a434ec99a59076a2696325d ; sys_platform == 'win32' and python_version == '3.11' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", - "torch @ https://download.pytorch.org/whl/xpu/torch-2.7.1%2Bxpu-cp312-cp312-win_amd64.whl#sha256=9b65dc8562521b60d77aa653132bc03a19da0291318fcf919faa3f03080d8f7e ; sys_platform == 'win32' and python_version == '3.12' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", - "torch @ https://download.pytorch.org/whl/xpu/torch-2.7.1%2Bxpu-cp313-cp313-win_amd64.whl#sha256=cd3669fee311bc3ee5501d696bf989226a6f2bf957d120a04881a07af05526d6 ; sys_platform == 'win32' and python_version == '3.13' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", - - "bitsandbytes @ https://github.com/bitsandbytes-foundation/bitsandbytes/releases/download/continuous-release_main/bitsandbytes-1.33.7.preview-py3-none-manylinux_2_24_x86_64.whl ; ('linux' in sys_platform) and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", - "bitsandbytes @ https://github.com/bitsandbytes-foundation/bitsandbytes/releases/download/continuous-release_main/bitsandbytes-1.33.7.preview-py3-none-win_amd64.whl ; (sys_platform == 'win32') and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", - - "torchvision @ https://download.pytorch.org/whl/xpu/torchvision-0.22.1%2Bxpu-cp310-cp310-manylinux_2_28_x86_64.whl#sha256=f8cdf6889c02b3166679eef661b68757ea7e99c314432c3d41dac3d2ed4a59d4 ; platform_system == 'Linux' and python_version == '3.10' and platform_machine == 'x86_64'", - "torchvision @ https://download.pytorch.org/whl/xpu/torchvision-0.22.1%2Bxpu-cp311-cp311-manylinux_2_28_x86_64.whl#sha256=f7d15b65d52809745992e0001c25034f33ac01f2dff5248614e07b5d009a59b7 ; platform_system == 'Linux' and python_version == '3.11' and platform_machine == 'x86_64'", - "torchvision @ https://download.pytorch.org/whl/xpu/torchvision-0.22.1%2Bxpu-cp312-cp312-manylinux_2_28_x86_64.whl#sha256=1ff1f98d70846352c7f56833bedab1a055ead27b11c120b8c719063ee0383554 ; platform_system == 'Linux' and python_version == '3.12' and platform_machine == 'x86_64'", - "torchvision @ https://download.pytorch.org/whl/xpu/torchvision-0.22.1%2Bxpu-cp313-cp313-manylinux_2_28_x86_64.whl#sha256=f46945344ea911a70309231eaaf3b80c96f6646ce5515dc89aa94f94144e310e ; platform_system == 'Linux' and python_version == '3.13' and platform_machine == 'x86_64'", - "torchvision @ https://download.pytorch.org/whl/xpu/torchvision-0.22.1%2Bxpu-cp310-cp310-win_amd64.whl#sha256=ecae9a02de769e2070d37388116beb407c3f0d60b8e65c1da1423f4eafee361a ; sys_platform == 'win32' and python_version == '3.10' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", - "torchvision @ https://download.pytorch.org/whl/xpu/torchvision-0.22.1%2Bxpu-cp311-cp311-win_amd64.whl#sha256=2914e62782431bebd6ad9a3b98a2b7311e448e84a7534bb7f35874b9279a17de ; sys_platform == 'win32' and python_version == '3.11' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", - "torchvision @ https://download.pytorch.org/whl/xpu/torchvision-0.22.1%2Bxpu-cp312-cp312-win_amd64.whl#sha256=5b462c156f4e2097e1e53649d3f298ce352fa4c5d1e6addd360375b10ebd6c67 ; sys_platform == 'win32' and python_version == '3.12' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", - "torchvision @ https://download.pytorch.org/whl/xpu/torchvision-0.22.1%2Bxpu-cp313-cp313-win_amd64.whl#sha256=fa87b3677cd1af67ce423004283c1bde80e3571f391182a3e89b485e18e3c70f ; sys_platform == 'win32' and python_version == '3.13' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", -] -intel-gpu-torch271 = [ - "unsloth[intelgputorch271]" -] -intelgputorch291 = [ - "unsloth_zoo[intelgpu]", - "unsloth[huggingfacenotorch]", - - "pytorch_triton_xpu @ https://download.pytorch.org/whl/pytorch_triton_xpu-3.5.0-cp310-cp310-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl#sha256=c169a1de14c19673b17c751290d467fa282fc90fa5da4314b2e5cdab1f553146 ; platform_system == 'Linux' and python_version == '3.10' and platform_machine == 'x86_64'", - "pytorch_triton_xpu @ https://download.pytorch.org/whl/pytorch_triton_xpu-3.5.0-cp311-cp311-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl#sha256=013d9dd5d6479bd22983161f462e61c8dbe1d82e6730624a7a8d5945507eaa61 ; platform_system == 'Linux' and python_version == '3.11' and platform_machine == 'x86_64'", - "pytorch_triton_xpu @ https://download.pytorch.org/whl/pytorch_triton_xpu-3.5.0-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl#sha256=afc8cabfbf7ed51fd278d1e0f88d6afc157b0201bad4b99d681e4d542f9e66d4 ; platform_system == 'Linux' and python_version == '3.12' and platform_machine == 'x86_64'", - "pytorch_triton_xpu @ https://download.pytorch.org/whl/pytorch_triton_xpu-3.5.0-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl#sha256=0d24c1716088f2764d0d24c64227732195b6a42706c3c5fc89eeb4904bfa0818 ; platform_system == 'Linux' and python_version == '3.13' and platform_machine == 'x86_64'", - "pytorch_triton_xpu @ https://download.pytorch.org/whl/pytorch_triton_xpu-3.5.0-cp310-cp310-win_amd64.whl#sha256=c83ab007311d9cfb6e809ee5a4587d99a9eef4be720b90da4f1aaa68b45139a0 ; sys_platform == 'win32' and python_version == '3.10' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", - "pytorch_triton_xpu @ https://download.pytorch.org/whl/pytorch_triton_xpu-3.5.0-cp311-cp311-win_amd64.whl#sha256=debf75348da8e8c7166b4d4a9b91d1508bb8d6581e339f79f7604b2e6746bacd ; sys_platform == 'win32' and python_version == '3.11' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", - "pytorch_triton_xpu @ https://download.pytorch.org/whl/pytorch_triton_xpu-3.5.0-cp312-cp312-win_amd64.whl#sha256=97337a47425f1963a723475bd61037460e84ba01db4f87a1d662c3718ff6c47e ; sys_platform == 'win32' and python_version == '3.12' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", - "pytorch_triton_xpu @ https://download.pytorch.org/whl/pytorch_triton_xpu-3.5.0-cp313-cp313-win_amd64.whl#sha256=2caf8138695f6abb023ecd02031a2611ba1bf8fff2f19802567cb2fadefe9e87 ; sys_platform == 'win32' and python_version == '3.13' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", - - "torch @ https://download.pytorch.org/whl/xpu/torch-2.9.1%2Bxpu-cp310-cp310-linux_x86_64.whl#sha256=fb7895c744132d6a8e56ce8434ae1d8355c9bda4e9f58832744ff742d6268eaf ; platform_system == 'Linux' and python_version == '3.10' and platform_machine == 'x86_64'", - "torch @ https://download.pytorch.org/whl/xpu/torch-2.9.1%2Bxpu-cp311-cp311-linux_x86_64.whl#sha256=da2604a9114a28de71ce654819424d20a246adf644d191ae160837df9731b79e ; platform_system == 'Linux' and python_version == '3.11' and platform_machine == 'x86_64'", - "torch @ https://download.pytorch.org/whl/xpu/torch-2.9.1%2Bxpu-cp312-cp312-linux_x86_64.whl#sha256=d5968d78d81c1d01efc1b3bf83d7da3d83161dcc3a9fcf91f500591db1c6c75d ; platform_system == 'Linux' and python_version == '3.12' and platform_machine == 'x86_64'", - "torch @ https://download.pytorch.org/whl/xpu/torch-2.9.1%2Bxpu-cp313-cp313-linux_x86_64.whl#sha256=b56d6b0d65863f370527e971dbfa046a5dd2a1f61cc95071db26c764f36e4dce ; platform_system == 'Linux' and python_version == '3.13' and platform_machine == 'x86_64'", - "torch @ https://download.pytorch.org/whl/xpu/torch-2.9.1%2Bxpu-cp310-cp310-win_amd64.whl#sha256=2f318fb6a4bf1101cc17f35a5371f7c1768b41fceed03628397834e85b3edfdd ; sys_platform == 'win32' and python_version == '3.10' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", - "torch @ https://download.pytorch.org/whl/xpu/torch-2.9.1%2Bxpu-cp311-cp311-win_amd64.whl#sha256=c9cedc3fb099366b2e6c563df6578e323564b1b5d40ac27be73c674755343a1d ; sys_platform == 'win32' and python_version == '3.11' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", - "torch @ https://download.pytorch.org/whl/xpu/torch-2.9.1%2Bxpu-cp312-cp312-win_amd64.whl#sha256=bee9623254d0f95a1ca115dbd17e9a9d966fdb8ae123e2ada4a9eb2fb8d38db8 ; sys_platform == 'win32' and python_version == '3.12' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", - "torch @ https://download.pytorch.org/whl/xpu/torch-2.9.1%2Bxpu-cp313-cp313-win_amd64.whl#sha256=cd5c857da52a63c121561b30b0979e69ade70b575fd74e389787bc7c1ee2ac11 ; sys_platform == 'win32' and python_version == '3.13' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", - - "bitsandbytes @ https://github.com/bitsandbytes-foundation/bitsandbytes/releases/download/continuous-release_main/bitsandbytes-1.33.7.preview-py3-none-manylinux_2_24_x86_64.whl ; ('linux' in sys_platform) and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", - "bitsandbytes @ https://github.com/bitsandbytes-foundation/bitsandbytes/releases/download/continuous-release_main/bitsandbytes-1.33.7.preview-py3-none-win_amd64.whl ; (sys_platform == 'win32') and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", - - "torchvision @ https://download.pytorch.org/whl/xpu/torchvision-0.24.1%2Bxpu-cp310-cp310-manylinux_2_28_x86_64.whl#sha256=cc5272da2cb4554edf059eedd6d1f5ef2859033b0fb79d5dcb8e99a0697f3325 ; platform_system == 'Linux' and python_version == '3.10' and platform_machine == 'x86_64'", - "torchvision @ https://download.pytorch.org/whl/xpu/torchvision-0.24.1%2Bxpu-cp311-cp311-manylinux_2_28_x86_64.whl#sha256=3c80d6a068c32fc4ebddb27953e03a0141bd0f10ca8730417cbc0e0748158285 ; platform_system == 'Linux' and python_version == '3.11' and platform_machine == 'x86_64'", - "torchvision @ https://download.pytorch.org/whl/xpu/torchvision-0.24.1%2Bxpu-cp312-cp312-manylinux_2_28_x86_64.whl#sha256=8cf640a867cf270b3fda7a10002c29d3fc2ad6dfbd76404a8cdd820489adb04c ; platform_system == 'Linux' and python_version == '3.12' and platform_machine == 'x86_64'", - "torchvision @ https://download.pytorch.org/whl/xpu/torchvision-0.24.1%2Bxpu-cp313-cp313-manylinux_2_28_x86_64.whl#sha256=d9c59ee5ae3d0560f02401c8dfd8054d50813a8dbb5d33a8777de7d02f6fcb7b ; platform_system == 'Linux' and python_version == '3.13' and platform_machine == 'x86_64'", - "torchvision @ https://download.pytorch.org/whl/xpu/torchvision-0.24.1%2Bxpu-cp310-cp310-win_amd64.whl#sha256=843ea7fcd8f5a22ebbc20d2d61d9eec7593821a0372eb8cabb73953d12ef6acf ; sys_platform == 'win32' and python_version == '3.10' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", - "torchvision @ https://download.pytorch.org/whl/xpu/torchvision-0.24.1%2Bxpu-cp311-cp311-win_amd64.whl#sha256=e5ff8a31d3c700f8dbac59697c8e32298a43ec059609ebc6ea7bab3eff6384e1 ; sys_platform == 'win32' and python_version == '3.11' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", - "torchvision @ https://download.pytorch.org/whl/xpu/torchvision-0.24.1%2Bxpu-cp312-cp312-win_amd64.whl#sha256=8bae6d4c042f8d20818da4a5aa9109c6fbd6ec11bc422be152ce8adf9a7095bf ; sys_platform == 'win32' and python_version == '3.12' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", - "torchvision @ https://download.pytorch.org/whl/xpu/torchvision-0.24.1%2Bxpu-cp313-cp313-win_amd64.whl#sha256=47059e290fc2a41ba78666ffcde102c436abf7ff8a34d200268b48c4fa0f9c45 ; sys_platform == 'win32' and python_version == '3.13' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", -] -intel-gpu-torch291 = [ - "unsloth[intelgputorch291]" -] -intelgputorch210 = [ - "unsloth_zoo[intelgpu]", - "unsloth[huggingfacenotorch]", - - "triton-xpu @ https://download.pytorch.org/whl/triton_xpu-3.6.0-cp310-cp310-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl ; platform_system == 'Linux' and python_version == '3.10' and platform_machine == 'x86_64'", - "triton-xpu @ https://download.pytorch.org/whl/triton_xpu-3.6.0-cp311-cp311-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl ; platform_system == 'Linux' and python_version == '3.11' and platform_machine == 'x86_64'", - "triton-xpu @ https://download.pytorch.org/whl/triton_xpu-3.6.0-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl ; platform_system == 'Linux' and python_version == '3.12' and platform_machine == 'x86_64'", - "triton-xpu @ https://download.pytorch.org/whl/triton_xpu-3.6.0-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl ; platform_system == 'Linux' and python_version == '3.13' and platform_machine == 'x86_64'", - "triton-xpu @ https://download.pytorch.org/whl/triton_xpu-3.6.0-cp310-cp310-win_amd64.whl ; sys_platform == 'win32' and python_version == '3.10' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", - "triton-xpu @ https://download.pytorch.org/whl/triton_xpu-3.6.0-cp311-cp311-win_amd64.whl ; sys_platform == 'win32' and python_version == '3.11' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", - "triton-xpu @ https://download.pytorch.org/whl/triton_xpu-3.6.0-cp312-cp312-win_amd64.whl ; sys_platform == 'win32' and python_version == '3.12' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", - "triton-xpu @ https://download.pytorch.org/whl/triton_xpu-3.6.0-cp313-cp313-win_amd64.whl ; sys_platform == 'win32' and python_version == '3.13' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", - - "torch @ https://download.pytorch.org/whl/xpu/torch-2.10.0%2Bxpu-cp310-cp310-linux_x86_64.whl#sha256=abb1d1ec1ac672bac0ff35420c965f2df0c636ef9d94e2a830e34578489d0a57 ; platform_system == 'Linux' and python_version == '3.10' and platform_machine == 'x86_64'", - "torch @ https://download.pytorch.org/whl/xpu/torch-2.10.0%2Bxpu-cp311-cp311-linux_x86_64.whl#sha256=71ad2f82da0f41eaec159f39fc85854e27c2391efa91b373e550648a6f4aaad3 ; platform_system == 'Linux' and python_version == '3.11' and platform_machine == 'x86_64'", - "torch @ https://download.pytorch.org/whl/xpu/torch-2.10.0%2Bxpu-cp312-cp312-linux_x86_64.whl#sha256=b473571d478912f92881cc13f15fa18f8463fb0fb8a068c96ed47a7d45a4da0a ; platform_system == 'Linux' and python_version == '3.12' and platform_machine == 'x86_64'", - "torch @ https://download.pytorch.org/whl/xpu/torch-2.10.0%2Bxpu-cp313-cp313-linux_x86_64.whl#sha256=3bc64a746ff25a93de140902c60c9e819d7413f5cea1e88d80999c27a5901e9c ; platform_system == 'Linux' and python_version == '3.13' and platform_machine == 'x86_64'", - "torch @ https://download.pytorch.org/whl/xpu/torch-2.10.0%2Bxpu-cp310-cp310-win_amd64.whl#sha256=ce50691ab3fb6301d9b7bb8b3834cf5fa7152a2b5f91fd24c5efdc601a25b780 ; sys_platform == 'win32' and python_version == '3.10' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", - "torch @ https://download.pytorch.org/whl/xpu/torch-2.10.0%2Bxpu-cp311-cp311-win_amd64.whl#sha256=cb9d37f21cb9fb7df67d62863f021c3144e8d8832b9ea8e8523ac308bc620ea1 ; sys_platform == 'win32' and python_version == '3.11' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", - "torch @ https://download.pytorch.org/whl/xpu/torch-2.10.0%2Bxpu-cp312-cp312-win_amd64.whl#sha256=3ad605be4728b6d3a28a44d07dd794b1a9e45551b0057815bf25eb2a6d6a56a7 ; sys_platform == 'win32' and python_version == '3.12' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", - "torch @ https://download.pytorch.org/whl/xpu/torch-2.10.0%2Bxpu-cp313-cp313-win_amd64.whl#sha256=2b4b56dd6c792aef82006904fa888692e3782e4ae5da27526801bad4898f05a5 ; sys_platform == 'win32' and python_version == '3.13' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", - - "bitsandbytes @ https://github.com/bitsandbytes-foundation/bitsandbytes/releases/download/continuous-release_main/bitsandbytes-1.33.7.preview-py3-none-manylinux_2_24_x86_64.whl ; ('linux' in sys_platform) and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", - "bitsandbytes @ https://github.com/bitsandbytes-foundation/bitsandbytes/releases/download/continuous-release_main/bitsandbytes-1.33.7.preview-py3-none-win_amd64.whl ; (sys_platform == 'win32') and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", - - "torchvision @ https://download.pytorch.org/whl/xpu/torchvision-0.25.0%2Bxpu-cp310-cp310-manylinux_2_28_x86_64.whl#sha256=7e1e7b170fcf7161c8499b67156c5a05462243626dc0974010791a0bab4378d3 ; platform_system == 'Linux' and python_version == '3.10' and platform_machine == 'x86_64'", - "torchvision @ https://download.pytorch.org/whl/xpu/torchvision-0.25.0%2Bxpu-cp311-cp311-manylinux_2_28_x86_64.whl#sha256=bd6add201bd7628af70437292e1447abb368e0b5f4ff9abd334ae435efd44792 ; platform_system == 'Linux' and python_version == '3.11' and platform_machine == 'x86_64'", - "torchvision @ https://download.pytorch.org/whl/xpu/torchvision-0.25.0%2Bxpu-cp312-cp312-manylinux_2_28_x86_64.whl#sha256=6ad2543496bc29e59d3dd614a94d09aa9870318aedb66045344fffddfedd2cf8 ; platform_system == 'Linux' and python_version == '3.12' and platform_machine == 'x86_64'", - "torchvision @ https://download.pytorch.org/whl/xpu/torchvision-0.25.0%2Bxpu-cp313-cp313-manylinux_2_28_x86_64.whl#sha256=80269f37865fcd8b57f20e4786efae2200bfa2b2727926c3c7acc82f0e7d3548 ; platform_system == 'Linux' and python_version == '3.13' and platform_machine == 'x86_64'", - "torchvision @ https://download.pytorch.org/whl/xpu/torchvision-0.25.0%2Bxpu-cp310-cp310-win_amd64.whl#sha256=6b9485ba85dcba4d196d6134d9c3332fb228fb2556416bf0450a64e8a472fcba ; sys_platform == 'win32' and python_version == '3.10' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", - "torchvision @ https://download.pytorch.org/whl/xpu/torchvision-0.25.0%2Bxpu-cp311-cp311-win_amd64.whl#sha256=36cbaedf10f6412af5c89afd9aeea474e6a56a0050348ada8fabe1ecaf6b879e ; sys_platform == 'win32' and python_version == '3.11' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", - "torchvision @ https://download.pytorch.org/whl/xpu/torchvision-0.25.0%2Bxpu-cp312-cp312-win_amd64.whl#sha256=738357d97468d75fe3d510ac37e65130f2787f81d9bbc1518898f7396dc3403f ; sys_platform == 'win32' and python_version == '3.12' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", - "torchvision @ https://download.pytorch.org/whl/xpu/torchvision-0.25.0%2Bxpu-cp313-cp313-win_amd64.whl#sha256=1c4b44b36a557f7381e3076fb8843366742238648441d607c8d049c6da0f8886 ; sys_platform == 'win32' and python_version == '3.13' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", -] -intel-gpu-torch210 = [ - "unsloth[intelgputorch210]" -] -intelgputorch2110 = [ - "unsloth_zoo[intelgpu]", - "unsloth[huggingfacenotorch]", - - "triton-xpu @ https://download.pytorch.org/whl/triton_xpu-3.7.0-cp310-cp310-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl#sha256=2a1841138750f708ec017becbf8d357526f3fa350deee6553be5735ad66160a3 ; platform_system == 'Linux' and python_version == '3.10' and platform_machine == 'x86_64'", - "triton-xpu @ https://download.pytorch.org/whl/triton_xpu-3.7.0-cp311-cp311-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl#sha256=e85378f1fc1ea002271de2a35475b75008fa554b86ef9d3bc55be9c513a63b51 ; platform_system == 'Linux' and python_version == '3.11' and platform_machine == 'x86_64'", - "triton-xpu @ https://download.pytorch.org/whl/triton_xpu-3.7.0-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl#sha256=a6663ebe43e3c0d560ff774708632d7a75208ee64a291c1724ed5c16a92d1c72 ; platform_system == 'Linux' and python_version == '3.12' and platform_machine == 'x86_64'", - "triton-xpu @ https://download.pytorch.org/whl/triton_xpu-3.7.0-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl#sha256=08c8d43b2831faf9d6799480df2b45dde58102257aebd810d07a2ce18cd4e5df ; platform_system == 'Linux' and python_version == '3.13' and platform_machine == 'x86_64'", - "triton-xpu @ https://download.pytorch.org/whl/triton_xpu-3.7.0-cp310-cp310-win_amd64.whl#sha256=90fb8f767950a4ffca627faa7f86d9c697237ea4352d7e23505c5c9ed8e72216 ; sys_platform == 'win32' and python_version == '3.10' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", - "triton-xpu @ https://download.pytorch.org/whl/triton_xpu-3.7.0-cp311-cp311-win_amd64.whl#sha256=aa7de82f4265089e74f25a2701b7532e5c47d74224d877b61da1d66156e3f0c1 ; sys_platform == 'win32' and python_version == '3.11' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", - "triton-xpu @ https://download.pytorch.org/whl/triton_xpu-3.7.0-cp312-cp312-win_amd64.whl#sha256=5ba3a31c6e1b259ad2d924e1b50f72a78c6ebd7eb4f364473bbf93e144734e80 ; sys_platform == 'win32' and python_version == '3.12' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", - "triton-xpu @ https://download.pytorch.org/whl/triton_xpu-3.7.0-cp313-cp313-win_amd64.whl#sha256=e8b4caba9b2399ea4c7f9a2777042564dea5d6f9e586a2dcb015a4ce20f000f7 ; sys_platform == 'win32' and python_version == '3.13' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", - - "torch @ https://download.pytorch.org/whl/xpu/torch-2.11.0%2Bxpu-cp310-cp310-linux_x86_64.whl ; platform_system == 'Linux' and python_version == '3.10' and platform_machine == 'x86_64'", - "torch @ https://download.pytorch.org/whl/xpu/torch-2.11.0%2Bxpu-cp311-cp311-linux_x86_64.whl ; platform_system == 'Linux' and python_version == '3.11' and platform_machine == 'x86_64'", - "torch @ https://download.pytorch.org/whl/xpu/torch-2.11.0%2Bxpu-cp312-cp312-linux_x86_64.whl ; platform_system == 'Linux' and python_version == '3.12' and platform_machine == 'x86_64'", - "torch @ https://download.pytorch.org/whl/xpu/torch-2.11.0%2Bxpu-cp313-cp313-linux_x86_64.whl ; platform_system == 'Linux' and python_version == '3.13' and platform_machine == 'x86_64'", - "torch @ https://download.pytorch.org/whl/xpu/torch-2.11.0%2Bxpu-cp310-cp310-win_amd64.whl ; sys_platform == 'win32' and python_version == '3.10' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", - "torch @ https://download.pytorch.org/whl/xpu/torch-2.11.0%2Bxpu-cp311-cp311-win_amd64.whl ; sys_platform == 'win32' and python_version == '3.11' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", - "torch @ https://download.pytorch.org/whl/xpu/torch-2.11.0%2Bxpu-cp312-cp312-win_amd64.whl ; sys_platform == 'win32' and python_version == '3.12' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", - "torch @ https://download.pytorch.org/whl/xpu/torch-2.11.0%2Bxpu-cp313-cp313-win_amd64.whl ; sys_platform == 'win32' and python_version == '3.13' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", - - "bitsandbytes @ https://github.com/bitsandbytes-foundation/bitsandbytes/releases/download/continuous-release_main/bitsandbytes-1.33.7.preview-py3-none-manylinux_2_24_x86_64.whl ; ('linux' in sys_platform) and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", - "bitsandbytes @ https://github.com/bitsandbytes-foundation/bitsandbytes/releases/download/continuous-release_main/bitsandbytes-1.33.7.preview-py3-none-win_amd64.whl ; (sys_platform == 'win32') and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", - - "torchvision @ https://download.pytorch.org/whl/xpu/torchvision-0.26.0%2Bxpu-cp310-cp310-manylinux_2_28_x86_64.whl#sha256=6e634354b752b7366e8ad16b84f3e7e5863776a7ab448bbabae4fd36668dee7a ; platform_system == 'Linux' and python_version == '3.10' and platform_machine == 'x86_64'", - "torchvision @ https://download.pytorch.org/whl/xpu/torchvision-0.26.0%2Bxpu-cp311-cp311-manylinux_2_28_x86_64.whl#sha256=293169899f562ce473a58836dd024f0b1e72a347400278287ab393d1b04991e4 ; platform_system == 'Linux' and python_version == '3.11' and platform_machine == 'x86_64'", - "torchvision @ https://download.pytorch.org/whl/xpu/torchvision-0.26.0%2Bxpu-cp312-cp312-manylinux_2_28_x86_64.whl#sha256=e204d14be6f0f84d5f0e6e9213556e80326c3ab682cac108bcbef340bf45297b ; platform_system == 'Linux' and python_version == '3.12' and platform_machine == 'x86_64'", - "torchvision @ https://download.pytorch.org/whl/xpu/torchvision-0.26.0%2Bxpu-cp313-cp313-manylinux_2_28_x86_64.whl#sha256=f134344006f0989a2d771554b7905fb05bd93d63b195e64626fde3495ec6f287 ; platform_system == 'Linux' and python_version == '3.13' and platform_machine == 'x86_64'", - "torchvision @ https://download.pytorch.org/whl/xpu/torchvision-0.26.0%2Bxpu-cp310-cp310-win_amd64.whl#sha256=7e52729cb9736c66dc79a7f42de6b31db93b9161d3357fd34cfa33f5fe32b8ea ; sys_platform == 'win32' and python_version == '3.10' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", - "torchvision @ https://download.pytorch.org/whl/xpu/torchvision-0.26.0%2Bxpu-cp311-cp311-win_amd64.whl#sha256=83a6130100c6b6750d8aa9fd29e5d0c53b1c85b1153b8ed4139aea54fc1892cc ; sys_platform == 'win32' and python_version == '3.11' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", - "torchvision @ https://download.pytorch.org/whl/xpu/torchvision-0.26.0%2Bxpu-cp312-cp312-win_amd64.whl#sha256=03788e0e5a5b85a2f09d11f0263d579fcb0cf5623d8810149be0e37836c2738c ; sys_platform == 'win32' and python_version == '3.12' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", - "torchvision @ https://download.pytorch.org/whl/xpu/torchvision-0.26.0%2Bxpu-cp313-cp313-win_amd64.whl#sha256=cb1da1d378ce440f7d1e0ed8cf21bd280d904ab25a55c9453f8377825818df74 ; sys_platform == 'win32' and python_version == '3.13' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", -] -intel-gpu-torch2110 = [ - "unsloth[intelgputorch2110]" -] -intelgputorch2120 = [ - "unsloth_zoo[intelgpu]", - "unsloth[huggingfacenotorch]", - - "triton-xpu @ https://download.pytorch.org/whl/triton_xpu-3.7.1-cp310-cp310-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl#sha256=844d981cb1b3948085e8cfa62c74de9f100259f6131959aa70be49123b88ae81 ; platform_system == 'Linux' and python_version == '3.10' and platform_machine == 'x86_64'", - "triton-xpu @ https://download.pytorch.org/whl/triton_xpu-3.7.1-cp311-cp311-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl#sha256=a16b1d00e94ad87d62af3512e390348b8656419598004100c56028bf494f086b ; platform_system == 'Linux' and python_version == '3.11' and platform_machine == 'x86_64'", - "triton-xpu @ https://download.pytorch.org/whl/triton_xpu-3.7.1-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl#sha256=4e46e71e077cf483404a4c17ce40d71c5f0e13a81459139d4346ca427b1dd455 ; platform_system == 'Linux' and python_version == '3.12' and platform_machine == 'x86_64'", - "triton-xpu @ https://download.pytorch.org/whl/triton_xpu-3.7.1-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl#sha256=4fdaed1bafc51d3a2834656a3420a6686a74ea226508765a49bf15d58ff3a930 ; platform_system == 'Linux' and python_version == '3.13' and platform_machine == 'x86_64'", - "triton-xpu @ https://download.pytorch.org/whl/triton_xpu-3.7.1-cp310-cp310-win_amd64.whl#sha256=2778b46b22e9fa0916398db299a125027a1b2331c1173b3dd2b9e2cab6263a31 ; sys_platform == 'win32' and python_version == '3.10' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", - "triton-xpu @ https://download.pytorch.org/whl/triton_xpu-3.7.1-cp311-cp311-win_amd64.whl#sha256=ad5b147d04ee0d40f3d4d32f85f5aa3a3beb6cd5799ca026d3d7f4afa3d9e24f ; sys_platform == 'win32' and python_version == '3.11' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", - "triton-xpu @ https://download.pytorch.org/whl/triton_xpu-3.7.1-cp312-cp312-win_amd64.whl#sha256=d9482063af2a308543f23333e32edd738ea87cbb33ade68afda9ae0fd704ccd9 ; sys_platform == 'win32' and python_version == '3.12' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", - "triton-xpu @ https://download.pytorch.org/whl/triton_xpu-3.7.1-cp313-cp313-win_amd64.whl#sha256=5d4d67f0deb1e851c01b293e602b8dcddad26ca2be61221cee3dc0e1aa0cdefd ; sys_platform == 'win32' and python_version == '3.13' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", - - "torch @ https://download.pytorch.org/whl/xpu/torch-2.12.0%2Bxpu-cp310-cp310-linux_x86_64.whl#sha256=e8923cd1fe560472904b1461b745d2f1826bb9c1bc0808225d5f28a450e4d553 ; platform_system == 'Linux' and python_version == '3.10' and platform_machine == 'x86_64'", - "torch @ https://download.pytorch.org/whl/xpu/torch-2.12.0%2Bxpu-cp311-cp311-linux_x86_64.whl#sha256=f7c082b2fc9b61def594d30ea57762dc4a8bc7111a9a9593953ed948de242e28 ; platform_system == 'Linux' and python_version == '3.11' and platform_machine == 'x86_64'", - "torch @ https://download.pytorch.org/whl/xpu/torch-2.12.0%2Bxpu-cp312-cp312-linux_x86_64.whl#sha256=f59decc04bec27862ed0197554a52370dbcba3e6892616d1fbce450e402bf2d5 ; platform_system == 'Linux' and python_version == '3.12' and platform_machine == 'x86_64'", - "torch @ https://download.pytorch.org/whl/xpu/torch-2.12.0%2Bxpu-cp313-cp313-linux_x86_64.whl#sha256=56f74e7c6c096e1a7ac215eb79ee590b764be3fbba8f4febc145bca47194a083 ; platform_system == 'Linux' and python_version == '3.13' and platform_machine == 'x86_64'", - "torch @ https://download.pytorch.org/whl/xpu/torch-2.12.0%2Bxpu-cp310-cp310-win_amd64.whl#sha256=b9779b71457b5a916ae052ed2467c10273cae4862d469b191359173b2038c53e ; sys_platform == 'win32' and python_version == '3.10' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", - "torch @ https://download.pytorch.org/whl/xpu/torch-2.12.0%2Bxpu-cp311-cp311-win_amd64.whl#sha256=7ef8e776c992e4e3ae007ebc108eb4f36b1d1dd9da97ecb308ab7fded89a2659 ; sys_platform == 'win32' and python_version == '3.11' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", - "torch @ https://download.pytorch.org/whl/xpu/torch-2.12.0%2Bxpu-cp312-cp312-win_amd64.whl#sha256=7f1d40febf2b8724adf4ff23866897d87478cc43de2a20f7776dc00be334c464 ; sys_platform == 'win32' and python_version == '3.12' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", - "torch @ https://download.pytorch.org/whl/xpu/torch-2.12.0%2Bxpu-cp313-cp313-win_amd64.whl#sha256=32770e2613df26e2c81ae64ea001b2ca12b8d152231285caff9b5f963a21ad75 ; sys_platform == 'win32' and python_version == '3.13' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", - - "bitsandbytes @ https://github.com/bitsandbytes-foundation/bitsandbytes/releases/download/continuous-release_main/bitsandbytes-1.33.7.preview-py3-none-manylinux_2_24_x86_64.whl ; ('linux' in sys_platform) and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", - "bitsandbytes @ https://github.com/bitsandbytes-foundation/bitsandbytes/releases/download/continuous-release_main/bitsandbytes-1.33.7.preview-py3-none-win_amd64.whl ; (sys_platform == 'win32') and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", - - "torchvision @ https://download.pytorch.org/whl/xpu/torchvision-0.27.0%2Bxpu-cp310-cp310-manylinux_2_28_x86_64.whl#sha256=0d517462caf6f5201c0d7c880f4ac431783c88fcc59b4587836da6c72a89509c ; platform_system == 'Linux' and python_version == '3.10' and platform_machine == 'x86_64'", - "torchvision @ https://download.pytorch.org/whl/xpu/torchvision-0.27.0%2Bxpu-cp311-cp311-manylinux_2_28_x86_64.whl#sha256=4b6feada86aa0bd606904b05898b33538106120d8ed706ba11d0011046534cb8 ; platform_system == 'Linux' and python_version == '3.11' and platform_machine == 'x86_64'", - "torchvision @ https://download.pytorch.org/whl/xpu/torchvision-0.27.0%2Bxpu-cp312-cp312-manylinux_2_28_x86_64.whl#sha256=e231819be0f87829c2344c909c1f0db9d6ae7d6faefe644a526a1a01d0c18d98 ; platform_system == 'Linux' and python_version == '3.12' and platform_machine == 'x86_64'", - "torchvision @ https://download.pytorch.org/whl/xpu/torchvision-0.27.0%2Bxpu-cp313-cp313-manylinux_2_28_x86_64.whl#sha256=8bc7d37515cea18af4c389d5fde58b1a9d76b015f2d87e4a7dc62ad50b1cc200 ; platform_system == 'Linux' and python_version == '3.13' and platform_machine == 'x86_64'", - "torchvision @ https://download.pytorch.org/whl/xpu/torchvision-0.27.0%2Bxpu-cp310-cp310-win_amd64.whl#sha256=65dbb041057dddfe369f29cfaab63f75563621779a23a7b1e2c0ff8a84d4376a ; sys_platform == 'win32' and python_version == '3.10' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", - "torchvision @ https://download.pytorch.org/whl/xpu/torchvision-0.27.0%2Bxpu-cp311-cp311-win_amd64.whl#sha256=df647445365924d69fe3bb2a15a7edfe5b63ef91e4ae69af11d93582985237a4 ; sys_platform == 'win32' and python_version == '3.11' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", - "torchvision @ https://download.pytorch.org/whl/xpu/torchvision-0.27.0%2Bxpu-cp312-cp312-win_amd64.whl#sha256=b0db3df0d0d154d18ba988ab420f1da2549f9372113ff54ff66e4ae3c7fe3bd0 ; sys_platform == 'win32' and python_version == '3.12' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", - "torchvision @ https://download.pytorch.org/whl/xpu/torchvision-0.27.0%2Bxpu-cp313-cp313-win_amd64.whl#sha256=c70850842068c43a0d50eaf139c25b6f6cc9b17a0dae70218c7e69edbee0bc80 ; sys_platform == 'win32' and python_version == '3.13' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", -] -intel-gpu-torch2120 = [ - "unsloth[intelgputorch2120]" -] -intel = [ - "unsloth[intelgputorch280]", -] -amd = [ - "unsloth[huggingfacenotorch]", - "bitsandbytes>=0.49.1 ; ('linux' in sys_platform) and (platform_machine == 'AMD64' or platform_machine == 'x86_64' or platform_machine == 'aarch64')", - "bitsandbytes>=0.49.1 ; (sys_platform == 'win32') and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", -] -rocm702-torch280 = [ - "unsloth[amd]", - - "triton @ https://repo.radeon.com/rocm/manylinux/rocm-rel-7.0.2/triton-3.4.0%2Brocm7.0.2.gitf9e5bf54-cp311-cp311-linux_x86_64.whl ; platform_system == 'Linux' and python_version == '3.11' and platform_machine == 'x86_64'", - "triton @ https://repo.radeon.com/rocm/manylinux/rocm-rel-7.0.2/triton-3.4.0%2Brocm7.0.2.gitf9e5bf54-cp312-cp312-linux_x86_64.whl ; platform_system == 'Linux' and python_version == '3.12' and platform_machine == 'x86_64'", - "triton @ https://repo.radeon.com/rocm/manylinux/rocm-rel-7.0.2/triton-3.4.0%2Brocm7.0.2.gitf9e5bf54-cp313-cp313-linux_x86_64.whl ; platform_system == 'Linux' and python_version == '3.13' and platform_machine == 'x86_64'", - - "torch @ https://repo.radeon.com/rocm/manylinux/rocm-rel-7.0.2/torch-2.8.0%2Brocm7.0.2.lw.git245bf6ed-cp311-cp311-linux_x86_64.whl ; platform_system == 'Linux' and python_version == '3.11' and platform_machine == 'x86_64'", - "torch @ https://repo.radeon.com/rocm/manylinux/rocm-rel-7.0.2/torch-2.8.0%2Brocm7.0.2.lw.git245bf6ed-cp312-cp312-linux_x86_64.whl ; platform_system == 'Linux' and python_version == '3.12' and platform_machine == 'x86_64'", - "torch @ https://repo.radeon.com/rocm/manylinux/rocm-rel-7.0.2/torch-2.8.0%2Brocm7.0.2.lw.git245bf6ed-cp313-cp313-linux_x86_64.whl ; platform_system == 'Linux' and python_version == '3.13' and platform_machine == 'x86_64'", - - "torchvision @ https://repo.radeon.com/rocm/manylinux/rocm-rel-7.0.2/torchvision-0.23.0%2Brocm7.0.2.git824e8c87-cp311-cp311-linux_x86_64.whl ; platform_system == 'Linux' and python_version == '3.11' and platform_machine == 'x86_64'", - "torchvision @ https://repo.radeon.com/rocm/manylinux/rocm-rel-7.0.2/torchvision-0.23.0%2Brocm7.0.2.git824e8c87-cp312-cp312-linux_x86_64.whl ; platform_system == 'Linux' and python_version == '3.12' and platform_machine == 'x86_64'", - "torchvision @ https://repo.radeon.com/rocm/manylinux/rocm-rel-7.0.2/torchvision-0.23.0%2Brocm7.0.2.git824e8c87-cp313-cp313-linux_x86_64.whl ; platform_system == 'Linux' and python_version == '3.13' and platform_machine == 'x86_64'", -] -rocm72-torch291 = [ - "unsloth[amd]", - - "triton @ https://repo.radeon.com/rocm/manylinux/rocm-rel-7.2/triton-3.5.1%2Brocm7.2.0.gita272dfa8-cp310-cp310-linux_x86_64.whl ; platform_system == 'Linux' and python_version == '3.10' and platform_machine == 'x86_64'", - "triton @ https://repo.radeon.com/rocm/manylinux/rocm-rel-7.2/triton-3.5.1%2Brocm7.2.0.gita272dfa8-cp311-cp311-linux_x86_64.whl ; platform_system == 'Linux' and python_version == '3.11' and platform_machine == 'x86_64'", - "triton @ https://repo.radeon.com/rocm/manylinux/rocm-rel-7.2/triton-3.5.1%2Brocm7.2.0.gita272dfa8-cp312-cp312-linux_x86_64.whl ; platform_system == 'Linux' and python_version == '3.12' and platform_machine == 'x86_64'", - "triton @ https://repo.radeon.com/rocm/manylinux/rocm-rel-7.2/triton-3.5.1%2Brocm7.2.0.gita272dfa8-cp313-cp313-linux_x86_64.whl ; platform_system == 'Linux' and python_version == '3.13' and platform_machine == 'x86_64'", - - "torch @ https://repo.radeon.com/rocm/manylinux/rocm-rel-7.2/torch-2.9.1%2Brocm7.2.0.lw.git7e1940d4-cp310-cp310-linux_x86_64.whl ; platform_system == 'Linux' and python_version == '3.10' and platform_machine == 'x86_64'", - "torch @ https://repo.radeon.com/rocm/manylinux/rocm-rel-7.2/torch-2.9.1%2Brocm7.2.0.lw.git7e1940d4-cp311-cp311-linux_x86_64.whl ; platform_system == 'Linux' and python_version == '3.11' and platform_machine == 'x86_64'", - "torch @ https://repo.radeon.com/rocm/manylinux/rocm-rel-7.2/torch-2.9.1%2Brocm7.2.0.lw.git7e1940d4-cp312-cp312-linux_x86_64.whl ; platform_system == 'Linux' and python_version == '3.12' and platform_machine == 'x86_64'", - "torch @ https://repo.radeon.com/rocm/manylinux/rocm-rel-7.2/torch-2.9.1%2Brocm7.2.0.lw.git7e1940d4-cp313-cp313-linux_x86_64.whl ; platform_system == 'Linux' and python_version == '3.13' and platform_machine == 'x86_64'", - "torch @ https://repo.radeon.com/rocm/windows/rocm-rel-7.2/torch-2.9.1%2Brocmsdk20260116-cp312-cp312-win_amd64.whl ; sys_platform == 'win32' and python_version == '3.12'", - - "torchvision @ https://repo.radeon.com/rocm/manylinux/rocm-rel-7.2/torchvision-0.24.0%2Brocm7.2.0.gitb919bd0c-cp310-cp310-linux_x86_64.whl ; platform_system == 'Linux' and python_version == '3.10' and platform_machine == 'x86_64'", - "torchvision @ https://repo.radeon.com/rocm/manylinux/rocm-rel-7.2/torchvision-0.24.0%2Brocm7.2.0.gitb919bd0c-cp311-cp311-linux_x86_64.whl ; platform_system == 'Linux' and python_version == '3.11' and platform_machine == 'x86_64'", - "torchvision @ https://repo.radeon.com/rocm/manylinux/rocm-rel-7.2/torchvision-0.24.0%2Brocm7.2.0.gitb919bd0c-cp312-cp312-linux_x86_64.whl ; platform_system == 'Linux' and python_version == '3.12' and platform_machine == 'x86_64'", - "torchvision @ https://repo.radeon.com/rocm/manylinux/rocm-rel-7.2/torchvision-0.24.0%2Brocm7.2.0.gitb919bd0c-cp313-cp313-linux_x86_64.whl ; platform_system == 'Linux' and python_version == '3.13' and platform_machine == 'x86_64'", - "torchvision @ https://repo.radeon.com/rocm/windows/rocm-rel-7.2/torchvision-0.24.1%2Brocmsdk20260116-cp312-cp312-win_amd64.whl ; sys_platform == 'win32' and python_version == '3.12'", -] -rocm711-torch291 = [ - "unsloth[amd]", - - "triton @ https://repo.radeon.com/rocm/manylinux/rocm-rel-7.1.1/triton-3.5.1%2Brocm7.1.1.gita272dfa8-cp310-cp310-linux_x86_64.whl ; platform_system == 'Linux' and python_version == '3.10' and platform_machine == 'x86_64'", - "triton @ https://repo.radeon.com/rocm/manylinux/rocm-rel-7.1.1/triton-3.5.1%2Brocm7.1.1.gita272dfa8-cp311-cp311-linux_x86_64.whl ; platform_system == 'Linux' and python_version == '3.11' and platform_machine == 'x86_64'", - "triton @ https://repo.radeon.com/rocm/manylinux/rocm-rel-7.1.1/triton-3.5.1%2Brocm7.1.1.gita272dfa8-cp312-cp312-linux_x86_64.whl ; platform_system == 'Linux' and python_version == '3.12' and platform_machine == 'x86_64'", - "triton @ https://repo.radeon.com/rocm/manylinux/rocm-rel-7.1.1/triton-3.5.1%2Brocm7.1.1.gita272dfa8-cp313-cp313-linux_x86_64.whl ; platform_system == 'Linux' and python_version == '3.13' and platform_machine == 'x86_64'", - - "torch @ https://repo.radeon.com/rocm/manylinux/rocm-rel-7.1.1/torch-2.9.1%2Brocm7.1.1.lw.git351ff442-cp310-cp310-linux_x86_64.whl ; platform_system == 'Linux' and python_version == '3.10' and platform_machine == 'x86_64'", - "torch @ https://repo.radeon.com/rocm/manylinux/rocm-rel-7.1.1/torch-2.9.1%2Brocm7.1.1.lw.git351ff442-cp311-cp311-linux_x86_64.whl ; platform_system == 'Linux' and python_version == '3.11' and platform_machine == 'x86_64'", - "torch @ https://repo.radeon.com/rocm/manylinux/rocm-rel-7.1.1/torch-2.9.1%2Brocm7.1.1.lw.git351ff442-cp312-cp312-linux_x86_64.whl ; platform_system == 'Linux' and python_version == '3.12' and platform_machine == 'x86_64'", - "torch @ https://repo.radeon.com/rocm/manylinux/rocm-rel-7.1.1/torch-2.9.1%2Brocm7.1.1.lw.git351ff442-cp313-cp313-linux_x86_64.whl ; platform_system == 'Linux' and python_version == '3.13' and platform_machine == 'x86_64'", - - "torchvision @ https://repo.radeon.com/rocm/manylinux/rocm-rel-7.1.1/torchvision-0.24.0%2Brocm7.1.1.gitb919bd0c-cp310-cp310-linux_x86_64.whl ; platform_system == 'Linux' and python_version == '3.10' and platform_machine == 'x86_64'", - "torchvision @ https://repo.radeon.com/rocm/manylinux/rocm-rel-7.1.1/torchvision-0.24.0%2Brocm7.1.1.gitb919bd0c-cp311-cp311-linux_x86_64.whl ; platform_system == 'Linux' and python_version == '3.11' and platform_machine == 'x86_64'", - "torchvision @ https://repo.radeon.com/rocm/manylinux/rocm-rel-7.1.1/torchvision-0.24.0%2Brocm7.1.1.gitb919bd0c-cp312-cp312-linux_x86_64.whl ; platform_system == 'Linux' and python_version == '3.12' and platform_machine == 'x86_64'", - "torchvision @ https://repo.radeon.com/rocm/manylinux/rocm-rel-7.1.1/torchvision-0.24.0%2Brocm7.1.1.gitb919bd0c-cp313-cp313-linux_x86_64.whl ; platform_system == 'Linux' and python_version == '3.13' and platform_machine == 'x86_64'", -] -rocm72-torch2100 = [ - "unsloth[amd]", - - "triton @ https://repo.radeon.com/rocm/manylinux/rocm-rel-7.2/triton-3.6.0%2Brocm7.2.0.gitba5c1517-cp310-cp310-linux_x86_64.whl ; platform_system == 'Linux' and python_version == '3.10' and platform_machine == 'x86_64'", - "triton @ https://repo.radeon.com/rocm/manylinux/rocm-rel-7.2/triton-3.6.0%2Brocm7.2.0.gitba5c1517-cp311-cp311-linux_x86_64.whl ; platform_system == 'Linux' and python_version == '3.11' and platform_machine == 'x86_64'", - "triton @ https://repo.radeon.com/rocm/manylinux/rocm-rel-7.2/triton-3.6.0%2Brocm7.2.0.gitba5c1517-cp312-cp312-linux_x86_64.whl ; platform_system == 'Linux' and python_version == '3.12' and platform_machine == 'x86_64'", - "triton @ https://repo.radeon.com/rocm/manylinux/rocm-rel-7.2/triton-3.6.0%2Brocm7.2.0.gitba5c1517-cp313-cp313-linux_x86_64.whl ; platform_system == 'Linux' and python_version == '3.13' and platform_machine == 'x86_64'", - - "torch @ https://repo.radeon.com/rocm/manylinux/rocm-rel-7.2/torch-2.10.0%2Brocm7.2.0.lw.gitb6ee5fde-cp310-cp310-linux_x86_64.whl ; platform_system == 'Linux' and python_version == '3.10' and platform_machine == 'x86_64'", - "torch @ https://repo.radeon.com/rocm/manylinux/rocm-rel-7.2/torch-2.10.0%2Brocm7.2.0.lw.gitb6ee5fde-cp311-cp311-linux_x86_64.whl ; platform_system == 'Linux' and python_version == '3.11' and platform_machine == 'x86_64'", - "torch @ https://repo.radeon.com/rocm/manylinux/rocm-rel-7.2/torch-2.10.0%2Brocm7.2.0.lw.gitb6ee5fde-cp312-cp312-linux_x86_64.whl ; platform_system == 'Linux' and python_version == '3.12' and platform_machine == 'x86_64'", - "torch @ https://repo.radeon.com/rocm/manylinux/rocm-rel-7.2/torch-2.10.0%2Brocm7.2.0.lw.gitb6ee5fde-cp313-cp313-linux_x86_64.whl ; platform_system == 'Linux' and python_version == '3.13' and platform_machine == 'x86_64'", - - "torchvision @ https://repo.radeon.com/rocm/manylinux/rocm-rel-7.2/torchvision-0.25.0%2Brocm7.2.0.git82df5f59-cp310-cp310-linux_x86_64.whl ; platform_system == 'Linux' and python_version == '3.10' and platform_machine == 'x86_64'", - "torchvision @ https://repo.radeon.com/rocm/manylinux/rocm-rel-7.2/torchvision-0.25.0%2Brocm7.2.0.git82df5f59-cp311-cp311-linux_x86_64.whl ; platform_system == 'Linux' and python_version == '3.11' and platform_machine == 'x86_64'", - "torchvision @ https://repo.radeon.com/rocm/manylinux/rocm-rel-7.2/torchvision-0.25.0%2Brocm7.2.0.git82df5f59-cp312-cp312-linux_x86_64.whl ; platform_system == 'Linux' and python_version == '3.12' and platform_machine == 'x86_64'", - "torchvision @ https://repo.radeon.com/rocm/manylinux/rocm-rel-7.2/torchvision-0.25.0%2Brocm7.2.0.git82df5f59-cp313-cp313-linux_x86_64.whl ; platform_system == 'Linux' and python_version == '3.13' and platform_machine == 'x86_64'", -] -rocm711-torch2100 = [ - "unsloth[amd]", - - "triton @ https://repo.radeon.com/rocm/manylinux/rocm-rel-7.1.1/triton-3.6.0%2Brocm7.1.1.gitba5c1517-cp310-cp310-linux_x86_64.whl ; platform_system == 'Linux' and python_version == '3.10' and platform_machine == 'x86_64'", - "triton @ https://repo.radeon.com/rocm/manylinux/rocm-rel-7.1.1/triton-3.6.0%2Brocm7.1.1.gitba5c1517-cp311-cp311-linux_x86_64.whl ; platform_system == 'Linux' and python_version == '3.11' and platform_machine == 'x86_64'", - "triton @ https://repo.radeon.com/rocm/manylinux/rocm-rel-7.1.1/triton-3.6.0%2Brocm7.1.1.gitba5c1517-cp312-cp312-linux_x86_64.whl ; platform_system == 'Linux' and python_version == '3.12' and platform_machine == 'x86_64'", - "triton @ https://repo.radeon.com/rocm/manylinux/rocm-rel-7.1.1/triton-3.6.0%2Brocm7.1.1.gitba5c1517-cp313-cp313-linux_x86_64.whl ; platform_system == 'Linux' and python_version == '3.13' and platform_machine == 'x86_64'", - - "torch @ https://repo.radeon.com/rocm/manylinux/rocm-rel-7.1.1/torch-2.10.0%2Brocm7.1.1.lw.gitd9556b05-cp310-cp310-linux_x86_64.whl ; platform_system == 'Linux' and python_version == '3.10' and platform_machine == 'x86_64'", - "torch @ https://repo.radeon.com/rocm/manylinux/rocm-rel-7.1.1/torch-2.10.0%2Brocm7.1.1.lw.gitd9556b05-cp311-cp311-linux_x86_64.whl ; platform_system == 'Linux' and python_version == '3.11' and platform_machine == 'x86_64'", - "torch @ https://repo.radeon.com/rocm/manylinux/rocm-rel-7.1.1/torch-2.10.0%2Brocm7.1.1.lw.gitd9556b05-cp312-cp312-linux_x86_64.whl ; platform_system == 'Linux' and python_version == '3.12' and platform_machine == 'x86_64'", - "torch @ https://repo.radeon.com/rocm/manylinux/rocm-rel-7.1.1/torch-2.10.0%2Brocm7.1.1.lw.gitd9556b05-cp313-cp313-linux_x86_64.whl ; platform_system == 'Linux' and python_version == '3.13' and platform_machine == 'x86_64'", - - "torchvision @ https://repo.radeon.com/rocm/manylinux/rocm-rel-7.1.1/torchvision-0.25.0%2Brocm7.1.1.git82df5f59-cp310-cp310-linux_x86_64.whl ; platform_system == 'Linux' and python_version == '3.10' and platform_machine == 'x86_64'", - "torchvision @ https://repo.radeon.com/rocm/manylinux/rocm-rel-7.1.1/torchvision-0.25.0%2Brocm7.1.1.git82df5f59-cp311-cp311-linux_x86_64.whl ; platform_system == 'Linux' and python_version == '3.11' and platform_machine == 'x86_64'", - "torchvision @ https://repo.radeon.com/rocm/manylinux/rocm-rel-7.1.1/torchvision-0.25.0%2Brocm7.1.1.git82df5f59-cp312-cp312-linux_x86_64.whl ; platform_system == 'Linux' and python_version == '3.12' and platform_machine == 'x86_64'", - "torchvision @ https://repo.radeon.com/rocm/manylinux/rocm-rel-7.1.1/torchvision-0.25.0%2Brocm7.1.1.git82df5f59-cp313-cp313-linux_x86_64.whl ; platform_system == 'Linux' and python_version == '3.13' and platform_machine == 'x86_64'", -] [project.urls] homepage = "https://unsloth.ai" @@ -1310,7 +756,6 @@ repository = "https://github.com/unslothai/unsloth" [tool.ruff] target-version = "py311" -line-length = 100 force-exclude = true extend-exclude = [ "*chat_templates.py", @@ -1343,4 +788,3 @@ ignore = [ # does NOT pick up the GPU-heavy tests under tests/python, tests/qlora, # etc. The CI security job runs `pytest tests/security` explicitly. testpaths = ["tests/security"] -pythonpath = ["."] diff --git a/scripts/check_frontend_dep_removal.py b/scripts/check_frontend_dep_removal.py index b95c4ca7f6..74d089220b 100644 --- a/scripts/check_frontend_dep_removal.py +++ b/scripts/check_frontend_dep_removal.py @@ -52,7 +52,9 @@ EXPECTED_NOISE_FILES = { } # File types where a quoted string can be a module specifier. -JS_LIKE_EXT = re.compile(r"\.(ts|tsx|js|jsx|mjs|cjs|html|htm|css|scss|sass|json|jsonc)$") +JS_LIKE_EXT = re.compile( + r"\.(ts|tsx|js|jsx|mjs|cjs|html|htm|css|scss|sass|json|jsonc)$" +) # Files where JS import patterns could be a real module reference (.mdx is # real ESM; .md code fences are not). SCRIPT_LIKE_EXT = re.compile(r"\.(ts|tsx|js|jsx|mjs|cjs|mdx)$") @@ -249,7 +251,9 @@ def classify(pkg: str, file: str, content: str) -> str | None: if is_script and re.search(rf"\bimport\(\s*['\"]{esc}{sub}['\"]\s*\)", content): return "dynamic_import" # require / require.resolve - if is_script and re.search(rf"\brequire(?:\.resolve)?\(\s*['\"]{esc}{sub}['\"]\s*\)", content): + if is_script and re.search( + rf"\brequire(?:\.resolve)?\(\s*['\"]{esc}{sub}['\"]\s*\)", content + ): return "require" # Re-exports: `export * from`, `export { x } from`, `export type { Foo } from`. if is_script and re.search( @@ -261,12 +265,16 @@ def classify(pkg: str, file: str, content: str) -> str | None: # HTML script / link. Match pkg as a complete path segment so # `/node_modules/foo-extra/...` is not treated as usage of `foo`. html_pkg = rf"{esc}(?:/[^'\"#?]*)?(?=['\"#?])" - if is_html and re.search(rf"]*src\s*=\s*['\"][^'\"]*/{html_pkg}", content): + if is_html and re.search( + rf"]*src\s*=\s*['\"][^'\"]*/{html_pkg}", content + ): return "html_script" if is_html and re.search(rf"]*href\s*=\s*['\"][^'\"]*/{html_pkg}", content): return "html_link" # TypeScript triple-slash - if is_ts and re.search(rf"///\s* str | None: if first in {"npx", "pnpx", "bunx"} and idx + 1 < len(words): idx += 1 continue - if first in {"pnpm", "yarn"} and idx + 2 < len(words) and words[idx + 1] in {"exec", "dlx"}: + if ( + first in {"pnpm", "yarn"} + and idx + 2 < len(words) + and words[idx + 1] in {"exec", "dlx"} + ): idx += 2 continue # 3. Wrapper bin (cross-env, dotenv): skip its flags and env prefixes. - bin_token = first.removeprefix("./node_modules/.bin/").removeprefix("node_modules/.bin/") + bin_token = first.removeprefix("./node_modules/.bin/").removeprefix( + "node_modules/.bin/" + ) if bin_token in _SCRIPT_WRAPPERS and bin_token not in seen_wrappers: seen_wrappers.add(bin_token) idx += 1 @@ -510,7 +524,9 @@ def _next_real_bin(words: list[str], idx: int) -> str | None: return None -def scripts_bin_refs(head_pkg: dict, bin_to_pkg: dict[str, str]) -> dict[str, list[str]]: +def scripts_bin_refs( + head_pkg: dict, bin_to_pkg: dict[str, str] +) -> dict[str, list[str]]: """Return `{package_name: ['scripts.X: cmd', ...]}` for every package referenced via its bin name in package.json scripts. @@ -566,7 +582,11 @@ def tsconfig_compiler_types_refs() -> set[str]: if not isinstance(t, str): continue # `vite/client` resolves to the `vite` package. - pkg = t.split("/", 1)[0] if not t.startswith("@") else "/".join(t.split("/", 2)[:2]) + pkg = ( + t.split("/", 1)[0] + if not t.startswith("@") + else "/".join(t.split("/", 2)[:2]) + ) out.add(pkg) return out @@ -704,7 +724,9 @@ _file_lines_cache: dict[str, list[str]] = {} def _read_file(path: str) -> list[str]: if path not in _file_lines_cache: try: - _file_lines_cache[path] = Path(path).read_text(errors = "replace").splitlines() + _file_lines_cache[path] = ( + Path(path).read_text(errors = "replace").splitlines() + ) except (OSError, UnicodeDecodeError): _file_lines_cache[path] = [] return _file_lines_cache[path] @@ -819,14 +841,18 @@ def find_types_runtime_usage(pkg: str, tsc_types: set[str]) -> list[Hit]: def main() -> int: - p = argparse.ArgumentParser(description = __doc__, formatter_class = argparse.RawTextHelpFormatter) + p = argparse.ArgumentParser( + description = __doc__, formatter_class = argparse.RawTextHelpFormatter + ) p.add_argument( "--base", default = "origin/main", help = "git ref to diff against (default: origin/main). " "Examples: HEAD~1, main, a-tag, a-sha.", ) - p.add_argument("--base-pkg", help = "optional override: read base package.json from this path") + p.add_argument( + "--base-pkg", help = "optional override: read base package.json from this path" + ) p.add_argument( "--base-lock", help = "optional override: read base package-lock.json from this path. " @@ -918,7 +944,9 @@ def main() -> int: print(f" - {w}") print() if missing_imports: - print(f"Imports without a matching package.json dep ({len(missing_imports)}):") + print( + f"Imports without a matching package.json dep ({len(missing_imports)}):" + ) for file, ln, spec in missing_imports[:20]: print(f" - {file}:{ln} imports '{spec}'") print() @@ -956,7 +984,9 @@ def main() -> int: return 1 return 0 - print(f"Checking {len(removed)} removed package(s) from studio/frontend/package.json") + print( + f"Checking {len(removed)} removed package(s) from studio/frontend/package.json" + ) print(f"Base: {args.base} Head: working tree") print() @@ -980,7 +1010,9 @@ def main() -> int: top = f"node_modules/{name}" top_path = top if top in reachable_paths else None nested = sorted( - p for p in reachable_paths if p != top and p.endswith(f"/node_modules/{name}") + p + for p in reachable_paths + if p != top and p.endswith(f"/node_modules/{name}") ) return top_path, nested @@ -1026,7 +1058,9 @@ def main() -> int: _print_hygiene() if failures: - print(f"FAIL: {len(failures)} removed package(s) still referenced and not resolvable") + print( + f"FAIL: {len(failures)} removed package(s) still referenced and not resolvable" + ) for name, _ in failures: print(f" - {name}") return 1 diff --git a/scripts/check_new_install_scripts.py b/scripts/check_new_install_scripts.py index 604d9b9f90..27e505c360 100644 --- a/scripts/check_new_install_scripts.py +++ b/scripts/check_new_install_scripts.py @@ -38,7 +38,9 @@ HIGH = "HIGH" class Finding: __slots__ = ("severity", "name", "version", "kind", "detail") - def __init__(self, severity: str, name: str, version: str, kind: str, detail: str) -> None: + def __init__( + self, severity: str, name: str, version: str, kind: str, detail: str + ) -> None: self.severity = severity self.name = name self.version = version @@ -161,7 +163,9 @@ def diff_new_install_scripts(base_lock: dict, head_lock: dict) -> list[Finding]: if key in base: continue # pre-existing install-script dep; not in scope name = head[key] - version = key[len(name) + 1 :] if key.startswith(name + "@") else "" + version = ( + key[len(name) + 1 :] if key.startswith(name + "@") else "" + ) scripts = _fetch_registry_scripts(name, version) if scripts: detail = "; ".join(f"{h}={cmd!r}" for h, cmd in scripts.items()) diff --git a/scripts/enforce_kwargs_spacing.py b/scripts/enforce_kwargs_spacing.py index dc4a6d6821..64f7eb757e 100755 --- a/scripts/enforce_kwargs_spacing.py +++ b/scripts/enforce_kwargs_spacing.py @@ -123,7 +123,9 @@ def remove_redundant_passes(text: str) -> tuple[str, bool]: lines = text.splitlines(keepends=True) changed = False - for node in sorted(redundant, key=lambda item: (item.lineno, item.col_offset), reverse=True): + for node in sorted( + redundant, key=lambda item: (item.lineno, item.col_offset), reverse=True + ): start = node.lineno - 1 end = (node.end_lineno or node.lineno) - 1 if start >= len(lines): @@ -181,7 +183,11 @@ def remove_blank_after_short_import(text: str) -> tuple[str, bool]: out: list[list[ast.stmt]] = [] for attr in ("body", "orelse", "finalbody"): val = getattr(node, attr, None) - if isinstance(val, list) and val and all(isinstance(s, ast.stmt) for s in val): + if ( + isinstance(val, list) + and val + and all(isinstance(s, ast.stmt) for s in val) + ): out.append(val) return out @@ -199,7 +205,9 @@ def remove_blank_after_short_import(text: str) -> tuple[str, bool]: j += 1 if j + 1 < len(suite): # an import block followed by another statement last_imp, nxt = suite[j], suite[j + 1] - gap = range((last_imp.end_lineno or last_imp.lineno) + 1, nxt.lineno) + gap = range( + (last_imp.end_lineno or last_imp.lineno) + 1, nxt.lineno + ) nums = [n for n in gap if 1 <= n <= len(lines)] if nums and all(lines[n - 1].strip() == "" for n in nums): drop.update(nums) @@ -211,7 +219,13 @@ def remove_blank_after_short_import(text: str) -> tuple[str, bool]: return "".join(kept), True -_STRING_TRIVIA = (tokenize.NL, tokenize.NEWLINE, tokenize.COMMENT, tokenize.INDENT, tokenize.DEDENT) +_STRING_TRIVIA = ( + tokenize.NL, + tokenize.NEWLINE, + tokenize.COMMENT, + tokenize.INDENT, + tokenize.DEDENT, +) _DEF_MIN_PARAMS_FOR_MULTILINE = 3 # signatures with < this many params stay one line diff --git a/scripts/lint_workflow_triggers.py b/scripts/lint_workflow_triggers.py index 0688f6c65c..fcdf71e891 100644 --- a/scripts/lint_workflow_triggers.py +++ b/scripts/lint_workflow_triggers.py @@ -29,7 +29,9 @@ from pathlib import Path try: import yaml except ImportError: - print("ERROR: PyYAML is required. Install with 'pip install pyyaml'", file = sys.stderr) + print( + "ERROR: PyYAML is required. Install with 'pip install pyyaml'", file = sys.stderr + ) sys.exit(2) REPO_ROOT = Path(__file__).resolve().parents[1] @@ -133,7 +135,9 @@ def main() -> int: ) if findings: - print("Workflow trigger lint failed with the following issues:", file = sys.stderr) + print( + "Workflow trigger lint failed with the following issues:", file = sys.stderr + ) for f in findings: print(f" - {f}", file = sys.stderr) return 1 diff --git a/scripts/lockfile_supply_chain_audit.py b/scripts/lockfile_supply_chain_audit.py index 66b48c094d..f9004ebe16 100644 --- a/scripts/lockfile_supply_chain_audit.py +++ b/scripts/lockfile_supply_chain_audit.py @@ -459,7 +459,9 @@ def audit_npm_lockfile(path: Path) -> list[Finding]: path = str(path), package = key, kind = "blocked-known-malicious", - detail = (f"{pkg_name}@{version} is on the BLOCKED_NPM_VERSIONS list"), + detail = ( + f"{pkg_name}@{version} is on the BLOCKED_NPM_VERSIONS list" + ), ) ) @@ -663,7 +665,9 @@ def main(argv: list[str] | None = None) -> int: "--cargo-lockfile", action = "append", default = None, - help = ("Path to a Cargo.lock (repeatable). Default: studio/src-tauri/Cargo.lock."), + help = ( + "Path to a Cargo.lock (repeatable). Default: studio/src-tauri/Cargo.lock." + ), ) parser.add_argument( "--strict", diff --git a/scripts/notebook_to_python.py b/scripts/notebook_to_python.py index 4b64123d6a..d4da846516 100644 --- a/scripts/notebook_to_python.py +++ b/scripts/notebook_to_python.py @@ -155,7 +155,9 @@ def convert_cell_to_python(source: str, *, allow_shell: bool = True) -> str: cmd_lines.append(lines[i].strip()) full_cmd = "\n".join(cmd_lines) - result.extend(_emit_shell_command(indent, full_cmd, allow_shell = allow_shell)) + result.extend( + _emit_shell_command(indent, full_cmd, allow_shell = allow_shell) + ) # %cd path -> os.chdir(path) elif stripped.startswith("%cd "): @@ -278,7 +280,9 @@ def convert_notebook_to_script( source_name = source output_filename = filename.replace(".ipynb", ".py") - output_filename = output_filename.replace("(", "").replace(")", "").replace("-", "_") + output_filename = ( + output_filename.replace("(", "").replace(")", "").replace("-", "_") + ) if output_dir: output_path = os.path.join(output_dir, output_filename) @@ -297,7 +301,9 @@ def convert_notebook_to_script( def main(): import argparse - class Formatter(argparse.ArgumentDefaultsHelpFormatter, argparse.RawDescriptionHelpFormatter): + class Formatter( + argparse.ArgumentDefaultsHelpFormatter, argparse.RawDescriptionHelpFormatter + ): pass parser = argparse.ArgumentParser( @@ -311,8 +317,12 @@ Examples: python notebook_to_python.py https://github.com/unslothai/notebooks/blob/main/nb/Oute_TTS_(1B).ipynb """, ) - parser.add_argument("notebooks", nargs = "+", help = "Notebook files or URLs to convert.") - parser.add_argument("-o", "--output", dest = "output_dir", default = ".", help = "Output directory.") + parser.add_argument( + "notebooks", nargs = "+", help = "Notebook files or URLs to convert." + ) + parser.add_argument( + "-o", "--output", dest = "output_dir", default = ".", help = "Output directory." + ) # Default True for backwards compat; pass --no-allow-shell for untrusted notebooks. parser.add_argument( "--allow-shell", diff --git a/scripts/notebook_validator.py b/scripts/notebook_validator.py index c1be7a63a4..0ac52072fa 100644 --- a/scripts/notebook_validator.py +++ b/scripts/notebook_validator.py @@ -87,7 +87,9 @@ COLAB_ORACLE_FILES: dict[str, str] = { "apt-list-gpu.txt": "colab_apt_list.gpu.txt", "os-info-gpu.txt": "colab_os_info.gpu.txt", } -COLAB_ORACLE_BASE_URL = "https://raw.githubusercontent.com/googlecolab/backend-info/main/" +COLAB_ORACLE_BASE_URL = ( + "https://raw.githubusercontent.com/googlecolab/backend-info/main/" +) # ----- Compat tables. PRs add rows as new releases land. ----- # @@ -187,7 +189,9 @@ def install_cells(nb: dict[str, Any]) -> list[tuple[int, str]]: if first and first[0].strip().startswith("%%capture"): out.append((i, src)) continue - if re.search(r"^[ \t]*!\s*(uv\s+)?pip\s+(install|uninstall)\b", src, re.MULTILINE): + if re.search( + r"^[ \t]*!\s*(uv\s+)?pip\s+(install|uninstall)\b", src, re.MULTILINE + ): out.append((i, src)) return out @@ -318,7 +322,9 @@ def parse_pip_line(line: str, line_no: int = 0) -> PipInvocation | None: if t in ("install", "uninstall"): continue packages.append(t) - return PipInvocation(tool = tool, flags = flags, packages = packages, raw = line, line_no = line_no) + return PipInvocation( + tool = tool, flags = flags, packages = packages, raw = line, line_no = line_no + ) def _glue_line_continuations(text: str) -> list[tuple[int, str]]: @@ -403,7 +409,9 @@ def pypi_metadata(name: str, version: str) -> dict[str, Any] | None: return data -def transitive_constraint(name: str, version: str, target: str) -> tuple[str | None, list[str]]: +def transitive_constraint( + name: str, version: str, target: str +) -> tuple[str | None, list[str]]: """Return (raw_specifier_string_or_None, list_of_(op,version) tuples) for the constraint that `name==version` places on `target`. """ @@ -477,7 +485,10 @@ def resolved_set(install_cell: str, colab: dict[str, str]) -> dict[str, str]: out[sp.name] = ver pinned.add(sp.name) elif op == "<=" and sp.name not in pinned: - if sp.name not in upper_bounds or cmp_versions(ver, upper_bounds[sp.name]) < 0: + if ( + sp.name not in upper_bounds + or cmp_versions(ver, upper_bounds[sp.name]) < 0 + ): upper_bounds[sp.name] = ver # Apply upper bounds where Colab's preinstall violates them. for name, ub in upper_bounds.items(): @@ -492,7 +503,9 @@ def resolved_set(install_cell: str, colab: dict[str, str]) -> dict[str, str]: # ----- Rules ----- # -def rule_inst_001_git_plus(install_cell: str, file: str, cell_idx: int) -> list[Finding]: +def rule_inst_001_git_plus( + install_cell: str, file: str, cell_idx: int +) -> list[Finding]: findings: list[Finding] = [] for inv in iter_pip_invocations(install_cell): if any("git+" in p for p in inv.packages) or "git+" in inv.raw: @@ -680,7 +693,9 @@ def rule_inst_005_transformers_tokenizers( _RE_DOUBLE_BANG = re.compile(r"^[ \t]*!{2,}\s*pip\b", re.MULTILINE) -def rule_inst_006_double_bang(install_cell: str, file: str, cell_idx: int) -> list[Finding]: +def rule_inst_006_double_bang( + install_cell: str, file: str, cell_idx: int +) -> list[Finding]: findings: list[Finding] = [] for m in _RE_DOUBLE_BANG.finditer(install_cell): line_no = install_cell.count("\n", 0, m.start()) + 1 @@ -771,7 +786,9 @@ POLICY_CLAUSES_DEFAULT = [ ] -def extract_policy_clauses(update_script: pathlib.Path) -> list[tuple[str, re.Pattern[str], Any]]: +def extract_policy_clauses( + update_script: pathlib.Path, +) -> list[tuple[str, re.Pattern[str], Any]]: """Best-effort scan of update_all_notebooks.py for canonical phrases; falls back to POLICY_CLAUSES_DEFAULT (which we use directly today). The permissive regexes avoid false positives on template rewords.""" @@ -831,7 +848,11 @@ def cmd_drift(args: argparse.Namespace) -> int: print(f"FAIL: {update_script} not found", file = sys.stderr) return 2 # Stash any pre-existing dirty state, run the updater, diff, restore. - head = subprocess.check_output(["git", "rev-parse", "HEAD"], cwd = nbdir).decode().strip() + head = ( + subprocess.check_output(["git", "rev-parse", "HEAD"], cwd = nbdir) + .decode() + .strip() + ) subprocess.run( ["git", "-C", str(nbdir), "stash", "--include-untracked"], check = False, @@ -932,7 +953,9 @@ def cmd_convert(args: argparse.Namespace) -> int: hint = proc.stderr[-200:].strip(), ) ) - print(f"converted {len(notebooks) - len(failed)}/{len(notebooks)} notebooks to {out}") + print( + f"converted {len(notebooks) - len(failed)}/{len(notebooks)} notebooks to {out}" + ) _emit(failed) return 0 if not failed else 1 @@ -942,7 +965,11 @@ def cmd_convert(args: argparse.Namespace) -> int: def cmd_lint(args: argparse.Namespace) -> int: nbdir = pathlib.Path(args.notebooks_dir).resolve() - colab_path = pathlib.Path(args.colab_pin).resolve() if args.colab_pin else COLAB_FALLBACK_FILE + colab_path = ( + pathlib.Path(args.colab_pin).resolve() + if args.colab_pin + else COLAB_FALLBACK_FILE + ) colab = parse_pip_freeze(colab_path) if not colab: print( @@ -982,9 +1009,13 @@ def cmd_lint(args: argparse.Namespace) -> int: first_cell = cells[0][0] if cells else None findings += rule_inst_003_peft_torchao(merged, oracle, rel, first_cell) findings += rule_inst_004_torchcodec_torch(merged, oracle, rel, first_cell) - findings += rule_inst_005_transformers_tokenizers(merged, oracle, rel, first_cell) + findings += rule_inst_005_transformers_tokenizers( + merged, oracle, rel, first_cell + ) if not args.no_pypi: - findings += rule_inst_002_no_deps_transitive(merged, oracle, rel, first_cell) + findings += rule_inst_002_no_deps_transitive( + merged, oracle, rel, first_cell + ) findings += scan_user_cells(nb, rel) _emit(findings) return 0 if not any(f.severity == "error" for f in findings) else 1 @@ -1159,7 +1190,9 @@ def cmd_colab_diff(args: argparse.Namespace) -> int: print(f"::warning::colab-diff: could not fetch {url}: {e}") continue if not snap_path.exists(): - print(f"::warning::colab-diff: no committed snapshot at {snap_path}; skipping") + print( + f"::warning::colab-diff: no committed snapshot at {snap_path}; skipping" + ) continue snapshot_text = snap_path.read_text(encoding = "utf-8", errors = "replace") parser = _COLAB_ORACLE_PARSERS[upstream_name] diff --git a/scripts/scan_npm_packages.py b/scripts/scan_npm_packages.py index fe90afa7e6..1db86338a8 100644 --- a/scripts/scan_npm_packages.py +++ b/scripts/scan_npm_packages.py @@ -768,7 +768,8 @@ def download_tarball( written += len(chunk) if written > max_bytes: return dest, ( - f"download exceeded cap {max_bytes} bytes " f"after {written} bytes" + f"download exceeded cap {max_bytes} bytes " + f"after {written} bytes" ) h.update(chunk) out.write(chunk) @@ -865,7 +866,11 @@ def safe_extract( # each gets its own cap (both are bounded). header = src.read(16) is_binary = _looks_binary(name, header) - file_cap = HARD_MAX_BINARY_FILE_BYTES if is_binary else HARD_MAX_TEXT_FILE_BYTES + file_cap = ( + HARD_MAX_BINARY_FILE_BYTES + if is_binary + else HARD_MAX_TEXT_FILE_BYTES + ) if declared > file_cap: return ( f"member {name!r} declared size {declared} > " @@ -988,7 +993,9 @@ def scan_package_json(pkg: PackageEntry, rel: str, text: str) -> list[Finding]: if isinstance(opt, dict): for k, v in opt.items(): if isinstance(v, str) and ( - v.startswith("github:") or v.startswith("git+") or v.startswith("git://") + v.startswith("github:") + or v.startswith("git+") + or v.startswith("git://") ): findings.append( Finding( @@ -1104,7 +1111,9 @@ def scan_text_blob(pkg: PackageEntry, rel: str, text: str) -> list[Finding]: filename = rel, pattern = "js-fetch-eval", evidence = _evidence(text, _JS_FETCH_EVAL), - detail = ("Function/eval against base64-decoded payload (obfuscated dropper shape)"), + detail = ( + "Function/eval against base64-decoded payload (obfuscated dropper shape)" + ), ) ) if _JS_ENV_TOKEN.search(text): @@ -1345,7 +1354,8 @@ def main(argv: list[str] | None = None) -> int: if hard_errors or blocking: if blocking: print( - f"\n[scan-npm] FAIL: {len(blocking)} finding(s) " f"at or above {threshold}", + f"\n[scan-npm] FAIL: {len(blocking)} finding(s) " + f"at or above {threshold}", file = sys.stderr, ) return 1 diff --git a/scripts/scan_packages.py b/scripts/scan_packages.py index 861b35617b..91bd6fa607 100644 --- a/scripts/scan_packages.py +++ b/scripts/scan_packages.py @@ -293,7 +293,9 @@ RE_CRYPTO_THEFT = re.compile( RE_PTH_IMPORT = re.compile(r"^\s*import\s+", re.MULTILINE) # openssl CLI invocations via subprocess (encrypted exfiltration) -RE_OPENSSL_CLI = re.compile(r"\bopenssl\s+(enc|rand|rsautl|pkeyutl|genrsa|dgst|s_client)\b") +RE_OPENSSL_CLI = re.compile( + r"\bopenssl\s+(enc|rand|rsautl|pkeyutl|genrsa|dgst|s_client)\b" +) # Write to /tmp then execute (staged dropper) RE_TEMP_EXEC = re.compile( @@ -1224,13 +1226,15 @@ def iter_archive_files(archive_path: str): # historically dereferenced them on extract. if member.issym() or member.islnk(): print( - f" [WARN] {path.name}: refused link member " f"{member.name!r}", + f" [WARN] {path.name}: refused link member " + f"{member.name!r}", file = sys.stderr, ) continue if member.isdev() or member.isfifo(): print( - f" [WARN] {path.name}: refused special member " f"{member.name!r}", + f" [WARN] {path.name}: refused special member " + f"{member.name!r}", file = sys.stderr, ) continue @@ -1432,7 +1436,9 @@ def download_packages( env = env, ) if proc.returncode != 0: - msg = f"pip download (with deps) failed: " f"{proc.stderr.strip()[:500]}" + msg = ( + f"pip download (with deps) failed: " f"{proc.stderr.strip()[:500]}" + ) print(f" [ERROR] {msg}", file = sys.stderr) download_errors.append(msg) except subprocess.TimeoutExpired: @@ -1473,7 +1479,10 @@ def download_packages( env = env, ) if proc.returncode != 0: - msg = f"pip download failed for {spec}: " f"{proc.stderr.strip()[:500]}" + msg = ( + f"pip download failed for {spec}: " + f"{proc.stderr.strip()[:500]}" + ) print(f" [ERROR] {msg}", file = sys.stderr) download_errors.append(msg) continue @@ -1499,7 +1508,9 @@ def _extract_pkg_name(spec: str) -> str: """Extract the package name from a pip spec string.""" m = _RE_NAME.match(spec) return ( - m.group(1) if m else spec.split("==")[0].split(">=")[0].split("<=")[0].split("[")[0].strip() + m.group(1) + if m + else spec.split("==")[0].split(">=")[0].split("<=")[0].split("[")[0].strip() ) @@ -1840,7 +1851,9 @@ def _run_fix(critical_pkgs: set[str], entries: list[dict], max_search: int) -> N if git_entries: for e in git_entries: src = e["source_file"] or "CLI" - print(f" [SKIP] {pkg_name} is a git URL dep in {src}, cannot auto-update") + print( + f" [SKIP] {pkg_name} is a git URL dep in {src}, cannot auto-update" + ) changes_summary.append(f" SKIP {pkg_name} (git URL)") continue @@ -1862,7 +1875,9 @@ def _run_fix(critical_pkgs: set[str], entries: list[dict], max_search: int) -> N shutil.rmtree(dl_dir, ignore_errors = True) if not current_ver: - print(f" [WARN] Cannot determine current version of {pkg_name}, skipping fix") + print( + f" [WARN] Cannot determine current version of {pkg_name}, skipping fix" + ) changes_summary.append(f" SKIP {pkg_name} (version unknown)") continue @@ -1879,7 +1894,9 @@ def _run_fix(critical_pkgs: set[str], entries: list[dict], max_search: int) -> N continue print(f" [OK] {pkg_name}: {current_ver} -> {safe_ver}") - changes_summary.append(f" FIX {pkg_name}=={current_ver} -> {pkg_name}=={safe_ver}") + changes_summary.append( + f" FIX {pkg_name}=={current_ver} -> {pkg_name}=={safe_ver}" + ) # Update all occurrences in requirements files file_updates: dict[str, dict[int, str]] = {} @@ -1926,7 +1943,9 @@ def _find_requirements_files(root: str) -> list[str]: dirnames[:] = [ d for d in dirnames - if not d.startswith(".") and d not in skip_dirs and not d.endswith(".egg-info") + if not d.startswith(".") + and d not in skip_dirs + and not d.endswith(".egg-info") ] dirname = os.path.basename(dirpath) for fname in sorted(filenames): @@ -1998,7 +2017,9 @@ def main() -> int: print(f" {f}") req_files.extend(found) else: - print(f" [WARN] No requirements files found in {scan_dir}/", file = sys.stderr) + print( + f" [WARN] No requirements files found in {scan_dir}/", file = sys.stderr + ) # Build unified entry list: list of dicts with source tracking entries: list[dict] = [] diff --git a/scripts/stamp_studio_release.py b/scripts/stamp_studio_release.py index 7dab35ea8a..112c83f125 100644 --- a/scripts/stamp_studio_release.py +++ b/scripts/stamp_studio_release.py @@ -42,7 +42,9 @@ def _atomic_write_text( REPO_ROOT = Path(__file__).resolve().parents[1] -BUILD_INFO_PATH = REPO_ROOT / "studio" / "backend" / "utils" / "_studio_release_build.py" +BUILD_INFO_PATH = ( + REPO_ROOT / "studio" / "backend" / "utils" / "_studio_release_build.py" +) BUILD_INFO_SUFFIX = "studio/backend/utils/_studio_release_build.py" VERSION_RE = re.compile(r"^v\d+\.\d+\.\d+(?:-[0-9A-Za-z.][0-9A-Za-z.-]*)?$") GIT_DESCRIBE_SUFFIX_RE = re.compile(r"-\d+-g[0-9A-Fa-f]+(?:-dirty)?$") diff --git a/scripts/sync_allow_scripts_pins.py b/scripts/sync_allow_scripts_pins.py index 1d9a075ab4..22beecaa75 100644 --- a/scripts/sync_allow_scripts_pins.py +++ b/scripts/sync_allow_scripts_pins.py @@ -74,7 +74,9 @@ def desired_key(name: str, versions: list[str]) -> str: return f"{name}@{' || '.join(versions)}" -def compute_renames(policy: dict, lock_versions: dict[str, list[str]]) -> dict[str, str]: +def compute_renames( + policy: dict, lock_versions: dict[str, list[str]] +) -> dict[str, str]: renames: dict[str, str] = {} for key in policy: name, rng = split_spec(key) @@ -93,7 +95,9 @@ def main(argv: list[str] | None = None) -> int: ap = argparse.ArgumentParser(description = __doc__) mode = ap.add_mutually_exclusive_group(required = True) mode.add_argument("--check", action = "store_true", help = "exit 1 if pins are stale") - mode.add_argument("--fix", action = "store_true", help = "rewrite package.json in place") + mode.add_argument( + "--fix", action = "store_true", help = "rewrite package.json in place" + ) ap.add_argument( "--dir", type = Path, @@ -105,20 +109,26 @@ def main(argv: list[str] | None = None) -> int: pkg_path = args.dir / "package.json" lock_path = args.dir / "package-lock.json" if not pkg_path.exists() or not lock_path.exists(): - print(f"sync-allow-scripts: nothing to do ({args.dir} has no package.json + lockfile)") + print( + f"sync-allow-scripts: nothing to do ({args.dir} has no package.json + lockfile)" + ) return 0 pkg = json.loads(pkg_path.read_text(encoding = "utf-8")) policy = pkg.get("allowScripts") if not isinstance(policy, dict) or not policy: - print("sync-allow-scripts: no allowScripts policy in package.json, nothing to do") + print( + "sync-allow-scripts: no allowScripts policy in package.json, nothing to do" + ) return 0 lock = json.loads(lock_path.read_text(encoding = "utf-8")) renames = compute_renames(policy, script_versions_from_lock(lock)) if not renames: - print(f"sync-allow-scripts: {len(policy)} allowScripts entries in sync with the lockfile") + print( + f"sync-allow-scripts: {len(policy)} allowScripts entries in sync with the lockfile" + ) return 0 for old, new in renames.items(): @@ -132,7 +142,9 @@ def main(argv: list[str] | None = None) -> int: return 1 pkg["allowScripts"] = {renames.get(k, k): v for k, v in policy.items()} - pkg_path.write_text(json.dumps(pkg, indent = 2, ensure_ascii = False) + "\n", encoding = "utf-8") + pkg_path.write_text( + json.dumps(pkg, indent = 2, ensure_ascii = False) + "\n", encoding = "utf-8" + ) print( f"sync-allow-scripts: re-pinned {len(renames)} entr{'y' if len(renames) == 1 else 'ies'} in {pkg_path}" ) diff --git a/scripts/verify_comment_only_diff.py b/scripts/verify_comment_only_diff.py index 23d06b85df..0ec49b5755 100644 --- a/scripts/verify_comment_only_diff.py +++ b/scripts/verify_comment_only_diff.py @@ -131,7 +131,8 @@ def _walk_yaml_diff( """Print a path-keyed summary of the first structural / scalar diff.""" if type(b) is not type(a): print( - f" type-diff at {prefix or '/'}: " f"{type(b).__name__} -> {type(a).__name__}", + f" type-diff at {prefix or '/'}: " + f"{type(b).__name__} -> {type(a).__name__}", ) return if isinstance(b, dict): diff --git a/scripts/verify_import_hoist.py b/scripts/verify_import_hoist.py index b8cb0573fe..6c16f439e4 100644 --- a/scripts/verify_import_hoist.py +++ b/scripts/verify_import_hoist.py @@ -161,7 +161,9 @@ class _Builder(ast.NodeVisitor): def _visit_stmt(self, node: ast.AST, scope: Scope) -> None: if isinstance(node, (ast.Import, ast.ImportFrom)): - star = isinstance(node, ast.ImportFrom) and any(a.name == "*" for a in node.names) + star = isinstance(node, ast.ImportFrom) and any( + a.name == "*" for a in node.names + ) if star: scope.star_import = True for alias in node.names: @@ -349,7 +351,9 @@ class _Builder(ast.NodeVisitor): self._bind_args(node.args, child) self._visit_expr(node.body, child) return - if isinstance(node, (ast.ListComp, ast.SetComp, ast.GeneratorExp, ast.DictComp)): + if isinstance( + node, (ast.ListComp, ast.SetComp, ast.GeneratorExp, ast.DictComp) + ): child = Scope("comp", f"{scope.qualname}.", scope) for i, gen in enumerate(node.generators): # first iterable evaluates in the enclosing scope @@ -606,7 +610,9 @@ def compare(before_src: str, after_src: str, path: str) -> list[tuple[str, str]] for scope, names in b["ambiguous"].items(): new = names - a["ambiguous"].get(scope, set()) for n in sorted(new): - findings.append(("WARN", f"{path}: AMBIGUOUS-BIND '{n}' import+non-import in {scope}")) + findings.append( + ("WARN", f"{path}: AMBIGUOUS-BIND '{n}' import+non-import in {scope}") + ) # 6. TARGET-MISSING (informational): a scope stopped resolving to an import # target. Real bugs are covered above; remaining cases are relocated code. @@ -618,7 +624,9 @@ def compare(before_src: str, after_src: str, path: str) -> list[tuple[str, str]] if t in added_module_targets else " [target not re-added here -> likely relocated/deleted]" ) - findings.append(("INFO", f"{path}: TARGET-MISSING {t} in scope {scope}{relocated}")) + findings.append( + ("INFO", f"{path}: TARGET-MISSING {t} in scope {scope}{relocated}") + ) return findings @@ -763,7 +771,9 @@ def audit_files(paths: list[str]) -> int: ok = n_err == 0 and n_fp == 0 print( "\nAUDIT:", - "ROBUST (no crashes, no false positives vs pyflakes)" if ok else "NEEDS WORK (see above)", + "ROBUST (no crashes, no false positives vs pyflakes)" + if ok + else "NEEDS WORK (see above)", ) return 0 if ok else 1 @@ -799,12 +809,18 @@ def main() -> int: blockers = [f for f in findings if f[0] == "BLOCKER"] warns = [f for f in findings if f[0] == "WARN"] infos = [f for f in findings if f[0] == "INFO"] - status = "CLEAN" if not blockers and not warns else ("BLOCKERS" if blockers else "WARNINGS") + status = ( + "CLEAN" + if not blockers and not warns + else ("BLOCKERS" if blockers else "WARNINGS") + ) print(f"\n=== {path}: {status} ===") for sev, m in blockers + warns + infos: print(f" [{sev}] {m}") any_blocker = any_blocker or bool(blockers) - print("\nOVERALL:", "FAIL (blockers found)" if any_blocker else "PASS (no blockers)") + print( + "\nOVERALL:", "FAIL (blockers found)" if any_blocker else "PASS (no blockers)" + ) return 1 if any_blocker else 0 diff --git a/studio/backend/auth/authentication.py b/studio/backend/auth/authentication.py index 9dd56489eb..8f2eb3b51d 100644 --- a/studio/backend/auth/authentication.py +++ b/studio/backend/auth/authentication.py @@ -108,7 +108,9 @@ def create_refresh_token(subject: str, *, desktop: bool = False) -> str: return token -def refresh_access_token(refresh_token: str) -> Tuple[Optional[str], Optional[str], bool]: +def refresh_access_token( + refresh_token: str, +) -> Tuple[Optional[str], Optional[str], bool]: """ Validate a refresh token and issue a new access token. @@ -135,7 +137,9 @@ def reload_secret() -> None: load_jwt_secret() -async def get_current_subject(credentials: HTTPAuthorizationCredentials = Depends(security)) -> str: +async def get_current_subject( + credentials: HTTPAuthorizationCredentials = Depends(security), +) -> str: """Validate JWT and require the password-change flow to be completed.""" return await _get_current_subject( credentials, diff --git a/studio/backend/auth/storage.py b/studio/backend/auth/storage.py index fa5b985513..90994e61ed 100644 --- a/studio/backend/auth/storage.py +++ b/studio/backend/auth/storage.py @@ -138,9 +138,13 @@ def get_connection() -> sqlite3.Connection: ); """ ) - api_key_columns = {row["name"] for row in conn.execute("PRAGMA table_info(api_keys)")} + api_key_columns = { + row["name"] for row in conn.execute("PRAGMA table_info(api_keys)") + } if "is_internal" not in api_key_columns: - conn.execute("ALTER TABLE api_keys ADD COLUMN is_internal INTEGER NOT NULL DEFAULT 0") + conn.execute( + "ALTER TABLE api_keys ADD COLUMN is_internal INTEGER NOT NULL DEFAULT 0" + ) conn.execute( """ CREATE TABLE IF NOT EXISTS app_secrets ( @@ -154,9 +158,13 @@ def get_connection() -> sqlite3.Connection: conn.execute( "ALTER TABLE auth_user ADD COLUMN must_change_password INTEGER NOT NULL DEFAULT 0" ) - refresh_columns = {row["name"] for row in conn.execute("PRAGMA table_info(refresh_tokens)")} + refresh_columns = { + row["name"] for row in conn.execute("PRAGMA table_info(refresh_tokens)") + } if "is_desktop" not in refresh_columns: - conn.execute("ALTER TABLE refresh_tokens ADD COLUMN is_desktop INTEGER NOT NULL DEFAULT 0") + conn.execute( + "ALTER TABLE refresh_tokens ADD COLUMN is_desktop INTEGER NOT NULL DEFAULT 0" + ) conn.commit() return conn diff --git a/studio/backend/cloudflare_tunnel.py b/studio/backend/cloudflare_tunnel.py index e5dba69452..c963734751 100644 --- a/studio/backend/cloudflare_tunnel.py +++ b/studio/backend/cloudflare_tunnel.py @@ -77,7 +77,9 @@ def _asset_name() -> Optional[Tuple[str, bool]]: def _cache_path() -> Optional[Path]: """studio_bin_root()/cloudflared(.exe), or None if the studio home is unresolvable.""" try: - from utils.paths.storage_roots import studio_bin_root # lazy: backend-only import + from utils.paths.storage_roots import ( + studio_bin_root, + ) # lazy: backend-only import except Exception: return None name = "cloudflared.exe" if sys.platform == "win32" else "cloudflared" @@ -262,7 +264,9 @@ class CloudflareTunnel: if self.url is None: self.error = "cloudflared exited before emitting a tunnel URL" elif not self.ready: - self.error = "cloudflared exited before the tunnel connection registered" + self.error = ( + "cloudflared exited before the tunnel connection registered" + ) self._url_event.set() self._ready_event.set() diff --git a/studio/backend/colab.py b/studio/backend/colab.py index ba46c52a6a..eb93121781 100644 --- a/studio/backend/colab.py +++ b/studio/backend/colab.py @@ -41,7 +41,12 @@ def get_colab_url(port: int = 8888) -> str: try: url = eval_js(f"google.colab.kernel.proxyPort({port})", timeout_sec = 10) # Valid proxy URL is https:// and embeds the port. - if url and isinstance(url, str) and url.startswith("https://") and str(port) in url: + if ( + url + and isinstance(url, str) + and url.startswith("https://") + and str(port) in url + ): return url.rstrip("/") except Exception as e: logger.info(f"Note: Could not get Colab URL (attempt {attempt + 1}/3: {e})") @@ -107,7 +112,9 @@ def _is_studio_healthy(port: int, timeout: float = 2.0) -> bool: """Return True if a Studio backend is already answering health checks on *port*.""" import urllib.request try: - with urllib.request.urlopen(f"http://localhost:{port}/api/health", timeout = timeout): + with urllib.request.urlopen( + f"http://localhost:{port}/api/health", timeout = timeout + ): return True except Exception: return False @@ -179,7 +186,9 @@ def start(port: int = 8888): # Fast path: Studio already running (cell re-run). Re-launching would collide on # the port, so just re-show the link and iframe. if _is_studio_healthy(port): - logger.info(f" Studio is already running on port {port} — reusing existing server.") + logger.info( + f" Studio is already running on port {port} — reusing existing server." + ) _show_and_embed(port) try: for _ in range(10000): @@ -202,7 +211,9 @@ def start(port: int = 8888): logger.info(" Starting server...") try: - app = run_server(host = "0.0.0.0", port = port, frontend_path = frontend_path, silent = True) + app = run_server( + host = "0.0.0.0", port = port, frontend_path = frontend_path, silent = True + ) except SystemExit as exc: logger.error(f"❌ Unsloth Studio failed to start: {exc}") return @@ -223,7 +234,9 @@ def start(port: int = 8888): server_ready = False for _ in range(40): try: - with urllib.request.urlopen(f"http://localhost:{actual_port}/api/health", timeout = 1): + with urllib.request.urlopen( + f"http://localhost:{actual_port}/api/health", timeout = 1 + ): server_ready = True break except Exception: diff --git a/studio/backend/core/data_recipe/huggingface.py b/studio/backend/core/data_recipe/huggingface.py index 7a1219b2c3..0f17f8dc7c 100644 --- a/studio/backend/core/data_recipe/huggingface.py +++ b/studio/backend/core/data_recipe/huggingface.py @@ -36,7 +36,9 @@ def _resolve_recipe_artifact_path(artifact_path: str) -> Path: if not resolved.exists(): raise RecipeDatasetPublishError("Execution artifacts are no longer available.") if not resolved.is_dir(): - raise RecipeDatasetPublishError("Execution artifact path is not a dataset folder.") + raise RecipeDatasetPublishError( + "Execution artifact path is not a dataset folder." + ) return resolved diff --git a/studio/backend/core/data_recipe/jobs/manager.py b/studio/backend/core/data_recipe/jobs/manager.py index 32523e469e..60a1740f3c 100644 --- a/studio/backend/core/data_recipe/jobs/manager.py +++ b/studio/backend/core/data_recipe/jobs/manager.py @@ -108,7 +108,9 @@ class Subscription: event_id = self._next_id body = json.dumps(event, separators = (",", ":"), ensure_ascii = False) event_type = event.get("type") or "message" - return (f"id: {event_id}\n" f"event: {event_type}\n" f"data: {body}\n\n").encode("utf-8") + return ( + f"id: {event_id}\n" f"event: {event_type}\n" f"data: {body}\n\n" + ).encode("utf-8") class JobManager: @@ -155,7 +157,9 @@ class JobManager: job_id = uuid.uuid4().hex self._job = Job(job_id = job_id, status = "pending", started_at = time.time()) self._job.progress_columns_total = llm_column_count - self._job.source_progress_estimated_total = _github_source_estimated_total(recipe) + self._job.source_progress_estimated_total = _github_source_estimated_total( + recipe + ) self._job.internal_api_key_id = internal_api_key_id self._events.clear() self._seq = 0 @@ -182,7 +186,9 @@ class JobManager: self._pump_thread = threading.Thread(target = self._pump_loop, daemon = True) self._pump_thread.start() - self._emit({"type": EVENT_JOB_ENQUEUED, "ts": time.time(), "job_id": job_id}) + self._emit( + {"type": EVENT_JOB_ENQUEUED, "ts": time.time(), "job_id": job_id} + ) return job_id def cancel(self, job_id: str) -> bool: @@ -193,7 +199,9 @@ class JobManager: if self._proc is None or not self._proc.is_alive(): return True self._job.status = "cancelling" - self._emit({"type": EVENT_JOB_CANCELLING, "ts": time.time(), "job_id": job_id}) + self._emit( + {"type": EVENT_JOB_CANCELLING, "ts": time.time(), "job_id": job_id} + ) try: self._proc.terminate() except (AttributeError, OSError): @@ -310,12 +318,16 @@ class JobManager: if not parquet_dir.exists(): return {"error": f"dataset path missing: {parquet_dir}"} - return self._load_dataset_page(parquet_dir = parquet_dir, limit = limit, offset = offset) + return self._load_dataset_page( + parquet_dir = parquet_dir, limit = limit, offset = offset + ) except Exception as exc: return {"error": f"dataset load failed: {exc}"} @staticmethod - def _load_dataset_page(*, parquet_dir: Path, limit: int, offset: int) -> dict[str, Any]: + def _load_dataset_page( + *, parquet_dir: Path, limit: int, offset: int + ) -> dict[str, Any]: dataset_page = JobManager._load_dataset_page_with_duckdb( parquet_dir = parquet_dir, limit = limit, @@ -478,7 +490,9 @@ class JobManager: self._job.error = self._job.error or "process exited" self._job.finished_at = time.time() event_type = ( - EVENT_JOB_CANCELLED if self._job.status == "cancelled" else EVENT_JOB_ERROR + EVENT_JOB_CANCELLED + if self._job.status == "cancelled" + else EVENT_JOB_ERROR ) self._emit( { diff --git a/studio/backend/core/data_recipe/jobs/parse.py b/studio/backend/core/data_recipe/jobs/parse.py index 3be830d0e4..c9e65b6097 100644 --- a/studio/backend/core/data_recipe/jobs/parse.py +++ b/studio/backend/core/data_recipe/jobs/parse.py @@ -119,7 +119,8 @@ def parse_log_message(msg: str) -> ParsedUpdate | None: page_items = page_items, rate_remaining = int(m.group("remaining")), message = ( - f"Scraping GitHub source: {repo} " f"{resource} page {page} (+{page_items})" + f"Scraping GitHub source: {repo} " + f"{resource} page {page} (+{page_items})" ), ), ) @@ -133,7 +134,9 @@ def parse_log_message(msg: str) -> ParsedUpdate | None: source = "github", status = "rate_limited", retry_after_sec = seconds, - message = ("Waiting for GitHub rate limit. Studio will resume automatically."), + message = ( + "Waiting for GitHub rate limit. Studio will resume automatically." + ), ), ) @@ -161,7 +164,9 @@ def parse_log_message(msg: str) -> ParsedUpdate | None: source = "github", status = "rate_limited", retry_after_sec = seconds, - message = ("Waiting for GitHub rate limit. Studio will resume automatically."), + message = ( + "Waiting for GitHub rate limit. Studio will resume automatically." + ), ), ) @@ -379,13 +384,15 @@ def _apply_source_progress(job: Job, progress: SourceProgress) -> None: count_key = f"{progress.repo}:{progress.resource}" if page_key not in job._source_seen_pages: job._source_seen_pages.add(page_key) - job._source_counts[count_key] = int(job._source_counts.get(count_key, 0)) + int( - page_items or 0 - ) + job._source_counts[count_key] = int( + job._source_counts.get(count_key, 0) + ) + int(page_items or 0) fetched_items = sum(job._source_counts.values()) if fetched_items <= 0: - fetched_items = progress.fetched_items or (previous.fetched_items if previous else None) + fetched_items = progress.fetched_items or ( + previous.fetched_items if previous else None + ) estimated_total = ( progress.estimated_total @@ -405,10 +412,14 @@ def _apply_source_progress(job: Job, progress: SourceProgress) -> None: repo = progress.repo or (previous.repo if previous else None), resource = progress.resource or (previous.resource if previous else None), page = ( - progress.page if progress.page is not None else (previous.page if previous else None) + progress.page + if progress.page is not None + else (previous.page if previous else None) ), page_items = ( - page_items if page_items is not None else (previous.page_items if previous else None) + page_items + if page_items is not None + else (previous.page_items if previous else None) ), fetched_items = fetched_items, estimated_total = estimated_total, @@ -439,7 +450,9 @@ def _compute_overall_progress(job: Job, column_progress: Progress) -> Progress: if len(job._column_done) == 0: done = current_done else: - sum_done = sum(max(0, min(value, total_rows)) for value in job._column_done.values()) + sum_done = sum( + max(0, min(value, total_rows)) for value in job._column_done.values() + ) done = int(sum_done / total_columns) prev_done = int(job.progress.done or 0) diff --git a/studio/backend/core/data_recipe/jobs/worker.py b/studio/backend/core/data_recipe/jobs/worker.py index 4073288cb3..e76353e4a7 100644 --- a/studio/backend/core/data_recipe/jobs/worker.py +++ b/studio/backend/core/data_recipe/jobs/worker.py @@ -60,7 +60,9 @@ def _slugify_run_name(value: str) -> str: return slug[:80].strip("-") -def _build_dataset_name(*, run_name: str | None, job_id: str, artifact_root: Path) -> str: +def _build_dataset_name( + *, run_name: str | None, job_id: str, artifact_root: Path +) -> str: fallback = f"recipe_{job_id}" slug = _slugify_run_name(run_name or "") base_name = f"recipe_{slug}" if slug else fallback @@ -72,7 +74,9 @@ def _build_dataset_name(*, run_name: str | None, job_id: str, artifact_root: Pat return candidate -def run_job_process(*, event_queue, recipe: dict[str, Any], run: dict[str, Any]) -> None: +def run_job_process( + *, event_queue, recipe: dict[str, Any], run: dict[str, Any] +) -> None: """Subprocess entrypoint. Sends events to `event_queue`.""" import os @@ -160,10 +164,14 @@ def run_job_process(*, event_queue, recipe: dict[str, Any], run: dict[str, Any]) } ) else: - results = designer.create(builder, num_records = rows, dataset_name = dataset_name) + results = designer.create( + builder, num_records = rows, dataset_name = dataset_name + ) analysis = to_jsonable(results.load_analysis().model_dump(mode = "json")) if merge_batches: - _merge_batches_to_single_parquet(results.artifact_storage.base_dataset_path) + _merge_batches_to_single_parquet( + results.artifact_storage.base_dataset_path + ) artifact_path = str(results.artifact_storage.base_dataset_path) event_queue.put( { diff --git a/studio/backend/core/data_recipe/local_callable_validators.py b/studio/backend/core/data_recipe/local_callable_validators.py index fbe847f9ce..455d81e0f8 100644 --- a/studio/backend/core/data_recipe/local_callable_validators.py +++ b/studio/backend/core/data_recipe/local_callable_validators.py @@ -133,7 +133,11 @@ def _parse_oxc_spec(*, column: dict[str, Any]) -> OxcLocalCallableValidatorSpec target_columns_raw = column.get("target_columns") target_columns = ( - [value.strip() for value in target_columns_raw if isinstance(value, str) and value.strip()] + [ + value.strip() + for value in target_columns_raw + if isinstance(value, str) and value.strip() + ] if isinstance(target_columns_raw, list) else [] ) @@ -173,7 +177,9 @@ def _parse_oxc_validation_marker(fn_name: str) -> tuple[str, str, str]: return "javascript", "syntax", "auto" code_lang = parts[0] if parts[0] in _OXC_LANG_TO_NODE_LANG else "javascript" mode = parts[1] if parts[1] in _OXC_VALIDATION_MODES else "syntax" - code_shape = parts[2] if len(parts) >= 3 and parts[2] in _OXC_CODE_SHAPES else "auto" + code_shape = ( + parts[2] if len(parts) >= 3 and parts[2] in _OXC_CODE_SHAPES else "auto" + ) return code_lang, mode, code_shape @@ -194,7 +200,10 @@ def _build_oxc_validation_function(lang: str, validation_mode: str, code_shape: code_values = ( ["" for _ in range(row_count)] if not code_column - else ["" if value is None else str(value) for value in df[code_column].tolist()] + else [ + "" if value is None else str(value) + for value in df[code_column].tolist() + ] ) results = _run_oxc_batch( @@ -210,9 +219,7 @@ def _build_oxc_validation_function(lang: str, validation_mode: str, code_shape: ) return pd.DataFrame(results) - _validator.__name__ = ( - f"{OXC_VALIDATION_FN_MARKER}_{node_lang}_{mode.replace('+', '_')}_{normalized_code_shape}" - ) + _validator.__name__ = f"{OXC_VALIDATION_FN_MARKER}_{node_lang}_{mode.replace('+', '_')}_{normalized_code_shape}" return _validator @@ -292,13 +299,21 @@ def _run_oxc_batch( warning_count_raw = item.get("warning_count") out.append( { - "is_valid": bool(is_valid_raw) if isinstance(is_valid_raw, bool) else False, - "error_count": int(error_count_raw) if isinstance(error_count_raw, int) else 0, + "is_valid": bool(is_valid_raw) + if isinstance(is_valid_raw, bool) + else False, + "error_count": int(error_count_raw) + if isinstance(error_count_raw, int) + else 0, "error_message": str(message_raw or ""), - "severity": str(severity_raw) if isinstance(severity_raw, str) else None, + "severity": str(severity_raw) + if isinstance(severity_raw, str) + else None, "code": str(code_raw) if isinstance(code_raw, str) else None, "labels": labels_raw if isinstance(labels_raw, list) else [], - "codeframe": str(codeframe_raw) if isinstance(codeframe_raw, str) else None, + "codeframe": str(codeframe_raw) + if isinstance(codeframe_raw, str) + else None, "warning_count": int(warning_count_raw) if isinstance(warning_count_raw, int) else 0, diff --git a/studio/backend/core/data_recipe/service.py b/studio/backend/core/data_recipe/service.py index 9d8ca5cfcc..a019cf1cda 100644 --- a/studio/backend/core/data_recipe/service.py +++ b/studio/backend/core/data_recipe/service.py @@ -22,7 +22,9 @@ def _encode_bytes_to_base64(value: bytes | bytearray) -> str: return base64.b64encode(bytes(value)).decode("utf-8") -def _load_image_file_to_base64(path_value: str, *, base_path: str | None = None) -> str | None: +def _load_image_file_to_base64( + path_value: str, *, base_path: str | None = None +) -> str | None: try: path = Path(path_value) candidates: list[Path] = [] @@ -117,7 +119,9 @@ def _apply_data_designer_image_context_patch() -> None: original_auto_resolve = ImageContext._auto_resolve_context_value - def _patched_auto_resolve(self: Any, context_value: Any, base_path: str | None) -> Any: + def _patched_auto_resolve( + self: Any, context_value: Any, base_path: str | None + ) -> Any: normalized = _normalize_image_context_value(context_value, base_path = base_path) return original_auto_resolve(self, normalized, base_path) @@ -159,7 +163,9 @@ def _recipe_has_llm_columns(recipe: dict[str, Any]) -> bool: return False -def _validate_recipe_runtime_support(recipe: dict[str, Any], model_providers: list[Any]) -> None: +def _validate_recipe_runtime_support( + recipe: dict[str, Any], model_providers: list[Any] +) -> None: if _recipe_has_llm_columns(recipe) and not model_providers: raise ValueError("Add a Provider connection block before running this recipe.") @@ -249,7 +255,9 @@ def build_config_builder(recipe: dict[str, Any]): if key not in {"model_providers", "mcp_providers"} } recipe_core = _strip_frontend_model_config_metadata(recipe_core) - recipe_core, oxc_local_callable_specs = split_oxc_local_callable_validators(recipe_core) + recipe_core, oxc_local_callable_specs = split_oxc_local_callable_validators( + recipe_core + ) builder = DataDesignerConfigBuilder.from_config({"data_designer": recipe_core}) register_oxc_local_callable_validators( builder = builder, @@ -320,10 +328,14 @@ def preview_recipe( dataset = [to_jsonable(row) for row in raw_rows] artifacts = ( - None if results.processor_artifacts is None else to_jsonable(results.processor_artifacts) + None + if results.processor_artifacts is None + else to_jsonable(results.processor_artifacts) ) analysis = ( - None if results.analysis is None else to_jsonable(results.analysis.model_dump(mode = "json")) + None + if results.analysis is None + else to_jsonable(results.analysis.model_dump(mode = "json")) ) return dataset, artifacts, analysis diff --git a/studio/backend/core/export/export.py b/studio/backend/core/export/export.py index b28b61f088..532bbdb8b2 100644 --- a/studio/backend/core/export/export.py +++ b/studio/backend/core/export/export.py @@ -58,7 +58,9 @@ def _apply_wsl_sudo_patch(): import unsloth_zoo.llama_cpp as llama_cpp_module def _wsl_do_we_need_sudo(system_type = "debian"): - logger.info("WSL detected — skipping sudo check (build deps pre-installed by setup.sh)") + logger.info( + "WSL detected — skipping sudo check (build deps pre-installed by setup.sh)" + ) return False llama_cpp_module.do_we_need_sudo = _wsl_do_we_need_sudo @@ -331,7 +333,9 @@ class ExportBackend: output_path: Optional[str] = None try: if _IS_MLX: - mlx_save_method = "merged_4bit" if format_type == "4-bit (FP4)" else "merged_16bit" + mlx_save_method = ( + "merged_4bit" if format_type == "4-bit (FP4)" else "merged_16bit" + ) else: if format_type == "4-bit (FP4)": save_method = "merged_4bit_forced" @@ -394,7 +398,9 @@ class ExportBackend: private = private, ) else: - hub_save_method = save_method if save_method is not None else "merged_16bit" + hub_save_method = ( + save_method if save_method is not None else "merged_16bit" + ) self.current_model.push_to_hub_merged( repo_id, self.current_tokenizer, @@ -498,7 +504,9 @@ class ExportBackend: else: # Base model name from request or model config base_model = ( - base_model_id or self.current_model.config._name_or_path or "unknown" + base_model_id + or self.current_model.config._name_or_path + or "unknown" ) hf_api = HfApi(token = hf_token) @@ -518,7 +526,9 @@ class ExportBackend: extra = "unsloth", ) card = ModelCard(content) - card.push_to_hub(repo_id, token = hf_token, commit_message = "Unsloth Model Card") + card.push_to_hub( + repo_id, token = hf_token, commit_message = "Unsloth Model Card" + ) if save_directory: hf_api.upload_folder( @@ -581,7 +591,9 @@ class ExportBackend: LLAMA_CPP_DEFAULT_DIR, _resolve_local_convert_script, # noqa: F401 ) - os.environ.setdefault("UNSLOTH_LLAMA_CPP_SCRIPTS_DIR", LLAMA_CPP_DEFAULT_DIR) + os.environ.setdefault( + "UNSLOTH_LLAMA_CPP_SCRIPTS_DIR", LLAMA_CPP_DEFAULT_DIR + ) except ImportError: if not _LLAMA_CPP_SCRIPTS_WARNING_EMITTED: logger.warning( @@ -609,12 +621,16 @@ class ExportBackend: cwd = os.getcwd() pre_existing_ggufs = set(glob.glob(os.path.join(cwd, "*.gguf"))) - pre_existing_subs = {d.name for d in Path(abs_save_dir).iterdir() if d.is_dir()} + pre_existing_subs = { + d.name for d in Path(abs_save_dir).iterdir() if d.is_dir() + } # Avoid clobbering an existing user-owned model/ directory. import uuid - _model_tmp = os.path.join(abs_save_dir, f"_tmp_model_{uuid.uuid4().hex[:8]}") + _model_tmp = os.path.join( + abs_save_dir, f"_tmp_model_{uuid.uuid4().hex[:8]}" + ) model_tmp_to_cleanup = _model_tmp self.current_model.save_pretrained_gguf( _model_tmp, @@ -623,11 +639,15 @@ class ExportBackend: ) # Relocate the .gguf that convert_to_gguf wrote to cwd (repo root). - new_ggufs = set(glob.glob(os.path.join(cwd, "*.gguf"))) - pre_existing_ggufs + new_ggufs = ( + set(glob.glob(os.path.join(cwd, "*.gguf"))) - pre_existing_ggufs + ) for src in sorted(new_ggufs): dest = os.path.join(abs_save_dir, os.path.basename(src)) shutil.move(src, dest) - logger.info(f"Relocated GGUF: {os.path.basename(src)} → {abs_save_dir}/") + logger.info( + f"Relocated GGUF: {os.path.basename(src)} → {abs_save_dir}/" + ) # Flatten GGUF files from subdirs created during this export. for sub in list(Path(abs_save_dir).iterdir()): @@ -647,7 +667,10 @@ class ExportBackend: if self.current_checkpoint: ckpt = Path(self.current_checkpoint) gguf_dir = ckpt.parent / f"{ckpt.name}_gguf" - if gguf_dir.is_dir() and gguf_dir.resolve() != Path(abs_save_dir).resolve(): + if ( + gguf_dir.is_dir() + and gguf_dir.resolve() != Path(abs_save_dir).resolve() + ): for src in gguf_dir.glob("*.gguf"): dest = os.path.join(abs_save_dir, src.name) shutil.move(str(src), dest) @@ -655,7 +678,9 @@ class ExportBackend: # Also relocate Ollama Modelfile if present modelfile = gguf_dir / "Modelfile" if modelfile.is_file(): - shutil.move(str(modelfile), os.path.join(abs_save_dir, "Modelfile")) + shutil.move( + str(modelfile), os.path.join(abs_save_dir, "Modelfile") + ) logger.info(f"Relocated Modelfile → {abs_save_dir}/") shutil.rmtree(str(gguf_dir), ignore_errors = True) logger.info(f"Cleaned up intermediate GGUF dir: {gguf_dir}") @@ -763,8 +788,12 @@ class ExportBackend: repo_type = "model", ) else: - self.current_model.push_to_hub(repo_id, token = hf_token, private = private) - self.current_tokenizer.push_to_hub(repo_id, token = hf_token, private = private) + self.current_model.push_to_hub( + repo_id, token = hf_token, private = private + ) + self.current_tokenizer.push_to_hub( + repo_id, token = hf_token, private = private + ) logger.info(f"Adapter pushed successfully to {repo_id}") return True, "LoRA adapter exported successfully", output_path diff --git a/studio/backend/core/export/orchestrator.py b/studio/backend/core/export/orchestrator.py index 20158d1891..13452699ef 100644 --- a/studio/backend/core/export/orchestrator.py +++ b/studio/backend/core/export/orchestrator.py @@ -276,7 +276,9 @@ class ExportOrchestrator: expected_type, ) - raise RuntimeError(f"Timeout waiting for '{expected_type}' response after {timeout}s") + raise RuntimeError( + f"Timeout waiting for '{expected_type}' response after {timeout}s" + ) def _drain_queue(self) -> list: """Drain all pending responses.""" @@ -326,7 +328,9 @@ class ExportOrchestrator: elif self._proc is not None: self._shutdown_subprocess(timeout = 2) - logger.info("Spawning fresh export subprocess for '%s'", checkpoint_path) + logger.info( + "Spawning fresh export subprocess for '%s'", checkpoint_path + ) self._spawn_subprocess(sub_config) try: @@ -438,7 +442,9 @@ class ExportOrchestrator: }, ) - def _run_export(self, export_type: str, params: dict) -> Tuple[bool, str, Optional[str]]: + def _run_export( + self, export_type: str, params: dict + ) -> Tuple[bool, str, Optional[str]]: """Send an export command and wait for the result. Returns ``(success, message, output_path)``. ``output_path`` is the on-disk @@ -500,7 +506,9 @@ class ExportOrchestrator: finally: self._export_active = False - def scan_checkpoints(self, outputs_dir: str = str(outputs_root())) -> List[Tuple[str, list]]: + def scan_checkpoints( + self, outputs_dir: str = str(outputs_root()) + ) -> List[Tuple[str, list]]: """Scan for checkpoints — runs locally, no ML imports.""" from utils.models.checkpoints import scan_checkpoints return scan_checkpoints(outputs_dir = outputs_dir) diff --git a/studio/backend/core/export/worker.py b/studio/backend/core/export/worker.py index fb2a893014..f9ab8ef98d 100644 --- a/studio/backend/core/export/worker.py +++ b/studio/backend/core/export/worker.py @@ -428,7 +428,9 @@ def run_export_process(*, cmd_queue: Any, resp_queue: Any, config: dict) -> None import transformers - logger.info("Export subprocess loaded transformers %s", transformers.__version__) + logger.info( + "Export subprocess loaded transformers %s", transformers.__version__ + ) except Exception as exc: _send_response( @@ -529,7 +531,9 @@ def run_export_process(*, cmd_queue: Any, resp_queue: Any, config: dict) -> None ) except Exception as exc: - logger.error("Error handling command '%s': %s", cmd_type, exc, exc_info = True) + logger.error( + "Error handling command '%s': %s", cmd_type, exc, exc_info = True + ) _send_response( resp_queue, { diff --git a/studio/backend/core/inference/anthropic_compat.py b/studio/backend/core/inference/anthropic_compat.py index 0307336dde..bcd1032fe9 100644 --- a/studio/backend/core/inference/anthropic_compat.py +++ b/studio/backend/core/inference/anthropic_compat.py @@ -36,7 +36,11 @@ def anthropic_tool_use_id(upstream_id = None) -> str: """Return an Anthropic-style tool_use id (prefix 'toolu_'). Reuses an upstream id only if it already starts with 'toolu_'; otherwise mints a fresh 'toolu_<24 hex>'.""" - if upstream_id and isinstance(upstream_id, str) and upstream_id.startswith("toolu_"): + if ( + upstream_id + and isinstance(upstream_id, str) + and upstream_id.startswith("toolu_") + ): return upstream_id return f"toolu_{uuid.uuid4().hex[:24]}" @@ -149,7 +153,9 @@ def anthropic_messages_to_openai( tc = b.get("content", "") if isinstance(tc, list): tc = " ".join( - p["text"] for p in tc if isinstance(p, dict) and p.get("type") == "text" + p["text"] + for p in tc + if isinstance(p, dict) and p.get("type") == "text" ) tool_results.append( { @@ -432,7 +438,9 @@ class AnthropicStreamEmitter: events.append(self._close_block()) # Reuse the id published in content_block_start; fall back to mapping # the raw id only if no tool_start preceded this end. - tool_use_id = self._open_tool_use_id or anthropic_tool_use_id(event.get("tool_call_id", "")) + tool_use_id = self._open_tool_use_id or anthropic_tool_use_id( + event.get("tool_call_id", "") + ) self._open_tool_call_id = None self._open_tool_use_id = None self._open_tool_args_sent = False diff --git a/studio/backend/core/inference/audio_codecs.py b/studio/backend/core/inference/audio_codecs.py index 93c7da72cb..fa5f4a9d8d 100644 --- a/studio/backend/core/inference/audio_codecs.py +++ b/studio/backend/core/inference/audio_codecs.py @@ -77,7 +77,9 @@ class AudioCodecManager: return from snac import SNAC - self._snac_model = SNAC.from_pretrained("hubertsiuzdak/snac_24khz").to(device).eval() + self._snac_model = ( + SNAC.from_pretrained("hubertsiuzdak/snac_24khz").to(device).eval() + ) logger.info("Loaded SNAC codec (24kHz)") def _load_bicodec( @@ -92,7 +94,9 @@ class AudioCodecManager: # Clone SparkAudio/Spark-TTS for the sparktts package (HF model repos # don't contain it) - spark_code_dir = os.path.join(os.path.dirname(model_repo_path or "."), "Spark-TTS") + spark_code_dir = os.path.join( + os.path.dirname(model_repo_path or "."), "Spark-TTS" + ) sparktts_pkg = os.path.join(spark_code_dir, "sparktts") if not os.path.isdir(sparktts_pkg): logger.info(f"Cloning SparkAudio/Spark-TTS to {spark_code_dir}...") @@ -175,7 +179,9 @@ class AudioCodecManager: # ── Decoders ───────────────────────────────────────────────── - def decode_snac(self, generated_ids: torch.Tensor, device: str) -> Tuple[bytes, int]: + def decode_snac( + self, generated_ids: torch.Tensor, device: str + ) -> Tuple[bytes, int]: """Decode SNAC tokens (Orpheus) into WAV bytes. Finds the START_OF_SPEECH (128257) marker, extracts codes after it, @@ -188,7 +194,9 @@ class AudioCodecManager: cropped = generated_ids[:, token_indices[1][-1] + 1 :] else: # Fall back to the entire output if the marker is missing - logger.warning("No START_OF_SPEECH token (128257) found — using full generated output") + logger.warning( + "No START_OF_SPEECH token (128257) found — using full generated output" + ) cropped = generated_ids row = cropped[0] @@ -214,7 +222,8 @@ class AudioCodecManager: layer_3.append(codes[7 * i + 6] - 24576) snac_codes = [ - torch.tensor(layer).unsqueeze(0).to(device) for layer in [layer_1, layer_2, layer_3] + torch.tensor(layer).unsqueeze(0).to(device) + for layer in [layer_1, layer_2, layer_3] ] with torch.no_grad(): @@ -241,12 +250,16 @@ class AudioCodecManager: f"BiCodec decode: {len(global_matches)} global tokens, {len(semantic_matches)} semantic tokens" ) if len(global_matches) < 10: - logger.info(f"BiCodec generated text (first 500 chars): {generated_text[:500]}") + logger.info( + f"BiCodec generated text (first 500 chars): {generated_text[:500]}" + ) if not semantic_matches: raise ValueError("No bicodec_semantic tokens found in generated output") - semantic_ids = torch.tensor([int(t) for t in semantic_matches]).long().unsqueeze(0) + semantic_ids = ( + torch.tensor([int(t) for t in semantic_matches]).long().unsqueeze(0) + ) # Speaker encoder expects exactly 32 global tokens (token_num=32); # pad with zeros or truncate. diff --git a/studio/backend/core/inference/chat_template_helpers.py b/studio/backend/core/inference/chat_template_helpers.py index b85e9c348a..42e5fdf65b 100644 --- a/studio/backend/core/inference/chat_template_helpers.py +++ b/studio/backend/core/inference/chat_template_helpers.py @@ -55,4 +55,6 @@ def apply_chat_template_for_generation( break if last_exc is not None: raise last_exc - raise RuntimeError("apply_chat_template_for_generation: no attempt produced a result") + raise RuntimeError( + "apply_chat_template_for_generation: no attempt produced a result" + ) diff --git a/studio/backend/core/inference/external_provider.py b/studio/backend/core/inference/external_provider.py index cae001c34d..b446fbff84 100644 --- a/studio/backend/core/inference/external_provider.py +++ b/studio/backend/core/inference/external_provider.py @@ -50,7 +50,9 @@ def _is_openai_family_cloud(base_url: Optional[str]) -> bool: return host == "api.openai.com" or host.endswith(".openai.azure.com") -_ANTHROPIC_4_7_SAMPLING_REMOVED = re.compile(r"^claude-(?:opus|sonnet|haiku)-4-7(?:[-.]|$)") +_ANTHROPIC_4_7_SAMPLING_REMOVED = re.compile( + r"^claude-(?:opus|sonnet|haiku)-4-7(?:[-.]|$)" +) _OPENAI_REASONING_SUMMARY_UNSUPPORTED = re.compile(r"^o3(?:[-.]|$)") _OPENAI_REASONING_STATUSES = {"in_progress", "completed", "incomplete"} @@ -109,7 +111,9 @@ _OPENAI_CITATION_MARKER = re.compile( ) -def _build_citation_lookup(url_citations: list[dict[str, Any]]) -> dict[str, tuple[int, str]]: +def _build_citation_lookup( + url_citations: list[dict[str, Any]], +) -> dict[str, tuple[int, str]]: """Map every known ``source_id`` alias to ``(citation_index, url)``. Accepts singular ``source_id`` and plural ``source_ids``. First-seen @@ -132,7 +136,9 @@ def _build_citation_lookup(url_citations: list[dict[str, Any]]) -> dict[str, tup return by_source -def _replace_openai_citation_markers(text: str, url_citations: list[dict[str, Any]]) -> str: +def _replace_openai_citation_markers( + text: str, url_citations: list[dict[str, Any]] +) -> str: """Rewrite `\\ue200cite\\ue202SOURCE_ID[\\ue202LOCATOR]\\ue201` markers into `[[N]](URL)` per resolvable id. Multi-source markers expand to one link per id; unresolved tokens drop. Idempotent on text without private-use @@ -336,7 +342,9 @@ def _anthropic_supports_compaction(model: str) -> bool: def _anthropic_supports_fast_mode(model: str) -> bool: # Require a family boundary ("" or "-") after the prefix so IDs like # "claude-opus-4-70" / "claude-opus-4-7b" don't match. - return any(model == p or model.startswith(f"{p}-") for p in _ANTHROPIC_FAST_MODE_PREFIXES) + return any( + model == p or model.startswith(f"{p}-") for p in _ANTHROPIC_FAST_MODE_PREFIXES + ) # Cap on ``cited_text`` forwarded in document_citations tool_events; bounds @@ -481,7 +489,8 @@ def _create_shared_http_client() -> httpx.AsyncClient: if "Unknown scheme for proxy URL" not in exc_str and "socksio" not in exc_str: raise logger.warning( - "Ignoring unsupported environment proxy for the shared HTTP client: %s", exc_str + "Ignoring unsupported environment proxy for the shared HTTP client: %s", + exc_str, ) return httpx.AsyncClient(trust_env = False) @@ -613,7 +622,9 @@ def _safe_fetch_image_for_gemini_sync( if rp_info is None: return None _rp, current_host, current_port = rp_info - ok2, reason2, pinned_ip = _validate_and_resolve_host(current_host, current_port) + ok2, reason2, pinned_ip = _validate_and_resolve_host( + current_host, current_port + ) if not ok2: logger.warning( "Gemini image fetch: refusing redirect host=%s reason=%s", @@ -633,9 +644,13 @@ def _safe_fetch_image_for_gemini_sync( with resp: status = getattr(resp, "status", None) or resp.getcode() if status != 200: - logger.info("Gemini image fetch: status=%s host=%s", status, current_host) + logger.info( + "Gemini image fetch: status=%s host=%s", status, current_host + ) return None - _hdr_mime = (resp.headers.get("content-type") or "").split(";")[0].strip().lower() + _hdr_mime = ( + (resp.headers.get("content-type") or "").split(";")[0].strip().lower() + ) # Declared non-image MIME is refused; missing MIME uses the caller's. if _hdr_mime and not _hdr_mime.startswith("image/"): logger.info( @@ -645,7 +660,9 @@ def _safe_fetch_image_for_gemini_sync( ) return None _final_mime_pre = _hdr_mime if _hdr_mime else fallback_mime - if not isinstance(_final_mime_pre, str) or not _final_mime_pre.startswith("image/"): + if not isinstance(_final_mime_pre, str) or not _final_mime_pre.startswith( + "image/" + ): logger.info( "Gemini image fetch: missing content-type and no image fallback host=%s", current_host, @@ -687,7 +704,9 @@ async def _safe_fetch_image_for_gemini( remaining per-request budget so over-budget URLs are rejected up front. """ import asyncio - return await asyncio.to_thread(_safe_fetch_image_for_gemini_sync, url, fallback_mime, max_bytes) + return await asyncio.to_thread( + _safe_fetch_image_for_gemini_sync, url, fallback_mime, max_bytes + ) # Synthetic-tool names stamped onto outbound _toolEvent.arguments so the @@ -764,10 +783,10 @@ class ExternalProviderClient: if self.provider_type == "gemini": _parsed_base = urlparse(self.base_url) if ( - _parsed_base.hostname or "" - ).lower() == "generativelanguage.googleapis.com" and _parsed_base.path.rstrip( - "/" - ) == "/v1beta/openai": + (_parsed_base.hostname or "").lower() + == "generativelanguage.googleapis.com" + and _parsed_base.path.rstrip("/") == "/v1beta/openai" + ): self.base_url = self.base_url[: -len("/openai")] self.api_key = api_key self._timeout = httpx.Timeout(timeout, connect = 10.0) @@ -984,7 +1003,9 @@ class ExternalProviderClient: else: body["thinking"] = {"type": "disabled"} elif self.provider_type == "mistral": - _apply_mistral_reasoning_controls(body, model, enable_thinking, reasoning_effort) + _apply_mistral_reasoning_controls( + body, model, enable_thinking, reasoning_effort + ) elif self.provider_type == "vllm" and enable_thinking is not None: # vLLM gates thinking via chat_template_kwargs.enable_thinking. tpl_kw = body.get("chat_template_kwargs") @@ -1025,7 +1046,9 @@ class ExternalProviderClient: and "web_search" in enabled_tools ): plugins = list(body.get("plugins") or []) - if not any(isinstance(p, dict) and p.get("id") == "web" for p in plugins): + if not any( + isinstance(p, dict) and p.get("id") == "web" for p in plugins + ): plugins.append({"id": "web"}) body["plugins"] = plugins logger.info( @@ -1072,7 +1095,9 @@ class ExternalProviderClient: response.status_code, error_text[:500], ) - yield _error_sse_line(response.status_code, error_text, self.provider_type) + yield _error_sse_line( + response.status_code, error_text, self.provider_type + ) return # Manual __anext__ (not `async for`) so we can close the @@ -1151,7 +1176,11 @@ class ExternalProviderClient: { "type": "tool_end", "tool_call_id": web_search_tool_id, - "result": ("\n---\n".join(blocks) if blocks else "(search complete)"), + "result": ( + "\n---\n".join(blocks) + if blocks + else "(search complete)" + ), } ) @@ -1199,14 +1228,18 @@ class ExternalProviderClient: # in particular returns 200 then surfaces the # failure as an SSE error event. if "error" in parsed: - event_counts["error"] = event_counts.get("error", 0) + 1 + event_counts["error"] = ( + event_counts.get("error", 0) + 1 + ) logger.warning( "%s SSE error event: %s", self.provider_type, parsed.get("error"), ) else: - event_counts["delta"] = event_counts.get("delta", 0) + 1 + event_counts["delta"] = ( + event_counts.get("delta", 0) + 1 + ) # OpenRouter (and most OAI-compat providers) # report the handling model in every chunk's # `model` field. Latch the first non-empty @@ -1232,13 +1265,20 @@ class ExternalProviderClient: ): if not isinstance(envelope, dict): continue - for ann in envelope.get("annotations") or []: + for ann in ( + envelope.get("annotations") + or [] + ): _record_or_url_citation(ann) yield line # Stream ended without [DONE] (some upstreams just close # the connection). Emit tool_end so the card doesn't stay # in "running" forever. - if web_search_active and web_search_tool_started and not web_search_tool_ended: + if ( + web_search_active + and web_search_tool_started + and not web_search_tool_ended + ): yield _build_web_search_tool_end() web_search_tool_ended = True except GeneratorExit: @@ -1315,7 +1355,9 @@ class ExternalProviderClient: # $web_search forbids thinking; sending the toggle would make the # server reject the request with 400. "thinking": {"type": "disabled"}, - "tools": [{"type": "builtin_function", "function": {"name": "$web_search"}}], + "tools": [ + {"type": "builtin_function", "function": {"name": "$web_search"}} + ], } if max_tokens is not None: body["max_tokens"] = max_tokens @@ -1365,7 +1407,9 @@ class ExternalProviderClient: response.status_code, error_text[:500], ) - yield _error_sse_line(response.status_code, error_text, self.provider_type) + yield _error_sse_line( + response.status_code, error_text, self.provider_type + ) return lines_gen = response.aiter_lines().__aiter__() @@ -1429,7 +1473,9 @@ class ExternalProviderClient: # call without the builtin tool. Mirrors the UX of every other # provider when web_search is on but the model didn't need it. search_calls = [ - tc for tc in tool_calls_acc.values() if tc["function"]["name"] == "$web_search" + tc + for tc in tool_calls_acc.values() + if tc["function"]["name"] == "$web_search" ] if not search_calls: logger.info( @@ -1453,7 +1499,9 @@ class ExternalProviderClient: response.status_code, error_text[:500], ) - yield _error_sse_line(response.status_code, error_text, self.provider_type) + yield _error_sse_line( + response.status_code, error_text, self.provider_type + ) return # Manual __anext__ loop instead of `async for` — see the # stream_chat_completion comment for the Python 3.13 + @@ -1554,7 +1602,9 @@ class ExternalProviderClient: response.status_code, error_text[:500], ) - yield _error_sse_line(response.status_code, error_text, self.provider_type) + yield _error_sse_line( + response.status_code, error_text, self.provider_type + ) return lines_gen = response.aiter_lines().__aiter__() @@ -1593,7 +1643,9 @@ class ExternalProviderClient: ): if not isinstance(envelope, dict): continue - for ann in envelope.get("annotations") or []: + for ann in ( + envelope.get("annotations") or [] + ): if isinstance(ann, dict): annotation_shapes.add( str(ann.get("type") or "?") @@ -1663,7 +1715,9 @@ class ExternalProviderClient: system = ( content if isinstance(content, str) - else "\n".join(p["text"] for p in content if p.get("type") == "text") + else "\n".join( + p["text"] for p in content if p.get("type") == "text" + ) ) continue @@ -1721,13 +1775,18 @@ class ExternalProviderClient: # https://platform.claude.com/docs/en/build-with-claude/compaction summary = part.get("content") or "" if isinstance(summary, str) and summary: - anthropic_parts.append({"type": "compaction", "content": summary}) + anthropic_parts.append( + {"type": "compaction", "content": summary} + ) elif part.get("type") == "image_url": url = part.get("image_url", {}).get("url", "") if url.startswith("data:"): # data:image/png;base64, -> split header and data header, _, b64data = url.partition(",") - media_type = header.split(";")[0].replace("data:", "") or "image/jpeg" + media_type = ( + header.split(";")[0].replace("data:", "") + or "image/jpeg" + ) anthropic_parts.append( { "type": "image", @@ -1802,7 +1861,9 @@ class ExternalProviderClient: # the same message. The native Messages API doesn't accept # OpenAI's top-level `tool_calls` field; the call lives inside a # content block `{type:"tool_use", id, name, input}`. - if msg.get("role") == "assistant" and isinstance(msg.get("tool_calls"), list): + if msg.get("role") == "assistant" and isinstance( + msg.get("tool_calls"), list + ): for _tc in msg["tool_calls"]: if not isinstance(_tc, dict): continue @@ -1811,7 +1872,9 @@ class ExternalProviderClient: continue _raw = _fn.get("arguments") or "{}" try: - _input = _json.loads(_raw) if isinstance(_raw, str) else _raw + _input = ( + _json.loads(_raw) if isinstance(_raw, str) else _raw + ) except Exception: _input = {"_raw": _raw} if not isinstance(_input, dict): @@ -1877,7 +1940,9 @@ class ExternalProviderClient: continue _raw = _fn.get("arguments") or "{}" try: - _input = _json.loads(_raw) if isinstance(_raw, str) else _raw + _input = ( + _json.loads(_raw) if isinstance(_raw, str) else _raw + ) except Exception: _input = {"_raw": _raw} if not isinstance(_input, dict): @@ -1959,12 +2024,16 @@ class ExternalProviderClient: last_msg["content"] = head thinking_spec = _anthropic_thinking_spec(model) allowed_efforts = ( - thinking_spec.efforts if thinking_spec else ("none", "low", "medium", "high") + thinking_spec.efforts + if thinking_spec + else ("none", "low", "medium", "high") ) effort = reasoning_effort if reasoning_effort in allowed_efforts else None # Claude 4.6 takes top-tier adaptive effort as "max" only ("xhigh" is # 4.7-only), so map "xhigh" -> "max" for 4.6 outbound requests. - if effort == "xhigh" and model.startswith(("claude-opus-4-6", "claude-sonnet-4-6")): + if effort == "xhigh" and model.startswith( + ("claude-opus-4-6", "claude-sonnet-4-6") + ): effort = "max" if effort is None: if enable_thinking is False: @@ -2015,12 +2084,17 @@ class ExternalProviderClient: and bool(tool_choice["function"].get("name")) ) _anthropic_hosted_builtins_allowed = ( - not _anthropic_tool_choice_disabled and not _anthropic_tool_choice_forced_function + not _anthropic_tool_choice_disabled + and not _anthropic_tool_choice_forced_function ) # Anthropic web_search (date-pinned per model family). # https://platform.claude.com/docs/en/agents-and-tools/tool-use/web-search-tool - if _anthropic_hosted_builtins_allowed and enabled_tools and "web_search" in enabled_tools: + if ( + _anthropic_hosted_builtins_allowed + and enabled_tools + and "web_search" in enabled_tools + ): anthropic_tools = list(body.get("tools") or []) anthropic_tools.append( { @@ -2034,7 +2108,9 @@ class ExternalProviderClient: # Anthropic web_fetch: only URLs already in conversation. Date-pinned. # https://platform.claude.com/docs/en/agents-and-tools/tool-use/web-fetch-tool web_fetch_enabled = bool( - _anthropic_hosted_builtins_allowed and enabled_tools and "web_fetch" in enabled_tools + _anthropic_hosted_builtins_allowed + and enabled_tools + and "web_fetch" in enabled_tools ) if web_fetch_enabled: anthropic_tools = list(body.get("tools") or []) @@ -2139,7 +2215,9 @@ class ExternalProviderClient: # Merge new beta flags onto whatever the registry contributed. existing_beta = request_headers.get("anthropic-beta", "").strip() beta_parts = ( - [p.strip() for p in existing_beta.split(",") if p.strip()] if existing_beta else [] + [p.strip() for p in existing_beta.split(",") if p.strip()] + if existing_beta + else [] ) if code_execution_enabled and _ANTHROPIC_CODE_EXECUTION_BETA not in beta_parts: beta_parts.append(_ANTHROPIC_CODE_EXECUTION_BETA) @@ -2171,7 +2249,10 @@ class ExternalProviderClient: # the id is expired / missing, emit container_invalidated so # the chat adapter clears the stored id and the next turn # falls back to auto-create. - if anthropic_code_exec_container_id and 400 <= response.status_code < 500: + if ( + anthropic_code_exec_container_id + and 400 <= response.status_code < 500 + ): lowered = error_text.lower() if "container" in lowered and ( "expired" in lowered @@ -2184,7 +2265,9 @@ class ExternalProviderClient: f"data: " f"{_json.dumps({'id': completion_id, 'object': 'chat.completion.chunk', 'choices': [{'index': 0, 'delta': {}, 'finish_reason': None}], '_toolEvent': {'type': 'container_invalidated'}})}" ) - yield _error_sse_line(response.status_code, error_text, self.provider_type) + yield _error_sse_line( + response.status_code, error_text, self.provider_type + ) return # NOTE: same manual __anext__ loop as stream_chat_completion — see comment there. @@ -2320,7 +2403,11 @@ class ExternalProviderClient: # Inline a short text preview so the source pill # carries usable context; skip for PDFs (body is # base64-encoded). - if media_type.startswith("text/") and isinstance(data, str) and data: + if ( + media_type.startswith("text/") + and isinstance(data, str) + and data + ): snippet = data[:240].strip() # Frontend parseSourcesFromResult only emits a source pill # when both `Title:` and `URL:` are present, so fall back to @@ -2368,7 +2455,9 @@ class ExternalProviderClient: if "lines" in inner and isinstance(inner.get("lines"), list): return "\n".join(str(line) for line in inner["lines"]) if "is_file_update" in inner: - return "Updated" if inner.get("is_file_update") else "Created" + return ( + "Updated" if inner.get("is_file_update") else "Created" + ) content_field = inner.get("content") if isinstance(content_field, str): return content_field @@ -2414,7 +2503,10 @@ class ExternalProviderClient: content_block = event.get("content_block") or {} block_type = content_block.get("type") block_name = content_block.get("name") - if block_type == "server_tool_use" and block_name == "web_search": + if ( + block_type == "server_tool_use" + and block_name == "web_search" + ): tool_use_id = content_block.get("id", "") or ( f"ws_{len(web_search_calls)}" ) @@ -2435,9 +2527,14 @@ class ExternalProviderClient: content = content_block.get("content") or [] current_result_block = { "tool_use_id": tool_use_id, - "results": list(content) if isinstance(content, list) else [], + "results": list(content) + if isinstance(content, list) + else [], } - elif block_type == "server_tool_use" and block_name == "web_fetch": + elif ( + block_type == "server_tool_use" + and block_name == "web_fetch" + ): tool_use_id = content_block.get("id", "") or ( f"wf_{len(web_fetch_calls)}" ) @@ -2464,7 +2561,9 @@ class ExternalProviderClient: f"ce_{len(code_execution_calls)}" ) kind = ( - "bash" if block_name == "bash_code_execution" else "text_editor" + "bash" + if block_name == "bash_code_execution" + else "text_editor" ) current_code_exec_use = { "id": tool_use_id, @@ -2539,7 +2638,9 @@ class ExternalProviderClient: if isinstance(cit, dict): key = _anthropic_citation_key(cit) idx_for_marker: Optional[int] = None - for idx, existing in enumerate(document_citations, start = 1): + for idx, existing in enumerate( + document_citations, start = 1 + ): if existing.get("_key") == key: idx_for_marker = idx break @@ -2588,7 +2689,9 @@ class ExternalProviderClient: "type": "tool_start", "tool_name": "web_search", "tool_call_id": tool_use_id, - "arguments": ({"query": query} if query else {}), + "arguments": ( + {"query": query} if query else {} + ), } ) current_server_tool_use = None @@ -2629,7 +2732,9 @@ class ExternalProviderClient: kind = current_code_exec_use["kind"] emit_args = {"kind": kind, **parsed_args} if tool_use_id in code_execution_calls: - code_execution_calls[tool_use_id]["arguments"] = emit_args + code_execution_calls[tool_use_id]["arguments"] = ( + emit_args + ) yield _emit_tool_event( { "type": "tool_start", @@ -2663,13 +2768,17 @@ class ExternalProviderClient: file_blocks = inner.get("content") if isinstance(file_blocks, list): for entry in file_blocks: - if isinstance(entry, dict) and entry.get("file_id"): + if isinstance(entry, dict) and entry.get( + "file_id" + ): code_execution_generated_files += 1 result_text = _format_code_execution_result( inner if isinstance(inner, dict) else {} ) if tool_use_id in code_execution_calls: - code_execution_calls[tool_use_id]["result"] = result_text + code_execution_calls[tool_use_id]["result"] = ( + result_text + ) yield _emit_tool_event( { "type": "tool_end", @@ -2748,7 +2857,10 @@ class ExternalProviderClient: c_in = 0 c_out = 0 for it in iterations: - if isinstance(it, dict) and it.get("type") == "compaction": + if ( + isinstance(it, dict) + and it.get("type") == "compaction" + ): c_in += int(it.get("input_tokens") or 0) c_out += int(it.get("output_tokens") or 0) if c_in or c_out: @@ -2760,14 +2872,18 @@ class ExternalProviderClient: # inbound id so reuse doesn't re-write it every turn. delta_obj = event.get("delta") or {} container_obj = delta_obj.get("container") - if isinstance(container_obj, dict) and latched_container_id is None: + if ( + isinstance(container_obj, dict) + and latched_container_id is None + ): probe = container_obj.get("id") if isinstance(probe, str) and probe: latched_container_id = probe if ( latched_container_id and not container_id_emitted - and latched_container_id != anthropic_code_exec_container_id + and latched_container_id + != anthropic_code_exec_container_id ): yield _emit_tool_event( { @@ -2806,7 +2922,9 @@ class ExternalProviderClient: "or remove the previous turn and try " "again._" ) - yield _emit_tool_event({"type": "anthropic_refusal"}) + yield _emit_tool_event( + {"type": "anthropic_refusal"} + ) if mapped is not None: chunk = { "id": completion_id, @@ -2834,8 +2952,13 @@ class ExternalProviderClient: for c in document_citations: entry = {k: v for k, v in c.items() if k != "_key"} cited = entry.get("cited_text") - if isinstance(cited, str) and len(cited) > _CITED_TEXT_MAX_LEN: - entry["cited_text"] = cited[:_CITED_TEXT_MAX_LEN] + "…" + if ( + isinstance(cited, str) + and len(cited) > _CITED_TEXT_MAX_LEN + ): + entry["cited_text"] = ( + cited[:_CITED_TEXT_MAX_LEN] + "…" + ) clean_cits.append(entry) yield _emit_tool_event( { @@ -2854,7 +2977,9 @@ class ExternalProviderClient: if usage_line: yield usage_line yield "data: [DONE]" - await response.aclose() # set PoolByteStream._closed=True FIRST + await ( + response.aclose() + ) # set PoolByteStream._closed=True FIRST break except GeneratorExit: await response.aclose() # set PoolByteStream._closed=True FIRST @@ -2862,21 +2987,31 @@ class ExternalProviderClient: raise finally: # Per-event-type counts + web_search summary for triage. - web_search_requested = bool(enabled_tools and "web_search" in enabled_tools) + web_search_requested = bool( + enabled_tools and "web_search" in enabled_tools + ) web_search_invocations = len(web_search_calls) total_results = sum( len(sc.get("results") or []) for sc in web_search_calls.values() ) - queries = [sc["query"] for sc in web_search_calls.values() if sc.get("query")] + queries = [ + sc["query"] + for sc in web_search_calls.values() + if sc.get("query") + ] # cache_read_input_tokens > 0 proves the cache_control marker # works (turn 1 shows cache_creation instead). code_execution_invocations = len(code_execution_calls) code_execution_results = sum( - 1 for c in code_execution_calls.values() if c.get("result") is not None + 1 + for c in code_execution_calls.values() + if c.get("result") is not None ) web_fetch_requested = web_fetch_enabled web_fetch_invocations = len(web_fetch_calls) - web_fetch_urls = [wf["url"] for wf in web_fetch_calls.values() if wf.get("url")] + web_fetch_urls = [ + wf["url"] for wf in web_fetch_calls.values() if wf.get("url") + ] logger.info( "Anthropic stream complete (model=%s, " "web_search_requested=%s, web_search_invocations=%s, " @@ -3064,7 +3199,10 @@ class ExternalProviderClient: if url.startswith("data:"): header, _, b64data = url.partition(",") media_type = ( - header.split(";")[0].replace("data:", "").strip().lower() + header.split(";")[0] + .replace("data:", "") + .strip() + .lower() or "image/jpeg" ) # Reject non-image data URLs (e.g. data:text/html); @@ -3079,7 +3217,10 @@ class ExternalProviderClient: # data: URLs share the same caps as fetched # URLs so inline payloads don't bypass them. _data_approx_bytes = (len(b64data) * 3) // 4 - if _remote_image_count >= _GEMINI_REMOTE_IMAGE_MAX_COUNT: + if ( + _remote_image_count + >= _GEMINI_REMOTE_IMAGE_MAX_COUNT + ): logger.info( "Gemini inlineData: per-request count cap %d reached, dropping image", _GEMINI_REMOTE_IMAGE_MAX_COUNT, @@ -3132,7 +3273,8 @@ class ExternalProviderClient: _guessed, _ = mimetypes.guess_type(_img_path) _media_type = ( _guessed - if isinstance(_guessed, str) and _guessed.startswith("image/") + if isinstance(_guessed, str) + and _guessed.startswith("image/") else "image/jpeg" ) if _is_youtube: @@ -3165,7 +3307,8 @@ class ExternalProviderClient: # budget is spent; pass the remainder so # over-budget URLs reject on Content-Length. _remaining_bytes = ( - _GEMINI_REMOTE_IMAGE_MAX_TOTAL_BYTES - _remote_image_total_bytes + _GEMINI_REMOTE_IMAGE_MAX_TOTAL_BYTES + - _remote_image_total_bytes ) if _remaining_bytes <= 0: logger.info( @@ -3210,7 +3353,9 @@ class ExternalProviderClient: if isinstance(_msg_extra, dict): _msg_g = _msg_extra.get("google") or {} if isinstance(_msg_g, dict): - _msg_sig = _msg_g.get("thought_signature") or _msg_g.get("thoughtSignature") + _msg_sig = _msg_g.get("thought_signature") or _msg_g.get( + "thoughtSignature" + ) if isinstance(_msg_sig, str) and _msg_sig: for _idx in range(len(parts) - 1, -1, -1): if "text" in parts[_idx]: @@ -3281,7 +3426,9 @@ class ExternalProviderClient: and isinstance(args, dict) and ( args.get("_server_tool") is True - or isinstance((args.get("google") or {}).get("native_part"), dict) + or isinstance( + (args.get("google") or {}).get("native_part"), dict + ) ) ) if _is_synthetic_server_builtin and not ( @@ -3321,9 +3468,9 @@ class ExternalProviderClient: # thoughtSignature only when one subpart exists; for # code+result, prefer executableCode and drop the # signature elsewhere. - _legacy_sig = _native_part.get("thoughtSignature") or _native_part.get( - "thought_signature" - ) + _legacy_sig = _native_part.get( + "thoughtSignature" + ) or _native_part.get("thought_signature") _legacy_subparts = [ _k for _k in ( @@ -3453,7 +3600,9 @@ class ExternalProviderClient: body: dict[str, Any] = {"contents": contents} if system_text_parts: - body["systemInstruction"] = {"parts": [{"text": "\n\n".join(system_text_parts)}]} + body["systemInstruction"] = { + "parts": [{"text": "\n\n".join(system_text_parts)}] + } # Generation config -- temperature / topP / topK / maxOutputTokens map # straight across. The frontend capability matrix restricts the sliders @@ -3490,12 +3639,16 @@ class ExternalProviderClient: and isinstance(tool_choice.get("function"), dict) and bool(tool_choice["function"].get("name")) ) - _hosted_builtins_allowed = not _tool_choice_disabled and not _tool_choice_forced_function + _hosted_builtins_allowed = ( + not _tool_choice_disabled and not _tool_choice_forced_function + ) # Image-tier models reject text-only tools and thinkingConfig regardless # of the pill (model-level constraint); the pill only controls image # output. Decouple the two so Images-off + Code/Search-on doesn't 400. image_tool_requested = bool( - _hosted_builtins_allowed and enabled_tools and "image_generation" in enabled_tools + _hosted_builtins_allowed + and enabled_tools + and "image_generation" in enabled_tools ) # Strict tool / thinking strip uses the model-id check. is_image_model_strict = is_image_picker_model @@ -3527,10 +3680,13 @@ class ExternalProviderClient: "gemini-pro-latest", ) _PRO_THINKING_PREFIXES = ("gemini-2.5-pro",) - is_gemini3_thinking = any(model_lc.startswith(p) for p in _GEMINI3_THINKING_PREFIXES) + is_gemini3_thinking = any( + model_lc.startswith(p) for p in _GEMINI3_THINKING_PREFIXES + ) is_gemini3_pro = any(model_lc.startswith(p) for p in _GEMINI3_PRO_PREFIXES) _is_pro_thinking_only = any( - model_lc == p or model_lc.startswith(p + "-") for p in _PRO_THINKING_PREFIXES + model_lc == p or model_lc.startswith(p + "-") + for p in _PRO_THINKING_PREFIXES ) effort_lc = (reasoning_effort or "").strip().lower() if not is_image_model_strict and is_gemini3_thinking: @@ -3607,7 +3763,8 @@ class ExternalProviderClient: ) google_search_allowed = ( - not is_image_model_strict or _gemini_image_model_allows_google_search(model_lc) + not is_image_model_strict + or _gemini_image_model_allows_google_search(model_lc) ) code_execution_allowed = not is_image_model_strict text_tools_allowed = not is_image_model_strict @@ -3660,7 +3817,9 @@ class ExternalProviderClient: } ) - def _resolve_local_schema_ref(root: Optional[dict[str, Any]], ref: str) -> Optional[Any]: + def _resolve_local_schema_ref( + root: Optional[dict[str, Any]], ref: str + ) -> Optional[Any]: # Walk a `#/foo/bar` JSON pointer against the schema root. Returns # None if the pointer doesn't resolve to a dict, so the caller can # fall back to the unresolved node. @@ -3703,7 +3862,9 @@ class ExternalProviderClient: **_target, **{k: v for k, v in node.items() if k != "$ref"}, } - return _sanitize_gemini_schema(_merged, root, _seen_refs | {_ref}) + return _sanitize_gemini_schema( + _merged, root, _seen_refs | {_ref} + ) cleaned: dict[str, Any] = {} _nullable_from_union = False _flattened_type: Optional[str] = None @@ -3719,7 +3880,9 @@ class ExternalProviderClient: # Preserve multi-type unions as anyOf; flattening to the # first non-null type silently drops the other branches # and changes the tool contract. - _union_any_of = [{"type": _t} for _t in _non_null if isinstance(_t, str)] + _union_any_of = [ + {"type": _t} for _t in _non_null if isinstance(_t, str) + ] for _k, _v in node.items(): if _k == "type" and isinstance(_v, list): # Handled below via _flattened_type. @@ -3748,10 +3911,15 @@ class ExternalProviderClient: _non_null_entries = [ _entry for _entry in _v - if not (isinstance(_entry, dict) and _entry.get("type") == "null") + if not ( + isinstance(_entry, dict) + and _entry.get("type") == "null" + ) ] if len(_non_null_entries) == 1 and _saw_null: - _inner = _sanitize_gemini_schema(_non_null_entries[0], root, _seen_refs) + _inner = _sanitize_gemini_schema( + _non_null_entries[0], root, _seen_refs + ) if isinstance(_inner, dict): for _ik, _iv in _inner.items(): cleaned.setdefault(_ik, _iv) @@ -3770,7 +3938,8 @@ class ExternalProviderClient: cleaned[_k] = _v if _union_any_of is not None and "anyOf" not in cleaned: cleaned["anyOf"] = [ - _sanitize_gemini_schema(_s, root, _seen_refs) for _s in _union_any_of + _sanitize_gemini_schema(_s, root, _seen_refs) + for _s in _union_any_of ] elif _flattened_type is not None: cleaned["type"] = _flattened_type @@ -3812,7 +3981,9 @@ class ExternalProviderClient: _mode = "NONE" elif _tc_lc in ("required", "any"): _mode = "ANY" - elif isinstance(tool_choice, dict) and tool_choice.get("type") == "function": + elif ( + isinstance(tool_choice, dict) and tool_choice.get("type") == "function" + ): _fn_pick = tool_choice.get("function") or {} _name = _fn_pick.get("name") if isinstance(_fn_pick, dict) else None if isinstance(_name, str) and _name: @@ -3862,7 +4033,9 @@ class ExternalProviderClient: } return f"data: {_json.dumps(chunk)}" - def _text_chunk(text: str, extra_content: Optional[dict[str, Any]] = None) -> str: + def _text_chunk( + text: str, extra_content: Optional[dict[str, Any]] = None + ) -> str: delta: dict[str, Any] = {"content": text} if extra_content: delta["extra_content"] = extra_content @@ -3944,7 +4117,9 @@ class ExternalProviderClient: response.status_code, error_text[:500], ) - yield _error_sse_line(response.status_code, error_text, self.provider_type) + yield _error_sse_line( + response.status_code, error_text, self.provider_type + ) return if web_search_active: @@ -3997,7 +4172,9 @@ class ExternalProviderClient: # promptFeedback.blockReason): surface as an error so the # client doesn't see an empty successful response. prompt_feedback = event.get("promptFeedback") - if isinstance(prompt_feedback, dict) and prompt_feedback.get("blockReason"): + if isinstance(prompt_feedback, dict) and prompt_feedback.get( + "blockReason" + ): block_reason = str(prompt_feedback.get("blockReason")) # Close out the synthetic web_search start so the UI # doesn't show a spinner stuck on "searching..." @@ -4048,7 +4225,9 @@ class ExternalProviderClient: u = web.get("uri") or "" if not u or not isinstance(u, str): continue - if any(c["url"] == u for c in web_search_citations): + if any( + c["url"] == u for c in web_search_citations + ): continue web_search_citations.append( { @@ -4060,7 +4239,9 @@ class ExternalProviderClient: content_obj = cand.get("content") or {} parts = ( - content_obj.get("parts") if isinstance(content_obj, dict) else None + content_obj.get("parts") + if isinstance(content_obj, dict) + else None ) if isinstance(parts, list): for part in parts: @@ -4099,7 +4280,10 @@ class ExternalProviderClient: if isinstance(fc, dict): fc_name = fc.get("name") or "" fc_args = fc.get("args") or {} - fc_id = fc.get("id") or f"call_{fc_name}_{time.time_ns()}" + fc_id = ( + fc.get("id") + or f"call_{fc_name}_{time.time_ns()}" + ) if fc_id in emitted_function_call_ids: continue emitted_function_call_ids.add(fc_id) @@ -4119,9 +4303,9 @@ class ExternalProviderClient: # Gemini 3 requires the part-level # thoughtSignature echoed next turn; stow # it on extra_content.google for replay. - thought_sig = part.get("thoughtSignature") or part.get( - "thought_signature" - ) + thought_sig = part.get( + "thoughtSignature" + ) or part.get("thought_signature") if isinstance(thought_sig, str) and thought_sig: tool_call_delta["extra_content"] = { "google": { @@ -4135,7 +4319,9 @@ class ExternalProviderClient: "choices": [ { "index": 0, - "delta": {"tool_calls": [tool_call_delta]}, + "delta": { + "tool_calls": [tool_call_delta] + }, "finish_reason": None, } ], @@ -4190,7 +4376,9 @@ class ExternalProviderClient: "kind": "code_execution", "language": ( ( - exec_code.get("language") + exec_code.get( + "language" + ) or "PYTHON" ).lower() ), @@ -4211,7 +4399,9 @@ class ExternalProviderClient: # outcomes as stderr so the UI surfaces # the error. if outcome and outcome != "OUTCOME_OK": - result_text = f"[{outcome}]\n{output}".rstrip() + result_text = ( + f"[{outcome}]\n{output}".rstrip() + ) else: result_text = output # Pair tool_end with the most recent @@ -4277,7 +4467,8 @@ class ExternalProviderClient: not is_image_model and last_code_exec_tool_id is not None and bool(enabled_tools) - and "code_execution" in (enabled_tools or []) + and "code_execution" + in (enabled_tools or []) ) if attached_to_code_exec: updated_result = ( @@ -4301,22 +4492,28 @@ class ExternalProviderClient: isinstance(_plot_thought_sig, str) and _plot_thought_sig ): - _plot_part_entry["thoughtSignature"] = ( - _plot_thought_sig - ) + _plot_part_entry[ + "thoughtSignature" + ] = _plot_thought_sig yield _emit_tool_event( { "type": "tool_end", - "tool_call_id": (last_code_exec_tool_id), + "tool_call_id": ( + last_code_exec_tool_id + ), "result": updated_result, "google": { "native_part": { - "parts": [_plot_part_entry], + "parts": [ + _plot_part_entry + ], }, }, } ) - last_code_exec_result_text = updated_result + last_code_exec_result_text = ( + updated_result + ) else: img_id = f"img_{time.time_ns()}" yield _emit_tool_event( @@ -4356,9 +4553,9 @@ class ExternalProviderClient: isinstance(_img_thought_sig, str) and _img_thought_sig ): - _img_part_entry["thoughtSignature"] = ( - _img_thought_sig - ) + _img_part_entry[ + "thoughtSignature" + ] = _img_thought_sig _img_native: dict[str, Any] = { "parts": [_img_part_entry], } @@ -4382,7 +4579,11 @@ class ExternalProviderClient: # End-of-stream order: web_search tool_end -> finish_reason -> # usage -> [DONE], matching the Anthropic/OpenAI helpers. - if web_search_active and web_search_tool_started and not web_search_tool_ended: + if ( + web_search_active + and web_search_tool_started + and not web_search_tool_ended + ): blocks: list[str] = [] for cit in web_search_citations: line_out = f"Title: {cit['title']}\nURL: {cit['url']}" @@ -4394,7 +4595,9 @@ class ExternalProviderClient: "type": "tool_end", "tool_call_id": web_search_tool_id, "result": ( - "\n---\n".join(blocks) if blocks else "(search complete)" + "\n---\n".join(blocks) + if blocks + else "(search complete)" ), } ) @@ -4428,19 +4631,25 @@ class ExternalProviderClient: # Gemini bills tool-call prompt slices separately via # `toolUsePromptTokenCount`. Fold into input so # total_tokens doesn't undercount tool turns. - tool_use_prompt_tokens = last_usage.get("toolUsePromptTokenCount") or 0 + tool_use_prompt_tokens = ( + last_usage.get("toolUsePromptTokenCount") or 0 + ) translated_usage = { "input_tokens": prompt_tokens + tool_use_prompt_tokens, "output_tokens": candidate_tokens + thought_tokens, "input_tokens_details": { - "cached_tokens": (last_usage.get("cachedContentTokenCount") or 0), + "cached_tokens": ( + last_usage.get("cachedContentTokenCount") or 0 + ), "tool_use_prompt_tokens": tool_use_prompt_tokens, }, "output_tokens_details": { "reasoning_tokens": thought_tokens, }, } - usage_line = _build_usage_chunk(completion_id, "openai", translated_usage) + usage_line = _build_usage_chunk( + completion_id, "openai", translated_usage + ) if usage_line: yield usage_line @@ -4610,9 +4819,13 @@ class ExternalProviderClient: elif _pt == "image_url": _u = _part.get("image_url", {}).get("url", "") if _u: - _asst_parts.append({"type": "input_image", "image_url": _u}) + _asst_parts.append( + {"type": "input_image", "image_url": _u} + ) if _asst_parts: - input_items.append({"role": "assistant", "content": _asst_parts}) + input_items.append( + {"role": "assistant", "content": _asst_parts} + ) for _tc in _tool_calls: if not isinstance(_tc, dict): @@ -4638,7 +4851,9 @@ class ExternalProviderClient: _is_server_builtin = True else: _g = _args_obj.get("google") - if isinstance(_g, dict) and isinstance(_g.get("native_part"), dict): + if isinstance(_g, dict) and isinstance( + _g.get("native_part"), dict + ): _is_server_builtin = True _call_id_out = _tc.get("id") or f"call_{time.time_ns()}" if _is_server_builtin: @@ -4674,7 +4889,9 @@ class ExternalProviderClient: if url: # Responses takes image_url as a flat string (both # https:// URLs and data: URLs are accepted). - translated_parts.append({"type": "input_image", "image_url": url}) + translated_parts.append( + {"type": "input_image", "image_url": url} + ) elif ( part_type == "reasoning" and role == "assistant" @@ -4832,7 +5049,11 @@ class ExternalProviderClient: # Server-side context compaction (OpenAI cloud only). # https://developers.openai.com/api/docs/guides/compaction - if is_openai_cloud and compaction_threshold is not None and compaction_threshold > 0: + if ( + is_openai_cloud + and compaction_threshold is not None + and compaction_threshold > 0 + ): body["context_management"] = [ { "type": "compaction", @@ -4908,10 +5129,13 @@ class ExternalProviderClient: and bool(tool_choice["function"].get("name")) ) _responses_hosted_builtins_allowed = ( - not _responses_tool_choice_none and not _responses_tool_choice_forced_function + not _responses_tool_choice_none + and not _responses_tool_choice_forced_function ) - if (enabled_tools or responses_user_function_tools) and not _responses_tool_choice_none: + if ( + enabled_tools or responses_user_function_tools + ) and not _responses_tool_choice_none: tools_array: list[dict[str, Any]] = list(responses_user_function_tools) if ( _responses_hosted_builtins_allowed @@ -4950,8 +5174,12 @@ class ExternalProviderClient: dict so the retry doesn't share state with the first attempt. """ attempt_body = dict(body) - if (enabled_tools or responses_user_function_tools) and not _responses_tool_choice_none: - tools_array_attempt: list[dict[str, Any]] = list(responses_user_function_tools) + if ( + enabled_tools or responses_user_function_tools + ) and not _responses_tool_choice_none: + tools_array_attempt: list[dict[str, Any]] = list( + responses_user_function_tools + ) if ( _responses_hosted_builtins_allowed and enabled_tools @@ -4966,8 +5194,13 @@ class ExternalProviderClient: } else: env_attempt = {"type": "container_auto"} - tools_array_attempt.append({"type": "shell", "environment": env_attempt}) - if _responses_hosted_builtins_allowed and image_generation_enabled_openai: + tools_array_attempt.append( + {"type": "shell", "environment": env_attempt} + ) + if ( + _responses_hosted_builtins_allowed + and image_generation_enabled_openai + ): tools_array_attempt.append(_openai_image_generation_tool()) if tools_array_attempt: attempt_body["tools"] = tools_array_attempt @@ -5023,7 +5256,9 @@ class ExternalProviderClient: retried = True attempt_container_id = None continue - yield _error_sse_line(response.status_code, error_text, self.provider_type) + yield _error_sse_line( + response.status_code, error_text, self.provider_type + ) return # NOTE: same manual __anext__ loop as stream_chat_completion — @@ -5119,7 +5354,9 @@ class ExternalProviderClient: # Unterminated: drop the whole tail, else the residual # ``cite`` would leak as plain text. return "" - rendered = _replace_openai_citation_markers(tail, all_url_citations) + rendered = _replace_openai_citation_markers( + tail, all_url_citations + ) # Scrub residual private-use bytes (e.g. a partial opener). for ch in ("", "", ""): rendered = rendered.replace(ch, "") @@ -5172,7 +5409,11 @@ class ExternalProviderClient: chunk_parts.append("(timeout)") if chunk_parts: parts.append("\n".join(chunk_parts)) - return "\n--- next command ---\n".join(parts) if parts else "(no output)" + return ( + "\n--- next command ---\n".join(parts) + if parts + else "(no output)" + ) def _record_url_citation(payload: dict[str, Any]) -> None: """Append a url_citation, deduped by URL: collect every @@ -5236,11 +5477,17 @@ class ExternalProviderClient: return existing summary_text = "" part = payload.get("part") - if isinstance(part, dict) and part.get("type") == "summary_text": + if ( + isinstance(part, dict) + and part.get("type") == "summary_text" + ): text = part.get("text") if isinstance(text, str): summary_text = text - elif payload.get("type") == "response.reasoning_summary_text.done": + elif ( + payload.get("type") + == "response.reasoning_summary_text.done" + ): text = payload.get("text") if isinstance(text, str): summary_text = text @@ -5252,9 +5499,14 @@ class ExternalProviderClient: "type": "summary_text", "text": summary_text, } - if isinstance(summary_index, int) and summary_index >= 0: + if ( + isinstance(summary_index, int) + and summary_index >= 0 + ): while len(summary) <= summary_index: - summary.append({"type": "summary_text", "text": ""}) + summary.append( + {"type": "summary_text", "text": ""} + ) summary[summary_index] = summary_part else: summary.append(summary_part) @@ -5269,7 +5521,9 @@ class ExternalProviderClient: if current_openai_response_id: arguments["openai_response_id"] = current_openai_response_id if last_openai_reasoning_replay_item: - arguments["openai_reasoning_item"] = last_openai_reasoning_replay_item + arguments["openai_reasoning_item"] = ( + last_openai_reasoning_replay_item + ) return arguments def _extract_reasoning_text(payload: Any) -> str: @@ -5328,7 +5582,9 @@ class ExternalProviderClient: # Flush any held-over partial marker; strip # private-use bytes so garbled glyphs don't leak. if pending_marker_tail: - flushed = _flush_pending_marker_tail(pending_marker_tail) + flushed = _flush_pending_marker_tail( + pending_marker_tail + ) pending_marker_tail = "" if flushed: if reasoning_open: @@ -5371,8 +5627,8 @@ class ExternalProviderClient: # Prepend any held-over tail so a marker # straddling two SSE events resolves cleanly. combined = pending_marker_tail + delta_text - head, pending_marker_tail = _split_pending_citation_tail( - combined + head, pending_marker_tail = ( + _split_pending_citation_tail(combined) ) if head: if reasoning_open: @@ -5394,7 +5650,9 @@ class ExternalProviderClient: ) ) if has_unresolved or pending_citation_segments: - pending_citation_segments.append(head_rewritten) + pending_citation_segments.append( + head_rewritten + ) elif head_rewritten: yield _chunk_with_text(head_rewritten) @@ -5413,14 +5671,24 @@ class ExternalProviderClient: elif event_type == "response.output_item.added": item = event.get("item", {}) - if isinstance(item, dict) and item.get("type") == "web_search_call": - item_id = item.get("id", "") or (f"ws_{len(web_search_calls)}") + if ( + isinstance(item, dict) + and item.get("type") == "web_search_call" + ): + item_id = item.get("id", "") or ( + f"ws_{len(web_search_calls)}" + ) web_search_calls.setdefault(item_id, {"query": ""}) # Register shell_call eagerly so out-of-order # output links back. Probe env.container_id to # emit container_ready before response.completed. - if isinstance(item, dict) and item.get("type") == "shell_call": - item_id = item.get("id", "") or (f"sc_{len(shell_calls)}") + if ( + isinstance(item, dict) + and item.get("type") == "shell_call" + ): + item_id = item.get("id", "") or ( + f"sc_{len(shell_calls)}" + ) shell_calls.setdefault( item_id, {"commands": [], "output": None}, @@ -5462,7 +5730,9 @@ class ExternalProviderClient: last_openai_reasoning_replay_item = ( _record_openai_reasoning_replay_item(item) ) - summary_text = _extract_reasoning_text(item.get("summary")) + summary_text = _extract_reasoning_text( + item.get("summary") + ) if summary_text and not reasoning_emitted: if not reasoning_open: summary_text = f"{summary_text}" @@ -5474,10 +5744,14 @@ class ExternalProviderClient: # tool_end here. Citations are aggregated and # the last call's result is overwritten at # response.completed. - item_id = item.get("id", "") or (f"ws_{len(web_search_calls)}") + item_id = item.get("id", "") or ( + f"ws_{len(web_search_calls)}" + ) action = item.get("action") query = ( - action.get("query", "") if isinstance(action, dict) else "" + action.get("query", "") + if isinstance(action, dict) + else "" ) web_search_calls[item_id] = {"query": query} yield _emit_tool_event( @@ -5485,12 +5759,16 @@ class ExternalProviderClient: "type": "tool_start", "tool_name": "web_search", "tool_call_id": item_id, - "arguments": ({"query": query} if query else {}), + "arguments": ( + {"query": query} if query else {} + ), } ) # Per-card text; last call gets overwritten # with citations at response.completed. - per_call_result = f"Searching: {query}" if query else "" + per_call_result = ( + f"Searching: {query}" if query else "" + ) yield _emit_tool_event( { "type": "tool_end", @@ -5503,10 +5781,14 @@ class ExternalProviderClient: # newline-separated string (the card renderer, # shared with Anthropic bash, wants a single # `command`). - item_id = item.get("id", "") or (f"sc_{len(shell_calls)}") + item_id = item.get("id", "") or ( + f"sc_{len(shell_calls)}" + ) action = item.get("action") or {} commands = ( - action.get("commands") if isinstance(action, dict) else None + action.get("commands") + if isinstance(action, dict) + else None ) or [] joined_command = ( "\n".join(str(c) for c in commands) @@ -5522,7 +5804,9 @@ class ExternalProviderClient: }, ) shell_calls[item_id]["commands"] = ( - list(commands) if isinstance(commands, list) else [] + list(commands) + if isinstance(commands, list) + else [] ) yield _emit_tool_event( { @@ -5538,14 +5822,19 @@ class ExternalProviderClient: # Fallback: output may be bundled on the # shell_call done event itself. embedded_output = item.get("output") - if isinstance(embedded_output, list) and embedded_output: + if ( + isinstance(embedded_output, list) + and embedded_output + ): shell_calls[item_id]["output"] = embedded_output shell_calls[item_id]["tool_end_emitted"] = True yield _emit_tool_event( { "type": "tool_end", "tool_call_id": item_id, - "result": _format_shell_output(embedded_output), + "result": _format_shell_output( + embedded_output + ), } ) elif item.get("type") == "shell_call_output": @@ -5553,11 +5842,15 @@ class ExternalProviderClient: # `id`, used as the tool_call_id on # tool_start. Match on call_id when present so # the matching card transitions to complete. - call_id = item.get("call_id") or item.get("id") or "" + call_id = ( + item.get("call_id") or item.get("id") or "" + ) output = item.get("output") or [] # Skip if bundled-output path already # finalised this card. - if shell_calls.get(call_id, {}).get("tool_end_emitted"): + if shell_calls.get(call_id, {}).get( + "tool_end_emitted" + ): continue if call_id in shell_calls: shell_calls[call_id]["output"] = output @@ -5577,7 +5870,9 @@ class ExternalProviderClient: raw_item_id = item.get("id") item_id = raw_item_id or f"img_{time.time_ns()}" prompt_in = ( - item.get("revised_prompt") or item.get("prompt") or "" + item.get("revised_prompt") + or item.get("prompt") + or "" ) done_arguments = _image_generation_arguments( prompt_in, @@ -5592,7 +5887,9 @@ class ExternalProviderClient: "arguments": done_arguments, } ) - b64 = item.get("result") or item.get("b64_json") or "" + b64 = ( + item.get("result") or item.get("b64_json") or "" + ) output_format = item.get("output_format") or "png" yield _emit_tool_event( { @@ -5642,7 +5939,9 @@ class ExternalProviderClient: "type": "function", "function": { "name": fn_name, - "arguments": (fn_args), + "arguments": ( + fn_args + ), }, } ], @@ -5655,10 +5954,17 @@ class ExternalProviderClient: ) saw_function_call = True - elif isinstance(event_type, str) and "reasoning" in event_type: - recorded_reasoning = _record_openai_reasoning_replay_item(event) + elif ( + isinstance(event_type, str) + and "reasoning" in event_type + ): + recorded_reasoning = ( + _record_openai_reasoning_replay_item(event) + ) if recorded_reasoning: - last_openai_reasoning_replay_item = recorded_reasoning + last_openai_reasoning_replay_item = ( + recorded_reasoning + ) reasoning_delta = _extract_reasoning_text(event) if reasoning_delta: if not reasoning_open: @@ -5668,14 +5974,18 @@ class ExternalProviderClient: reasoning_emitted = True elif event_type == "response.completed": - completed_usage = (event.get("response") or {}).get("usage") + completed_usage = (event.get("response") or {}).get( + "usage" + ) if isinstance(completed_usage, dict): last_usage = completed_usage # Flush any unterminated citation tail; by now all # annotations are recorded, else private-use bytes # are stripped. if pending_marker_tail: - flushed = _flush_pending_marker_tail(pending_marker_tail) + flushed = _flush_pending_marker_tail( + pending_marker_tail + ) pending_marker_tail = "" if flushed: if reasoning_open: @@ -5713,7 +6023,8 @@ class ExternalProviderClient: if ( latched_container_id and not container_id_emitted - and latched_container_id != openai_code_exec_container_id + and latched_container_id + != openai_code_exec_container_id ): yield _emit_tool_event( { @@ -5728,7 +6039,9 @@ class ExternalProviderClient: last_id = list(web_search_calls.keys())[-1] blocks: list[str] = [] for cit in all_url_citations: - line = f"Title: {cit['title']}\nURL: {cit['url']}" + line = ( + f"Title: {cit['title']}\nURL: {cit['url']}" + ) if cit.get("snippet"): line += f"\nSnippet: {cit['snippet']}" blocks.append(line) @@ -5762,7 +6075,9 @@ class ExternalProviderClient: "index": 0, "delta": {}, "finish_reason": ( - "tool_calls" if saw_function_call else "stop" + "tool_calls" + if saw_function_call + else "stop" ), } ], @@ -5780,13 +6095,17 @@ class ExternalProviderClient: yield usage_line elif event_type == "response.incomplete": - incomplete_usage = (event.get("response") or {}).get("usage") + incomplete_usage = (event.get("response") or {}).get( + "usage" + ) if isinstance(incomplete_usage, dict): last_usage = incomplete_usage # Same flush as response.completed -- truncated # streams can leave a half-marker in the buffer. if pending_marker_tail: - flushed = _flush_pending_marker_tail(pending_marker_tail) + flushed = _flush_pending_marker_tail( + pending_marker_tail + ) pending_marker_tail = "" if flushed: if reasoning_open: @@ -5811,7 +6130,9 @@ class ExternalProviderClient: last_id = list(web_search_calls.keys())[-1] blocks = [] for cit in all_url_citations: - line = f"Title: {cit['title']}\nURL: {cit['url']}" + line = ( + f"Title: {cit['title']}\nURL: {cit['url']}" + ) if cit.get("snippet"): line += f"\nSnippet: {cit['snippet']}" blocks.append(line) @@ -5864,7 +6185,9 @@ class ExternalProviderClient: elif event_type in ("response.failed", "error"): # Surface the failure to the client; the outer # route emits [DONE] as part of its cleanup. - error_payload = event.get("response", {}).get("error", {}) or { + error_payload = event.get("response", {}).get( + "error", {} + ) or { "message": event.get("message", "Unknown error"), "code": event.get("code"), } @@ -5880,11 +6203,15 @@ class ExternalProviderClient: raise finally: # Per-turn tool summary for triage. - web_search_requested = bool(enabled_tools and "web_search" in enabled_tools) + web_search_requested = bool( + enabled_tools and "web_search" in enabled_tools + ) web_search_invocations = len(web_search_calls) total_citations = len(all_url_citations) queries = [ - sc["query"] for sc in web_search_calls.values() if sc.get("query") + sc["query"] + for sc in web_search_calls.values() + if sc.get("query") ] # On /v1/responses cached tokens live at # usage.input_tokens_details.cached_tokens (not @@ -5897,7 +6224,9 @@ class ExternalProviderClient: code_execution_requested = code_execution_enabled_openai code_execution_invocations = len(shell_calls) code_execution_results = sum( - 1 for sc in shell_calls.values() if sc.get("output") is not None + 1 + for sc in shell_calls.values() + if sc.get("output") is not None ) logger.info( "OpenAI Responses stream complete (model=%s, " @@ -6039,7 +6368,9 @@ class ExternalProviderClient: if ( isinstance(methods, list) and methods - and not any(m in methods for m in ("generateContent", "streamGenerateContent")) + and not any( + m in methods for m in ("generateContent", "streamGenerateContent") + ) ): continue base_id = entry.get("baseModelId") @@ -6134,11 +6465,17 @@ class ExternalProviderClient: logger.info( "openai_container_list.response count=%s items=%s", len(result), - [{"id": c.get("id"), "status": c.get("status")} for c in result if isinstance(c, dict)], + [ + {"id": c.get("id"), "status": c.get("status")} + for c in result + if isinstance(c, dict) + ], ) return result - async def create_openai_container(self, name: str, ttl_minutes: int) -> dict[str, Any]: + async def create_openai_container( + self, name: str, ttl_minutes: int + ) -> dict[str, Any]: """ POST /v1/containers with ``expires_after.anchor="last_active_at"``. ``ttl_minutes`` is the idle timeout — every API call touching the @@ -6251,7 +6588,9 @@ def _error_sse_line(status_code: int, message: str, provider_type: str) -> str: def _build_usage_chunk( - completion_id: str, provider: Literal["anthropic", "openai"], last_usage: Optional[dict] + completion_id: str, + provider: Literal["anthropic", "openai"], + last_usage: Optional[dict], ) -> Optional[str]: """Build an OpenAI ``include_usage``-style SSE chunk carrying upstream prompt-cache accounting back to the client. diff --git a/studio/backend/core/inference/inference.py b/studio/backend/core/inference/inference.py index 2b9517692f..a250d3d738 100644 --- a/studio/backend/core/inference/inference.py +++ b/studio/backend/core/inference/inference.py @@ -260,7 +260,9 @@ class InferenceBackend: if config.is_audio: audio_type = config.audio_type adapter_info = " (LoRA adapter)" if config.is_lora else "" - logger.info(f"Loading audio ({audio_type}) model{adapter_info}: {model_name}") + logger.info( + f"Loading audio ({audio_type}) model{adapter_info}: {model_name}" + ) log_gpu_memory(f"Before loading {model_name}") if audio_type == "csm": @@ -294,7 +296,9 @@ class InferenceBackend: from huggingface_hub import snapshot_download local_dir = base_path.split("/")[-1] - repo_path = snapshot_download(base_path, local_dir = local_dir) + repo_path = snapshot_download( + base_path, local_dir = local_dir + ) abs_repo_path = os.path.abspath(repo_path) logger.info( @@ -405,7 +409,9 @@ class InferenceBackend: ) # Reject CPU/disk offload for audio models too - raise_if_offloaded(self.models[model_name]["model"], device_map, "Inference") + raise_if_offloaded( + self.models[model_name]["model"], device_map, "Inference" + ) self.models[model_name]["context_length"] = runtime_context_length( self.models[model_name].get("model"), max_seq_length, @@ -418,7 +424,9 @@ class InferenceBackend: return True model_type = "vision" if config.is_vision else "text" - adapter_info = " (LoRA adapter)" if self.models[model_name]["is_lora"] else "" + adapter_info = ( + " (LoRA adapter)" if self.models[model_name]["is_lora"] else "" + ) logger.info(f"Loading {model_type} model{adapter_info}: {model_name}") log_gpu_memory(f"Before loading {model_name}") @@ -442,10 +450,13 @@ class InferenceBackend: from transformers import ProcessorMixin if not ( - isinstance(processor, ProcessorMixin) or hasattr(processor, "image_processor") + isinstance(processor, ProcessorMixin) + or hasattr(processor, "image_processor") ): # LoRA adapters: use base model. Local merged exports: read base from export_metadata.json. - processor_source = config.base_model if config.is_lora else config.identifier + processor_source = ( + config.base_model if config.is_lora else config.identifier + ) if not config.is_lora and config.is_local: _meta_path = Path(config.path) / "export_metadata.json" try: @@ -466,7 +477,9 @@ class InferenceBackend: token = hf_token if hf_token and hf_token.strip() else None, trust_remote_code = trust_remote_code, ) - logger.info(f"Loaded {type(processor).__name__} from {processor_source}") + logger.info( + f"Loaded {type(processor).__name__} from {processor_source}" + ) self.models[model_name]["model"] = model self.models[model_name]["tokenizer"] = processor @@ -489,7 +502,9 @@ class InferenceBackend: self.models[model_name]["model"] = model self.models[model_name]["tokenizer"] = tokenizer - raise_if_offloaded(self.models[model_name]["model"], device_map, "Inference") + raise_if_offloaded( + self.models[model_name]["model"], device_map, "Inference" + ) self.models[model_name]["context_length"] = runtime_context_length( self.models[model_name].get("model"), max_seq_length, @@ -537,7 +552,11 @@ class InferenceBackend: import sys as _sys from utils.cache_cleanup import clear_unsloth_compiled_cache - _preserve = ["Unsloth*Trainer.py"] if _sys.platform in ("win32", "darwin") else None + _preserve = ( + ["Unsloth*Trainer.py"] + if _sys.platform in ("win32", "darwin") + else None + ) clear_unsloth_compiled_cache(preserve_patterns = _preserve) logger.info(f"Model '{model_name}' successfully unloaded.") @@ -604,9 +623,13 @@ class InferenceBackend: base_model_name = lora_config.base_model # 1. Load the base model if not already in memory - if base_model_name not in self.models or not self.models[base_model_name].get("model"): + if base_model_name not in self.models or not self.models[ + base_model_name + ].get("model"): logger.info(f"Base model '{base_model_name}' not loaded, loading now.") - base_config = ModelConfig.from_ui_selection(base_model_name, None, is_lora = False) + base_config = ModelConfig.from_ui_selection( + base_model_name, None, is_lora = False + ) if not self.load_model( base_config, max_seq_length, @@ -642,7 +665,9 @@ class InferenceBackend: logger.error(traceback.format_exc()) return False, None, None - def load_adapter(self, base_model_name: str, adapter_path: str, adapter_name: str) -> bool: + def load_adapter( + self, base_model_name: str, adapter_path: str, adapter_name: str + ) -> bool: """Load an adapter onto the model only if not already attached.""" model = self.models[base_model_name].get("model") @@ -713,12 +738,16 @@ class InferenceBackend: ) model.base_model.disable_adapter_layers() else: - logger.info(f"Compare mode: model '{base}' is not a PeftModel, already base") + logger.info( + f"Compare mode: model '{base}' is not a PeftModel, already base" + ) elif use_adapter is True: # Re-enable LoRA layers -> adapter output. if isinstance(model, (PeftModel, PeftModelForCausalLM)): - logger.info(f"Compare mode: enabling adapters on '{base}' for LoRA generation") + logger.info( + f"Compare mode: enabling adapters on '{base}' for LoRA generation" + ) model.base_model.enable_adapter_layers() else: logger.warning("use_adapter=true but model is not a PeftModel") @@ -726,11 +755,15 @@ class InferenceBackend: elif isinstance(use_adapter, str): # Enable adapters and set the named one active. if isinstance(model, (PeftModel, PeftModelForCausalLM)): - logger.info(f"Compare mode: enabling adapter '{use_adapter}' on '{base}'") + logger.info( + f"Compare mode: enabling adapter '{use_adapter}' on '{base}'" + ) model.base_model.enable_adapter_layers() self.set_active_adapter(base, use_adapter) else: - logger.warning(f"use_adapter='{use_adapter}' but model is not a PeftModel") + logger.warning( + f"use_adapter='{use_adapter}' but model is not a PeftModel" + ) def generate_with_adapter_control( self, @@ -902,7 +935,8 @@ class InferenceBackend: processor = model_info.get("processor") has_image_processing = processor is not None and ( - isinstance(processor, ProcessorMixin) or hasattr(processor, "image_processor") + isinstance(processor, ProcessorMixin) + or hasattr(processor, "image_processor") ) if has_image_processing: yield from self._generate_vision_response( @@ -954,7 +988,9 @@ class InferenceBackend: # Step 2: format with tokenizer.apply_chat_template(). if system_prompt: - template_messages = [{"role": "system", "content": system_prompt}] + messages + template_messages = [ + {"role": "system", "content": system_prompt} + ] + messages else: template_messages = messages try: @@ -1069,7 +1105,9 @@ class InferenceBackend: else: # Text-only path for a vision model formatted_prompt = self.format_chat_prompt(messages, system_prompt) - inputs = raw_tokenizer(formatted_prompt, return_tensors = "pt").to(model.device) + inputs = raw_tokenizer(formatted_prompt, return_tensors = "pt").to( + model.device + ) # Stream with TextIteratorStreamer + background thread try: @@ -1356,7 +1394,9 @@ class InferenceBackend: timeout = 0.2, ) except Exception as e: - logger.warning(f"HarmonyTextStreamer init failed, falling back: {e}") + logger.warning( + f"HarmonyTextStreamer init failed, falling back: {e}" + ) streamer = TextIteratorStreamer( tokenizer, skip_prompt = True, @@ -1452,7 +1492,9 @@ class InferenceBackend: cancel_event.set() thread.join(timeout = 10) if thread.is_alive(): - logger.warning("Generation thread did not exit after cancel/join timeout") + logger.warning( + "Generation thread did not exit after cancel/join timeout" + ) if err.get("msg"): yield f"Error: {err['msg']}" @@ -1527,12 +1569,21 @@ class InferenceBackend: raise RuntimeError(f"Unknown audio_type: {audio_type}") def _generate_snac( - self, model, tokenizer, text, temperature, top_p, max_new_tokens, repetition_penalty + self, + model, + tokenizer, + text, + temperature, + top_p, + max_new_tokens, + repetition_penalty, ): """Generate audio using SNAC codec (Orpheus).""" device = model.device start_token = torch.tensor([[128259]], device = device) # START_OF_HUMAN - end_tokens = torch.tensor([[128009, 128260]], device = device) # EOT, END_OF_HUMAN + end_tokens = torch.tensor( + [[128009, 128260]], device = device + ) # EOT, END_OF_HUMAN text_ids = tokenizer(text, return_tensors = "pt").input_ids.to(device) input_ids = torch.cat([start_token, text_ids, end_tokens], dim = 1) attention_mask = torch.ones_like(input_ids) @@ -1556,12 +1607,20 @@ class InferenceBackend: inputs = processor( f"[{speaker_id}]{text}", add_special_tokens = True, return_tensors = "pt" ).to(model.device) - audio_values = model.generate(**inputs, max_new_tokens = max_new_tokens, output_audio = True) + audio_values = model.generate( + **inputs, max_new_tokens = max_new_tokens, output_audio = True + ) return self._audio_codec_manager.decode_csm(audio_values) - def _generate_bicodec(self, model, tokenizer, text, temperature, top_k, max_new_tokens): + def _generate_bicodec( + self, model, tokenizer, text, temperature, top_k, max_new_tokens + ): """Generate audio using BiCodec (Spark-TTS).""" - prompt = "<|task_tts|><|start_content|>" + text + "<|end_content|><|start_global_token|>" + prompt = ( + "<|task_tts|><|start_content|>" + + text + + "<|end_content|><|start_global_token|>" + ) inputs = tokenizer([prompt], return_tensors = "pt").to(model.device) generated = model.generate( **inputs, @@ -1632,7 +1691,9 @@ class InferenceBackend: def __init__(self, penalty: float): self.penalty_last_n = 64 if not isinstance(penalty, float) or penalty <= 0: - raise ValueError(f"`penalty` has to be a positive float, but is {penalty}") + raise ValueError( + f"`penalty` has to be a positive float, but is {penalty}" + ) self.penalty = penalty @torch.no_grad() @@ -1657,8 +1718,12 @@ class InferenceBackend: ) return scores - generation_utils.RepetitionPenaltyLogitsProcessor = RepetitionPenaltyLogitsProcessorPatch - logger.info("Patched RepetitionPenaltyLogitsProcessor with 64-token window for OuteTTS") + generation_utils.RepetitionPenaltyLogitsProcessor = ( + RepetitionPenaltyLogitsProcessorPatch + ) + logger.info( + "Patched RepetitionPenaltyLogitsProcessor with 64-token window for OuteTTS" + ) def _apply_chat_template_for_generation( self, @@ -1700,7 +1765,9 @@ class InferenceBackend: logger.error("Tokenizer not loaded for active model") return "" - chat_template_info = self.models[self.active_model_name].get("chat_template_info", {}) + chat_template_info = self.models[self.active_model_name].get( + "chat_template_info", {} + ) tokenizer = self.models[self.active_model_name]["tokenizer"] tokenizer = getattr(tokenizer, "tokenizer", tokenizer) @@ -1717,7 +1784,9 @@ class InferenceBackend: if role in ["system", "user", "assistant"] and content.strip(): if role == last_role: - logger.debug(f"Skipping consecutive {role} message to maintain alternation") + logger.debug( + f"Skipping consecutive {role} message to maintain alternation" + ) continue if role == "user": @@ -1733,7 +1802,9 @@ class InferenceBackend: continue if chat_messages and chat_messages[-1]["role"] == "assistant": - logger.debug("Removing final assistant message to ensure proper alternation") + logger.debug( + "Removing final assistant message to ensure proper alternation" + ) chat_messages.pop() logger.info(f"Sending {len(chat_messages)} messages to tokenizer:") @@ -1748,7 +1819,10 @@ class InferenceBackend: return formatted_prompt except Exception as e: error_msg = str(e).lower() - if "chat_template is not set" in error_msg or "no template argument" in error_msg: + if ( + "chat_template is not set" in error_msg + or "no template argument" in error_msg + ): logger.info( f"Base model detected - no built-in chat template available, using fallback formatting" ) @@ -1759,7 +1833,9 @@ class InferenceBackend: ) if chat_template_info.get("has_template", False): - logger.info("Falling back to manual template formatting based on detected patterns") + logger.info( + "Falling back to manual template formatting based on detected patterns" + ) template_type = chat_template_info.get("format_type", "generic") manual_prompt = self._format_chat_manual( chat_messages, @@ -1772,7 +1848,9 @@ class InferenceBackend: logger.info("Using generic chat formatting for base model") return self._format_generic_template(chat_messages, {}) - def _format_chat_manual(self, messages: list, template_type: str, special_tokens: dict) -> str: + def _format_chat_manual( + self, messages: list, template_type: str, special_tokens: dict + ) -> str: """Manual chat-formatting fallback when the tokenizer template fails. Args: @@ -1802,7 +1880,9 @@ class InferenceBackend: for msg in messages: role = msg["role"] content = msg["content"] - formatted += f"<|start_header_id|>{role}<|end_header_id|>\n\n{content}<|eot_id|>" + formatted += ( + f"<|start_header_id|>{role}<|end_header_id|>\n\n{content}<|eot_id|>" + ) formatted += "<|start_header_id|>assistant<|end_header_id|>\n\n" return formatted @@ -1831,7 +1911,10 @@ class InferenceBackend: formatted += f"[INST] {user_content} [/INST]" - if i + 1 < len(conversation) and conversation[i + 1]["role"] == "assistant": + if ( + i + 1 < len(conversation) + and conversation[i + 1]["role"] == "assistant" + ): formatted += f" {conversation[i + 1]['content']}" i += 2 else: @@ -1965,7 +2048,9 @@ class InferenceBackend: return text.strip() def _load_chat_template_info(self, model_name: str): - if model_name not in self.models or not self.models[model_name].get("tokenizer"): + if model_name not in self.models or not self.models[model_name].get( + "tokenizer" + ): return tokenizer = self.models[model_name]["tokenizer"] @@ -1983,7 +2068,9 @@ class InferenceBackend: # Exact match first model_name_lower = model_name.lower() if model_name_lower in MODEL_TO_TEMPLATE_MAPPER: - chat_template_info["template_name"] = MODEL_TO_TEMPLATE_MAPPER[model_name_lower] + chat_template_info["template_name"] = MODEL_TO_TEMPLATE_MAPPER[ + model_name_lower + ] logger.info( f"Detected template '{chat_template_info['template_name']}' for {model_name} from mapper" ) @@ -1991,13 +2078,17 @@ class InferenceBackend: # Partial match (for variants like model_name-bnb-4bit) for key in MODEL_TO_TEMPLATE_MAPPER: if key in model_name_lower or model_name_lower in key: - chat_template_info["template_name"] = MODEL_TO_TEMPLATE_MAPPER[key] + chat_template_info["template_name"] = MODEL_TO_TEMPLATE_MAPPER[ + key + ] logger.info( f"Detected template '{chat_template_info['template_name']}' for {model_name} (partial match)" ) break except Exception as e: - logger.warning(f"Could not detect template from mapper for {model_name}: {e}") + logger.warning( + f"Could not detect template from mapper for {model_name}: {e}" + ) try: if hasattr(tokenizer, "chat_template") and tokenizer.chat_template: @@ -2006,7 +2097,10 @@ class InferenceBackend: template_str = tokenizer.chat_template.lower() - if "start_header_id" in template_str and "end_header_id" in template_str: + if ( + "start_header_id" in template_str + and "end_header_id" in template_str + ): chat_template_info["format_type"] = "llama3" elif "[inst]" in template_str and "[/inst]" in template_str: chat_template_info["format_type"] = "mistral" @@ -2033,7 +2127,9 @@ class InferenceBackend: chat_template_info["special_tokens"] = special_tokens else: - logger.info(f"No chat template found for {model_name}, will use generic formatting") + logger.info( + f"No chat template found for {model_name}, will use generic formatting" + ) except Exception as e: logger.error(f"Error loading chat template info for {model_name}: {e}") @@ -2045,7 +2141,9 @@ class InferenceBackend: f"Chat template loaded for {model_name}: {chat_template_info['format_type']} format" ) else: - logger.info(f"No built-in chat template for {model_name}, will use generic formatting") + logger.info( + f"No built-in chat template for {model_name}, will use generic formatting" + ) def get_current_model(self) -> Optional[str]: """Currently active model name.""" diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index 16f6376aab..220f75cc25 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -140,7 +140,10 @@ _FINAL_ANSWER_SIGNAL = re.compile( def _is_short_intent_without_action(text: str) -> bool: stripped = text.strip() - return 0 < len(stripped) < _REPROMPT_MAX_CHARS and _INTENT_SIGNAL.search(stripped) is not None + return ( + 0 < len(stripped) < _REPROMPT_MAX_CHARS + and _INTENT_SIGNAL.search(stripped) is not None + ) def _should_suppress_forced_no_tool_output(text: str) -> bool: @@ -285,7 +288,9 @@ def _fetch_swa_entry_from_hf(repo_id: str) -> Optional[object]: return period lt = src.get("layer_types") if isinstance(lt, list) and lt: - return _period_from_layer_types(lt) or ["full" not in str(t).lower() for t in lt] + return _period_from_layer_types(lt) or [ + "full" not in str(t).lower() for t in lt + ] return None @@ -305,11 +310,15 @@ def _swa_entry_from_config_obj(cfg) -> Optional[object]: return period lt = getattr(src, "layer_types", None) if isinstance(lt, list) and lt: - return _period_from_layer_types(lt) or ["full" not in str(t).lower() for t in lt] + return _period_from_layer_types(lt) or [ + "full" not in str(t).lower() for t in lt + ] return None -_SWA_PATTERN_SOURCE_RE = re.compile(r"sliding_window_pattern\s*(?::\s*[\w\[\], ]*)?\s*=\s*(\d+)") +_SWA_PATTERN_SOURCE_RE = re.compile( + r"sliding_window_pattern\s*(?::\s*[\w\[\], ]*)?\s*=\s*(\d+)" +) def _resolve_swa_entry_from_transformers(arch: str) -> Optional[object]: @@ -508,7 +517,9 @@ def detect_reasoning_flags( return flags -def _is_mtp_model_name(model_identifier: Optional[str], gguf_path: Optional[str] = None) -> bool: +def _is_mtp_model_name( + model_identifier: Optional[str], gguf_path: Optional[str] = None +) -> bool: """Name-based MTP detector. Fallback for the metadata signal.""" for cand in (model_identifier, Path(gguf_path).name if gguf_path else None): if cand and "-mtp" in cand.lower(): @@ -1214,7 +1225,9 @@ class LlamaCppBackend: _capability_cache: dict[tuple[str, int], dict[str, object]] = {} @classmethod - def probe_server_capabilities(cls, binary: Optional[str] = None) -> dict[str, object]: + def probe_server_capabilities( + cls, binary: Optional[str] = None + ) -> dict[str, object]: """Parse `llama-server --help` for feature flags. Returns {found, mtp_token, supports_mtp, ngram_mod_flavor, supports_ngram_mod, spec_draft_n_max_flag}. @@ -1287,7 +1300,9 @@ class LlamaCppBackend: # first non-flag token so flag references inside # descriptions are ignored. for tok in re.split(r"[,\s]+", stripped): - if tok.startswith("--") and re.match(r"--[A-Za-z][A-Za-z0-9_-]*$", tok): + if tok.startswith("--") and re.match( + r"--[A-Za-z][A-Za-z0-9_-]*$", tok + ): current_flags.append(tok) elif tok.startswith("-") and len(tok) > 1: # short alias like -fa; keep scanning aliases. @@ -1376,7 +1391,11 @@ class LlamaCppBackend: if m: prefix, _, num_total = m.group(1), m.group(2), m.group(3) sibling_pat = re.compile( - r"^" + re.escape(prefix) + r"-\d{5}-of-" + re.escape(num_total) + r"\.gguf$" + r"^" + + re.escape(prefix) + + r"-\d{5}-of-" + + re.escape(num_total) + + r"\.gguf$" ) for sibling in main.parent.iterdir(): if sibling != main and sibling_pat.match(sibling.name): @@ -1399,7 +1418,10 @@ class LlamaCppBackend: return False for _i in range(torch.cuda.device_count()): try: - _arch = getattr(torch.cuda.get_device_properties(_i), "gcnArchName", "") or "" + _arch = ( + getattr(torch.cuda.get_device_properties(_i), "gcnArchName", "") + or "" + ) except Exception: continue if _arch.split(":")[0].strip().lower() in {"gfx1150", "gfx1151"}: @@ -1450,7 +1472,9 @@ class LlamaCppBackend: names_by_id: dict[int, str] = {} for ordinal in range(count): try: - name = (torch.cuda.get_device_properties(ordinal).name or "").lower() + name = ( + torch.cuda.get_device_properties(ordinal).name or "" + ).lower() except Exception: continue pid = ( @@ -1545,7 +1569,9 @@ class LlamaCppBackend: # `if x.strip()` filters trailing-comma masks ("0,1,"). # Empty mask (CVD="") yields an empty set -> all GPUs # filtered out, per codebase convention. - allowed = set(int(x.strip()) for x in cvd.split(",") if x.strip()) + allowed = set( + int(x.strip()) for x in cvd.split(",") if x.strip() + ) except ValueError: pass gpus: list[tuple[int, int]] = [] @@ -1751,7 +1777,9 @@ class LlamaCppBackend: return out @staticmethod - def _build_windows_path_dirs(binary_dir: str, prefix: str, cuda_path: str) -> list[str]: + def _build_windows_path_dirs( + binary_dir: str, prefix: str, cuda_path: str + ) -> list[str]: """Ordered PATH entries prepended so llama-server.exe resolves cudart / cublas DLLs: binary_dir, pip nvidia wheels, CUDA_PATH/bin, .../bin/x64. Extracted so test_windows_gpu_detection_mock tests the real logic. #5106.""" @@ -1835,7 +1863,9 @@ class LlamaCppBackend: ) def _kv_heads_for_layer(self, layer_idx: int, fallback: int) -> int: - if self._n_kv_heads_by_layer is not None and layer_idx < len(self._n_kv_heads_by_layer): + if self._n_kv_heads_by_layer is not None and layer_idx < len( + self._n_kv_heads_by_layer + ): return self._n_kv_heads_by_layer[layer_idx] return fallback @@ -1911,7 +1941,10 @@ class LlamaCppBackend: # Path 2: Hybrid Mamba/Attention (Qwen3.5-27B, Qwen3.5-35B-A3B) # Only 1 in N layers is attention; the rest are Mamba (no KV cache). - if self._ssm_inner_size is not None and self._full_attention_interval is not None: + if ( + self._ssm_inner_size is not None + and self._full_attention_interval is not None + ): fai = self._full_attention_interval n_attn = -(-n_layers // fai) if fai > 0 else n_layers # ceiling division if key_len is not None and val_len is not None: @@ -1939,7 +1972,9 @@ class LlamaCppBackend: # --swa-full caches full context like non-SWA (per-slot cells = # per_slot_ctx, collapsing to constant n_ctx total); otherwise SWA # caches 2*sliding_window per slot, clamped at per-slot ctx. - swa_cells_per_slot = per_slot_ctx if swa_full else min(n_ctx, 2 * swa, per_slot_ctx) + swa_cells_per_slot = ( + per_slot_ctx if swa_full else min(n_ctx, 2 * swa, per_slot_ctx) + ) key_len_swa = self._kv_key_length_swa or key_len val_len_swa = self._kv_value_length_swa or val_len if self._sliding_window_pattern is not None: @@ -1956,7 +1991,10 @@ class LlamaCppBackend: ) if is_swa: swa_bytes_per_slot += ( - swa_cells_per_slot * layer_n_kv * (key_len_swa + val_len_swa) * bpe + swa_cells_per_slot + * layer_n_kv + * (key_len_swa + val_len_swa) + * bpe ) if ctx_checkpoints > 0 and not swa_full: checkpoint_extra_per_slot += ( @@ -1968,7 +2006,10 @@ class LlamaCppBackend: ) else: global_bytes += n_ctx * layer_n_kv * (key_len + val_len) * bpe - return int(global_bytes + slots * (swa_bytes_per_slot + checkpoint_extra_per_slot)) + return int( + global_bytes + + slots * (swa_bytes_per_slot + checkpoint_extra_per_slot) + ) n_global = max(1, n_layers_kv // 4) n_swa = n_layers_kv - n_global kv_per_token = n_kv * (key_len + val_len) * bpe @@ -1980,7 +2021,9 @@ class LlamaCppBackend: if ctx_checkpoints > 0 and not swa_full else 0.0 ) - return int(global_bytes + slots * (swa_bytes_per_slot + checkpoint_extra_per_slot)) + return int( + global_bytes + slots * (swa_bytes_per_slot + checkpoint_extra_per_slot) + ) # Path 4: Standard GQA with explicit key/value dimensions if key_len is not None and val_len is not None: @@ -2042,7 +2085,9 @@ class LlamaCppBackend: # can override outright (tensor-parallel mode passes a fatter margin), so # only compute a default when none was supplied. if budget_frac is None: - budget_frac = _CTX_FIT_VRAM_FRACTION - (_MTP_VRAM_RESERVE_FRAC if mtp_engaged else 0.0) + budget_frac = _CTX_FIT_VRAM_FRACTION - ( + _MTP_VRAM_RESERVE_FRAC if mtp_engaged else 0.0 + ) budget_bytes = available_mib * 1024 * 1024 * budget_frac model_footprint = model_size_bytes @@ -2101,7 +2146,9 @@ class LlamaCppBackend: files = list_repo_files(hf_repo, token = hf_token) gguf_files = [ - f for f in files if f.endswith(".gguf") and not _is_companion_gguf_path(f) + f + for f in files + if f.endswith(".gguf") and not _is_companion_gguf_path(f) ] if not gguf_files: return None @@ -2317,7 +2364,10 @@ class LlamaCppBackend: if vtype == 8: # STRING slen = struct.unpack(" {self._llama_log_path}") except OSError as e: logger.debug(f"Could not open diffusion runner log file: {e}") @@ -2707,9 +2782,15 @@ class LlamaCppBackend: prefix = m.group(1) total = m.group(3) sibling_pat = re.compile( - r"^" + re.escape(prefix) + r"-\d{5}-of-" + re.escape(total) + r"\.gguf$" + r"^" + + re.escape(prefix) + + r"-\d{5}-of-" + + re.escape(total) + + r"\.gguf$" ) - gguf_extra_shards = [f for f in gguf_files[1:] if sibling_pat.match(f)] + gguf_extra_shards = [ + f for f in gguf_files[1:] if sibling_pat.match(f) + ] except Exception as e: logger.warning(f"Could not list repo files: {e}") @@ -2722,13 +2803,17 @@ class LlamaCppBackend: try: from utils.models.model_config import _iter_hf_cache_snapshots boundary = re.compile( - r"(? %s from local HF cache", @@ -2915,7 +3002,9 @@ class LlamaCppBackend: try: from utils.models.model_config import _iter_hf_cache_snapshots for snap in _iter_hf_cache_snapshots(hf_repo): - rel_files = [p.relative_to(snap).as_posix() for p in snap.rglob("*.gguf")] + rel_files = [ + p.relative_to(snap).as_posix() for p in snap.rglob("*.gguf") + ] target = pick(rel_files) if target is not None: logger.info("Resolved %s %s from local HF cache", label, target) @@ -2990,7 +3079,8 @@ class LlamaCppBackend: mtp_files = sorted( f for f in candidates - if f.lower().endswith(".gguf") and Path(f).name.lower().startswith("mtp-") + if f.lower().endswith(".gguf") + and Path(f).name.lower().startswith("mtp-") ) return mtp_files[0] if mtp_files else None @@ -3037,7 +3127,9 @@ class LlamaCppBackend: logger.debug(f"Could not size mmproj {launch_mmproj_path}: {e}") return 0 - def _resolve_launch_mtp_path(self, *, mtp_draft_path: Optional[str]) -> Optional[str]: + def _resolve_launch_mtp_path( + self, *, mtp_draft_path: Optional[str] + ) -> Optional[str]: """Return mtp_draft_path iff it exists on disk, else None. No family check needed: the drafter is only ever auto-resolved from @@ -3206,7 +3298,9 @@ class LlamaCppBackend: ) free_by_idx = {idx: free for idx, free in usable_gpus} pool_mib = sum(free_by_idx.values()) - kv_budget_b = (pool_mib - len(gpu_indices) * reserve_mib) * 1024 * 1024 - model_size + kv_budget_b = ( + pool_mib - len(gpu_indices) * reserve_mib + ) * 1024 * 1024 - model_size if mtp_engaged: # MTP keeps a draft model + its own KV cache on GPU. kv_budget_b -= 2 * 1024**3 @@ -3220,7 +3314,9 @@ class LlamaCppBackend: # Weights + buffers exceed the pool -> floor; the load then # falls back to layer split. return ctx_floor - kv_at = self._estimate_kv_cache_bytes(ctx, cache_type_kv, n_parallel = n_parallel) + kv_at = self._estimate_kv_cache_bytes( + ctx, cache_type_kv, n_parallel = n_parallel + ) if kv_at <= kv_budget_b: return ctx return max(ctx_floor, int(ctx * kv_budget_b / kv_at)) @@ -3230,13 +3326,17 @@ class LlamaCppBackend: # max_available_ctx is the hardware ceiling for the UI bound, sized from # the native context independent of an explicit small -c (which only # caps effective_ctx). - max_ctx_target = max_target_ctx if (max_target_ctx and max_target_ctx > 0) else target_ctx + max_ctx_target = ( + max_target_ctx if (max_target_ctx and max_target_ctx > 0) else target_ctx + ) max_available_ctx = _fit_ctx(max_ctx_target) effective_ctx = min(_fit_ctx(target_ctx), max_available_ctx) min_free_mib = min(free_by_idx.values()) kv_bytes = ( - self._estimate_kv_cache_bytes(effective_ctx, cache_type_kv, n_parallel = n_parallel) + self._estimate_kv_cache_bytes( + effective_ctx, cache_type_kv, n_parallel = n_parallel + ) if (self._can_estimate_kv() and effective_ctx > 0) else 0 ) @@ -3312,7 +3412,9 @@ class LlamaCppBackend: try: log_dir = _swa_cache_path().parent / "logs" / "llama-server" log_dir.mkdir(parents = True, exist_ok = True) - self._llama_log_path = log_dir / f"llama-{int(time.time())}-port-{self._port}.log" + self._llama_log_path = ( + log_dir / f"llama-{int(time.time())}-port-{self._port}.log" + ) self._llama_log_fh = open( self._llama_log_path, "w", @@ -3446,9 +3548,9 @@ class LlamaCppBackend: # Re-derive after a retried probe (_mmproj_has_audio persists). from utils.models.model_config import is_audio_input_type - self._has_audio_input = bool(is_audio_input_type(self._audio_type)) or bool( - self._mmproj_has_audio - ) + self._has_audio_input = bool( + is_audio_input_type(self._audio_type) + ) or bool(self._mmproj_has_audio) if not self._healthy: return False return True @@ -3477,7 +3579,11 @@ class LlamaCppBackend: hf_token = hf_token, ) # Auto-download mmproj for vision models unless opted out. - if is_vision and not mmproj_path and not extra_args_disable_mmproj(extra_args): + if ( + is_vision + and not mmproj_path + and not extra_args_disable_mmproj(extra_args) + ): mmproj_path = self._download_mmproj( hf_repo = hf_repo, hf_token = hf_token, @@ -3582,7 +3688,8 @@ class LlamaCppBackend: if ( tensor_parallel and cache_type_kv - and cache_type_kv.strip().lower() not in self._TENSOR_PARALLEL_KV_TYPES + and cache_type_kv.strip().lower() + not in self._TENSOR_PARALLEL_KV_TYPES ): logger.info( "Tensor parallelism requires a non-quantized KV cache; " @@ -3600,15 +3707,21 @@ class LlamaCppBackend: strip_split_mode = False, ) if ctx_override is not None and ctx_override > 0: - logger.info(f"User --ctx-size {ctx_override} honored; skipping auto-reduce") + logger.info( + f"User --ctx-size {ctx_override} honored; skipping auto-reduce" + ) if cache_override is not None: - logger.info(f"User --cache-type-k/-v {cache_override} honored for KV estimate") + logger.info( + f"User --cache-type-k/-v {cache_override} honored for KV estimate" + ) if split_mode_override is not None: logger.info( f"User --split-mode {split_mode_override} honored; " "reconciled into tensor_parallel state" ) - effective_ctx = requested_ctx if requested_ctx > 0 else (self._context_length or 0) + effective_ctx = ( + requested_ctx if requested_ctx > 0 else (self._context_length or 0) + ) max_available_ctx = self._context_length or effective_ctx gpus: list[tuple[int, int]] = [] # Keep fit-budget and launch-flag mmproj resolution in sync. @@ -3631,7 +3744,9 @@ class LlamaCppBackend: gguf_size = self._get_gguf_size_bytes(model_path) # Include GPU-loaded mmproj in the fit budget (#5825). mmproj_size = ( - self._mmproj_vram_bytes(launch_mmproj_path) if effective_is_vision else 0 + self._mmproj_vram_bytes(launch_mmproj_path) + if effective_is_vision + else 0 ) model_size = gguf_size + mmproj_size gpus = self._get_gpu_free_memory() @@ -3762,7 +3877,9 @@ class LlamaCppBackend: # bounds), independent of the currently requested context. native_ctx_for_cap = self._context_length or effective_ctx if native_ctx_for_cap > 0: - ranked_for_cap = sorted(gpus, key = lambda g: g[1], reverse = True) + ranked_for_cap = sorted( + gpus, key = lambda g: g[1], reverse = True + ) best_cap = 0 for n_gpus in range(1, len(ranked_for_cap) + 1): subset = ranked_for_cap[:n_gpus] @@ -3779,7 +3896,9 @@ class LlamaCppBackend: capped, cache_type_kv, n_parallel = n_parallel ) total_mib = (model_size + kv) / (1024 * 1024) - if total_mib <= pool_mib * (_CTX_FIT_VRAM_FRACTION - _mtp_reserve): + if total_mib <= pool_mib * ( + _CTX_FIT_VRAM_FRACTION - _mtp_reserve + ): best_cap = max(best_cap, capped) if best_cap > 0: max_available_ctx = best_cap @@ -3793,8 +3912,11 @@ class LlamaCppBackend: # Honor the requested context verbatim. If it fits, # pin GPUs and skip --fit; else ship -c --fit # on and let llama-server flex -ngl (CPU offload). - requested_total = model_size + self._estimate_kv_cache_bytes( - effective_ctx, cache_type_kv, n_parallel = n_parallel + requested_total = ( + model_size + + self._estimate_kv_cache_bytes( + effective_ctx, cache_type_kv, n_parallel = n_parallel + ) ) gpu_indices, use_fit = self._select_gpus( requested_total, gpus, usable_fraction = _pin_fraction @@ -3841,7 +3963,9 @@ class LlamaCppBackend: ) total_mib = (model_size + kv) / (1024 * 1024) if total_mib <= pool_mib * pin_fraction: - gpu_indices = sorted(idx for idx, _ in subset) + gpu_indices = sorted( + idx for idx, _ in subset + ) use_fit = False break @@ -3858,7 +3982,9 @@ class LlamaCppBackend: if use_fit and not explicit_ctx: # Weights don't fit on any subset; default UI to 4096 # so the slider isn't on an unusable native ctx. - effective_ctx = min(4096, effective_ctx) if effective_ctx > 0 else 4096 + effective_ctx = ( + min(4096, effective_ctx) if effective_ctx > 0 else 4096 + ) if effective_ctx < original_ctx: kv_est = self._estimate_kv_cache_bytes( @@ -3874,7 +4000,9 @@ class LlamaCppBackend: effective_ctx, cache_type_kv, n_parallel = n_parallel ) mmproj_note = ( - f"mmproj: {mmproj_size / (1024**3):.1f} GB, " if mmproj_size else "" + f"mmproj: {mmproj_size / (1024**3):.1f} GB, " + if mmproj_size + else "" ) logger.info( f"GGUF size: {gguf_size / (1024**3):.1f} GB, " @@ -3939,7 +4067,9 @@ class LlamaCppBackend: # so we don't inherit llama-server's internal default, which # has varied (hardware concurrency incl. hyperthreads on some # builds). - cmd.extend(["--threads", str(n_threads if n_threads is not None else -1)]) + cmd.extend( + ["--threads", str(n_threads if n_threads is not None else -1)] + ) # Enable Jinja chat template rendering cmd.extend(["--jinja"]) @@ -4024,7 +4154,9 @@ class LlamaCppBackend: self._supports_reasoning = flags["supports_reasoning"] self._reasoning_style = flags["reasoning_style"] self._reasoning_always_on = flags["reasoning_always_on"] - self._supports_preserve_thinking = flags["supports_preserve_thinking"] + self._supports_preserve_thinking = flags[ + "supports_preserve_thinking" + ] self._supports_tools = flags["supports_tools"] self._chat_template_file = tempfile.NamedTemporaryFile( @@ -4037,7 +4169,9 @@ class LlamaCppBackend: self._chat_template_file.write(chat_template_override) self._chat_template_file.close() cmd.extend(["--chat-template-file", self._chat_template_file.name]) - logger.info(f"Using custom chat template file: {self._chat_template_file.name}") + logger.info( + f"Using custom chat template file: {self._chat_template_file.name}" + ) # Default thinking mode for reasoning models. Qwen3.5/3.6 below # 9B disable thinking by default; 9B+ enable it. Always-on @@ -4077,7 +4211,9 @@ class LlamaCppBackend: if _os.getenv("UNSLOTH_DIRECT_STREAM", "0") == "1": self._api_key = _secrets.token_urlsafe(32) cmd.extend(["--api-key", self._api_key]) - logger.info("llama-server started with --api-key for direct streaming") + logger.info( + "llama-server started with --api-key for direct streaming" + ) else: self._api_key = None @@ -4086,7 +4222,9 @@ class LlamaCppBackend: # validated by the route via validate_extra_args(). if extra_args: cmd.extend(str(a) for a in extra_args) - logger.info(f"Appending user extra args to llama-server: {list(extra_args)}") + logger.info( + f"Appending user extra args to llama-server: {list(extra_args)}" + ) _log_cmd = list(cmd) if "--api-key" in _log_cmd: @@ -4106,7 +4244,9 @@ class LlamaCppBackend: # shared system RAM. setdefault so a user value wins. if self._amd_apu_wants_unified_memory(): env.setdefault("GGML_CUDA_ENABLE_UNIFIED_MEMORY", "1") - logger.info("AMD unified-memory APU: set GGML_CUDA_ENABLE_UNIFIED_MEMORY=1") + logger.info( + "AMD unified-memory APU: set GGML_CUDA_ENABLE_UNIFIED_MEMORY=1" + ) # DC NVIDIA GPUs: FP32 accum (+ P2P / launch queues for multi-GPU). # See _apply_datacenter_env; opt out with UNSLOTH_DISABLE_DC_TUNING=1. @@ -4131,9 +4271,13 @@ class LlamaCppBackend: # searches /rocblas/library/ which doesn't exist # -> silent crash on the first GEMM. ROCBLAS_TENSILE_LIBPATH # repoints that search at the ROCm install. - _hip_path = os.environ.get("HIP_PATH", os.environ.get("ROCM_PATH", "")) + _hip_path = os.environ.get( + "HIP_PATH", os.environ.get("ROCM_PATH", "") + ) if _hip_path: - _rocblas_lib = os.path.join(_hip_path, "bin", "rocblas", "library") + _rocblas_lib = os.path.join( + _hip_path, "bin", "rocblas", "library" + ) if os.path.isdir(_rocblas_lib): env.setdefault("ROCBLAS_TENSILE_LIBPATH", _rocblas_lib) else: @@ -4204,7 +4348,9 @@ class LlamaCppBackend: lib_dirs.append(cuda_lib) existing_ld = env.get("LD_LIBRARY_PATH", "") new_ld = ":".join(lib_dirs) - env["LD_LIBRARY_PATH"] = f"{new_ld}:{existing_ld}" if existing_ld else new_ld + env["LD_LIBRARY_PATH"] = ( + f"{new_ld}:{existing_ld}" if existing_ld else new_ld + ) # Pin to selected GPU(s). On ROCm, narrowing only # CUDA_VISIBLE_DEVICES leaves an AMD child seeing the full @@ -4230,7 +4376,9 @@ class LlamaCppBackend: # clear any inherited ROCR mask so it can't double up. env.pop("ROCR_VISIBLE_DEVICES", None) except Exception as e: - logger.debug("Failed to set ROCm visibility env vars for child: %s", e) + logger.debug( + "Failed to set ROCm visibility env vars for child: %s", e + ) # Captured before any text-only fallback strips it from cmd. launched_with_mmproj = "--mmproj" in cmd @@ -4281,7 +4429,9 @@ class LlamaCppBackend: encoding = "utf-8", buffering = 1, ) - logger.info(f"llama-server stdout/stderr -> {self._llama_log_path}") + logger.info( + f"llama-server stdout/stderr -> {self._llama_log_path}" + ) except OSError as e: # Best-effort; never block the load on logging. logger.debug(f"Could not open llama-server log file: {e}") @@ -4303,9 +4453,14 @@ class LlamaCppBackend: if self._wait_for_health(timeout = 600.0): return True _startup_crashed = ( - self._process.poll() is not None and self._process.returncode != 0 + self._process.poll() is not None + and self._process.returncode != 0 ) - if _spawn_attempt == 0 and _fit_retry_allowed and _startup_crashed: + if ( + _spawn_attempt == 0 + and _fit_retry_allowed + and _startup_crashed + ): logger.warning( "llama-server crashed during startup (exit code %s) " "with the default memory-fit step enabled; Studio " @@ -4347,7 +4502,9 @@ class LlamaCppBackend: ) self._reconcile_effective_ctx_with_server() self._max_context_length = ( - max_available_ctx if max_available_ctx > 0 else self._effective_context_length + max_available_ctx + if max_available_ctx > 0 + else self._effective_context_length ) healthy = _spawn_and_wait(cmd) @@ -4360,7 +4517,11 @@ class LlamaCppBackend: # _requested_spec_mode so a duplicate /load doesn't thrash. The # cancel check stops an /unload-killed attempt respawning. _spec_requested_mtp = any("mtp" in str(t).lower() for t in spec_flags) - if not healthy and _spec_requested_mtp and not self._cancel_event.is_set(): + if ( + not healthy + and _spec_requested_mtp + and not self._cancel_event.is_set() + ): # Blame the binary only when the output shows MTP itself # failing (unknown arch / draft or context build); an # unrelated crash (e.g. OOM) gets a neutral message. @@ -4385,9 +4546,7 @@ class LlamaCppBackend: "binary_outdated" if _arch_unsupported else "runtime_error" ) else: - _retry_reason = ( - "retrying without speculative decoding in case MTP is the cause" - ) + _retry_reason = "retrying without speculative decoding in case MTP is the cause" self._spec_fallback_reason = "runtime_error" _drafter = ( Path(launch_mtp_draft_path).name @@ -4601,7 +4760,9 @@ class LlamaCppBackend: # The sub-3B regression is an embedded-head cost; a separate drafter # (Gemma) is a cheap standalone model that wins below 3B, so exempt it. _mtp_too_small = ( - _mtp_size_b is not None and _mtp_size_b < _MTP_MIN_SIZE_B and not bool(mtp_draft_path) + _mtp_size_b is not None + and _mtp_size_b < _MTP_MIN_SIZE_B + and not bool(mtp_draft_path) ) if user_owns_spec_type: @@ -4823,7 +4984,9 @@ class LlamaCppBackend: # Reconcile a user --split-mode in extras (load_model does the same), so # an extras-driven tensor load isn't seen as a mismatch that needlessly # kills/reloads a healthy server. - if self._tensor_parallel != resolve_tensor_parallel(extra_args, tensor_parallel): + if self._tensor_parallel != resolve_tensor_parallel( + extra_args, tensor_parallel + ): return False # Compare on the canonical requested mode. With --spec-type in @@ -5245,7 +5408,11 @@ class LlamaCppBackend: @staticmethod def _ctx_integrity_flags( - n_parallel: int, use_fit: bool, requested_ctx: int, effective_ctx: int, caps: dict + n_parallel: int, + use_fit: bool, + requested_ctx: int, + effective_ctx: int, + caps: dict, ) -> list[str]: """Flags that keep the per-request window equal to the advertised ctx. @@ -5259,7 +5426,12 @@ class LlamaCppBackend: flags: list[str] = [] if n_parallel > 1 and caps.get("supports_kv_unified"): flags.append("--kv-unified") - if use_fit and requested_ctx > 0 and effective_ctx > 0 and caps.get("supports_fit_ctx"): + if ( + use_fit + and requested_ctx > 0 + and effective_ctx > 0 + and caps.get("supports_fit_ctx") + ): flags.extend(["--fit-ctx", str(effective_ctx)]) return flags @@ -5290,7 +5462,10 @@ class LlamaCppBackend: actual_n_ctx = self._query_server_n_ctx() if not actual_n_ctx or actual_n_ctx <= 0: return - if self._effective_context_length and actual_n_ctx < self._effective_context_length: + if ( + self._effective_context_length + and actual_n_ctx < self._effective_context_length + ): logger.warning( "llama-server allocated a smaller per-request context than " f"requested ({self._effective_context_length} -> {actual_n_ctx}; " @@ -5304,7 +5479,9 @@ class LlamaCppBackend: # ── Message building (OpenAI format) ────────────────────────── @staticmethod - def _parse_tool_calls_from_text(content: str, *, allow_incomplete: bool = True) -> list[dict]: + def _parse_tool_calls_from_text( + content: str, *, allow_incomplete: bool = True + ) -> list[dict]: """Thin wrapper around the shared parser in tool_call_parser so safetensors and llama_cpp pick up the same fixes.""" return _shared_parse_tool_calls_from_text( @@ -5313,7 +5490,9 @@ class LlamaCppBackend: ) @staticmethod - def _build_openai_messages(messages: list[dict], image_b64: Optional[str] = None) -> list[dict]: + def _build_openai_messages( + messages: list[dict], image_b64: Optional[str] = None + ) -> list[dict]: """Build OpenAI-format messages, optionally injecting an image_url part into the last user message for vision models. As-is if no image.""" if not image_b64: @@ -5348,7 +5527,9 @@ class LlamaCppBackend: cancel_event: Optional[threading.Event] = None, stall_timeout_s: float = _DEFAULT_STREAM_STALL_TIMEOUT_S, first_token_deadline: Optional[float] = None, - post_first_chunk_read_timeout_s: Optional[float] = _DEFAULT_STREAM_STALL_TIMEOUT_S, + post_first_chunk_read_timeout_s: Optional[ + float + ] = _DEFAULT_STREAM_STALL_TIMEOUT_S, ) -> Generator[str, None, None]: """Iterate a stream while polling cancel and stall timeouts.""" text_iter = response.iter_text() @@ -5363,11 +5544,16 @@ class LlamaCppBackend: if last_chunk_at is None: remaining_s = first_token_deadline - time.monotonic() if remaining_s <= 0: - raise httpx.ReadTimeout("The model did not produce a first token in time.") + raise httpx.ReadTimeout( + "The model did not produce a first token in time." + ) LlamaCppBackend._set_stream_read_timeout(response, remaining_s) chunk = next(text_iter) if chunk: - if last_chunk_at is None and post_first_chunk_read_timeout_s is not None: + if ( + last_chunk_at is None + and post_first_chunk_read_timeout_s is not None + ): LlamaCppBackend._set_stream_read_timeout( response, post_first_chunk_read_timeout_s, @@ -5382,11 +5568,15 @@ class LlamaCppBackend: if now >= first_token_deadline: raise elif now - last_chunk_at >= stall_timeout_s: - raise httpx.ReadTimeout("The model stopped producing tokens mid-response.") + raise httpx.ReadTimeout( + "The model stopped producing tokens mid-response." + ) continue @staticmethod - def _set_stream_read_timeout(response: "httpx.Response", read_timeout_s: float) -> None: + def _set_stream_read_timeout( + response: "httpx.Response", read_timeout_s: float + ) -> None: """Lower only post-header stream reads; keep prefill timeout long.""" try: timeout_ext = response.request.extensions.get("timeout") @@ -5451,13 +5641,17 @@ class LlamaCppBackend: LlamaCppBackend._shutdown_active_httpx_sockets(client) return except Exception as e: - logger.debug(f"Error closing request in cancel watcher: {e}") + logger.debug( + f"Error closing request in cancel watcher: {e}" + ) _cancel_closed.wait(timeout = 0.1) return watcher = None if cancel_event is not None: - watcher = threading.Thread(target = _cancel_watcher, daemon = True, name = "prefill-cancel") + watcher = threading.Thread( + target = _cancel_watcher, daemon = True, name = "prefill-cancel" + ) watcher.start() try: @@ -5560,7 +5754,9 @@ class LlamaCppBackend: try: # Prefill can use the long first-token timeout; body reads are lowered after headers. stream_timeout = httpx.Timeout(connect = 10, read = 0.5, write = 10, pool = 10) - _auth_headers = {"Authorization": f"Bearer {self._api_key}"} if self._api_key else None + _auth_headers = ( + {"Authorization": f"Bearer {self._api_key}"} if self._api_key else None + ) with httpx.Client( timeout = stream_timeout, limits = httpx.Limits(max_keepalive_connections = 0) ) as client: @@ -5655,7 +5851,9 @@ class LlamaCppBackend: cumulative += token yield cumulative except json.JSONDecodeError: - logger.debug(f"Skipping malformed SSE line: {line[:100]}") + logger.debug( + f"Skipping malformed SSE line: {line[:100]}" + ) if _stream_done: break # exit outer for if _metadata_usage or _metadata_timings or _metadata_finish_reason: @@ -5723,7 +5921,11 @@ class LlamaCppBackend: # Forced first-pass RAG so a doc question doesn't lose to web_search. Emits # the same tool card + citations a real call would. - _auto = None if confirm_tool_calls else build_rag_autoinject(conversation, rag_scope) + _auto = ( + None + if confirm_tool_calls + else build_rag_autoinject(conversation, rag_scope) + ) if _auto: for _ev in _auto["events"]: yield _ev @@ -5759,7 +5961,9 @@ class LlamaCppBackend: def _tool_succeeded(tool_name: str) -> bool: key_prefix = f"{tool_name}:" return any( - record.executed and not record.is_error and record.key.startswith(key_prefix) + record.executed + and not record.is_error + and record.key.startswith(key_prefix) for record in tool_controller.history ) @@ -5823,7 +6027,9 @@ class LlamaCppBackend: try: _auth_headers = ( - {"Authorization": f"Bearer {self._api_key}"} if self._api_key else None + {"Authorization": f"Bearer {self._api_key}"} + if self._api_key + else None ) # ── Speculative buffer state machine ────────────────── @@ -5861,7 +6067,9 @@ class LlamaCppBackend: timeout = stream_timeout, limits = httpx.Limits(max_keepalive_connections = 0), ) as client: - first_token_deadline = time.monotonic() + _DEFAULT_FIRST_TOKEN_TIMEOUT_S + first_token_deadline = ( + time.monotonic() + _DEFAULT_FIRST_TOKEN_TIMEOUT_S + ) with self._stream_with_retry( client, url, @@ -5959,20 +6167,23 @@ class LlamaCppBackend: tool_calls_acc[idx]["id"] = tc_d["id"] func = tc_d.get("function", {}) if func.get("name"): - tool_calls_acc[idx]["function"]["name"] += func[ + tool_calls_acc[idx]["function"][ "name" - ] + ] += func["name"] if func.get("arguments"): - tool_calls_acc[idx]["function"]["arguments"] += ( - func["arguments"] - ) - current_name = tool_calls_acc[idx]["function"].get( - "name", "" - ) + tool_calls_acc[idx]["function"][ + "arguments" + ] += func["arguments"] + current_name = tool_calls_acc[idx][ + "function" + ].get("name", "") fallback_id = f"call_{idx}" - current_id = tool_calls_acc[idx].get("id", fallback_id) + current_id = tool_calls_acc[idx].get( + "id", fallback_id + ) already_started = ( - current_id in provisional_render_html_tool_call_ids + current_id + in provisional_render_html_tool_call_ids ) has_real_id = current_id != fallback_id if ( @@ -5980,7 +6191,9 @@ class LlamaCppBackend: and not _tool_succeeded("render_html") and any( ( - (tool.get("function") or {}).get("name") + ( + tool.get("function") or {} + ).get("name") == "render_html" ) for tool in active_tools @@ -6071,7 +6284,9 @@ class LlamaCppBackend: # route sends it before tool_start. if reasoning_accum: cumulative_display += "" - cumulative_display += reasoning_accum + cumulative_display += ( + reasoning_accum + ) cumulative_display += "" cumulative_display += content_buffer cleaned = _strip_tool_markup_streaming( @@ -6087,7 +6302,9 @@ class LlamaCppBackend: } detect_state = _S_DRAINING elif ( - is_prefix and len(stripped_buf) < _MAX_BUFFER_CHARS + is_prefix + and len(stripped_buf) + < _MAX_BUFFER_CHARS ): pass # keep buffering else: @@ -6097,7 +6314,9 @@ class LlamaCppBackend: # during BUFFERING. if reasoning_accum: cumulative_display += "" - cumulative_display += reasoning_accum + cumulative_display += ( + reasoning_accum + ) cumulative_display += "" cumulative_display += content_buffer cleaned = _strip_tool_markup( @@ -6112,14 +6331,18 @@ class LlamaCppBackend: } except json.JSONDecodeError: - logger.debug(f"Skipping malformed SSE line: {line[:100]}") + logger.debug( + f"Skipping malformed SSE line: {line[:100]}" + ) if _stream_done: break # exit outer for # ── Resolve BUFFERING at stream end ── if detect_state == _S_BUFFERING: stripped_buf = content_buffer.lstrip() - if stripped_buf and any(s in stripped_buf for s in _tool_xml_signals): + if stripped_buf and any( + s in stripped_buf for s in _tool_xml_signals + ): detect_state = _S_DRAINING elif content_accum or reasoning_accum: detect_state = _S_STREAMING @@ -6202,10 +6425,15 @@ class LlamaCppBackend: available_tool_names = [ (tool.get("function") or {}).get("name") for tool in active_tools - if isinstance(tool, dict) and isinstance(tool.get("function"), dict) + if isinstance(tool, dict) + and isinstance(tool.get("function"), dict) ] - available_tool_names = [name for name in available_tool_names if name] - tool_hint = " or ".join(available_tool_names) or "an available tool" + available_tool_names = [ + name for name in available_tool_names if name + ] + tool_hint = ( + " or ".join(available_tool_names) or "an available tool" + ) _forced_tool_call_pending = True conversation.append( { @@ -6219,8 +6447,13 @@ class LlamaCppBackend: } ) # Accumulate tokens and timing from this iteration. - _fu_r = _backfill_usage_from_timings(_iter_usage, _iter_timings) or {} - _accumulated_completion_tokens += _fu_r.get("completion_tokens", 0) + _fu_r = ( + _backfill_usage_from_timings(_iter_usage, _iter_timings) + or {} + ) + _accumulated_completion_tokens += _fu_r.get( + "completion_tokens", 0 + ) _it_r = _iter_timings or {} _accumulated_predicted_ms += _it_r.get("predicted_ms", 0) _accumulated_predicted_n += _it_r.get("predicted_n", 0) @@ -6250,17 +6483,27 @@ class LlamaCppBackend: # Content was already streamed. Yield metadata. yield {"type": "status", "text": ""} - _fu = _backfill_usage_from_timings(_iter_usage, _iter_timings) or {} + _fu = ( + _backfill_usage_from_timings(_iter_usage, _iter_timings) + or {} + ) _fc = _fu.get("completion_tokens", 0) _fp = _fu.get("prompt_tokens", 0) _tc = _fc + _accumulated_completion_tokens - if _iter_usage or _iter_timings or _accumulated_completion_tokens: + if ( + _iter_usage + or _iter_timings + or _accumulated_completion_tokens + ): _mt = dict(_iter_timings) if _iter_timings else {} if _accumulated_predicted_ms or _accumulated_predicted_n: _mt["predicted_ms"] = ( - _mt.get("predicted_ms", 0) + _accumulated_predicted_ms + _mt.get("predicted_ms", 0) + + _accumulated_predicted_ms + ) + _tn = ( + _mt.get("predicted_n", 0) + _accumulated_predicted_n ) - _tn = _mt.get("predicted_n", 0) + _accumulated_predicted_n _mt["predicted_n"] = _tn _tms = _mt["predicted_ms"] if _tms > 0: @@ -6297,9 +6540,16 @@ class LlamaCppBackend: tool_calls = [ tool_calls_acc[i] for i in sorted(tool_calls_acc) - if (tool_calls_acc[i].get("function", {}).get("name", "").strip()) + if ( + tool_calls_acc[i] + .get("function", {}) + .get("name", "") + .strip() + ) ] or None - if not tool_calls and any(s in content_accum for s in _tool_xml_signals): + if not tool_calls and any( + s in content_accum for s in _tool_xml_signals + ): tool_calls = self._parse_tool_calls_from_text( content_accum, allow_incomplete = auto_heal_tool_calls, @@ -6322,20 +6572,32 @@ class LlamaCppBackend: yield {"type": "status", "text": ""} if content_accum: # Strip leaked tool-call XML before yielding. - content_accum = _strip_tool_markup(content_accum, final = True) + content_accum = _strip_tool_markup( + content_accum, final = True + ) if content_accum: yield {"type": "content", "text": content_accum} - _fu = _backfill_usage_from_timings(_iter_usage, _iter_timings) or {} + _fu = ( + _backfill_usage_from_timings(_iter_usage, _iter_timings) + or {} + ) _fc = _fu.get("completion_tokens", 0) _fp = _fu.get("prompt_tokens", 0) _tc = _fc + _accumulated_completion_tokens - if _iter_usage or _iter_timings or _accumulated_completion_tokens: + if ( + _iter_usage + or _iter_timings + or _accumulated_completion_tokens + ): _mt = dict(_iter_timings) if _iter_timings else {} if _accumulated_predicted_ms or _accumulated_predicted_n: _mt["predicted_ms"] = ( - _mt.get("predicted_ms", 0) + _accumulated_predicted_ms + _mt.get("predicted_ms", 0) + + _accumulated_predicted_ms + ) + _tn = ( + _mt.get("predicted_n", 0) + _accumulated_predicted_n ) - _tn = _mt.get("predicted_n", 0) + _accumulated_predicted_n _mt["predicted_n"] = _tn _tms = _mt["predicted_ms"] if _tms > 0: @@ -6397,7 +6659,9 @@ class LlamaCppBackend: break if not assistant_appended: - assistant_msg["tool_calls"] = [decision.as_assistant_tool_call()] + assistant_msg["tool_calls"] = [ + decision.as_assistant_tool_call() + ] conversation.append(assistant_msg) assistant_appended = True else: @@ -6408,7 +6672,9 @@ class LlamaCppBackend: needs_confirm = bool(confirm_tool_calls) approval_id = new_approval_id() if needs_confirm else "" decision_slot = ( - begin_tool_decision(session_id, approval_id) if needs_confirm else None + begin_tool_decision(session_id, approval_id) + if needs_confirm + else None ) start_event = decision.tool_start_event() start_event["approval_id"] = approval_id @@ -6451,7 +6717,9 @@ class LlamaCppBackend: if decision_slot is not None: abort_tool_decision(decision_slot, approval_id) - _effective_timeout = None if tool_call_timeout >= 9999 else tool_call_timeout + _effective_timeout = ( + None if tool_call_timeout >= 9999 else tool_call_timeout + ) # RAG: cap paraphrased KB re-searches that slip past the dup guard. if ( decision.tool_name == "search_knowledge_base" @@ -6478,7 +6746,10 @@ class LlamaCppBackend: # Clear tool status badge before next generation/final pass. yield {"type": "status", "text": ""} - if tool_controller.force_final_answer or not tool_controller.active_tools(): + if ( + tool_controller.force_final_answer + or not tool_controller.active_tools() + ): _append_budget_exhausted_nudge = False break continue @@ -6548,7 +6819,9 @@ class LlamaCppBackend: try: stream_timeout = httpx.Timeout(connect = 10, read = 0.5, write = 10, pool = 10) - _auth_headers = {"Authorization": f"Bearer {self._api_key}"} if self._api_key else None + _auth_headers = ( + {"Authorization": f"Bearer {self._api_key}"} if self._api_key else None + ) with httpx.Client( timeout = stream_timeout, limits = httpx.Limits(max_keepalive_connections = 0) ) as client: @@ -6586,7 +6859,9 @@ class LlamaCppBackend: cumulative += "" yield { "type": "content", - "text": _strip_tool_markup(cumulative, final = True), + "text": _strip_tool_markup( + cumulative, final = True + ), } else: cumulative = reasoning_text @@ -6634,27 +6909,35 @@ class LlamaCppBackend: _last_emitted = cleaned yield {"type": "content", "text": cleaned} except json.JSONDecodeError: - logger.debug(f"Skipping malformed SSE line: {line[:100]}") + logger.debug( + f"Skipping malformed SSE line: {line[:100]}" + ) if _stream_done: break # exit outer for _final_usage = _metadata_usage or {} _final_completion = _final_usage.get("completion_tokens", 0) _final_prompt = _final_usage.get("prompt_tokens", 0) - _total_completion = _final_completion + _accumulated_completion_tokens + _total_completion = ( + _final_completion + _accumulated_completion_tokens + ) if _metadata_usage or _metadata_timings or _metadata_finish_reason: - _merged_timings = dict(_metadata_timings) if _metadata_timings else {} + _merged_timings = ( + dict(_metadata_timings) if _metadata_timings else {} + ) if _accumulated_predicted_ms or _accumulated_predicted_n: _merged_timings["predicted_ms"] = ( - _merged_timings.get("predicted_ms", 0) + _accumulated_predicted_ms + _merged_timings.get("predicted_ms", 0) + + _accumulated_predicted_ms ) _total_predicted_n = ( - _merged_timings.get("predicted_n", 0) + _accumulated_predicted_n + _merged_timings.get("predicted_n", 0) + + _accumulated_predicted_n ) _merged_timings["predicted_n"] = _total_predicted_n _total_predicted_ms = _merged_timings["predicted_ms"] if _total_predicted_ms > 0: - _merged_timings["predicted_per_second"] = _total_predicted_n / ( - _total_predicted_ms / 1000.0 + _merged_timings["predicted_per_second"] = ( + _total_predicted_n / (_total_predicted_ms / 1000.0) ) yield { "type": "metadata", @@ -6703,7 +6986,9 @@ class LlamaCppBackend: continue if not isinstance(block, dict): return True - if block.get("type") == "text" and isinstance(block.get("text"), str): + if block.get("type") == "text" and isinstance( + block.get("text"), str + ): continue if isinstance(block.get("text"), str): continue @@ -6714,7 +6999,9 @@ class LlamaCppBackend: if _has_non_text_content(system): return True for msg in messages or []: - if isinstance(msg, dict) and _has_non_text_content(msg.get("content", "")): + if isinstance(msg, dict) and _has_non_text_content( + msg.get("content", "") + ): return True return False @@ -6725,7 +7012,9 @@ class LlamaCppBackend: parts = [] for block in content: if isinstance(block, dict): - if block.get("type") == "text" and isinstance(block.get("text"), str): + if block.get("type") == "text" and isinstance( + block.get("text"), str + ): parts.append(block["text"]) elif isinstance(block.get("text"), str): parts.append(block["text"]) @@ -6742,7 +7031,9 @@ class LlamaCppBackend: system_text = _block_text(system) try: - _auth_headers = {"Authorization": f"Bearer {self._api_key}"} if self._api_key else None + _auth_headers = ( + {"Authorization": f"Bearer {self._api_key}"} if self._api_key else None + ) with httpx.Client(timeout = 10, headers = _auth_headers) as client: def _tokenize(text: str) -> int: @@ -6757,7 +7048,9 @@ class LlamaCppBackend: tokens = r.json().get("tokens", []) if not isinstance(tokens, list): if strict: - raise RuntimeError("llama-server tokenizer returned invalid tokens") + raise RuntimeError( + "llama-server tokenizer returned invalid tokens" + ) return 0 return len(tokens) @@ -6826,7 +7119,9 @@ class LlamaCppBackend: """Codec name on match, None on non-audio, raises on transport/JSON errors.""" if not self.is_loaded: return None - _auth_headers = {"Authorization": f"Bearer {self._api_key}"} if self._api_key else None + _auth_headers = ( + {"Authorization": f"Bearer {self._api_key}"} if self._api_key else None + ) with httpx.Client(timeout = 10, headers = _auth_headers) as client: def _detok(tid: int) -> str: @@ -6847,7 +7142,9 @@ class LlamaCppBackend: return r.json().get("tokens", []) # Codec-specific tokens (not generic ones that non-audio models may have) - if "")) == 1 and len(_tok("<|audio_eos|>")) == 1: return "csm" @@ -6856,7 +7153,10 @@ class LlamaCppBackend: # Gemma 3n: ; Gemma 4: <|audio|> (not csm's <|AUDIO|>). if len(_tok("")) == 1 or len(_tok("<|audio|>")) == 1: return "audio_vlm" - if len(_tok("<|bicodec_semantic_0|>")) == 1 and len(_tok("<|bicodec_global_0|>")) == 1: + if ( + len(_tok("<|bicodec_semantic_0|>")) == 1 + and len(_tok("<|bicodec_global_0|>")) == 1 + ): return "bicodec" if len(_tok("<|c1_0|>")) == 1 and len(_tok("<|c2_0|>")) == 1: return "dac" @@ -6900,10 +7200,14 @@ class LlamaCppBackend: from huggingface_hub import snapshot_download import os - repo_path = snapshot_download("unsloth/Spark-TTS-0.5B", local_dir = "Spark-TTS-0.5B") + repo_path = snapshot_download( + "unsloth/Spark-TTS-0.5B", local_dir = "Spark-TTS-0.5B" + ) model_repo_path = os.path.abspath(repo_path) - LlamaCppBackend._codec_mgr.load_codec(audio_type, device, model_repo_path = model_repo_path) + LlamaCppBackend._codec_mgr.load_codec( + audio_type, device, model_repo_path = model_repo_path + ) logger.info(f"Loaded audio codec for GGUF TTS: {audio_type}") def generate_audio_response( @@ -6941,11 +7245,17 @@ class LlamaCppBackend: if need_ids: payload["n_probs"] = 1 - _auth_headers = {"Authorization": f"Bearer {self._api_key}"} if self._api_key else None - with httpx.Client(timeout = httpx.Timeout(300, connect = 10), headers = _auth_headers) as client: + _auth_headers = ( + {"Authorization": f"Bearer {self._api_key}"} if self._api_key else None + ) + with httpx.Client( + timeout = httpx.Timeout(300, connect = 10), headers = _auth_headers + ) as client: resp = client.post(f"{self.base_url}/completion", json = payload) if resp.status_code != 200: - raise RuntimeError(f"llama-server returned {resp.status_code}: {resp.text}") + raise RuntimeError( + f"llama-server returned {resp.status_code}: {resp.text}" + ) data = resp.json() token_ids = ( diff --git a/studio/backend/core/inference/llama_server_args.py b/studio/backend/core/inference/llama_server_args.py index 69a86fa3ba..69d8442433 100644 --- a/studio/backend/core/inference/llama_server_args.py +++ b/studio/backend/core/inference/llama_server_args.py @@ -124,7 +124,9 @@ def is_managed_flag(flag: str) -> bool: # from inherited extras so they can't last-wins-override an Apply that # re-sets the same field. _CONTEXT_FLAGS: frozenset[str] = frozenset({"-c", "--ctx-size"}) -_CACHE_FLAGS: frozenset[str] = frozenset({"-ctk", "--cache-type-k", "-ctv", "--cache-type-v"}) +_CACHE_FLAGS: frozenset[str] = frozenset( + {"-ctk", "--cache-type-k", "-ctv", "--cache-type-v"} +) _SPEC_FLAGS: frozenset[str] = frozenset( { "--spec-default", @@ -170,11 +172,17 @@ _TENSOR_SPLIT_FLAGS: frozenset[str] = frozenset({"-ts", "--tensor-split"}) _SPLIT_SHADOWING_FLAGS: frozenset[str] = _SPLIT_MODE_FLAGS | _TENSOR_SPLIT_FLAGS _SHADOWING_FLAGS: frozenset[str] = ( - _CONTEXT_FLAGS | _CACHE_FLAGS | _SPEC_FLAGS | _TEMPLATE_FLAGS | _SPLIT_SHADOWING_FLAGS + _CONTEXT_FLAGS + | _CACHE_FLAGS + | _SPEC_FLAGS + | _TEMPLATE_FLAGS + | _SPLIT_SHADOWING_FLAGS ) # Shadowing flags that take no value -- strip the flag only, not the next token. -_BOOLEAN_SHADOWING_FLAGS: frozenset[str] = frozenset({"--spec-default", "--jinja", "--no-jinja"}) +_BOOLEAN_SHADOWING_FLAGS: frozenset[str] = frozenset( + {"--spec-default", "--jinja", "--no-jinja"} +) def parse_ctx_override(args: Optional[Iterable[str]]) -> Optional[int]: @@ -201,16 +209,22 @@ def parse_ctx_override(args: Optional[Iterable[str]]) -> Optional[int]: i += 1 else: if i + 1 >= n or _flag_name(tokens[i + 1]) is not None: - raise ValueError(f"llama-server flag '{flag}' requires an integer value") + raise ValueError( + f"llama-server flag '{flag}' requires an integer value" + ) raw_value = tokens[i + 1] i += 2 try: value = int(str(raw_value).strip()) except ValueError as exc: - raise ValueError(f"llama-server flag '{flag}' requires an integer value") from exc + raise ValueError( + f"llama-server flag '{flag}' requires an integer value" + ) from exc if value < 0: - raise ValueError(f"llama-server flag '{flag}' requires a non-negative integer value") + raise ValueError( + f"llama-server flag '{flag}' requires a non-negative integer value" + ) override = value return override @@ -226,7 +240,9 @@ def resolve_requested_ctx(args: Optional[Iterable[str]], fallback_n_ctx: int) -> return override if override is not None else fallback_n_ctx -def _last_flag_value(args: Optional[Iterable[str]], flags: frozenset[str]) -> Optional[str]: +def _last_flag_value( + args: Optional[Iterable[str]], flags: frozenset[str] +) -> Optional[str]: """Return the last-wins string value among ``flags`` in extras, or None. Handles both ``--flag=value`` and ``--flag value`` forms and raises if a @@ -295,7 +311,9 @@ def parse_split_mode_override(args: Optional[Iterable[str]]) -> Optional[str]: return _last_flag_value(args, _SPLIT_MODE_FLAGS) -def resolve_tensor_parallel(args: Optional[Iterable[str]], fallback_tensor_parallel: bool) -> bool: +def resolve_tensor_parallel( + args: Optional[Iterable[str]], fallback_tensor_parallel: bool +) -> bool: """Return the tensor-parallel state load_model should treat as requested. A user-supplied ``--split-mode`` in extras last-wins-overrides the diff --git a/studio/backend/core/inference/mcp_client.py b/studio/backend/core/inference/mcp_client.py index 5a36d90c5d..1d75598cd7 100644 --- a/studio/backend/core/inference/mcp_client.py +++ b/studio/backend/core/inference/mcp_client.py @@ -217,7 +217,9 @@ def _client( auth = OAuth(mcp_url = url, token_storage = _oauth_store()) transport_cls = ( - SSETransport if infer_transport_type_from_url(url) == "sse" else StreamableHttpTransport + SSETransport + if infer_transport_type_from_url(url) == "sse" + else StreamableHttpTransport ) return Client(transport_cls(url = url, headers = headers or None, auth = auth)) @@ -252,7 +254,9 @@ _probe_cooloff_until: dict[str, float] = {} # endpoint/auth used to probe it (url, headers, oauth) or whether it's used at # all (is_enabled). A rename does not. The update route's eviction and # get_enabled_mcp_tools' mid-probe guard both key off this so they can't drift. -TOOL_CACHE_INVALIDATING_FIELDS = frozenset({"url", "headers_json", "use_oauth", "is_enabled"}) +TOOL_CACHE_INVALIDATING_FIELDS = frozenset( + {"url", "headers_json", "use_oauth", "is_enabled"} +) def get_cached_tools(server_id: str) -> Optional[list[dict]]: @@ -265,7 +269,11 @@ def cache_tools(server_id: str, tools: list[dict]) -> None: def record_probe_failure(server_id: str, use_oauth: bool = False) -> None: - cooloff = OAUTH_FAILED_PROBE_COOLOFF_SECONDS if use_oauth else FAILED_PROBE_COOLOFF_SECONDS + cooloff = ( + OAUTH_FAILED_PROBE_COOLOFF_SECONDS + if use_oauth + else FAILED_PROBE_COOLOFF_SECONDS + ) _probe_cooloff_until[server_id] = time.monotonic() + cooloff diff --git a/studio/backend/core/inference/mcp_config_import.py b/studio/backend/core/inference/mcp_config_import.py index 534ed84c1d..3c316e0378 100644 --- a/studio/backend/core/inference/mcp_config_import.py +++ b/studio/backend/core/inference/mcp_config_import.py @@ -60,14 +60,19 @@ def _enabled_from_spec(label: str, spec: dict) -> tuple[Optional[bool], Optional return not disabled, None -def _parse_entry(name: str, spec: object) -> tuple[Optional[ParsedMcpEntry], Optional[str]]: +def _parse_entry( + name: str, spec: object +) -> tuple[Optional[ParsedMcpEntry], Optional[str]]: label = str(name).strip() if not label: return None, "Server entry has an empty name." if not isinstance(spec, dict): return None, f"{label}: entry must be an object." if _has_variable_reference(spec): - return None, f"{label}: VS Code variable references are not supported by import." + return ( + None, + f"{label}: VS Code variable references are not supported by import.", + ) is_enabled, error = _enabled_from_spec(label, spec) if error: @@ -91,8 +96,13 @@ def _parse_entry(name: str, spec: object) -> tuple[Optional[ParsedMcpEntry], Opt if sandbox_enabled is not None and not isinstance(sandbox_enabled, bool): return None, f"{label}: 'sandboxEnabled' must be true or false." if sandbox_enabled: - return None, f"{label}: sandboxed stdio servers cannot be preserved by import." - unsupported = [field for field in _UNSUPPORTED_STDIO_FIELDS if spec.get(field) is not None] + return ( + None, + f"{label}: sandboxed stdio servers cannot be preserved by import.", + ) + unsupported = [ + field for field in _UNSUPPORTED_STDIO_FIELDS if spec.get(field) is not None + ] if unsupported: return None, f"{label}: import cannot preserve {', '.join(unsupported)}." if spec.get("oauth") is not None: @@ -104,7 +114,10 @@ def _parse_entry(name: str, spec: object) -> tuple[Optional[ParsedMcpEntry], Opt if env is not None and not isinstance(env, dict): return None, f"{label}: 'env' must be an object." if _has_null_value(env): - return None, f"{label}: null environment values are not supported by import." + return ( + None, + f"{label}: null environment values are not supported by import.", + ) url = join_stdio_command([command, *(str(a) for a in args)]) headers = _coerce_str_dict(env) if env else None return ParsedMcpEntry(label, url, headers, True, is_enabled = is_enabled), None @@ -120,12 +133,21 @@ def _parse_entry(name: str, spec: object) -> tuple[Optional[ParsedMcpEntry], Opt field for field in _UNSUPPORTED_TIMEOUT_FIELDS if spec.get(field) is not None ] if unsupported_timeout: - return None, f"{label}: import cannot preserve {', '.join(unsupported_timeout)}." + return ( + None, + f"{label}: import cannot preserve {', '.join(unsupported_timeout)}.", + ) url_infers_sse = url.rstrip("/").endswith("/sse") if entry_type == "sse" and not url_infers_sse: - return None, f"{label}: explicit SSE transport cannot be preserved for this URL." + return ( + None, + f"{label}: explicit SSE transport cannot be preserved for this URL.", + ) if entry_type in _HTTP_REMOTE_TYPES and url_infers_sse: - return None, f"{label}: explicit HTTP transport cannot be preserved for this URL." + return ( + None, + f"{label}: explicit HTTP transport cannot be preserved for this URL.", + ) oauth_raw = spec.get("oauth") if oauth_raw is not None and not isinstance(oauth_raw, dict): return None, f"{label}: 'oauth' must be an object." diff --git a/studio/backend/core/inference/mlx_inference.py b/studio/backend/core/inference/mlx_inference.py index 5c7799152f..f3f3eb36d5 100644 --- a/studio/backend/core/inference/mlx_inference.py +++ b/studio/backend/core/inference/mlx_inference.py @@ -296,7 +296,9 @@ class MLXInferenceBackend: elif isinstance(content, list): # Prepend image if not already present has_image = any( - p.get("type") == "image" for p in content if isinstance(p, dict) + p.get("type") == "image" + for p in content + if isinstance(p, dict) ) if not has_image: content.insert(0, {"type": "image"}) @@ -366,7 +368,9 @@ class MLXInferenceBackend: preserve_thinking = preserve_thinking, ) if prompt is None: - raise RuntimeError("apply_chat_template returned None — tokenizer may be incompatible") + raise RuntimeError( + "apply_chat_template returned None — tokenizer may be incompatible" + ) sampler = make_sampler( temp = temperature, @@ -513,7 +517,9 @@ class MLXInferenceBackend: **vlm_kwargs, ): final_response = response - token_text = response.text if hasattr(response, "text") else str(response) + token_text = ( + response.text if hasattr(response, "text") else str(response) + ) cumulative += token_text yield cumulative if cancel_event and cancel_event.is_set(): diff --git a/studio/backend/core/inference/orchestrator.py b/studio/backend/core/inference/orchestrator.py index 4ccac2912e..c57bd7c87e 100644 --- a/studio/backend/core/inference/orchestrator.py +++ b/studio/backend/core/inference/orchestrator.py @@ -89,7 +89,9 @@ class InferenceOrchestrator: atexit.register(self._cleanup) logger.info("InferenceOrchestrator initialized (subprocess mode)") - threading.Thread(target = self._fetch_top_models, daemon = True, name = "top-models").start() + threading.Thread( + target = self._fetch_top_models, daemon = True, name = "top-models" + ).start() # ------------------------------------------------------------------ # Default models (top GGUFs fetched dynamically from HF) @@ -129,12 +131,14 @@ class InferenceOrchestrator: if resp.status_code == 200: models = resp.json() # Top 40 GGUFs (deep pool for frontend infinite scroll) - gguf_ids = [m["id"] for m in models if m.get("id", "").upper().endswith("-GGUF")][ - :40 - ] + gguf_ids = [ + m["id"] for m in models if m.get("id", "").upper().endswith("-GGUF") + ][:40] # Top 40 non-GGUF hub models hub_ids = [ - m["id"] for m in models if not m.get("id", "").upper().endswith("-GGUF") + m["id"] + for m in models + if not m.get("id", "").upper().endswith("-GGUF") ][:40] if gguf_ids: self._top_gguf_cache = gguf_ids @@ -272,7 +276,8 @@ class InferenceOrchestrator: "Try a smaller model, lower context length, or close other GPU-heavy apps." ) return ( - f"{message}{suffix} " f"Details: pid={pid}, signal={sig_name}, exitcode={exitcode}." + f"{message}{suffix} " + f"Details: pid={pid}, signal={sig_name}, exitcode={exitcode}." ) return f"{message} Details: pid={pid}, exitcode={exitcode}." @@ -355,7 +360,8 @@ class InferenceOrchestrator: ) raise RuntimeError( - f"Timeout waiting for '{expected_type}' response " f"(no activity for {timeout}s)" + f"Timeout waiting for '{expected_type}' response " + f"(no activity for {timeout}s)" ) def _drain_queue(self) -> list: @@ -733,9 +739,13 @@ class InferenceOrchestrator: # without re-entering the subprocess. _tpl_info = model_info.get("chat_template_info") if isinstance(_tpl_info, dict): - self.models[self.active_model_name]["chat_template_info"] = _tpl_info + self.models[self.active_model_name]["chat_template_info"] = ( + _tpl_info + ) self.loading_models.discard(model_name) - logger.info("Model '%s' loaded successfully in subprocess", model_name) + logger.info( + "Model '%s' loaded successfully in subprocess", model_name + ) return True else: error = resp.get("error", "Failed to load model") @@ -1166,7 +1176,9 @@ class InferenceOrchestrator: if resp is None: if not self._ensure_subprocess_alive(): - raise RuntimeError(self._subprocess_crash_message("audio generation")) + raise RuntimeError( + self._subprocess_crash_message("audio generation") + ) continue rtype = resp.get("type", "") @@ -1258,7 +1270,9 @@ class InferenceOrchestrator: # numpy array -> list for mp.Queue serialization audio_data = ( - audio_array.tolist() if hasattr(audio_array, "tolist") else list(audio_array) + audio_array.tolist() + if hasattr(audio_array, "tolist") + else list(audio_array) ) cmd = { @@ -1288,7 +1302,10 @@ class InferenceOrchestrator: if resp is None: if not self._ensure_subprocess_alive(): - yield ("Error: " + self._subprocess_crash_message("audio input generation")) + yield ( + "Error: " + + self._subprocess_crash_message("audio input generation") + ) return continue diff --git a/studio/backend/core/inference/pricing.py b/studio/backend/core/inference/pricing.py index 3b611d3596..980d3f98ef 100644 --- a/studio/backend/core/inference/pricing.py +++ b/studio/backend/core/inference/pricing.py @@ -103,7 +103,9 @@ def _lookup(provider: str, model: str) -> Optional[dict[str, float]]: return None -def calculate_cost(provider: str, model: str, usage: dict[str, Any]) -> dict[str, float]: +def calculate_cost( + provider: str, model: str, usage: dict[str, Any] +) -> dict[str, float]: """Return a per-turn USD cost breakdown (per-bucket + total). Unknown model -> ``priced`` False and USD fields 0.0 (token counts still report). @@ -131,7 +133,8 @@ def calculate_cost(provider: str, model: str, usage: dict[str, Any]) -> dict[str # Clamp >=0 so corrupted payloads can't produce a negative bill. cache_creation = max(0, int(usage.get("cache_creation_input_tokens") or 0)) cache_read_native_present = ( - "cache_read_input_tokens" in usage and usage.get("cache_read_input_tokens") is not None + "cache_read_input_tokens" in usage + and usage.get("cache_read_input_tokens") is not None ) cache_read = max(0, int(usage.get("cache_read_input_tokens") or 0)) # Fall back to mirrored prompt_tokens_details only when native @@ -214,10 +217,14 @@ def calculate_cost(provider: str, model: str, usage: dict[str, Any]) -> dict[str if cc_5m + cc_1h == 0 and cache_creation > 0: # No breakdown -- assume default 5m pool. cc_5m = cache_creation - out["cache_write_usd"] = (cc_5m / 1_000_000.0) * base * ANTHROPIC_CACHE_5M_WRITE_MULT + ( + out["cache_write_usd"] = ( + cc_5m / 1_000_000.0 + ) * base * ANTHROPIC_CACHE_5M_WRITE_MULT + ( cc_1h / 1_000_000.0 ) * base * ANTHROPIC_CACHE_1H_WRITE_MULT - out["cache_read_usd"] = (cache_read / 1_000_000.0) * base * ANTHROPIC_CACHE_READ_MULT + out["cache_read_usd"] = ( + (cache_read / 1_000_000.0) * base * ANTHROPIC_CACHE_READ_MULT + ) # Server-tool surcharges. srv = usage.get("server_tool_use") or {} if isinstance(srv, dict): @@ -234,7 +241,9 @@ def calculate_cost(provider: str, model: str, usage: dict[str, Any]) -> dict[str if cache_read > 0: non_cached_input = max(0, input_tokens - cache_read) out["input_usd"] = (non_cached_input / 1_000_000.0) * base - out["cache_read_usd"] = (cache_read / 1_000_000.0) * base * OPENAI_CACHE_READ_MULT + out["cache_read_usd"] = ( + (cache_read / 1_000_000.0) * base * OPENAI_CACHE_READ_MULT + ) # OpenAI server-tool surcharges arrive under `openai_tool_use` # (normalised by the SSE finaliser from output items). srv = usage.get("openai_tool_use") or {} diff --git a/studio/backend/core/inference/safetensors_agentic.py b/studio/backend/core/inference/safetensors_agentic.py index 3b6a393f3d..906fb0e2a3 100644 --- a/studio/backend/core/inference/safetensors_agentic.py +++ b/studio/backend/core/inference/safetensors_agentic.py @@ -92,7 +92,9 @@ def _detect_render_html_tool_start(content: str) -> bool: if not function_match and tool_call_index < 0: return False - if function_match and (tool_call_index < 0 or function_match.start() < tool_call_index): + if function_match and ( + tool_call_index < 0 or function_match.start() < tool_call_index + ): return function_match.group(1) == "render_html" if tool_call_index >= 0: @@ -182,7 +184,9 @@ def run_safetensors_tool_loop( # Forced first-pass RAG (mirrors the GGUF loop) so doc Qs don't lose to web_search. from core.inference.tools import build_rag_autoinject - _auto = None if confirm_tool_calls else build_rag_autoinject(conversation, rag_scope) + _auto = ( + None if confirm_tool_calls else build_rag_autoinject(conversation, rag_scope) + ) if _auto: for _ev in _auto["events"]: yield _ev @@ -201,7 +205,9 @@ def run_safetensors_tool_loop( def _tool_succeeded(tool_name: str) -> bool: key_prefix = f"{tool_name}:" return any( - record.executed and not record.is_error and record.key.startswith(key_prefix) + record.executed + and not record.is_error + and record.key.startswith(key_prefix) for record in tool_controller.history ) @@ -226,7 +232,9 @@ def run_safetensors_tool_loop( final_attempt_done = True active_tools = [] - tool_protocol_active = not final_attempt_done and (unrestricted_tools or bool(active_tools)) + tool_protocol_active = not final_attempt_done and ( + unrestricted_tools or bool(active_tools) + ) tool_xml_signals = TOOL_XML_SIGNALS if tool_protocol_active else () detect_state = _state_buffering @@ -515,11 +523,15 @@ def run_safetensors_tool_loop( conversation.append(assistant_msg) assistant_appended = True else: - assistant_msg.setdefault("tool_calls", []).append(decision.as_assistant_tool_call()) + assistant_msg.setdefault("tool_calls", []).append( + decision.as_assistant_tool_call() + ) needs_confirm = bool(confirm_tool_calls) approval_id = new_approval_id() if needs_confirm else "" - decision_slot = begin_tool_decision(session_id, approval_id) if needs_confirm else None + decision_slot = ( + begin_tool_decision(session_id, approval_id) if needs_confirm else None + ) start_event = decision.tool_start_event() start_event["approval_id"] = approval_id start_event["awaiting_confirmation"] = needs_confirm diff --git a/studio/backend/core/inference/tool_call_parser.py b/studio/backend/core/inference/tool_call_parser.py index 8d5d45269e..938abe3f01 100644 --- a/studio/backend/core/inference/tool_call_parser.py +++ b/studio/backend/core/inference/tool_call_parser.py @@ -194,7 +194,11 @@ def parse_tool_calls_from_text( for idx, fm in enumerate(func_starts): func_name = fm.group(1) body_start = fm.end() - next_func = func_starts[idx + 1].start() if idx + 1 < len(func_starts) else len(content) + next_func = ( + func_starts[idx + 1].start() + if idx + 1 < len(func_starts) + else len(content) + ) end_tag = _TC_END_TAG_RE.search(content[body_start:]) if end_tag: body_end = body_start + end_tag.start() diff --git a/studio/backend/core/inference/tool_loop_controller.py b/studio/backend/core/inference/tool_loop_controller.py index cb751ede3d..7918ba0b21 100644 --- a/studio/backend/core/inference/tool_loop_controller.py +++ b/studio/backend/core/inference/tool_loop_controller.py @@ -287,7 +287,9 @@ class ToolLoopController: self._restrict_to_allowed = tools is not None self._tools = [copy.deepcopy(dict(tool)) for tool in (tools or [])] self._allowed_tool_names = { - name for name in (_tool_name_from_schema(tool) for tool in self._tools) if name + name + for name in (_tool_name_from_schema(tool) for tool in self._tools) + if name } self._auto_heal_tool_calls = auto_heal_tool_calls self._one_shot_tools = one_shot_tools @@ -364,7 +366,9 @@ class ToolLoopController: noop_result = noop, ) - def record_result(self, decision: ToolCallDecision, result: Any) -> ToolCallCompletion: + def record_result( + self, decision: ToolCallDecision, result: Any + ) -> ToolCallCompletion: """Record a real tool execution and return model/frontend payload helpers.""" result_text = result if isinstance(result, str) else str(result) failed = is_tool_error(result_text) diff --git a/studio/backend/core/inference/tools.py b/studio/backend/core/inference/tools.py index 43c9610282..7e4c2663d7 100644 --- a/studio/backend/core/inference/tools.py +++ b/studio/backend/core/inference/tools.py @@ -125,7 +125,9 @@ _BLOCKED_COMMANDS = ( ) -_SHELL_SEPARATORS = frozenset({";", "&&", "||", "|", "&", "\n", "(", ")", "`", "{", "}"}) +_SHELL_SEPARATORS = frozenset( + {";", "&&", "||", "|", "&", "\n", "(", ")", "`", "{", "}"} +) # Bash keywords starting a new command position (then $cmd, do $cmd, etc.). _SHELL_KEYWORDS_AS_SEP = frozenset({"then", "do", "else", "elif"}) # Wrappers whose next non-flag argument is the command Bash will exec. @@ -248,7 +250,9 @@ def _find_blocked_commands(command: str) -> set[str]: tok_lower = token.lower() # Match -c exactly, or combined flags ending in c (e.g. -lc, -xc) is_unix_c = tok_lower == "-c" or ( - tok_lower.startswith("-") and tok_lower.endswith("c") and not tok_lower.startswith("--") + tok_lower.startswith("-") + and tok_lower.endswith("c") + and not tok_lower.startswith("--") ) is_win_c = tok_lower == "/c" if not (is_unix_c or is_win_c) or i < 1 or i + 1 >= len(tokens): @@ -352,11 +356,18 @@ def _sandbox_preexec(): except (ValueError, OSError, AttributeError): pass try: - _resource.setrlimit(_resource.RLIMIT_FSIZE, (100 * 1024 * 1024, 100 * 1024 * 1024)) + _resource.setrlimit( + _resource.RLIMIT_FSIZE, (100 * 1024 * 1024, 100 * 1024 * 1024) + ) except (ValueError, OSError): pass try: - as_bytes = int(os.environ.get("UNSLOTH_STUDIO_SANDBOX_AS_GB", "8")) * 1024 * 1024 * 1024 + as_bytes = ( + int(os.environ.get("UNSLOTH_STUDIO_SANDBOX_AS_GB", "8")) + * 1024 + * 1024 + * 1024 + ) _resource.setrlimit(_resource.RLIMIT_AS, (as_bytes, as_bytes)) except (ValueError, OSError, AttributeError): pass @@ -371,7 +382,9 @@ def _sandbox_preexec(): # when the parent's hard cap is below the request. nofile = int(os.environ.get("UNSLOTH_STUDIO_SANDBOX_NOFILE", "16384")) _soft_cur, hard_cur = _resource.getrlimit(_resource.RLIMIT_NOFILE) - target = nofile if hard_cur == _resource.RLIM_INFINITY else min(nofile, hard_cur) + target = ( + nofile if hard_cur == _resource.RLIM_INFINITY else min(nofile, hard_cur) + ) _resource.setrlimit(_resource.RLIMIT_NOFILE, (target, target)) except (ValueError, OSError, AttributeError): pass @@ -404,7 +417,9 @@ def _get_project_workdir(session_id: str) -> str | None: from storage.studio_db import ensure_chat_project_workspace project = ensure_chat_project_workspace(project_id) except Exception: - logger.warning("Failed to resolve project sandbox for %s", session_id, exc_info = True) + logger.warning( + "Failed to resolve project sandbox for %s", session_id, exc_info = True + ) return None if not project: return None @@ -435,7 +450,9 @@ def _get_workdir(session_id: str | None = None) -> str: workdir = project_workdir elif session_id and _SESSION_ID_RE.match(session_id): workdir = os.path.join(sandbox_root, session_id) - if not os.path.realpath(workdir).startswith(os.path.realpath(sandbox_root) + os.sep): + if not os.path.realpath(workdir).startswith( + os.path.realpath(sandbox_root) + os.sep + ): workdir = os.path.join(sandbox_root, "_invalid") elif session_id: workdir = os.path.join(sandbox_root, "_invalid") @@ -614,7 +631,9 @@ def _mcp_specs_for_server(server: dict, mcp_tools: list[dict]) -> list[dict]: continue # Duplicate tool names would also 400 OpenAI; drop dupes. if name in seen_names: - logger.warning("Skipping duplicate MCP tool '%s' on '%s'.", raw_name, display) + logger.warning( + "Skipping duplicate MCP tool '%s' on '%s'.", raw_name, display + ) continue seen_names.add(name) specs.append( @@ -623,7 +642,8 @@ def _mcp_specs_for_server(server: dict, mcp_tools: list[dict]) -> list[dict]: "function": { "name": name, "description": f"[{display}] {tool.get('description') or ''}".strip(), - "parameters": tool.get("inputSchema") or {"type": "object", "properties": {}}, + "parameters": tool.get("inputSchema") + or {"type": "object", "properties": {}}, }, } ) @@ -643,7 +663,9 @@ async def get_enabled_mcp_tools() -> list[dict]: # server gets re-probed -- and blocks the send for the full timeout -- on # every message. uncached = [ - s for s in servers if get_cached_tools(s["id"]) is None and not in_failure_cooloff(s["id"]) + s + for s in servers + if get_cached_tools(s["id"]) is None and not in_failure_cooloff(s["id"]) ] if uncached: results = await asyncio.gather( @@ -731,7 +753,9 @@ def execute_tool( ``rag_scope``: hidden per-request RAG context the model never sees; consumed by ``search_knowledge_base``. """ - logger.info(f"execute_tool: name={name}, session_id={session_id}, timeout={timeout}") + logger.info( + f"execute_tool: name={name}, session_id={session_id}, timeout={timeout}" + ) effective_timeout = _EXEC_TIMEOUT if timeout is _TIMEOUT_UNSET else timeout if name == "search_knowledge_base": return _search_knowledge_base(arguments, rag_scope) @@ -765,9 +789,13 @@ def execute_tool( timeout = effective_timeout, ) if name == "python": - return _python_exec(arguments.get("code", ""), cancel_event, effective_timeout, session_id) + return _python_exec( + arguments.get("code", ""), cancel_event, effective_timeout, session_id + ) if name == "terminal": - return _bash_exec(arguments.get("command", ""), cancel_event, effective_timeout, session_id) + return _bash_exec( + arguments.get("command", ""), cancel_event, effective_timeout, session_id + ) return f"Unknown tool: {name}" @@ -875,7 +903,9 @@ def _last_user_text(conversation: list[dict]) -> str: return "" -def build_rag_autoinject(conversation: list[dict], rag_scope: dict | None) -> dict | None: +def build_rag_autoinject( + conversation: list[dict], rag_scope: dict | None +) -> dict | None: """Pre-retrieve the latest user turn; if a hit clears the cosine floor return ``{"events": [...], "messages": [...]}`` to splice into the loop, else ``None``. Toggle via ``rag_scope.autoinject`` (else env ``RAG_AUTOINJECT``); floor via @@ -971,7 +1001,9 @@ def build_rag_autoinject(conversation: list[dict], rag_scope: dict | None) -> di "content": text, }, ] - logger.info("RAG auto-inject: %d passage(s) >= %.2f for %r", len(sources), floor, query[:80]) + logger.info( + "RAG auto-inject: %d passage(s) >= %.2f for %r", len(sources), floor, query[:80] + ) return {"events": events, "messages": messages} @@ -1138,7 +1170,9 @@ def _fetch_page_text( resp = opener.open(req, timeout = timeout) except _HTTPError as e: if e.code not in (301, 302, 303, 307, 308): - return f"Failed to fetch URL: HTTP {e.code} {getattr(e, 'reason', '')}" + return ( + f"Failed to fetch URL: HTTP {e.code} {getattr(e, 'reason', '')}" + ) location = e.headers.get("Location") if not location: return "Failed to fetch URL: redirect missing Location header." @@ -1398,7 +1432,9 @@ def _check_signal_escape_patterns(code: str): if func_name: if func_name in ("signal.signal", "signal"): if len(node.args) >= 1: - if _ast_name_matches(node.args[0], ("SIGALRM", "signal.SIGALRM")): + if _ast_name_matches( + node.args[0], ("SIGALRM", "signal.SIGALRM") + ): signal_tampering.append( { "type": "signal_handler_override", @@ -1408,7 +1444,9 @@ def _check_signal_escape_patterns(code: str): ) elif func_name in ("signal.setitimer", "setitimer"): if len(node.args) >= 1: - if _ast_name_matches(node.args[0], ("ITIMER_REAL", "signal.ITIMER_REAL")): + if _ast_name_matches( + node.args[0], ("ITIMER_REAL", "signal.ITIMER_REAL") + ): signal_tampering.append( { "type": "timer_manipulation", @@ -1461,7 +1499,9 @@ def _check_signal_escape_patterns(code: str): else: has_opaque_kwargs = True - cmd_kw_values = [v for k, v in expanded_kwargs.items() if k in _CMD_KWARGS] + cmd_kw_values = [ + v for k, v in expanded_kwargs.items() if k in _CMD_KWARGS + ] all_call_args = list(node.args) + cmd_kw_values blocked_in_args = _check_args_for_blocked(all_call_args) @@ -1471,7 +1511,9 @@ def _check_signal_escape_patterns(code: str): { "type": "shell_escape_dynamic", "line": node.lineno, - "description": (f"{shell_func}() called with dynamic **kwargs"), + "description": ( + f"{shell_func}() called with dynamic **kwargs" + ), } ) elif blocked_in_args: @@ -1502,7 +1544,8 @@ def _check_signal_escape_patterns(code: str): ) shell_node = expanded_kwargs.get("shell") shell_safe = shell_node is None or ( - isinstance(shell_node, ast.Constant) and shell_node.value is False + isinstance(shell_node, ast.Constant) + and shell_node.value is False ) # Dynamic shell-exec args (chr/format/concat bypasses). if ( @@ -1515,10 +1558,15 @@ def _check_signal_escape_patterns(code: str): if _extract_string_from_node(n) is not None: return True if isinstance(n, (ast.List, ast.Tuple)): - return all(_extract_string_from_node(e) is not None for e in n.elts) + return all( + _extract_string_from_node(e) is not None + for e in n.elts + ) return False - has_non_literal = any(not _is_safe_literal(a) for a in all_call_args) + has_non_literal = any( + not _is_safe_literal(a) for a in all_call_args + ) if has_non_literal: shell_escapes.append( { @@ -1775,7 +1823,9 @@ def _check_signal_escape_patterns(code: str): "/etc/sudoers", "/etc/ssh/", ) - _SENSITIVE_FILE_RE = re.compile(r"^/proc/(?:self|\d+)/(?:environ|cmdline|task/\d+/environ)$") + _SENSITIVE_FILE_RE = re.compile( + r"^/proc/(?:self|\d+)/(?:environ|cmdline|task/\d+/environ)$" + ) def _normalize_host(host: str) -> str: if not host: @@ -1818,9 +1868,15 @@ def _check_signal_escape_patterns(code: str): return True if kw.arg == "data": v = kw.value - if isinstance(v, ast.Call) and isinstance(v.func, ast.Name) and v.func.id == "open": + if ( + isinstance(v, ast.Call) + and isinstance(v.func, ast.Name) + and v.func.id == "open" + ): return True - if isinstance(v, ast.Constant) and isinstance(v.value, (bytes, bytearray)): + if isinstance(v, ast.Constant) and isinstance( + v.value, (bytes, bytearray) + ): return True return False @@ -1951,7 +2007,9 @@ def _check_signal_escape_patterns(code: str): """Whether the path argument resolves to a sandbox-local literal.""" if node is None: return False - if isinstance(node, ast.Constant) and isinstance(node.value, (bytes, bytearray)): + if isinstance(node, ast.Constant) and isinstance( + node.value, (bytes, bytearray) + ): return True # inline bytes, no file access if isinstance(node, ast.Constant) and isinstance(node.value, str): return _is_safe_relative_path(node.value) @@ -2037,7 +2095,11 @@ def _check_signal_escape_patterns(code: str): ) # Direct sock.connect((host, port)) bypasses the FQ-prefix branch. - if isinstance(node.func, ast.Attribute) and node.func.attr == "connect" and node.args: + if ( + isinstance(node.func, ast.Attribute) + and node.func.attr == "connect" + and node.args + ): a0 = node.args[0] host_lit = None if isinstance(a0, ast.Tuple) and a0.elts: @@ -2074,7 +2136,9 @@ def _check_signal_escape_patterns(code: str): { "type": "upload_blocked", "line": getattr(node, "lineno", -1), - "description": ("Blocked: file upload disallowed in sandbox"), + "description": ( + "Blocked: file upload disallowed in sandbox" + ), } ) @@ -2175,18 +2239,28 @@ def _check_code_safety(code: str) -> str | None: if info.get("error"): return None - reasons = [item.get("description", "") for item in info.get("signal_tampering", [])] - shell_reasons = [item.get("description", "") for item in info.get("shell_escapes", [])] + reasons = [ + item.get("description", "") for item in info.get("signal_tampering", []) + ] + shell_reasons = [ + item.get("description", "") for item in info.get("shell_escapes", []) + ] exception_reasons = [ item.get("description", "") for item in info.get("exception_catching", []) ] - network_reasons = [item.get("description", "") for item in info.get("network_calls", [])] + network_reasons = [ + item.get("description", "") for item in info.get("network_calls", []) + ] file_reasons = [ item.get("description", "") for item in info.get("sensitive_file_reads", []) ] all_reasons = [ r - for r in reasons + shell_reasons + exception_reasons + network_reasons + file_reasons + for r in reasons + + shell_reasons + + exception_reasons + + network_reasons + + file_reasons if r ] if all_reasons: @@ -2266,7 +2340,9 @@ def _python_exec( except OSError: pass try: - fd, tmp_path = tempfile.mkstemp(suffix = ".py", prefix = "studio_exec_", dir = workdir) + fd, tmp_path = tempfile.mkstemp( + suffix = ".py", prefix = "studio_exec_", dir = workdir + ) with os.fdopen(fd, "w") as f: f.write(code) diff --git a/studio/backend/core/inference/worker.py b/studio/backend/core/inference/worker.py index cc79654087..792031f811 100644 --- a/studio/backend/core/inference/worker.py +++ b/studio/backend/core/inference/worker.py @@ -89,7 +89,9 @@ def _build_model_config(config: dict): return mc -def _get_hf_download_state(model_names: list[str] | None = None) -> tuple[int, bool] | None: +def _get_hf_download_state( + model_names: list[str] | None = None, +) -> tuple[int, bool] | None: """Return (total_bytes, has_incomplete) for the HF Hub cache, or None on error. With *model_names*, only those models' ``blobs/`` dirs are checked (faster); @@ -240,10 +242,14 @@ def _handle_load(backend, config: dict, resp_queue: Any) -> None: adapter_cfg = json.load(f) training_method = adapter_cfg.get("unsloth_training_method") if training_method == "lora" and load_in_4bit: - logger.info("adapter_config.json says lora — setting load_in_4bit=False") + logger.info( + "adapter_config.json says lora — setting load_in_4bit=False" + ) load_in_4bit = False elif training_method == "qlora" and not load_in_4bit: - logger.info("adapter_config.json says qlora — setting load_in_4bit=True") + logger.info( + "adapter_config.json says qlora — setting load_in_4bit=True" + ) load_in_4bit = True elif not training_method: if ( @@ -513,7 +519,9 @@ def _handle_generate_audio(backend, cmd: dict, resp_queue: Any) -> None: ) -def _handle_generate_audio_input(backend, cmd: dict, resp_queue: Any, cancel_event) -> None: +def _handle_generate_audio_input( + backend, cmd: dict, resp_queue: Any, cancel_event +) -> None: """Handle audio input generation (ASR/Whisper) — streams text tokens back.""" request_id = cmd.get("request_id", "") @@ -548,7 +556,9 @@ def _handle_generate_audio_input(backend, cmd: dict, resp_queue: Any, cancel_eve for text_chunk in generator: if cancel_event.is_set(): - logger.info("Audio input generation cancelled for request %s", request_id) + logger.info( + "Audio input generation cancelled for request %s", request_id + ) break _send_response( @@ -615,7 +625,9 @@ def _handle_unload(backend, cmd: dict, resp_queue: Any) -> None: ) -def run_inference_process(*, cmd_queue: Any, resp_queue: Any, cancel_event, config: dict) -> None: +def run_inference_process( + *, cmd_queue: Any, resp_queue: Any, cancel_event, config: dict +) -> None: """Subprocess entrypoint. Persistent — runs the command loop until shutdown. Args: @@ -625,7 +637,9 @@ def run_inference_process(*, cmd_queue: Any, resp_queue: Any, cancel_event, conf config: Initial configuration dict with model info. """ os.environ["TOKENIZERS_PARALLELISM"] = "false" - os.environ["PYTHONWARNINGS"] = "ignore" # Suppress warnings at C-level before imports + os.environ["PYTHONWARNINGS"] = ( + "ignore" # Suppress warnings at C-level before imports + ) if config.get("disable_xet"): os.environ["HF_HUB_DISABLE_XET"] = "1" @@ -930,7 +944,9 @@ def run_inference_process(*, cmd_queue: Any, resp_queue: Any, cancel_event, conf ) except Exception as exc: - logger.error("Error handling command '%s': %s", cmd_type, exc, exc_info = True) + logger.error( + "Error handling command '%s': %s", cmd_type, exc, exc_info = True + ) _send_response( resp_queue, { diff --git a/studio/backend/core/rag/captioner.py b/studio/backend/core/rag/captioner.py index be8e341064..9a3708e57b 100644 --- a/studio/backend/core/rag/captioner.py +++ b/studio/backend/core/rag/captioner.py @@ -26,14 +26,18 @@ def vision_endpoint() -> tuple[str, str] | None: try: from routes.inference import get_llama_cpp_backend backend = get_llama_cpp_backend() - if getattr(backend, "is_loaded", False) and getattr(backend, "is_vision", False): + if getattr(backend, "is_loaded", False) and getattr( + backend, "is_vision", False + ): return backend.base_url, "local" except Exception: # noqa: BLE001 - never let discovery break ingestion return None return None -def _caption_one(base_url: str, model: str, image_bytes: bytes, timeout: float) -> str | None: +def _caption_one( + base_url: str, model: str, image_bytes: bytes, timeout: float +) -> str | None: import httpx data_url = "data:image/png;base64," + base64.b64encode(image_bytes).decode("ascii") diff --git a/studio/backend/core/rag/chunking.py b/studio/backend/core/rag/chunking.py index c64acb4c60..72228ebd97 100644 --- a/studio/backend/core/rag/chunking.py +++ b/studio/backend/core/rag/chunking.py @@ -27,7 +27,9 @@ class Chunk: page_char_end: int -def _split(text: str, seps: tuple[str, ...], max_tokens: int, count: TokenCounter) -> list[str]: +def _split( + text: str, seps: tuple[str, ...], max_tokens: int, count: TokenCounter +) -> list[str]: """Recursively split into pieces each <= max_tokens (best effort). Pieces rejoin to ``text`` exactly, so offsets are a running length.""" if count(text) <= max_tokens: @@ -41,7 +43,9 @@ def _split(text: str, seps: tuple[str, ...], max_tokens: int, count: TokenCounte out: list[str] = [] for p in parts: out.extend( - [p] if count(p) <= max_tokens else _split(p, seps[i + 1 :], max_tokens, count) + [p] + if count(p) <= max_tokens + else _split(p, seps[i + 1 :], max_tokens, count) ) return [p for p in out if p] n = max(1, max_tokens * 4) @@ -49,7 +53,11 @@ def _split(text: str, seps: tuple[str, ...], max_tokens: int, count: TokenCounte def _merge( - pieces: list[str], starts: list[int], max_tokens: int, overlap: int, count: TokenCounter + pieces: list[str], + starts: list[int], + max_tokens: int, + overlap: int, + count: TokenCounter, ) -> list[tuple[str, int, int]]: """Greedy-merge pieces into <= max_tokens chunks with token overlap. ``starts[i]`` is ``pieces[i]``'s page char offset; returns @@ -106,7 +114,9 @@ def chunk_pages( for piece in pieces: starts.append(cursor) cursor += len(piece) - for text, char_start, char_end in _merge(pieces, starts, max_tokens, overlap, count): + for text, char_start, char_end in _merge( + pieces, starts, max_tokens, overlap, count + ): out.append( Chunk( text = text, diff --git a/studio/backend/core/rag/config.py b/studio/backend/core/rag/config.py index 993423683c..b6ea9b083f 100644 --- a/studio/backend/core/rag/config.py +++ b/studio/backend/core/rag/config.py @@ -31,7 +31,9 @@ CAPTION_TIMEOUT_S = float(os.environ.get("RAG_CAPTION_TIMEOUT_S", "30")) EMBED_BACKEND = os.environ.get("RAG_EMBED_BACKEND", "auto") # llama-server backend only. F16 over Q8_0: faster (no per-block dequant for this # tiny model) and exact vs fp32, for ~30MB more on disk. -EMBED_GGUF_REPO = os.environ.get("RAG_EMBED_GGUF_REPO", "unsloth/bge-small-en-v1.5-GGUF") +EMBED_GGUF_REPO = os.environ.get( + "RAG_EMBED_GGUF_REPO", "unsloth/bge-small-en-v1.5-GGUF" +) EMBED_GGUF_VARIANT = os.environ.get("RAG_EMBED_GGUF_VARIANT", "F16") EMBED_DEVICE = os.environ.get("RAG_EMBED_DEVICE", "auto") # "auto" | "gpu" | "cpu" EMBED_HOST = os.environ.get("RAG_EMBED_HOST", "127.0.0.1") diff --git a/studio/backend/core/rag/embed_llama_server.py b/studio/backend/core/rag/embed_llama_server.py index 2fadbcc6c6..b2b2871cd5 100644 --- a/studio/backend/core/rag/embed_llama_server.py +++ b/studio/backend/core/rag/embed_llama_server.py @@ -123,7 +123,9 @@ class LlamaServerBackend: repo = config.EMBED_GGUF_REPO token = os.environ.get("HF_TOKEN") or None - files = [f for f in list_repo_files(repo, token = token) if f.lower().endswith(".gguf")] + files = [ + f for f in list_repo_files(repo, token = token) if f.lower().endswith(".gguf") + ] files = [f for f in files if "mmproj" not in f.lower()] if not files: raise RuntimeError(f"no .gguf file found in embedder repo {repo!r}") @@ -161,7 +163,9 @@ class LlamaServerBackend: gpus = LlamaCppBackend._get_gpu_free_memory() # [(idx, free_mib)], honors CVD return any(free >= LlamaServerBackend._MIN_GPU_FREE_MIB for _, free in gpus) - def _build_cmd(self, binary: str, model_path: str, port: int, *, use_gpu: bool) -> list[str]: + def _build_cmd( + self, binary: str, model_path: str, port: int, *, use_gpu: bool + ) -> list[str]: # No --embd-normalize (not in every build; we normalize in Python to match # the ST path). --fit off: don't auto-resize ctx/offload to device memory. cmd = [ @@ -204,8 +208,12 @@ class LlamaServerBackend: arch = platform.machine() lib_dirs = [binary_dir] for pattern in ( - os.path.join(sys.prefix, "lib", "python*", "site-packages", "nvidia", "cu*", "lib"), - os.path.join(sys.prefix, "lib", "python*", "site-packages", "nvidia", "cudnn", "lib"), + os.path.join( + sys.prefix, "lib", "python*", "site-packages", "nvidia", "cu*", "lib" + ), + os.path.join( + sys.prefix, "lib", "python*", "site-packages", "nvidia", "cudnn", "lib" + ), ): lib_dirs.extend(d for d in glob.glob(pattern) if os.path.isdir(d)) for cuda_lib in ( @@ -382,7 +390,9 @@ class LlamaServerBackend: raise RuntimeError( f"llama-server embedder POST {path} -> {e.response.status_code}: {body}" ) from e - raise RuntimeError(f"llama-server embedder POST {path} failed after retry") from last_exc + raise RuntimeError( + f"llama-server embedder POST {path} failed after retry" + ) from last_exc def encode( self, diff --git a/studio/backend/core/rag/ingestion.py b/studio/backend/core/rag/ingestion.py index c0c9a9f656..c048a05f20 100644 --- a/studio/backend/core/rag/ingestion.py +++ b/studio/backend/core/rag/ingestion.py @@ -108,7 +108,9 @@ def _run( stored_path, max_figures = config.CAPTION_MAX_IMAGES ) except Exception: - logger.warning("figure rendering failed for job %s", job_id, exc_info = True) + logger.warning( + "figure rendering failed for job %s", job_id, exc_info = True + ) figures = [] if figures: _progress(conn, job_id, "captioning", 0.2) @@ -139,12 +141,16 @@ def _run( from . import locators regions = locators.pdf_regions_for_chunks(stored_path, pages, chunks) except Exception: - logger.warning("pdf region location failed for job %s", job_id, exc_info = True) + logger.warning( + "pdf region location failed for job %s", job_id, exc_info = True + ) regions = None _progress(conn, job_id, "storing", 0.9) store.add_chunks(conn, scope, document_id, chunks, vectors, regions) - store.set_document_status(conn, document_id, "completed", num_chunks = len(chunks)) + store.set_document_status( + conn, document_id, "completed", num_chunks = len(chunks) + ) _set_job(conn, job_id, status = "completed", stage = "done", progress = 1.0) _emit(job_id, {"type": "complete", "num_chunks": len(chunks)}) @@ -267,7 +273,9 @@ def get_job_status(job_id: str) -> dict | None: """Read the persisted ingestion job row (status / stage / progress / error).""" conn = rag_db.get_connection() try: - row = conn.execute("SELECT * FROM ingestion_jobs WHERE id=?", (job_id,)).fetchone() + row = conn.execute( + "SELECT * FROM ingestion_jobs WHERE id=?", (job_id,) + ).fetchone() return dict(row) if row else None finally: conn.close() diff --git a/studio/backend/core/rag/locators.py b/studio/backend/core/rag/locators.py index 57c0487486..56e3727aa3 100644 --- a/studio/backend/core/rag/locators.py +++ b/studio/backend/core/rag/locators.py @@ -118,7 +118,9 @@ def _rects_from_words(page_words: list, indices: list[int], pw: float, ph: float return out -def _regions_for_match(doc: Any, page_text: str, match: LocatorMatch) -> list[dict[str, Any]]: +def _regions_for_match( + doc: Any, page_text: str, match: LocatorMatch +) -> list[dict[str, Any]]: try: if match.page_index < 0 or match.page_index >= len(doc): return [] @@ -145,7 +147,9 @@ def _regions_for_match(doc: Any, page_text: str, match: LocatorMatch) -> list[di return [] -def pdf_regions_for_chunks(pdf_path: Path, pages: list, chunks: list) -> list[list[dict[str, Any]]]: +def pdf_regions_for_chunks( + pdf_path: Path, pages: list, chunks: list +) -> list[list[dict[str, Any]]]: """Region rects per chunk (parallel to ``chunks``), keyed off each chunk's ``source_page_index`` / ``page_char_start`` / ``page_char_end``. Non-PDFs and failures yield [], never an exception.""" diff --git a/studio/backend/core/rag/retrieval.py b/studio/backend/core/rag/retrieval.py index fe6a033a52..85b0ac8e78 100644 --- a/studio/backend/core/rag/retrieval.py +++ b/studio/backend/core/rag/retrieval.py @@ -27,7 +27,10 @@ def retrieve_lexical( k: int | None = None, ) -> list[Hit]: k = k or config.TOP_K_LEXICAL - return [Hit(cid, s, lexical_score = s) for cid, s in store.search_lexical(conn, scope, query, k)] + return [ + Hit(cid, s, lexical_score = s) + for cid, s in store.search_lexical(conn, scope, query, k) + ] def retrieve_dense( @@ -40,7 +43,9 @@ def retrieve_dense( ) -> list[Hit]: k = k or config.TOP_K_DENSE vec = embeddings.encode([query], model_name = model_name, normalize = True)[0] - return [Hit(cid, s, dense_score = s) for cid, s in store.search_dense(conn, scope, vec, k)] + return [ + Hit(cid, s, dense_score = s) for cid, s in store.search_dense(conn, scope, vec, k) + ] def _rrf(rankings: list[list[Hit]], rrf_k: int, top_k: int) -> list[Hit]: @@ -48,13 +53,19 @@ def _rrf(rankings: list[list[Hit]], rrf_k: int, top_k: int) -> list[Hit]: best: dict[str, Hit] = {} for ranking in rankings: for rank, hit in enumerate(ranking): - fused[hit.chunk_id] = fused.get(hit.chunk_id, 0.0) + 1.0 / (rrf_k + rank + 1) + fused[hit.chunk_id] = fused.get(hit.chunk_id, 0.0) + 1.0 / ( + rrf_k + rank + 1 + ) cur = best.get(hit.chunk_id) if cur is None: - best[hit.chunk_id] = Hit(hit.chunk_id, 0.0, hit.lexical_score, hit.dense_score) + best[hit.chunk_id] = Hit( + hit.chunk_id, 0.0, hit.lexical_score, hit.dense_score + ) else: cur.lexical_score = ( - cur.lexical_score if cur.lexical_score is not None else hit.lexical_score + cur.lexical_score + if cur.lexical_score is not None + else hit.lexical_score ) cur.dense_score = ( cur.dense_score if cur.dense_score is not None else hit.dense_score @@ -85,7 +96,9 @@ def retrieve_hybrid( if mode == "dense": return retrieve_dense(conn, scope, query, k, model_name = model_name) lexical = retrieve_lexical(conn, scope, query, config.TOP_K_LEXICAL) - dense = retrieve_dense(conn, scope, query, config.TOP_K_DENSE, model_name = model_name) + dense = retrieve_dense( + conn, scope, query, config.TOP_K_DENSE, model_name = model_name + ) return _rrf([lexical, dense], config.RRF_K, k) diff --git a/studio/backend/core/rag/store.py b/studio/backend/core/rag/store.py index 7d58931e53..0b8a3e8ab1 100644 --- a/studio/backend/core/rag/store.py +++ b/studio/backend/core/rag/store.py @@ -89,7 +89,10 @@ def delete_kb(conn: sqlite3.Connection, kb_id: str) -> None: """Delete a knowledge base and every document (+ chunks) under it.""" scope = kb_scope(kb_id) doc_ids = [ - r["id"] for r in conn.execute("SELECT id FROM documents WHERE scope=?", (scope,)).fetchall() + r["id"] + for r in conn.execute( + "SELECT id FROM documents WHERE scope=?", (scope,) + ).fetchall() ] for doc_id in doc_ids: delete_document(conn, doc_id) @@ -170,7 +173,9 @@ def document_by_hash(conn: sqlite3.Connection, scope: str, sha256: str) -> str | return row["id"] if row else None -def failed_documents_by_hash(conn: sqlite3.Connection, scope: str, sha256: str) -> list[dict]: +def failed_documents_by_hash( + conn: sqlite3.Connection, scope: str, sha256: str +) -> list[dict]: rows = conn.execute( "SELECT id, stored_path FROM documents WHERE scope=? AND sha256=? AND status='failed'", (scope, sha256), diff --git a/studio/backend/core/rag/tool.py b/studio/backend/core/rag/tool.py index ccb1b47e63..90c49a726b 100644 --- a/studio/backend/core/rag/tool.py +++ b/studio/backend/core/rag/tool.py @@ -167,10 +167,14 @@ def search_for_autoinject( mode = mode, ) strong = [ - h for h in hits if h.dense_score is not None and h.dense_score >= min_dense_score + h + for h in hits + if h.dense_score is not None and h.dense_score >= min_dense_score ][:k] if not strong and hits and mode == "lexical": - probe = retrieval.retrieve_dense(conn, scope, query, 1, model_name = model_name) + probe = retrieval.retrieve_dense( + conn, scope, query, 1, model_name = model_name + ) if ( probe and probe[0].dense_score is not None diff --git a/studio/backend/core/tool_healing.py b/studio/backend/core/tool_healing.py index 973520d5cd..121f568aad 100644 --- a/studio/backend/core/tool_healing.py +++ b/studio/backend/core/tool_healing.py @@ -83,7 +83,9 @@ def parse_tool_calls_from_text(content: str) -> list[dict]: }, } if isinstance(tc["function"]["arguments"], dict): - tc["function"]["arguments"] = json.dumps(tc["function"]["arguments"]) + tc["function"]["arguments"] = json.dumps( + tc["function"]["arguments"] + ) tool_calls.append(tc) except (json.JSONDecodeError, ValueError): pass @@ -99,7 +101,11 @@ def parse_tool_calls_from_text(content: str) -> list[dict]: func_name = fm.group(1) body_start = fm.end() # Boundaries: next - next_func = func_starts[idx + 1].start() if idx + 1 < len(func_starts) else len(content) + next_func = ( + func_starts[idx + 1].start() + if idx + 1 < len(func_starts) + else len(content) + ) end_tag = _TC_END_TAG_RE.search(content[body_start:]) if end_tag: body_end = body_start + end_tag.start() diff --git a/studio/backend/core/training/s3_dataset.py b/studio/backend/core/training/s3_dataset.py index 3d05d19c75..f20521ab73 100644 --- a/studio/backend/core/training/s3_dataset.py +++ b/studio/backend/core/training/s3_dataset.py @@ -159,7 +159,9 @@ def prepare_s3_dataset_download( bucket/prefix contains no supported dataset files. """ if not boto3_available(): - raise RuntimeError("S3 dataset loading requires boto3. Install it with: pip install boto3") + raise RuntimeError( + "S3 dataset loading requires boto3. Install it with: pip install boto3" + ) bucket = s3_config.get("bucket") if not bucket: @@ -193,7 +195,9 @@ def prepare_s3_dataset_download( local_path = _unique_local_path(target_dir, filename, used_paths) download_kwargs = {} if cancel_callback is not None: - download_kwargs["Callback"] = lambda _bytes: _raise_if_cancelled(cancel_callback) + download_kwargs["Callback"] = lambda _bytes: _raise_if_cancelled( + cancel_callback + ) client.download_file(bucket, key, local_path, **download_kwargs) _raise_if_cancelled(cancel_callback) local_files.append(local_path) diff --git a/studio/backend/core/training/trainer.py b/studio/backend/core/training/trainer.py index 018b66ea92..d45dfc6fb5 100644 --- a/studio/backend/core/training/trainer.py +++ b/studio/backend/core/training/trainer.py @@ -20,7 +20,9 @@ os.environ["TOKENIZERS_PARALLELISM"] = "false" # UNSLOTH_COMPILE_LOCATION via PYTHONPATH lets any subprocess find them. # Do NOT import unsloth_zoo.compiler here -- it triggers heavy torch/triton imports. if sys.platform in ("win32", "darwin"): - _compile_cache = os.environ.get("UNSLOTH_COMPILE_LOCATION", "unsloth_compiled_cache") + _compile_cache = os.environ.get( + "UNSLOTH_COMPILE_LOCATION", "unsloth_compiled_cache" + ) if not os.path.isabs(_compile_cache): _compile_cache = os.path.abspath(_compile_cache) os.environ["UNSLOTH_COMPILE_LOCATION"] = _compile_cache @@ -133,10 +135,16 @@ class UnslothTrainer: self.is_cpt = False # True for Continued Pretraining self.is_vlm = False self.is_audio = False - self.is_audio_vlm = False # Multimodal model (e.g. Gemma 3N) trained on audio data + self.is_audio_vlm = ( + False # Multimodal model (e.g. Gemma 3N) trained on audio data + ) self._audio_type = None # 'csm', 'whisper', 'snac', 'bicodec', 'dac' - self._cuda_audio_used = False # Set once after audio CUDA preprocessing; never cleared - self._spark_tts_repo_dir = None # Downloaded Spark-TTS repo path (for BiCodecTokenizer) + self._cuda_audio_used = ( + False # Set once after audio CUDA preprocessing; never cleared + ) + self._spark_tts_repo_dir = ( + None # Downloaded Spark-TTS repo path (for BiCodecTokenizer) + ) self.model_name = None # Training metrics tracking @@ -193,7 +201,11 @@ class UnslothTrainer: self._cuda_audio_used = False # --- Detect VLM --- - vision = is_vision_model(model_name, hf_token = hf_token) if not self.is_audio else False + vision = ( + is_vision_model(model_name, hf_token = hf_token) + if not self.is_audio + else False + ) self.is_vlm = not self.is_audio_vlm and vision and is_dataset_image logger.info( @@ -272,7 +284,9 @@ class UnslothTrainer: if total_steps > 0: steps_remaining = total_steps - current_step if steps_remaining > 0: - eta_seconds = (elapsed_seconds / current_step) * steps_remaining + eta_seconds = ( + elapsed_seconds / current_step + ) * steps_remaining num_tokens = getattr(state, "num_input_tokens_seen", None) @@ -300,7 +314,9 @@ class UnslothTrainer: return _ProgressCallback() - def _calculate_total_steps(self, num_samples, batch_size, grad_accum, num_epochs, max_steps): + def _calculate_total_steps( + self, num_samples, batch_size, grad_accum, num_epochs, max_steps + ): """Calculate total training steps from dataset size and training params.""" if max_steps and max_steps > 0: return max_steps @@ -321,7 +337,9 @@ class UnslothTrainer: size, lr, warmup, fp16/bf16, etc.) with per-branch overrides via extra_args. """ batch_size = training_args.get("batch_size", 2) - gradient_accumulation_steps = training_args.get("gradient_accumulation_steps", 4) + gradient_accumulation_steps = training_args.get( + "gradient_accumulation_steps", 4 + ) warmup_steps_val = training_args.get("warmup_steps", 5) max_steps_val = training_args.get("max_steps", 0) learning_rate = training_args.get("learning_rate", 2e-4) @@ -389,7 +407,9 @@ class UnslothTrainer: elif self.should_stop: msg = f"{label} training cancelled" if label else "Training cancelled" logger.info(f"\n{msg}.\n") - self._update_progress(is_training = False, status_message = "Training cancelled.") + self._update_progress( + is_training = False, status_message = "Training cancelled." + ) else: self.trainer.save_model() self.tokenizer.save_pretrained(output_dir) @@ -417,7 +437,9 @@ class UnslothTrainer: ] # Spark-TTS path is relative to the downloaded repo if self._spark_tts_repo_dir: - spark_code_dir = os.path.join(os.path.dirname(self._spark_tts_repo_dir), "Spark-TTS") + spark_code_dir = os.path.join( + os.path.dirname(self._spark_tts_repo_dir), "Spark-TTS" + ) audio_paths.append(spark_code_dir) removed_paths = [] @@ -471,7 +493,11 @@ class UnslothTrainer: # Hardcoded fallback audio_col = next((c for c in cols if c.lower() in ("audio", "speech")), None) text_col = next( - (c for c in cols if c.lower() in ("text", "sentence", "transcript", "transcription")), + ( + c + for c in cols + if c.lower() in ("text", "sentence", "transcript", "transcription") + ), None, ) @@ -501,7 +527,9 @@ class UnslothTrainer: ) -> bool: """Load model for training (supports both text and vision models)""" self.load_in_4bit = load_in_4bit # For training_meta.json - self.trust_remote_code = trust_remote_code # For AutoProcessor etc. used during training + self.trust_remote_code = ( + trust_remote_code # For AutoProcessor etc. used during training + ) try: if self.model is not None: del self.model @@ -535,14 +563,18 @@ class UnslothTrainer: # Remove stale compiled cache so the new model gets a fresh one from utils.cache_cleanup import clear_unsloth_compiled_cache - _preserve = ["Unsloth*Trainer.py"] if sys.platform in ("win32", "darwin") else None + _preserve = ( + ["Unsloth*Trainer.py"] if sys.platform in ("win32", "darwin") else None + ) clear_unsloth_compiled_cache(preserve_patterns = _preserve) # Detect audio model type dynamically (config.json + tokenizer) self._audio_type = detect_audio_type(model_name, hf_token) # audio_vlm is detected as an audio_type now; handle separately if self._audio_type == "audio_vlm": self.is_audio = False - self.is_audio_vlm = is_dataset_audio # Only use audio VLM path if dataset has audio + self.is_audio_vlm = ( + is_dataset_audio # Only use audio VLM path if dataset has audio + ) self._audio_type = None else: self.is_audio = self._audio_type is not None @@ -552,7 +584,11 @@ class UnslothTrainer: self._cuda_audio_used = False # VLM: vision model + image dataset (mutually exclusive with audio) - vision = is_vision_model(model_name, hf_token = hf_token) if not self.is_audio else False + vision = ( + is_vision_model(model_name, hf_token = hf_token) + if not self.is_audio + else False + ) self.is_vlm = not self.is_audio_vlm and vision and is_dataset_image self.model_name = model_name self.max_seq_length = max_seq_length @@ -560,7 +596,9 @@ class UnslothTrainer: logger.info( f"Audio type: {self._audio_type}, is_audio: {self.is_audio}, is_audio_vlm: {self.is_audio_vlm}" ) - logger.info(f"Dataset has images: {is_dataset_image}, audio: {is_dataset_audio}") + logger.info( + f"Dataset has images: {is_dataset_image}, audio: {is_dataset_audio}" + ) logger.info(f"Using VLM path: {self.is_vlm}") # Reset training state for new run @@ -574,8 +612,12 @@ class UnslothTrainer: ) # Update UI with loading message - model_display = model_name.split("/")[-1] if "/" in model_name else model_name - model_type_label = "audio" if self.is_audio else ("vision" if self.is_vlm else "text") + model_display = ( + model_name.split("/")[-1] if "/" in model_name else model_name + ) + model_type_label = ( + "audio" if self.is_audio else ("vision" if self.is_vlm else "text") + ) self._update_progress( status_message = f"Loading {model_type_label} model... {model_display}" ) @@ -628,9 +670,12 @@ class UnslothTrainer: # (incl. FORCE_FLOAT32) is honored -- T4/V100 must NOT be coerced to # float16. Derive ROCm inline since hardware.IS_ROCM may be unset here. _is_rocm = ( - bool(getattr(torch.version, "hip", None)) or "rocm" in torch.__version__.lower() + bool(getattr(torch.version, "hip", None)) + or "rocm" in torch.__version__.lower() + ) + _auto_dtype = ( + torch.float16 if (_is_rocm and not is_bfloat16_supported()) else None ) - _auto_dtype = torch.float16 if (_is_rocm and not is_bfloat16_supported()) else None # Branch based on model type if self._audio_type == "csm": @@ -687,7 +732,9 @@ class UnslothTrainer: token = hf_token, trust_remote_code = trust_remote_code, ) - logger.info(f"Loaded {self._audio_type} audio model (FastLanguageModel)") + logger.info( + f"Loaded {self._audio_type} audio model (FastLanguageModel)" + ) elif self._audio_type == "bicodec": # Spark-TTS: download full repo (sparktts + BiCodec weights), then @@ -708,7 +755,9 @@ class UnslothTrainer: llm_path = f"{local_dir}/LLM" repo_path = snapshot_download(hf_repo, local_dir = local_dir) - self._spark_tts_repo_dir = os.path.abspath(repo_path) # Absolute for sys.path + self._spark_tts_repo_dir = os.path.abspath( + repo_path + ) # Absolute for sys.path llm_path = os.path.join(self._spark_tts_repo_dir, "LLM") self.model, self.tokenizer = FastModel.from_pretrained( @@ -771,15 +820,21 @@ class UnslothTrainer: from transformers import ProcessorMixin tok = self.tokenizer - has_image_proc = isinstance(tok, ProcessorMixin) or hasattr(tok, "image_processor") - logger.info(f"\n[VLM Diagnostic] FastVisionModel returned: {type(tok).__name__}") + has_image_proc = isinstance(tok, ProcessorMixin) or hasattr( + tok, "image_processor" + ) + logger.info( + f"\n[VLM Diagnostic] FastVisionModel returned: {type(tok).__name__}" + ) logger.info( f"[VLM Diagnostic] Is ProcessorMixin: {isinstance(tok, ProcessorMixin)}" ) logger.info( f"[VLM Diagnostic] Has image_processor: {hasattr(tok, 'image_processor')}" ) - logger.info(f"[VLM Diagnostic] Usable as vision processor: {has_image_proc}\n") + logger.info( + f"[VLM Diagnostic] Usable as vision processor: {has_image_proc}\n" + ) else: # Load text model - returns (model, tokenizer) self.model, self.tokenizer = FastLanguageModel.from_pretrained( @@ -899,7 +954,9 @@ class UnslothTrainer: # Full finetuning - skip PEFT entirely if not use_lora: - self._update_progress(status_message = "Full finetuning mode - no LoRA adapters") + self._update_progress( + status_message = "Full finetuning mode - no LoRA adapters" + ) logger.info("Full finetuning mode - training all parameters\n") return True @@ -926,7 +983,10 @@ class UnslothTrainer: # Normalize gradient_checkpointing to True, False, or "unsloth" if isinstance(use_gradient_checkpointing, str): use_gradient_checkpointing = use_gradient_checkpointing.strip().lower() - if use_gradient_checkpointing == "" or use_gradient_checkpointing == "unsloth": + if ( + use_gradient_checkpointing == "" + or use_gradient_checkpointing == "unsloth" + ): use_gradient_checkpointing = "unsloth" elif use_gradient_checkpointing in ("true", "1", "yes"): use_gradient_checkpointing = True @@ -954,14 +1014,14 @@ class UnslothTrainer: # Check expected attributes if not hasattr(self.model, "config"): - error_msg = ( - "Model does not have config attribute - model may not be loaded correctly" - ) + error_msg = "Model does not have config attribute - model may not be loaded correctly" logger.error(error_msg) self._update_progress(error = error_msg) return False - logger.info(f"Configuring LoRA adapters (r={lora_r}, alpha={lora_alpha})...\n") + logger.info( + f"Configuring LoRA adapters (r={lora_r}, alpha={lora_alpha})...\n" + ) logger.info( f"Gradient checkpointing: {use_gradient_checkpointing} (type: {type(use_gradient_checkpointing).__name__})\n" ) @@ -976,8 +1036,12 @@ class UnslothTrainer: logger.info(f" - Target modules: {target_modules}") if self.is_audio_vlm: logger.info(f" - Finetune vision layers: {finetune_vision_layers}") - logger.info(f" - Finetune language layers: {finetune_language_layers}") - logger.info(f" - Finetune attention modules: {finetune_attention_modules}") + logger.info( + f" - Finetune language layers: {finetune_language_layers}" + ) + logger.info( + f" - Finetune attention modules: {finetune_attention_modules}" + ) logger.info(f" - Finetune MLP modules: {finetune_mlp_modules}") logger.info() @@ -990,7 +1054,9 @@ class UnslothTrainer: use_gradient_checkpointing = use_gradient_checkpointing, random_state = 3407, use_rslora = use_rslora, - loftq_config = {"loftq_bits": 4, "loftq_iter": 1} if use_loftq else None, + loftq_config = {"loftq_bits": 4, "loftq_iter": 1} + if use_loftq + else None, ) # Audio VLM models support VLM-style layer selection if self.is_audio_vlm: @@ -1020,7 +1086,9 @@ class UnslothTrainer: use_gradient_checkpointing = use_gradient_checkpointing, random_state = 3407, use_rslora = use_rslora, - loftq_config = {"loftq_bits": 4, "loftq_iter": 1} if use_loftq else None, + loftq_config = {"loftq_bits": 4, "loftq_iter": 1} + if use_loftq + else None, task_type = None, ) @@ -1039,7 +1107,9 @@ class UnslothTrainer: use_gradient_checkpointing = use_gradient_checkpointing, random_state = 3407, use_rslora = use_rslora, - loftq_config = {"loftq_bits": 4, "loftq_iter": 1} if use_loftq else None, + loftq_config = {"loftq_bits": 4, "loftq_iter": 1} + if use_loftq + else None, ) elif self.is_vlm: @@ -1047,7 +1117,9 @@ class UnslothTrainer: logger.info(f"Vision model LoRA configuration:") logger.info(f" - Finetune vision layers: {finetune_vision_layers}") logger.info(f" - Finetune language layers: {finetune_language_layers}") - logger.info(f" - Finetune attention modules: {finetune_attention_modules}") + logger.info( + f" - Finetune attention modules: {finetune_attention_modules}" + ) logger.info(f" - Finetune MLP modules: {finetune_mlp_modules}\n") self.model = FastVisionModel.get_peft_model( @@ -1064,7 +1136,9 @@ class UnslothTrainer: use_gradient_checkpointing = use_gradient_checkpointing, random_state = 3407, use_rslora = use_rslora, - loftq_config = {"loftq_bits": 4, "loftq_iter": 1} if use_loftq else None, + loftq_config = {"loftq_bits": 4, "loftq_iter": 1} + if use_loftq + else None, modules_to_save = modules_to_save, ) else: @@ -1084,7 +1158,9 @@ class UnslothTrainer: use_gradient_checkpointing = use_gradient_checkpointing, random_state = 3407, use_rslora = use_rslora, - loftq_config = {"loftq_bits": 4, "loftq_iter": 1} if use_loftq else None, + loftq_config = {"loftq_bits": 4, "loftq_iter": 1} + if use_loftq + else None, modules_to_save = modules_to_save, ) @@ -1102,7 +1178,9 @@ class UnslothTrainer: import sys error_details = ( - f"{type(e).__name__}: {str(e)}" if str(e) else f"{type(e).__name__} (no message)" + f"{type(e).__name__}: {str(e)}" + if str(e) + else f"{type(e).__name__} (no message)" ) full_traceback = traceback.format_exc() logger.error(f"Error preparing model: {error_details}") @@ -1168,7 +1246,9 @@ class UnslothTrainer: kwargs.pop("task_ids", None) # Only keep recognized TransformersKwargs - clean_kwargs = {k: v for k, v in kwargs.items() if k in _TRANSFORMERS_KWARGS} + clean_kwargs = { + k: v for k, v in kwargs.items() if k in _TRANSFORMERS_KWARGS + } if input_ids is not None and input_ids.ndim == 2: merged = self._merge_input_ids_with_input_values( @@ -1193,7 +1273,9 @@ class UnslothTrainer: backbone_hidden_states = backbone_outputs[0] slice_indices = ( - slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep + slice(-logits_to_keep, None) + if isinstance(logits_to_keep, int) + else logits_to_keep ) backbone_logits = self.lm_head(backbone_hidden_states[:, slice_indices, :]) @@ -1211,7 +1293,9 @@ class UnslothTrainer: ) train_mask = ~(labels[:, :, 1:] == -100).all(dim = -1) - depth_decoder_input_ids = labels[train_mask][..., : self.config.num_codebooks - 1] + depth_decoder_input_ids = labels[train_mask][ + ..., : self.config.num_codebooks - 1 + ] depth_decoder_input_ids = nn.functional.pad( depth_decoder_input_ids, (1, 0), value = 0 ) @@ -1225,9 +1309,9 @@ class UnslothTrainer: # Scale num_items_in_batch for the depth decoder's 31 codebooks. dd_kwargs = clean_kwargs.copy() if "num_items_in_batch" in dd_kwargs: - dd_kwargs["num_items_in_batch"] = dd_kwargs["num_items_in_batch"] * ( - self.config.num_codebooks - 1 - ) + dd_kwargs["num_items_in_batch"] = dd_kwargs[ + "num_items_in_batch" + ] * (self.config.num_codebooks - 1) depth_decoder_outputs = self.depth_decoder( input_ids = depth_decoder_input_ids, @@ -1264,10 +1348,14 @@ class UnslothTrainer: depth_decoder_outputs.logits if depth_decoder_outputs else None ), depth_decoder_past_key_values = ( - depth_decoder_outputs.past_key_values if depth_decoder_outputs else None + depth_decoder_outputs.past_key_values + if depth_decoder_outputs + else None ), depth_decoder_hidden_states = ( - depth_decoder_outputs.hidden_states if depth_decoder_outputs else None + depth_decoder_outputs.hidden_states + if depth_decoder_outputs + else None ), depth_decoder_attentions = ( depth_decoder_outputs.attentions if depth_decoder_outputs else None @@ -1307,11 +1395,17 @@ class UnslothTrainer: speaker_key = resolved["speaker_col"] if audio_col is None: - raise ValueError(f"No audio column found in dataset. Columns: {dataset.column_names}") + raise ValueError( + f"No audio column found in dataset. Columns: {dataset.column_names}" + ) if text_col is None: - raise ValueError(f"No text column found in dataset. Columns: {dataset.column_names}") + raise ValueError( + f"No text column found in dataset. Columns: {dataset.column_names}" + ) if speaker_key is None: - logger.info("No speaker found, adding default 'source' of 0 for all examples\n") + logger.info( + "No speaker found, adding default 'source' of 0 for all examples\n" + ) dataset = dataset.add_column("source", ["0"] * len(dataset)) speaker_key = "source" @@ -1391,11 +1485,14 @@ class UnslothTrainer: ) if not processed_examples: - raise ValueError(f"No valid examples after CSM preprocessing (skipped {skipped})") + raise ValueError( + f"No valid examples after CSM preprocessing (skipped {skipped})" + ) result_dataset = Dataset.from_list(processed_examples) logger.info( - f"CSM preprocessing complete: {len(result_dataset)} examples " f"({skipped} skipped)\n" + f"CSM preprocessing complete: {len(result_dataset)} examples " + f"({skipped} skipped)\n" ) return result_dataset @@ -1557,7 +1654,9 @@ class UnslothTrainer: # --- Encode audio with SNAC (notebook 122-142) --- waveform = ( - torch.from_numpy(audio_data["array"]).unsqueeze(0).to(dtype = torch.float32) + torch.from_numpy(audio_data["array"]) + .unsqueeze(0) + .to(dtype = torch.float32) ) if resample_transform is not None: waveform = resample_transform(waveform) @@ -1571,11 +1670,21 @@ class UnslothTrainer: for i in range(codes[0].shape[1]): all_codes.append(codes[0][0][i].item() + AUDIO_OFFSET) all_codes.append(codes[1][0][2 * i].item() + AUDIO_OFFSET + 4096) - all_codes.append(codes[2][0][4 * i].item() + AUDIO_OFFSET + (2 * 4096)) - all_codes.append(codes[2][0][(4 * i) + 1].item() + AUDIO_OFFSET + (3 * 4096)) - all_codes.append(codes[1][0][(2 * i) + 1].item() + AUDIO_OFFSET + (4 * 4096)) - all_codes.append(codes[2][0][(4 * i) + 2].item() + AUDIO_OFFSET + (5 * 4096)) - all_codes.append(codes[2][0][(4 * i) + 3].item() + AUDIO_OFFSET + (6 * 4096)) + all_codes.append( + codes[2][0][4 * i].item() + AUDIO_OFFSET + (2 * 4096) + ) + all_codes.append( + codes[2][0][(4 * i) + 1].item() + AUDIO_OFFSET + (3 * 4096) + ) + all_codes.append( + codes[1][0][(2 * i) + 1].item() + AUDIO_OFFSET + (4 * 4096) + ) + all_codes.append( + codes[2][0][(4 * i) + 2].item() + AUDIO_OFFSET + (5 * 4096) + ) + all_codes.append( + codes[2][0][(4 * i) + 3].item() + AUDIO_OFFSET + (6 * 4096) + ) if len(all_codes) == 0: skipped += 1 @@ -1631,7 +1740,9 @@ class UnslothTrainer: # Progress update every 100 examples if (idx + 1) % 100 == 0: - self._update_progress(status_message = f"Encoding audio... {idx + 1}/{len(dataset)}") + self._update_progress( + status_message = f"Encoding audio... {idx + 1}/{len(dataset)}" + ) # Free SNAC model from GPU logger.info("Freeing SNAC codec model from GPU...\n") @@ -1643,11 +1754,14 @@ class UnslothTrainer: self._cuda_audio_used = True if not processed_examples: - raise ValueError(f"No valid examples after SNAC preprocessing (skipped {skipped})") + raise ValueError( + f"No valid examples after SNAC preprocessing (skipped {skipped})" + ) result_dataset = Dataset.from_list(processed_examples) logger.info( - f"SNAC preprocessing complete: {len(result_dataset)} examples " f"({skipped} skipped)\n" + f"SNAC preprocessing complete: {len(result_dataset)} examples " + f"({skipped} skipped)\n" ) return result_dataset @@ -1670,7 +1784,9 @@ class UnslothTrainer: # sparktts lives in the SparkAudio/Spark-TTS GitHub repo, not the HF model # repo. Clone if needed. - spark_code_dir = os.path.join(os.path.dirname(self._spark_tts_repo_dir), "Spark-TTS") + spark_code_dir = os.path.join( + os.path.dirname(self._spark_tts_repo_dir), "Spark-TTS" + ) sparktts_pkg = os.path.join(spark_code_dir, "sparktts") if not os.path.isdir(sparktts_pkg): self._update_progress(status_message = "Cloning Spark-TTS code repo...") @@ -1737,7 +1853,9 @@ class UnslothTrainer: return_tensors = "pt", padding = True, ) - input_values = processed.input_values.to(audio_tokenizer.feature_extractor.device) + input_values = processed.input_values.to( + audio_tokenizer.feature_extractor.device + ) model_output = audio_tokenizer.feature_extractor(input_values) if model_output.hidden_states is None: @@ -1786,8 +1904,12 @@ class UnslothTrainer: ref_wav_np = audio_tokenizer.get_ref_clip(audio_array) # Prepare tensors - audio_tensor = torch.from_numpy(audio_array).unsqueeze(0).float().to(device) - ref_wav_tensor = torch.from_numpy(ref_wav_np).unsqueeze(0).float().to(device) + audio_tensor = ( + torch.from_numpy(audio_array).unsqueeze(0).float().to(device) + ) + ref_wav_tensor = ( + torch.from_numpy(ref_wav_np).unsqueeze(0).float().to(device) + ) # Extract wav2vec2 features feat = extract_wav2vec2_features(audio_tensor) @@ -1799,10 +1921,15 @@ class UnslothTrainer: } # BiCodec tokenize - semantic_token_ids, global_token_ids = audio_tokenizer.model.tokenize(batch) + semantic_token_ids, global_token_ids = audio_tokenizer.model.tokenize( + batch + ) global_tokens = "".join( - [f"<|bicodec_global_{i}|>" for i in global_token_ids.squeeze().cpu().numpy()] + [ + f"<|bicodec_global_{i}|>" + for i in global_token_ids.squeeze().cpu().numpy() + ] ) semantic_tokens = "".join( [ @@ -1858,7 +1985,9 @@ class UnslothTrainer: self._cuda_audio_used = True if not processed_examples: - raise ValueError(f"No valid examples after BiCodec preprocessing (skipped {skipped})") + raise ValueError( + f"No valid examples after BiCodec preprocessing (skipped {skipped})" + ) result_dataset = Dataset.from_list(processed_examples) logger.info( @@ -1946,7 +2075,9 @@ class UnslothTrainer: logger.info("Cast audio column to 24kHz\n") # Load Whisper for word timings - self._update_progress(status_message = "Loading Whisper model for word timings...") + self._update_progress( + status_message = "Loading Whisper model for word timings..." + ) logger.info("Loading Whisper model for word timings...\n") import whisper @@ -1965,7 +2096,9 @@ class UnslothTrainer: prompt_processor = PromptProcessor(model_tokenizer_path) self._update_progress(status_message = "Preprocessing audio with OuteTTS...") - logger.info(f"DAC preprocessing: audio_col='{audio_col}', text_col='{text_col}'\n") + logger.info( + f"DAC preprocessing: audio_col='{audio_col}', text_col='{text_col}'\n" + ) processed_examples = [] skipped = 0 @@ -2005,7 +2138,9 @@ class UnslothTrainer: tmp.flush() tmp_path = tmp.name try: - whisper_result = whisper_model.transcribe(tmp_path, word_timestamps = True) + whisper_result = whisper_model.transcribe( + tmp_path, word_timestamps = True + ) finally: Path(tmp_path).unlink(missing_ok = True) @@ -2066,11 +2201,14 @@ class UnslothTrainer: self._cuda_audio_used = True if not processed_examples: - raise ValueError(f"No valid examples after DAC preprocessing (skipped {skipped})") + raise ValueError( + f"No valid examples after DAC preprocessing (skipped {skipped})" + ) result_dataset = HFDataset.from_list(processed_examples) logger.info( - f"DAC preprocessing complete: {len(result_dataset)} examples " f"({skipped} skipped)\n" + f"DAC preprocessing complete: {len(result_dataset)} examples " + f"({skipped} skipped)\n" ) sample = result_dataset[0]["text"] logger.info(f"Sample text (first 200 chars): {sample[:200]}...\n") @@ -2101,7 +2239,9 @@ class UnslothTrainer: ) # Cast audio to 16kHz (Whisper's expected sample rate) - dataset = dataset.cast_column(audio_col, Audio(sampling_rate = WHISPER_SAMPLE_RATE)) + dataset = dataset.cast_column( + audio_col, Audio(sampling_rate = WHISPER_SAMPLE_RATE) + ) # Train/eval split (notebook does dataset.train_test_split) eval_dataset_raw = None @@ -2128,7 +2268,11 @@ class UnslothTrainer: try: audio_data = example.get(audio_col) text = example.get(text_col) - if audio_data is None or audio_data.get("array") is None or not text: + if ( + audio_data is None + or audio_data.get("array") is None + or not text + ): skipped += 1 continue @@ -2146,7 +2290,9 @@ class UnslothTrainer: } ) except Exception as e: - logger.warning(f"Error processing Whisper {split_name} example {idx}: {e}") + logger.warning( + f"Error processing Whisper {split_name} example {idx}: {e}" + ) skipped += 1 continue @@ -2161,7 +2307,9 @@ class UnslothTrainer: return processed train_data = process_split(dataset, "train") - eval_data = process_split(eval_dataset_raw, "eval") if eval_dataset_raw else None + eval_data = ( + process_split(eval_dataset_raw, "eval") if eval_dataset_raw else None + ) if not train_data: raise ValueError("No valid examples after Whisper preprocessing") @@ -2196,7 +2344,9 @@ class UnslothTrainer: if candidates: all_files.extend(str(c) for c in candidates) continue - raise ValueError(f"No supported data files in directory: {file_path_obj}") + raise ValueError( + f"No supported data files in directory: {file_path_obj}" + ) else: all_files.append(str(file_path_obj)) return all_files @@ -2244,7 +2394,9 @@ class UnslothTrainer: try: dataset = None eval_dataset = None - has_separate_eval_source = False # True if eval comes from a separate HF split + has_separate_eval_source = ( + False # True if eval comes from a separate HF split + ) eval_enabled = eval_steps is not None and eval_steps > 0 raw_text_mode = is_cpt or format_type == "raw" @@ -2402,7 +2554,9 @@ class UnslothTrainer: if eval_dataset is not None: has_separate_eval_source = True else: - logger.info("Eval disabled (eval_steps <= 0), skipping eval split detection\n") + logger.info( + "Eval disabled (eval_steps <= 0), skipping eval split detection\n" + ) if dataset is None: raise ValueError("No dataset provided") @@ -2411,7 +2565,11 @@ class UnslothTrainer: if dataset_slice_start is not None or dataset_slice_end is not None: total_rows = len(dataset) start = dataset_slice_start if dataset_slice_start is not None else 0 - end = dataset_slice_end if dataset_slice_end is not None else total_rows - 1 + end = ( + dataset_slice_end + if dataset_slice_end is not None + else total_rows - 1 + ) # Clamp to valid range start = max(0, min(start, total_rows - 1)) end = max(start, min(end, total_rows - 1)) @@ -2442,11 +2600,15 @@ class UnslothTrainer: return (train_data, eval_data) elif self._audio_type == "snac": - processed = self._preprocess_snac_dataset(dataset, custom_format_mapping) + processed = self._preprocess_snac_dataset( + dataset, custom_format_mapping + ) return (processed, None) elif self._audio_type == "bicodec": - processed = self._preprocess_bicodec_dataset(dataset, custom_format_mapping) + processed = self._preprocess_bicodec_dataset( + dataset, custom_format_mapping + ) return ({"dataset": processed, "final_format": "audio_bicodec"}, None) elif self._audio_type == "dac": @@ -2496,7 +2658,9 @@ class UnslothTrainer: return (dataset_info, eval_dataset) elif self.is_audio_vlm: - formatted = self._format_audio_vlm_dataset(dataset, custom_format_mapping) + formatted = self._format_audio_vlm_dataset( + dataset, custom_format_mapping + ) return (formatted, None) # ========== FORMAT FIRST ========== @@ -2532,7 +2696,9 @@ class UnslothTrainer: self._update_progress( status_message = f"Dataset ready ({final_n:,} samples, {detected} format)" ) - logger.info(f"Dataset formatted successfully ({final_n} samples, {detected})\n") + logger.info( + f"Dataset formatted successfully ({final_n} samples, {detected})\n" + ) # ========== THEN SPLIT ========== if has_separate_eval_source and eval_dataset is not None: @@ -2735,7 +2901,9 @@ class UnslothTrainer: # Store training parameters for metrics calculation self.batch_size = training_args.get("batch_size", 2) self.max_seq_length = training_args.get("max_seq_length", 2048) - self.gradient_accumulation_steps = training_args.get("gradient_accumulation_steps", 4) + self.gradient_accumulation_steps = training_args.get( + "gradient_accumulation_steps", 4 + ) # Set training start time self.training_start_time = time.time() @@ -2743,10 +2911,14 @@ class UnslothTrainer: self._update_progress(is_training = True, error = None) # Setup logging - if training_args.get("enable_wandb", False) and training_args.get("wandb_token"): + if training_args.get("enable_wandb", False) and training_args.get( + "wandb_token" + ): os.environ["WANDB_API_KEY"] = training_args["wandb_token"] import wandb - wandb.init(project = training_args.get("wandb_project", "unsloth-training")) + wandb.init( + project = training_args.get("wandb_project", "unsloth-training") + ) # Create output directory output_dir = str(resolve_output_dir(training_args.get("output_dir"))) @@ -2782,7 +2954,9 @@ class UnslothTrainer: training_args.get("num_epochs", 3), training_args.get("max_steps", 0), ) - self._update_progress(total_steps = total, status_message = "Starting CSM training...") + self._update_progress( + total_steps = total, status_message = "Starting CSM training..." + ) logger.info(f"CSM training config: {config}\n") self.trainer.train( resume_from_checkpoint = training_args.get("resume_from_checkpoint") @@ -2821,7 +2995,9 @@ class UnslothTrainer: training_args.get("num_epochs", 3), training_args.get("max_steps", 0), ) - self._update_progress(total_steps = total, status_message = "Starting SNAC training...") + self._update_progress( + total_steps = total, status_message = "Starting SNAC training..." + ) logger.info(f"SNAC training config: {config}\n") self.trainer.train( resume_from_checkpoint = training_args.get("resume_from_checkpoint") @@ -2847,7 +3023,9 @@ class UnslothTrainer: trainer_kwargs = { "model": self.model, "train_dataset": dataset, - "data_collator": DataCollatorSpeechSeq2SeqWithPadding(processor = self.tokenizer), + "data_collator": DataCollatorSpeechSeq2SeqWithPadding( + processor = self.tokenizer + ), "processing_class": self.tokenizer.feature_extractor, "args": Seq2SeqTrainingArguments(**config), } @@ -2886,12 +3064,16 @@ class UnslothTrainer: # ========== DATA COLLATOR SELECTION ========== model_name_lower = self.model_name.lower() - is_deepseek_ocr = "deepseek" in model_name_lower and "ocr" in model_name_lower + is_deepseek_ocr = ( + "deepseek" in model_name_lower and "ocr" in model_name_lower + ) logger.info("Configuring data collator...\n") dataset_final_format = ( - str(dataset.get("final_format", "")).lower() if isinstance(dataset, dict) else "" + str(dataset.get("final_format", "")).lower() + if isinstance(dataset, dict) + else "" ) raw_text_mode = dataset_final_format == "raw_text" @@ -2929,7 +3111,9 @@ class UnslothTrainer: image_size = 640, base_size = 1024, crop_mode = True, - train_on_responses_only = training_args.get("train_on_completions", False), + train_on_responses_only = training_args.get( + "train_on_completions", False + ), ) logger.info("DeepSeek OCR data collator configured successfully\n") @@ -2959,7 +3143,9 @@ class UnslothTrainer: texts.append(text) audios.append(example[audio_col_name]["array"]) - batch = processor(text = texts, audio = audios, return_tensors = "pt", padding = True) + batch = processor( + text = texts, audio = audios, return_tensors = "pt", padding = True + ) # Labels = input_ids with special tokens masked labels = batch["input_ids"].clone() @@ -2987,9 +3173,13 @@ class UnslothTrainer: FastVisionModel.for_training(self.model) vision_image_size = training_args.get("vision_image_size") if vision_image_size is None: - data_collator = UnslothVisionDataCollator(self.model, self.tokenizer) + data_collator = UnslothVisionDataCollator( + self.model, self.tokenizer + ) else: - logger.info(f"Vision image resize: {vision_image_size} (max dimension)\n") + logger.info( + f"Vision image resize: {vision_image_size} (max dimension)\n" + ) data_collator = UnslothVisionDataCollator( self.model, self.tokenizer, @@ -3009,8 +3199,12 @@ class UnslothTrainer: config_args = { "per_device_train_batch_size": training_args.get("batch_size", 2), - "gradient_accumulation_steps": training_args.get("gradient_accumulation_steps", 4), - "num_train_epochs": training_args.get("num_epochs", 3), # Default to epochs + "gradient_accumulation_steps": training_args.get( + "gradient_accumulation_steps", 4 + ), + "num_train_epochs": training_args.get( + "num_epochs", 3 + ), # Default to epochs "learning_rate": lr_value, "fp16": not is_bfloat16_supported(), "bf16": is_bfloat16_supported(), @@ -3102,7 +3296,9 @@ class UnslothTrainer: logger.info(f"Configuring {label} model training parameters\n") # Provided values or vision defaults optim_value = training_args.get("optim", "adamw_torch_fused") - lr_scheduler_type_value = training_args.get("lr_scheduler_type", "cosine") + lr_scheduler_type_value = training_args.get( + "lr_scheduler_type", "cosine" + ) config_args.update( { "optim": optim_value, @@ -3157,7 +3353,9 @@ class UnslothTrainer: # Audio VLM (e.g. Gemma 3N): raw Dataset from # _format_audio_vlm_dataset, processing_class=processor.tokenizer. # Raw-text runs go to the text path below. - train_dataset = dataset if isinstance(dataset, Dataset) else dataset["dataset"] + train_dataset = ( + dataset if isinstance(dataset, Dataset) else dataset["dataset"] + ) processing_class = ( self.tokenizer.tokenizer if hasattr(self.tokenizer, "tokenizer") @@ -3176,7 +3374,9 @@ class UnslothTrainer: elif self.is_vlm and not raw_text_mode: # Image VLM: dataset is a dict wrapper from # format_and_template_dataset. Raw-text runs go to the text path below. - train_dataset = dataset["dataset"] if isinstance(dataset, dict) else dataset + train_dataset = ( + dataset["dataset"] if isinstance(dataset, dict) else dataset + ) trainer_kwargs = { "model": self.model, "train_dataset": train_dataset, @@ -3263,7 +3463,9 @@ class UnslothTrainer: ) if is_cpt: - logger.info("CPT mode: skipping train_on_responses_only — training on all tokens\n") + logger.info( + "CPT mode: skipping train_on_responses_only — training on all tokens\n" + ) elif raw_text_mode: logger.info( "Raw-text mode: skipping train_on_responses_only — training on all tokens\n" @@ -3288,12 +3490,16 @@ class UnslothTrainer: logger.info(f"Detected template: {template_name}\n") if template_name in TEMPLATE_TO_RESPONSES_MAPPER: - instruction_part = TEMPLATE_TO_RESPONSES_MAPPER[template_name][ - "instruction" + instruction_part = TEMPLATE_TO_RESPONSES_MAPPER[ + template_name + ]["instruction"] + response_part = TEMPLATE_TO_RESPONSES_MAPPER[template_name][ + "response" ] - response_part = TEMPLATE_TO_RESPONSES_MAPPER[template_name]["response"] - logger.info(f"Instruction marker: {instruction_part[:50]}...\n") + logger.info( + f"Instruction marker: {instruction_part[:50]}...\n" + ) logger.info(f"Response marker: {response_part[:50]}...\n") else: logger.info( @@ -3301,7 +3507,9 @@ class UnslothTrainer: ) train_on_responses_enabled = False else: - logger.info(f"No template mapping found for model: {self.model_name}\n") + logger.info( + f"No template mapping found for model: {self.model_name}\n" + ) train_on_responses_enabled = False except Exception as e: @@ -3335,7 +3543,11 @@ class UnslothTrainer: filtered_len = len(self.trainer.train_dataset) original_len = len(dataset["dataset"]) dropped = original_len - filtered_len - drop_pct = round(100 * dropped / original_len, 1) if original_len > 0 else 0 + drop_pct = ( + round(100 * dropped / original_len, 1) + if original_len > 0 + else 0 + ) if filtered_len == 0 or drop_pct > 30: max_seq = training_args.get("max_seq_length", 2048) @@ -3394,7 +3606,10 @@ class UnslothTrainer: self.trainer.add_callback(self._create_progress_callback()) num_samples = None - if hasattr(self.trainer, "train_dataset") and self.trainer.train_dataset is not None: + if ( + hasattr(self.trainer, "train_dataset") + and self.trainer.train_dataset is not None + ): try: num_samples = len(self.trainer.train_dataset) except TypeError: @@ -3404,7 +3619,9 @@ class UnslothTrainer: ) if num_samples is None: - num_samples = len(dataset["dataset"] if isinstance(dataset, dict) else dataset) + num_samples = len( + dataset["dataset"] if isinstance(dataset, dict) else dataset + ) batch_size = training_args.get("batch_size", 2) total_steps = self._calculate_total_steps( @@ -3415,9 +3632,13 @@ class UnslothTrainer: training_args.get("max_steps", 0), ) # ========== START TRAINING ========== - self._update_progress(total_steps = total_steps, status_message = "Starting training...") + self._update_progress( + total_steps = total_steps, status_message = "Starting training..." + ) logger.info("Starting training...\n") - self.trainer.train(resume_from_checkpoint = training_args.get("resume_from_checkpoint")) + self.trainer.train( + resume_from_checkpoint = training_args.get("resume_from_checkpoint") + ) # ========== SAVE MODEL ========== self._finalize_training(output_dir) @@ -3456,7 +3677,9 @@ class UnslothTrainer: method = "lora" config["unsloth_training_method"] = method - logger.info(f"Patching adapter_config.json with unsloth_training_method='{method}'") + logger.info( + f"Patching adapter_config.json with unsloth_training_method='{method}'" + ) with open(config_path, "w") as f: json.dump(config, f, indent = 2) @@ -3470,7 +3693,9 @@ class UnslothTrainer: self.should_stop = True self.save_on_stop = save stop_msg = ( - "Stopping training and saving checkpoint..." if save else "Cancelling training..." + "Stopping training and saving checkpoint..." + if save + else "Cancelling training..." ) self._update_progress(status_message = stop_msg) @@ -3513,7 +3738,9 @@ def _ensure_deepseek_ocr_installed(): pass try: - logger.info("DeepSeek OCR module not found. Auto-installing from HuggingFace...") + logger.info( + "DeepSeek OCR module not found. Auto-installing from HuggingFace..." + ) logger.info("\n Downloading DeepSeek OCR module from HuggingFace...\n") from huggingface_hub import snapshot_download @@ -3526,7 +3753,9 @@ def _ensure_deepseek_ocr_installed(): # Download to project root as 'deepseek_ocr' folder local_dir = os.path.join(parent_dir, "deepseek_ocr") - snapshot_download("unsloth/DeepSeek-OCR", local_dir = local_dir, local_dir_use_symlinks = False) + snapshot_download( + "unsloth/DeepSeek-OCR", local_dir = local_dir, local_dir_use_symlinks = False + ) if parent_dir not in sys.path: sys.path.insert(0, parent_dir) diff --git a/studio/backend/core/training/training.py b/studio/backend/core/training/training.py index 6dd42976c7..2dac25f200 100644 --- a/studio/backend/core/training/training.py +++ b/studio/backend/core/training/training.py @@ -42,7 +42,9 @@ _HF_TMP_CHECKPOINT_RE = re.compile(r"^tmp-checkpoint-\d+$") def _sanitize_db_config(config: dict[str, Any]) -> dict[str, Any]: db_config = { - k: v for k, v in config.items() if k not in {"hf_token", "wandb_token", "s3_config"} + k: v + for k, v in config.items() + if k not in {"hf_token", "wandb_token", "s3_config"} } s3_config = config.get("s3_config") if hasattr(s3_config, "model_dump"): @@ -202,7 +204,9 @@ class TrainingBackend: if self._pump_thread is not None and self._pump_thread.is_alive(): self._pump_thread.join(timeout = 5.0) if self._pump_thread.is_alive(): - logger.warning("Previous pump thread did not exit within 5s — refusing to start") + logger.warning( + "Previous pump thread did not exit within 5s — refusing to start" + ) return False self._pump_thread = None @@ -254,7 +258,9 @@ class TrainingBackend: "train_on_completions": kwargs.get("train_on_completions", False), "finetune_vision_layers": kwargs.get("finetune_vision_layers", True), "finetune_language_layers": kwargs.get("finetune_language_layers", True), - "finetune_attention_modules": kwargs.get("finetune_attention_modules", True), + "finetune_attention_modules": kwargs.get( + "finetune_attention_modules", True + ), "finetune_mlp_modules": kwargs.get("finetune_mlp_modules", True), "enable_wandb": kwargs.get("enable_wandb", False), "wandb_token": kwargs.get("wandb_token"), @@ -367,7 +373,9 @@ class TrainingBackend: pass # Update progress immediately for responsive UI. self._progress.status_message = ( - "Stopping training and saving checkpoint..." if save else "Cancelling training..." + "Stopping training and saving checkpoint..." + if save + else "Cancelling training..." ) return True @@ -375,7 +383,9 @@ class TrainingBackend: """Force-kill the training subprocess so state can be reset immediately.""" with self._lock: if self._proc is not None and self._proc.is_alive(): - logger.info("Force-terminating training subprocess (pid=%s)", self._proc.pid) + logger.info( + "Force-terminating training subprocess (pid=%s)", self._proc.pid + ) self._proc.terminate() proc = self._proc cancelled = self._cancel_requested @@ -529,7 +539,8 @@ class TrainingBackend: else: self._progress.is_training = False self._progress.error = ( - self._progress.error or "Training process exited unexpectedly" + self._progress.error + or "Training process exited unexpectedly" ) self._ensure_db_run_created() @@ -563,7 +574,9 @@ class TrainingBackend: except (TypeError, ValueError): logger.debug("Could not convert loss to float: %s", _raw_loss) _safe_loss = None - _loss_is_nonfinite = _safe_loss is not None and not math.isfinite(_safe_loss) + _loss_is_nonfinite = _safe_loss is not None and not math.isfinite( + _safe_loss + ) if _loss_is_nonfinite: # Drop the value rather than laundering it back to the last # finite loss; clients see loss=None at this step so the NaN @@ -579,7 +592,9 @@ class TrainingBackend: try: _safe_lr = float(_raw_lr) if _raw_lr is not None else None except (TypeError, ValueError): - logger.debug("Could not convert learning_rate to float: %s", _raw_lr) + logger.debug( + "Could not convert learning_rate to float: %s", _raw_lr + ) _safe_lr = None if _safe_lr is not None and not math.isfinite(_safe_lr): _safe_lr = None @@ -591,7 +606,9 @@ class TrainingBackend: self._progress.loss = None if _safe_lr is not None: self._progress.learning_rate = _safe_lr - self._progress.total_steps = event.get("total_steps", self._progress.total_steps) + self._progress.total_steps = event.get( + "total_steps", self._progress.total_steps + ) self._progress.elapsed_seconds = event.get("elapsed_seconds") self._progress.eta_seconds = event.get("eta_seconds") self._progress.grad_norm = event.get("grad_norm") @@ -635,7 +652,9 @@ class TrainingBackend: try: eval_loss = float(eval_loss) except (TypeError, ValueError): - logger.debug("Could not convert eval_loss to float: %s", eval_loss) + logger.debug( + "Could not convert eval_loss to float: %s", eval_loss + ) eval_loss = None if step > 0 and eval_loss is not None and math.isfinite(eval_loss): self.eval_loss_history.append(eval_loss) @@ -665,9 +684,12 @@ class TrainingBackend: "job_id": self.current_job_id, "model_name": self._db_config["model_name"], "dataset_name": self._db_config.get("hf_dataset") - or next(iter(self._db_config.get("local_datasets") or []), "unknown"), + or next( + iter(self._db_config.get("local_datasets") or []), "unknown" + ), "config_json": _json.dumps(self._db_config), - "started_at": self._db_started_at or datetime.now(timezone.utc).isoformat(), + "started_at": self._db_started_at + or datetime.now(timezone.utc).isoformat(), "total_steps": event.get("total_steps"), } elif ( @@ -745,7 +767,9 @@ class TrainingBackend: elif db_action == "update_total_steps": try: from storage.studio_db import update_run_total_steps - update_run_total_steps(db_action_kwargs["job_id"], db_action_kwargs["total_steps"]) + update_run_total_steps( + db_action_kwargs["job_id"], db_action_kwargs["total_steps"] + ) self._db_total_steps_set = True except Exception: logger.warning("Failed to update total_steps in DB", exc_info = True) @@ -772,12 +796,15 @@ class TrainingBackend: model_name = self._db_config["model_name"], dataset_name = dataset_name, config_json = _json.dumps(self._db_config), - started_at = self._db_started_at or datetime.now(timezone.utc).isoformat(), + started_at = self._db_started_at + or datetime.now(timezone.utc).isoformat(), total_steps = self._progress.total_steps or None, ) self._db_run_created = True except Exception: - logger.warning("Failed to create DB run record for early failure", exc_info = True) + logger.warning( + "Failed to create DB run record for early failure", exc_info = True + ) def _finalize_run_in_db( self, @@ -800,7 +827,10 @@ class TrainingBackend: ended_at = datetime.now(timezone.utc).isoformat(), final_step = self._progress.step, final_loss = self._progress.loss - if (self._progress.loss is not None and math.isfinite(self._progress.loss)) + if ( + self._progress.loss is not None + and math.isfinite(self._progress.loss) + ) else None, duration_seconds = self._progress.elapsed_seconds, loss_sparkline = _json.dumps(sparkline), @@ -809,11 +839,17 @@ class TrainingBackend: ) self._run_finalized = True except Exception: - logger.warning("Failed to finalize run in DB (status=%s)", status, exc_info = True) + logger.warning( + "Failed to finalize run in DB (status=%s)", status, exc_info = True + ) def _flush_metrics_to_db(self) -> None: """Flush buffered metrics to the database and update live progress.""" - if not self._metric_buffer or not self.current_job_id or not self._db_run_created: + if ( + not self._metric_buffer + or not self.current_job_id + or not self._db_run_created + ): return # Cap buffer to bound memory growth. if len(self._metric_buffer) > 500: @@ -833,7 +869,10 @@ class TrainingBackend: id = self.current_job_id, step = self._progress.step, loss = self._progress.loss - if (self._progress.loss is not None and math.isfinite(self._progress.loss)) + if ( + self._progress.loss is not None + and math.isfinite(self._progress.loss) + ) else None, duration_seconds = self._progress.elapsed_seconds, ) @@ -951,7 +990,9 @@ class TrainingBackend: else: title = "Training Loss" - ax.set_title(title, fontsize = 11, fontweight = "bold", pad = 10, color = style["text"]) + ax.set_title( + title, fontsize = 11, fontweight = "bold", pad = 10, color = style["text"] + ) ax.grid(True, alpha = 0.4, linestyle = "--", color = style["grid_color"]) ax.tick_params(colors = style["text"], which = "both") ax.spines["top"].set_visible(False) diff --git a/studio/backend/core/training/worker.py b/studio/backend/core/training/worker.py index 1120744a2d..8768bb9b5d 100644 --- a/studio/backend/core/training/worker.py +++ b/studio/backend/core/training/worker.py @@ -36,7 +36,8 @@ from typing import Any, Callable if sys.platform.startswith("linux") and "HSA_ENABLE_DXG_DETECTION" not in os.environ: try: if os.path.exists("/dev/dxg") and any( - os.path.exists(_p + "/librocdxg.so") for _p in ("/opt/rocm/lib", "/opt/rocm/lib64") + os.path.exists(_p + "/librocdxg.so") + for _p in ("/opt/rocm/lib", "/opt/rocm/lib64") ): os.environ["HSA_ENABLE_DXG_DETECTION"] = "1" except Exception: @@ -54,7 +55,9 @@ from utils.wheel_utils import ( ) -def _output_dir_from_resume_checkpoint(resume_from_checkpoint: str | None) -> str | None: +def _output_dir_from_resume_checkpoint( + resume_from_checkpoint: str | None, +) -> str | None: if not resume_from_checkpoint: return None path = Path(resume_from_checkpoint) @@ -117,7 +120,9 @@ if sys.platform == "win32": try: if os.path.isdir(_default_root): - for _ver in sorted(os.listdir(_default_root), key = _ver_key, reverse = True): + for _ver in sorted( + os.listdir(_default_root), key = _ver_key, reverse = True + ): _bin = os.path.join(_default_root, _ver, "bin") if os.path.isdir(_bin): _candidates.append(_bin) @@ -258,7 +263,9 @@ def _install_package_wheel_first( "(this may take several minutes)..." ) else: - pypi_status_message = f"Installing {display_name} from PyPI for faster training..." + pypi_status_message = ( + f"Installing {display_name} from PyPI for faster training..." + ) _send_status(event_queue, pypi_status_message) @@ -343,7 +350,8 @@ def _install_package_wheel_first( ) _send_status( event_queue, - f"{display_name} installation timed out after " f"{_run_kwargs.get('timeout')}s", + f"{display_name} installation timed out after " + f"{_run_kwargs.get('timeout')}s", ) return False @@ -461,7 +469,9 @@ def _ensure_flash_linear_attention_unconditional(event_queue: Any) -> bool: if os.getenv(_FLA_SKIP_ENV) == "1": return False if sys.platform == "win32": - logger.info("Skipping flash-linear-attention install: no prebuilt wheel for Windows") + logger.info( + "Skipping flash-linear-attention install: no prebuilt wheel for Windows" + ) return False if sys.version_info < _FLA_MIN_PYTHON: logger.info( @@ -536,7 +546,9 @@ def _ensure_flash_linear_attention_unconditional(event_queue: Any) -> bool: ) except _sp.TimeoutExpired: logger.warning("flash-linear-attention install timed out; continuing") - _send_status(event_queue, "flash-linear-attention install timed out; continuing") + _send_status( + event_queue, "flash-linear-attention install timed out; continuing" + ) return False if result.returncode != 0: @@ -727,7 +739,10 @@ def _rocm_classify_unified_memory(props: Any) -> tuple[str, bool]: # Arch attrs absent — fall back to device-name matching. dev_lower = (getattr(props, "name", "") or "").lower() is_unified = ( - "890m" in dev_lower or "880m" in dev_lower or "8060s" in dev_lower or "8050s" in dev_lower + "890m" in dev_lower + or "880m" in dev_lower + or "8060s" in dev_lower + or "8050s" in dev_lower ) return gcn_arch, is_unified @@ -770,7 +785,9 @@ def _run_pip(cmd: list[str], event_queue: Any, label: str) -> bool: _send_status(event_queue, f"{label} install timed out; continuing") return False if result.returncode != 0: - logger.warning("%s install failed (continuing without it):\n%s", label, result.stdout) + logger.warning( + "%s install failed (continuing without it):\n%s", label, result.stdout + ) _send_status(event_queue, f"{label} install failed; continuing") return False return True @@ -871,7 +888,9 @@ def _ensure_tilelang_backend(event_queue: Any, model_name: str) -> None: # UNSLOTH_STUDIO_SKIP_FAST_PATH_HOOKS=1 falls back to the substring path. -def _rebind_in_already_imported_modules(*, attr_name: str, old_obj: Any, new_obj: Any) -> int: +def _rebind_in_already_imported_modules( + *, attr_name: str, old_obj: Any, new_obj: Any +) -> int: """Rebind `attr_name -> new_obj` in every module that imported `old_obj`. `from X import Y` creates a local binding that reassigning X.Y won't reach. @@ -944,7 +963,9 @@ def _install_fast_path_hooks(event_queue: Any, model_name: str) -> None: try: ok = bool(install_fn(event_queue)) except Exception as exc: - logger.warning("%s install raised: %s; falling back to torch", gate_name, exc) + logger.warning( + "%s install raised: %s; falling back to torch", gate_name, exc + ) ok = False logger.info("%s hook done; available=%s", gate_name, ok) # post_available_fn handles "gate already True but ancillary kernel broken" @@ -953,7 +974,9 @@ def _install_fast_path_hooks(event_queue: Any, model_name: str) -> None: try: post_available_fn(event_queue) except Exception as exc: - logger.warning("%s post-available step raised: %s; continuing", gate_name, exc) + logger.warning( + "%s post-available step raised: %s; continuing", gate_name, exc + ) state["installed"] = True return ok @@ -964,7 +987,9 @@ def _install_fast_path_hooks(event_queue: Any, model_name: str) -> None: def _fla_install(eq: Any) -> bool: # FLA alone ~2.35x; +tilelang adds ~26%. tilelang is GDN-only (Qwen3.5 family). if not _ensure_flash_linear_attention_unconditional(eq): - logger.info("FLA install did not produce an importable runtime; skipping TileLang") + logger.info( + "FLA install did not produce an importable runtime; skipping TileLang" + ) return False if _model_wants_tilelang(model_name): _ensure_tilelang_backend_unconditional(eq) @@ -979,7 +1004,10 @@ def _install_fast_path_hooks(event_queue: Any, model_name: str) -> None: # FLA imports; repair tilelang if missing or on the broken tvm-ffi list. if not _model_wants_tilelang(model_name): return - if _installed_tvm_ffi_version() not in _TVM_FFI_BROKEN_VERSIONS and _tilelang_importable(): + if ( + _installed_tvm_ffi_version() not in _TVM_FFI_BROKEN_VERSIONS + and _tilelang_importable() + ): return _ensure_tilelang_backend_unconditional(eq) @@ -995,7 +1023,9 @@ def _install_fast_path_hooks(event_queue: Any, model_name: str) -> None: pypi_version = _CAUSAL_CONV1D_PACKAGE_VERSION, filename_prefix = "causal_conv1d", release_tag = _CAUSAL_CONV1D_RELEASE_TAG, - release_base_url = ("https://github.com/Dao-AILab/causal-conv1d/releases/download"), + release_base_url = ( + "https://github.com/Dao-AILab/causal-conv1d/releases/download" + ), ) return bool(ok) @@ -1015,7 +1045,9 @@ def _install_fast_path_hooks(event_queue: Any, model_name: str) -> None: rebound = _rebind_in_already_imported_modules( attr_name = gate_name, old_obj = original, new_obj = wrapped ) - logger.info("Installed fast-path hook on %s (rebound %d modules)", gate_name, rebound) + logger.info( + "Installed fast-path hook on %s (rebound %d modules)", gate_name, rebound + ) def _should_try_runtime_flash_attn_install(max_seq_length: int) -> bool: @@ -1171,7 +1203,10 @@ def _resize_mlx_vlm_images( image_layout = None, ): if isinstance(value, list): - return [_resize_mlx_vlm_image(image, resize, image_layout = image_layout) for image in value] + return [ + _resize_mlx_vlm_image(image, resize, image_layout = image_layout) + for image in value + ] return _resize_mlx_vlm_image(value, resize, image_layout = image_layout) @@ -1255,7 +1290,8 @@ def _normalize_mlx_studio_optimizer(value): except KeyError: supported = ", ".join(sorted(_MLX_STUDIO_OPTIM_MAP)) raise ValueError( - f"Unsupported optimizer for MLX training: {value!r}. " f"Supported values: {supported}." + f"Unsupported optimizer for MLX training: {value!r}. " + f"Supported values: {supported}." ) @@ -1277,7 +1313,9 @@ def _resolve_mlx_local_dataset_files(file_paths: list) -> list[str]: all_files: list[str] = [] for dataset_file in file_paths or []: file_path = ( - dataset_file if os.path.isabs(dataset_file) else str(resolve_dataset_path(dataset_file)) + dataset_file + if os.path.isabs(dataset_file) + else str(resolve_dataset_path(dataset_file)) ) file_path_obj = Path(file_path) @@ -1409,7 +1447,9 @@ def _run_mlx_training(event_queue, stop_queue, config): raise NotImplementedError(message) optim_name = _normalize_mlx_studio_optimizer(config.get("optim", "adamw_8bit")) - lr_scheduler_type = _normalize_mlx_studio_scheduler(config.get("lr_scheduler_type", "linear")) + lr_scheduler_type = _normalize_mlx_studio_scheduler( + config.get("lr_scheduler_type", "linear") + ) # ── 1. Load model ── # Force text-only for non-image datasets even on vision-capable models @@ -1489,9 +1529,15 @@ def _run_mlx_training(event_queue, stop_queue, config): finetune_language = config.get("finetune_language_layers", True) finetune_attention = config.get("finetune_attention_modules", True) finetune_mlp = config.get("finetune_mlp_modules", True) - finetune_vision = config.get("finetune_vision_layers", False) if is_vlm else False + finetune_vision = ( + config.get("finetune_vision_layers", False) if is_vlm else False + ) - if (finetune_attention or finetune_mlp) and not finetune_language and not finetune_vision: + if ( + (finetune_attention or finetune_mlp) + and not finetune_language + and not finetune_vision + ): finetune_language = True peft_kwargs["finetune_language_layers"] = finetune_language @@ -1524,7 +1570,9 @@ def _run_mlx_training(event_queue, stop_queue, config): if len(file_paths) == 1: p = Path(file_paths[0]) - if p.is_dir() and ((p / "dataset_info.json").exists() or (p / "state.json").exists()): + if p.is_dir() and ( + (p / "dataset_info.json").exists() or (p / "state.json").exists() + ): return load_from_disk(str(p)) all_files = _resolve_mlx_local_dataset_files(file_paths) if not all_files: @@ -1612,7 +1660,9 @@ def _run_mlx_training(event_queue, stop_queue, config): ) else: errors = vlm_info.get("errors", []) - raise ValueError(f"VLM dataset format conversion failed: {'; '.join(errors)}") + raise ValueError( + f"VLM dataset format conversion failed: {'; '.join(errors)}" + ) if eval_dataset is not None: ev_info = format_and_template_dataset( eval_dataset, @@ -1757,7 +1807,11 @@ def _run_mlx_training(event_queue, stop_queue, config): ) template_name = MODEL_TO_TEMPLATE_MAPPER.get(model_name.lower()) - markers = TEMPLATE_TO_RESPONSES_MAPPER.get(template_name) if template_name else None + markers = ( + TEMPLATE_TO_RESPONSES_MAPPER.get(template_name) + if template_name + else None + ) if markers: trainer = train_on_responses_only( trainer, @@ -1849,7 +1903,11 @@ def _run_mlx_training(event_queue, stop_queue, config): "train/tokens_per_sec": tok_s, "train/peak_gb": peak_gb, "train/num_tokens": num_tokens, - **({"train/grad_norm": grad_norm} if grad_norm is not None else {}), + **( + {"train/grad_norm": grad_norm} + if grad_norm is not None + else {} + ), }, step = step, ) @@ -1872,7 +1930,9 @@ def _run_mlx_training(event_queue, stop_queue, config): _send("progress", step = step, eval_loss = eval_loss) if wandb_run is not None: try: - wandb_run.log({"eval/loss": eval_loss, "eval/perplexity": perplexity}, step = step) + wandb_run.log( + {"eval/loss": eval_loss, "eval/perplexity": perplexity}, step = step + ) except Exception: pass if tb_writer is not None: @@ -2174,7 +2234,9 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) -> if os.path.isfile(os.path.join(_scripts_dir, "hipInfo.exe")): import shutil as _shutil if not _shutil.which("hipinfo.exe"): - os.environ["PATH"] = _scripts_dir + os.pathsep + os.environ.get("PATH", "") + os.environ["PATH"] = ( + _scripts_dir + os.pathsep + os.environ.get("PATH", "") + ) # BNB picks a rocm DLL from torch.version.hip, but AMD's Windows BNB # wheel may ship a DLL whose suffix doesn't match. Detect the actual @@ -2215,7 +2277,9 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) -> # so later import fixes can still redetect or opt out. DLL # with unparsable name -> seeded value or "72". if _found_rocm_bnb: - _bnb_rocm_ver = _bnb_rocm_ver or os.environ.get("BNB_ROCM_VERSION") or "72" + _bnb_rocm_ver = ( + _bnb_rocm_ver or os.environ.get("BNB_ROCM_VERSION") or "72" + ) os.environ["BNB_ROCM_VERSION"] = _bnb_rocm_ver os.environ["UNSLOTH_BNB_ROCM_VERSION_SOURCE"] = "detected" logger.info( @@ -2229,7 +2293,9 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) -> # the rocm version embedded in torch.__version__ when version.hip is # unset (AMD SDK / Radeon wheels). def _hip_ver_at_least(major: int, minor: int) -> bool: - _hip_str = getattr(getattr(_torch_for_rocm, "version", None), "hip", None) + _hip_str = getattr( + getattr(_torch_for_rocm, "version", None), "hip", None + ) if not _hip_str: # Try the standard "+rocmX.Y.Z" embedded version first. _ver_match = re.search(r"rocm(\d+)\.(\d+)", _build_version_for_rocm) @@ -2321,7 +2387,9 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) -> if prev < self.shape[0]: a_tail = self[prev:].contiguous() b_tail = ( - mat2[-1].contiguous() if mat2.dim() == 3 else mat2.contiguous() + mat2[-1].contiguous() + if mat2.dim() == 3 + else mat2.contiguous() ) pieces.append(_t.mm(a_tail, b_tail)) result = ( @@ -2491,7 +2559,11 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) -> def _on_progress(progress: TrainingProgress): has_train_loss = progress.step > 0 and progress.loss is not None has_eval_loss = progress.eval_loss is not None - if (progress.step == 0 and progress.total_steps > 0) or has_train_loss or has_eval_loss: + if ( + (progress.step == 0 and progress.total_steps > 0) + or has_train_loss + or has_eval_loss + ): event_queue.put( { "type": "progress", @@ -2601,12 +2673,15 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) -> if dataset is None or trainer.should_stop: if trainer.should_stop: - event_queue.put({"type": "complete", "output_dir": None, "ts": time.time()}) + event_queue.put( + {"type": "complete", "output_dir": None, "ts": time.time()} + ) else: event_queue.put( { "type": "error", - "error": trainer.training_progress.error or "Failed to load dataset", + "error": trainer.training_progress.error + or "Failed to load dataset", "stack": "", "ts": time.time(), } @@ -2627,7 +2702,9 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) -> desc = getattr(bar, "desc", "") or "" if total > 0 and n > 0 and desc: pct = min(int(n * 100 / total), 100) - _send_status(event_queue, f"{desc.strip()} {pct}% ({n:,}/{total:,})") + _send_status( + event_queue, f"{desc.strip()} {pct}% ({n:,}/{total:,})" + ) except (AttributeError, ReferenceError): pass _tqdm_stop.wait(3) @@ -2655,7 +2732,9 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) -> ) if not success or trainer.should_stop: if trainer.should_stop: - event_queue.put({"type": "complete", "output_dir": None, "ts": time.time()}) + event_queue.put( + {"type": "complete", "output_dir": None, "ts": time.time()} + ) else: error_msg = trainer.training_progress.error or "Failed to load model" event_queue.put( @@ -2696,7 +2775,9 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) -> lora_r = config.get("lora_r", 128), lora_alpha = config.get("lora_alpha", 32), lora_dropout = config.get("lora_dropout", 0.0), - use_gradient_checkpointing = config.get("gradient_checkpointing", "unsloth"), + use_gradient_checkpointing = config.get( + "gradient_checkpointing", "unsloth" + ), use_rslora = config.get("use_rslora", False), use_loftq = config.get("use_loftq", False), ) @@ -2706,13 +2787,17 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) -> use_lora = True, finetune_vision_layers = config.get("finetune_vision_layers", True), finetune_language_layers = config.get("finetune_language_layers", True), - finetune_attention_modules = config.get("finetune_attention_modules", True), + finetune_attention_modules = config.get( + "finetune_attention_modules", True + ), finetune_mlp_modules = config.get("finetune_mlp_modules", True), target_modules = config.get("target_modules"), lora_r = config.get("lora_r", 16), lora_alpha = config.get("lora_alpha", 16), lora_dropout = config.get("lora_dropout", 0.0), - use_gradient_checkpointing = config.get("gradient_checkpointing", "unsloth"), + use_gradient_checkpointing = config.get( + "gradient_checkpointing", "unsloth" + ), use_rslora = config.get("use_rslora", False), use_loftq = config.get("use_loftq", False), ) @@ -2722,12 +2807,15 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) -> if not success or trainer.should_stop: if trainer.should_stop: - event_queue.put({"type": "complete", "output_dir": None, "ts": time.time()}) + event_queue.put( + {"type": "complete", "output_dir": None, "ts": time.time()} + ) else: event_queue.put( { "type": "error", - "error": trainer.training_progress.error or "Failed to prepare model", + "error": trainer.training_progress.error + or "Failed to prepare model", "stack": "", "ts": time.time(), } @@ -2783,7 +2871,9 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) -> ensure_dir(Path(tensorboard_dir)) # Start training directly — no inner thread, we ARE the subprocess. - dataset_display = config.get("hf_dataset", "") or config.get("uploaded_file", "") or "" + dataset_display = ( + config.get("hf_dataset", "") or config.get("uploaded_file", "") or "" + ) _send_status( event_queue, f'Training "{model_name}"' @@ -2807,7 +2897,9 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) -> weight_decay = config.get("weight_decay", 0.001), random_seed = config.get("random_seed", 3407), packing = config.get("packing", False), - train_on_completions = False if is_cpt else config.get("train_on_completions", False), + train_on_completions = False + if is_cpt + else config.get("train_on_completions", False), enable_wandb = config.get("enable_wandb", False), wandb_project = config.get("wandb_project", "unsloth-training"), wandb_token = config.get("wandb_token"), @@ -3068,7 +3160,9 @@ def _run_embedding_training(event_queue: Any, stop_queue: Any, config: dict) -> if candidates: all_files.extend(str(c) for c in candidates) continue - raise ValueError(f"No supported data files in directory: {file_path_obj}") + raise ValueError( + f"No supported data files in directory: {file_path_obj}" + ) else: all_files.append(file_path) @@ -3182,7 +3276,9 @@ def _run_embedding_training(event_queue: Any, stop_queue: Any, config: dict) -> resume_from_checkpoint ) if not output_dir: - output_dir = str(resolve_output_dir(f"{model_name.replace('/', '_')}_{int(time.time())}")) + output_dir = str( + resolve_output_dir(f"{model_name.replace('/', '_')}_{int(time.time())}") + ) output_dir = str(resolve_output_dir(output_dir)) num_epochs = config.get("num_epochs", 2) diff --git a/studio/backend/hub/routes/datasets.py b/studio/backend/hub/routes/datasets.py index edf4f36ac0..40a595b022 100644 --- a/studio/backend/hub/routes/datasets.py +++ b/studio/backend/hub/routes/datasets.py @@ -61,14 +61,17 @@ async def list_cached_datasets(current_subject: str = Depends(get_current_subjec @router.delete("/cached", response_model = DeleteCachedDatasetResponse) async def delete_cached_dataset( - repo_id: str = Body(..., embed = True), current_subject: str = Depends(get_current_subject) + repo_id: str = Body(..., embed = True), + current_subject: str = Depends(get_current_subject), ): return await cache_inventory.delete_cached_dataset_response(repo_id) @router.get("/download-progress", response_model = DownloadProgressResponse) async def get_dataset_download_progress( - repo_id: str = Query(..., description = "HuggingFace dataset repo ID, e.g. 'unsloth/LaTeX_OCR'"), + repo_id: str = Query( + ..., description = "HuggingFace dataset repo ID, e.g. 'unsloth/LaTeX_OCR'" + ), expected_bytes: int = Query(0, description = "Expected total download size in bytes"), hf_token: Optional[str] = Depends(get_hf_token), current_subject: str = Depends(get_current_subject), @@ -89,9 +92,12 @@ async def download_dataset( return await downloads.download_dataset_response(body, hf_token) -@router.post("/download/cancel", response_model = CancelDatasetDownloadResponse, status_code = 202) +@router.post( + "/download/cancel", response_model = CancelDatasetDownloadResponse, status_code = 202 +) async def cancel_dataset_download( - body: CancelDatasetDownloadRequest, current_subject: str = Depends(get_current_subject) + body: CancelDatasetDownloadRequest, + current_subject: str = Depends(get_current_subject), ): return await downloads.cancel_dataset_download_response(body) diff --git a/studio/backend/hub/routes/inventory.py b/studio/backend/hub/routes/inventory.py index fcfdd0ad14..696ed810f4 100644 --- a/studio/backend/hub/routes/inventory.py +++ b/studio/backend/hub/routes/inventory.py @@ -130,7 +130,9 @@ async def cancel_download_model( @router.get("/download-status", response_model = DownloadJobStatus) async def get_download_status( repo_id: str = Query(..., description = "HuggingFace repo ID"), - gguf_variant: str = Query("", description = "Quantization variant (empty for safetensors)"), + gguf_variant: str = Query( + "", description = "Quantization variant (empty for safetensors)" + ), current_subject: str = Depends(get_current_subject), ): return await downloads.get_download_status_response(repo_id, gguf_variant) @@ -147,7 +149,9 @@ async def get_active_downloads( @router.get("/transport-status", response_model = TransportStatusResponse) async def get_model_transport_status( repo_id: str = Query(..., description = "HuggingFace repo ID"), - gguf_variant: str = Query("", description = "Quantization variant (empty for safetensors)"), + gguf_variant: str = Query( + "", description = "Quantization variant (empty for safetensors)" + ), hf_token: Optional[str] = Depends(get_hf_token), current_subject: str = Depends(get_current_subject), ): diff --git a/studio/backend/hub/schemas/downloads.py b/studio/backend/hub/schemas/downloads.py index dccd0c6733..aa64ceb835 100644 --- a/studio/backend/hub/schemas/downloads.py +++ b/studio/backend/hub/schemas/downloads.py @@ -7,7 +7,9 @@ from pydantic import BaseModel, Field from typing import List, Literal, Optional -DownloadJobState = Literal["idle", "running", "cancelling", "cancelled", "complete", "error"] +DownloadJobState = Literal[ + "idle", "running", "cancelling", "cancelled", "complete", "error" +] class DownloadModelRequest(BaseModel): diff --git a/studio/backend/hub/schemas/inventory.py b/studio/backend/hub/schemas/inventory.py index c333c7ca89..9b681385a4 100644 --- a/studio/backend/hub/schemas/inventory.py +++ b/studio/backend/hub/schemas/inventory.py @@ -17,13 +17,19 @@ ModelRuntime = Literal["llama_cpp", "transformers", "adapter", "unknown"] class GgufVariantDetail(BaseModel): """A single GGUF quantization variant in a HuggingFace repo.""" - filename: str = Field(..., description = "GGUF filename (e.g., 'gemma-3-4b-it-Q4_K_M.gguf')") - quant: str = Field(..., description = "Quantization label or internal GGUF variant key") + filename: str = Field( + ..., description = "GGUF filename (e.g., 'gemma-3-4b-it-Q4_K_M.gguf')" + ) + quant: str = Field( + ..., description = "Quantization label or internal GGUF variant key" + ) display_label: Optional[str] = Field( None, description = "Optional user-facing label when quant is an internal key" ) size_bytes: int = Field(0, description = "File size in bytes") - download_size_bytes: int = Field(0, description = "Total bytes needed to download this variant") + download_size_bytes: int = Field( + 0, description = "Total bytes needed to download this variant" + ) downloaded: bool = Field( False, description = "Whether this variant is already in the local HF cache" ) @@ -132,7 +138,9 @@ class LocalModelInfo(BaseModel): class LocalModelListResponse(BaseModel): """Response schema for listing local/cached models.""" - models_dir: str = Field(..., description = "Directory scanned for custom local models") + models_dir: str = Field( + ..., description = "Directory scanned for custom local models" + ) hf_cache_dir: Optional[str] = Field( None, description = "HF cache root that was scanned", diff --git a/studio/backend/hub/services/__init__.py b/studio/backend/hub/services/__init__.py index e86fcb6f46..be6333260b 100644 --- a/studio/backend/hub/services/__init__.py +++ b/studio/backend/hub/services/__init__.py @@ -12,7 +12,9 @@ from fastapi import HTTPException from hub.utils.hf_cache_state import resolve_destructive_case_matches -def resolve_destructive_repo_ids(repo_id: str, candidates: Iterable[str], *, noun: str) -> set[str]: +def resolve_destructive_repo_ids( + repo_id: str, candidates: Iterable[str], *, noun: str +) -> set[str]: """Cache-dir repo ids a destructive op on *repo_id* may target. Refuses with 409 on ambiguous case-only matches so a delete never removes diff --git a/studio/backend/hub/services/datasets/cache_inventory.py b/studio/backend/hub/services/datasets/cache_inventory.py index a180c9df58..b2d3aba8c3 100644 --- a/studio/backend/hub/services/datasets/cache_inventory.py +++ b/studio/backend/hub/services/datasets/cache_inventory.py @@ -194,7 +194,9 @@ def _hf_datasets_cache_roots() -> list[Path]: if hf_home: _add(Path(hf_home).expanduser() / "datasets") - xdg_cache = Path(os.environ.get("XDG_CACHE_HOME", Path.home() / ".cache")).expanduser() + xdg_cache = Path( + os.environ.get("XDG_CACHE_HOME", Path.home() / ".cache") + ).expanduser() _add(xdg_cache / "huggingface" / "datasets") return roots @@ -284,7 +286,11 @@ def _scan_hf_dataset_caches() -> list[dict]: rev_id = getattr(rev, "commit_hash", None) or str(id(rev)) for f in rev.files: blob_path = getattr(f, "blob_path", None) - key = str(blob_path) if blob_path else f"{rev_id}:{f.file_name}" + key = ( + str(blob_path) + if blob_path + else f"{rev_id}:{f.file_name}" + ) unique_blobs[key] = int(f.size_on_disk or 0) total_size = sum(unique_blobs.values()) key = repo_info.repo_id.lower() @@ -320,7 +326,9 @@ def _scan_hf_dataset_caches() -> list[dict]: existing = seen_lower.get(key) if _prefer_dataset_cache_row(row, existing): seen_lower[key] = row - elif existing is not None and bool(existing.get("partial")) == bool(row.get("partial")): + elif existing is not None and bool(existing.get("partial")) == bool( + row.get("partial") + ): existing["size_bytes"] = max(existing["size_bytes"], row["size_bytes"]) existing["cache_path"] = existing.get("cache_path") or row.get("cache_path") if ( @@ -332,7 +340,9 @@ def _scan_hf_dataset_caches() -> list[dict]: for row in _scan_processed_dataset_caches(): key = row["repo_id"].lower() existing = seen_lower.get(key) - if existing is None or (bool(existing.get("partial")) and not bool(row.get("partial"))): + if existing is None or ( + bool(existing.get("partial")) and not bool(row.get("partial")) + ): seen_lower[key] = row else: existing["size_bytes"] = max(existing["size_bytes"], row["size_bytes"]) @@ -366,7 +376,9 @@ async def delete_cached_dataset_response(repo_id: str) -> dict: if not _is_valid_repo_id(repo_id): raise HTTPException(status_code = 400, detail = "Invalid repo_id format") - repo_key = await asyncio.to_thread(resolve_cached_repo_id_case, repo_id, repo_type = "dataset") + repo_key = await asyncio.to_thread( + resolve_cached_repo_id_case, repo_id, repo_type = "dataset" + ) if not downloads.registry.begin_delete(repo_key): raise HTTPException( status_code = 400, @@ -401,7 +413,9 @@ def _delete_cached_dataset_blocking(repo_id: str) -> dict: if str(repo_info.repo_id) not in matched_repo_ids: continue try: - strategy = hf_cache.delete_revisions(*(rev.commit_hash for rev in repo_info.revisions)) + strategy = hf_cache.delete_revisions( + *(rev.commit_hash for rev in repo_info.revisions) + ) strategy.execute() deleted = True except Exception as exc: @@ -430,7 +444,9 @@ def _delete_cached_dataset_blocking(repo_id: str) -> dict: cache_purged = purge_repo_cache_dirs("dataset", repo_id) partial_purged = purge_partial_repo("dataset", repo_id) state_purged = download_manifest.purge_all_state_for_repo("dataset", repo_id) > 0 - if not (deleted or processed_deleted or cache_purged or partial_purged or state_purged): + if not ( + deleted or processed_deleted or cache_purged or partial_purged or state_purged + ): raise HTTPException(status_code = 404, detail = "Dataset not found in cache") return {"status": "deleted", "repo_id": repo_id} diff --git a/studio/backend/hub/services/datasets/downloads.py b/studio/backend/hub/services/datasets/downloads.py index ac90be8c6f..778ac65559 100644 --- a/studio/backend/hub/services/datasets/downloads.py +++ b/studio/backend/hub/services/datasets/downloads.py @@ -37,9 +37,7 @@ from hub.utils.snapshot_filters import ( logger = get_logger(__name__) -_dataset_size_cache: "OrderedDict[str, tuple[int, frozenset[str], bool, str, float]]" = ( - OrderedDict() -) +_dataset_size_cache: "OrderedDict[str, tuple[int, frozenset[str], bool, str, float]]" = OrderedDict() _dataset_size_neg_cache: "OrderedDict[tuple[str, str], float]" = OrderedDict() _DATASET_SIZE_CACHE_MAX = 256 _DATASET_SIZE_POS_TTL = 60.0 @@ -88,7 +86,9 @@ def get_dataset_snapshot_metadata_cached( ) total = total_size_for_siblings(info.siblings) hashes = blob_hashes_for_siblings(info.siblings) - restricted = bool(getattr(info, "private", False) or getattr(info, "gated", False)) + restricted = bool( + getattr(info, "private", False) or getattr(info, "gated", False) + ) except Exception: with _dataset_size_cache_lock: _dataset_size_neg_cache[cache_key] = time.monotonic() @@ -132,7 +132,9 @@ async def get_dataset_download_progress_response( ) -def _dataset_status(key: str, *, repo_id: Optional[str] = None) -> DatasetDownloadJobStatus: +def _dataset_status( + key: str, *, repo_id: Optional[str] = None +) -> DatasetDownloadJobStatus: state, error, generation = download_lifecycle.idle_status( _registry, key, @@ -154,7 +156,9 @@ async def download_dataset_response( detail = f"Invalid repo_id: {repo_id!r}", ) # Canonicalize so two different-cased paste-ins share one job + cache dir. - repo_id = await asyncio.to_thread(resolve_cached_repo_id_case, repo_id, repo_type = "dataset") + repo_id = await asyncio.to_thread( + resolve_cached_repo_id_case, repo_id, repo_type = "dataset" + ) key = _download_job_key(repo_id) transport = download_lifecycle.resolve_transport(body.use_xet) @@ -211,7 +215,9 @@ async def cancel_dataset_download_response(body: CancelDatasetDownloadRequest) - status_code = 400, detail = f"Invalid repo_id: {repo_id!r}", ) - repo_id = await asyncio.to_thread(resolve_cached_repo_id_case, repo_id, repo_type = "dataset") + repo_id = await asyncio.to_thread( + resolve_cached_repo_id_case, repo_id, repo_type = "dataset" + ) key = _download_job_key(repo_id) state = download_lifecycle.cancel_worker( @@ -224,21 +230,29 @@ async def cancel_dataset_download_response(body: CancelDatasetDownloadRequest) - return {"repo_id": repo_id, "state": state} -async def get_dataset_download_status_response(repo_id: str) -> DatasetDownloadJobStatus: +async def get_dataset_download_status_response( + repo_id: str, +) -> DatasetDownloadJobStatus: """Return the latest state of a background dataset download job.""" repo_id = repo_id.strip() if not _is_valid_repo_id(repo_id): return DatasetDownloadJobStatus(state = "idle") - repo_id = await asyncio.to_thread(resolve_cached_repo_id_case, repo_id, repo_type = "dataset") + repo_id = await asyncio.to_thread( + resolve_cached_repo_id_case, repo_id, repo_type = "dataset" + ) return _dataset_status(_download_job_key(repo_id), repo_id = repo_id) -async def get_active_dataset_downloads_response(repo_id: str = "") -> ActiveDownloadsResponse: +async def get_active_dataset_downloads_response( + repo_id: str = "", +) -> ActiveDownloadsResponse: repo_id = repo_id.strip() if repo_id and not _is_valid_repo_id(repo_id): return ActiveDownloadsResponse(downloads = []) canonical_repo_id = ( - await asyncio.to_thread(resolve_cached_repo_id_case, repo_id, repo_type = "dataset") + await asyncio.to_thread( + resolve_cached_repo_id_case, repo_id, repo_type = "dataset" + ) if repo_id else None ) @@ -260,7 +274,9 @@ async def get_dataset_transport_status_response(repo_id: str) -> dict: return {"has_partial": False, "last_transport": None, "resumable": False} return { "has_partial": has_active_incomplete_blobs("dataset", repo_id), - "last_transport": download_registry.read_active_transport_marker("dataset", repo_id), + "last_transport": download_registry.read_active_transport_marker( + "dataset", repo_id + ), "resumable": download_registry.is_resumable_partial("dataset", repo_id), } diff --git a/studio/backend/hub/services/datasets/formatting.py b/studio/backend/hub/services/datasets/formatting.py index 1c78d80e21..d239f2b4b3 100644 --- a/studio/backend/hub/services/datasets/formatting.py +++ b/studio/backend/hub/services/datasets/formatting.py @@ -177,10 +177,14 @@ def _repo_file_matches_split(path: str, split: str) -> bool: def _select_tier1_repo_file( files: list[str], *, subset: Optional[str], train_split: str ) -> Optional[str]: - data_files = sorted(f for f in files if any(f.lower().endswith(ext) for ext in DATA_EXTS)) + data_files = sorted( + f for f in files if any(f.lower().endswith(ext) for ext in DATA_EXTS) + ) if not data_files: return None - tabular_files = [f for f in data_files if any(f.lower().endswith(ext) for ext in _TABULAR_EXTS)] + tabular_files = [ + f for f in data_files if any(f.lower().endswith(ext) for ext in _TABULAR_EXTS) + ] candidates = tabular_files or data_files if subset: candidates = [f for f in candidates if _repo_file_matches_label(f, subset)] @@ -405,7 +409,9 @@ def check_format_response( processed = format_dataset_preview(preview_slice) preview_samples = _serialize_preview_rows(processed) except Exception as e: - logger.warning(f"Processed preview generation failed (non-fatal): {e}") + logger.warning( + f"Processed preview generation failed (non-fatal): {e}" + ) preview_samples = _serialize_preview_rows(preview_slice) else: preview_samples = _serialize_preview_rows(preview_slice) @@ -416,7 +422,9 @@ def check_format_response( if image_col and image_col in (result.get("columns") or []): try: sample_val = preview_slice[0][image_col] - if isinstance(sample_val, str) and sample_val.startswith(("http://", "https://")): + if isinstance(sample_val, str) and sample_val.startswith( + ("http://", "https://") + ): url_warning = ( "This dataset contains image URLs instead of embedded images. " "Images will be downloaded during training, which may be slow for large datasets." @@ -483,7 +491,8 @@ def ai_assist_mapping_response( from hub.utils.llm_assist import llm_conversion_advisor truncated = [ - {col: str(s.get(col, ""))[:200] for col in request.columns} for s in request.samples[:5] + {col: str(s.get(col, ""))[:200] for col in request.columns} + for s in request.samples[:5] ] result = llm_conversion_advisor( diff --git a/studio/backend/hub/services/datasets/local.py b/studio/backend/hub/services/datasets/local.py index 8d48c4f735..56b7731cd9 100644 --- a/studio/backend/hub/services/datasets/local.py +++ b/studio/backend/hub/services/datasets/local.py @@ -223,7 +223,9 @@ def _stream_file_preview_slice(path: Path, preview_size: int): return Dataset.from_list(rows), None -def _load_local_preview_slice(*, dataset_path: Path, train_split: str, preview_size: int): +def _load_local_preview_slice( + *, dataset_path: Path, train_split: str, preview_size: int +): # Non-streaming loads take the cached builder lock; use the EACCES-safe wrapper. from utils.datasets.cache_safe import load_dataset_cache_safe as load_dataset @@ -258,7 +260,9 @@ def _load_local_preview_slice(*, dataset_path: Path, train_split: str, preview_s # Parquet/Arrow give a cheap exact total_rows via len()+select; JSON/CSV # carry no such metadata, so stream them and report total_rows=None. if suffix == ".parquet": - dataset = load_dataset("parquet", data_files = str(dataset_path), split = train_split) + dataset = load_dataset( + "parquet", data_files = str(dataset_path), split = train_split + ) total_rows = len(dataset) preview_slice = dataset.select(range(min(preview_size, total_rows))) return preview_slice, total_rows @@ -272,7 +276,9 @@ def _load_local_preview_slice(*, dataset_path: Path, train_split: str, preview_s ) return preview - raise HTTPException(status_code = 400, detail = f"Unsupported file format: {dataset_path.suffix}") + raise HTTPException( + status_code = 400, detail = f"Unsupported file format: {dataset_path.suffix}" + ) def _sanitize_filename(filename: str) -> str: @@ -285,7 +291,10 @@ def _sanitize_filename(filename: str) -> str: def _upload_too_large(size_bytes: int) -> HTTPException: return HTTPException( status_code = 413, - detail = (f"Upload is too large " f"({size_bytes:,} bytes; max {LOCAL_UPLOAD_MAX_BYTES:,})."), + detail = ( + f"Upload is too large " + f"({size_bytes:,} bytes; max {LOCAL_UPLOAD_MAX_BYTES:,})." + ), ) diff --git a/studio/backend/hub/services/download_lifecycle.py b/studio/backend/hub/services/download_lifecycle.py index 8256b00252..801223ca45 100644 --- a/studio/backend/hub/services/download_lifecycle.py +++ b/studio/backend/hub/services/download_lifecycle.py @@ -26,8 +26,12 @@ def backend_dir() -> Path: def resolve_transport(use_xet: bool) -> str: - transport = download_registry.TRANSPORT_XET if use_xet else download_registry.TRANSPORT_HTTP - unavailable_reason = download_registry.download_transport_unavailable_reason(transport) + transport = ( + download_registry.TRANSPORT_XET if use_xet else download_registry.TRANSPORT_HTTP + ) + unavailable_reason = download_registry.download_transport_unavailable_reason( + transport + ) if unavailable_reason is not None: raise HTTPException(status_code = 400, detail = unavailable_reason) return transport @@ -49,7 +53,9 @@ def spawn_worker( shared ``.incomplete`` (e.g. bundled mmproj) is never deleted. """ cwd = backend_dir() - mode = download_registry.TRANSPORT_XET if use_xet else download_registry.TRANSPORT_HTTP + mode = ( + download_registry.TRANSPORT_XET if use_xet else download_registry.TRANSPORT_HTTP + ) env = os.environ.copy() if protected_blob_hashes: env["UNSLOTH_PROTECTED_BLOB_HASHES"] = ",".join(sorted(protected_blob_hashes)) @@ -73,7 +79,9 @@ def spawn_worker( if hf_token: env["HF_TOKEN"] = hf_token existing_path = env.get("PYTHONPATH", "") - env["PYTHONPATH"] = f"{cwd}{os.pathsep}{existing_path}" if existing_path else str(cwd) + env["PYTHONPATH"] = ( + f"{cwd}{os.pathsep}{existing_path}" if existing_path else str(cwd) + ) return subprocess.Popen( [ sys.executable, @@ -216,7 +224,9 @@ def finalize_worker_exit( f"{label}: {stderr_text}" ) else: - logger.info(f"{log_prefix} worker diagnostics for {label}: {stderr_text}") + logger.info( + f"{log_prefix} worker diagnostics for {label}: {stderr_text}" + ) logger.info(f"{log_prefix} complete: {label}") # Defensive cleanup: the canonical clear is at download-start; this # catches the rare case where that failed but the download succeeded. @@ -422,13 +432,18 @@ def idle_status( def active_download_refs( - registry: download_registry.DownloadRegistry, repo_id: Optional[str], *, with_variant: bool + registry: download_registry.DownloadRegistry, + repo_id: Optional[str], + *, + with_variant: bool, ) -> list[ActiveDownload]: downloads: list[ActiveDownload] = [] for ref in registry.active_job_refs(repo_id): metadata = ref.metadata if with_variant: - ref_repo_id = metadata.repo_id if metadata is not None else ref.key.split("::", 1)[0] + ref_repo_id = ( + metadata.repo_id if metadata is not None else ref.key.split("::", 1)[0] + ) if metadata is not None: variant = metadata.variant else: diff --git a/studio/backend/hub/services/models/cache_inventory.py b/studio/backend/hub/services/models/cache_inventory.py index a961a6ae9d..ba5e962c6d 100644 --- a/studio/backend/hub/services/models/cache_inventory.py +++ b/studio/backend/hub/services/models/cache_inventory.py @@ -39,7 +39,9 @@ from hub.services.models.common import ( logger = get_logger(__name__) -_repo_size_cache: "OrderedDict[tuple[str, str], tuple[int, frozenset[str], float]]" = OrderedDict() +_repo_size_cache: "OrderedDict[tuple[str, str], tuple[int, frozenset[str], float]]" = ( + OrderedDict() +) _repo_size_neg_cache: "OrderedDict[tuple[str, str], float]" = OrderedDict() _REPO_SIZE_CACHE_MAX = 256 _REPO_SIZE_POS_TTL = 60.0 @@ -168,7 +170,9 @@ def _scan_cached_gguf() -> list[dict]: continue repo_id = repo_info.repo_id total_size = _repo_gguf_size_bytes(repo_info) - has_variant_state, variant_state_size = _gguf_variant_state_summary(repo_id) + has_variant_state, variant_state_size = _gguf_variant_state_summary( + repo_id + ) if total_size == 0 and not has_variant_state: continue partial = hf_cache_scan.is_gguf_repo_partial( @@ -239,7 +243,9 @@ def _repo_non_gguf_model_payload(repo_info) -> _CachedNonGgufPayload: has_transformers_safetensors = False has_checkpoint = False - def _record_blob(target: dict[str, int], file_obj, rev_id: str, file_name: str) -> None: + def _record_blob( + target: dict[str, int], file_obj, rev_id: str, file_name: str + ) -> None: blob_path = getattr(file_obj, "blob_path", None) size = int(file_obj.size_on_disk or 0) key = str(blob_path) if blob_path else f"{rev_id}:{file_name}" @@ -367,7 +373,9 @@ def _cached_model_local_metadata(repo_path: Path) -> dict: result["library_name"] = library_name.strip() tags = card.get("tags") if isinstance(tags, list): - clean_tags = [tag.strip() for tag in tags if isinstance(tag, str) and tag.strip()] + clean_tags = [ + tag.strip() for tag in tags if isinstance(tag, str) and tag.strip() + ] if clean_tags: result["tags"] = clean_tags return result diff --git a/studio/backend/hub/services/models/common.py b/studio/backend/hub/services/models/common.py index f381bffe9c..fe33fd00ee 100644 --- a/studio/backend/hub/services/models/common.py +++ b/studio/backend/hub/services/models/common.py @@ -84,7 +84,9 @@ def _is_model_directory(d: Path) -> bool: return False try: - has_config = (d / "config.json").exists() or (d / "adapter_config.json").exists() + has_config = (d / "config.json").exists() or ( + d / "adapter_config.json" + ).exists() if not has_config: return False return any(_is_weight_file(f) for f in d.iterdir() if f.is_file()) @@ -193,7 +195,9 @@ def _apply_format_aware_partial( continue # GGUF row-level transport is ambiguous (variants may differ); per-variant # detail lives on GgufVariantDetail.partial_transport via the variants endpoint. - partial_transport = None if row.model_format == "gguf" else snapshot_partial_transport + partial_transport = ( + None if row.model_format == "gguf" else snapshot_partial_transport + ) rewritten.append( row.model_copy( update = { @@ -217,7 +221,9 @@ def _weight_basename(name: str) -> str: def _is_adapter_weight_name(name: str) -> bool: lower = _weight_basename(name) - return lower.startswith("adapter_model") and lower.endswith((".safetensors", ".bin")) + return lower.startswith("adapter_model") and lower.endswith( + (".safetensors", ".bin") + ) def _is_transformers_safetensors_weight_name(name: str) -> bool: @@ -267,7 +273,9 @@ def _classify_non_gguf_model_format( has_checkpoint_weights: bool, trusted_hf_cache_repo: bool = False, ) -> Optional[ModelFormat]: - if has_safetensors and (has_config or (trusted_hf_cache_repo and has_transformers_safetensors)): + if has_safetensors and ( + has_config or (trusted_hf_cache_repo and has_transformers_safetensors) + ): return "safetensors" if has_adapter_config and has_adapter_weights: return "adapter" @@ -278,7 +286,9 @@ def _classify_non_gguf_model_format( def _is_main_gguf_filename(name: str) -> bool: return ( - _is_gguf_filename(name) and not _is_mmproj_filename(name) and not _is_mtp_drafter_path(name) + _is_gguf_filename(name) + and not _is_mmproj_filename(name) + and not _is_mtp_drafter_path(name) ) @@ -445,7 +455,8 @@ def _local_model_info( ), load_id = load_id, model_id = model_id, - display_name = display_name or (scan_path.stem if scan_path.is_file() else scan_path.name), + display_name = display_name + or (scan_path.stem if scan_path.is_file() else scan_path.name), path = str(load_path), size_bytes = max(0, int(size_bytes or 0)), source = source, @@ -520,12 +531,17 @@ def _classify_local_path( (scan_path / "adapter_config.json").is_file() if scan_path.is_dir() else False ) adapter_config = _read_adapter_config(scan_path) if has_adapter_config else {} - adapter_base_model = _clean_optional_string(adapter_config.get("base_model_name_or_path")) + adapter_base_model = _clean_optional_string( + adapter_config.get("base_model_name_or_path") + ) adapter_type = _clean_optional_string(adapter_config.get("peft_type")) - training_method = _clean_optional_string(adapter_config.get("unsloth_training_method")) + training_method = _clean_optional_string( + adapter_config.get("unsloth_training_method") + ) has_adapter_weights = any(_is_adapter_weight_file(f) for f in files) has_safetensors = any( - f.suffix.lower() == ".safetensors" and not _is_adapter_weight_file(f) for f in files + f.suffix.lower() == ".safetensors" and not _is_adapter_weight_file(f) + for f in files ) has_transformers_safetensors = any( _is_transformers_safetensors_weight_file(f) and not _is_adapter_weight_file(f) @@ -554,7 +570,9 @@ def _classify_local_path( if f.suffix.lower() == ".safetensors" and not _is_adapter_weight_file(f) ) else: - size_bytes = _sum_file_sizes(f for f in files if _is_checkpoint_weight_file(f)) + size_bytes = _sum_file_sizes( + f for f in files if _is_checkpoint_weight_file(f) + ) rows.append( _local_model_info( scan_path = scan_path, diff --git a/studio/backend/hub/services/models/deletion.py b/studio/backend/hub/services/models/deletion.py index ecc9f8426d..0e7ed6f8df 100644 --- a/studio/backend/hub/services/models/deletion.py +++ b/studio/backend/hub/services/models/deletion.py @@ -75,7 +75,9 @@ def _path_exists_or_symlink(path: Path) -> bool: return False -def _repo_file_matches(target_repo, predicate) -> list[tuple[Path, Optional[Path], str]]: +def _repo_file_matches( + target_repo, predicate +) -> list[tuple[Path, Optional[Path], str]]: matches: list[tuple[Path, Optional[Path], str]] = [] for rev in getattr(target_repo, "revisions", ()): for f in getattr(rev, "files", ()): @@ -121,7 +123,11 @@ def _delete_gguf_variant_from_repos( completed_hashes: set[str] = set() for target_repo in target_repos: - repo_dir = Path(target_repo.repo_path) if getattr(target_repo, "repo_path", None) else None + repo_dir = ( + Path(target_repo.repo_path) + if getattr(target_repo, "repo_path", None) + else None + ) matched = _repo_file_matches( target_repo, lambda name: _is_main_gguf_filename(name) @@ -233,7 +239,10 @@ def _loaded_id_matches_repo(loaded_id: str, repo_id: str) -> bool: def _loaded_repo_variant_blocks_delete( - loaded_id: str, repo_id: str, delete_variant: Optional[str], loaded_variant: Optional[str] + loaded_id: str, + repo_id: str, + delete_variant: Optional[str], + loaded_variant: Optional[str], ) -> bool: if not _loaded_id_matches_repo(loaded_id, repo_id): return False @@ -256,7 +265,9 @@ def _llama_cpp_blocks_delete(repo_id: str, variant: Optional[str]) -> bool: from routes.inference import get_llama_cpp_backend backend = get_llama_cpp_backend() except Exception as e: - logger.debug(f"llama.cpp backend unavailable during delete guard for {repo_id}: {e}") + logger.debug( + f"llama.cpp backend unavailable during delete guard for {repo_id}: {e}" + ) return False loaded_id = backend.model_identifier loaded_variant = getattr(backend, "hf_variant", None) @@ -283,7 +294,9 @@ def _inference_backend_blocks_delete(repo_id: str) -> bool: from core.inference import get_inference_backend backend = get_inference_backend() except Exception as e: - logger.debug(f"Inference backend unavailable during delete guard for {repo_id}: {e}") + logger.debug( + f"Inference backend unavailable during delete guard for {repo_id}: {e}" + ) return False active_name = backend.active_model_name return bool(active_name) and _loaded_id_matches_repo(active_name, repo_id) @@ -316,7 +329,9 @@ async def delete_cached_model_response( _inference_backend_blocks_delete(repo_id) ) except Exception as e: - logger.warning(f"Load-state verification failed for {repo_id}; refusing delete: {e}") + logger.warning( + f"Load-state verification failed for {repo_id}; refusing delete: {e}" + ) raise HTTPException( status_code = 503, detail = _LOAD_STATE_UNVERIFIABLE_DETAIL, @@ -327,7 +342,9 @@ async def delete_cached_model_response( detail = "Unload the model before deleting", ) - repo_key = await asyncio.to_thread(resolve_cached_repo_id_case, repo_id, repo_type = "model") + repo_key = await asyncio.to_thread( + resolve_cached_repo_id_case, repo_id, repo_type = "model" + ) if not downloads.registry.begin_delete(repo_key, variant): detail = ( f"Cancel the {variant} download before deleting it." @@ -336,7 +353,9 @@ async def delete_cached_model_response( ) raise HTTPException(status_code = 400, detail = detail) try: - return await asyncio.to_thread(_delete_cached_model_blocking, repo_id, variant, hf_token) + return await asyncio.to_thread( + _delete_cached_model_blocking, repo_id, variant, hf_token + ) finally: downloads.registry.end_delete(repo_key, variant) cache_inventory.invalidate_hf_cache_scans() @@ -375,18 +394,22 @@ def _delete_cached_model_blocking( if not target_entries: if variant is None: - cache_purged = purge_repo_cache_dirs("model", repo_id) or purge_partial_repo( + cache_purged = purge_repo_cache_dirs( "model", repo_id + ) or purge_partial_repo("model", repo_id) + state_purged = ( + download_manifest.purge_all_state_for_repo("model", repo_id) > 0 ) - state_purged = download_manifest.purge_all_state_for_repo("model", repo_id) > 0 if cache_purged or state_purged: return {"status": "deleted", "repo_id": repo_id} if variant: - incomplete_result = gguf_variants.delete_variant_incomplete_blobs_result( - repo_id, - variant, - hf_token, - companions = not sibling_active, + incomplete_result = ( + gguf_variants.delete_variant_incomplete_blobs_result( + repo_id, + variant, + hf_token, + companions = not sibling_active, + ) ) if incomplete_result.unresolved: raise HTTPException( @@ -422,7 +445,9 @@ def _delete_cached_model_blocking( deleted_revisions = False for hf_cache, repo_info in target_entries: revision_hashes = [ - rev.commit_hash for rev in repo_info.revisions if getattr(rev, "commit_hash", None) + rev.commit_hash + for rev in repo_info.revisions + if getattr(rev, "commit_hash", None) ] if not revision_hashes: continue diff --git a/studio/backend/hub/services/models/downloads.py b/studio/backend/hub/services/models/downloads.py index db95b82c95..da0d42ed12 100644 --- a/studio/backend/hub/services/models/downloads.py +++ b/studio/backend/hub/services/models/downloads.py @@ -78,7 +78,9 @@ def _spawn_download_worker( ) -async def download_model_response(body: DownloadModelRequest, hf_token: Optional[str] = None): +async def download_model_response( + body: DownloadModelRequest, hf_token: Optional[str] = None +): """Start a background download for a HuggingFace model.""" repo_id = body.repo_id.strip() if not _is_valid_repo_id(repo_id): @@ -87,7 +89,9 @@ async def download_model_response(body: DownloadModelRequest, hf_token: Optional detail = f"Invalid repo_id: {repo_id!r}", ) # Canonicalize so two different-cased paste-ins share one job + cache dir. - repo_id = await asyncio.to_thread(resolve_cached_repo_id_case, repo_id, repo_type = "model") + repo_id = await asyncio.to_thread( + resolve_cached_repo_id_case, repo_id, repo_type = "model" + ) variant = (body.gguf_variant or "").strip() or None if variant is not None and not _is_valid_gguf_variant(variant): @@ -200,7 +204,9 @@ async def cancel_download_model_response(body: CancelDownloadRequest): status_code = 400, detail = f"Invalid repo_id: {repo_id!r}", ) - repo_id = await asyncio.to_thread(resolve_cached_repo_id_case, repo_id, repo_type = "model") + repo_id = await asyncio.to_thread( + resolve_cached_repo_id_case, repo_id, repo_type = "model" + ) variant = (body.gguf_variant or "").strip() or None if variant is not None and not _is_valid_gguf_variant(variant): raise HTTPException( @@ -219,12 +225,16 @@ async def cancel_download_model_response(body: CancelDownloadRequest): return {"job_key": key, "state": state} -async def get_download_status_response(repo_id: str, gguf_variant: str = "") -> DownloadJobStatus: +async def get_download_status_response( + repo_id: str, gguf_variant: str = "" +) -> DownloadJobStatus: """Return the latest state of a background download job.""" repo_id = repo_id.strip() if not _is_valid_repo_id(repo_id): return DownloadJobStatus(state = "idle") - repo_id = await asyncio.to_thread(resolve_cached_repo_id_case, repo_id, repo_type = "model") + repo_id = await asyncio.to_thread( + resolve_cached_repo_id_case, repo_id, repo_type = "model" + ) variant = (gguf_variant or "").strip() or None key = _download_job_key(repo_id, variant) return _job_status(key, repo_id = repo_id, variant = variant) @@ -249,7 +259,9 @@ async def get_active_downloads_response(repo_id: str = "") -> ActiveDownloadsRes ) -def _variant_transport_status(repo_id: str, variant: str, hf_token: Optional[str]) -> dict: +def _variant_transport_status( + repo_id: str, variant: str, hf_token: Optional[str] +) -> dict: incomplete_hashes = download_registry.incomplete_blob_hashes( "model", repo_id, @@ -281,13 +293,16 @@ def _variant_transport_status(repo_id: str, variant: str, hf_token: Optional[str variant, ) has_matching_incomplete = bool( - incomplete_hashes and variant_hashes and incomplete_hashes.intersection(variant_hashes) + incomplete_hashes + and variant_hashes + and incomplete_hashes.intersection(variant_hashes) ) return { "has_partial": has_partial, "last_transport": last_transport, "resumable": ( - has_matching_incomplete and last_transport == download_registry.TRANSPORT_HTTP + has_matching_incomplete + and last_transport == download_registry.TRANSPORT_HTTP ), } @@ -315,7 +330,9 @@ async def get_model_transport_status_response( return _variant_transport_status(repo_id, variant, hf_token) return { "has_partial": has_active_incomplete_blobs("model", repo_id), - "last_transport": download_registry.read_active_transport_marker("model", repo_id), + "last_transport": download_registry.read_active_transport_marker( + "model", repo_id + ), "resumable": download_registry.is_resumable_partial("model", repo_id), } @@ -359,7 +376,9 @@ async def get_gguf_download_progress_response( if manifest is not None: return ( sum(max(0, int(file.size or 0)) for file in manifest.expected_files), - frozenset(file.sha256 for file in manifest.expected_files if file.sha256), + frozenset( + file.sha256 for file in manifest.expected_files if file.sha256 + ), ) return ( expected_total, diff --git a/studio/backend/hub/services/models/folder_browser.py b/studio/backend/hub/services/models/folder_browser.py index 9b0b46509b..0132790500 100644 --- a/studio/backend/hub/services/models/folder_browser.py +++ b/studio/backend/hub/services/models/folder_browser.py @@ -267,7 +267,12 @@ def _browse_relative_parts(requested_path: str, root: Path) -> Optional[list[str parts = [part for part in rel_text.split(os.sep) if part not in ("", ".")] altsep = os.altsep for part in parts: - if part == ".." or "\x00" in part or os.sep in part or (altsep and altsep in part): + if ( + part == ".." + or "\x00" in part + or os.sep in part + or (altsep and altsep in part) + ): return None return parts @@ -460,7 +465,9 @@ def browse_folders_response( # Parent is None at the FS root and when it would step outside the sandbox, # so the up-row never 403s on click. parent: Optional[str] - if target.parent == target or not _is_path_inside_allowlist(target.parent, allowed_roots): + if target.parent == target or not _is_path_inside_allowlist( + target.parent, allowed_roots + ): parent = None else: parent = str(target.parent) diff --git a/studio/backend/hub/services/models/gguf_variants.py b/studio/backend/hub/services/models/gguf_variants.py index 0b4671c6d2..0d9d02150d 100644 --- a/studio/backend/hub/services/models/gguf_variants.py +++ b/studio/backend/hub/services/models/gguf_variants.py @@ -49,9 +49,7 @@ from hub.utils.gguf_plan import ( logger = get_logger(__name__) -_VARIANT_HASH_CACHE: "OrderedDict[tuple[str, str, str, bool], tuple[frozenset[str], float]]" = ( - OrderedDict() -) +_VARIANT_HASH_CACHE: "OrderedDict[tuple[str, str, str, bool], tuple[frozenset[str], float]]" = OrderedDict() _VARIANT_REQUIREMENT_CACHE: "OrderedDict[tuple[str, str, str], tuple[_GgufVariantRequirement, float]]" = OrderedDict() _VARIANT_REQUIREMENT_NEG_CACHE: "OrderedDict[tuple[str, str], float]" = OrderedDict() _VARIANT_HASH_MAX = 512 @@ -120,7 +118,9 @@ def _variant_requirement_neg_cache_clear(key: tuple[str, str]) -> None: _VARIANT_REQUIREMENT_NEG_CACHE.pop(key, None) -def _variant_hash_cache_get(key: tuple[str, str, str, bool]) -> Optional[frozenset[str]]: +def _variant_hash_cache_get( + key: tuple[str, str, str, bool], +) -> Optional[frozenset[str]]: with _VARIANT_HASH_LOCK: cached = _VARIANT_HASH_CACHE.get(key) if cached is None: @@ -133,7 +133,9 @@ def _variant_hash_cache_get(key: tuple[str, str, str, bool]) -> Optional[frozens return hashes -def _variant_hash_cache_set(key: tuple[str, str, str, bool], hashes: frozenset[str]) -> None: +def _variant_hash_cache_set( + key: tuple[str, str, str, bool], hashes: frozenset[str] +) -> None: with _VARIANT_HASH_LOCK: _VARIANT_HASH_CACHE[key] = (hashes, time.monotonic()) _VARIANT_HASH_CACHE.move_to_end(key) @@ -141,7 +143,9 @@ def _variant_hash_cache_set(key: tuple[str, str, str, bool], hashes: frozenset[s _VARIANT_HASH_CACHE.popitem(last = False) -def _variant_requirement_cache_get(key: tuple[str, str, str]) -> Optional[_GgufVariantRequirement]: +def _variant_requirement_cache_get( + key: tuple[str, str, str], +) -> Optional[_GgufVariantRequirement]: with _VARIANT_HASH_LOCK: cached = _VARIANT_REQUIREMENT_CACHE.get(key) if cached is None: @@ -155,7 +159,9 @@ def _variant_requirement_cache_get(key: tuple[str, str, str]) -> Optional[_GgufV def _variant_requirement_cache_set_many( - repo_id: str, hf_token: Optional[str], requirements: dict[str, _GgufVariantRequirement] + repo_id: str, + hf_token: Optional[str], + requirements: dict[str, _GgufVariantRequirement], ) -> None: with _VARIANT_HASH_LOCK: now = time.monotonic() @@ -167,7 +173,9 @@ def _variant_requirement_cache_set_many( _VARIANT_REQUIREMENT_CACHE.popitem(last = False) -def _build_gguf_variant_requirements(siblings: list) -> dict[str, _GgufVariantRequirement]: +def _build_gguf_variant_requirements( + siblings: list, +) -> dict[str, _GgufVariantRequirement]: return build_gguf_variant_plans(siblings) @@ -278,7 +286,11 @@ def gguf_variant_blob_hashes( if requirement is None and allow_remote: requirement = gguf_variant_requirements(repo_id, variant, hf_token) if requirement is not None: - hashes = requirement.required_hashes if include_companions else requirement.main_hashes + hashes = ( + requirement.required_hashes + if include_companions + else requirement.main_hashes + ) if hashes: _variant_hash_cache_set(key, hashes) return hashes @@ -300,7 +312,9 @@ def delete_variant_incomplete_blobs_result( # With a sibling still downloading, ``companions=False`` keeps a shared mmproj # from being unlinked out from under it; the repo's last delete reclaims it. target_hashes = ( - gguf_variant_blob_hashes(repo_id, variant, hf_token, include_companions = companions) + gguf_variant_blob_hashes( + repo_id, variant, hf_token, include_companions = companions + ) | extra_hashes ) if not target_hashes: @@ -310,7 +324,9 @@ def delete_variant_incomplete_blobs_result( incomplete_blob_hashes = set(), variant_blob_hashes = frozenset(), ) - has_repo_partials = bool(download_registry.incomplete_blob_hashes("model", repo_id)) + has_repo_partials = bool( + download_registry.incomplete_blob_hashes("model", repo_id) + ) return VariantIncompleteDeleteResult( deleted = 0, unresolved = has_variant_partial_state and has_repo_partials, @@ -440,7 +456,9 @@ async def get_gguf_variants_response( ) try: - variants, has_vision, siblings = list_gguf_variants(repo_id, hf_token = hf_token) + variants, has_vision, siblings = list_gguf_variants( + repo_id, hf_token = hf_token + ) except Exception: cached = list_gguf_variants_from_hf_cache(repo_id) if cached is not None: @@ -555,17 +573,25 @@ async def get_gguf_variants_response( partial_quants: set[str] = set() partial_quant_transports: dict[str, Optional[str]] = {} try: - incomplete_hashes = download_registry.incomplete_blob_hashes("model", repo_id) + incomplete_hashes = download_registry.incomplete_blob_hashes( + "model", repo_id + ) except Exception as e: - logger.warning(f"Failed to compute partial GGUF variants for {repo_id}: {e}") + logger.warning( + f"Failed to compute partial GGUF variants for {repo_id}: {e}" + ) incomplete_hashes = set() - scan_snapshot_dir = hf_cache_scan.resolve_snapshot_dir_for_scan("model", repo_id) + scan_snapshot_dir = hf_cache_scan.resolve_snapshot_dir_for_scan( + "model", repo_id + ) # Manifest + marker + main incomplete-blob check: catches variants whose # download was cancelled or whose expected shards are missing/undersized. for variant in variants: try: requirement = requirements_by_quant.get(variant.quant.lower()) - variant_hashes = requirement.main_hashes if requirement is not None else None + variant_hashes = ( + requirement.main_hashes if requirement is not None else None + ) if variant_hashes is None and incomplete_hashes: variant_hashes = gguf_variant_blob_hashes( repo_id, @@ -581,13 +607,16 @@ async def get_gguf_variants_response( variant_blob_hashes = variant_hashes, ): partial_quants.add(variant.quant) - partial_quant_transports[variant.quant] = _partial_transport_for_variant( - repo_id, - variant.quant, + partial_quant_transports[variant.quant] = ( + _partial_transport_for_variant( + repo_id, + variant.quant, + ) ) except Exception as e: logger.warning( - f"Manifest-based partial check failed for " f"{repo_id}/{variant.quant}: {e}" + f"Manifest-based partial check failed for " + f"{repo_id}/{variant.quant}: {e}" ) if incomplete_hashes: for variant in variants: @@ -597,7 +626,8 @@ async def get_gguf_variants_response( # companion_hashes adds the MTP drafter (mmproj_hashes covers # every mmproj precision in the repo, not just the planned one). if ( - (requirement.mmproj_hashes | requirement.companion_hashes) & incomplete_hashes + (requirement.mmproj_hashes | requirement.companion_hashes) + & incomplete_hashes ) and _filenames_cached( requirement.main_filenames, requirement.main_size_bytes, @@ -617,11 +647,15 @@ async def get_gguf_variants_response( display_label = v.display_label, size_bytes = v.size_bytes, download_size_bytes = ( - requirement.download_size_bytes if requirement is not None else v.size_bytes + requirement.download_size_bytes + if requirement is not None + else v.size_bytes ), downloaded = _is_fully_downloaded(v) and not is_partial, partial = is_partial, - partial_transport = (partial_quant_transports.get(v.quant) if is_partial else None), + partial_transport = ( + partial_quant_transports.get(v.quant) if is_partial else None + ), ) return GgufVariantsResponse( diff --git a/studio/backend/hub/services/models/local_inventory.py b/studio/backend/hub/services/models/local_inventory.py index 6e6a28b919..d19f774b6b 100644 --- a/studio/backend/hub/services/models/local_inventory.py +++ b/studio/backend/hub/services/models/local_inventory.py @@ -98,7 +98,9 @@ def _is_model_directory_for_scan(path: Path, *, entry_limit: int | None) -> bool if entry_limit is None: return _is_model_directory(path) try: - has_config = (path / "config.json").exists() or (path / "adapter_config.json").exists() + has_config = (path / "config.json").exists() or ( + path / "adapter_config.json" + ).exists() except OSError: return False return has_config and _has_immediate_model_weight(path) @@ -151,7 +153,9 @@ def _scan_models_dir( break try: is_dir = child.is_dir() - is_gguf_file = not is_dir and child.suffix.lower() == ".gguf" and child.is_file() + is_gguf_file = ( + not is_dir and child.suffix.lower() == ".gguf" and child.is_file() + ) if not is_dir and not is_gguf_file: continue has_model_files = is_gguf_file or _has_immediate_model_signal(child) @@ -190,7 +194,9 @@ def _hf_repo_dir_has_content(repo_dir: Path) -> bool: return False -def _scan_hf_cache(cache_dir: Path, *, entry_limit: int | None = None) -> List[LocalModelInfo]: +def _scan_hf_cache( + cache_dir: Path, *, entry_limit: int | None = None +) -> List[LocalModelInfo]: if not cache_dir.exists() or not cache_dir.is_dir(): return [] @@ -228,7 +234,9 @@ def _scan_hf_cache(cache_dir: Path, *, entry_limit: int | None = None) -> List[L repo_dir, ) gguf_partial = hf_cache_scan.is_gguf_repo_partial(model_id, repo_dir) - has_gguf_variant_state, gguf_variant_state_size = _gguf_variant_state_summary(model_id) + has_gguf_variant_state, gguf_variant_state_size = _gguf_variant_state_summary( + model_id + ) snapshot_partial_transport = ( hf_cache_scan.partial_transport_for( "model", @@ -311,7 +319,9 @@ def _scan_hf_cache(cache_dir: Path, *, entry_limit: int | None = None) -> List[L return found -def _scan_lmstudio_dir(lm_dir: Path, *, entry_limit: int | None = None) -> List[LocalModelInfo]: +def _scan_lmstudio_dir( + lm_dir: Path, *, entry_limit: int | None = None +) -> List[LocalModelInfo]: """Scan an LM Studio models dir (``publisher/model-name`` folders of GGUFs, or top-level standalone GGUFs).""" if not lm_dir.exists() or not lm_dir.is_dir(): return [] @@ -427,7 +437,9 @@ def _resolve_allowed_models_dir(models_dir: str, allowed_roots: list[Path]) -> P if not models_dir or not models_dir.strip(): raise ValueError("Directory not allowed") - requested = Path(os.path.realpath(os.path.expanduser(normalize_path(models_dir.strip())))) + requested = Path( + os.path.realpath(os.path.expanduser(normalize_path(models_dir.strip()))) + ) if any(path_is_same_or_child(requested, root) for root in allowed_roots): return requested @@ -510,7 +522,9 @@ async def _collect_models_from_default_sources( and hf_default.resolve() != hf_cache_dir.resolve() and hf_default.resolve() != legacy_hf.resolve() ): - local_models += await _scan_source("default HF cache", _scan_hf_cache, hf_default) + local_models += await _scan_source( + "default HF cache", _scan_hf_cache, hf_default + ) for lm_dir in lm_dirs: local_models += await _scan_source("LM Studio", _scan_lmstudio_dir, lm_dir) @@ -612,7 +626,9 @@ def _dedupe_local_models(local_models: List[LocalModelInfo]) -> list[LocalModelI ) -async def list_local_models_response(models_dir: str = "./models") -> LocalModelListResponse: +async def list_local_models_response( + models_dir: str = "./models", +) -> LocalModelListResponse: """List local model candidates from every supported on-device source.""" hf_cache_dir = _resolve_hf_cache_dir() legacy_hf = legacy_hf_cache_dir() diff --git a/studio/backend/hub/services/models/ollama.py b/studio/backend/hub/services/models/ollama.py index 96a4114620..aed76be7a1 100644 --- a/studio/backend/hub/services/models/ollama.py +++ b/studio/backend/hub/services/models/ollama.py @@ -124,7 +124,9 @@ def _ollama_links_dir(ollama_dir: Path) -> Optional[Path]: return None -def _make_ollama_blob_link(link_dir: Path, link_name: str, target: Path) -> Optional[str]: +def _make_ollama_blob_link( + link_dir: Path, link_name: str, target: Path +) -> Optional[str]: """Create a .gguf-named link to an Ollama blob: tries symlink then hardlink, skips the model if neither works (a full multi-GB copy would block the API). Idempotent.""" try: link_dir.mkdir(parents = True, exist_ok = True) @@ -137,7 +139,9 @@ def _make_ollama_blob_link(link_dir: Path, link_name: str, target: Path) -> Opti return None link_path = _contained_link_path(link_dir, link_name) if link_path is None: - logger.warning("Refusing unsafe Ollama link name %r under %s", link_name, link_dir) + logger.warning( + "Refusing unsafe Ollama link name %r under %s", link_name, link_dir + ) return None try: resolved = target.resolve() @@ -232,7 +236,9 @@ def _ollama_model_info_from_manifest( model_type = cfg.get("model_type", "") file_type = cfg.get("file_type", "") except (json.JSONDecodeError, OSError) as e: - logger.debug("Could not parse Ollama config blob %s: %s", config_blob, e) + logger.debug( + "Could not parse Ollama config blob %s: %s", config_blob, e + ) layers = manifest.get("layers") or [] if not isinstance(layers, list): @@ -260,11 +266,17 @@ def _ollama_model_info_from_manifest( model_blob = candidate if materialize_links and model_link_dir is not None: link_name = f"{safe_name}-{tag}{quant}.gguf" - gguf_link_path = _make_ollama_blob_link(model_link_dir, link_name, candidate) + gguf_link_path = _make_ollama_blob_link( + model_link_dir, link_name, candidate + ) elif materialize_links and media == "application/vnd.ollama.image.projector": candidate = _ollama_blob_path(blobs_dir, digest) - if candidate is not None and _safe_is_file(candidate) and model_link_dir is not None: + if ( + candidate is not None + and _safe_is_file(candidate) + and model_link_dir is not None + ): mmproj_name = f"{safe_name}-{tag}-mmproj.gguf" _make_ollama_blob_link(model_link_dir, mmproj_name, candidate) diff --git a/studio/backend/hub/tests/test_dataset_services.py b/studio/backend/hub/tests/test_dataset_services.py index 4890714cd0..d0705be28e 100644 --- a/studio/backend/hub/tests/test_dataset_services.py +++ b/studio/backend/hub/tests/test_dataset_services.py @@ -198,7 +198,9 @@ def test_delete_cached_dataset_absent_everywhere_raises_404(monkeypatch): def test_check_format_rejects_invalid_path_as_400(): with pytest.raises(HTTPException) as exc_info: - formatting.check_format_response(CheckFormatRequest(dataset_name = "../../etc/passwd")) + formatting.check_format_response( + CheckFormatRequest(dataset_name = "../../etc/passwd") + ) assert exc_info.value.status_code == 400 @@ -286,7 +288,9 @@ def test_dataset_claim_register_cancel_uses_registry_marker_owner(monkeypatch): ) result = asyncio.run( - downloads.download_dataset_response(SimpleNamespace(repo_id = "Org/Data", use_xet = False)) + downloads.download_dataset_response( + SimpleNamespace(repo_id = "Org/Data", use_xet = False) + ) ) assert result["state"] == "cancelled" @@ -328,7 +332,9 @@ def test_upload_dataset_response_writes_non_empty_file(monkeypatch, tmp_path): payload = b'{"text":"hello"}\n' monkeypatch.setattr(local, "DATASET_UPLOAD_DIR", tmp_path) - response = asyncio.run(local.upload_dataset_response(_Upload("../train.jsonl", payload))) + response = asyncio.run( + local.upload_dataset_response(_Upload("../train.jsonl", payload)) + ) stored_path = Path(response.stored_path) assert response.filename == "train.jsonl" diff --git a/studio/backend/hub/tests/test_model_services.py b/studio/backend/hub/tests/test_model_services.py index 44ab4b80d0..4130e9f517 100644 --- a/studio/backend/hub/tests/test_model_services.py +++ b/studio/backend/hub/tests/test_model_services.py @@ -184,7 +184,9 @@ def test_make_ollama_blob_link_refuses_escaping_name(tmp_path): blob.parent.mkdir(parents = True) blob.write_bytes(b"weights") - escaped = ollama._make_ollama_blob_link(link_dir, "model-tag-../../../pwned.gguf", blob) + escaped = ollama._make_ollama_blob_link( + link_dir, "model-tag-../../../pwned.gguf", blob + ) assert escaped is None assert not list(tmp_path.rglob("pwned.gguf")) @@ -200,7 +202,9 @@ def test_cached_gguf_scan_dedupes_and_excludes_mmproj_only(monkeypatch, tmp_path [_file("Q4_K_M.gguf", 300), _file("Q8_0.gguf", 200)], tmp_path / "large", ) - mmproj_only = _repo("Org/VisionAdapter", [_file("mmproj-F16.gguf", 900)], tmp_path / "mmproj") + mmproj_only = _repo( + "Org/VisionAdapter", [_file("mmproj-F16.gguf", 900)], tmp_path / "mmproj" + ) monkeypatch.setattr( cache_inventory, "all_hf_cache_scans", @@ -241,7 +245,9 @@ def test_cached_gguf_scan_preserves_partial_flag(monkeypatch, tmp_path): assert row["capabilities"]["can_chat"] is False -def test_cached_gguf_scan_includes_variant_state_without_completed_gguf(monkeypatch, tmp_path): +def test_cached_gguf_scan_includes_variant_state_without_completed_gguf( + monkeypatch, tmp_path +): monkeypatch.setattr(state_dir, "cache_root", lambda: tmp_path / "state") repo_path = tmp_path / "hub" / "models--Org--PartialGguf" repo_path.mkdir(parents = True) @@ -257,7 +263,9 @@ def test_cached_gguf_scan_includes_variant_state_without_completed_gguf(monkeypa [download_manifest.ExpectedFile(path = "model-Q4_K_M.gguf", size = 4096)], "http", ) - assert download_manifest.write_cancel_marker("model", "Org/PartialGguf", "Q4_K_M", "http") + assert download_manifest.write_cancel_marker( + "model", "Org/PartialGguf", "Q4_K_M", "http" + ) monkeypatch.setattr( cache_inventory, "all_hf_cache_scans", @@ -366,7 +374,9 @@ def test_gguf_variant_blob_hashes_skip_missing_rfilename(monkeypatch): monkeypatch.setattr( gguf_variants, "_fetch_gguf_variant_requirements", - lambda _repo_id, _hf_token = None: gguf_variants._build_gguf_variant_requirements(siblings), + lambda _repo_id, _hf_token = None: gguf_variants._build_gguf_variant_requirements( + siblings + ), ) result = gguf_variants.gguf_variant_blob_hashes("Org/Malformed", "Q4_K_M", None) @@ -410,7 +420,9 @@ def test_download_gguf_variant_purges_only_main_quant_hashes(monkeypatch, tmp_pa ), ) monkeypatch.setattr( - hf_download, "_verify_completed_download", lambda *args, **kwargs: verified.append(args) + hf_download, + "_verify_completed_download", + lambda *args, **kwargs: verified.append(args), ) monkeypatch.setattr( download_registry, @@ -425,7 +437,8 @@ def test_download_gguf_variant_purges_only_main_quant_hashes(monkeypatch, tmp_pa sys.modules, "huggingface_hub", SimpleNamespace( - snapshot_download = lambda **kwargs: snapshot_calls.append(kwargs) or str(tmp_path) + snapshot_download = lambda **kwargs: snapshot_calls.append(kwargs) + or str(tmp_path) ), ) @@ -441,12 +454,20 @@ def test_download_gguf_variant_purges_only_main_quant_hashes(monkeypatch, tmp_pa }, ) ] - assert [file.path for file in written[0][3]] == ["model-Q4_K_M.gguf", "mmproj-F16.gguf"] - assert snapshot_calls[0]["allow_patterns"] == ["model-Q4_K_M.gguf", "mmproj-F16.gguf"] + assert [file.path for file in written[0][3]] == [ + "model-Q4_K_M.gguf", + "mmproj-F16.gguf", + ] + assert snapshot_calls[0]["allow_patterns"] == [ + "model-Q4_K_M.gguf", + "mmproj-F16.gguf", + ] assert verified == [("model", "Org/Vision", "Q4_K_M", str(tmp_path))] -def test_download_gguf_variant_manifest_resume_purges_only_main_quant_hashes(monkeypatch, tmp_path): +def test_download_gguf_variant_manifest_resume_purges_only_main_quant_hashes( + monkeypatch, tmp_path +): prepare_calls = [] snapshot_calls = [] @@ -484,12 +505,15 @@ def test_download_gguf_variant_manifest_resume_purges_only_main_quant_hashes(mon "prepare_cache_for_transport", lambda *args, **kwargs: prepare_calls.append((args, kwargs)) or 0, ) - monkeypatch.setattr(hf_download, "_verify_completed_download", lambda *_args, **_kwargs: None) + monkeypatch.setattr( + hf_download, "_verify_completed_download", lambda *_args, **_kwargs: None + ) monkeypatch.setitem( sys.modules, "huggingface_hub", SimpleNamespace( - snapshot_download = lambda **kwargs: snapshot_calls.append(kwargs) or str(tmp_path) + snapshot_download = lambda **kwargs: snapshot_calls.append(kwargs) + or str(tmp_path) ), ) @@ -505,10 +529,15 @@ def test_download_gguf_variant_manifest_resume_purges_only_main_quant_hashes(mon }, ) ] - assert snapshot_calls[0]["allow_patterns"] == ["model-Q4_K_M.gguf", "mmproj-F16.gguf"] + assert snapshot_calls[0]["allow_patterns"] == [ + "model-Q4_K_M.gguf", + "mmproj-F16.gguf", + ] -def test_download_snapshot_recovers_manifest_after_metadata_fallback(monkeypatch, tmp_path): +def test_download_snapshot_recovers_manifest_after_metadata_fallback( + monkeypatch, tmp_path +): metadata_calls = [] written = [] cleared = [] @@ -518,11 +547,15 @@ def test_download_snapshot_recovers_manifest_after_metadata_fallback(monkeypatch metadata_calls.append(True) if len(metadata_calls) == 1: raise RuntimeError("metadata down") - return SimpleNamespace(siblings = [SimpleNamespace(rfilename = "config.json", size = 12)]) + return SimpleNamespace( + siblings = [SimpleNamespace(rfilename = "config.json", size = 12)] + ) monkeypatch.setattr(hf_download, "_model_info_with_retry", _metadata) monkeypatch.setattr( - hf_download, "_verify_completed_download", lambda *args, **kwargs: verified.append(args) + hf_download, + "_verify_completed_download", + lambda *args, **kwargs: verified.append(args), ) monkeypatch.setattr( download_registry, "prepare_cache_for_transport", lambda *_args, **_kwargs: 0 @@ -561,7 +594,9 @@ def test_download_dataset_continues_without_metadata_manifest(monkeypatch, tmp_p monkeypatch.setattr(hf_download, "_dataset_info_with_retry", _metadata) monkeypatch.setattr( - hf_download, "_verify_completed_download", lambda *args, **kwargs: verified.append(args) + hf_download, + "_verify_completed_download", + lambda *args, **kwargs: verified.append(args), ) monkeypatch.setattr( download_registry, "prepare_cache_for_transport", lambda *_args, **_kwargs: 0 @@ -579,7 +614,8 @@ def test_download_dataset_continues_without_metadata_manifest(monkeypatch, tmp_p sys.modules, "huggingface_hub", SimpleNamespace( - snapshot_download = lambda **kwargs: snapshot_calls.append(kwargs) or str(tmp_path) + snapshot_download = lambda **kwargs: snapshot_calls.append(kwargs) + or str(tmp_path) ), ) @@ -613,13 +649,17 @@ def test_download_snapshot_fails_when_metadata_unavailable_and_partial_remains( monkeypatch.setattr(hf_download, "_model_info_with_retry", _metadata) monkeypatch.setattr( - hf_download, "_verify_completed_download", lambda *args, **kwargs: verified.append(args) + hf_download, + "_verify_completed_download", + lambda *args, **kwargs: verified.append(args), ) monkeypatch.setattr( download_registry, "prepare_cache_for_transport", lambda *_args, **_kwargs: 0 ) monkeypatch.setattr(download_manifest, "clear_cancel_marker", lambda *_args: None) - monkeypatch.setattr(download_manifest, "read_manifest", lambda *_args, **_kwargs: None) + monkeypatch.setattr( + download_manifest, "read_manifest", lambda *_args, **_kwargs: None + ) monkeypatch.setattr( download_manifest, "write_manifest", lambda *args: written.append(args) or True ) @@ -719,7 +759,9 @@ def test_gguf_download_progress_fallback_logs_warning(monkeypatch): assert kwargs == {} -def test_gguf_progress_counts_completed_mmproj_with_expected_bytes(monkeypatch, tmp_path): +def test_gguf_progress_counts_completed_mmproj_with_expected_bytes( + monkeypatch, tmp_path +): """A finished mmproj companion keeps counting toward progress once the caller supplies expected bytes; resolving the variant requirement credits it.""" entry = tmp_path / "models--Org--Model-GGUF" @@ -1138,7 +1180,9 @@ def test_gguf_progress_scoped_hashes_exclude_sibling_quant(monkeypatch, tmp_path assert result["downloaded_bytes"] == 5 -def test_gguf_progress_unknown_hashes_does_not_count_foreign_blobs(monkeypatch, tmp_path): +def test_gguf_progress_unknown_hashes_does_not_count_foreign_blobs( + monkeypatch, tmp_path +): # With a variant's hashes unresolved (metadata flaked, no manifest), the # shared blobs/ dir's FINALIZED blobs must NOT be counted wholesale: a cached # sibling quant (``siblinghash``) alongside is the "instant ~900 MB" bug. @@ -1190,7 +1234,9 @@ def test_gguf_progress_unknown_hashes_does_not_count_foreign_blobs(monkeypatch, assert result["complete_on_disk"] is False -def test_gguf_progress_unknown_hashes_drops_unscoped_incomplete_blob(monkeypatch, tmp_path): +def test_gguf_progress_unknown_hashes_drops_unscoped_incomplete_blob( + monkeypatch, tmp_path +): # With hashes unresolved, an .incomplete in the shared blobs/ dir can't be # attributed to this variant (it may be a concurrent sibling's active write), # so it is dropped, mirroring the finalized-blob guard. In production the @@ -1239,7 +1285,9 @@ def test_gguf_progress_unknown_hashes_drops_unscoped_incomplete_blob(monkeypatch assert result["completed_bytes"] == 0 # finalized sibling still ignored -def test_gguf_progress_unknown_hashes_no_backward_dip_when_variant_finalizes(monkeypatch, tmp_path): +def test_gguf_progress_unknown_hashes_no_backward_dip_when_variant_finalizes( + monkeypatch, tmp_path +): # Regression for the two-variant dip: with hashes unresolved, the first quant # finalizes while the sibling still writes its .incomplete. The sibling's # bytes used to leak into this numerator, dipping the bar ~99% -> ~78% for @@ -1310,7 +1358,9 @@ def test_hf_cache_model_file_probe_is_bounded(monkeypatch, tmp_path): model.write_bytes(b"weights") entries = [first, second, model] - monkeypatch.setattr(model_common.Path, "rglob", lambda _self, _pattern: iter(entries)) + monkeypatch.setattr( + model_common.Path, "rglob", lambda _self, _pattern: iter(entries) + ) monkeypatch.setattr(model_common, "_HF_CACHE_MODEL_FILE_PROBE_LIMIT", 2) bounded = model_common._iter_hf_cache_model_files(snapshot) @@ -1333,7 +1383,9 @@ def test_download_state_lookup_is_repo_case_insensitive(monkeypatch, tmp_path): None, [download_manifest.ExpectedFile(path = "config.json", size = 12)], ) - assert download_manifest.write_cancel_marker("model", "Owner/Repo", "Q4_K_M", "http") + assert download_manifest.write_cancel_marker( + "model", "Owner/Repo", "Q4_K_M", "http" + ) manifest = download_manifest.read_manifest("model", "owner/repo", None) @@ -1366,7 +1418,9 @@ def test_hf_cache_scan_fallback_row_uses_local_model_info_alias(monkeypatch, tmp blobs_dir = repo_dir / "blobs" blobs_dir.mkdir(parents = True) (blobs_dir / "blob").write_bytes(b"content") - monkeypatch.setattr(local_inventory, "_classify_local_path", lambda *_args, **_kwargs: []) + monkeypatch.setattr( + local_inventory, "_classify_local_path", lambda *_args, **_kwargs: [] + ) monkeypatch.setattr( local_inventory.hf_cache_scan, "is_snapshot_partial", @@ -1405,8 +1459,12 @@ def test_hf_cache_scan_uses_gguf_partial_row_for_variant_state(monkeypatch, tmp_ [download_manifest.ExpectedFile(path = "model-Q4_K_M.gguf", size = 8192)], "http", ) - assert download_manifest.write_cancel_marker("model", "Org/PartialGguf", "Q4_K_M", "http") - monkeypatch.setattr(local_inventory, "_classify_local_path", lambda *_args, **_kwargs: []) + assert download_manifest.write_cancel_marker( + "model", "Org/PartialGguf", "Q4_K_M", "http" + ) + monkeypatch.setattr( + local_inventory, "_classify_local_path", lambda *_args, **_kwargs: [] + ) monkeypatch.setattr( local_inventory.hf_cache_scan, "is_snapshot_partial", @@ -1443,7 +1501,9 @@ def test_model_download_job_helpers_preserve_idle_shape(): assert status.error is None -def test_gguf_repo_partial_treats_completed_disk_variant_as_clean(monkeypatch, tmp_path): +def test_gguf_repo_partial_treats_completed_disk_variant_as_clean( + monkeypatch, tmp_path +): monkeypatch.setattr(state_dir, "cache_root", lambda: tmp_path / "state") snapshot = tmp_path / "cache" / "models--Org--Repo" / "snapshots" / "abc" snapshot.mkdir(parents = True) @@ -1545,13 +1605,17 @@ def test_variant_partial_accepts_variant_filtered_legacy_hashes(monkeypatch, tmp ) -def test_gguf_variants_partial_marker_overrides_size_only_downloaded(monkeypatch, tmp_path): +def test_gguf_variants_partial_marker_overrides_size_only_downloaded( + monkeypatch, tmp_path +): async def _run_inline(fn, *args, **kwargs): return fn(*args, **kwargs) monkeypatch.setattr(state_dir, "cache_root", lambda: tmp_path / "state") monkeypatch.setattr(gguf_variants.asyncio, "to_thread", _run_inline) - assert download_manifest.write_cancel_marker("model", "Org/PartialRepo", "Q4_K_M", "http") + assert download_manifest.write_cancel_marker( + "model", "Org/PartialRepo", "Q4_K_M", "http" + ) snapshot = tmp_path / "cache" / "models--Org--PartialRepo" / "snapshots" / "rev0" snapshot.mkdir(parents = True) (snapshot / "model-Q4_K_M.gguf").write_bytes(b"x" * 100) @@ -1881,7 +1945,9 @@ def test_finalize_worker_exit_never_kills_a_healthy_worker(monkeypatch, tmp_path ) -def test_prepare_cache_for_transport_purges_only_requested_hashes(monkeypatch, tmp_path): +def test_prepare_cache_for_transport_purges_only_requested_hashes( + monkeypatch, tmp_path +): root = tmp_path / "hub" blobs = root / "models--Org--Repo" / "blobs" blobs.mkdir(parents = True) @@ -1910,7 +1976,9 @@ def _vision_cache_root(monkeypatch, tmp_path): return blobs -def test_prepare_cache_for_transport_purges_cross_transport_companion(monkeypatch, tmp_path): +def test_prepare_cache_for_transport_purges_cross_transport_companion( + monkeypatch, tmp_path +): blobs = _vision_cache_root(monkeypatch, tmp_path) companion = frozenset({"shared-mmproj"}) @@ -1940,7 +2008,9 @@ def test_prepare_cache_for_transport_purges_cross_transport_companion(monkeypatc assert not (blobs / "shared-mmproj.incomplete").exists() -def test_prepare_cache_for_transport_preserves_same_transport_companion(monkeypatch, tmp_path): +def test_prepare_cache_for_transport_preserves_same_transport_companion( + monkeypatch, tmp_path +): blobs = _vision_cache_root(monkeypatch, tmp_path) companion = frozenset({"shared-mmproj"}) @@ -1995,7 +2065,9 @@ def test_prepare_cache_for_transport_protects_peer_companion(monkeypatch, tmp_pa assert (blobs / "shared-mmproj.incomplete").exists() -def test_model_download_records_completed_baseline_for_new_gguf_variant(monkeypatch, tmp_path): +def test_model_download_records_completed_baseline_for_new_gguf_variant( + monkeypatch, tmp_path +): async def _run_inline(fn, *args, **kwargs): return fn(*args, **kwargs) @@ -2010,7 +2082,9 @@ def test_model_download_records_completed_baseline_for_new_gguf_variant(monkeypa downloads.gguf_variants, "gguf_variant_blob_hashes", lambda _repo, _variant, _token = None, include_companions = True, **_kwargs: ( - frozenset({"mainhash", "mmprojhash"}) if include_companions else frozenset({"mainhash"}) + frozenset({"mainhash", "mmprojhash"}) + if include_companions + else frozenset({"mainhash"}) ), ) monkeypatch.setattr( @@ -2053,7 +2127,9 @@ def test_model_download_records_completed_baseline_for_new_gguf_variant(monkeypa registry = _Registry() monkeypatch.setattr(downloads, "_registry", registry) - monkeypatch.setattr(downloads, "_spawn_download_worker", lambda *_args, **_kwargs: _Proc()) + monkeypatch.setattr( + downloads, "_spawn_download_worker", lambda *_args, **_kwargs: _Proc() + ) asyncio.run( downloads.download_model_response( @@ -2062,7 +2138,9 @@ def test_model_download_records_completed_baseline_for_new_gguf_variant(monkeypa ) assert registry.claim_kwargs["blob_hashes"] == frozenset({"mainhash"}) - assert registry.claim_kwargs["progress_blob_hashes"] == frozenset({"mainhash", "mmprojhash"}) + assert registry.claim_kwargs["progress_blob_hashes"] == frozenset( + {"mainhash", "mmprojhash"} + ) assert registry.claim_kwargs["completed_baseline_bytes"] == 30 @@ -2096,7 +2174,9 @@ def test_gguf_model_download_skips_completed_baseline_for_variant_resume_state( downloads.gguf_variants, "gguf_variant_blob_hashes", lambda _repo, _variant, _token = None, include_companions = True, **_kwargs: ( - frozenset({"mainhash", "mmprojhash"}) if include_companions else frozenset({"mainhash"}) + frozenset({"mainhash", "mmprojhash"}) + if include_companions + else frozenset({"mainhash"}) ), ) monkeypatch.setattr( @@ -2139,7 +2219,9 @@ def test_gguf_model_download_skips_completed_baseline_for_variant_resume_state( registry = _Registry() monkeypatch.setattr(downloads, "_registry", registry) - monkeypatch.setattr(downloads, "_spawn_download_worker", lambda *_args, **_kwargs: _Proc()) + monkeypatch.setattr( + downloads, "_spawn_download_worker", lambda *_args, **_kwargs: _Proc() + ) asyncio.run( downloads.download_model_response( @@ -2153,7 +2235,9 @@ def test_gguf_model_download_skips_completed_baseline_for_variant_resume_state( def test_model_idle_status_uses_cancel_marker_after_restart(monkeypatch, tmp_path): monkeypatch.setattr(state_dir, "cache_root", lambda: tmp_path) monkeypatch.setattr(downloads, "_registry", download_registry.DownloadRegistry()) - assert download_manifest.write_cancel_marker("model", "Owner/Repo", "Q4_K_M", "http") + assert download_manifest.write_cancel_marker( + "model", "Owner/Repo", "Q4_K_M", "http" + ) status = asyncio.run(downloads.get_download_status_response("owner/repo", "Q4_K_M")) @@ -2405,7 +2489,9 @@ def test_model_download_watcher_invalidates_hf_cache_scan(monkeypatch): "_spawn_download_worker", lambda *_args, **_kwargs: object(), ) - monkeypatch.setattr(downloads.download_lifecycle.threading, "Thread", _ImmediateThread) + monkeypatch.setattr( + downloads.download_lifecycle.threading, "Thread", _ImmediateThread + ) monkeypatch.setattr( downloads.hf_cache_scan, "invalidate_hf_cache_scans", @@ -2512,7 +2598,10 @@ def test_two_concurrent_same_repo_variants_both_complete(monkeypatch, tmp_path): while time.monotonic() < deadline: s4 = registry.get_job(key_q4).state s8 = registry.get_job(key_q8).state - if s4 in download_registry.TERMINAL_STATES and s8 in download_registry.TERMINAL_STATES: + if ( + s4 in download_registry.TERMINAL_STATES + and s8 in download_registry.TERMINAL_STATES + ): break time.sleep(0.02) @@ -2628,7 +2717,9 @@ def test_snapshot_progress_filters_stale_blobs(monkeypatch, tmp_path): assert result["expected_bytes"] == 140 -def test_snapshot_progress_confirms_complete_only_with_verified_snapshot(monkeypatch, tmp_path): +def test_snapshot_progress_confirms_complete_only_with_verified_snapshot( + monkeypatch, tmp_path +): entry = tmp_path / "models--Org--Model" blobs = entry / "blobs" snap = entry / "snapshots" / "rev0" @@ -2691,7 +2782,9 @@ def test_expected_files_from_snapshot_dir_records_relative_paths_and_sizes(tmp_p assert all(f.sha256 is None for f in files) -def test_snapshot_progress_complete_with_manifest_synthesized_from_disk(monkeypatch, tmp_path): +def test_snapshot_progress_complete_with_manifest_synthesized_from_disk( + monkeypatch, tmp_path +): """A finished snapshot whose only manifest was synthesized from on-disk files still verifies as complete, so a refresh finalizes it instead of capping at 99% and evicting it as gone.""" @@ -2848,7 +2941,9 @@ def test_download_snapshot_writes_manifest_for_xet(monkeypatch, tmp_path): ), ) monkeypatch.setattr( - hf_download, "_verify_completed_download", lambda *args, **kwargs: verified.append(args) + hf_download, + "_verify_completed_download", + lambda *args, **kwargs: verified.append(args), ) monkeypatch.setattr( download_registry, "prepare_cache_for_transport", lambda *_args, **_kwargs: 0 @@ -2883,7 +2978,9 @@ def test_download_gguf_variant_writes_manifest_for_xet(monkeypatch, tmp_path): ), ) monkeypatch.setattr( - hf_download, "_verify_completed_download", lambda *args, **kwargs: verified.append(args) + hf_download, + "_verify_completed_download", + lambda *args, **kwargs: verified.append(args), ) monkeypatch.setattr( download_registry, "prepare_cache_for_transport", lambda *_args, **_kwargs: 0 @@ -2918,7 +3015,9 @@ def test_download_dataset_writes_manifest_for_xet(monkeypatch, tmp_path): ), ) monkeypatch.setattr( - hf_download, "_verify_completed_download", lambda *args, **kwargs: verified.append(args) + hf_download, + "_verify_completed_download", + lambda *args, **kwargs: verified.append(args), ) monkeypatch.setattr( download_registry, "prepare_cache_for_transport", lambda *_args, **_kwargs: 0 @@ -2956,7 +3055,9 @@ def test_dataset_status_includes_generation(monkeypatch): lambda repo_id, **_kwargs: repo_id, ) - result = asyncio.run(dataset_downloads.get_dataset_download_status_response("Org/Data")) + result = asyncio.run( + dataset_downloads.get_dataset_download_status_response("Org/Data") + ) assert result.state == "running" assert result.generation == 4 diff --git a/studio/backend/hub/utils/dataset_cache.py b/studio/backend/hub/utils/dataset_cache.py index 1a7a90d8a5..4ac737fe76 100644 --- a/studio/backend/hub/utils/dataset_cache.py +++ b/studio/backend/hub/utils/dataset_cache.py @@ -50,7 +50,9 @@ def _matches_label(snapshot: Path, path: Path, label: str) -> bool: return label in rel -def dataset_snapshot_from_cache_path(local_path: Optional[str], repo_id: str) -> Optional[Path]: +def dataset_snapshot_from_cache_path( + local_path: Optional[str], repo_id: str +) -> Optional[Path]: if not local_path or not repo_id: return None try: @@ -115,7 +117,9 @@ def cached_dataset_candidates( ) -> list[Path]: try: files = [ - p for p in snapshot.rglob("*") if p.is_file() and p.name.lower().endswith(extensions) + p + for p in snapshot.rglob("*") + if p.is_file() and p.name.lower().endswith(extensions) ] except OSError: return [] @@ -127,7 +131,9 @@ def cached_dataset_candidates( def score(path: Path) -> tuple[int, int, str]: rel = _rel_lower(snapshot, path) - subset_match = bool(subset_lower and _matches_label(snapshot, path, subset_lower)) + subset_match = bool( + subset_lower and _matches_label(snapshot, path, subset_lower) + ) split_match = bool(split_lower and split_label_matches(rel, split_lower)) location_rank = 3 if split_match and (not subset_lower or subset_match): diff --git a/studio/backend/hub/utils/dataset_format.py b/studio/backend/hub/utils/dataset_format.py index df02035365..db57bd202d 100644 --- a/studio/backend/hub/utils/dataset_format.py +++ b/studio/backend/hub/utils/dataset_format.py @@ -23,7 +23,10 @@ def _column_names(dataset, sample: Optional[dict] = None) -> list[str]: def _keyword_in_column(keyword: str, col_name: str) -> bool: - return re.search(r"\b" + re.escape(keyword) + r"\b", col_name, re.IGNORECASE) is not None + return ( + re.search(r"\b" + re.escape(keyword) + r"\b", col_name, re.IGNORECASE) + is not None + ) def _unknown_dataset_format( @@ -165,14 +168,19 @@ def detect_custom_format_heuristic(dataset): def has_keyword(col_name, keywords): col_lower = col_name.lower() col_normalized = col_lower.replace("_", "").replace("-", "").replace(" ", "") - return any(keyword in col_lower or keyword in col_normalized for keyword in keywords) + return any( + keyword in col_lower or keyword in col_normalized for keyword in keywords + ) def is_metadata(col_name): col_lower = col_name.lower() if col_lower in metadata_exact_match or col_lower in metadata_prefix_patterns: return True for pattern in metadata_prefix_patterns: - if col_lower.startswith(pattern.split("_")[0] + "_") and col_lower != pattern: + if ( + col_lower.startswith(pattern.split("_")[0] + "_") + and col_lower != pattern + ): if "_" in col_lower: prefix = col_lower.split("_")[0] if prefix in ["generation", "pass", "inference"]: @@ -181,7 +189,11 @@ def detect_custom_format_heuristic(dataset): def get_priority_score(col_name): col_lower = col_name.lower() - return sum(score for pattern, score in priority_patterns.items() if pattern in col_lower) + return sum( + score + for pattern, score in priority_patterns.items() + if pattern in col_lower + ) def get_content_length(col_name): try: @@ -195,7 +207,9 @@ def detect_custom_format_heuristic(dataset): score = 10 if role_type == "user": col_lower = col_name.lower() - if "task" in col_lower and not any(kw in col_lower for kw in user_words_high_priority): + if "task" in col_lower and not any( + kw in col_lower for kw in user_words_high_priority + ): score -= 15 score += get_priority_score(col_name) if role_type in ["assistant", "user"]: @@ -219,12 +233,19 @@ def detect_custom_format_heuristic(dataset): return score content_columns = [col for col in all_columns if not is_metadata(col)] - assistant_potential = [col for col in content_columns if has_keyword(col, assistant_words)] + assistant_potential = [ + col for col in content_columns if has_keyword(col, assistant_words) + ] user_potential = [col for col in content_columns if has_keyword(col, user_words)] assistant_candidates = [ (col, score) for col in assistant_potential - if (score := score_column(col, assistant_words, "assistant", len(assistant_potential))) > 0 + if ( + score := score_column( + col, assistant_words, "assistant", len(assistant_potential) + ) + ) + > 0 ] if assistant_candidates: assistant_candidates.sort(key = lambda item: item[1], reverse = True) @@ -399,7 +420,9 @@ def detect_multimodal_dataset(dataset): audio_columns.append(col_name) modality_types.add("audio") if audio_columns: - multimodal_columns = [col for col in multimodal_columns if col not in set(audio_columns)] + multimodal_columns = [ + col for col in multimodal_columns if col not in set(audio_columns) + ] detected_text_col = None if audio_columns: @@ -454,7 +477,9 @@ def detect_vlm_dataset_structure(dataset): and isinstance(content[0], dict) and "type" in content[0] ): - has_index = any("index" in item for item in content if isinstance(item, dict)) + has_index = any( + "index" in item for item in content if isinstance(item, dict) + ) if has_index and "images" in column_names: return { "format": "vlm_messages_llava", @@ -463,7 +488,9 @@ def detect_vlm_dataset_structure(dataset): "image_column": "images", "text_column": None, } - has_image = any("image" in item for item in content if isinstance(item, dict)) + has_image = any( + "image" in item for item in content if isinstance(item, dict) + ) if has_image: return { "format": "vlm_messages", @@ -555,9 +582,9 @@ def detect_vlm_dataset_structure(dataset): image_candidates = [] for col in column_names: value = sample[col] - if any(_keyword_in_column(keyword, col) for keyword in image_keywords) or _is_image_value( - value - ): + if any( + _keyword_in_column(keyword, col) for keyword in image_keywords + ) or _is_image_value(value): if hasattr(value, "size") and hasattr(value, "mode"): score = 100 elif isinstance(value, dict) and ("bytes" in value or "path" in value): @@ -726,7 +753,9 @@ def _standardize_sharegpt_row(row: dict[str, Any], chat_column: str) -> dict[str if not isinstance(message, dict): continue role = message.get("role") or message.get("from") - content = message.get("content") if "content" in message else message.get("value") + content = ( + message.get("content") if "content" in message else message.get("value") + ) messages.append( { "role": _ROLE_MAP.get(str(role), str(role or "user")), diff --git a/studio/backend/hub/utils/download_manifest.py b/studio/backend/hub/utils/download_manifest.py index 5366689296..7db83ab3f8 100644 --- a/studio/backend/hub/utils/download_manifest.py +++ b/studio/backend/hub/utils/download_manifest.py @@ -59,7 +59,9 @@ _LEGACY_MARKER_VERSION = 1 # Verbatim phrase the worker emits on a degraded completion and the download # lifecycle escalates to a warning log. Shared so the emit and match stay coupled. -MANIFEST_DEGRADED_MARKER = "completed without a manifest so partial detection is degraded" +MANIFEST_DEGRADED_MARKER = ( + "completed without a manifest so partial detection is degraded" +) @dataclass(frozen = True) @@ -473,14 +475,18 @@ def _iter_variant_state_files( yield _variant_from_state_file(entry, variant), entry -def iter_variant_manifests(repo_type: RepoType, repo_id: str) -> Iterator[tuple[str, Path]]: +def iter_variant_manifests( + repo_type: RepoType, repo_id: str +) -> Iterator[tuple[str, Path]]: """Yield (variant, manifest_path) for every variant-keyed manifest written for this repo. Used by is_gguf_repo_partial to enumerate all variants present on disk so the all-variants-broken gate can run.""" yield from _iter_variant_state_files(manifests_dir(), repo_type, repo_id) -def iter_variant_markers(repo_type: RepoType, repo_id: str) -> Iterator[tuple[str, Path]]: +def iter_variant_markers( + repo_type: RepoType, repo_id: str +) -> Iterator[tuple[str, Path]]: """Yield (variant, marker_path) for every variant-keyed cancel marker. Companion to iter_variant_manifests: catches variants cancelled before download-start ever wrote a manifest (very early failures).""" diff --git a/studio/backend/hub/utils/download_registry.py b/studio/backend/hub/utils/download_registry.py index 777d63e1b5..d2d4874e56 100644 --- a/studio/backend/hub/utils/download_registry.py +++ b/studio/backend/hub/utils/download_registry.py @@ -113,7 +113,9 @@ def _worker_breadcrumb_path(key: str) -> Optional[Path]: return parent / f"{safe}.json" -def write_worker_breadcrumb(key: str, pid: int, metadata: Optional["DownloadMetadata"]) -> None: +def write_worker_breadcrumb( + key: str, pid: int, metadata: Optional["DownloadMetadata"] +) -> None: """Record a live worker's PID so a restarted backend can reap it. Best effort: a write failure only forfeits boot-time reaping for this worker, still covered by the worker's own parent-death watchdog.""" @@ -366,7 +368,9 @@ def _iter_active_snapshot_dirs(repo_type: str, repo_id: str) -> Iterator[Path]: yield snapshot -def _manifest_verifies_against_active_cache(repo_type: str, repo_id: str, manifest) -> bool: +def _manifest_verifies_against_active_cache( + repo_type: str, repo_id: str, manifest +) -> bool: from hub.utils import download_manifest for snapshot_dir in _iter_active_snapshot_dirs(repo_type, repo_id): if download_manifest.verify_against_disk(manifest, snapshot_dir).ok: @@ -374,7 +378,9 @@ def _manifest_verifies_against_active_cache(repo_type: str, repo_id: str, manife return False -def _manifest_has_active_incomplete_blobs(repo_type: str, repo_id: str, manifest) -> bool: +def _manifest_has_active_incomplete_blobs( + repo_type: str, repo_id: str, manifest +) -> bool: if not getattr(manifest, "variant", None): return has_active_incomplete_blobs(repo_type, repo_id) expected_hashes = frozenset( @@ -383,7 +389,9 @@ def _manifest_has_active_incomplete_blobs(repo_type: str, repo_id: str, manifest if not expected_hashes: return has_active_incomplete_blobs(repo_type, repo_id) return bool( - incomplete_blob_hashes(repo_type, repo_id, active_only = True).intersection(expected_hashes) + incomplete_blob_hashes(repo_type, repo_id, active_only = True).intersection( + expected_hashes + ) ) @@ -398,7 +406,9 @@ def _is_transport_marker_file(path: Path) -> bool: # Matches ".transport", its tmps, and variant-scoped ".transport.gguf-*". # Real HF cache entries (blobs/refs/snapshots/.no_exist) never start with # ".transport.". - return path.name == TRANSPORT_MARKER_NAME or path.name.startswith(f"{TRANSPORT_MARKER_NAME}.") + return path.name == TRANSPORT_MARKER_NAME or path.name.startswith( + f"{TRANSPORT_MARKER_NAME}." + ) def _companion_marker_path(entry: Path) -> Path: @@ -516,9 +526,13 @@ def prepare_cache_for_transport( total_purged += _purge_incomplete_blobs(entry, only_blob_hashes, protected) else: if _read_marker(entry, variant) != mode: - total_purged += _purge_incomplete_blobs(entry, only_blob_hashes, protected) + total_purged += _purge_incomplete_blobs( + entry, only_blob_hashes, protected + ) if companion_blob_hashes and _read_companion_marker(entry) != mode: - total_purged += _purge_incomplete_blobs(entry, companion_blob_hashes, protected) + total_purged += _purge_incomplete_blobs( + entry, companion_blob_hashes, protected + ) _write_marker(entry, mode, variant) if has_companion: _write_companion_marker(entry, mode) @@ -562,7 +576,9 @@ def purge_empty_marker_dir( contents = list(entry.iterdir()) except OSError: continue - if not contents or not all(_is_transport_marker_file(item) for item in contents): + if not contents or not all( + _is_transport_marker_file(item) for item in contents + ): continue own_name = _marker_path(entry, variant).name own_markers = [ @@ -635,7 +651,9 @@ def incomplete_blob_hashes( return out -def completed_blob_bytes(repo_type: str, repo_id: str, blob_hashes: frozenset[str]) -> int: +def completed_blob_bytes( + repo_type: str, repo_id: str, blob_hashes: frozenset[str] +) -> int: """Sum finalized blob bytes for *blob_hashes* in the active HF cache root. A worker only writes to the active ``HF_HUB_CACHE`` root, so a baseline must @@ -658,7 +676,9 @@ def completed_blob_bytes(repo_type: str, repo_id: str, blob_hashes: frozenset[st return total -def existing_blob_bytes(repo_type: str, repo_id: str, blob_hashes: frozenset[str]) -> int: +def existing_blob_bytes( + repo_type: str, repo_id: str, blob_hashes: frozenset[str] +) -> int: """Bytes already on disk (finalized + ``.incomplete``) for *blob_hashes* in the active HF cache root. A blob is in exactly one state, so summing both candidate names never double-counts. Used to size what a (possibly resumed) @@ -1046,7 +1066,9 @@ class DownloadRegistry: return (metadata.variant or "").strip().lower() or None return variant_from_key(key) - def _delete_blocked_by_active_locked(self, repo_id: str, variant: Optional[str]) -> bool: + def _delete_blocked_by_active_locked( + self, repo_id: str, variant: Optional[str] + ) -> bool: """Whether an active download conflicts with deleting *repo_id*/*variant*. A whole-repo delete (``variant is None``) conflicts with any active @@ -1107,7 +1129,9 @@ class DownloadRegistry: if repo_key: candidate_keys = list(self._repo_active.get(repo_key, set())) else: - candidate_keys = [key for active in self._repo_active.values() for key in active] + candidate_keys = [ + key for active in self._repo_active.values() for key in active + ] refs: list[ActiveDownloadRef] = [] for key in candidate_keys: job = self._jobs.get(key) @@ -1231,7 +1255,9 @@ class DownloadRegistry: try: proc.wait(timeout = max(0.0, deadline - time.monotonic())) except subprocess.TimeoutExpired: - logger.warning(f"shutdown: {kind} worker for {key} did not exit after kill") + logger.warning( + f"shutdown: {kind} worker for {key} did not exit after kill" + ) except Exception: pass # Mark only genuinely interrupted workers (rc != 0, or None on wait diff --git a/studio/backend/hub/utils/gguf.py b/studio/backend/hub/utils/gguf.py index c9b3ce97e3..42a3802437 100644 --- a/studio/backend/hub/utils/gguf.py +++ b/studio/backend/hub/utils/gguf.py @@ -143,7 +143,9 @@ def pick_best_gguf(filenames: list[str]) -> Optional[str]: gguf_files = [ name for name in filenames - if is_gguf_filename(name) and not is_mmproj_filename(name) and not is_mtp_drafter_path(name) + if is_gguf_filename(name) + and not is_mmproj_filename(name) + and not is_mtp_drafter_path(name) ] if not gguf_files: return None @@ -246,7 +248,9 @@ def iter_hf_cache_snapshots(repo_id: str): yield from snapshots -def list_gguf_variants_from_hf_cache(repo_id: str) -> Optional[tuple[list[GgufVariantInfo], bool]]: +def list_gguf_variants_from_hf_cache( + repo_id: str, +) -> Optional[tuple[list[GgufVariantInfo], bool]]: for snapshot in iter_hf_cache_snapshots(repo_id): variants, has_vision = list_local_gguf_variants(str(snapshot)) if variants or has_vision: @@ -372,7 +376,9 @@ def list_gguf_variants( has_vision = True continue quant = extract_quant_label(filename) - quant_totals[quant] = quant_totals.get(quant, 0) + int(getattr(sibling, "size", 0) or 0) + quant_totals[quant] = quant_totals.get(quant, 0) + int( + getattr(sibling, "size", 0) or 0 + ) quant_first_file.setdefault(quant, filename) for quant, total_size in quant_totals.items(): diff --git a/studio/backend/hub/utils/gguf_plan.py b/studio/backend/hub/utils/gguf_plan.py index 60782fa109..e894c1a34a 100644 --- a/studio/backend/hub/utils/gguf_plan.py +++ b/studio/backend/hub/utils/gguf_plan.py @@ -60,7 +60,9 @@ def expected_file_from_sibling(sibling) -> Optional[ExpectedFile]: def is_companion_gguf_path(path: str) -> bool: """Companion (non-main) GGUF downloaded alongside a variant: the vision mmproj or the separate MTP drafter (Gemma 4).""" - return is_gguf_filename(path) and (is_mmproj_filename(path) or is_mtp_drafter_path(path)) + return is_gguf_filename(path) and ( + is_mmproj_filename(path) or is_mtp_drafter_path(path) + ) def is_main_gguf_variant_path(path: str, variant: str) -> bool: @@ -81,7 +83,9 @@ def _gguf_rfilename(sibling) -> Optional[str]: def mmproj_siblings(siblings: Sequence) -> list: - return [s for s in siblings if (name := _gguf_rfilename(s)) and is_mmproj_filename(name)] + return [ + s for s in siblings if (name := _gguf_rfilename(s)) and is_mmproj_filename(name) + ] def preferred_mmproj_sibling(siblings: Sequence) -> Optional[object]: @@ -89,7 +93,11 @@ def preferred_mmproj_sibling(siblings: Sequence) -> Optional[object]: if not candidates: return None return next( - (s for s in candidates if extract_quant_label(getattr(s, "rfilename")).upper() == "F16"), + ( + s + for s in candidates + if extract_quant_label(getattr(s, "rfilename")).upper() == "F16" + ), candidates[0], ) @@ -106,7 +114,8 @@ def preferred_mtp_sibling(siblings: Sequence) -> Optional[object]: ( s for s in siblings - if (name := _gguf_rfilename(s)) and name.lower().rsplit("/", 1)[-1].startswith("mtp-") + if (name := _gguf_rfilename(s)) + and name.lower().rsplit("/", 1)[-1].startswith("mtp-") ), key = lambda s: getattr(s, "rfilename"), ) @@ -121,11 +130,17 @@ def build_gguf_variant_plans(siblings: Sequence) -> dict[str, GgufVariantPlan]: for s in all_mmproj if isinstance(getattr(s, "rfilename", None), str) ) - all_mmproj_hashes = frozenset(h for h in (sibling_sha256(s) for s in all_mmproj) if h) + all_mmproj_hashes = frozenset( + h for h in (sibling_sha256(s) for s in all_mmproj) if h + ) companion = preferred_mmproj_sibling(siblings) - companion_expected = expected_file_from_sibling(companion) if companion is not None else None + companion_expected = ( + expected_file_from_sibling(companion) if companion is not None else None + ) mtp_sibling = preferred_mtp_sibling(siblings) - mtp_expected = expected_file_from_sibling(mtp_sibling) if mtp_sibling is not None else None + mtp_expected = ( + expected_file_from_sibling(mtp_sibling) if mtp_sibling is not None else None + ) companions_expected = tuple( file for file in (companion_expected, mtp_expected) if file is not None ) @@ -167,11 +182,17 @@ def plan_from_expected_files( all_mmproj_hashes: frozenset[str] | None = None, ) -> GgufVariantPlan: expected = tuple(expected_files) - main_files = tuple(file for file in expected if is_main_gguf_variant_path(file.path, variant)) - companion_files = tuple(file for file in expected if is_companion_gguf_path(file.path)) + main_files = tuple( + file for file in expected if is_main_gguf_variant_path(file.path, variant) + ) + companion_files = tuple( + file for file in expected if is_companion_gguf_path(file.path) + ) # Manifest-resume fallback for the mmproj fields below: companion_files # also holds the MTP drafter, so keep an mmproj-only view. - mmproj_files = tuple(file for file in companion_files if is_mmproj_filename(file.path)) + mmproj_files = tuple( + file for file in companion_files if is_mmproj_filename(file.path) + ) main_hashes = frozenset(file.sha256 for file in main_files if file.sha256) companion_hashes = frozenset(file.sha256 for file in companion_files if file.sha256) required_hashes = frozenset(file.sha256 for file in expected if file.sha256) diff --git a/studio/backend/hub/utils/hf_cache_state.py b/studio/backend/hub/utils/hf_cache_state.py index a1ac372abb..80d12cbb2c 100644 --- a/studio/backend/hub/utils/hf_cache_state.py +++ b/studio/backend/hub/utils/hf_cache_state.py @@ -66,7 +66,9 @@ def repo_cache_dir_name(repo_type: str, repo_id: str) -> str: return f"{repo_type}s--{repo_id.replace('/', '--')}" -def resolve_destructive_case_matches(target: str, candidates: Iterable[str]) -> Optional[set[str]]: +def resolve_destructive_case_matches( + target: str, candidates: Iterable[str] +) -> Optional[set[str]]: values = list(candidates) exact = {candidate for candidate in values if candidate == target} if exact: @@ -176,7 +178,9 @@ def iter_destructive_repo_cache_dirs(repo_type: str, repo_id: str) -> Iterator[P folded_target = target.lower() for root in hf_cache_roots(): try: - entries = [entry for entry in root.iterdir() if entry.name.lower() == folded_target] + entries = [ + entry for entry in root.iterdir() if entry.name.lower() == folded_target + ] except OSError: continue matched_names = resolve_destructive_case_matches( diff --git a/studio/backend/hub/utils/inventory_scan.py b/studio/backend/hub/utils/inventory_scan.py index 0f7ce6fe34..8f777ac3da 100644 --- a/studio/backend/hub/utils/inventory_scan.py +++ b/studio/backend/hub/utils/inventory_scan.py @@ -346,7 +346,11 @@ def _completed_gguf_variants(snapshot_dir: Optional[Path]) -> set[str]: except OSError: continue rel = path.relative_to(snapshot_dir).as_posix() - if not is_gguf_filename(rel) or is_mmproj_filename(rel) or is_mtp_drafter_path(rel): + if ( + not is_gguf_filename(rel) + or is_mmproj_filename(rel) + or is_mtp_drafter_path(rel) + ): continue quant = extract_quant_label(rel) split = _GGUF_SPLIT_RE.search(path.name) @@ -409,7 +413,8 @@ def is_snapshot_partial( state_applies = _state_applies_to_repo_cache_dir(repo_cache_dir) return _compose_partial( - lambda: state_applies and download_manifest.has_cancel_marker(repo_type, repo_id, None), + lambda: state_applies + and download_manifest.has_cancel_marker(repo_type, repo_id, None), lambda: _snapshot_legacy_partial(repo_type, repo_id, repo_cache_dir), lambda: _manifest_partial( repo_type, @@ -441,7 +446,8 @@ def is_variant_partial( state_applies = _state_applies_to_repo_cache_dir(repo_cache_dir) return _compose_partial( - lambda: state_applies and download_manifest.has_cancel_marker("model", repo_id, variant), + lambda: state_applies + and download_manifest.has_cancel_marker("model", repo_id, variant), lambda: bool( incomplete_blob_hashes and variant_blob_hashes diff --git a/studio/backend/hub/utils/llm_assist.py b/studio/backend/hub/utils/llm_assist.py index 00edb204c5..fcf5bbbdce 100644 --- a/studio/backend/hub/utils/llm_assist.py +++ b/studio/backend/hub/utils/llm_assist.py @@ -61,7 +61,9 @@ def _parse_json_response(text: str) -> Optional[dict[str, Any]]: return parsed if isinstance(parsed, dict) else None -def _generate_with_backend(backend, messages: list[dict[str, str]], max_tokens: int) -> str: +def _generate_with_backend( + backend, messages: list[dict[str, str]], max_tokens: int +) -> str: cumulative = "" for chunk in backend.generate_chat_completion( messages = messages, @@ -159,7 +161,9 @@ def _run_multi_pass_advisor( return None repo = os.environ.get("UNSLOTH_HELPER_MODEL_REPO", DEFAULT_HELPER_MODEL_REPO) - variant = os.environ.get("UNSLOTH_HELPER_MODEL_VARIANT", DEFAULT_HELPER_MODEL_VARIANT) + variant = os.environ.get( + "UNSLOTH_HELPER_MODEL_VARIANT", DEFAULT_HELPER_MODEL_VARIANT + ) backend = None try: from core.inference.llama_cpp import LlamaCppBackend @@ -179,7 +183,9 @@ def _run_multi_pass_advisor( samples_text = _sample_text(columns, samples) metadata_text = ( - json.dumps(dataset_metadata, indent = 2, default = str)[:500] if dataset_metadata else "N/A" + json.dumps(dataset_metadata, indent = 2, default = str)[:500] + if dataset_metadata + else "N/A" ) card_excerpt = (dataset_card or "")[:1200] or "N/A" hints = _target_hints(model_name, model_type) @@ -284,7 +290,9 @@ def _run_multi_pass_advisor( system_prompt = "" if not pass1.get("is_conversational"): user_cols = [col for col, role in column_roles.items() if role == "user"] - assistant_cols = [col for col, role in column_roles.items() if role == "assistant"] + assistant_cols = [ + col for col, role in column_roles.items() if role == "assistant" + ] prompt_raw = _generate_with_backend( backend, [ diff --git a/studio/backend/hub/utils/paths.py b/studio/backend/hub/utils/paths.py index afcb0b41dc..e107ae1258 100644 --- a/studio/backend/hub/utils/paths.py +++ b/studio/backend/hub/utils/paths.py @@ -189,7 +189,8 @@ def is_valid_repo_id(repo_id: str) -> bool: if len(segments) not in (1, 2): return False return all( - segment not in ("", ".", "..") and _VALID_REPO_ID_SEGMENT.fullmatch(segment) is not None + segment not in ("", ".", "..") + and _VALID_REPO_ID_SEGMENT.fullmatch(segment) is not None for segment in segments ) @@ -276,7 +277,9 @@ def _hf_hub_cache_dir() -> Path: from huggingface_hub.constants import HF_HUB_CACHE return Path(HF_HUB_CACHE) except Exception as exc: - logger.debug("Could not read huggingface_hub HF_HUB_CACHE, using default: %s", exc) + logger.debug( + "Could not read huggingface_hub HF_HUB_CACHE, using default: %s", exc + ) return Path.home() / ".cache" / "huggingface" / "hub" @@ -403,7 +406,9 @@ def resolve_dataset_path(path_value: str) -> Path: return path except ValueError: continue - raise ValueError(f"dataset path must be relative or under a dataset root: {raw!r}") + raise ValueError( + f"dataset path must be relative or under a dataset root: {raw!r}" + ) parts = [part for part in Path(normalized).parts if part not in ("", ".")] if parts[:2] == ["assets", "datasets"]: diff --git a/studio/backend/hub/utils/state_dir.py b/studio/backend/hub/utils/state_dir.py index a304477a3d..e7f39799e3 100644 --- a/studio/backend/hub/utils/state_dir.py +++ b/studio/backend/hub/utils/state_dir.py @@ -80,7 +80,9 @@ def repo_cache_basename(repo_type: RepoType, repo_id: str) -> str: # wrong filename and a misclassified scanner row (the Literal only guards # statically; dynamic/JSON-sourced values slip past it). if repo_type not in _VALID_REPO_TYPES: - raise ValueError(f"repo_type must be one of {_VALID_REPO_TYPES}, got {repo_type!r}") + raise ValueError( + f"repo_type must be one of {_VALID_REPO_TYPES}, got {repo_type!r}" + ) return f"{repo_type}s--{repo_id.replace('/', '--')}".lower() diff --git a/studio/backend/hub/workers/hf_download.py b/studio/backend/hub/workers/hf_download.py index 42a8ca52b3..21da3a255e 100644 --- a/studio/backend/hub/workers/hf_download.py +++ b/studio/backend/hub/workers/hf_download.py @@ -183,14 +183,17 @@ def _hf_token_arg(hf_token: str | None) -> HfTokenArg: def _retry_metadata_fetch(repo_id: str, fetch, *, label: str): - for attempt, timeout in enumerate((_METADATA_REQUEST_TIMEOUT, _METADATA_RETRY_TIMEOUT)): + for attempt, timeout in enumerate( + (_METADATA_REQUEST_TIMEOUT, _METADATA_RETRY_TIMEOUT) + ): try: return fetch(timeout) except Exception as e: if attempt == 1: raise print( - f"{label} request failed for {repo_id} " f"({type(e).__name__}: {e}); retrying.", + f"{label} request failed for {repo_id} " + f"({type(e).__name__}: {e}); retrying.", file = sys.stderr, ) time.sleep(_METADATA_RETRY_DELAY) @@ -439,7 +442,9 @@ def _recover_manifest_after_download( ) sys.exit(1) - fallback_files = download_manifest.expected_files_from_snapshot_dir(Path(snapshot_path)) + fallback_files = download_manifest.expected_files_from_snapshot_dir( + Path(snapshot_path) + ) if fallback_files and download_manifest.write_manifest( repo_type, repo_id, @@ -515,7 +520,9 @@ def _download_snapshot(repo_id: str, hf_token: str | None, mode: str) -> None: snapshot_path, mode, fetch_info = lambda: _model_info_with_retry(repo_id, hf_token), - expected_files_from_info = lambda recovered: _snapshot_download_plan(recovered)[1], + expected_files_from_info = lambda recovered: _snapshot_download_plan( + recovered + )[1], ) _verify_completed_download( "model", @@ -538,12 +545,15 @@ def _gguf_variant_target_plan( file = sys.stderr, ) raise RuntimeError( - f"Metadata unavailable while resolving GGUF variant '{variant}' " f"for {repo_id}" + f"Metadata unavailable while resolving GGUF variant '{variant}' " + f"for {repo_id}" ) from e return build_gguf_variant_plans(list(info.siblings)).get(variant.lower()) -def _download_gguf_variant(repo_id: str, variant: str, hf_token: str | None, mode: str) -> None: +def _download_gguf_variant( + repo_id: str, variant: str, hf_token: str | None, mode: str +) -> None: from huggingface_hub import snapshot_download from hub.utils.download_registry import prepare_cache_for_transport from hub.utils.hf_cache_state import has_active_incomplete_blobs diff --git a/studio/backend/main.py b/studio/backend/main.py index 064e261753..21eb800839 100644 --- a/studio/backend/main.py +++ b/studio/backend/main.py @@ -47,7 +47,9 @@ if sys.platform == "win32": try: if os.path.isdir(_default_root): - for _ver in sorted(os.listdir(_default_root), key = _ver_key, reverse = True): + for _ver in sorted( + os.listdir(_default_root), key = _ver_key, reverse = True + ): _bin = os.path.join(_default_root, _ver, "bin") if os.path.isdir(_bin): candidates.append(_bin) @@ -106,7 +108,9 @@ if sys.platform == "win32": _all_vers_main: list[str] = [] for _pkg_dir in _bnb_spec.submodule_search_locations: - for _dll in _glob.glob(os.path.join(_pkg_dir, "libbitsandbytes_rocm*.dll")): + for _dll in _glob.glob( + os.path.join(_pkg_dir, "libbitsandbytes_rocm*.dll") + ): _found_rocm_bnb = True _km = _re_bnb.search( r"libbitsandbytes_rocm(\d+)\.dll", os.path.basename(_dll) @@ -124,7 +128,9 @@ if sys.platform == "win32": # (HIP SDK on a CUDA/CPU box) must not force a ROCm backend onto a # non-ROCm bitsandbytes, which raises at import. DLL unparsable -> "72". if _found_rocm_bnb: - _bnb_rocm_ver_final = _bnb_rocm_ver or os.environ.get("BNB_ROCM_VERSION") or "72" + _bnb_rocm_ver_final = ( + _bnb_rocm_ver or os.environ.get("BNB_ROCM_VERSION") or "72" + ) os.environ["BNB_ROCM_VERSION"] = _bnb_rocm_ver_final os.environ["UNSLOTH_BNB_ROCM_VERSION_SOURCE"] = "detected" _logging.getLogger(__name__).info( @@ -167,7 +173,9 @@ try: configure_cpu_threads() except ValueError as exc: _raw = os.environ.get("UNSLOTH_CPU_THREADS") - raise SystemExit(f"Error: Invalid UNSLOTH_CPU_THREADS value {_raw!r}: {exc}") from None + raise SystemExit( + f"Error: Invalid UNSLOTH_CPU_THREADS value {_raw!r}: {exc}" + ) from None # Anaconda/conda-forge Python: seed platform._sys_version_cache before any # library import triggers attrs -> rich -> structlog -> platform crash. @@ -219,7 +227,9 @@ def _read_studio_install_id() -> str: /api/health emits "" and the launcher accepts any healthy backend. Carries no install-path info (matters when Studio runs -H 0.0.0.0).""" try: - token = (_STUDIO_ROOT_RESOLVED / "share" / "studio_install_id").read_text().strip() + token = ( + (_STUDIO_ROOT_RESOLVED / "share" / "studio_install_id").read_text().strip() + ) except (OSError, ValueError): return "" return token if _STUDIO_INSTALL_ID_RE.fullmatch(token) else "" @@ -311,7 +321,9 @@ def get_unsloth_version() -> str: except PackageNotFoundError: pass - version_file = _Path(__file__).resolve().parents[2] / "unsloth" / "models" / "_utils.py" + version_file = ( + _Path(__file__).resolve().parents[2] / "unsloth" / "models" / "_utils.py" + ) try: for line in version_file.read_text(encoding = "utf-8").splitlines(): if line.startswith("__version__ = "): @@ -417,7 +429,9 @@ async def lifespan(app: FastAPI): print(f"WARNING: {_msg}", flush = True) except Exception as _probe_exc: import structlog as _structlog - _structlog.get_logger(__name__).debug("llama.cpp startup probes failed: %s", _probe_exc) + _structlog.get_logger(__name__).debug( + "llama.cpp startup probes failed: %s", _probe_exc + ) from storage.studio_db import cleanup_orphaned_runs @@ -425,7 +439,9 @@ async def lifespan(app: FastAPI): cleanup_orphaned_runs() except Exception as exc: import structlog - structlog.get_logger(__name__).warning("cleanup_orphaned_runs failed at startup: %s", exc) + structlog.get_logger(__name__).warning( + "cleanup_orphaned_runs failed at startup: %s", exc + ) _start_helper_precache_if_enabled() @@ -528,7 +544,9 @@ def _build_csp(script_nonce: "str | None" = None) -> str: "https://*.googleusercontent.com wss://*.googleusercontent.com" ) else: - connect_src = "'self' https://huggingface.co https://datasets-server.huggingface.co" + connect_src = ( + "'self' https://huggingface.co https://datasets-server.huggingface.co" + ) return ( "default-src 'self'; " @@ -630,7 +648,11 @@ async def _send_411(send) -> None: async def _send_413(send, total_bytes: int, max_bytes: int) -> None: payload = _json_for_413.dumps( - {"detail": (f"Request body too large ({total_bytes:,} bytes; max {max_bytes:,}).")}, + { + "detail": ( + f"Request body too large ({total_bytes:,} bytes; max {max_bytes:,})." + ) + }, ).encode("utf-8") await send( { @@ -814,7 +836,9 @@ app.include_router(data_recipe_router, prefix = "/api/data-recipe", tags = ["dat app.include_router(llama_router, prefix = "/api/llama", tags = ["llama"]) app.include_router(export_router, prefix = "/api/export", tags = ["export"]) app.include_router(rag_router, prefix = "/api/rag", tags = ["rag"]) -app.include_router(training_history_router, prefix = "/api/train", tags = ["training-history"]) +app.include_router( + training_history_router, prefix = "/api/train", tags = ["training-history"] +) app.include_router(hub_inventory_router, prefix = "/api/hub", tags = ["hub"]) app.include_router(hub_datasets_router, prefix = "/api/hub/datasets", tags = ["hub"]) @@ -856,7 +880,9 @@ async def health_check(request: Request): from auth.authentication import get_current_subject as _gcs from fastapi.security import HTTPAuthorizationCredentials - creds = HTTPAuthorizationCredentials(scheme = "Bearer", credentials = auth.split(" ", 1)[1]) + creds = HTTPAuthorizationCredentials( + scheme = "Bearer", credentials = auth.split(" ", 1)[1] + ) # Must await: a bare coroutine is truthy and would skip the auth check subject = await _gcs(creds) except HTTPException: @@ -892,12 +918,16 @@ def studio_update_status(_current_subject: str = Depends(get_current_subject)): "/api/studio/download-transport-capabilities", response_model = TransportCapabilities, ) -def studio_download_transport_capabilities(_current_subject: str = Depends(get_current_subject)): +def studio_download_transport_capabilities( + _current_subject: str = Depends(get_current_subject), +): return asdict(get_download_transport_capabilities()) @app.post("/api/shutdown") -async def shutdown_server(request: Request, current_subject: str = Depends(get_current_subject)): +async def shutdown_server( + request: Request, current_subject: str = Depends(get_current_subject) +): """Gracefully shut down the Unsloth Studio server. Called by the frontend quit dialog so users can stop the server from the UI diff --git a/studio/backend/models/auth.py b/studio/backend/models/auth.py index f451f8d9dd..1d183615ae 100644 --- a/studio/backend/models/auth.py +++ b/studio/backend/models/auth.py @@ -24,13 +24,17 @@ class DesktopLoginRequest(BaseModel): class RefreshTokenRequest(BaseModel): """Refresh token payload to obtain new access + refresh tokens.""" - refresh_token: str = Field(..., description = "Refresh token from a previous login or refresh") + refresh_token: str = Field( + ..., description = "Refresh token from a previous login or refresh" + ) class AuthStatusResponse(BaseModel): """Indicate whether the seeded admin auth flow is ready.""" - initialized: bool = Field(..., description = "True if the auth database contains a login user") + initialized: bool = Field( + ..., description = "True if the auth database contains a login user" + ) default_username: str = Field( "unsloth", description = "Default admin username for first-boot UI prefill.", @@ -71,7 +75,9 @@ class ApiKeyResponse(BaseModel): id: int name: str - key_prefix: str = Field(..., description = "First 8 characters after sk-unsloth- for display") + key_prefix: str = Field( + ..., description = "First 8 characters after sk-unsloth- for display" + ) created_at: str last_used_at: Optional[str] = None expires_at: Optional[str] = None diff --git a/studio/backend/models/data_recipe.py b/studio/backend/models/data_recipe.py index e6f27e64df..c100b2b5f9 100644 --- a/studio/backend/models/data_recipe.py +++ b/studio/backend/models/data_recipe.py @@ -101,7 +101,9 @@ class SeedInspectUploadRequest(BaseModel): if not self.block_id: raise ValueError("block_id is required when using file_ids") if self.file_names is None or len(self.file_ids) != len(self.file_names): - raise ValueError("file_names must be provided and same length as file_ids") + raise ValueError( + "file_names must be provided and same length as file_ids" + ) if has_legacy: if not self.filename: raise ValueError("filename is required when using content_base64") diff --git a/studio/backend/models/export.py b/studio/backend/models/export.py index 584f82dea0..56cb53a371 100644 --- a/studio/backend/models/export.py +++ b/studio/backend/models/export.py @@ -21,7 +21,11 @@ def _validate_save_directory(value: str) -> str: if any(ch in raw for ch in ("\r", "\n")): raise ValueError("save_directory may not contain control characters") path = Path(raw).expanduser() - path_parts = (*path.parts, *PureWindowsPath(raw).parts, *raw.replace("\\", "/").split("/")) + path_parts = ( + *path.parts, + *PureWindowsPath(raw).parts, + *raw.replace("\\", "/").split("/"), + ) if any(len(part) > 255 for part in path_parts if part not in ("", ".", "/", "\\")): raise ValueError("save_directory path components must be <= 255 characters") if ( diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py index 04c63c2247..e73fd37a73 100644 --- a/studio/backend/models/inference.py +++ b/studio/backend/models/inference.py @@ -26,7 +26,9 @@ class LoadRequest(BaseModel): native_path_lease: Optional[str] = Field( None, description = "Frontend-visible signed native path grant" ) - hf_token: Optional[str] = Field(None, description = "HuggingFace token for gated models") + hf_token: Optional[str] = Field( + None, description = "HuggingFace token for gated models" + ) max_seq_length: int = Field( 0, ge = 0, @@ -49,7 +51,9 @@ class LoadRequest(BaseModel): @field_validator("chat_template_override") @classmethod - def normalize_blank_chat_template_override(cls, value: Optional[str]) -> Optional[str]: + def normalize_blank_chat_template_override( + cls, value: Optional[str] + ) -> Optional[str]: if value is not None and value.strip() == "": return None return value @@ -121,7 +125,9 @@ class ValidateModelRequest(BaseModel): native_path_lease: Optional[str] = Field( None, description = "Frontend-visible signed native path grant" ) - hf_token: Optional[str] = Field(None, description = "HuggingFace token for gated models") + hf_token: Optional[str] = Field( + None, description = "HuggingFace token for gated models" + ) gguf_variant: Optional[str] = Field( None, description = "GGUF quantization variant (e.g. 'Q4_K_M')" ) @@ -136,7 +142,9 @@ class ValidateModelResponse(BaseModel): valid: bool = Field(..., description = "Whether the model identifier looks valid") message: str = Field(..., description = "Human-readable validation message") identifier: Optional[str] = Field(None, description = "Resolved model identifier") - display_name: Optional[str] = Field(None, description = "Display name derived from identifier") + display_name: Optional[str] = Field( + None, description = "Display name derived from identifier" + ) is_gguf: bool = Field(False, description = "Whether this is a GGUF model (llama.cpp)") is_lora: bool = Field(False, description = "Whether this is a LoRA adapter") is_vision: bool = Field(False, description = "Whether this is a vision-capable model") @@ -154,10 +162,16 @@ class GenerateRequest(BaseModel): temperature: float = Field(0.6, ge = 0.0, le = 2.0, description = "Sampling temperature") top_p: float = Field(0.95, ge = 0.0, le = 1.0, description = "Top-p sampling") top_k: int = Field(20, ge = -1, le = 100, description = "Top-k sampling") - max_new_tokens: int = Field(2048, ge = 1, le = 4096, description = "Maximum tokens to generate") - repetition_penalty: float = Field(1.0, ge = 1.0, le = 2.0, description = "Repetition penalty") + max_new_tokens: int = Field( + 2048, ge = 1, le = 4096, description = "Maximum tokens to generate" + ) + repetition_penalty: float = Field( + 1.0, ge = 1.0, le = 2.0, description = "Repetition penalty" + ) presence_penalty: float = Field(0.0, ge = 0.0, le = 2.0, description = "Presence penalty") - image_base64: Optional[str] = Field(None, description = "Base64 encoded image for vision models") + image_base64: Optional[str] = Field( + None, description = "Base64 encoded image for vision models" + ) class LoadResponse(BaseModel): @@ -168,13 +182,19 @@ class LoadResponse(BaseModel): display_name: str = Field(..., description = "Display name of the model") is_vision: bool = Field(False, description = "Whether model is a vision model") is_lora: bool = Field(False, description = "Whether model is a LoRA adapter") - is_gguf: bool = Field(False, description = "Whether model is a GGUF model (llama.cpp)") + is_gguf: bool = Field( + False, description = "Whether model is a GGUF model (llama.cpp)" + ) is_diffusion: bool = Field( False, description = "Whether model is a block-diffusion model (DiffusionGemma)" ) is_audio: bool = Field(False, description = "Whether model is a TTS audio model") - audio_type: Optional[str] = Field(None, description = "Audio codec type: snac, csm, bicodec, dac") - has_audio_input: bool = Field(False, description = "Whether model accepts audio input (ASR)") + audio_type: Optional[str] = Field( + None, description = "Audio codec type: snac, csm, bicodec, dac" + ) + has_audio_input: bool = Field( + False, description = "Whether model accepts audio input (ASR)" + ) inference: dict = Field( ..., description = "Inference parameters (temperature, top_p, top_k, min_p)" ) @@ -274,7 +294,9 @@ class LoadProgressResponse(BaseModel): 0, description = "Total bytes across all GGUF shards for the active model.", ) - fraction: float = Field(0.0, description = "bytes_loaded / bytes_total, clamped to 0..1.") + fraction: float = Field( + 0.0, description = "bytes_loaded / bytes_total, clamped to 0..1." + ) class InferenceStatusResponse(BaseModel): @@ -287,17 +309,34 @@ class InferenceStatusResponse(BaseModel): None, description = "Loadable identifier for the active model.", ) - is_vision: bool = Field(False, description = "Whether the active model is a vision model") - is_gguf: bool = Field(False, description = "Whether the active model is a GGUF model (llama.cpp)") - is_diffusion: bool = Field( - False, description = "Whether the active model is a block-diffusion model (DiffusionGemma)" + is_vision: bool = Field( + False, description = "Whether the active model is a vision model" + ) + is_gguf: bool = Field( + False, description = "Whether the active model is a GGUF model (llama.cpp)" + ) + is_diffusion: bool = Field( + False, + description = "Whether the active model is a block-diffusion model (DiffusionGemma)", + ) + gguf_variant: Optional[str] = Field( + None, description = "GGUF quantization variant (e.g. Q4_K_M)" + ) + is_audio: bool = Field( + False, description = "Whether the active model is a TTS audio model" + ) + audio_type: Optional[str] = Field( + None, description = "Audio codec type: snac, csm, bicodec, dac" + ) + has_audio_input: bool = Field( + False, description = "Whether model accepts audio input (ASR)" + ) + loading: List[str] = Field( + default_factory = list, description = "Models currently being loaded" + ) + loaded: List[str] = Field( + default_factory = list, description = "Models currently loaded" ) - gguf_variant: Optional[str] = Field(None, description = "GGUF quantization variant (e.g. Q4_K_M)") - is_audio: bool = Field(False, description = "Whether the active model is a TTS audio model") - audio_type: Optional[str] = Field(None, description = "Audio codec type: snac, csm, bicodec, dac") - has_audio_input: bool = Field(False, description = "Whether model accepts audio input (ASR)") - loading: List[str] = Field(default_factory = list, description = "Models currently being loaded") - loaded: List[str] = Field(default_factory = list, description = "Models currently loaded") inference: Optional[Dict[str, Any]] = Field( None, description = "Recommended inference parameters for the active model" ) @@ -322,7 +361,9 @@ class InferenceStatusResponse(BaseModel): supports_tools: bool = Field( False, description = "Whether the active model supports tool calling" ) - context_length: Optional[int] = Field(None, description = "Context length of the active model") + context_length: Optional[int] = Field( + None, description = "Context length of the active model" + ) max_context_length: Optional[int] = Field( None, description = "Maximum context length currently available for the active model", @@ -651,7 +692,9 @@ class ChatCompletionRequest(BaseModel): parallel_tool_calls: Optional[bool] = Field( None, description = "Whether to enable parallel function calling during tool use." ) - seed: Optional[int] = Field(None, description = "Best-effort deterministic sampling seed.") + seed: Optional[int] = Field( + None, description = "Best-effort deterministic sampling seed." + ) stream_options: Optional[dict] = Field( None, description = 'Streaming options, e.g. {"include_usage": true} to emit a final usage chunk.', @@ -659,7 +702,9 @@ class ChatCompletionRequest(BaseModel): # ── Unsloth extensions (ignored by standard OpenAI clients) ── top_k: int = Field(20, ge = -1, le = 100, description = "[x-unsloth] Top-k sampling") - min_p: float = Field(0.01, ge = 0.0, le = 1.0, description = "[x-unsloth] Min-p sampling threshold") + min_p: float = Field( + 0.01, ge = 0.0, le = 1.0, description = "[x-unsloth] Min-p sampling threshold" + ) repetition_penalty: float = Field( 1.0, ge = 1.0, le = 2.0, description = "[x-unsloth] Repetition penalty" ) @@ -935,7 +980,9 @@ class ChatCompletionRequest(BaseModel): if not tc_id: continue function = tc.get("function") - function_name = function.get("name") if isinstance(function, dict) else None + function_name = ( + function.get("name") if isinstance(function, dict) else None + ) if msg.name and function_name == msg.name: name_match = (tc_id, asst_idx, tc_idx) break @@ -1030,7 +1077,9 @@ class ChoiceDelta(BaseModel): content: Optional[str] = None -OpenAIFinishReason = Literal["stop", "length", "tool_calls", "content_filter", "function_call"] +OpenAIFinishReason = Literal[ + "stop", "length", "tool_calls", "content_filter", "function_call" +] class ChunkChoice(BaseModel): @@ -1183,13 +1232,17 @@ class ResponsesFunctionCallInputItem(BaseModel): """ type: Literal["function_call"] - id: Optional[str] = Field(None, description = "Item id assigned by the server (e.g. fc_...)") + id: Optional[str] = Field( + None, description = "Item id assigned by the server (e.g. fc_...)" + ) call_id: str = Field( ..., description = "Correlation id matching a function_call_output on the next turn.", ) name: str - arguments: str = Field(..., description = "JSON string of the arguments the model produced.") + arguments: str = Field( + ..., description = "JSON string of the arguments the model produced." + ) status: Optional[Literal["in_progress", "completed", "incomplete"]] = None @@ -1276,7 +1329,9 @@ class ResponsesRequest(BaseModel): default = [], description = "Input text or list of messages / function_call / function_call_output items", ) - instructions: Optional[str] = Field(None, description = "System / developer instructions") + instructions: Optional[str] = Field( + None, description = "System / developer instructions" + ) temperature: Optional[float] = Field(None, ge = 0.0, le = 2.0) top_p: Optional[float] = Field(None, ge = 0.0, le = 1.0) max_output_tokens: Optional[int] = Field(None, ge = 1) @@ -1364,7 +1419,9 @@ class ResponsesOutputFunctionCall(BaseModel): id: str = Field(default_factory = lambda: f"fc_{uuid.uuid4().hex[:12]}") call_id: str name: str - arguments: str = Field(..., description = "JSON string of the arguments the model produced.") + arguments: str = Field( + ..., description = "JSON string of the arguments the model produced." + ) status: Literal["completed", "in_progress", "incomplete"] = "completed" @@ -1477,12 +1534,16 @@ def _merge_anthropic_system(system: Any, additions: list[str]) -> Any: if not additions: return system - addition_blocks = [{"type": "text", "text": text} for text in additions if text.strip()] + addition_blocks = [ + {"type": "text", "text": text} for text in additions if text.strip() + ] if not addition_blocks: return system if system is None: - return addition_blocks[0]["text"] if len(addition_blocks) == 1 else addition_blocks + return ( + addition_blocks[0]["text"] if len(addition_blocks) == 1 else addition_blocks + ) if isinstance(system, str): return "\n\n".join([system, *[block["text"] for block in addition_blocks]]) if isinstance(system, list): @@ -1561,7 +1622,9 @@ class AnthropicMessagesRequest(BaseModel): normalized = dict(data) normalized["messages"] = normalized_messages - normalized["system"] = _merge_anthropic_system(normalized.get("system"), system_additions) + normalized["system"] = _merge_anthropic_system( + normalized.get("system"), system_additions + ) return normalized @@ -1587,7 +1650,9 @@ class AnthropicResponseToolUseBlock(BaseModel): input: dict -AnthropicResponseBlock = Union[AnthropicResponseTextBlock, AnthropicResponseToolUseBlock] +AnthropicResponseBlock = Union[ + AnthropicResponseTextBlock, AnthropicResponseToolUseBlock +] class AnthropicMessagesResponse(BaseModel): diff --git a/studio/backend/models/mcp_servers.py b/studio/backend/models/mcp_servers.py index 606c2423bf..cf062efda7 100644 --- a/studio/backend/models/mcp_servers.py +++ b/studio/backend/models/mcp_servers.py @@ -53,5 +53,7 @@ class McpServerImportRequest(BaseModel): class McpServerImportResult(BaseModel): created: list[McpServerResponse] = Field(default_factory = list) - skipped: list[str] = Field(default_factory = list) # display names skipped as duplicates + skipped: list[str] = Field( + default_factory = list + ) # display names skipped as duplicates errors: list[str] = Field(default_factory = list) diff --git a/studio/backend/models/models.py b/studio/backend/models/models.py index ff00363ff8..7204c1c761 100644 --- a/studio/backend/models/models.py +++ b/studio/backend/models/models.py @@ -12,7 +12,9 @@ ModelType = Literal["text", "vision", "audio", "embeddings"] class CheckpointInfo(BaseModel): """Information about a discovered checkpoint directory.""" - display_name: str = Field(..., description = "User-friendly checkpoint name (folder name)") + display_name: str = Field( + ..., description = "User-friendly checkpoint name (folder name)" + ) path: str = Field(..., description = "Full path to the checkpoint directory") loss: Optional[float] = Field(None, description = "Training loss at this checkpoint") @@ -61,23 +63,33 @@ class ModelDetails(BaseModel): None, description = "Model identifier (alias for id, for backward compatibility)" ) name: Optional[str] = Field(None, description = "Display name for the model") - config: Optional[Dict[str, Any]] = Field(None, description = "Model configuration dictionary") + config: Optional[Dict[str, Any]] = Field( + None, description = "Model configuration dictionary" + ) is_vision: bool = Field(False, description = "Whether model is a vision model") is_embedding: bool = Field( False, description = "Whether model is an embedding/sentence-transformer model" ) is_lora: bool = Field(False, description = "Whether model is a LoRA adapter") - is_gguf: bool = Field(False, description = "Whether model is a GGUF model (llama.cpp format)") + is_gguf: bool = Field( + False, description = "Whether model is a GGUF model (llama.cpp format)" + ) is_mlx: bool = Field( False, description = "Whether model is served via the MLX backend (Apple Silicon)" ) is_audio: bool = Field(False, description = "Whether model is a TTS audio model") - audio_type: Optional[str] = Field(None, description = "Audio codec type: snac, csm, bicodec, dac") - has_audio_input: bool = Field(False, description = "Whether model accepts audio input (ASR)") + audio_type: Optional[str] = Field( + None, description = "Audio codec type: snac, csm, bicodec, dac" + ) + has_audio_input: bool = Field( + False, description = "Whether model accepts audio input (ASR)" + ) model_type: Optional[ModelType] = Field( None, description = "Collapsed model modality: text, vision, audio, or embeddings" ) - base_model: Optional[str] = Field(None, description = "Base model if this is a LoRA adapter") + base_model: Optional[str] = Field( + None, description = "Base model if this is a LoRA adapter" + ) max_position_embeddings: Optional[int] = Field( None, description = "Maximum context length supported by the model" ) @@ -90,7 +102,9 @@ class LoRAInfo(BaseModel): """LoRA adapter or exported model information""" display_name: str = Field(..., description = "Display name for the LoRA") - adapter_path: str = Field(..., description = "Path to the LoRA adapter or exported model") + adapter_path: str = Field( + ..., description = "Path to the LoRA adapter or exported model" + ) base_model: Optional[str] = Field(None, description = "Base model identifier") source: Optional[str] = Field(None, description = "'training' or 'exported'") export_type: Optional[str] = Field( @@ -101,21 +115,29 @@ class LoRAInfo(BaseModel): class LoRAScanResponse(BaseModel): """Response schema for scanning trained LoRA adapters""" - loras: List[LoRAInfo] = Field(default_factory = list, description = "List of found LoRA adapters") + loras: List[LoRAInfo] = Field( + default_factory = list, description = "List of found LoRA adapters" + ) outputs_dir: str = Field(..., description = "Directory that was scanned") class ModelListResponse(BaseModel): """Response schema for listing models""" - models: List[ModelDetails] = Field(default_factory = list, description = "List of models") - default_models: List[str] = Field(default_factory = list, description = "List of default model IDs") + models: List[ModelDetails] = Field( + default_factory = list, description = "List of models" + ) + default_models: List[str] = Field( + default_factory = list, description = "List of default model IDs" + ) class GgufVariantDetail(BaseModel): """A single GGUF quantization variant in a HuggingFace repo.""" - filename: str = Field(..., description = "GGUF filename (e.g., 'gemma-3-4b-it-Q4_K_M.gguf')") + filename: str = Field( + ..., description = "GGUF filename (e.g., 'gemma-3-4b-it-Q4_K_M.gguf')" + ) quant: str = Field(..., description = "Quantization label (e.g., 'Q4_K_M')") size_bytes: int = Field(0, description = "File size in bytes") downloaded: bool = Field( @@ -161,7 +183,9 @@ class LocalModelInfo(BaseModel): class LocalModelListResponse(BaseModel): """Response schema for listing local/cached models.""" - models_dir: str = Field(..., description = "Directory scanned for custom local models") + models_dir: str = Field( + ..., description = "Directory scanned for custom local models" + ) hf_cache_dir: Optional[str] = Field( None, description = "HF cache root that was scanned", @@ -179,7 +203,9 @@ class LocalModelListResponse(BaseModel): class AddScanFolderRequest(BaseModel): """Request body for adding a custom scan folder.""" - path: str = Field(..., description = "Absolute or relative directory path to scan for models") + path: str = Field( + ..., description = "Absolute or relative directory path to scan for models" + ) class ScanFolderInfo(BaseModel): diff --git a/studio/backend/models/providers.py b/studio/backend/models/providers.py index 5a75246c07..7823d4f22c 100644 --- a/studio/backend/models/providers.py +++ b/studio/backend/models/providers.py @@ -14,7 +14,9 @@ from pydantic import BaseModel, Field class ProviderRegistryEntry(BaseModel): """A supported provider type with its default configuration.""" - provider_type: str = Field(..., description = "Provider identifier (e.g. 'openai', 'mistral')") + provider_type: str = Field( + ..., description = "Provider identifier (e.g. 'openai', 'mistral')" + ) display_name: str = Field(..., description = "Human-readable provider name") base_url: str = Field(..., description = "Default API base URL") default_models: list[str] = Field( @@ -42,7 +44,9 @@ class ProviderCreate(BaseModel): """Request to create a saved provider configuration.""" provider_type: str = Field(..., description = "Provider type from the registry") - display_name: str = Field(..., description = "User-chosen label (e.g. 'My OpenAI Key')") + display_name: str = Field( + ..., description = "User-chosen label (e.g. 'My OpenAI Key')" + ) base_url: Optional[str] = Field( None, description = "Custom base URL (overrides registry default). Omit to use the default.", @@ -54,7 +58,9 @@ class ProviderUpdate(BaseModel): display_name: Optional[str] = Field(None, description = "New display name") base_url: Optional[str] = Field(None, description = "New base URL") - is_enabled: Optional[bool] = Field(None, description = "Enable or disable this provider") + is_enabled: Optional[bool] = Field( + None, description = "Enable or disable this provider" + ) class ProviderResponse(BaseModel): @@ -77,7 +83,9 @@ class ProviderModelInfo(BaseModel): id: str = Field(..., description = "Model ID as expected by the provider API") display_name: str = Field("", description = "Human-readable model name") - context_length: Optional[int] = Field(None, description = "Maximum context length in tokens") + context_length: Optional[int] = Field( + None, description = "Maximum context length in tokens" + ) owned_by: Optional[str] = Field(None, description = "Model owner/organization") diff --git a/studio/backend/models/responses.py b/studio/backend/models/responses.py index cc420c0f69..f8fc586a86 100644 --- a/studio/backend/models/responses.py +++ b/studio/backend/models/responses.py @@ -21,10 +21,16 @@ class TrainingStopResponse(BaseModel): class TrainingMetricsResponse(BaseModel): """Response for training metrics history""" - loss_history: List[float] = Field(default_factory = list, description = "Loss values per step") - lr_history: List[float] = Field(default_factory = list, description = "Learning rate per step") + loss_history: List[float] = Field( + default_factory = list, description = "Loss values per step" + ) + lr_history: List[float] = Field( + default_factory = list, description = "Learning rate per step" + ) step_history: List[int] = Field(default_factory = list, description = "Step numbers") - grad_norm_history: List[float] = Field(default_factory = list, description = "Gradient norm values") + grad_norm_history: List[float] = Field( + default_factory = list, description = "Gradient norm values" + ) grad_norm_step_history: List[int] = Field( default_factory = list, description = "Step numbers for gradient norm values" ) diff --git a/studio/backend/models/training.py b/studio/backend/models/training.py index ca04591178..801d24d7f0 100644 --- a/studio/backend/models/training.py +++ b/studio/backend/models/training.py @@ -37,7 +37,9 @@ class S3Config(BaseModel): bucket: str = Field(..., description = "S3 bucket name") region: str = Field("us-east-1", description = "AWS region") - prefix: Optional[str] = Field(None, description = "Optional path prefix within bucket") + prefix: Optional[str] = Field( + None, description = "Optional path prefix within bucket" + ) access_key_id: Optional[str] = Field( None, alias = "accessKeyId", @@ -58,7 +60,9 @@ class S3Config(BaseModel): def _check_credentials(self) -> "S3Config": # Require either IAM role auth or a full key pair so credentials are # never half-configured. - if not self.use_iam_role and not (self.access_key_id and self.secret_access_key): + if not self.use_iam_role and not ( + self.access_key_id and self.secret_access_key + ): raise ValueError( "s3_config requires either use_iam_role=True or both " "access_key_id and secret_access_key" @@ -77,7 +81,9 @@ def _parse_lr(v: Any) -> float: except (TypeError, ValueError): raise ValueError(f"learning_rate must be parseable as float (got {v!r})") if not (lr > 0.0): - raise ValueError(f"learning_rate must be > 0 (got {lr!r}); typical range is 1e-6 .. 1e-3") + raise ValueError( + f"learning_rate must be > 0 (got {lr!r}); typical range is 1e-6 .. 1e-3" + ) if lr >= _MAX_LR_VALUE: raise ValueError( f"learning_rate must be < 1.0 (got {lr!r}); " @@ -93,9 +99,11 @@ class TrainingStartRequest(BaseModel): model_name: str = Field( ..., description = "Model identifier (e.g., 'unsloth/llama-3-8b-bnb-4bit')" ) - training_type: Literal["LoRA/QLoRA", "Full Finetuning", "Continued Pretraining"] = Field( - ..., - description = "Training type: 'LoRA/QLoRA', 'Full Finetuning', or 'Continued Pretraining'", + training_type: Literal["LoRA/QLoRA", "Full Finetuning", "Continued Pretraining"] = ( + Field( + ..., + description = "Training type: 'LoRA/QLoRA', 'Full Finetuning', or 'Continued Pretraining'", + ) ) hf_token: Optional[str] = Field(None, description = "HuggingFace token") load_in_4bit: bool = Field(True, description = "Load model in 4-bit quantization") @@ -110,7 +118,9 @@ class TrainingStartRequest(BaseModel): ) # Dataset parameters - hf_dataset: Optional[str] = Field(None, description = "HuggingFace dataset identifier") + hf_dataset: Optional[str] = Field( + None, description = "HuggingFace dataset identifier" + ) local_datasets: List[str] = Field( default_factory = list, description = "List of local dataset paths" ) @@ -120,8 +130,12 @@ class TrainingStartRequest(BaseModel): format_type: str = Field(..., description = "Dataset format type") subset: Optional[str] = None train_split: Optional[str] = Field("train", description = "Training split name") - eval_split: Optional[str] = Field(None, description = "Eval split name. None = auto-detect") - eval_steps: float = Field(0.00, description = "Fraction of total steps between evals (0-1)") + eval_split: Optional[str] = Field( + None, description = "Eval split name. None = auto-detect" + ) + eval_steps: float = Field( + 0.00, description = "Fraction of total steps between evals (0-1)" + ) dataset_slice_start: Optional[int] = Field( None, description = "Inclusive start row index for dataset slicing" ) @@ -150,7 +164,9 @@ class TrainingStartRequest(BaseModel): if v is None: raise ValueError("batch_size is required") if v < 1 or v > _MAX_BATCH_SIZE: - raise ValueError(f"batch_size must be in [1, {_MAX_BATCH_SIZE}] (got {v!r})") + raise ValueError( + f"batch_size must be in [1, {_MAX_BATCH_SIZE}] (got {v!r})" + ) return v @field_validator("gradient_accumulation_steps") @@ -160,7 +176,8 @@ class TrainingStartRequest(BaseModel): return 1 if v < 1 or v > _MAX_GRAD_ACCUM: raise ValueError( - f"gradient_accumulation_steps must be in [1, {_MAX_GRAD_ACCUM}] " f"(got {v!r})" + f"gradient_accumulation_steps must be in [1, {_MAX_GRAD_ACCUM}] " + f"(got {v!r})" ) return v @@ -181,14 +198,18 @@ class TrainingStartRequest(BaseModel): if v is None: return v if not isinstance(v, int) or v < 0 or v > _MAX_STEPS: - raise ValueError(f"max_steps must be a non-negative int <= {_MAX_STEPS} (got {v!r})") + raise ValueError( + f"max_steps must be a non-negative int <= {_MAX_STEPS} (got {v!r})" + ) return v @field_validator("max_seq_length") @classmethod def _check_max_seq_length(cls, v: int) -> int: if v is None or v < 1 or v > _MAX_SEQ_LENGTH: - raise ValueError(f"max_seq_length must be in [1, {_MAX_SEQ_LENGTH}] (got {v!r})") + raise ValueError( + f"max_seq_length must be in [1, {_MAX_SEQ_LENGTH}] (got {v!r})" + ) return v @field_validator("vision_image_size", mode = "before") @@ -231,7 +252,8 @@ class TrainingStartRequest(BaseModel): return v if not isinstance(v, int) or v < 0 or v > _MAX_STEPS: raise ValueError( - f"warmup_steps must be a non-negative int <= {_MAX_STEPS} " f"(got {v!r})" + f"warmup_steps must be a non-negative int <= {_MAX_STEPS} " + f"(got {v!r})" ) return v @@ -267,7 +289,9 @@ class TrainingStartRequest(BaseModel): except (TypeError, ValueError): raise ValueError(f"weight_decay must be a number (got {v!r})") if wd < 0 or wd > 10.0: - raise ValueError(f"weight_decay must be in [0, 10] (got {wd!r}); typical 0..0.1") + raise ValueError( + f"weight_decay must be in [0, 10] (got {wd!r}); typical 0..0.1" + ) return wd @field_validator("lora_r") @@ -285,7 +309,9 @@ class TrainingStartRequest(BaseModel): if v is None: return 16 if v < 1 or v > _MAX_LORA_ALPHA: - raise ValueError(f"lora_alpha must be in [1, {_MAX_LORA_ALPHA}] (got {v!r})") + raise ValueError( + f"lora_alpha must be in [1, {_MAX_LORA_ALPHA}] (got {v!r})" + ) return v @field_validator("lora_dropout") @@ -314,7 +340,9 @@ class TrainingStartRequest(BaseModel): num_epochs: int = Field(1, description = "Number of training epochs") learning_rate: str = Field("2e-4", description = "Learning rate") batch_size: int = Field(1, description = "Batch size") - gradient_accumulation_steps: int = Field(1, description = "Gradient accumulation steps") + gradient_accumulation_steps: int = Field( + 1, description = "Gradient accumulation steps" + ) warmup_steps: Optional[int] = Field(None, description = "Warmup steps") warmup_ratio: Optional[float] = Field(None, description = "Warmup ratio") max_steps: Optional[int] = Field(None, description = "Maximum training steps") @@ -342,19 +370,31 @@ class TrainingStartRequest(BaseModel): lora_r: int = Field(16, description = "LoRA rank") lora_alpha: int = Field(16, description = "LoRA alpha") lora_dropout: float = Field(0.0, description = "LoRA dropout") - target_modules: List[str] = Field(default_factory = list, description = "Target modules for LoRA") - gradient_checkpointing: str = Field("", description = "Gradient checkpointing setting") + target_modules: List[str] = Field( + default_factory = list, description = "Target modules for LoRA" + ) + gradient_checkpointing: str = Field( + "", description = "Gradient checkpointing setting" + ) use_rslora: bool = Field(False, description = "Use RSLoRA") use_loftq: bool = Field(False, description = "Use LoftQ") train_on_completions: bool = Field(False, description = "Train on completions only") # Vision-specific LoRA parameters finetune_vision_layers: bool = Field(False, description = "Finetune vision layers") - finetune_language_layers: bool = Field(False, description = "Finetune language layers") - finetune_attention_modules: bool = Field(False, description = "Finetune attention modules") + finetune_language_layers: bool = Field( + False, description = "Finetune language layers" + ) + finetune_attention_modules: bool = Field( + False, description = "Finetune attention modules" + ) finetune_mlp_modules: bool = Field(False, description = "Finetune MLP modules") - is_dataset_image: bool = Field(False, description = "Whether the dataset contains image data") - is_dataset_audio: bool = Field(False, description = "Whether the dataset contains audio data") + is_dataset_image: bool = Field( + False, description = "Whether the dataset contains image data" + ) + is_dataset_audio: bool = Field( + False, description = "Whether the dataset contains audio data" + ) is_embedding: bool = Field( False, description = "Whether model is an embedding/sentence-transformer model" ) @@ -385,7 +425,9 @@ class TrainingStartRequest(BaseModel): def _check_steps_or_epochs(self) -> "TrainingStartRequest": # Each accepts 0 as "use the other"; both 0 means nothing to train. if (self.max_steps is None or self.max_steps == 0) and self.num_epochs == 0: - raise ValueError("Either num_epochs or max_steps must be > 0; both cannot be 0.") + raise ValueError( + "Either num_epochs or max_steps must be > 0; both cannot be 0." + ) return self @@ -412,7 +454,9 @@ class TrainingStatus(BaseModel): "error", "stopped", ] = Field(..., description = "Current phase of training pipeline") - is_training_running: bool = Field(..., description = "True if training loop is actively running") + is_training_running: bool = Field( + ..., description = "True if training loop is actively running" + ) eval_enabled: bool = Field( False, description = "True if evaluation dataset is configured for this training run", @@ -437,7 +481,9 @@ class TrainingProgress(BaseModel): total_steps: int = Field(..., description = "Total training steps") loss: Optional[float] = Field(None, description = "Current loss value") learning_rate: Optional[float] = Field(None, description = "Current learning rate") - progress_percent: float = Field(..., description = "Progress percentage (0.0 to 100.0)") + progress_percent: float = Field( + ..., description = "Progress percentage (0.0 to 100.0)" + ) epoch: Optional[float] = Field(None, description = "Current epoch") elapsed_seconds: Optional[float] = Field( None, description = "Time elapsed since training started" @@ -446,7 +492,9 @@ class TrainingProgress(BaseModel): grad_norm: Optional[float] = Field( None, description = "L2 norm of gradients, computed before gradient clipping" ) - num_tokens: Optional[int] = Field(None, description = "Total number of tokens processed so far") + num_tokens: Optional[int] = Field( + None, description = "Total number of tokens processed so far" + ) eval_loss: Optional[float] = Field( None, description = "Eval loss from the most recent evaluation step" ) diff --git a/studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/scraper.py b/studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/scraper.py index 637193e8b3..adf8210575 100644 --- a/studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/scraper.py +++ b/studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/scraper.py @@ -90,13 +90,18 @@ def _read_jsonl(path: Path, max_rows: int | None = None): def _flatten_issue_row(r: dict, repo: str, include_comments: bool, max_c: int) -> dict: - labels = [l.get("name") for l in (r.get("labels", {}) or {}).get("nodes", []) if l.get("name")] + labels = [ + l.get("name") + for l in (r.get("labels", {}) or {}).get("nodes", []) + if l.get("name") + ] comments_nodes = (r.get("comments") or {}).get("nodes") or [] comments_text = "" if include_comments and comments_nodes: kept = comments_nodes[:max_c] comments_text = "\n\n".join( - f"[{(c.get('author') or {}).get('login', '?')}]: {c.get('body') or ''}" for c in kept + f"[{(c.get('author') or {}).get('login', '?')}]: {c.get('body') or ''}" + for c in kept ) return { "item_type": "issue", @@ -115,13 +120,18 @@ def _flatten_issue_row(r: dict, repo: str, include_comments: bool, max_c: int) - def _flatten_pr_row(r: dict, repo: str, include_comments: bool, max_c: int) -> dict: - labels = [l.get("name") for l in (r.get("labels", {}) or {}).get("nodes", []) if l.get("name")] + labels = [ + l.get("name") + for l in (r.get("labels", {}) or {}).get("nodes", []) + if l.get("name") + ] comments_nodes = (r.get("comments") or {}).get("nodes") or [] comments_text = "" if include_comments and comments_nodes: kept = comments_nodes[:max_c] comments_text = "\n\n".join( - f"[{(c.get('author') or {}).get('login', '?')}]: {c.get('body') or ''}" for c in kept + f"[{(c.get('author') or {}).get('login', '?')}]: {c.get('body') or ''}" + for c in kept ) return { "item_type": "pull", @@ -194,8 +204,14 @@ def scrape(cfg: ScrapeConfig, base_dir: Path): scraper.scrape_prs() if "commits" in cfg.item_types: default_ref = repo_meta.get("defaultBranchRef") or {} - default_branch = default_ref.get("name") if isinstance(default_ref, dict) else None - branch = f"refs/heads/{default_branch}" if default_branch else "refs/heads/main" + default_branch = ( + default_ref.get("name") if isinstance(default_ref, dict) else None + ) + branch = ( + f"refs/heads/{default_branch}" + if default_branch + else "refs/heads/main" + ) scraper.scrape_commits(branch = branch) finally: scraper.close() @@ -205,12 +221,16 @@ def scrape(cfg: ScrapeConfig, base_dir: Path): if "issues" in cfg.item_types: for row in _read_jsonl(repo_dir / "issues.jsonl", read_cap): all_rows.append( - _flatten_issue_row(row, repo, cfg.include_comments, cfg.max_comments_per_item) + _flatten_issue_row( + row, repo, cfg.include_comments, cfg.max_comments_per_item + ) ) if "pulls" in cfg.item_types: for row in _read_jsonl(repo_dir / "pull_requests.jsonl", read_cap): all_rows.append( - _flatten_pr_row(row, repo, cfg.include_comments, cfg.max_comments_per_item) + _flatten_pr_row( + row, repo, cfg.include_comments, cfg.max_comments_per_item + ) ) if "commits" in cfg.item_types: for row in _read_jsonl(repo_dir / "commits.jsonl", read_cap): diff --git a/studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/scraper_impl/gh_client.py b/studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/scraper_impl/gh_client.py index 0ba3394ca3..8b65483375 100644 --- a/studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/scraper_impl/gh_client.py +++ b/studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/scraper_impl/gh_client.py @@ -60,7 +60,9 @@ class GitHubClient: token_source: str | None = None, ): if token: - self._token_source = token_source or "explicit token argument (recipe-level field)" + self._token_source = ( + token_source or "explicit token argument (recipe-level field)" + ) elif os.environ.get("GH_TOKEN"): self._token_source = "GH_TOKEN environment variable" token = os.environ["GH_TOKEN"] @@ -70,7 +72,9 @@ class GitHubClient: else: raise RuntimeError("GH_TOKEN or GITHUB_TOKEN not set in environment") self.session = requests.Session() - self.session.headers.update({**BASE_HEADERS, "Authorization": f"Bearer {token}"}) + self.session.headers.update( + {**BASE_HEADERS, "Authorization": f"Bearer {token}"} + ) self.min_remaining_graphql = min_remaining_graphql self.min_remaining_rest = min_remaining_rest self.graphql_remaining: Optional[int] = None @@ -205,7 +209,9 @@ class GitHubClient: errs = data["errors"] for e in errs: if e.get("type") == "RATE_LIMITED": - self._sleep_until((self.graphql_reset or int(time.time()) + 60)) + self._sleep_until( + (self.graphql_reset or int(time.time()) + 60) + ) break else: # No rate-limit error: log and return partial @@ -237,7 +243,9 @@ class GitHubClient: last_err = None for attempt in range(max_retries): try: - r = self.session.request(method, url, params = params, json = json_body, timeout = 120) + r = self.session.request( + method, url, params = params, json = json_body, timeout = 120 + ) self.calls_rest += 1 rem = r.headers.get("X-RateLimit-Remaining") rst = r.headers.get("X-RateLimit-Reset") @@ -261,7 +269,9 @@ class GitHubClient: if r.status_code in (403, 429): retry_after = _retry_after_seconds(r.headers.get("Retry-After")) if retry_after is not None: - log.warning("Secondary rate limit on REST. Sleep %ds.", retry_after) + log.warning( + "Secondary rate limit on REST. Sleep %ds.", retry_after + ) time.sleep(retry_after + 2) continue # Primary rate limit @@ -291,7 +301,9 @@ class GitHubClient: while True: r = self.rest("GET", url, params = params if url == path else None) if r.status_code != 200: - log.error("REST paginate got %s at %s: %s", r.status_code, url, r.text[:200]) + log.error( + "REST paginate got %s at %s: %s", r.status_code, url, r.text[:200] + ) return items = r.json() if isinstance(items, dict): diff --git a/studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/scraper_impl/scraper.py b/studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/scraper_impl/scraper.py index a7ddaef5fe..4fe399847f 100644 --- a/studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/scraper_impl/scraper.py +++ b/studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/scraper_impl/scraper.py @@ -86,7 +86,11 @@ class RepoScraper: return counter >= lim def _log_rate(self, where: str, data: Dict[str, Any]) -> None: - rl = data.get("data", {}).get("rateLimit") if isinstance(data.get("data"), dict) else None + rl = ( + data.get("data", {}).get("rateLimit") + if isinstance(data.get("data"), dict) + else None + ) if rl: log.debug( "[%s] rate cost=%s remaining=%s resetAt=%s", @@ -98,7 +102,9 @@ class RepoScraper: # ----- repo meta ----- def scrape_repo_meta(self) -> Dict[str, Any]: - data = self.client.graphql(Q.REPO_META_QUERY, {"owner": self.owner, "name": self.name}) + data = self.client.graphql( + Q.REPO_META_QUERY, {"owner": self.owner, "name": self.name} + ) self._log_rate("repo_meta", data) repo = data.get("data", {}).get("repository") or {} repo["_fetchedAt"] = ts() @@ -143,7 +149,11 @@ class RepoScraper: self._paginate_issue_comments( it["number"], it["comments"]["pageInfo"]["endCursor"] ) - if it.get("timelineItems", {}).get("pageInfo", {}).get("hasNextPage"): + if ( + it.get("timelineItems", {}) + .get("pageInfo", {}) + .get("hasNextPage") + ): self._paginate_issue_timeline( it["number"], it["timelineItems"]["pageInfo"]["endCursor"], @@ -249,16 +259,30 @@ class RepoScraper: num = pr["number"] if not self.light: if pr.get("comments", {}).get("pageInfo", {}).get("hasNextPage"): - self._paginate_pr_comments(num, pr["comments"]["pageInfo"]["endCursor"]) - if pr.get("timelineItems", {}).get("pageInfo", {}).get("hasNextPage"): + self._paginate_pr_comments( + num, pr["comments"]["pageInfo"]["endCursor"] + ) + if ( + pr.get("timelineItems", {}) + .get("pageInfo", {}) + .get("hasNextPage") + ): self._paginate_pr_timeline( num, pr["timelineItems"]["pageInfo"]["endCursor"] ) if pr.get("commits", {}).get("pageInfo", {}).get("hasNextPage"): - self._paginate_pr_commits(num, pr["commits"]["pageInfo"]["endCursor"]) + self._paginate_pr_commits( + num, pr["commits"]["pageInfo"]["endCursor"] + ) if pr.get("files", {}).get("pageInfo", {}).get("hasNextPage"): - self._paginate_pr_files(num, pr["files"]["pageInfo"]["endCursor"]) - if pr.get("reviewThreads", {}).get("pageInfo", {}).get("hasNextPage"): + self._paginate_pr_files( + num, pr["files"]["pageInfo"]["endCursor"] + ) + if ( + pr.get("reviewThreads", {}) + .get("pageInfo", {}) + .get("hasNextPage") + ): self._paginate_pr_review_threads( num, pr["reviewThreads"]["pageInfo"]["endCursor"] ) @@ -316,7 +340,9 @@ class RepoScraper: "after": cur, } data = self.client.graphql(Q.PR_TIMELINE_QUERY, vars_) - item = ((data.get("data") or {}).get("repository") or {}).get("pullRequest") or {} + item = ((data.get("data") or {}).get("repository") or {}).get( + "pullRequest" + ) or {} tl = item.get("timelineItems") or {} for ev in tl.get("nodes") or []: ev["_owner"] = self.owner @@ -339,7 +365,9 @@ class RepoScraper: "after": cur, } data = self.client.graphql(Q.PR_COMMITS_QUERY, vars_) - item = ((data.get("data") or {}).get("repository") or {}).get("pullRequest") or {} + item = ((data.get("data") or {}).get("repository") or {}).get( + "pullRequest" + ) or {} cc = item.get("commits") or {} for c in cc.get("nodes") or []: c["_owner"] = self.owner @@ -362,7 +390,9 @@ class RepoScraper: "after": cur, } data = self.client.graphql(Q.PR_FILES_QUERY, vars_) - item = ((data.get("data") or {}).get("repository") or {}).get("pullRequest") or {} + item = ((data.get("data") or {}).get("repository") or {}).get( + "pullRequest" + ) or {} ff = item.get("files") or {} for f in ff.get("nodes") or []: f["_owner"] = self.owner @@ -387,7 +417,9 @@ class RepoScraper: "after": cur, } data = self.client.graphql(Q.PR_REVIEW_THREADS_QUERY, vars_) - item = ((data.get("data") or {}).get("repository") or {}).get("pullRequest") or {} + item = ((data.get("data") or {}).get("repository") or {}).get( + "pullRequest" + ) or {} rt = item.get("reviewThreads") or {} for th in rt.get("nodes") or []: th["_owner"] = self.owner @@ -427,7 +459,9 @@ class RepoScraper: d["_fetchedAt"] = ts() num = d["number"] if d.get("comments", {}).get("pageInfo", {}).get("hasNextPage"): - self._paginate_discussion_comments(num, d["comments"]["pageInfo"]["endCursor"]) + self._paginate_discussion_comments( + num, d["comments"]["pageInfo"]["endCursor"] + ) # paginate replies per comment if needed for c in d.get("comments", {}).get("nodes", []) or []: if c.get("replies", {}).get("pageInfo", {}).get("hasNextPage"): @@ -465,7 +499,9 @@ class RepoScraper: "after": cur, } data = self.client.graphql(Q.DISCUSSION_COMMENTS_QUERY, vars_) - disc = ((data.get("data") or {}).get("repository") or {}).get("discussion") or {} + disc = ((data.get("data") or {}).get("repository") or {}).get( + "discussion" + ) or {} cc = disc.get("comments") or {} for c in cc.get("nodes") or []: c["_owner"] = self.owner @@ -475,7 +511,9 @@ class RepoScraper: info = cc.get("pageInfo") or {} cur = info.get("endCursor") if info.get("hasNextPage") else None - def _paginate_discussion_replies(self, comment_id: str, after: str, disc_number: int) -> None: + def _paginate_discussion_replies( + self, comment_id: str, after: str, disc_number: int + ) -> None: cur = after while cur: vars_ = { @@ -610,8 +648,12 @@ def setup_logging(log_file: Path) -> None: def main(): ap = argparse.ArgumentParser() - ap.add_argument("--base-dir", default = "/mnt/disks/unslothai/ubuntu/workspace_34/github_scraper") - ap.add_argument("--repos", nargs = "+", default = ["unslothai/unsloth", "unslothai/unsloth-zoo"]) + ap.add_argument( + "--base-dir", default = "/mnt/disks/unslothai/ubuntu/workspace_34/github_scraper" + ) + ap.add_argument( + "--repos", nargs = "+", default = ["unslothai/unsloth", "unslothai/unsloth-zoo"] + ) ap.add_argument("--trial", action = "store_true", help = "Small trial run") ap.add_argument( "--only", @@ -684,9 +726,15 @@ def main(): if not only or "commits" in only: default_ref = repo_meta.get("defaultBranchRef") or {} default_branch = ( - default_ref.get("name") if isinstance(default_ref, dict) else None + default_ref.get("name") + if isinstance(default_ref, dict) + else None + ) + branch = ( + f"refs/heads/{default_branch}" + if default_branch + else "refs/heads/main" ) - branch = f"refs/heads/{default_branch}" if default_branch else "refs/heads/main" scraper.scrape_commits(branch = branch) finally: scraper.close() diff --git a/studio/backend/plugins/data-designer-unstructured-seed/src/data_designer_unstructured_seed/chunking.py b/studio/backend/plugins/data-designer-unstructured-seed/src/data_designer_unstructured_seed/chunking.py index ee7d8727fc..a79a0059b5 100644 --- a/studio/backend/plugins/data-designer-unstructured-seed/src/data_designer_unstructured_seed/chunking.py +++ b/studio/backend/plugins/data-designer-unstructured-seed/src/data_designer_unstructured_seed/chunking.py @@ -42,7 +42,9 @@ def build_unstructured_preview_rows( try: import pandas as pd except ImportError as exc: # pragma: no cover - raise RuntimeError(f"pandas is required for unstructured seed processing: {exc}") from exc + raise RuntimeError( + f"pandas is required for unstructured seed processing: {exc}" + ) from exc dataframe = pd.read_parquet(parquet_path).head(count) return [ @@ -69,7 +71,9 @@ def build_multi_file_preview_rows( return _round_robin_preview(rows, preview_size) -def _round_robin_preview(rows: list[dict[str, str]], preview_size: int) -> list[dict[str, str]]: +def _round_robin_preview( + rows: list[dict[str, str]], preview_size: int +) -> list[dict[str, str]]: """Pick preview rows round-robin across source files so each is represented.""" if not rows or preview_size <= 0: return [] @@ -133,7 +137,9 @@ def materialize_unstructured_seed_dataset( try: import pandas as pd except ImportError as exc: # pragma: no cover - raise RuntimeError(f"pandas is required for unstructured seed processing: {exc}") from exc + raise RuntimeError( + f"pandas is required for unstructured seed processing: {exc}" + ) from exc tmp_path = _CACHE_DIR / f"{key}.tmp.parquet" pd.DataFrame(rows).to_parquet(tmp_path, index = False) @@ -192,7 +198,9 @@ def normalize_unstructured_text(text: str) -> str: return re.sub(r"\n{3,}", "\n\n", normalized).strip() -def split_text_into_chunks(*, text: str, chunk_size: int, chunk_overlap: int) -> list[str]: +def split_text_into_chunks( + *, text: str, chunk_size: int, chunk_overlap: int +) -> list[str]: if not text: return [] if chunk_size <= 0: @@ -246,7 +254,9 @@ def _to_int(value: Any, fallback: int) -> int: return parsed -def _compute_cache_key(*, source_path: Path, chunk_size: int, chunk_overlap: int) -> str: +def _compute_cache_key( + *, source_path: Path, chunk_size: int, chunk_overlap: int +) -> str: stat = source_path.stat() payload = "|".join( [ diff --git a/studio/backend/routes/auth.py b/studio/backend/routes/auth.py index 3ac85380fe..a971766c78 100644 --- a/studio/backend/routes/auth.py +++ b/studio/backend/routes/auth.py @@ -220,7 +220,9 @@ async def auth_status() -> AuthStatusResponse: return AuthStatusResponse( initialized = storage.is_initialized(), default_username = storage.DEFAULT_ADMIN_USERNAME, - requires_password_change = storage.requires_password_change(storage.DEFAULT_ADMIN_USERNAME) + requires_password_change = storage.requires_password_change( + storage.DEFAULT_ADMIN_USERNAME + ) if storage.is_initialized() else True, ) @@ -237,7 +239,10 @@ async def login(payload: AuthLoginRequest, request: Request) -> Token: status_code = status.HTTP_429_TOO_MANY_REQUESTS, # IP not interpolated into the body; behind a proxy/NAT it's # misleading or an info leak. - detail = (f"Too many failed login attempts. " f"Try again in {blocked_for} seconds."), + detail = ( + f"Too many failed login attempts. " + f"Try again in {blocked_for} seconds." + ), headers = {"Retry-After": str(blocked_for)}, ) @@ -273,7 +278,8 @@ async def login(payload: AuthLoginRequest, request: Request) -> Token: @router.post("/logout", status_code = status.HTTP_204_NO_CONTENT) async def logout( - request: Request, current_subject: str = Depends(get_current_subject_allow_password_change) + request: Request, + current_subject: str = Depends(get_current_subject_allow_password_change), ) -> Response: """Revoke refresh tokens for the subject; the access token is stateless and expires on its own.""" try: @@ -322,7 +328,9 @@ async def refresh(payload: RefreshTokenRequest) -> Token: access_token = new_access_token, refresh_token = new_refresh_token, token_type = "bearer", - must_change_password = False if is_desktop else storage.requires_password_change(username), + must_change_password = False + if is_desktop + else storage.requires_password_change(username), ) @@ -408,7 +416,9 @@ async def create_api_key( @router.get("/api-keys", response_model = ApiKeyListResponse) -async def list_api_keys(current_subject: str = Depends(get_current_subject)) -> ApiKeyListResponse: +async def list_api_keys( + current_subject: str = Depends(get_current_subject), +) -> ApiKeyListResponse: """List all API keys for the authenticated user (raw keys are never exposed).""" rows = storage.list_api_keys(current_subject) return ApiKeyListResponse( @@ -417,7 +427,9 @@ async def list_api_keys(current_subject: str = Depends(get_current_subject)) -> @router.delete("/api-keys/{key_id}") -async def revoke_api_key(key_id: int, current_subject: str = Depends(get_current_subject)) -> dict: +async def revoke_api_key( + key_id: int, current_subject: str = Depends(get_current_subject) +) -> dict: """Revoke (soft-delete) an API key.""" if not storage.revoke_api_key(current_subject, key_id): raise HTTPException( diff --git a/studio/backend/routes/chat_history.py b/studio/backend/routes/chat_history.py index 2ea572e30e..89e7cebd5b 100644 --- a/studio/backend/routes/chat_history.py +++ b/studio/backend/routes/chat_history.py @@ -163,7 +163,9 @@ class ChatSettingsPayload(BaseModel): inferenceParams: Optional[ChatInferenceSettings] = None customPresets: Optional[list[ChatPreset]] = None activePreset: Optional[str] = None - activePresetSource: Optional[Literal["builtin-default", "custom", "modified"]] = None + activePresetSource: Optional[Literal["builtin-default", "custom", "modified"]] = ( + None + ) autoTitle: Optional[bool] = None reasoningEffort: Optional[ Literal["none", "minimal", "low", "medium", "high", "max", "xhigh"] @@ -222,7 +224,9 @@ async def list_threads( @router.post("/threads", response_model = ChatThread) -async def save_thread(payload: ChatThread, current_subject: str = Depends(get_current_subject)): +async def save_thread( + payload: ChatThread, current_subject: str = Depends(get_current_subject) +): if payload.projectId and get_chat_project(payload.projectId) is None: raise HTTPException( status_code = 404, @@ -232,7 +236,9 @@ async def save_thread(payload: ChatThread, current_subject: str = Depends(get_cu @router.get("/threads/{thread_id}", response_model = ChatThread) -async def get_thread(thread_id: str, current_subject: str = Depends(get_current_subject)): +async def get_thread( + thread_id: str, current_subject: str = Depends(get_current_subject) +): thread = get_chat_thread(thread_id) if thread is None: raise HTTPException(status_code = 404, detail = f"Thread {thread_id} not found") @@ -273,7 +279,8 @@ async def delete_threads( @router.get("/projects", response_model = ChatProjectListResponse) async def list_projects( - include_archived: bool = Query(False), current_subject: str = Depends(get_current_subject) + include_archived: bool = Query(False), + current_subject: str = Depends(get_current_subject), ): return ChatProjectListResponse( projects = [ @@ -284,12 +291,16 @@ async def list_projects( @router.post("/projects", response_model = ChatProject) -async def save_project(payload: ChatProject, current_subject: str = Depends(get_current_subject)): +async def save_project( + payload: ChatProject, current_subject: str = Depends(get_current_subject) +): return ChatProject(**upsert_chat_project(payload.model_dump())) @router.get("/projects/{project_id}", response_model = ChatProject) -async def get_project(project_id: str, current_subject: str = Depends(get_current_subject)): +async def get_project( + project_id: str, current_subject: str = Depends(get_current_subject) +): project = ensure_chat_project_workspace(project_id) if project is None: raise HTTPException( @@ -360,12 +371,16 @@ async def delete_project( finally: conn.close() except Exception: # noqa: BLE001 - source cleanup must not block project deletion - logger.warning("failed to delete RAG sources for project %s", project_id, exc_info = True) + logger.warning( + "failed to delete RAG sources for project %s", project_id, exc_info = True + ) return ChatProject(**project) @router.get("/threads/{thread_id}/messages", response_model = ChatMessageListResponse) -async def get_thread_messages(thread_id: str, current_subject: str = Depends(get_current_subject)): +async def get_thread_messages( + thread_id: str, current_subject: str = Depends(get_current_subject) +): if get_chat_thread(thread_id) is None: raise HTTPException(status_code = 404, detail = f"Thread {thread_id} not found") return ChatMessageListResponse( @@ -375,7 +390,8 @@ async def get_thread_messages(thread_id: str, current_subject: str = Depends(get @router.post("/messages:batch", response_model = ChatMessagesBatchResponse) async def batch_thread_messages( - payload: ChatMessagesBatchRequest, current_subject: str = Depends(get_current_subject) + payload: ChatMessagesBatchRequest, + current_subject: str = Depends(get_current_subject), ): """One round-trip per sidebar/search rebuild instead of N. Unknown thread ids return empty lists.""" by_thread: dict[str, list[ChatMessage]] = {tid: [] for tid in payload.threadIds} @@ -429,10 +445,14 @@ async def replace_thread_messages( payload: ChatMessageSyncRequest, current_subject: str = Depends(get_current_subject), ): - mismatched_ids = [message.id for message in payload.messages if message.threadId != thread_id] + mismatched_ids = [ + message.id for message in payload.messages if message.threadId != thread_id + ] if mismatched_ids: preview = ", ".join(mismatched_ids[:5]) - suffix = "" if len(mismatched_ids) <= 5 else f" (+{len(mismatched_ids) - 5} more)" + suffix = ( + "" if len(mismatched_ids) <= 5 else f" (+{len(mismatched_ids) - 5} more)" + ) raise HTTPException( status_code = 400, detail = f"Message threadId mismatch: {preview}{suffix}", @@ -477,7 +497,8 @@ async def get_import_ledger(current_subject: str = Depends(get_current_subject)) @router.post("/import-ledger", response_model = ChatImportLedgerRecordResponse) async def record_import_ledger( - payload: ChatImportLedgerRecordRequest, current_subject: str = Depends(get_current_subject) + payload: ChatImportLedgerRecordRequest, + current_subject: str = Depends(get_current_subject), ): """Mark each legacy thread id as imported. Idempotent.""" accepted, inserted = upsert_chat_legacy_imports(payload.threadIds) diff --git a/studio/backend/routes/data_recipe/jobs.py b/studio/backend/routes/data_recipe/jobs.py index 59714380da..04203bf0b2 100644 --- a/studio/backend/routes/data_recipe/jobs.py +++ b/studio/backend/routes/data_recipe/jobs.py @@ -151,7 +151,9 @@ def _ensure_selected_local_model_loaded( ) -> None: model_loaded, active_model, active_variant = _loaded_local_model_identity() if not model_loaded: - raise ValueError("No model loaded in Chat. Load a model first, then run the recipe.") + raise ValueError( + "No model loaded in Chat. Load a model first, then run the recipe." + ) selection = _single_used_local_model_selection(recipe, local_provider_names) if selection is None: @@ -161,7 +163,9 @@ def _ensure_selected_local_model_loaded( variant_matches = not gguf_variant or active_variant == gguf_variant if active_model.lower() != target.lower() or not variant_matches: selected = f"{target} ({gguf_variant})" if gguf_variant else target - active = f"{active_model} ({active_variant})" if active_variant else active_model + active = ( + f"{active_model} ({active_variant})" if active_variant else active_model + ) raise ValueError( "Selected local model is not loaded. " f"Selected {selected}; active {active or 'none'}. " @@ -190,7 +194,9 @@ def _inject_local_structured_response_format( for mc in model_configs: if not isinstance(mc, dict): continue - if mc.get("provider") in local_provider_names and isinstance(mc.get("alias"), str): + if mc.get("provider") in local_provider_names and isinstance( + mc.get("alias"), str + ): alias_to_local_mc[mc["alias"]] = mc if not alias_to_local_mc: @@ -286,12 +292,18 @@ def _inject_local_providers(recipe: dict[str, Any], request: Request) -> Optiona # Only gate on model-loaded if a local provider is reachable from an LLM # column via a model_config. Orphan model_config nodes shouldn't block runs; # the recipe never calls /v1 for them. - local_names = {providers[i].get("name") for i in local_indices if providers[i].get("name")} + local_names = { + providers[i].get("name") for i in local_indices if providers[i].get("name") + } used_aliases = _used_llm_model_aliases(recipe) referenced_providers = { mc.get("provider") for mc in recipe.get("model_configs", []) - if (isinstance(mc, dict) and mc.get("provider") and mc.get("alias") in used_aliases) + if ( + isinstance(mc, dict) + and mc.get("provider") + and mc.get("alias") in used_aliases + ) } token = "" @@ -367,7 +379,9 @@ def _normalize_run_name(value: Any) -> str | None: if value is None: return None if not isinstance(value, str): - raise HTTPException(status_code = 400, detail = "invalid run_name: must be a string") + raise HTTPException( + status_code = 400, detail = "invalid run_name: must be a string" + ) trimmed = value.strip() if not trimmed: return None @@ -529,7 +543,9 @@ def publish_job_dataset(job_id: str, payload: PublishDatasetRequest): description = payload.description.strip() hf_token = payload.hf_token.strip() if isinstance(payload.hf_token, str) else None artifact_path = ( - payload.artifact_path.strip() if isinstance(payload.artifact_path, str) else None + payload.artifact_path.strip() + if isinstance(payload.artifact_path, str) + else None ) if not repo_id: @@ -540,7 +556,10 @@ def publish_job_dataset(job_id: str, payload: PublishDatasetRequest): mgr = get_job_manager() status = mgr.get_status(job_id) if status is not None: - if status.get("status") != "completed" or status.get("execution_type") != "full": + if ( + status.get("status") != "completed" + or status.get("execution_type") != "full" + ): raise HTTPException( status_code = 409, detail = "Only completed full runs can be published.", diff --git a/studio/backend/routes/data_recipe/mcp.py b/studio/backend/routes/data_recipe/mcp.py index 78a39877cc..2c79d323f3 100644 --- a/studio/backend/routes/data_recipe/mcp.py +++ b/studio/backend/routes/data_recipe/mcp.py @@ -69,7 +69,9 @@ def list_mcp_tools(payload: McpToolsListRequest) -> McpToolsListResponse: provider = built[0] try: tools = mcp_io.list_tools(provider, timeout_sec = payload.timeout_sec) - tool_names = sorted({tool.name for tool in tools if getattr(tool, "name", "")}) + tool_names = sorted( + {tool.name for tool in tools if getattr(tool, "name", "")} + ) for tool_name in tool_names: tool_to_providers[tool_name].append(provider.name) providers.append( diff --git a/studio/backend/routes/data_recipe/seed.py b/studio/backend/routes/data_recipe/seed.py index 8fb034ea4e..12b5765138 100644 --- a/studio/backend/routes/data_recipe/seed.py +++ b/studio/backend/routes/data_recipe/seed.py @@ -63,7 +63,9 @@ _SAFE_ID_RE = re.compile(r"^[a-zA-Z0-9_-]+$") def _validate_safe_id(value: str, label: str) -> str: if not value or not _SAFE_ID_RE.match(value): - raise HTTPException(400, f"Invalid {label}: must be alphanumeric/dash/underscore only") + raise HTTPException( + 400, f"Invalid {label}: must be alphanumeric/dash/underscore only" + ) return value @@ -73,7 +75,8 @@ def _serialize_preview_value(value: Any) -> Any: def _serialize_preview_rows(rows: list[dict[str, Any]]) -> list[dict[str, Any]]: return [ - {str(key): _serialize_preview_value(value) for key, value in row.items()} for row in rows + {str(key): _serialize_preview_value(value) for key, value in row.items()} + for row in rows ] @@ -194,7 +197,9 @@ def _decode_base64_payload(content_base64: str) -> bytes: raise HTTPException(status_code = 400, detail = "invalid base64 payload") from exc -def _read_preview_rows_from_local_file(path: Path, preview_size: int) -> list[dict[str, Any]]: +def _read_preview_rows_from_local_file( + path: Path, preview_size: int +) -> list[dict[str, Any]]: try: import pandas as pd except ImportError as exc: @@ -292,7 +297,9 @@ def _read_preview_rows_from_multi_files( for fid, fname in zip(file_ids, file_names): extracted = block_dir / f"{fid}.extracted.txt" if not extracted.exists(): - raise HTTPException(404, f"Extracted text not found for file: {fname} (id: {fid})") + raise HTTPException( + 404, f"Extracted text not found for file: {fname} (id: {fid})" + ) file_entries.append((extracted, fname)) return build_multi_file_preview_rows( @@ -372,7 +379,9 @@ def inspect_seed_dataset(payload: SeedInspectRequest) -> SeedInspectResponse: ) from exc if not preview_rows: - raise HTTPException(status_code = 422, detail = "dataset appears empty or unreadable") + raise HTTPException( + status_code = 422, detail = "dataset appears empty or unreadable" + ) preview_rows = _serialize_preview_rows(preview_rows) columns = _extract_columns(preview_rows) @@ -381,7 +390,9 @@ def inspect_seed_dataset(payload: SeedInspectRequest) -> SeedInspectResponse: else: resolved_path = _resolve_seed_hf_path(dataset_name, data_files, split) if not resolved_path: - raise HTTPException(status_code = 422, detail = "unable to resolve seed dataset path") + raise HTTPException( + status_code = 422, detail = "unable to resolve seed dataset path" + ) return SeedInspectResponse( dataset_name = dataset_name, @@ -500,7 +511,9 @@ async def upload_unstructured_file( try: meta_path = block_dir / f"{file_id}.meta.json" meta_path.write_text( - json.dumps({"original_filename": original_filename, "size_bytes": size_bytes}), + json.dumps( + {"original_filename": original_filename, "size_bytes": size_bytes} + ), encoding = "utf-8", ) except OSError: @@ -625,7 +638,9 @@ def inspect_seed_upload(payload: SeedInspectUploadRequest) -> SeedInspectRespons int(payload.preview_size), ) if not preview_rows: - raise HTTPException(status_code = 422, detail = "dataset appears empty or unreadable") + raise HTTPException( + status_code = 422, detail = "dataset appears empty or unreadable" + ) columns = _extract_columns(preview_rows) return SeedInspectResponse( diff --git a/studio/backend/routes/data_recipe/validate.py b/studio/backend/routes/data_recipe/validate.py index ffe36ba69e..5c2469ef76 100644 --- a/studio/backend/routes/data_recipe/validate.py +++ b/studio/backend/routes/data_recipe/validate.py @@ -21,9 +21,7 @@ from utils.utils import safe_error_detail, safe_curated_detail, log_and_http_err logger = get_logger(__name__) router = APIRouter() -_GITHUB_VALIDATE_NOTE = ( - "Recipe shape is valid. GitHub access and rate limits are checked when the run starts." -) +_GITHUB_VALIDATE_NOTE = "Recipe shape is valid. GitHub access and rate limits are checked when the run starts." _GITHUB_ITEM_TYPES = {"issues", "pulls", "commits"} @@ -46,17 +44,23 @@ def _validate_github_seed_static(source: dict[str, Any]) -> list[ValidateError]: else: for repo in repos: if not isinstance(repo, str) or not repo.strip() or "/" not in repo: - errors.append(ValidateError(message = "GitHub repos must be owner/name strings.")) + errors.append( + ValidateError(message = "GitHub repos must be owner/name strings.") + ) break item_types = source.get("item_types") if not isinstance(item_types, list) or not item_types: - errors.append(ValidateError(message = "GitHub seed requires at least one item type.")) + errors.append( + ValidateError(message = "GitHub seed requires at least one item type.") + ) else: invalid_items = [item for item in item_types if item not in _GITHUB_ITEM_TYPES] if invalid_items: errors.append( - ValidateError(message = "GitHub item types must be issues, pulls, or commits.") + ValidateError( + message = "GitHub item types must be issues, pulls, or commits." + ) ) try: diff --git a/studio/backend/routes/datasets.py b/studio/backend/routes/datasets.py index 46319ca2ba..5ed391e352 100644 --- a/studio/backend/routes/datasets.py +++ b/studio/backend/routes/datasets.py @@ -261,7 +261,11 @@ def _select_best_hf_preview_candidate( def _select_hf_preview_file( - repo_files: list[str], *, metadata: dict | None, subset: str | None, split: str | None + repo_files: list[str], + *, + metadata: dict | None, + subset: str | None, + split: str | None, ) -> str | None: normalized_repo_files = [_normalize_hf_repo_path(path) for path in repo_files] repo_file_set = set(normalized_repo_files) @@ -272,13 +276,21 @@ def _select_hf_preview_file( if path in repo_file_set and _is_hf_preview_data_file(path) ] if metadata_candidates: - return _select_best_hf_preview_candidate(metadata_candidates, subset = subset, split = split) + return _select_best_hf_preview_candidate( + metadata_candidates, subset = subset, split = split + ) - data_candidates = [path for path in normalized_repo_files if _is_hf_preview_data_file(path)] - return _select_best_hf_preview_candidate(data_candidates, subset = subset, split = split) + data_candidates = [ + path for path in normalized_repo_files if _is_hf_preview_data_file(path) + ] + return _select_best_hf_preview_candidate( + data_candidates, subset = subset, split = split + ) -def _download_hf_metadata(*, repo_id: str, repo_files: list[str], token: str | None) -> dict | None: +def _download_hf_metadata( + *, repo_id: str, repo_files: list[str], token: str | None +) -> dict | None: metadata_file = next( ( path @@ -411,7 +423,9 @@ def _build_local_dataset_items() -> list[LocalDatasetItem]: return items -def _load_local_preview_slice(*, dataset_path: Path, train_split: str, preview_size: int): +def _load_local_preview_slice( + *, dataset_path: Path, train_split: str, preview_size: int +): # Non-streaming loads take the cached builder lock; use the EACCES-safe wrapper. from utils.datasets.cache_safe import load_dataset_cache_safe as load_dataset @@ -447,7 +461,9 @@ def _load_local_preview_slice(*, dataset_path: Path, train_split: str, preview_s elif dataset_path.suffix == ".csv": dataset = load_dataset("csv", data_files = str(dataset_path), split = train_split) elif dataset_path.suffix == ".parquet": - dataset = load_dataset("parquet", data_files = str(dataset_path), split = train_split) + dataset = load_dataset( + "parquet", data_files = str(dataset_path), split = train_split + ) else: raise HTTPException( status_code = 400, detail = f"Unsupported file format: {dataset_path.suffix}" @@ -524,7 +540,9 @@ def list_local_datasets( @router.get("/download-progress") async def get_dataset_download_progress( - repo_id: str = Query(..., description = "HuggingFace dataset repo ID, e.g. 'unsloth/LaTeX_OCR'"), + repo_id: str = Query( + ..., description = "HuggingFace dataset repo ID, e.g. 'unsloth/LaTeX_OCR'" + ), current_subject: str = Depends(get_current_subject), ): """Return download progress for a HuggingFace dataset repo. @@ -599,7 +617,9 @@ async def get_dataset_download_progress( @router.post("/check-format", response_model = CheckFormatResponse) -def check_format(request: CheckFormatRequest, current_subject: str = Depends(get_current_subject)): +def check_format( + request: CheckFormatRequest, current_subject: str = Depends(get_current_subject) +): """Check if a dataset requires manual column mapping. HuggingFace strategy: @@ -721,7 +741,9 @@ def check_format(request: CheckFormatRequest, current_subject: str = Depends(get processed = format_result["dataset"] preview_samples = _serialize_preview_rows(processed) except Exception as e: - logger.warning(f"Processed preview generation failed (non-fatal): {e}") + logger.warning( + f"Processed preview generation failed (non-fatal): {e}" + ) preview_samples = _serialize_preview_rows(preview_slice) else: preview_samples = _serialize_preview_rows(preview_slice) @@ -732,7 +754,9 @@ def check_format(request: CheckFormatRequest, current_subject: str = Depends(get if image_col and image_col in (result.get("columns") or []): try: sample_val = preview_slice[0][image_col] - if isinstance(sample_val, str) and sample_val.startswith(("http://", "https://")): + if isinstance(sample_val, str) and sample_val.startswith( + ("http://", "https://") + ): url_warning = ( "This dataset contains image URLs instead of embedded images. " "Images will be downloaded during training, which may be slow for large datasets." @@ -782,7 +806,8 @@ def ai_assist_mapping( # Truncate sample values for the LLM prompt. truncated = [ - {col: str(s.get(col, ""))[:200] for col in request.columns} for s in request.samples[:5] + {col: str(s.get(col, ""))[:200] for col in request.columns} + for s in request.samples[:5] ] result = llm_conversion_advisor( diff --git a/studio/backend/routes/export.py b/studio/backend/routes/export.py index f7b3a56a71..5f6657551b 100644 --- a/studio/backend/routes/export.py +++ b/studio/backend/routes/export.py @@ -80,7 +80,9 @@ async def load_checkpoint( break await asyncio.sleep(0.5) else: - logger.warning("Training subprocess did not exit within 30s, proceeding anyway") + logger.warning( + "Training subprocess did not exit within 30s, proceeding anyway" + ) except Exception as e: logger.warning("Could not stop training: %s", e) @@ -193,7 +195,8 @@ def _export_details(output_path: Optional[str]) -> Optional[Dict[str, Any]]: @router.post("/export/merged", response_model = ExportOperationResponse) async def export_merged_model( - request: ExportMergedModelRequest, current_subject: str = Depends(get_current_subject) + request: ExportMergedModelRequest, + current_subject: str = Depends(get_current_subject), ): """Export a merged PEFT model (16-bit or 4-bit), optionally pushing to Hub. @@ -306,7 +309,8 @@ async def export_gguf( @router.post("/export/lora", response_model = ExportOperationResponse) async def export_lora_adapter( - request: ExportLoRAAdapterRequest, current_subject: str = Depends(get_current_subject) + request: ExportLoRAAdapterRequest, + current_subject: str = Depends(get_current_subject), ): """Export only the LoRA adapter (if the loaded model is PEFT). diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index e9353cd803..7f639ec94e 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -99,7 +99,9 @@ def _loaded_chat_template() -> Optional[str]: return None -def _template_raise_message(error_text: str, chat_template: Optional[str]) -> Optional[str]: +def _template_raise_message( + error_text: str, chat_template: Optional[str] +) -> Optional[str]: """A chat-template raise_exception message to surface, but only when it appears verbatim in chat_template (simple substring check), so we never leak arbitrary llama-server text. Anchors on llama.cpp's "Jinja Exception:" prefix.""" @@ -139,9 +141,7 @@ def _friendly_error(exc: Exception) -> str: # WriteError, PoolTimeout, ...) means the llama-server subprocess is # unreachable -- crashed or still coming up. if isinstance(exc, httpx.RequestError): - return ( - "Lost connection to the model server. It may have crashed -- try reloading the model." - ) + return "Lost connection to the model server. It may have crashed -- try reloading the model." msg = str(exc) m = _re.search( r"request \((\d+) tokens?\) exceeds the available context size \((\d+) tokens?\)", @@ -154,9 +154,7 @@ def _friendly_error(exc: Exception) -> str: f"or shorten the conversation." ) if "Lost connection to llama-server" in msg: - return ( - "Lost connection to the model server. It may have crashed -- try reloading the model." - ) + return "Lost connection to the model server. It may have crashed -- try reloading the model." template_msg = _template_raise_message(msg, _loaded_chat_template()) if template_msg: return f"An internal error occurred: {template_msg}" @@ -226,7 +224,9 @@ def _raise_unsupported_openai_parameter(param: str, message: str) -> None: def _raise_unsupported_n(path_label: str) -> None: - _raise_unsupported_openai_parameter("n", f"n > 1 is not supported for {path_label}.") + _raise_unsupported_openai_parameter( + "n", f"n > 1 is not supported for {path_label}." + ) def _openai_stream_error_chunk(exc) -> dict: @@ -281,7 +281,10 @@ def _overflow_truncation_requested(payload) -> bool: requested = getattr(payload, "context_overflow", None) if requested is not None: return requested == "truncate_middle" - return os.environ.get("UNSLOTH_CONTEXT_OVERFLOW", "").strip().lower() == "truncate_middle" + return ( + os.environ.get("UNSLOTH_CONTEXT_OVERFLOW", "").strip().lower() + == "truncate_middle" + ) def _parse_overflow_counts(err_text: str): @@ -386,7 +389,9 @@ def _clip_long_contents(messages: list, target_est: int) -> int: if sum(_estimate_message_tokens(m) for m in messages) <= target_est: return clipped content = msg.get("content") - if not isinstance(content, str) or len(content) <= 2 * keep + len(_CLIP_MARKER): + if not isinstance(content, str) or len(content) <= 2 * keep + len( + _CLIP_MARKER + ): continue msg["content"] = content[:keep] + _CLIP_MARKER + content[-keep:] clipped += 1 @@ -402,7 +407,9 @@ def _apply_overflow_truncation(body: dict, err_text: str) -> bool: total_est = sum(_estimate_message_tokens(m) for m in messages) if counts: n_prompt, n_ctx = counts - keep_ratio = min(0.95, (_OVERFLOW_PROMPT_TARGET_FRACTION * n_ctx) / max(1, n_prompt)) + keep_ratio = min( + 0.95, (_OVERFLOW_PROMPT_TARGET_FRACTION * n_ctx) / max(1, n_prompt) + ) # Scale the server-token target into char-estimate units. target_est = int(total_est * keep_ratio) else: @@ -414,7 +421,10 @@ def _apply_overflow_truncation(body: dict, err_text: str) -> bool: if dropped: body["messages"] = new_messages clipped = 0 - if sum(_estimate_message_tokens(m) for m in body.get("messages") or []) > target_est: + if ( + sum(_estimate_message_tokens(m) for m in body.get("messages") or []) + > target_est + ): clipped = _clip_long_contents(body.get("messages") or [], target_est) if not dropped and not clipped: return False @@ -455,7 +465,9 @@ def _drop_parallel_tool_call_deltas(chunk) -> bool: delta = ch.get("delta") or {} tcs = delta.get("tool_calls") if isinstance(tcs, list): - kept = [tc for tc in tcs if isinstance(tc, dict) and (tc.get("index") or 0) == 0] + kept = [ + tc for tc in tcs if isinstance(tc, dict) and (tc.get("index") or 0) == 0 + ] if len(kept) != len(tcs): delta["tool_calls"] = kept changed = True @@ -515,7 +527,9 @@ def _openai_stream_usage_chunk( prompt_tokens = _prompt_tokens, completion_tokens = _completion_tokens, total_tokens = _total_tokens, - prompt_tokens_details = _prompt_tokens_details(_usage.get("prompt_tokens_details")), + prompt_tokens_details = _prompt_tokens_details( + _usage.get("prompt_tokens_details") + ), ), timings = stream_timings, ) @@ -586,7 +600,8 @@ def _classify_llama_generation_error(exc: Exception) -> Optional[bool]: msg = str(exc) msg_l = msg.lower() if "n_ctx" in msg_l or ( - "context" in msg_l and any(t in msg_l for t in ("exceed", "length", "window", "too long")) + "context" in msg_l + and any(t in msg_l for t in ("exceed", "length", "window", "too long")) ): return True if _re.search(r"llama-server returned (4\d\d)", msg): @@ -695,7 +710,9 @@ def _set_stream_response_read_timeout( pass -async def _preheader_cancelled(cancel_event = None, request: Optional[Request] = None) -> bool: +async def _preheader_cancelled( + cancel_event = None, request: Optional[Request] = None +) -> bool: if cancel_event is not None and cancel_event.is_set(): return True if request is not None and await request.is_disconnected(): @@ -705,7 +722,9 @@ async def _preheader_cancelled(cancel_event = None, request: Optional[Request] = return False -async def _wait_preheader_cancel(cancel_event = None, request: Optional[Request] = None) -> None: +async def _wait_preheader_cancel( + cancel_event = None, request: Optional[Request] = None +) -> None: while not await _preheader_cancelled(cancel_event, request): await asyncio.sleep(0.05) @@ -782,15 +801,21 @@ async def _aiter_llama_stream_items( if waiting_first_item: remaining_s = first_token_deadline - time.monotonic() if remaining_s <= 0: - raise httpx.ReadTimeout("The model did not produce a first token in time.") + raise httpx.ReadTimeout( + "The model did not produce a first token in time." + ) if response is not None: _set_stream_response_read_timeout(response, remaining_s) - item = await asyncio.wait_for(async_iter.__anext__(), timeout = remaining_s) + item = await asyncio.wait_for( + async_iter.__anext__(), timeout = remaining_s + ) else: item = await async_iter.__anext__() except asyncio.TimeoutError as exc: if waiting_first_item: - raise httpx.ReadTimeout("The model did not produce a first token in time.") from exc + raise httpx.ReadTimeout( + "The model did not produce a first token in time." + ) from exc raise except StopAsyncIteration: return @@ -1005,7 +1030,9 @@ async def artifact_preview_frame( await get_current_subject(creds) csp = ( - _ARTIFACT_PREVIEW_FRAME_NETWORK_CSP if allow_network else _ARTIFACT_PREVIEW_FRAME_STRICT_CSP + _ARTIFACT_PREVIEW_FRAME_NETWORK_CSP + if allow_network + else _ARTIFACT_PREVIEW_FRAME_STRICT_CSP ) return Response( content = _ARTIFACT_PREVIEW_FRAME_HTML, @@ -1095,7 +1122,9 @@ _PENDING_CANCEL_TTL_S = 30.0 def _prune_pending(now: float) -> None: - for k in [k for k, ts in _PENDING_CANCELS.items() if now - ts > _PENDING_CANCEL_TTL_S]: + for k in [ + k for k, ts in _PENDING_CANCELS.items() if now - ts > _PENDING_CANCEL_TTL_S + ]: _PENDING_CANCELS.pop(k, None) @@ -1240,7 +1269,9 @@ def _build_tool_action_nudge(*, tools: list[dict], model_name: str) -> str: compact_web_tip = model_size_b is not None and model_size_b < 9 tool_tip_parts: list[str] = [] if has_web: - tool_tip_parts.append(_TOOL_WEB_COMPACT_TIP if compact_web_tip else _TOOL_WEB_EXPANDED_TIP) + tool_tip_parts.append( + _TOOL_WEB_COMPACT_TIP if compact_web_tip else _TOOL_WEB_EXPANDED_TIP + ) if has_code: tool_tip_parts.append(_TOOL_CODE_TIP) if has_artifact: @@ -1329,7 +1360,9 @@ def _normalise_settings_str(value: Optional[str]) -> Optional[str]: return value -def _should_strip_split_mode(request: LoadRequest, backend_extra: Optional[list[str]]) -> bool: +def _should_strip_split_mode( + request: LoadRequest, backend_extra: Optional[list[str]] +) -> bool: """Whether an inherited --split-mode should be stripped on reload. The binary Tensor Parallelism toggle can't carry --split-mode's row/none/ @@ -1474,7 +1507,9 @@ def _resolve_model_identifier_for_request( status_code = 400, detail = redact_native_paths(str(exc)), ) from exc - display_label = grant.display_label or Path(request.model_path).name or "Native model" + display_label = ( + grant.display_label or Path(request.model_path).name or "Native model" + ) return str(grant.canonical_path), display_label, True @@ -1568,7 +1603,9 @@ async def load_model( inference_config = load_inference_config(llama_backend.model_identifier) _gguf_audio = ( - llama_backend._audio_type if hasattr(llama_backend, "_audio_type") else None + llama_backend._audio_type + if hasattr(llama_backend, "_audio_type") + else None ) _gguf_is_audio = getattr(llama_backend, "_is_audio", False) return LoadResponse( @@ -1608,7 +1645,9 @@ async def load_model( backend.active_model_name and backend.active_model_name.lower() == model_identifier.lower() ): - logger.info(f"Model already loaded (Unsloth): {model_log_label}, skipping reload") + logger.info( + f"Model already loaded (Unsloth): {model_log_label}, skipping reload" + ) inference_config = load_inference_config(backend.active_model_name) _model_info = backend.models.get(backend.active_model_name, {}) _chat_template = None @@ -1625,7 +1664,9 @@ async def load_model( _sf_reasoning_style = _sf_flags["reasoning_style"] return LoadResponse( status = "already_loaded", - model = model_log_label if native_grant_backed else backend.active_model_name, + model = model_log_label + if native_grant_backed + else backend.active_model_name, display_name = model_log_label if native_grant_backed else backend.active_model_name, @@ -1644,7 +1685,9 @@ async def load_model( reasoning_always_on = _sf_flags["reasoning_always_on"], supports_preserve_thinking = _sf_flags["supports_preserve_thinking"], supports_tools = _sf_flags["supports_tools"], - context_length = _positive_int_or_none(_model_info.get("context_length")), + context_length = _positive_int_or_none( + _model_info.get("context_length") + ), chat_template = _chat_template, ) @@ -1703,12 +1746,16 @@ async def load_model( request_variant = (request.gguf_variant or "").lower() stored_variant = (source[1] or "").lower() if source else "" same_model = bool( - source and source[0] and source[0].lower() == model_identifier.lower() + source + and source[0] + and source[0].lower() == model_identifier.lower() ) if request.gguf_variant: variant_mismatch = request_variant != stored_variant else: - variant_mismatch = bool(stored_variant and resolved_variant != stored_variant) + variant_mismatch = bool( + stored_variant and resolved_variant != stored_variant + ) same_source = same_model and not variant_mismatch if not same_source: logger.info( @@ -1735,7 +1782,8 @@ async def load_model( strip_context = "max_seq_length" in fields_set, strip_cache = "cache_type_kv" in fields_set, strip_spec = ( - "speculative_type" in fields_set or "spec_draft_n_max" in fields_set + "speculative_type" in fields_set + or "spec_draft_n_max" in fields_set ), strip_template = ( "chat_template_override" in fields_set @@ -1793,7 +1841,9 @@ async def load_model( if native_grant_backed: if config.gguf_mmproj_file: _validate_native_gguf_companion( - config.gguf_mmproj_file, config.gguf_file, "vision companion" + config.gguf_mmproj_file, + config.gguf_file, + "vision companion", ) if config.gguf_mtp_file: # The drafter is optional (unlike mmproj for a vision @@ -1803,7 +1853,9 @@ async def load_model( config.gguf_mtp_file, config.gguf_file, "MTP drafter" ) except HTTPException as exc: - logger.warning("Dropping MTP drafter for native load: %s", exc.detail) + logger.warning( + "Dropping MTP drafter for native load: %s", exc.detail + ) config.gguf_mtp_file = None _source_load_kwargs = dict( gguf_path = config.gguf_file, @@ -1856,7 +1908,9 @@ async def load_model( # Audio detection moved into load_model under _serial_load_lock (#5642). _gguf_audio = llama_backend._audio_type _gguf_is_audio = llama_backend._is_audio - llama_backend._native_display_label = model_log_label if native_grant_backed else None + llama_backend._native_display_label = ( + model_log_label if native_grant_backed else None + ) llama_backend._native_grant_backed = bool(native_grant_backed) if _gguf_is_audio: logger.info(f"GGUF model detected as audio: audio_type={_gguf_audio}") @@ -1866,7 +1920,9 @@ async def load_model( return LoadResponse( status = "loaded", model = model_log_label if native_grant_backed else config.identifier, - display_name = model_log_label if native_grant_backed else config.display_name, + display_name = model_log_label + if native_grant_backed + else config.display_name, is_vision = llama_backend.is_vision, is_lora = False, is_gguf = True, @@ -1875,7 +1931,9 @@ async def load_model( audio_type = _gguf_audio, has_audio_input = llama_backend._has_audio_input, inference = inference_config, - requires_trust_remote_code = bool(inference_config.get("trust_remote_code", False)), + requires_trust_remote_code = bool( + inference_config.get("trust_remote_code", False) + ), context_length = llama_backend.context_length, max_context_length = llama_backend.max_context_length, native_context_length = llama_backend.native_context_length, @@ -1905,7 +1963,9 @@ async def load_model( from core.export import get_export_backend exp_backend = get_export_backend() if exp_backend.current_checkpoint: - logger.info("Shutting down export subprocess to free GPU memory for inference") + logger.info( + "Shutting down export subprocess to free GPU memory for inference" + ) exp_backend._shutdown_subprocess() exp_backend.current_checkpoint = None exp_backend.is_vision = False @@ -1975,7 +2035,9 @@ async def load_model( # Check if YAML says this model needs trust_remote_code. if not request.trust_remote_code: model_defaults = load_model_defaults(config.identifier) - yaml_trust = model_defaults.get("inference", {}).get("trust_remote_code", False) + yaml_trust = model_defaults.get("inference", {}).get( + "trust_remote_code", False + ) if yaml_trust: raise HTTPException( status_code = 400, @@ -2011,7 +2073,9 @@ async def load_model( return LoadResponse( status = "loaded", model = model_log_label if native_grant_backed else config.identifier, - display_name = model_log_label if native_grant_backed else config.display_name, + display_name = model_log_label + if native_grant_backed + else config.display_name, is_vision = config.is_vision, is_lora = config.is_lora, is_gguf = False, @@ -2019,7 +2083,9 @@ async def load_model( audio_type = config.audio_type, has_audio_input = config.has_audio_input, inference = inference_config, - requires_trust_remote_code = bool(inference_config.get("trust_remote_code", False)), + requires_trust_remote_code = bool( + inference_config.get("trust_remote_code", False) + ), supports_reasoning = _sf_flags["supports_reasoning"], reasoning_style = _sf_flags["reasoning_style"], reasoning_always_on = _sf_flags["reasoning_always_on"], @@ -2149,7 +2215,9 @@ async def validate_model( @router.post("/unload", response_model = UnloadResponse) -async def unload_model(request: UnloadRequest, current_subject: str = Depends(get_current_subject)): +async def unload_model( + request: UnloadRequest, current_subject: str = Depends(get_current_subject) +): """ Unload a model from memory. Routes to the correct backend (llama-server for GGUF, Unsloth otherwise). @@ -2159,7 +2227,9 @@ async def unload_model(request: UnloadRequest, current_subject: str = Depends(ge llama_backend = get_llama_cpp_backend() if llama_backend.is_active and ( llama_backend.model_identifier == request.model_path - or is_registered_native_path_label(llama_backend.model_identifier, request.model_path) + or is_registered_native_path_label( + llama_backend.model_identifier, request.model_path + ) or not llama_backend.is_loaded ): llama_backend.unload_model() @@ -2178,7 +2248,9 @@ async def unload_model(request: UnloadRequest, current_subject: str = Depends(ge @studio_router.post("/cancel") -async def cancel_inference(request: Request, current_subject: str = Depends(get_current_subject)): +async def cancel_inference( + request: Request, current_subject: str = Depends(get_current_subject) +): """Cancel in-flight inference requests. Body (JSON, at least one key required): @@ -2417,13 +2489,17 @@ async def get_status(current_subject: str = Depends(get_current_subject)): has_audio_input = model_info.get("has_audio_input", False) chat_template_info = model_info.get("chat_template_info", {}) chat_template = ( - chat_template_info.get("template") if isinstance(chat_template_info, dict) else None + chat_template_info.get("template") + if isinstance(chat_template_info, dict) + else None ) # Non-GGUF: classify from the loaded template. _sf_flags = _detect_safetensors_features(backend, chat_template) inference_config = ( - load_inference_config(backend.active_model_name) if backend.active_model_name else None + load_inference_config(backend.active_model_name) + if backend.active_model_name + else None ) return InferenceStatusResponse( @@ -2504,7 +2580,9 @@ async def generate_audio( _, chat_messages, _ = _extract_content_parts(payload.messages) if not chat_messages: raise HTTPException(status_code = 400, detail = "No messages provided.") - last_user_msg = next((m for m in reversed(chat_messages) if m["role"] == "user"), None) + last_user_msg = next( + (m for m in reversed(chat_messages) if m["role"] == "user"), None + ) if not last_user_msg: raise HTTPException(status_code = 400, detail = "No user message found.") text = last_user_msg["content"] @@ -2529,7 +2607,9 @@ async def generate_audio( raise HTTPException(status_code = 400, detail = "No model loaded.") model_info = backend.models.get(backend.active_model_name, {}) if not model_info.get("is_audio"): - raise HTTPException(status_code = 400, detail = "Active model is not an audio model.") + raise HTTPException( + status_code = 400, detail = "Active model is not an audio model." + ) model_name = backend.active_model_name gen = lambda: backend.generate_audio_response( text = text, @@ -2543,7 +2623,9 @@ async def generate_audio( ) try: - wav_bytes, sample_rate = await asyncio.get_event_loop().run_in_executor(None, gen) + wav_bytes, sample_rate = await asyncio.get_event_loop().run_in_executor( + None, gen + ) except Exception as e: logger.error(f"Audio generation error: {e}", exc_info = True) raise HTTPException(status_code = 500, detail = safe_error_detail(e)) @@ -2627,7 +2709,9 @@ def _sniff_audio_container(raw: bytes) -> Optional[str]: return "wav" # mp3: ID3 tag, or an MPEG audio frame sync (no other accepted format leads # with 0xFF, so the simple sync check doesn't collide). - if raw[:3] == b"ID3" or (len(raw) >= 2 and raw[0] == 0xFF and (raw[1] & 0xE0) == 0xE0): + if raw[:3] == b"ID3" or ( + len(raw) >= 2 and raw[0] == 0xFF and (raw[1] & 0xE0) == 0xE0 + ): return "mp3" return None @@ -2641,7 +2725,9 @@ def _mono_f32_to_wav_bytes(arr: np.ndarray, sample_rate: int) -> bytes: import io import wave - arr = np.nan_to_num(np.asarray(arr, dtype = np.float32).flatten(), posinf = 0.0, neginf = 0.0) + arr = np.nan_to_num( + np.asarray(arr, dtype = np.float32).flatten(), posinf = 0.0, neginf = 0.0 + ) if arr.size == 0: raise ValueError("decoded audio is empty") peak = float(np.abs(arr).max()) @@ -2658,7 +2744,9 @@ def _mono_f32_to_wav_bytes(arr: np.ndarray, sample_rate: int) -> bytes: return buf.getvalue() -def _resample_mono_linear(arr: np.ndarray, source_rate: int, target_rate: int) -> np.ndarray: +def _resample_mono_linear( + arr: np.ndarray, source_rate: int, target_rate: int +) -> np.ndarray: """Small numpy-only resampler for upload size limiting.""" if source_rate <= 0 or target_rate <= 0 or source_rate == target_rate: return arr @@ -2671,7 +2759,9 @@ def _resample_mono_linear(arr: np.ndarray, source_rate: int, target_rate: int) - return np.interp(target_x, source_x, arr).astype(np.float32) -def _fit_transcoded_audio_to_wav_cap(arr: np.ndarray, sample_rate: int) -> tuple[np.ndarray, int]: +def _fit_transcoded_audio_to_wav_cap( + arr: np.ndarray, sample_rate: int +) -> tuple[np.ndarray, int]: """Downsample only when needed so transcoded WAV stays within the upload cap.""" if sample_rate <= 0: raise ValueError("decoded audio has an invalid sample rate") @@ -2731,7 +2821,9 @@ def _decode_audio_mono(raw: bytes) -> tuple[np.ndarray, int]: if arr.ndim > 1: arr = arr.mean(axis = 1) if sr > 0 and len(arr) > sr * _MAX_AUDIO_SECONDS: - raise ValueError(f"decoded audio exceeds the {_MAX_AUDIO_SECONDS // 60}-minute limit") + raise ValueError( + f"decoded audio exceeds the {_MAX_AUDIO_SECONDS // 60}-minute limit" + ) return arr, sr @@ -2798,7 +2890,9 @@ def _extract_content_parts(messages: list) -> tuple[str, list[dict], "Optional[s system_parts.append(msg.content) elif isinstance(msg.content, list): # Unlikely but handle: join text parts - system_parts.append("\n".join(p.text for p in msg.content if p.type == "text")) + system_parts.append( + "\n".join(p.text for p in msg.content if p.type == "text") + ) continue # ── User / assistant messages ───────────────────────── @@ -2817,7 +2911,9 @@ def _extract_content_parts(messages: list) -> tuple[str, list[dict], "Optional[s # data:image/png;base64, -> extract first_image_b64 = url.split(",", 1)[1] if "," in url else None else: - logger.warning(f"Remote image URLs not yet supported: {url[:80]}...") + logger.warning( + f"Remote image URLs not yet supported: {url[:80]}..." + ) combined_text = "\n".join(text_parts) if text_parts else "" chat_messages.append({"role": msg.role, "content": combined_text}) @@ -2980,7 +3076,11 @@ def _build_external_messages( # (some providers reject empty assistant turns). Preserve assistant # turns whose only payload is tool_calls so multi-turn # function-call loops round-trip. - if msg.role == "assistant" and not msg.content.strip() and not msg.tool_calls: + if ( + msg.role == "assistant" + and not msg.content.strip() + and not msg.tool_calls + ): continue out: dict[str, Any] = {"role": msg.role, "content": msg.content} if msg.role == "assistant" and msg.tool_calls: @@ -3033,7 +3133,9 @@ def _build_external_messages( "image_url": {"url": part.image_url.url}, } ) - elif part.type == "reasoning" and openai and msg.role == "assistant": + elif ( + part.type == "reasoning" and openai and msg.role == "assistant" + ): reasoning: dict[str, Any] = { "type": "reasoning", "id": part.id, @@ -3043,7 +3145,9 @@ def _build_external_messages( reasoning["status"] = part.status parts.append(reasoning) elif ( - part.type == "image_generation_call" and openai and msg.role == "assistant" + part.type == "image_generation_call" + and openai + and msg.role == "assistant" ): # ExternalProviderClient maps this onto a top-level # Responses input item after the current user prompt, @@ -3110,7 +3214,11 @@ def _build_external_messages( if p.status: reasoning["status"] = p.status preserved.append(reasoning) - elif p.type == "image_generation_call" and openai and msg.role == "assistant": + elif ( + p.type == "image_generation_call" + and openai + and msg.role == "assistant" + ): image_ref = {"type": "image_generation_call", "id": p.id} if getattr(p, "response_id", None): image_ref["response_id"] = p.response_id @@ -3135,7 +3243,9 @@ def _build_external_messages( _entry_content = entry.get("content") _has_text = ( isinstance(_entry_content, str) and _entry_content.strip() - ) or (isinstance(_entry_content, list) and len(_entry_content) > 0) + ) or ( + isinstance(_entry_content, list) and len(_entry_content) > 0 + ) if not _has_text: continue if msg.role == "tool": @@ -3290,7 +3400,9 @@ async def _proxy_to_external_provider( # ── OpenAI shell-tool container management ─────────────────────── -def _resolve_openai_cloud_client(body: OpenAIContainerRequest) -> ExternalProviderClient: +def _resolve_openai_cloud_client( + body: OpenAIContainerRequest, +) -> ExternalProviderClient: """ Decrypt the API key + validate the base URL points at OpenAI cloud, then build an ExternalProviderClient for the three container CRUD endpoints @@ -3333,7 +3445,9 @@ def _summarize_container(raw: dict) -> OpenAIContainerSummary: return OpenAIContainerSummary( id = str(raw.get("id") or ""), name = raw.get("name"), - created_at = raw.get("created_at") if isinstance(raw.get("created_at"), int) else None, + created_at = raw.get("created_at") + if isinstance(raw.get("created_at"), int) + else None, last_active_at = raw.get("last_active_at") if isinstance(raw.get("last_active_at"), int) else None, @@ -3678,7 +3792,9 @@ async def openai_chat_completions( id = completion_id, created = created, model = model_name, - choices = [ChunkChoice(delta = ChoiceDelta(), finish_reason = "stop")], + choices = [ + ChunkChoice(delta = ChoiceDelta(), finish_reason = "stop") + ], ) yield f"data: {final_chunk.model_dump_json(exclude_none = True)}\n\n" yield "data: [DONE]\n\n" @@ -3686,7 +3802,9 @@ async def openai_chat_completions( cancel_event.set() raise except Exception as e: - logger.error(f"Error during audio input streaming: {e}", exc_info = True) + logger.error( + f"Error during audio input streaming: {e}", exc_info = True + ) yield f"data: {json.dumps({'error': {'message': _friendly_error(e), 'type': 'server_error'}})}\n\n" finally: _tracker.__exit__(None, None, None) @@ -3794,7 +3912,9 @@ async def openai_chat_completions( ) # ── Parse messages (handles multimodal content parts) ───── - system_prompt, chat_messages, extracted_image_b64 = _extract_content_parts(payload.messages) + system_prompt, chat_messages, extracted_image_b64 = _extract_content_parts( + payload.messages + ) if not chat_messages: raise HTTPException( @@ -3869,7 +3989,9 @@ async def openai_chat_completions( tools_to_use = [] elif payload.enabled_tools is not None: tools_to_use = [ - t for t in ALL_TOOLS if t["function"]["name"] in payload.enabled_tools + t + for t in ALL_TOOLS + if t["function"]["name"] in payload.enabled_tools ] else: tools_to_use = ALL_TOOLS @@ -3877,7 +3999,9 @@ async def openai_chat_completions( # Drop the RAG tool without a scope: nothing to search over. if not payload.rag_scope: tools_to_use = [ - t for t in tools_to_use if t["function"]["name"] != "search_knowledge_base" + t + for t in tools_to_use + if t["function"]["name"] != "search_knowledge_base" ] if _mcp_allowed: @@ -3911,7 +4035,9 @@ async def openai_chat_completions( ) # Nudge the model to ground in attached documents instead of memory. - _tool_names = {(t.get("function") or {}).get("name") for t in (tools_to_use or [])} + _tool_names = { + (t.get("function") or {}).get("name") for t in (tools_to_use or []) + } _rag_active = "search_knowledge_base" in _tool_names and payload.rag_scope if _rag_active: _rag_nudge = ( @@ -3923,7 +4049,11 @@ async def openai_chat_completions( ) # Prefix the date when the tool nudge is empty (RAG-only tool set). _date_line = f"The current date is {_date.today().isoformat()}." - _nudge = _date_line + " " + _rag_nudge if not _nudge else _nudge + " " + _rag_nudge + _nudge = ( + _date_line + " " + _rag_nudge + if not _nudge + else _nudge + " " + _rag_nudge + ) if _nudge: # Append nudge to system prompt (preserve user's prompt) @@ -3931,15 +4061,21 @@ async def openai_chat_completions( system_prompt = system_prompt.rstrip() + "\n\n" + _nudge else: system_prompt = _nudge - gguf_messages = _set_or_prepend_system_message(gguf_messages, system_prompt) + gguf_messages = _set_or_prepend_system_message( + gguf_messages, system_prompt + ) _gguf_auto_heal_tool_calls = ( - payload.auto_heal_tool_calls if payload.auto_heal_tool_calls is not None else True + payload.auto_heal_tool_calls + if payload.auto_heal_tool_calls is not None + else True ) # ── Strip stale tool-call XML from conversation history ─ for _msg in gguf_messages: - if _msg.get("role") == "assistant" and isinstance(_msg.get("content"), str): + if _msg.get("role") == "assistant" and isinstance( + _msg.get("content"), str + ): _msg["content"] = _strip_tool_xml_for_display( _msg["content"], auto_heal_tool_calls = _gguf_auto_heal_tool_calls, @@ -4200,7 +4336,11 @@ async def openai_chat_completions( else: logger.warning( "gguf_stream_chunks: unexpected dict event: %s", - {k: v for k, v in cumulative.items() if k != "timings"}, + { + k: v + for k, v in cumulative.items() + if k != "timings" + }, ) continue new_text = cumulative[len(prev_text) :] @@ -4300,10 +4440,16 @@ async def openai_chat_completions( # The prompt is shared across all n choices, so count its # tokens ONCE (OpenAI bills only generated tokens for each # extra choice). Only completion_tokens accumulates. - _prompt_tokens = completion_usage.get("prompt_tokens") or _prompt_tokens - _sum_completion += completion_usage.get("completion_tokens") or 0 + _prompt_tokens = ( + completion_usage.get("prompt_tokens") or _prompt_tokens + ) + _sum_completion += ( + completion_usage.get("completion_tokens") or 0 + ) if _prompt_details is None: - _prompt_details = completion_usage.get("prompt_tokens_details") + _prompt_details = completion_usage.get( + "prompt_tokens_details" + ) response = ChatCompletion( id = completion_id, @@ -4386,12 +4532,16 @@ async def openai_chat_completions( # the GGUF path). _sf_is_gptoss = False try: - _sf_is_gptoss = bool(hasattr(backend, "_is_gpt_oss_model") and backend._is_gpt_oss_model()) + _sf_is_gptoss = bool( + hasattr(backend, "_is_gpt_oss_model") and backend._is_gpt_oss_model() + ) except Exception: _sf_is_gptoss = False _sf_tool_budget = ( - payload.max_tool_calls_per_message if payload.max_tool_calls_per_message is not None else 25 + payload.max_tool_calls_per_message + if payload.max_tool_calls_per_message is not None + else 25 ) # Match the GGUF path: mcp_enabled also opens the tool loop on its own @@ -4424,7 +4574,9 @@ async def openai_chat_completions( # Drop the RAG tool unless the request carries a retrieval scope. if not payload.rag_scope: _sf_tools_to_use = [ - t for t in _sf_tools_to_use if t["function"]["name"] != "search_knowledge_base" + t + for t in _sf_tools_to_use + if t["function"]["name"] != "search_knowledge_base" ] if _sf_mcp_allowed: @@ -4453,7 +4605,9 @@ async def openai_chat_completions( ) # RAG nudge, mirroring the GGUF path. - _sf_tool_names = {(t.get("function") or {}).get("name") for t in (_sf_tools_to_use or [])} + _sf_tool_names = { + (t.get("function") or {}).get("name") for t in (_sf_tools_to_use or []) + } _sf_rag_active = "search_knowledge_base" in _sf_tool_names and payload.rag_scope if _sf_rag_active: _sf_rag_nudge = ( @@ -4479,7 +4633,9 @@ async def openai_chat_completions( _sf_system_prompt = _sf_nudge _sf_auto_heal_tool_calls = ( - payload.auto_heal_tool_calls if payload.auto_heal_tool_calls is not None else True + payload.auto_heal_tool_calls + if payload.auto_heal_tool_calls is not None + else True ) # Strip stale tool-call XML from prior assistant turns. @@ -5009,10 +5165,14 @@ def _openai_model_objects() -> list[dict]: _ctx = _positive_int_or_none(getattr(llama_backend, "context_length", None)) if _ctx is not None: entry["context_length"] = _ctx - _max_ctx = _positive_int_or_none(getattr(llama_backend, "max_context_length", None)) + _max_ctx = _positive_int_or_none( + getattr(llama_backend, "max_context_length", None) + ) if _max_ctx is not None: entry["max_context_length"] = _max_ctx - _native_ctx = _positive_int_or_none(getattr(llama_backend, "native_context_length", None)) + _native_ctx = _positive_int_or_none( + getattr(llama_backend, "native_context_length", None) + ) if _native_ctx is not None: entry["native_context_length"] = _native_ctx models.append(entry) @@ -5055,7 +5215,9 @@ async def openai_list_models(current_subject: str = Depends(get_current_subject) @router.get("/models/{model_id:path}") -async def openai_retrieve_model(model_id: str, current_subject: str = Depends(get_current_subject)): +async def openai_retrieve_model( + model_id: str, current_subject: str = Depends(get_current_subject) +): """ OpenAI-compatible single-model retrieval endpoint (``GET /v1/models/{id}``). @@ -5083,7 +5245,9 @@ async def openai_retrieve_model(model_id: str, current_subject: str = Depends(ge @router.post("/completions") -async def openai_completions(request: Request, current_subject: str = Depends(get_current_subject)): +async def openai_completions( + request: Request, current_subject: str = Depends(get_current_subject) +): """ OpenAI-compatible text completions endpoint (non-chat). @@ -5116,20 +5280,26 @@ async def openai_completions(request: Request, current_subject: str = Depends(ge # separator) so _cmpl_stream_event_out can rewrite the cmpl- id and # honor stream_options.include_usage per event, while keeping SSE # framing and token bytes intact. - _include_usage = bool((body.get("stream_options") or {}).get("include_usage")) + _include_usage = bool( + (body.get("stream_options") or {}).get("include_usage") + ) client = httpx.AsyncClient(timeout = _llama_streaming_generation_timeout()) resp = None bytes_iter = None try: req = client.build_request("POST", target_url, json = body) first_token_deadline = time.monotonic() + _DEFAULT_FIRST_TOKEN_TIMEOUT_S - resp = await _send_stream_with_preheader_cancel(client, req, request = request) + resp = await _send_stream_with_preheader_cancel( + client, req, request = request + ) if resp is None: return if resp.status_code != 200: err_bytes = await resp.aread() err_text = err_bytes.decode("utf-8", errors = "replace") - raise RuntimeError(f"llama-server returned {resp.status_code}: {err_text}") + raise RuntimeError( + f"llama-server returned {resp.status_code}: {err_text}" + ) bytes_iter = resp.aiter_bytes() buffer = b"" async for chunk in _aiter_llama_stream_items( @@ -5197,7 +5367,9 @@ async def openai_completions(request: Request, current_subject: str = Depends(ge @router.post("/embeddings") -async def openai_embeddings(request: Request, current_subject: str = Depends(get_current_subject)): +async def openai_embeddings( + request: Request, current_subject: str = Depends(get_current_subject) +): """ OpenAI-compatible embeddings endpoint. @@ -5217,7 +5389,9 @@ async def openai_embeddings(request: Request, current_subject: str = Depends(get target_url = f"{llama_backend.base_url}/v1/embeddings" async with httpx.AsyncClient() as client: - resp = await client.post(target_url, json = body, timeout = _DEFAULT_FIRST_TOKEN_TIMEOUT_S) + resp = await client.post( + target_url, json = body, timeout = _DEFAULT_FIRST_TOKEN_TIMEOUT_S + ) return Response( content = resp.content, status_code = resp.status_code, @@ -5230,7 +5404,9 @@ async def openai_embeddings(request: Request, current_subject: str = Depends(get # ===================================================================== -def _translate_responses_tools_to_chat(tools: Optional[list[dict]]) -> Optional[list[dict]]: +def _translate_responses_tools_to_chat( + tools: Optional[list[dict]], +) -> Optional[list[dict]]: """Translate Responses-shape function tools to the Chat Completions nested shape. Responses uses a flat shape per tool entry:: @@ -5322,7 +5498,15 @@ def _responses_tool_output_text(output: Union[str, list]) -> str: _RESPONSES_THINK_OPEN = "" _RESPONSES_THINK_CLOSE = "" -_RESPONSES_REASONING_EFFORTS = {"none", "minimal", "low", "medium", "high", "max", "xhigh"} +_RESPONSES_REASONING_EFFORTS = { + "none", + "minimal", + "low", + "medium", + "high", + "max", + "xhigh", +} def _coerce_responses_reasoning_text(value: Any) -> str: @@ -5380,10 +5564,14 @@ class _ResponsesReasoningExtractor: close_idx = self._buffer.find(_RESPONSES_THINK_CLOSE) if close_idx != -1: reasoning_parts.append(self._buffer[:close_idx]) - self._buffer = self._buffer[close_idx + len(_RESPONSES_THINK_CLOSE) :] + self._buffer = self._buffer[ + close_idx + len(_RESPONSES_THINK_CLOSE) : + ] self._in_reasoning = False continue - keep = _responses_marker_holdback(self._buffer, (_RESPONSES_THINK_CLOSE,)) + keep = _responses_marker_holdback( + self._buffer, (_RESPONSES_THINK_CLOSE,) + ) if keep == len(self._buffer): break reasoning_parts.append(self._buffer[:-keep] if keep else self._buffer) @@ -5449,10 +5637,15 @@ def _responses_should_parse_think_markers( return False if chat_req.enable_thinking is True: return True - return chat_req.enable_thinking is None and chat_req.reasoning_effort not in (None, "none") + return chat_req.enable_thinking is None and chat_req.reasoning_effort not in ( + None, + "none", + ) -def _responses_reasoning_output_item(reasoning_text: str, item_id: Optional[str] = None) -> dict: +def _responses_reasoning_output_item( + reasoning_text: str, item_id: Optional[str] = None +) -> dict: kwargs: dict[str, Any] = { "status": "completed", "summary": [], @@ -5696,12 +5889,16 @@ async def _responses_non_streaming( if choices: msg = choices[0].get("message", {}) or {} raw_content = msg.get("content", "") or "" - raw_text = raw_content if isinstance(raw_content, str) else json.dumps(raw_content) + raw_text = ( + raw_content if isinstance(raw_content, str) else json.dumps(raw_content) + ) llama_backend = get_llama_cpp_backend() reasoning_text, text = _extract_responses_reasoning( raw_text, msg.get("reasoning_content"), - parse_think_markers = _responses_should_parse_think_markers(chat_req, llama_backend), + parse_think_markers = _responses_should_parse_think_markers( + chat_req, llama_backend + ), ) tool_calls = msg.get("tool_calls") or [] @@ -5799,7 +5996,8 @@ async def _responses_stream( # Direct pass-through bypasses the openai_chat_completions image gate. if not llama_backend.is_vision and any( - isinstance(m.content, list) and any(isinstance(p, ImageContentPart) for p in m.content) + isinstance(m.content, list) + and any(isinstance(p, ImageContentPart) for p in m.content) for m in messages ): raise HTTPException( @@ -5819,10 +6017,20 @@ async def _responses_stream( input_tokens = 0 output_tokens = 0 extractor = _ResponsesReasoningExtractor( - parse_think_markers = _responses_should_parse_think_markers(chat_req, llama_backend) + parse_think_markers = _responses_should_parse_think_markers( + chat_req, llama_backend + ) ) - reasoning_state: dict[str, Any] = {"output_index": None, "item_id": None, "opened": False} - message_state: dict[str, Any] = {"output_index": None, "item_id": None, "opened": False} + reasoning_state: dict[str, Any] = { + "output_index": None, + "item_id": None, + "opened": False, + } + message_state: dict[str, Any] = { + "output_index": None, + "item_id": None, + "opened": False, + } # Per-tool-call state keyed by Chat Completions `tool_calls[].index`, # stable across chunks for the same call. Values: # {output_index, item_id, call_id, name, arguments, opened} @@ -5920,7 +6128,9 @@ async def _responses_stream( "id": reasoning_state["item_id"], "status": "completed", "summary": [], - "content": [{"type": "reasoning_text", "text": full_reasoning}], + "content": [ + {"type": "reasoning_text", "text": full_reasoning} + ], }, ) ) @@ -6011,7 +6221,9 @@ async def _responses_stream( req = client.build_request("POST", target_url, json = body) first_token_deadline = time.monotonic() + _DEFAULT_FIRST_TOKEN_TIMEOUT_S try: - resp = await _send_stream_with_preheader_cancel(client, req, request = request) + resp = await _send_stream_with_preheader_cancel( + client, req, request = request + ) if resp is None: return except httpx.RequestError as e: @@ -6173,7 +6385,9 @@ async def _responses_stream( "output_index": st["output_index"], "delta": arg_delta, } - yield _sse("response.function_call_arguments.delta", args_delta_event) + yield _sse( + "response.function_call_arguments.delta", args_delta_event + ) elif arg_delta: # Buffer args until we can open the item (some models # send id/name in the same chunk as the first arg delta; @@ -6186,7 +6400,9 @@ async def _responses_stream( output_tokens = usage.get("completion_tokens", output_tokens) except Exception as e: logger.error("responses stream error: %s", e) - status_code = 400 if _classify_llama_generation_error(e) is not None else 500 + status_code = ( + 400 if _classify_llama_generation_error(e) is not None else 500 + ) yield _sse( "response.failed", _failed_response_payload(e, status_code), @@ -6254,10 +6470,16 @@ async def _responses_stream( close_items: list[tuple[int, str, dict[str, Any]]] = [] if reasoning_state["opened"]: - close_items.append((reasoning_state["output_index"], "reasoning", reasoning_state)) + close_items.append( + (reasoning_state["output_index"], "reasoning", reasoning_state) + ) if message_state["opened"]: - close_items.append((message_state["output_index"], "message", message_state)) - close_items.extend((st["output_index"], "tool", st) for st in tool_call_state.values()) + close_items.append( + (message_state["output_index"], "message", message_state) + ) + close_items.extend( + (st["output_index"], "tool", st) for st in tool_call_state.values() + ) for _, kind, st in sorted(close_items, key = lambda item: item[0]): if kind == "reasoning": @@ -6291,7 +6513,9 @@ async def _responses_stream( "id": st["item_id"], "status": "completed", "summary": [], - "content": [{"type": "reasoning_text", "text": full_reasoning}], + "content": [ + {"type": "reasoning_text", "text": full_reasoning} + ], }, }, ) @@ -6315,7 +6539,11 @@ async def _responses_stream( "item_id": st["item_id"], "output_index": st["output_index"], "content_index": 0, - "part": {"type": "output_text", "text": full_text, "annotations": []}, + "part": { + "type": "output_text", + "text": full_text, + "annotations": [], + }, }, ) yield _sse( @@ -6329,7 +6557,11 @@ async def _responses_stream( "status": "completed", "role": "assistant", "content": [ - {"type": "output_text", "text": full_text, "annotations": []} + { + "type": "output_text", + "text": full_text, + "annotations": [], + } ], }, }, @@ -6475,7 +6707,9 @@ def _anthropic_requested_studio_tools(tools: Optional[list]) -> set[str]: def _select_anthropic_server_tools( - all_tools: list[dict], requested_studio_tools: set[str], enabled_tools: Optional[list[str]] + all_tools: list[dict], + requested_studio_tools: set[str], + enabled_tools: Optional[list[str]], ) -> list[dict]: """Select Studio tools requested through Anthropic tools and extensions.""" if not requested_studio_tools and enabled_tools is None: @@ -6488,7 +6722,9 @@ def _select_anthropic_server_tools( return [tool for tool in all_tools if tool["function"]["name"] in selected_names] -def _normalize_anthropic_openai_images(openai_messages: list[dict], is_vision: bool) -> bool: +def _normalize_anthropic_openai_images( + openai_messages: list[dict], is_vision: bool +) -> bool: """Enforce the vision guard on translated Anthropic messages and normalize any base64-data-URL ``image_url`` parts to PNG. @@ -6603,7 +6839,11 @@ def _set_or_prepend_system_message( # Drop existing system/developer turns so the backend never sees duplicate # or conflicting system instructions, then prepend the resolved prompt. - others = [dict(msg) for msg in safe_messages if msg.get("role") not in ("system", "developer")] + others = [ + dict(msg) + for msg in safe_messages + if msg.get("role") not in ("system", "developer") + ] return [{"role": "system", "content": system_prompt}, *others] @@ -6661,7 +6901,9 @@ async def anthropic_messages( # Enforce vision guard + re-encode embedded images to PNG so the Anthropic # endpoint matches /v1/chat/completions. - _has_image = _normalize_anthropic_openai_images(openai_messages, llama_backend.is_vision) + _has_image = _normalize_anthropic_openai_images( + openai_messages, llama_backend.is_vision + ) temperature = payload.temperature if payload.temperature is not None else 0.6 top_p = payload.top_p if payload.top_p is not None else 0.95 @@ -6670,7 +6912,9 @@ async def anthropic_messages( repetition_penalty = ( payload.repetition_penalty if payload.repetition_penalty is not None else 1.0 ) - presence_penalty = payload.presence_penalty if payload.presence_penalty is not None else 0.0 + presence_penalty = ( + payload.presence_penalty if payload.presence_penalty is not None else 0.0 + ) stop = payload.stop_sequences or None # Translate Anthropic tool_choice to OpenAI format for llama-server. Falls @@ -6748,7 +6992,9 @@ async def anthropic_messages( and not _has_image ) client_tools = ( - not server_tools and len(openai_client_tools) > 0 and llama_backend.supports_tools + not server_tools + and len(openai_client_tools) > 0 + and llama_backend.supports_tools ) # Anthropic tool_choice.disable_parallel_tool_use caps the response to a @@ -7037,7 +7283,9 @@ async def _anthropic_plain_stream( # makes blocking HTTP calls to llama-server, so run it off the event loop. input_tokens = 0 if llama_backend is not None and openai_messages is not None: - input_tokens = await asyncio.to_thread(llama_backend.count_chat_tokens, openai_messages) + input_tokens = await asyncio.to_thread( + llama_backend.count_chat_tokens, openai_messages + ) async def _stream(): emitter = AnthropicStreamEmitter() @@ -7073,7 +7321,9 @@ async def _anthropic_plain_stream( yield _error_event return - stop_reason = openai_finish_to_anthropic_stop(captured_finish_reason, had_tool_calls = False) + stop_reason = openai_finish_to_anthropic_stop( + captured_finish_reason, had_tool_calls = False + ) for line in emitter.finish(stop_reason = stop_reason, stop_sequence = None): yield line @@ -7152,14 +7402,18 @@ async def _anthropic_tool_non_streaming( prev_text = clean if new: ends_on_tool_use = False - if content_blocks and isinstance(content_blocks[-1], AnthropicResponseTextBlock): + if content_blocks and isinstance( + content_blocks[-1], AnthropicResponseTextBlock + ): content_blocks[-1].text += new else: content_blocks.append(AnthropicResponseTextBlock(text = new)) elif etype == "tool_start": tool_call_id = event["tool_call_id"] arguments = event.get("arguments", {}) - existing_tool_block = tool_blocks_by_id.get(tool_call_id) if tool_call_id else None + existing_tool_block = ( + tool_blocks_by_id.get(tool_call_id) if tool_call_id else None + ) if existing_tool_block is not None: if arguments or not existing_tool_block.input: existing_tool_block.input = arguments @@ -7246,7 +7500,9 @@ async def _anthropic_plain_non_streaming(run_gen, message_id, model_name): if full_text: content_blocks.append(AnthropicResponseTextBlock(text = full_text)) - stop_reason = openai_finish_to_anthropic_stop(captured_finish_reason, had_tool_calls = False) + stop_reason = openai_finish_to_anthropic_stop( + captured_finish_reason, had_tool_calls = False + ) resp = AnthropicMessagesResponse( id = message_id, @@ -7299,7 +7555,9 @@ def _build_passthrough_payload( if stream and stream_options is not None: body["stream_options"] = stream_options body["max_tokens"] = ( - max_tokens if max_tokens is not None else (backend_ctx or _DEFAULT_MAX_TOKENS_FLOOR) + max_tokens + if max_tokens is not None + else (backend_ctx or _DEFAULT_MAX_TOKENS_FLOOR) ) # Normalize stop the same way the non-passthrough path does (the passthrough # was previously the one path that forwarded an empty stop string verbatim). @@ -7450,7 +7708,9 @@ async def _anthropic_passthrough_stream( # blocks during llama-server prefill, so the in-loop cancel # check is unreachable until the first SSE chunk arrives. # The watcher closes `resp` on cancel, raising in aiter_lines. - cancel_watcher = asyncio.create_task(_await_cancel_then_close(cancel_event, resp)) + cancel_watcher = asyncio.create_task( + _await_cancel_then_close(cancel_event, resp) + ) lines_iter = resp.aiter_lines() async for raw_line in _aiter_llama_stream_items( lines_iter, @@ -7607,7 +7867,9 @@ async def _anthropic_passthrough_non_streaming( ) ) - stop_reason = openai_finish_to_anthropic_stop(finish_reason, had_tool_calls = bool(tool_calls)) + stop_reason = openai_finish_to_anthropic_stop( + finish_reason, had_tool_calls = bool(tool_calls) + ) usage = data.get("usage") or {} resp_obj = AnthropicMessagesResponse( @@ -7702,7 +7964,9 @@ def _strip_provider_synthetic_tool_history(messages: list[dict]) -> list[dict]: if args_obj.get("_server_tool") is True: is_synthetic = True google = args_obj.get("google") - if isinstance(google, dict) and isinstance(google.get("native_part"), dict): + if isinstance(google, dict) and isinstance( + google.get("native_part"), dict + ): is_synthetic = True if is_synthetic: tc_id = tc.get("id") @@ -7756,7 +8020,9 @@ def _openai_messages_for_passthrough(payload) -> list[dict]: content part so vision + function-calling requests work transparently. """ messages = _strip_provider_synthetic_tool_history( - _drop_empty_assistant_sentinels([m.model_dump(exclude_none = True) for m in payload.messages]) + _drop_empty_assistant_sentinels( + [m.model_dump(exclude_none = True) for m in payload.messages] + ) ) if not payload.image_base64: @@ -7806,7 +8072,9 @@ def _openai_messages_for_gguf_chat(payload, is_vision: bool) -> tuple[list[dict] attached to its original turn. """ messages = _strip_provider_synthetic_tool_history( - _drop_empty_assistant_sentinels([m.model_dump(exclude_none = True) for m in payload.messages]) + _drop_empty_assistant_sentinels( + [m.model_dump(exclude_none = True) for m in payload.messages] + ) ) has_message_image = any( isinstance(msg.get("content"), list) @@ -7934,7 +8202,9 @@ async def _openai_passthrough_stream( ) resp = None _truncate_budget = ( - _OVERFLOW_TRUNCATE_MAX_RETRIES if _overflow_truncation_requested(payload) else 0 + _OVERFLOW_TRUNCATE_MAX_RETRIES + if _overflow_truncation_requested(payload) + else 0 ) while True: try: @@ -8014,7 +8284,9 @@ async def _openai_passthrough_stream( # recovers from. Run a tiny watcher that closes `resp` as soon as # cancel fires, unblocking the iterator with a RemoteProtocolError # caught in the except clause below. - cancel_watcher = asyncio.create_task(_await_cancel_then_close(cancel_event, resp)) + cancel_watcher = asyncio.create_task( + _await_cancel_then_close(cancel_event, resp) + ) try: lines_iter = resp.aiter_lines() async for raw_line in _aiter_llama_stream_items( @@ -8032,7 +8304,10 @@ async def _openai_passthrough_stream( # deltas with index>=1 so only the first call streams. Only # lines carrying tool_calls are reparsed; everything else is # relayed byte-for-byte. - if payload.parallel_tool_calls is False and '"tool_calls"' in raw_line: + if ( + payload.parallel_tool_calls is False + and '"tool_calls"' in raw_line + ): raw_line = _cap_parallel_tool_calls_sse_line(raw_line) # Relay verbatim to preserve llama-server's native id, # finish_reason, delta.tool_calls, and usage chunks. @@ -8111,7 +8386,9 @@ async def _openai_passthrough_non_streaming(llama_backend, payload, model_name): # llama-server subprocess crashed / starting / unreachable. Surface the # same friendly message the sync chat path emits so operators don't see # a bare 500 with no diagnostic. - logger.error("openai passthrough non-streaming: upstream unreachable: %s", e) + logger.error( + "openai passthrough non-streaming: upstream unreachable: %s", e + ) raise HTTPException( status_code = 502, detail = _friendly_error(e), diff --git a/studio/backend/routes/llama.py b/studio/backend/routes/llama.py index 3aae6f4209..2dac899824 100644 --- a/studio/backend/routes/llama.py +++ b/studio/backend/routes/llama.py @@ -31,7 +31,9 @@ class LlamaUpdateJob(BaseModel): from_tag: Optional[str] = None to_tag: Optional[str] = None error: Optional[str] = None - progress: Optional[float] = Field(None, description = "0..1 while running, 1 on success.") + progress: Optional[float] = Field( + None, description = "0..1 while running, 1 on success." + ) started_at: Optional[str] = None finished_at: Optional[str] = None @@ -42,10 +44,12 @@ class LlamaUpdateStatusResponse(BaseModel): description = "True when the install came from an Unsloth prebuilt (has a marker).", ) update_available: bool = Field( - False, description = "True when the latest release is genuinely newer than the install." + False, + description = "True when the latest release is genuinely newer than the install.", ) stale: bool = Field( - False, description = "Update available AND install older than the staleness threshold." + False, + description = "Update available AND install older than the staleness threshold.", ) installed_tag: Optional[str] = None latest_tag: Optional[str] = None @@ -53,7 +57,8 @@ class LlamaUpdateStatusResponse(BaseModel): installed_at_utc: Optional[str] = None age_days: Optional[int] = None source_build: bool = Field( - False, description = "True when there is no marker (source build) but a prebuilt is offered." + False, + description = "True when there is no marker (source build) but a prebuilt is offered.", ) job: LlamaUpdateJob = Field(default_factory = LlamaUpdateJob) diff --git a/studio/backend/routes/mcp_servers.py b/studio/backend/routes/mcp_servers.py index 37d99a222e..0bb18a3a6f 100644 --- a/studio/backend/routes/mcp_servers.py +++ b/studio/backend/routes/mcp_servers.py @@ -152,7 +152,9 @@ def _changes_from_payload(payload: McpServerUpdate) -> dict: if "display_name" in sent: name = (payload.display_name or "").strip() if not name: - raise HTTPException(status_code = 400, detail = "display_name must not be empty") + raise HTTPException( + status_code = 400, detail = "display_name must not be empty" + ) changes["display_name"] = name if "url" in sent: changes["url"] = _validate_url(payload.url or "") @@ -161,11 +163,15 @@ def _changes_from_payload(payload: McpServerUpdate) -> dict: changes["headers_json"] = json.dumps(headers) if headers else None if "is_enabled" in sent: if payload.is_enabled is None: - raise HTTPException(status_code = 400, detail = "is_enabled must be true or false") + raise HTTPException( + status_code = 400, detail = "is_enabled must be true or false" + ) changes["is_enabled"] = payload.is_enabled if "use_oauth" in sent: if payload.use_oauth is None: - raise HTTPException(status_code = 400, detail = "use_oauth must be true or false") + raise HTTPException( + status_code = 400, detail = "use_oauth must be true or false" + ) changes["use_oauth"] = payload.use_oauth # stdio is OAuth-less: drop a stale OAuth flag when switching to a command. if "url" in changes and is_stdio(changes["url"]): @@ -198,7 +204,8 @@ async def update_mcp_server( # fastmcp keys tokens by URL and would otherwise let a re-pointed server # silently inherit the old account's credentials. if bool(old.get("use_oauth")) and ( - ("url" in changes and changes["url"] != old["url"]) or changes.get("use_oauth") is False + ("url" in changes and changes["url"] != old["url"]) + or changes.get("use_oauth") is False ): await clear_oauth_tokens_async(old["url"]) mcp_servers_db.update_server(server_id, changes) @@ -211,7 +218,9 @@ async def update_mcp_server( @router.delete("/{server_id}", status_code = 204) -async def delete_mcp_server(server_id: str, current_subject: str = Depends(get_current_subject)): +async def delete_mcp_server( + server_id: str, current_subject: str = Depends(get_current_subject) +): old = mcp_servers_db.get_server(server_id) if not old: raise HTTPException(status_code = 404, detail = "MCP server not found") @@ -231,7 +240,9 @@ async def refresh_mcp_server_tools( # Refresh uses the stored address, so re-check the stdio gate here too: a # stdio row from a desktop DB must not spawn on a hosted/network host. if is_stdio(server["url"]) and not stdio_mcp_enabled(): - raise HTTPException(status_code = 400, detail = "stdio MCP servers are disabled on this host") + raise HTTPException( + status_code = 400, detail = "stdio MCP servers are disabled on this host" + ) use_oauth = bool(server.get("use_oauth")) try: diff --git a/studio/backend/routes/models.py b/studio/backend/routes/models.py index a2f2eca81b..1ae6732bd6 100644 --- a/studio/backend/routes/models.py +++ b/studio/backend/routes/models.py @@ -204,7 +204,9 @@ def _is_model_directory(d: Path) -> bool: return False try: - has_config = (d / "config.json").exists() or (d / "adapter_config.json").exists() + has_config = (d / "config.json").exists() or ( + d / "adapter_config.json" + ).exists() if not has_config: return False return any(_is_weight_file(f) for f in d.iterdir() if f.is_file()) @@ -212,7 +214,9 @@ def _is_model_directory(d: Path) -> bool: return False -def _scan_models_dir(models_dir: Path, *, limit: int | None = None) -> List[LocalModelInfo]: +def _scan_models_dir( + models_dir: Path, *, limit: int | None = None +) -> List[LocalModelInfo]: if not models_dir.exists() or not models_dir.is_dir(): return [] @@ -471,7 +475,9 @@ def _ollama_links_dir(ollama_dir: Path) -> Optional[Path]: return None -def _scan_ollama_dir(ollama_dir: Path, limit: Optional[int] = None) -> List[LocalModelInfo]: +def _scan_ollama_dir( + ollama_dir: Path, limit: Optional[int] = None +) -> List[LocalModelInfo]: """Scan an Ollama models directory for downloaded models. Ollama uses a content-addressable layout @@ -546,7 +552,9 @@ def _scan_ollama_dir(ollama_dir: Path, limit: Optional[int] = None) -> List[Loca if tmp_path.is_symlink() or tmp_path.exists(): tmp_path.unlink() except OSError as cleanup_err: - logger.debug("Could not clean up tmp path %s: %s", tmp_path, cleanup_err) + logger.debug( + "Could not clean up tmp path %s: %s", tmp_path, cleanup_err + ) return None try: @@ -563,7 +571,11 @@ def _scan_ollama_dir(ollama_dir: Path, limit: Optional[int] = None) -> List[Loca repo_parts = list(parts[1:-1]) tag = parts[-1] - if host == "registry.ollama.ai" and repo_parts and repo_parts[0] == "library": + if ( + host == "registry.ollama.ai" + and repo_parts + and repo_parts[0] == "library" + ): repo_name = "/".join(repo_parts[1:]) elif host == "registry.ollama.ai": repo_name = "/".join(repo_parts) @@ -620,7 +632,9 @@ def _scan_ollama_dir(ollama_dir: Path, limit: Optional[int] = None) -> List[Loca candidate = blobs_dir / digest.replace(":", "-") if candidate.is_file(): link_name = f"{safe_name}-{tag}{quant}.gguf" - gguf_link_path = _make_link(model_link_dir, link_name, candidate) + gguf_link_path = _make_link( + model_link_dir, link_name, candidate + ) elif media == "application/vnd.ollama.image.projector": candidate = blobs_dir / digest.replace(":", "-") @@ -752,7 +766,10 @@ async def list_local_models( + _scan_hf_cache(folder_path) + _scan_lmstudio_dir(folder_path) ) - if not any(p in (".studio_links", "ollama_links") for p in Path(m.path).parts) + if not any( + p in (".studio_links", "ollama_links") + for p in Path(m.path).parts + ) ] custom_models = _generic if len(custom_models) < _MAX_MODELS_PER_FOLDER: @@ -763,7 +780,9 @@ async def list_local_models( except OSError as e: logger.warning("Skipping unreadable scan folder %s: %s", folder_path, e) continue - local_models += [m.model_copy(update = {"source": "custom"}) for m in custom_models] + local_models += [ + m.model_copy(update = {"source": "custom"}) for m in custom_models + ] # Deduplicate, but always keep custom folder entries (keyed by # (id, source)) so they show in the "Custom Folders" UI section @@ -1122,7 +1141,9 @@ def _match_browse_child(current: Path, name: str) -> Optional[Path]: detail = f"Permission denied reading {current.name}", ) from None except OSError as exc: - logger.warning("browse-folders: could not read %s: %s", current, exc, exc_info = True) + logger.warning( + "browse-folders: could not read %s: %s", current, exc, exc_info = True + ) raise HTTPException( status_code = 500, detail = f"Could not read {os.path.basename(str(current))}", @@ -1260,7 +1281,9 @@ async def browse_folders( detail = f"Permission denied reading {os.path.basename(str(target))}", ) except OSError as exc: - logger.warning("browse-folders: could not read %s: %s", target, exc, exc_info = True) + logger.warning( + "browse-folders: could not read %s: %s", target, exc, exc_info = True + ) raise HTTPException( status_code = 500, detail = f"Could not read {os.path.basename(str(target))}", @@ -1313,7 +1336,9 @@ async def browse_folders( # sandbox (else the up-row would 403 on click); users can still hop # to other allowed roots via the suggestion chips. parent: Optional[str] - if target.parent == target or not _is_path_inside_allowlist(target.parent, allowed_roots): + if target.parent == target or not _is_path_inside_allowlist( + target.parent, allowed_roots + ): parent = None else: parent = str(target.parent) @@ -1454,12 +1479,16 @@ def _get_max_position_embeddings(config) -> Optional[int]: """Extract max_position_embeddings from a config, with text_config fallback.""" if hasattr(config, "max_position_embeddings"): return config.max_position_embeddings - if hasattr(config, "text_config") and hasattr(config.text_config, "max_position_embeddings"): + if hasattr(config, "text_config") and hasattr( + config.text_config, "max_position_embeddings" + ): return config.text_config.max_position_embeddings return None -def _get_model_size_bytes(model_name: str, hf_token: Optional[str] = None) -> Optional[int]: +def _get_model_size_bytes( + model_name: str, hf_token: Optional[str] = None +) -> Optional[int]: """Total size of model weight files from HF Hub.""" try: from huggingface_hub import HfApi @@ -1472,7 +1501,9 @@ def _get_model_size_bytes(model_name: str, hf_token: Optional[str] = None) -> Op weight_exts = (".safetensors", ".bin", ".pt", ".pth", ".gguf") total = 0 for sibling in info.siblings: - if sibling.rfilename and any(sibling.rfilename.endswith(ext) for ext in weight_exts): + if sibling.rfilename and any( + sibling.rfilename.endswith(ext) for ext in weight_exts + ): if sibling.size is not None: total += sibling.size @@ -1652,10 +1683,14 @@ def _loaded_model_matches_deleted_path(active_model: str, deleted_path: Path) -> ) active_lower = active_model.lower() target_lower = str(deleted_path).lower() - return active_lower == target_lower or active_lower.startswith(f"{target_lower}{os.sep}") + return active_lower == target_lower or active_lower.startswith( + f"{target_lower}{os.sep}" + ) -def _loading_model_matches_deleted_path(loading_model: object, deleted_path: Path) -> bool: +def _loading_model_matches_deleted_path( + loading_model: object, deleted_path: Path +) -> bool: if not loading_model: return False return _loaded_model_matches_deleted_path(str(loading_model), deleted_path) @@ -1874,7 +1909,9 @@ async def delete_finetuned_model( except HTTPException: raise except Exception as e: - logger.warning("Could not check inference backend loaded model before delete: %s", e) + logger.warning( + "Could not check inference backend loaded model before delete: %s", e + ) raise HTTPException( status_code = 503, detail = "Could not verify model load status before deleting", @@ -1946,7 +1983,9 @@ async def delete_finetuned_model( @router.get("/loras/{lora_path:path}/base-model", response_model = LoRABaseModelResponse) -async def get_lora_base_model(lora_path: str, current_subject: str = Depends(get_current_subject)): +async def get_lora_base_model( + lora_path: str, current_subject: str = Depends(get_current_subject) +): """ Get the base model for a LoRA adapter. @@ -1979,7 +2018,9 @@ async def get_lora_base_model(lora_path: str, current_subject: str = Depends(get @router.get("/check-vision/{model_name:path}", response_model = VisionCheckResponse) -async def check_vision_model(model_name: str, current_subject: str = Depends(get_current_subject)): +async def check_vision_model( + model_name: str, current_subject: str = Depends(get_current_subject) +): """ Check if a model is a vision model. @@ -2020,7 +2061,9 @@ async def check_embedding_model( logger.info(f"Checking if embedding model: {model_name}") is_embedding = is_embedding_model(model_name, hf_token = hf_token) - logger.info(f"Embedding check result for {model_name}: is_embedding={is_embedding}") + logger.info( + f"Embedding check result for {model_name}: is_embedding={is_embedding}" + ) return EmbeddingCheckResponse( model_name = model_name, is_embedding = is_embedding, @@ -2041,7 +2084,9 @@ async def get_gguf_variants( repo_id: str = Query( ..., description = "HuggingFace repo ID (e.g. 'unsloth/gemma-3-4b-it-GGUF')" ), - hf_token: Optional[str] = Query(None, description = "HuggingFace token for private repos"), + hf_token: Optional[str] = Query( + None, description = "HuggingFace token for private repos" + ), current_subject: str = Depends(get_current_subject), ): """List GGUF quantization variants for a HF repo or local directory. @@ -2200,7 +2245,11 @@ async def get_gguf_download_progress( break total_progress_bytes = downloaded_bytes + in_progress_bytes - progress = min(total_progress_bytes / expected_bytes, 0.99) if expected_bytes > 0 else 0 + progress = ( + min(total_progress_bytes / expected_bytes, 0.99) + if expected_bytes > 0 + else 0 + ) # Report 1.0 only when all bytes are in completed files. if expected_bytes > 0 and downloaded_bytes >= expected_bytes: progress = 1.0 @@ -2486,7 +2535,9 @@ async def list_cached_gguf(current_subject: str = Depends(get_current_subject)): } # Keep the newest timestamp across duplicate caches; # attach only when known so absent rows sort as oldest. - lm = max(last_modified, (existing or {}).get("last_modified", 0.0)) + lm = max( + last_modified, (existing or {}).get("last_modified", 0.0) + ) if lm > 0: row["last_modified"] = lm seen_lower[key] = row @@ -2527,7 +2578,9 @@ async def list_cached_models(current_subject: str = Depends(get_current_subject) if _repo_has_gguf_files(repo_info): continue total_size = sum( - (f.size_on_disk or 0) for rev in repo_info.revisions for f in rev.files + (f.size_on_disk or 0) + for rev in repo_info.revisions + for f in rev.files ) if total_size == 0: continue @@ -2556,7 +2609,9 @@ async def list_cached_models(current_subject: str = Depends(get_current_subject) } # Keep the newest timestamp across duplicate caches; # attach only when known so absent rows sort as oldest. - lm = max(last_modified, (existing or {}).get("last_modified", 0.0)) + lm = max( + last_modified, (existing or {}).get("last_modified", 0.0) + ) if lm > 0: row["last_modified"] = lm seen_lower[key] = row diff --git a/studio/backend/routes/prompts.py b/studio/backend/routes/prompts.py index df81008766..d3fc9eb356 100644 --- a/studio/backend/routes/prompts.py +++ b/studio/backend/routes/prompts.py @@ -69,7 +69,9 @@ def remove_entry(entry_id: str, current_subject: str = Depends(get_current_subje @router.post("/entries/bulk") -def bulk_entries(req: BulkEntriesRequest, current_subject: str = Depends(get_current_subject)): +def bulk_entries( + req: BulkEntriesRequest, current_subject: str = Depends(get_current_subject) +): count = bulk_upsert_prompt_entries([e.model_dump() for e in req.entries]) return {"count": count} @@ -96,6 +98,8 @@ def remove_list(list_id: str, current_subject: str = Depends(get_current_subject @router.post("/lists/bulk") -def bulk_lists(req: BulkListsRequest, current_subject: str = Depends(get_current_subject)): +def bulk_lists( + req: BulkListsRequest, current_subject: str = Depends(get_current_subject) +): count = bulk_upsert_prompt_lists([l.model_dump() for l in req.lists]) return {"count": count} diff --git a/studio/backend/routes/providers.py b/studio/backend/routes/providers.py index 5a55c9b0bb..865ae7c957 100644 --- a/studio/backend/routes/providers.py +++ b/studio/backend/routes/providers.py @@ -206,7 +206,9 @@ async def test_provider( try: api_key = decrypt_api_key(payload.encrypted_api_key) except Exception as exc: - logger.warning("Failed to decrypt API key (%s): %s", type(exc).__name__, exc) + logger.warning( + "Failed to decrypt API key (%s): %s", type(exc).__name__, exc + ) raise HTTPException( status_code = 400, detail = "Failed to decrypt API key. The public key may have changed — try refreshing the page.", @@ -305,7 +307,9 @@ async def list_provider_models( try: api_key = decrypt_api_key(payload.encrypted_api_key) except Exception as exc: - logger.warning("Failed to decrypt API key (%s): %s", type(exc).__name__, exc) + logger.warning( + "Failed to decrypt API key (%s): %s", type(exc).__name__, exc + ) raise HTTPException( status_code = 400, detail = "Failed to decrypt API key. The public key may have changed — try refreshing the page.", @@ -350,7 +354,9 @@ async def list_provider_models( if allow_prefixes is not None: prefix_tuple = tuple(str(p) for p in allow_prefixes if str(p)) if prefix_tuple: - models = [m for m in models if m.get("id", "").startswith(prefix_tuple)] + models = [ + m for m in models if m.get("id", "").startswith(prefix_tuple) + ] allowlist = info.get("model_id_allowlist") if allowlist is not None: models = [m for m in models if allowlist.match(m.get("id", ""))] diff --git a/studio/backend/routes/rag.py b/studio/backend/routes/rag.py index 8d23240fd5..85262e3299 100644 --- a/studio/backend/routes/rag.py +++ b/studio/backend/routes/rag.py @@ -183,7 +183,9 @@ def update_knowledge_base( params.append(payload.description or None) if sets: params.append(kb_id) - conn.execute(f"UPDATE knowledge_bases SET {', '.join(sets)} WHERE id=?", params) + conn.execute( + f"UPDATE knowledge_bases SET {', '.join(sets)} WHERE id=?", params + ) conn.commit() return {"ok": True} finally: @@ -191,7 +193,9 @@ def update_knowledge_base( @router.delete("/knowledge-bases/{kb_id}") -def delete_knowledge_base(kb_id: str, subject: str = Depends(get_current_subject)) -> dict: +def delete_knowledge_base( + kb_id: str, subject: str = Depends(get_current_subject) +) -> dict: _require_rag() conn = rag_db.get_connection() try: @@ -249,7 +253,9 @@ async def upload_thread_document( @router.get("/threads/{thread_id}/documents") -def list_thread_documents(thread_id: str, subject: str = Depends(get_current_subject)) -> dict: +def list_thread_documents( + thread_id: str, subject: str = Depends(get_current_subject) +) -> dict: _require_rag() conn = rag_db.get_connection() try: @@ -283,7 +289,9 @@ async def upload_project_document( @router.get("/projects/{project_id}/documents") -def list_project_documents(project_id: str, subject: str = Depends(get_current_subject)) -> dict: +def list_project_documents( + project_id: str, subject: str = Depends(get_current_subject) +) -> dict: _require_rag() conn = rag_db.get_connection() try: @@ -294,7 +302,9 @@ def list_project_documents(project_id: str, subject: str = Depends(get_current_s @router.delete("/documents/{document_id}") -def delete_document(document_id: str, subject: str = Depends(get_current_subject)) -> dict: +def delete_document( + document_id: str, subject: str = Depends(get_current_subject) +) -> dict: _require_rag() conn = rag_db.get_connection() try: @@ -325,7 +335,9 @@ def job_status(job_id: str, subject: str = Depends(get_current_subject)) -> dict @router.get("/jobs/{job_id}/events") -def job_events(job_id: str, subject: str = Depends(get_current_subject)) -> StreamingResponse: +def job_events( + job_id: str, subject: str = Depends(get_current_subject) +) -> StreamingResponse: _require_rag() def gen(): @@ -355,7 +367,9 @@ def search(payload: SearchRequest, subject: str = Depends(get_current_subject)) if payload.thread_id: scopes.append(store.thread_scope(payload.thread_id)) if not scopes: - raise HTTPException(status_code = 400, detail = "Provide kb_id, project_id, or thread_id") + raise HTTPException( + status_code = 400, detail = "Provide kb_id, project_id, or thread_id" + ) scope = scopes[0] if len(scopes) == 1 else scopes conn = rag_db.get_connection() @@ -365,7 +379,9 @@ def search(payload: SearchRequest, subject: str = Depends(get_current_subject)) elif payload.mode == "dense": hits = retrieval.retrieve_dense(conn, scope, payload.query, payload.top_k) else: - hits = retrieval.retrieve_hybrid(conn, scope, payload.query, k = payload.top_k) + hits = retrieval.retrieve_hybrid( + conn, scope, payload.query, k = payload.top_k + ) hits = retrieval.filter_min_score(hits, payload.min_score) rows = store.chunks_by_id(conn, [h.chunk_id for h in hits]) results = [] @@ -470,7 +486,9 @@ def preview_target( @router.get("/documents/{document_id}/file-url") -def document_file_url(document_id: str, subject: str = Depends(get_current_subject)) -> dict: +def document_file_url( + document_id: str, subject: str = Depends(get_current_subject) +) -> dict: """Mint a short-lived signed URL for the source file.""" _require_rag() conn = rag_db.get_connection() diff --git a/studio/backend/routes/settings.py b/studio/backend/routes/settings.py index 516502c1e5..2746b6bda4 100644 --- a/studio/backend/routes/settings.py +++ b/studio/backend/routes/settings.py @@ -68,7 +68,9 @@ def _helper_precache_response(enabled: bool | None = None) -> HelperPrecacheResp @router.get("/upload-limit", response_model = UploadLimitResponse) -def get_upload_limit(current_subject: str = Depends(get_current_subject)) -> UploadLimitResponse: +def get_upload_limit( + current_subject: str = Depends(get_current_subject), +) -> UploadLimitResponse: return _upload_limit_response(get_upload_limit_mb()) diff --git a/studio/backend/routes/training.py b/studio/backend/routes/training.py index 281f03bcaf..fe09493e3c 100644 --- a/studio/backend/routes/training.py +++ b/studio/backend/routes/training.py @@ -69,7 +69,9 @@ router = APIRouter() logger = get_logger(__name__) -def _validate_local_dataset_paths(paths: list[str], label: str = "Local dataset") -> list[str]: +def _validate_local_dataset_paths( + paths: list[str], label: str = "Local dataset" +) -> list[str]: """Resolve and validate a list of local dataset paths. Returns validated absolute paths.""" validated = [] missing = [] @@ -102,7 +104,9 @@ async def get_hardware_utilization(current_subject: str = Depends(get_current_su @router.get("/hardware/visible") -async def get_visible_hardware_utilization(current_subject: str = Depends(get_current_subject)): +async def get_visible_hardware_utilization( + current_subject: str = Depends(get_current_subject), +): from utils.hardware import get_visible_gpu_utilization return get_visible_gpu_utilization() @@ -151,7 +155,9 @@ async def start_training( # Job ID; start_training() sets it on the backend only after the old # pump thread is dead. - job_id = f"job_{datetime.now().strftime('%Y%m%d_%H%M%S')}_{_uuid.uuid4().hex[:8]}" + job_id = ( + f"job_{datetime.now().strftime('%Y%m%d_%H%M%S')}_{_uuid.uuid4().hex[:8]}" + ) # Validate dataset paths if provided. if request.local_datasets: @@ -165,7 +171,9 @@ async def start_training( resume_output_dir: Optional[str] = None if request.resume_from_checkpoint: try: - resume_output_dir = normalize_resume_output_dir(request.resume_from_checkpoint) + resume_output_dir = normalize_resume_output_dir( + request.resume_from_checkpoint + ) except ValueError as e: # Deliberate user-facing validation message. validation_message = str(e) @@ -223,7 +231,9 @@ async def start_training( "lora_r": request.lora_r, "lora_alpha": request.lora_alpha, "lora_dropout": request.lora_dropout, - "target_modules": request.target_modules if request.target_modules else None, + "target_modules": request.target_modules + if request.target_modules + else None, "gradient_checkpointing": request.gradient_checkpointing.strip() if request.gradient_checkpointing and request.gradient_checkpointing.strip() else "unsloth", @@ -253,9 +263,13 @@ async def start_training( # YAML model defaults directly so models that need it always get it. if not training_kwargs["trust_remote_code"]: model_defaults = load_model_defaults(request.model_name) - yaml_trust = model_defaults.get("training", {}).get("trust_remote_code", False) + yaml_trust = model_defaults.get("training", {}).get( + "trust_remote_code", False + ) if yaml_trust: - logger.info(f"YAML config sets trust_remote_code=True for {request.model_name}") + logger.info( + f"YAML config sets trust_remote_code=True for {request.model_name}" + ) training_kwargs["trust_remote_code"] = True # Free GPU memory: shut down any running inference/export subprocesses @@ -278,7 +292,9 @@ async def start_training( from core.export import get_export_backend exp_backend = get_export_backend() if exp_backend.current_checkpoint: - logger.info("Shutting down export subprocess to free GPU memory for training") + logger.info( + "Shutting down export subprocess to free GPU memory for training" + ) exp_backend._shutdown_subprocess() exp_backend.current_checkpoint = None exp_backend.is_vision = False @@ -372,10 +388,14 @@ async def reset_training(current_subject: str = Depends(get_current_subject)): if is_active: if backend._cancel_requested: # Cancel (save=False) requested — force-terminate to reset immediately. - logger.info("Force-terminating subprocess for immediate reset (cancel path)") + logger.info( + "Force-terminating subprocess for immediate reset (cancel path)" + ) backend.force_terminate() else: - logger.warning("Rejected reset while training active: is_active=%s", is_active) + logger.warning( + "Rejected reset while training active: is_active=%s", is_active + ) raise HTTPException( status_code = 409, detail = "Training is still running. Stop training and wait for it to finish before resetting.", @@ -441,7 +461,9 @@ async def get_training_status(current_subject: str = Depends(get_current_subject msg_lower = status_message.lower() if "loading" in msg_lower or "importing" in msg_lower: phase = "loading_model" - elif any(k in msg_lower for k in ["preparing", "initializing", "configuring"]): + elif any( + k in msg_lower for k in ["preparing", "initializing", "configuring"] + ): phase = "configuring" else: phase = "training" @@ -580,10 +602,14 @@ async def stream_training_progress( if step < 0 or total == 0: progress_percent = 0.0 else: - progress_percent = float(step) / float(total) * 100.0 if total > 0 else 0.0 + progress_percent = ( + float(step) / float(total) * 100.0 if total > 0 else 0.0 + ) # Pull values from the progress object if available. - elapsed_seconds = getattr(progress, "elapsed_seconds", None) if progress else None + elapsed_seconds = ( + getattr(progress, "elapsed_seconds", None) if progress else None + ) eta_seconds = getattr(progress, "eta_seconds", None) if progress else None grad_norm = grad_norm_override if grad_norm is None and progress: @@ -639,15 +665,25 @@ async def stream_training_progress( } for i, step_val in enumerate(backend.step_history): if step_val > resume_from_step: - loss_val = backend.loss_history[i] if i < len(backend.loss_history) else None - lr_val = backend.lr_history[i] if i < len(backend.lr_history) else None + loss_val = ( + backend.loss_history[i] + if i < len(backend.loss_history) + else None + ) + lr_val = ( + backend.lr_history[i] if i < len(backend.lr_history) else None + ) tp_replay = getattr( getattr(backend, "trainer", None), "training_progress", None ) total_replay = ( - getattr(tp_replay, "total_steps", step_val) if tp_replay else step_val + getattr(tp_replay, "total_steps", step_val) + if tp_replay + else step_val + ) + epoch_replay = ( + getattr(tp_replay, "epoch", None) if tp_replay else None ) - epoch_replay = getattr(tp_replay, "epoch", None) if tp_replay else None payload = build_progress( step_val, loss_val, @@ -657,7 +693,9 @@ async def stream_training_progress( progress = tp_replay, grad_norm_override = grad_norm_by_step.get(step_val), ) - yield format_sse(payload.model_dump_json(), event = "progress", event_id = step_val) + yield format_sse( + payload.model_dump_json(), event = "progress", event_id = step_val + ) replayed += 1 if replayed: logger.info(f"SSE reconnect: replayed {replayed} missed steps") @@ -677,14 +715,18 @@ async def stream_training_progress( epoch = initial_epoch, progress = tp, ) - yield format_sse(initial_progress.model_dump_json(), event = "progress", event_id = 0) + yield format_sse( + initial_progress.model_dump_json(), event = "progress", event_id = 0 + ) # If not active, send final state and exit if not is_active: _live = (getattr(tp, "step", 0) or 0) if tp else 0 if backend.step_history or _live > 0: final_step = backend.step_history[-1] if backend.step_history else 0 - final_loss = backend.loss_history[-1] if backend.loss_history else None + final_loss = ( + backend.loss_history[-1] if backend.loss_history else None + ) final_lr = backend.lr_history[-1] if backend.lr_history else None # Histories skip non-finite steps; report the live step with # loss=None instead of the last finite pair. @@ -692,7 +734,9 @@ async def stream_training_progress( final_step = _live final_loss = getattr(tp, "loss", None) final_lr = getattr(tp, "learning_rate", final_lr) - final_total_steps = getattr(tp, "total_steps", final_step) if tp else final_step + final_total_steps = ( + getattr(tp, "total_steps", final_step) if tp else final_step + ) final_epoch = getattr(tp, "epoch", None) if tp else None payload = build_progress( final_step, @@ -707,7 +751,9 @@ async def stream_training_progress( ) else: yield format_sse( - build_progress(-1, None, None, 0, progress = tp).model_dump_json(), + build_progress( + -1, None, None, 0, progress = tp + ).model_dump_json(), event = "complete", event_id = 0, ) @@ -720,11 +766,17 @@ async def stream_training_progress( while backend.is_training_active(): try: - tp_inner = getattr(getattr(backend, "trainer", None), "training_progress", None) + tp_inner = getattr( + getattr(backend, "trainer", None), "training_progress", None + ) live_step = (getattr(tp_inner, "step", 0) or 0) if tp_inner else 0 if backend.step_history or live_step > 0: - current_step = backend.step_history[-1] if backend.step_history else 0 - current_loss = backend.loss_history[-1] if backend.loss_history else None + current_step = ( + backend.step_history[-1] if backend.step_history else 0 + ) + current_loss = ( + backend.loss_history[-1] if backend.loss_history else None + ) current_lr = backend.lr_history[-1] if backend.lr_history else None # Histories skip non-finite steps; follow the live progress # step and report its loss (None until it recovers). @@ -733,9 +785,13 @@ async def stream_training_progress( current_loss = getattr(tp_inner, "loss", None) current_lr = getattr(tp_inner, "learning_rate", current_lr) current_total_steps = ( - getattr(tp_inner, "total_steps", current_step) if tp_inner else current_step + getattr(tp_inner, "total_steps", current_step) + if tp_inner + else current_step + ) + current_epoch = ( + getattr(tp_inner, "epoch", None) if tp_inner else None ) - current_epoch = getattr(tp_inner, "epoch", None) if tp_inner else None # Only send if the step changed. if current_step != last_step: @@ -782,7 +838,9 @@ async def stream_training_progress( "training_progress", None, ) - prep_total = getattr(tp_prep, "total_steps", 0) if tp_prep else 0 + prep_total = ( + getattr(tp_prep, "total_steps", 0) if tp_prep else 0 + ) preparing_payload = build_progress( 0, None, @@ -802,7 +860,9 @@ async def stream_training_progress( tp_timeout = getattr( getattr(backend, "trainer", None), "training_progress", None ) - timeout_payload = build_progress(last_step, None, None, 0, progress = tp_timeout) + timeout_payload = build_progress( + last_step, None, None, 0, progress = tp_timeout + ) yield format_sse( timeout_payload.model_dump_json(), event = "error", @@ -814,7 +874,9 @@ async def stream_training_progress( except Exception as e: logger.error(f"Error in progress stream: {e}", exc_info = True) - tp_error = getattr(getattr(backend, "trainer", None), "training_progress", None) + tp_error = getattr( + getattr(backend, "trainer", None), "training_progress", None + ) error_payload = build_progress(0, None, None, 0, progress = tp_error) yield format_sse( error_payload.model_dump_json(), @@ -835,7 +897,9 @@ async def stream_training_progress( final_step = _final_live_step final_loss = getattr(final_tp, "loss", None) final_lr = getattr(final_tp, "learning_rate", final_lr) - final_total_steps = getattr(final_tp, "total_steps", final_step) if final_tp else final_step + final_total_steps = ( + getattr(final_tp, "total_steps", final_step) if final_tp else final_step + ) final_epoch = getattr(final_tp, "epoch", None) if final_tp else None final_payload = build_progress( final_step, diff --git a/studio/backend/routes/training_history.py b/studio/backend/routes/training_history.py index 1560c72767..37fcdd6062 100644 --- a/studio/backend/routes/training_history.py +++ b/studio/backend/routes/training_history.py @@ -42,13 +42,18 @@ async def list_training_runs( """List training runs, newest first.""" result = list_runs(limit = limit, offset = offset) return TrainingRunListResponse( - runs = [TrainingRunSummary(**{**r, "can_resume": can_resume_run(r)}) for r in result["runs"]], + runs = [ + TrainingRunSummary(**{**r, "can_resume": can_resume_run(r)}) + for r in result["runs"] + ], total = result["total"], ) @router.get("/runs/{run_id}", response_model = TrainingRunDetailResponse) -async def get_training_run_detail(run_id: str, current_subject: str = Depends(get_current_subject)): +async def get_training_run_detail( + run_id: str, current_subject: str = Depends(get_current_subject) +): """Get a single training run with full config and metrics.""" run = get_run(run_id) if run is None: @@ -103,13 +108,17 @@ async def update_training_run( @router.delete("/runs/{run_id}", response_model = TrainingRunDeleteResponse) -async def delete_training_run(run_id: str, current_subject: str = Depends(get_current_subject)): +async def delete_training_run( + run_id: str, current_subject: str = Depends(get_current_subject) +): """Delete a training run and its metrics (CASCADE).""" run = get_run(run_id) if run is None: raise HTTPException(status_code = 404, detail = f"Run {run_id} not found") if run["status"] == "running": - raise HTTPException(status_code = 409, detail = "Cannot delete a running training run") + raise HTTPException( + status_code = 409, detail = "Cannot delete a running training run" + ) logger.info("Deleting training run %s", run_id) delete_run(run_id) return TrainingRunDeleteResponse( diff --git a/studio/backend/run.py b/studio/backend/run.py index 84991a71bb..1cf875bc22 100644 --- a/studio/backend/run.py +++ b/studio/backend/run.py @@ -26,7 +26,9 @@ try: configure_cpu_threads() except ValueError as exc: configured = os.environ.get("UNSLOTH_CPU_THREADS") - raise SystemExit(f"Error: Invalid UNSLOTH_CPU_THREADS value {configured!r}: {exc}") from None + raise SystemExit( + f"Error: Invalid UNSLOTH_CPU_THREADS value {configured!r}: {exc}" + ) from None # Anaconda/conda-forge Python: seed platform._sys_version_cache before imports # that trigger attrs -> rich -> structlog -> platform crash. @@ -91,7 +93,9 @@ def _install_uvicorn_startup_log_rewrite(bind_host: str, display_host: str) -> N import re rewrite_host = ( - bind_host in ("0.0.0.0", "::") and bool(display_host) and display_host != bind_host + bind_host in ("0.0.0.0", "::") + and bool(display_host) + and display_host != bind_host ) new_suffix = "(To stop: press Ctrl+C -- on macOS, Control+C not Command+C)" old_suffix_re = re.compile(r"\(Press CTRL\+C to quit\)") @@ -175,7 +179,9 @@ def _localhost_ipv6_mismatch_url(bind_host: str, port: int) -> "str | None": return None try: - addr_info = socket.getaddrinfo("localhost", port, socket.AF_UNSPEC, socket.SOCK_STREAM) + addr_info = socket.getaddrinfo( + "localhost", port, socket.AF_UNSPEC, socket.SOCK_STREAM + ) except Exception: return None @@ -367,7 +373,8 @@ def _verify_global_reachability(display_host: str, port: int) -> None: flush = True, ) print( - f"{dim} ssh -L {port}:localhost:{port} " f"@{display_host}{reset}", + f"{dim} ssh -L {port}:localhost:{port} " + f"@{display_host}{reset}", flush = True, ) print( @@ -520,7 +527,9 @@ def _find_free_port( candidate = start + offset if _is_port_free(host, candidate): return candidate - raise RuntimeError(f"Could not find a free port in range {start}-{start + max_attempts - 1}") + raise RuntimeError( + f"Could not find a free port in range {start}-{start + max_attempts - 1}" + ) from utils.paths.storage_roots import studio_root as _studio_root @@ -679,7 +688,9 @@ def _iter_frontend_fallback_candidates() -> "list[Path]": continue # Tolerate single/multi-line dict literals; [^}]* rejects nested # dicts, which the setuptools editable template never emits. - m = re.search(r"^MAPPING\s*(?::[^=]*)?=\s*(\{[^}]*\})", src, re.M | re.S) + m = re.search( + r"^MAPPING\s*(?::[^=]*)?=\s*(\{[^}]*\})", src, re.M | re.S + ) if not m: continue try: @@ -882,7 +893,9 @@ def run_server( print("=" * 50) if blocker: pid, name = blocker - print(f"Port {original_port} is already in use by " f"{name} (PID {pid}).") + print( + f"Port {original_port} is already in use by " f"{name} (PID {pid})." + ) else: print(f"Port {original_port} is already in use.") print(f"Unsloth Studio will use port {port} instead.") @@ -1029,7 +1042,9 @@ def run_server( global _cloudflare_url _cloudflare_url = None app.state.cloudflare_url = None - _cloudflare_enabled = cloudflare and host == "0.0.0.0" and not api_only and not _IS_COLAB + _cloudflare_enabled = ( + cloudflare and host == "0.0.0.0" and not api_only and not _IS_COLAB + ) if _cloudflare_enabled: try: # best-effort: any failure must not block startup from cloudflare_tunnel import start_studio_tunnel, stop_studio_tunnel @@ -1127,7 +1142,9 @@ if __name__ == "__main__": sys.stderr.write("=" * 60 + "\n") traceback.print_exc(file = sys.stderr) sys.stderr.write("\n") - sys.stderr.write("If a package is missing, try re-running: unsloth studio setup\n") + sys.stderr.write( + "If a package is missing, try re-running: unsloth studio setup\n" + ) sys.stderr.flush() sys.exit(1) diff --git a/studio/backend/state/tool_approvals.py b/studio/backend/state/tool_approvals.py index f66226b61d..414f1e4a29 100644 --- a/studio/backend/state/tool_approvals.py +++ b/studio/backend/state/tool_approvals.py @@ -106,7 +106,9 @@ def request_tool_decision( ): """Register and wait in one call (when the slot is not needed early).""" slot = begin_tool_decision(session_id, approval_id) - return wait_tool_decision(slot, approval_id, cancel_event = cancel_event, timeout = timeout) + return wait_tool_decision( + slot, approval_id, cancel_event = cancel_event, timeout = timeout + ) def resolve_tool_decision( diff --git a/studio/backend/state/tool_policy.py b/studio/backend/state/tool_policy.py index 9b0fc7d6cb..9343a39806 100644 --- a/studio/backend/state/tool_policy.py +++ b/studio/backend/state/tool_policy.py @@ -21,7 +21,9 @@ def get_tool_policy() -> Optional[bool]: def set_tool_policy(value: Optional[bool]) -> None: if value is not None and not isinstance(value, bool): - raise TypeError(f"tool_policy must be Optional[bool], got {type(value).__name__}") + raise TypeError( + f"tool_policy must be Optional[bool], got {type(value).__name__}" + ) global _tool_policy _tool_policy = value diff --git a/studio/backend/storage/mcp_servers_db.py b/studio/backend/storage/mcp_servers_db.py index 6482bae140..ba2e4eea7c 100644 --- a/studio/backend/storage/mcp_servers_db.py +++ b/studio/backend/storage/mcp_servers_db.py @@ -29,9 +29,13 @@ def _ensure_schema(conn: sqlite3.Connection) -> None: """ ) # Backfill use_oauth for pre-existing DBs. - cols = {r["name"] for r in conn.execute("PRAGMA table_info(mcp_servers)").fetchall()} + cols = { + r["name"] for r in conn.execute("PRAGMA table_info(mcp_servers)").fetchall() + } if "use_oauth" not in cols: - conn.execute("ALTER TABLE mcp_servers ADD COLUMN use_oauth INTEGER NOT NULL DEFAULT 0") + conn.execute( + "ALTER TABLE mcp_servers ADD COLUMN use_oauth INTEGER NOT NULL DEFAULT 0" + ) def get_connection() -> sqlite3.Connection: diff --git a/studio/backend/storage/providers_db.py b/studio/backend/storage/providers_db.py index 07165cbe70..c0809089c4 100644 --- a/studio/backend/storage/providers_db.py +++ b/studio/backend/storage/providers_db.py @@ -59,7 +59,9 @@ def get_connection() -> sqlite3.Connection: return conn -def create_provider(id: str, provider_type: str, display_name: str, base_url: str) -> None: +def create_provider( + id: str, provider_type: str, display_name: str, base_url: str +) -> None: """Insert a new provider configuration.""" now = datetime.now(timezone.utc).isoformat() conn = get_connection() @@ -137,7 +139,9 @@ def list_providers() -> list[dict]: """List all provider configurations, ordered by creation time.""" conn = get_connection() try: - rows = conn.execute("SELECT * FROM llm_providers ORDER BY created_at").fetchall() + rows = conn.execute( + "SELECT * FROM llm_providers ORDER BY created_at" + ).fetchall() return [dict(row) for row in rows] finally: conn.close() diff --git a/studio/backend/storage/studio_db.py b/studio/backend/storage/studio_db.py index 85cfacbc27..7ad8868b14 100644 --- a/studio/backend/storage/studio_db.py +++ b/studio/backend/storage/studio_db.py @@ -85,7 +85,9 @@ def _delete_project_workspace(project: dict) -> None: try: root_resolved = root.resolve(strict = False) except (OSError, RuntimeError, ValueError): - logger.warning("Skipping project workspace delete for invalid path %r", root_path) + logger.warning( + "Skipping project workspace delete for invalid path %r", root_path + ) return project_id = str(project["id"]) @@ -149,7 +151,9 @@ def _ensure_schema(conn: sqlite3.Connection) -> None: ) """ ) - existing_cols = {row[1] for row in conn.execute("PRAGMA table_info(training_runs)").fetchall()} + existing_cols = { + row[1] for row in conn.execute("PRAGMA table_info(training_runs)").fetchall() + } if "display_name" not in existing_cols: conn.execute("ALTER TABLE training_runs ADD COLUMN display_name TEXT") conn.execute( @@ -169,7 +173,9 @@ def _ensure_schema(conn: sqlite3.Connection) -> None: ) """ ) - conn.execute("CREATE INDEX IF NOT EXISTS idx_metrics_run_id ON training_metrics(run_id)") + conn.execute( + "CREATE INDEX IF NOT EXISTS idx_metrics_run_id ON training_metrics(run_id)" + ) # Windows: COLLATE NOCASE so C:\Models and c:\models dedup. Elsewhere keep # case-sensitive BINARY so /Models and /models stay distinct. collation = "COLLATE NOCASE" if platform.system() == "Windows" else "" @@ -226,9 +232,13 @@ def _ensure_schema(conn: sqlite3.Connection) -> None: if "project_id" not in chat_thread_cols: conn.execute("ALTER TABLE chat_threads ADD COLUMN project_id TEXT") if "openai_code_exec_container_id" not in chat_thread_cols: - conn.execute("ALTER TABLE chat_threads ADD COLUMN openai_code_exec_container_id TEXT") + conn.execute( + "ALTER TABLE chat_threads ADD COLUMN openai_code_exec_container_id TEXT" + ) if "anthropic_code_exec_container_id" not in chat_thread_cols: - conn.execute("ALTER TABLE chat_threads ADD COLUMN anthropic_code_exec_container_id TEXT") + conn.execute( + "ALTER TABLE chat_threads ADD COLUMN anthropic_code_exec_container_id TEXT" + ) conn.execute( """ CREATE TABLE IF NOT EXISTS chat_messages ( @@ -246,7 +256,9 @@ def _ensure_schema(conn: sqlite3.Connection) -> None: conn.execute( "CREATE INDEX IF NOT EXISTS idx_chat_threads_model_type_created_at ON chat_threads(model_type, created_at)" ) - conn.execute("CREATE INDEX IF NOT EXISTS idx_chat_threads_pair_id ON chat_threads(pair_id)") + conn.execute( + "CREATE INDEX IF NOT EXISTS idx_chat_threads_pair_id ON chat_threads(pair_id)" + ) conn.execute( "CREATE INDEX IF NOT EXISTS idx_chat_threads_project_id ON chat_threads(project_id)" ) @@ -336,7 +348,9 @@ def _prompt_entry_from_row(row: sqlite3.Row) -> dict: def list_prompt_entries() -> list[dict]: conn = get_connection() try: - rows = conn.execute("SELECT * FROM prompt_entries ORDER BY created_at DESC").fetchall() + rows = conn.execute( + "SELECT * FROM prompt_entries ORDER BY created_at DESC" + ).fetchall() return [_prompt_entry_from_row(r) for r in rows] finally: conn.close() @@ -391,7 +405,10 @@ def bulk_upsert_prompt_entries(entries: list[dict]) -> int: text = excluded.text, updated_at = excluded.updated_at """, - [(e["id"], e["name"], e["text"], e["createdAt"], e["updatedAt"]) for e in entries], + [ + (e["id"], e["name"], e["text"], e["createdAt"], e["updatedAt"]) + for e in entries + ], ) conn.commit() return len(entries) @@ -412,7 +429,9 @@ def _prompt_list_from_row(row: sqlite3.Row) -> dict: def list_prompt_lists_db() -> list[dict]: conn = get_connection() try: - rows = conn.execute("SELECT * FROM prompt_lists ORDER BY created_at DESC").fetchall() + rows = conn.execute( + "SELECT * FROM prompt_lists ORDER BY created_at DESC" + ).fetchall() return [_prompt_list_from_row(r) for r in rows] finally: conn.close() @@ -679,7 +698,9 @@ def list_runs(limit: int = 50, offset: int = 0) -> dict: try: run["loss_sparkline"] = json.loads(sparkline) except (json.JSONDecodeError, TypeError): - logger.debug("Failed to parse loss_sparkline for run %s", run.get("id")) + logger.debug( + "Failed to parse loss_sparkline for run %s", run.get("id") + ) run["loss_sparkline"] = None runs.append(run) return {"runs": runs, "total": total} @@ -755,7 +776,9 @@ def get_resumable_run_by_output_dir(output_dir: str) -> Optional[dict]: try: run["loss_sparkline"] = json.loads(sparkline) except (json.JSONDecodeError, TypeError): - logger.debug("Failed to parse loss_sparkline for output_dir %s", output_dir) + logger.debug( + "Failed to parse loss_sparkline for output_dir %s", output_dir + ) run["loss_sparkline"] = None return run finally: @@ -1289,7 +1312,9 @@ def _parse_chat_setting_json(key: str, value_json: str) -> tuple[bool, Any]: return False, None -def _load_chat_settings_for_merge(conn: sqlite3.Connection) -> tuple[dict[str, Any], set[str]]: +def _load_chat_settings_for_merge( + conn: sqlite3.Connection, +) -> tuple[dict[str, Any], set[str]]: rows = conn.execute("SELECT key, value_json FROM chat_settings").fetchall() current: dict[str, Any] = {} corrupt: set[str] = set() @@ -1425,8 +1450,12 @@ def sync_chat_messages( m.get("parentId"), m["role"], json.dumps(m.get("content", [])), - json.dumps(m.get("attachments")) if m.get("attachments") is not None else None, - json.dumps(m.get("metadata")) if m.get("metadata") is not None else None, + json.dumps(m.get("attachments")) + if m.get("attachments") is not None + else None, + json.dumps(m.get("metadata")) + if m.get("metadata") is not None + else None, int(m["createdAt"]), ) for m in messages @@ -1506,7 +1535,9 @@ def list_chat_messages_for_threads(thread_ids: list[str]) -> list[dict]: def get_app_setting(key: str, fallback = None): conn = get_connection() try: - row = conn.execute("SELECT value_json FROM app_settings WHERE key = ?", (key,)).fetchone() + row = conn.execute( + "SELECT value_json FROM app_settings WHERE key = ?", (key,) + ).fetchone() if row is None: return fallback return _json_loads(row["value_json"], fallback) @@ -1531,7 +1562,9 @@ def upsert_app_settings(settings: dict[str, Any]) -> dict[str, Any]: [(key, json.dumps(value), now) for key, value in settings.items()], ) conn.commit() - rows = conn.execute("SELECT key, value_json FROM app_settings ORDER BY key").fetchall() + rows = conn.execute( + "SELECT key, value_json FROM app_settings ORDER BY key" + ).fetchall() return {row["key"]: _json_loads(row["value_json"], None) for row in rows} finally: conn.close() @@ -1540,7 +1573,9 @@ def upsert_app_settings(settings: dict[str, Any]) -> dict[str, Any]: def list_chat_settings() -> dict[str, Any]: conn = get_connection() try: - rows = conn.execute("SELECT key, value_json FROM chat_settings ORDER BY key").fetchall() + rows = conn.execute( + "SELECT key, value_json FROM chat_settings ORDER BY key" + ).fetchall() settings: dict[str, Any] = {} for row in rows: settings[row["key"]] = _json_loads(row["value_json"], None) @@ -1571,7 +1606,9 @@ def upsert_chat_settings(settings: dict[str, Any]) -> dict[str, Any]: conn.close() -def _deep_merge_settings(current: dict[str, Any], updates: dict[str, Any]) -> dict[str, Any]: +def _deep_merge_settings( + current: dict[str, Any], updates: dict[str, Any] +) -> dict[str, Any]: merged = dict(current) for key, value in updates.items(): current_value = merged.get(key) @@ -1592,7 +1629,9 @@ def upsert_chat_settings_merge(updates: dict[str, Any]) -> dict[str, Any]: conn.execute("BEGIN IMMEDIATE") current, corrupt = _load_chat_settings_for_merge(conn) unsafe_partial_keys = [ - key for key, value in updates.items() if key in corrupt and isinstance(value, dict) + key + for key, value in updates.items() + if key in corrupt and isinstance(value, dict) ] if unsafe_partial_keys: conn.commit() @@ -1633,7 +1672,9 @@ def list_chat_legacy_imports() -> list[str]: """Return the legacy_thread_id of every thread already imported.""" conn = get_connection() try: - rows = conn.execute("SELECT legacy_thread_id FROM chat_legacy_imports").fetchall() + rows = conn.execute( + "SELECT legacy_thread_id FROM chat_legacy_imports" + ).fetchall() return [row[0] for row in rows] finally: conn.close() diff --git a/studio/backend/tests/test_anthropic_citations_edge.py b/studio/backend/tests/test_anthropic_citations_edge.py index f89e5ebda8..18222c76fa 100644 --- a/studio/backend/tests/test_anthropic_citations_edge.py +++ b/studio/backend/tests/test_anthropic_citations_edge.py @@ -79,7 +79,8 @@ def _capture( client = _make_client() try: async for line in client.stream_chat_completion( - messages = messages or [{"role": "user", "content": "what color is grass?"}], + messages = messages + or [{"role": "user", "content": "what color is grass?"}], model = "claude-opus-4-7", max_tokens = 64, ): @@ -157,7 +158,10 @@ def _citation_payload(body: str) -> dict: except json.JSONDecodeError: continue tool_event = payload.get("_toolEvent") if isinstance(payload, dict) else None - if isinstance(tool_event, dict) and tool_event.get("type") == "document_citations": + if ( + isinstance(tool_event, dict) + and tool_event.get("type") == "document_citations" + ): return tool_event raise AssertionError("document_citations event not parsed out of SSE body") diff --git a/studio/backend/tests/test_anthropic_code_execution.py b/studio/backend/tests/test_anthropic_code_execution.py index 22bbb19125..f334f76595 100644 --- a/studio/backend/tests/test_anthropic_code_execution.py +++ b/studio/backend/tests/test_anthropic_code_execution.py @@ -174,7 +174,9 @@ def test_no_code_execution_tool_when_pill_off(monkeypatch): # Pill off: no code_execution variant on the wire. assert all("code_execution" not in (t.get("type") or "") for t in tools) # Beta header must omit code-execution when the tool is off (opt-in only). - assert "code-execution-2025-08-25" not in captured["headers"].get("anthropic-beta", "") + assert "code-execution-2025-08-25" not in captured["headers"].get( + "anthropic-beta", "" + ) def test_bash_code_execution_emits_tool_start_and_end(monkeypatch): @@ -247,7 +249,11 @@ def test_bash_code_execution_emits_tool_start_and_end(monkeypatch): assert start["tool_name"] == "code_execution" assert start["tool_call_id"] == "srvtoolu_1" # `_server_tool: True` marks a provider-side synthetic tool card. - assert start["arguments"] == {"kind": "bash", "command": "ls -la", "_server_tool": True} + assert start["arguments"] == { + "kind": "bash", + "command": "ls -la", + "_server_tool": True, + } assert end["type"] == "tool_end" assert end["tool_call_id"] == "srvtoolu_1" diff --git a/studio/backend/tests/test_anthropic_compaction.py b/studio/backend/tests/test_anthropic_compaction.py index 1528eebe8b..bdbf5787db 100644 --- a/studio/backend/tests/test_anthropic_compaction.py +++ b/studio/backend/tests/test_anthropic_compaction.py @@ -122,9 +122,13 @@ def test_supported_model_attaches_compaction_block_and_beta(monkeypatch): def test_threshold_clamped_to_50k_minimum(monkeypatch): # Below-min values get clamped UP so we don't 400 upstream. captured = _capture(monkeypatch, "claude-opus-4-7", 60_000) - assert captured["body"]["context_management"]["edits"][0]["trigger"]["value"] == 60_000 + assert ( + captured["body"]["context_management"]["edits"][0]["trigger"]["value"] == 60_000 + ) captured = _capture(monkeypatch, "claude-opus-4-7", 1) - assert captured["body"]["context_management"]["edits"][0]["trigger"]["value"] == 50_000 + assert ( + captured["body"]["context_management"]["edits"][0]["trigger"]["value"] == 50_000 + ) # ── beta header merge with code execution ──────────────────────────── @@ -202,7 +206,9 @@ def test_chat_completion_request_accepts_sub_50k_compaction_threshold(): # ── usage.iterations[] surfaces compaction tokens ────────────────── -def test_message_delta_iterations_array_aggregates_compaction_tokens(monkeypatch, capsys): +def test_message_delta_iterations_array_aggregates_compaction_tokens( + monkeypatch, capsys +): # On mid-stream compaction the message_delta usage carries # `iterations: [{type:"compaction", ...}, ...]`. Top-level tokens only cover # the `message` iteration, so the helper folds compaction totals into @@ -521,7 +527,9 @@ def test_build_external_messages_passes_compaction_for_anthropic_only(): } ) ] - out = _build_external_messages(msgs, supports_vision = True, provider_type = "anthropic") + out = _build_external_messages( + msgs, supports_vision = True, provider_type = "anthropic" + ) assert len(out) == 1 parts = out[0]["content"] assert parts[0] == {"type": "compaction", "content": "prior summary"} @@ -547,7 +555,9 @@ def test_build_external_messages_strips_compaction_for_non_anthropic_providers() ) ] for provider in ("openai", "deepseek", "mistral", "gemini", "kimi", "openrouter"): - out = _build_external_messages(msgs, supports_vision = True, provider_type = provider) + out = _build_external_messages( + msgs, supports_vision = True, provider_type = provider + ) assert len(out) == 1, (provider, out) parts = out[0]["content"] types = [p.get("type") for p in parts if isinstance(p, dict)] @@ -595,10 +605,14 @@ def test_build_external_messages_non_vision_anthropic_keeps_compaction(): } ) ] - out = _build_external_messages(msgs, supports_vision = False, provider_type = "anthropic") + out = _build_external_messages( + msgs, supports_vision = False, provider_type = "anthropic" + ) parts = out[0]["content"] assert {"type": "compaction", "content": "prior summary"} in parts # Non-anthropic + non-vision -> compaction stripped, text collapsed # back to a string. - out2 = _build_external_messages(msgs, supports_vision = False, provider_type = "deepseek") + out2 = _build_external_messages( + msgs, supports_vision = False, provider_type = "deepseek" + ) assert out2[0]["content"] == "answer", out2 diff --git a/studio/backend/tests/test_anthropic_fast_mode_edge.py b/studio/backend/tests/test_anthropic_fast_mode_edge.py index dd69d77590..1da5884368 100644 --- a/studio/backend/tests/test_anthropic_fast_mode_edge.py +++ b/studio/backend/tests/test_anthropic_fast_mode_edge.py @@ -270,7 +270,9 @@ def test_refusal_notice_appears_before_content_filter_chunk(monkeypatch): """The notice content delta must precede the finish_reason chunk.""" _, lines = _capture(monkeypatch, sse = _refusal_sse(), model = "claude-opus-4-7") notice_idx = next(i for i, l in enumerate(lines) if "stopped by Anthropic" in l) - filter_idx = next(i for i, l in enumerate(lines) if '"finish_reason": "content_filter"' in l) + filter_idx = next( + i for i, l in enumerate(lines) if '"finish_reason": "content_filter"' in l + ) assert notice_idx < filter_idx, (notice_idx, filter_idx, lines) @@ -421,7 +423,9 @@ def test_usage_speed_propagates_to_final_usage_chunk_fast(monkeypatch): def test_usage_speed_propagates_to_final_usage_chunk_standard(monkeypatch): _, lines = _capture(monkeypatch, sse = _fast_speed_sse(speed = "standard")) parsed = [ - json.loads(l[len("data: ") :]) for l in lines if l.startswith("data: ") and '"usage"' in l + json.loads(l[len("data: ") :]) + for l in lines + if l.startswith("data: ") and '"usage"' in l ] speeds = [p["usage"].get("speed") for p in parsed if "usage" in p] assert "standard" in speeds, parsed @@ -431,7 +435,9 @@ def test_usage_speed_absent_when_anthropic_does_not_report(monkeypatch): """Studio must not invent ``usage.speed`` when upstream omits it.""" _, lines = _capture(monkeypatch) parsed = [ - json.loads(l[len("data: ") :]) for l in lines if l.startswith("data: ") and '"usage"' in l + json.loads(l[len("data: ") :]) + for l in lines + if l.startswith("data: ") and '"usage"' in l ] for p in parsed: usage = p.get("usage") or {} diff --git a/studio/backend/tests/test_anthropic_messages.py b/studio/backend/tests/test_anthropic_messages.py index f8d3f44d4b..a10c060f70 100644 --- a/studio/backend/tests/test_anthropic_messages.py +++ b/studio/backend/tests/test_anthropic_messages.py @@ -68,7 +68,10 @@ class TestToolActionNudge: assert nudge.startswith("The current date is ") assert "Tools are available when they materially improve" in nudge assert "prefer using tools rather than answering from memory" not in nudge - assert "fetch its full content by calling web_search with the url parameter" in nudge + assert ( + "fetch its full content by calling web_search with the url parameter" + in nudge + ) assert "Use code execution for math" in nudge assert "render_html" not in nudge @@ -418,7 +421,10 @@ class TestAnthropicMessagesToOpenAI: ] result = anthropic_messages_to_openai(msgs) parts = result[0]["content"] - assert parts[1] == {"type": "image_url", "image_url": {"url": "https://x/y.png"}} + assert parts[1] == { + "type": "image_url", + "image_url": {"url": "https://x/y.png"}, + } def test_image_only_user_message_emits_no_text_part(self): msgs = [ @@ -564,7 +570,9 @@ class TestAnthropicToolsToOpenAI: assert [tool["function"]["name"] for tool in result] == ["web_search", "python"] def test_pydantic_model_input(self): - tool = AnthropicTool(name = "test", description = "desc", input_schema = {"type": "object"}) + tool = AnthropicTool( + name = "test", description = "desc", input_schema = {"type": "object"} + ) result = anthropic_tools_to_openai([tool]) assert result[0]["function"]["name"] == "test" @@ -664,8 +672,12 @@ class TestAnthropicStreamEmitter: } ) - first_payloads = [json.loads(event.split("data: ")[1]) for event in first_events] - second_payloads = [json.loads(event.split("data: ")[1]) for event in second_events] + first_payloads = [ + json.loads(event.split("data: ")[1]) for event in first_events + ] + second_payloads = [ + json.loads(event.split("data: ")[1]) for event in second_events + ] tool_starts = [ payload @@ -681,7 +693,9 @@ class TestAnthropicStreamEmitter: "index": tool_starts[0]["index"], "delta": { "type": "input_json_delta", - "partial_json": json.dumps({"code": ""}), + "partial_json": json.dumps( + {"code": ""} + ), }, } ] @@ -845,7 +859,9 @@ class TestAnthropicToolNonStreaming: response = asyncio.run(_anthropic_tool_non_streaming(_run_gen, "msg_1", "m")) body = json.loads(response.body) - tool_blocks = [block for block in body["content"] if block["type"] == "tool_use"] + tool_blocks = [ + block for block in body["content"] if block["type"] == "tool_use" + ] assert len(tool_blocks) == 1 assert tool_blocks[0]["type"] == "tool_use" @@ -948,14 +964,26 @@ class TestAnthropicPassthroughEmitter: events1 = e.feed_chunk( { "choices": [ - {"delta": {"tool_calls": [{"index": 0, "function": {"arguments": '{"cmd'}}]}} + { + "delta": { + "tool_calls": [ + {"index": 0, "function": {"arguments": '{"cmd'}} + ] + } + } ] } ) events2 = e.feed_chunk( { "choices": [ - {"delta": {"tool_calls": [{"index": 0, "function": {"arguments": '": "ls"}'}}]}} + { + "delta": { + "tool_calls": [ + {"index": 0, "function": {"arguments": '": "ls"}'}} + ] + } + } ] } ) @@ -1440,7 +1468,9 @@ class TestAnthropicMessagesToolRouting: assert exc.value.status_code == 400 assert "Mixing Anthropic server tools" in exc.value.detail - def test_mixed_rejected_when_client_tool_name_collides_with_server_alias(self, monkeypatch): + def test_mixed_rejected_when_client_tool_name_collides_with_server_alias( + self, monkeypatch + ): # Regression: a client tool sharing a name with a mapped server tool # (e.g. a custom "web_search") must still trigger the mixed-mode 400; # otherwise the post-name filter drops the client tool and silently @@ -1499,7 +1529,9 @@ class TestAnthropicMessagesToolRouting: assert exc.value.status_code == 400 assert "name" in exc.value.detail - def test_alias_named_client_tool_without_schema_rejected_with_400(self, monkeypatch): + def test_alias_named_client_tool_without_schema_rejected_with_400( + self, monkeypatch + ): # Regression: a typo'd client tool whose name collides with a Studio # alias (e.g. a custom "python" tool missing input_schema) must # surface a 400, not silently switch into Studio's built-in python @@ -1553,7 +1585,10 @@ class TestAnthropicMessagesToolRouting: with pytest.raises(HTTPException) as exc: _drive(anthropic_messages(payload, request = None, current_subject = "t")) assert exc.value.status_code == 400 - assert "confirm_tool_calls is not supported" in exc.value.detail["error"]["message"] + assert ( + "confirm_tool_calls is not supported" + in exc.value.detail["error"]["message"] + ) assert backend.calls == [] def test_per_request_enable_tools_false_blocks_server_tool_alias(self, monkeypatch): diff --git a/studio/backend/tests/test_anthropic_web_fetch.py b/studio/backend/tests/test_anthropic_web_fetch.py index bf24175256..2935446447 100644 --- a/studio/backend/tests/test_anthropic_web_fetch.py +++ b/studio/backend/tests/test_anthropic_web_fetch.py @@ -178,7 +178,9 @@ def test_no_web_fetch_tool_when_pill_off(monkeypatch): _drive(run()) tools = captured["body"].get("tools") or [] - assert all(t.get("type") not in ("web_fetch_20250910", "web_fetch_20260209") for t in tools) + assert all( + t.get("type") not in ("web_fetch_20250910", "web_fetch_20260209") for t in tools + ) # ── SSE translation ───────────────────────────────────────────────── @@ -246,7 +248,9 @@ def test_web_fetch_success_emits_tool_start_and_end(monkeypatch): client = _make_client() return await _collect( client._stream_anthropic( - messages = [{"role": "user", "content": "Fetch https://example.com/article"}], + messages = [ + {"role": "user", "content": "Fetch https://example.com/article"} + ], model = "claude-opus-4-7", temperature = 0.7, top_p = 0.95, @@ -264,7 +268,10 @@ def test_web_fetch_success_emits_tool_start_and_end(monkeypatch): assert start["tool_call_id"] == "srvtoolu_wf1" # `_server_tool: True` marks this a provider-side synthetic tool card # for the frontend's history serializer. - assert start["arguments"] == {"url": "https://example.com/article", "_server_tool": True} + assert start["arguments"] == { + "url": "https://example.com/article", + "_server_tool": True, + } assert end["type"] == "tool_end" assert end["tool_call_id"] == "srvtoolu_wf1" # The source pill uses Title / URL / snippet as parseSourcesFromResult expects. diff --git a/studio/backend/tests/test_audio_token_detection.py b/studio/backend/tests/test_audio_token_detection.py index 02d4790519..d07af172ba 100644 --- a/studio/backend/tests/test_audio_token_detection.py +++ b/studio/backend/tests/test_audio_token_detection.py @@ -18,7 +18,9 @@ def _classify(tokens: list[str]) -> str | None: def test_gemma3n_audio_soft_token_is_audio_vlm(): - assert _classify(["", "", ""]) == "audio_vlm" + assert ( + _classify(["", "", ""]) == "audio_vlm" + ) def test_gemma4_pipe_audio_token_is_audio_vlm(): diff --git a/studio/backend/tests/test_cached_gguf_routes.py b/studio/backend/tests/test_cached_gguf_routes.py index 5a4ca68ab4..ec0d5d3a74 100644 --- a/studio/backend/tests/test_cached_gguf_routes.py +++ b/studio/backend/tests/test_cached_gguf_routes.py @@ -65,7 +65,9 @@ def test_iter_gguf_paths_matches_extension_case_insensitively(tmp_path): assert result == ["Q4_K_M.gguf", "Q8_0.GGUF"] -def test_list_cached_gguf_includes_non_suffix_repo_when_cache_contains_gguf(monkeypatch, tmp_path): +def test_list_cached_gguf_includes_non_suffix_repo_when_cache_contains_gguf( + monkeypatch, tmp_path +): repo = _repo( "HauhauCS/Gemma-4-E4B-Uncensored-HauhauCS-Aggressive", [_file("Q4_K_M.gguf", 5_000), _file("README.md", 10)], @@ -127,7 +129,9 @@ def test_list_cached_gguf_skips_repos_without_positive_gguf_size(monkeypatch, tm assert result["cached"] == [] -def test_list_cached_gguf_keeps_largest_duplicate_repo_across_scans(monkeypatch, tmp_path): +def test_list_cached_gguf_keeps_largest_duplicate_repo_across_scans( + monkeypatch, tmp_path +): smaller = _repo( "Org/Dupe", [_file("Q4_K_M.gguf", 2_000)], @@ -188,7 +192,9 @@ def test_list_cached_gguf_dedupes_shared_blobs_across_revisions(monkeypatch, tmp ] -def test_list_cached_models_skips_non_suffix_repo_when_gguf_files_exist(monkeypatch, tmp_path): +def test_list_cached_models_skips_non_suffix_repo_when_gguf_files_exist( + monkeypatch, tmp_path +): mixed = _repo( "Org/MixedRepo", [ @@ -209,7 +215,9 @@ def test_list_cached_models_skips_non_suffix_repo_when_gguf_files_exist(monkeypa assert result["cached"] == [] -def test_list_cached_gguf_includes_mixed_repo_with_gguf_and_safetensors(monkeypatch, tmp_path): +def test_list_cached_gguf_includes_mixed_repo_with_gguf_and_safetensors( + monkeypatch, tmp_path +): """Mixed repo still surfaces in cached-gguf as a GGUF download.""" mixed = _repo( "Org/MixedRepo", @@ -263,7 +271,9 @@ def test_list_cached_gguf_handles_none_size_on_disk(monkeypatch, tmp_path): ] -def test_list_cached_gguf_skips_malformed_repo_without_wiping_response(monkeypatch, tmp_path): +def test_list_cached_gguf_skips_malformed_repo_without_wiping_response( + monkeypatch, tmp_path +): """One repo raising during classification must not poison the response.""" class _ExplodingRepo: @@ -344,7 +354,9 @@ def test_list_cached_models_includes_repo_with_only_mmproj_gguf(monkeypatch, tmp assert result["cached"] == [{"repo_id": "Org/MmprojAux", "size_bytes": 15_000}] -def test_list_cached_gguf_includes_vision_repo_with_main_gguf_and_mmproj(monkeypatch, tmp_path): +def test_list_cached_gguf_includes_vision_repo_with_main_gguf_and_mmproj( + monkeypatch, tmp_path +): """A vision GGUF repo (main weight + mmproj) is a GGUF repo; reported size is the main weight only, since mmproj is filtered at classification.""" vision_repo = _repo( @@ -426,7 +438,9 @@ def test_all_hf_cache_scans_survives_inaccessible_aux_cache(monkeypatch, tmp_pat ] -def test_list_cached_gguf_sorts_newest_first_grouping_by_latest_quant(monkeypatch, tmp_path): +def test_list_cached_gguf_sorts_newest_first_grouping_by_latest_quant( + monkeypatch, tmp_path +): """Downloaded is ordered newest-first, and a multi-quant repo is placed by its most recently downloaded quant (``last_modified`` = newest quant).""" older = _repo( @@ -459,13 +473,20 @@ def test_list_cached_gguf_sorts_newest_first_grouping_by_latest_quant(monkeypatc def test_list_cached_gguf_dedupe_keeps_newest_timestamp(monkeypatch, tmp_path): """Same repo in two caches with equal size keeps the newest last_modified, regardless of scan order.""" - older = _repo("org/dupe", [_gfile("dupe-Q4_K_M.gguf", 5_000, 1_000.0)], tmp_path / "a") - newer = _repo("org/dupe", [_gfile("dupe-Q4_K_M.gguf", 5_000, 9_000.0)], tmp_path / "b") + older = _repo( + "org/dupe", [_gfile("dupe-Q4_K_M.gguf", 5_000, 1_000.0)], tmp_path / "a" + ) + newer = _repo( + "org/dupe", [_gfile("dupe-Q4_K_M.gguf", 5_000, 9_000.0)], tmp_path / "b" + ) for scans in ([older, newer], [newer, older]): # both orders monkeypatch.setattr( models_route, "_all_hf_cache_scans", - lambda s = scans: [SimpleNamespace(repos = [s[0]]), SimpleNamespace(repos = [s[1]])], + lambda s = scans: [ + SimpleNamespace(repos = [s[0]]), + SimpleNamespace(repos = [s[1]]), + ], ) result = asyncio.run(models_route.list_cached_gguf(current_subject = "t")) assert len(result["cached"]) == 1 @@ -479,17 +500,23 @@ def test_gguf_variants_mmproj_does_not_mark_quant_downloaded(monkeypatch, tmp_pa import huggingface_hub.constants as hf_constants variants = [ - SimpleNamespace(filename = "model-Q4_K_M.gguf", quant = "Q4_K_M", size_bytes = 10_000), + SimpleNamespace( + filename = "model-Q4_K_M.gguf", quant = "Q4_K_M", size_bytes = 10_000 + ), SimpleNamespace(filename = "model-F16.gguf", quant = "F16", size_bytes = 20_000), ] monkeypatch.setattr( - models_route, "list_gguf_variants", lambda repo_id, hf_token = None: (variants, True) + models_route, + "list_gguf_variants", + lambda repo_id, hf_token = None: (variants, True), ) monkeypatch.setattr(hf_constants, "HF_HUB_CACHE", str(tmp_path)) snap = tmp_path / "models--org--repo" / "snapshots" / "rev" snap.mkdir(parents = True) - (snap / "model-Q4_K_M.gguf").write_bytes(b"x" * 10_000) # real weight, fully present + (snap / "model-Q4_K_M.gguf").write_bytes( + b"x" * 10_000 + ) # real weight, fully present (snap / "mmproj-F16.gguf").write_bytes(b"y" * 20_000) # mmproj adapter, label "F16" result = asyncio.run( diff --git a/studio/backend/tests/test_chat_history_routes.py b/studio/backend/tests/test_chat_history_routes.py index e3aeb4eb22..18a8bdf7ef 100644 --- a/studio/backend/tests/test_chat_history_routes.py +++ b/studio/backend/tests/test_chat_history_routes.py @@ -109,13 +109,16 @@ def test_chat_inference_settings_covers_frontend_persisted_fields(): pytest.skip("frontend runtime.ts not present") with open(runtime_ts, encoding = "utf-8") as fh: - block = re.search(r"interface InferenceParams \{(.*?)\n\}", fh.read(), re.DOTALL) + block = re.search( + r"interface InferenceParams \{(.*?)\n\}", fh.read(), re.DOTALL + ) assert block, "InferenceParams interface not found in runtime.ts" persisted = set(re.findall(r"^\s*(\w+)\??:", block.group(1), re.M)) - {"checkpoint"} backend = set(chat_history.ChatInferenceSettings.model_fields) assert persisted == backend, ( - f"schema drift: frontend-only {persisted - backend}, " f"backend-only {backend - persisted}" + f"schema drift: frontend-only {persisted - backend}, " + f"backend-only {backend - persisted}" ) diff --git a/studio/backend/tests/test_chat_history_storage.py b/studio/backend/tests/test_chat_history_storage.py index bc74ec172a..951a5bf5b3 100644 --- a/studio/backend/tests/test_chat_history_storage.py +++ b/studio/backend/tests/test_chat_history_storage.py @@ -231,16 +231,22 @@ def test_settings_merge_atomic_under_concurrency(tmp_path, monkeypatch): def test_settings_merge_preserves_nested_keys(tmp_path, monkeypatch): _reset_studio_db(tmp_path, monkeypatch) - studio_db.upsert_chat_settings_merge({"inferenceParams": {"temperature": 0.5, "topP": 0.8}}) + studio_db.upsert_chat_settings_merge( + {"inferenceParams": {"temperature": 0.5, "topP": 0.8}} + ) studio_db.upsert_chat_settings_merge({"inferenceParams": {"temperature": 0.9}}) params = studio_db.list_chat_settings()["inferenceParams"] assert params == {"temperature": 0.9, "topP": 0.8} -def test_settings_merge_quarantines_corrupt_json_and_rejects_partial_patch(tmp_path, monkeypatch): +def test_settings_merge_quarantines_corrupt_json_and_rejects_partial_patch( + tmp_path, monkeypatch +): _reset_studio_db(tmp_path, monkeypatch) - studio_db.upsert_chat_settings_merge({"inferenceParams": {"temperature": 0.5, "topP": 0.8}}) + studio_db.upsert_chat_settings_merge( + {"inferenceParams": {"temperature": 0.5, "topP": 0.8}} + ) conn = studio_db.get_connection() try: conn.execute( @@ -288,7 +294,9 @@ def test_settings_merge_replaces_corrupt_scalar_after_quarantine(tmp_path, monke assert settings["autoTitle"] is True conn = studio_db.get_connection() try: - quarantined = conn.execute("SELECT key, reason FROM chat_settings_quarantine").fetchall() + quarantined = conn.execute( + "SELECT key, reason FROM chat_settings_quarantine" + ).fetchall() finally: conn.close() assert [(row["key"], row["reason"]) for row in quarantined] == [ @@ -344,7 +352,11 @@ def test_legacy_imports_records_and_lists(tmp_path, monkeypatch): ) assert accepted == 3 assert inserted == 3 - assert set(studio_db.list_chat_legacy_imports()) == {"legacy-a", "legacy-b", "legacy-c"} + assert set(studio_db.list_chat_legacy_imports()) == { + "legacy-a", + "legacy-b", + "legacy-c", + } def test_legacy_imports_is_idempotent(tmp_path, monkeypatch): @@ -358,7 +370,11 @@ def test_legacy_imports_is_idempotent(tmp_path, monkeypatch): assert (accepted1, inserted1) == (2, 2) # legacy-b is already in the ledger, only legacy-c is genuinely new. assert (accepted2, inserted2) == (2, 1) - assert set(studio_db.list_chat_legacy_imports()) == {"legacy-a", "legacy-b", "legacy-c"} + assert set(studio_db.list_chat_legacy_imports()) == { + "legacy-a", + "legacy-b", + "legacy-c", + } def test_legacy_imports_dedups_input(tmp_path, monkeypatch): diff --git a/studio/backend/tests/test_cloudflare_tunnel.py b/studio/backend/tests/test_cloudflare_tunnel.py index 8042240b64..82d0cdabfa 100644 --- a/studio/backend/tests/test_cloudflare_tunnel.py +++ b/studio/backend/tests/test_cloudflare_tunnel.py @@ -201,7 +201,9 @@ def test_cache_path_uses_exe_on_windows(monkeypatch, tmp_path): def test_ensure_windows_downloads_exe(monkeypatch, tmp_path): cached = tmp_path / "cloudflared.exe" monkeypatch.setattr(ct, "find_cloudflared", lambda: None) - monkeypatch.setattr(ct, "_asset_name", lambda: ("cloudflared-windows-amd64.exe", False)) + monkeypatch.setattr( + ct, "_asset_name", lambda: ("cloudflared-windows-amd64.exe", False) + ) monkeypatch.setattr(ct, "_cache_path", lambda: cached) monkeypatch.setattr(ct.sys, "platform", "win32") @@ -212,7 +214,9 @@ def test_ensure_windows_downloads_exe(monkeypatch, tmp_path): monkeypatch.setattr(ct, "_download", fake_download) # chmod is skipped on Windows; would raise on a path that does not exist yet. - monkeypatch.setattr(ct.os, "chmod", lambda *a, **k: pytest.fail("chmod called on win32")) + monkeypatch.setattr( + ct.os, "chmod", lambda *a, **k: pytest.fail("chmod called on win32") + ) assert ct.ensure_cloudflared() == str(cached) assert cached.read_bytes() == b"MZ" @@ -220,7 +224,9 @@ def test_ensure_windows_downloads_exe(monkeypatch, tmp_path): def test_ensure_macos_extracts_tgz_and_chmods(monkeypatch, tmp_path): cached = tmp_path / "cloudflared" monkeypatch.setattr(ct, "find_cloudflared", lambda: None) - monkeypatch.setattr(ct, "_asset_name", lambda: ("cloudflared-darwin-arm64.tgz", True)) + monkeypatch.setattr( + ct, "_asset_name", lambda: ("cloudflared-darwin-arm64.tgz", True) + ) monkeypatch.setattr(ct, "_cache_path", lambda: cached) monkeypatch.setattr(ct.sys, "platform", "darwin") @@ -328,7 +334,9 @@ def test_start_after_stop_does_not_spawn(monkeypatch): def poll(self): return 0 - monkeypatch.setattr(ct.subprocess, "Popen", lambda *a, **k: (spawned.append(a), _FakeProc())[1]) + monkeypatch.setattr( + ct.subprocess, "Popen", lambda *a, **k: (spawned.append(a), _FakeProc())[1] + ) t.stop() # proc is None -> no-op terminate, but marks the tunnel stopped t.start() # must short-circuit before Popen assert spawned == [] @@ -662,7 +670,10 @@ def test_start_studio_tunnel_aborts_retry_on_concurrent_shutdown(monkeypatch): def _func_param_defaults(source, func_name): tree = ast.parse(source) for node in ast.walk(tree): - if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)) and node.name == func_name: + if ( + isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)) + and node.name == func_name + ): args = node.args.args defaults = node.args.defaults offset = len(args) - len(defaults) @@ -736,9 +747,14 @@ def _run_print_cloudflare_line(monkeypatch, *, cloudflare_url, public_reachable) def test_cloudflare_line_reworded_when_public_unreachable(monkeypatch): out = _run_print_cloudflare_line( - monkeypatch, cloudflare_url = "https://x.trycloudflare.com", public_reachable = False + monkeypatch, + cloudflare_url = "https://x.trycloudflare.com", + public_reachable = False, + ) + assert ( + "Use the secure link access via Cloudflare instead: https://x.trycloudflare.com" + in out ) - assert "Use the secure link access via Cloudflare instead: https://x.trycloudflare.com" in out def test_cloudflare_line_default_wording_when_reachable(monkeypatch): @@ -759,5 +775,7 @@ def test_cloudflare_line_default_wording_when_unknown(monkeypatch): def test_cloudflare_line_prints_nothing_without_tunnel(monkeypatch): - out = _run_print_cloudflare_line(monkeypatch, cloudflare_url = None, public_reachable = False) + out = _run_print_cloudflare_line( + monkeypatch, cloudflare_url = None, public_reachable = False + ) assert out == "" diff --git a/studio/backend/tests/test_context_overflow_truncation.py b/studio/backend/tests/test_context_overflow_truncation.py index 4f4c240934..ddcb88f801 100644 --- a/studio/backend/tests/test_context_overflow_truncation.py +++ b/studio/backend/tests/test_context_overflow_truncation.py @@ -53,7 +53,10 @@ def _tool_turn(i: int, result_chars: int = 400) -> list[dict]: { "id": f"call_{i}", "type": "function", - "function": {"name": "read", "arguments": f'{{"filePath":"/f{i}"}}'}, + "function": { + "name": "read", + "arguments": f'{{"filePath":"/f{i}"}}', + }, } ], }, @@ -109,7 +112,10 @@ def test_truncation_never_orphans_tool_results(): new, dropped = _truncate_middle_messages(msgs, keep_ratio = 0.4) assert dropped > 0 surviving_call_ids = { - tc["id"] for m in new if m.get("role") == "assistant" for tc in (m.get("tool_calls") or []) + tc["id"] + for m in new + if m.get("role") == "assistant" + for tc in (m.get("tool_calls") or []) } for m in new: if m.get("role") == "tool": @@ -266,8 +272,12 @@ class _FakeEmptyBackend: def test_v1_models_exposes_real_context_window(monkeypatch): - monkeypatch.setattr(routes_mod, "get_llama_cpp_backend", lambda: _FakeLlamaBackend()) - monkeypatch.setattr(routes_mod, "get_inference_backend", lambda: _FakeEmptyBackend()) + monkeypatch.setattr( + routes_mod, "get_llama_cpp_backend", lambda: _FakeLlamaBackend() + ) + monkeypatch.setattr( + routes_mod, "get_inference_backend", lambda: _FakeEmptyBackend() + ) models = _openai_model_objects() assert len(models) == 1 entry = models[0] diff --git a/studio/backend/tests/test_cpu_threads.py b/studio/backend/tests/test_cpu_threads.py index 2930c9f081..28f6902004 100644 --- a/studio/backend/tests/test_cpu_threads.py +++ b/studio/backend/tests/test_cpu_threads.py @@ -63,7 +63,9 @@ def test_cpu_thread_cap_is_opt_in(raw): # Anything that is not a positive integer raises a clear ValueError. -@pytest.mark.parametrize("raw", ["zero", "0", "-3", "1.5", "abc", "8a", "0x4", "1e3", "4 0"]) +@pytest.mark.parametrize( + "raw", ["zero", "0", "-3", "1.5", "abc", "8a", "0x4", "1e3", "4 0"] +) def test_cpu_thread_cap_requires_positive_integer(raw): with pytest.raises(ValueError, match = "must be a positive integer"): configure_cpu_threads({"UNSLOTH_CPU_THREADS": raw}) diff --git a/studio/backend/tests/test_datacenter_gpu_tuning.py b/studio/backend/tests/test_datacenter_gpu_tuning.py index fd9b291e8a..f1d2e41cc3 100644 --- a/studio/backend/tests/test_datacenter_gpu_tuning.py +++ b/studio/backend/tests/test_datacenter_gpu_tuning.py @@ -119,7 +119,9 @@ def test_is_datacenter_gpu_masked_host_physical_ids(monkeypatch): def test_is_datacenter_gpu_masked_host_reordered(monkeypatch): # Reordered mask preserves order: ordinal 0 -> physical 7, 1 -> 4, ... monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "7,4,5,6") - monkeypatch.setitem(sys.modules, "torch", _fake_torch(["NVIDIA H100 80GB HBM3"] * 4)) + monkeypatch.setitem( + sys.modules, "torch", _fake_torch(["NVIDIA H100 80GB HBM3"] * 4) + ) assert LlamaCppBackend._is_datacenter_gpu([7, 4]) is True @@ -220,7 +222,9 @@ def test_apply_env_multi_dc_gpu_sets_all(monkeypatch): def test_apply_env_none_indices_uses_visible_count(monkeypatch): # None on a 2x DC box -> multi-GPU flags applied. monkeypatch.delenv("UNSLOTH_DISABLE_DC_TUNING", raising = False) - monkeypatch.setitem(sys.modules, "torch", _fake_torch(["NVIDIA H100", "NVIDIA H100"])) + monkeypatch.setitem( + sys.modules, "torch", _fake_torch(["NVIDIA H100", "NVIDIA H100"]) + ) env: dict = {} assert LlamaCppBackend._apply_datacenter_env(env, None) is True assert env["GGML_CUDA_P2P"] == "1" @@ -229,7 +233,9 @@ def test_apply_env_none_indices_uses_visible_count(monkeypatch): def test_apply_env_consumer_gpu_is_noop(monkeypatch): monkeypatch.delenv("UNSLOTH_DISABLE_DC_TUNING", raising = False) - monkeypatch.setitem(sys.modules, "torch", _fake_torch(["NVIDIA GeForce RTX 4090"] * 2)) + monkeypatch.setitem( + sys.modules, "torch", _fake_torch(["NVIDIA GeForce RTX 4090"] * 2) + ) env: dict = {} assert LlamaCppBackend._apply_datacenter_env(env, [0, 1]) is False assert env == {} diff --git a/studio/backend/tests/test_dataset_upload_limits.py b/studio/backend/tests/test_dataset_upload_limits.py index dc6030edbf..0991059318 100644 --- a/studio/backend/tests/test_dataset_upload_limits.py +++ b/studio/backend/tests/test_dataset_upload_limits.py @@ -40,7 +40,9 @@ def isolate_upload_dir(tmp_path, monkeypatch): def test_dataset_upload_under_configured_cap_succeeds(isolate_upload_dir): upload = FakeUploadFile("sample.csv", [b"a,b\n1,2\n"]) response = asyncio.run( - datasets_route.upload_dataset(cast(UploadFile, upload), current_subject = "test-user") + datasets_route.upload_dataset( + cast(UploadFile, upload), current_subject = "test-user" + ) ) stored = Path(response.stored_path) assert response.filename == "sample.csv" @@ -56,7 +58,9 @@ def test_dataset_upload_over_configured_cap_removes_partial_file(isolate_upload_ ) with pytest.raises(HTTPException) as exc: asyncio.run( - datasets_route.upload_dataset(cast(UploadFile, upload), current_subject = "test-user") + datasets_route.upload_dataset( + cast(UploadFile, upload), current_subject = "test-user" + ) ) assert exc.value.status_code == 413 assert "Maximum is 1MB" in exc.value.detail diff --git a/studio/backend/tests/test_desktop_auth.py b/studio/backend/tests/test_desktop_auth.py index b4b39ba1a9..96fd3cdd70 100644 --- a/studio/backend/tests/test_desktop_auth.py +++ b/studio/backend/tests/test_desktop_auth.py @@ -51,8 +51,12 @@ def auth_client(): def data_recipe_jobs_module(): - route_path = Path(__file__).resolve().parents[1] / "routes" / "data_recipe" / "jobs.py" - spec = importlib.util.spec_from_file_location("_desktop_data_recipe_jobs", route_path) + route_path = ( + Path(__file__).resolve().parents[1] / "routes" / "data_recipe" / "jobs.py" + ) + spec = importlib.util.spec_from_file_location( + "_desktop_data_recipe_jobs", route_path + ) jobs_route = importlib.util.module_from_spec(spec) assert spec.loader is not None spec.loader.exec_module(jobs_route) @@ -261,7 +265,9 @@ def test_consume_refresh_token_concurrent_only_one_succeeds(tmp_path, monkeypatc results = list(pool.map(attempt, range(workers))) successes = [r for r in results if r is not None] - assert len(successes) == 1, f"expected exactly one consumer to win, got {len(successes)}" + assert ( + len(successes) == 1 + ), f"expected exactly one consumer to win, got {len(successes)}" assert successes[0] == (storage.DEFAULT_ADMIN_USERNAME, False) @@ -279,7 +285,9 @@ def test_desktop_session_uses_real_admin_identity_for_api_keys(): seed_user(must_change_password = True) raw = storage.create_desktop_secret() client = auth_client() - token = client.post("/api/auth/desktop-login", json = {"secret": raw}).json()["access_token"] + token = client.post("/api/auth/desktop-login", json = {"secret": raw}).json()[ + "access_token" + ] response = client.post( "/api/auth/api-keys", @@ -313,7 +321,9 @@ def test_local_recipe_token_authenticates_as_admin_for_desktop_user(loaded_local scheme = "Bearer", credentials = local_token, ) - assert asyncio.run(get_current_subject(credentials)) == storage.DEFAULT_ADMIN_USERNAME + assert ( + asyncio.run(get_current_subject(credentials)) == storage.DEFAULT_ADMIN_USERNAME + ) def test_local_recipe_token_authenticates_as_admin_for_web_user(loaded_local_model): @@ -333,7 +343,9 @@ def test_local_recipe_token_authenticates_as_admin_for_web_user(loaded_local_mod scheme = "Bearer", credentials = local_token, ) - assert asyncio.run(get_current_subject(credentials)) == storage.DEFAULT_ADMIN_USERNAME + assert ( + asyncio.run(get_current_subject(credentials)) == storage.DEFAULT_ADMIN_USERNAME + ) def test_desktop_login_rejects_invalid_secret(): @@ -529,9 +541,12 @@ if result.exit_code != 0: """ ).fetchone() app_secrets = { - row["key"]: row["value"] for row in conn.execute("SELECT key, value FROM app_secrets") + row["key"]: row["value"] + for row in conn.execute("SELECT key, value FROM app_secrets") + } + refresh_columns = { + row["name"] for row in conn.execute("PRAGMA table_info(refresh_tokens)") } - refresh_columns = {row["name"] for row in conn.execute("PRAGMA table_info(refresh_tokens)")} finally: conn.close() @@ -623,7 +638,9 @@ def test_update_password_clears_desktop_secret(): raw = storage.create_desktop_secret() assert storage.validate_desktop_secret(raw) == storage.DEFAULT_ADMIN_USERNAME - changed = storage.update_password(storage.DEFAULT_ADMIN_USERNAME, "new-admin-password") + changed = storage.update_password( + storage.DEFAULT_ADMIN_USERNAME, "new-admin-password" + ) assert changed is True assert storage.validate_desktop_secret(raw) is None @@ -639,7 +656,11 @@ def test_update_password_on_unknown_user_leaves_desktop_secret_intact(): def test_desktop_auth_provision_has_bounded_timeout(): rs_path = ( - Path(__file__).resolve().parents[3] / "studio" / "src-tauri" / "src" / "desktop_auth.rs" + Path(__file__).resolve().parents[3] + / "studio" + / "src-tauri" + / "src" + / "desktop_auth.rs" ) src = rs_path.read_text() start = src.index("async fn provision_desktop_auth(") diff --git a/studio/backend/tests/test_detect_mmproj_file.py b/studio/backend/tests/test_detect_mmproj_file.py index 64dd8ebd90..1449c4af45 100644 --- a/studio/backend/tests/test_detect_mmproj_file.py +++ b/studio/backend/tests/test_detect_mmproj_file.py @@ -132,7 +132,10 @@ def test_family_token_mistral_does_not_match_ministral(): assert _detect_family_token("Ministral-3-8B-Instruct-2512-BF16.gguf") == "ministral" assert _detect_family_token("Mistral-7B-Instruct-v0.3.gguf") == "mistral" assert _detect_family_token("Magistral-Small-2506-BF16.gguf") == "magistral" - assert _detect_family_token("Devstral-Small-2-24B-Instruct-2512-BF16.gguf") == "devstral" + assert ( + _detect_family_token("Devstral-Small-2-24B-Instruct-2512-BF16.gguf") + == "devstral" + ) def test_family_token_picks_leftmost_when_multiple_present(): @@ -157,7 +160,9 @@ def test_family_token_new_families_recognised(): def test_blocks_cross_family_for_new_token_pair(tmp_path: Path): """Nemotron weight + lone Gemma projector returns None.""" - model = _touch(tmp_path / "NVIDIA-Nemotron-3-Nano-Omni-30B-A3B-Reasoning-MXFP4_MOE.gguf") + model = _touch( + tmp_path / "NVIDIA-Nemotron-3-Nano-Omni-30B-A3B-Reasoning-MXFP4_MOE.gguf" + ) _touch(tmp_path / "gemma-4-26B-A4B-it.mmproj-q8_0.gguf") assert detect_mmproj_file(str(model)) is None diff --git a/studio/backend/tests/test_export_absolute_paths.py b/studio/backend/tests/test_export_absolute_paths.py index f333a590c2..6979a1c7df 100644 --- a/studio/backend/tests/test_export_absolute_paths.py +++ b/studio/backend/tests/test_export_absolute_paths.py @@ -154,7 +154,9 @@ def _install_lightweight_backend_stubs(monkeypatch): monkeypatch.setitem(sys.modules, "utils.models", utils_models) utils_model_config = types.ModuleType("utils.models.model_config") - utils_model_config._pick_best_gguf = lambda variants: variants[0] if variants else None + utils_model_config._pick_best_gguf = ( + lambda variants: variants[0] if variants else None + ) utils_model_config._extract_quant_label = lambda value: value utils_model_config.is_audio_input_type = lambda *args, **kwargs: None monkeypatch.setitem( @@ -268,7 +270,9 @@ def _install_export_backend_stubs(monkeypatch): def test_gguf_export_cleans_temp_dir_when_post_processing_fails(tmp_path, monkeypatch): _install_export_backend_stubs(monkeypatch) - export_mod = _load_module("test_core_export_backend", "core/export/export.py", monkeypatch) + export_mod = _load_module( + "test_core_export_backend", "core/export/export.py", monkeypatch + ) cwd = tmp_path / "cwd" save_dir = tmp_path / "export" @@ -310,7 +314,9 @@ def test_save_directory_validator_rejects_windows_parent_segments(monkeypatch): def test_save_directory_validator_allows_deep_absolute_paths(monkeypatch, tmp_path): _install_pydantic_stub(monkeypatch) - export_models = _load_module("test_models_export_deep_path", "models/export.py", monkeypatch) + export_models = _load_module( + "test_models_export_deep_path", "models/export.py", monkeypatch + ) deep_path = tmp_path for index in range(40): @@ -331,7 +337,9 @@ def test_save_directory_validator_rejects_long_path_component(monkeypatch, tmp_p export_models._validate_save_directory(str(tmp_path / ("a" * 256))) -def test_export_write_dir_accepts_external_absolute_but_read_dir_rejects(tmp_path, monkeypatch): +def test_export_write_dir_accepts_external_absolute_but_read_dir_rejects( + tmp_path, monkeypatch +): storage_roots = _load_module( "test_storage_roots_accept_external", "utils/paths/storage_roots.py", @@ -365,7 +373,10 @@ def test_export_write_dir_accepts_expanded_home_path(tmp_path, monkeypatch): else: monkeypatch.setenv("HOME", str(home)) - assert storage_roots.resolve_export_write_dir("~/exports/model") == home / "exports" / "model" + assert ( + storage_roots.resolve_export_write_dir("~/exports/model") + == home / "exports" / "model" + ) def test_resolve_export_write_dir_rejects_backslash_parent_segment(): @@ -378,7 +389,9 @@ def test_resolve_export_write_dir_rejects_backslash_parent_segment(): storage_roots.resolve_export_write_dir(r"exports\..\outside") -def test_export_write_dir_handles_non_native_windows_absolute_as_relative(tmp_path, monkeypatch): +def test_export_write_dir_handles_non_native_windows_absolute_as_relative( + tmp_path, monkeypatch +): storage_roots = _load_module( "test_storage_roots_non_native_windows_path", "utils/paths/storage_roots.py", diff --git a/studio/backend/tests/test_external_provider_proxy_env.py b/studio/backend/tests/test_external_provider_proxy_env.py index f17b655908..4fa482674e 100644 --- a/studio/backend/tests/test_external_provider_proxy_env.py +++ b/studio/backend/tests/test_external_provider_proxy_env.py @@ -32,7 +32,9 @@ def test_shared_http_client_ignores_unsupported_proxy_scheme(monkeypatch): def __init__(self, **kwargs): calls.append(kwargs) if kwargs.get("trust_env") is not False: - raise ValueError("Unknown scheme for proxy URL URL('socks4://127.0.0.1:12345')") + raise ValueError( + "Unknown scheme for proxy URL URL('socks4://127.0.0.1:12345')" + ) monkeypatch.setattr(ep_mod.httpx, "AsyncClient", FakeAsyncClient) diff --git a/studio/backend/tests/test_external_provider_usage_chunk.py b/studio/backend/tests/test_external_provider_usage_chunk.py index 3ec9133718..520a4a2169 100644 --- a/studio/backend/tests/test_external_provider_usage_chunk.py +++ b/studio/backend/tests/test_external_provider_usage_chunk.py @@ -183,7 +183,11 @@ def _usage_chunks(lines: list[str]) -> list[dict]: parsed = json.loads(payload) except json.JSONDecodeError: continue - if isinstance(parsed, dict) and "usage" in parsed and parsed.get("choices") == []: + if ( + isinstance(parsed, dict) + and "usage" in parsed + and parsed.get("choices") == [] + ): out.append(parsed["usage"]) return out @@ -239,7 +243,9 @@ def test_custom_provider_test_endpoint_probes_chat_completion(monkeypatch): from pathlib import Path module_path = Path(__file__).resolve().parents[1] / "routes" / "providers.py" - spec = importlib.util.spec_from_file_location("_providers_route_under_test", module_path) + spec = importlib.util.spec_from_file_location( + "_providers_route_under_test", module_path + ) assert spec is not None assert spec.loader is not None providers_route = importlib.util.module_from_spec(spec) @@ -289,7 +295,9 @@ def test_custom_provider_test_endpoint_requires_model_id(monkeypatch): from pathlib import Path module_path = Path(__file__).resolve().parents[1] / "routes" / "providers.py" - spec = importlib.util.spec_from_file_location("_providers_route_under_test", module_path) + spec = importlib.util.spec_from_file_location( + "_providers_route_under_test", module_path + ) assert spec is not None assert spec.loader is not None providers_route = importlib.util.module_from_spec(spec) @@ -375,9 +383,13 @@ def test_anthropic_stream_emits_usage_chunk_before_done(monkeypatch): # Usage chunk must come before [DONE]. data_lines = [ln for ln in lines if ln.startswith("data:")] - done_idx = next(i for i, ln in enumerate(data_lines) if ln.strip().endswith("[DONE]")) + done_idx = next( + i for i, ln in enumerate(data_lines) if ln.strip().endswith("[DONE]") + ) usage_idx = next( - i for i, ln in enumerate(data_lines) if '"usage":' in ln and '"choices": []' in ln + i + for i, ln in enumerate(data_lines) + if '"usage":' in ln and '"choices": []' in ln ) assert usage_idx < done_idx diff --git a/studio/backend/tests/test_frontend_resolution.py b/studio/backend/tests/test_frontend_resolution.py index c3e0524a30..f24d0b0c60 100644 --- a/studio/backend/tests/test_frontend_resolution.py +++ b/studio/backend/tests/test_frontend_resolution.py @@ -129,7 +129,13 @@ def test_resolver_falls_back_to_windows_layout_site_packages(tmp_path, monkeypat alongside the POSIX `lib/python*/site-packages`.""" studio_home = tmp_path / "studio_home" sp_dist = ( - studio_home / "unsloth_studio" / "Lib" / "site-packages" / "studio" / "frontend" / "dist" + studio_home + / "unsloth_studio" + / "Lib" + / "site-packages" + / "studio" + / "frontend" + / "dist" ) sp_dist.mkdir(parents = True) (sp_dist / "index.html").write_text("", encoding = "utf-8") diff --git a/studio/backend/tests/test_gemini_provider.py b/studio/backend/tests/test_gemini_provider.py index 85ceb04d27..89f091f399 100644 --- a/studio/backend/tests/test_gemini_provider.py +++ b/studio/backend/tests/test_gemini_provider.py @@ -542,7 +542,9 @@ def test_finish_reason_swaps_to_tool_calls_when_function_call_emitted(monkeypatc { "content": { "role": "model", - "parts": [{"functionCall": {"name": "lookup", "args": {"k": "v"}}}], + "parts": [ + {"functionCall": {"name": "lookup", "args": {"k": "v"}}} + ], }, "finishReason": "STOP", } @@ -626,7 +628,9 @@ def test_thought_signature_emitted_in_tool_call_delta(monkeypatch): deltas = [ tc for c in chunks - for tc in (c.get("choices", [{}])[0].get("delta", {}) or {}).get("tool_calls", []) + for tc in (c.get("choices", [{}])[0].get("delta", {}) or {}).get( + "tool_calls", [] + ) ] assert deltas, chunks sig = deltas[0].get("extra_content", {}).get("google", {}).get("thought_signature") @@ -680,7 +684,10 @@ def test_image_generation_tool_on_image_model_drops_text_tools(monkeypatch): ], ) assert "tools" not in captured["body"], captured["body"] - assert captured["body"]["generationConfig"].get("responseModalities") == ["TEXT", "IMAGE"] + assert captured["body"]["generationConfig"].get("responseModalities") == [ + "TEXT", + "IMAGE", + ] def test_prompt_feedback_block_reason_surfaces_as_error(monkeypatch): @@ -694,7 +701,9 @@ def test_prompt_feedback_block_reason_surfaces_as_error(monkeypatch): chunks = _parse_chunks(_collect(monkeypatch, sse)) error_chunks = [c for c in chunks if "error" in c] assert error_chunks, chunks - assert "SAFETY" in (error_chunks[0].get("error", {}).get("message") or ""), error_chunks + assert "SAFETY" in ( + error_chunks[0].get("error", {}).get("message") or "" + ), error_chunks def test_usage_chunk_includes_thoughts_tokens(monkeypatch): @@ -782,7 +791,10 @@ def test_image_model_sets_response_modalities(monkeypatch): model = "gemini-2.5-flash-image", enabled_tools = ["image_generation"], ) - assert captured["body"]["generationConfig"]["responseModalities"] == ["TEXT", "IMAGE"] + assert captured["body"]["generationConfig"]["responseModalities"] == [ + "TEXT", + "IMAGE", + ] def test_image_generation_tool_sets_response_modalities_on_image_model(monkeypatch): @@ -795,7 +807,10 @@ def test_image_generation_tool_sets_response_modalities_on_image_model(monkeypat model = "gemini-2.5-flash-image", enabled_tools = ["image_generation"], ) - assert captured["body"]["generationConfig"]["responseModalities"] == ["TEXT", "IMAGE"] + assert captured["body"]["generationConfig"]["responseModalities"] == [ + "TEXT", + "IMAGE", + ] def test_image_response_emits_image_b64_tool_event(monkeypatch): @@ -983,7 +998,9 @@ def test_parallel_function_calls_get_distinct_tool_call_indices(monkeypatch): ) ] assert len(tool_call_chunks) == 2, tool_call_chunks - indices = [c["choices"][0]["delta"]["tool_calls"][0]["index"] for c in tool_call_chunks] + indices = [ + c["choices"][0]["delta"]["tool_calls"][0]["index"] for c in tool_call_chunks + ] assert indices == [0, 1], indices @@ -1030,7 +1047,10 @@ def test_function_call_ids_forwarded_into_gemini_function_call_part(monkeypatch) call_ids = [p["functionCall"]["id"] for p in assistant_parts if "functionCall" in p] assert call_ids == ["call_alpha", "call_beta"], assistant_parts response_ids = [ - p["functionResponse"]["id"] for c in contents for p in c["parts"] if "functionResponse" in p + p["functionResponse"]["id"] + for c in contents + for p in c["parts"] + if "functionResponse" in p ] assert response_ids == ["call_alpha", "call_beta"], contents @@ -1108,7 +1128,9 @@ def test_code_execution_parts_translate_to_code_execution_tool_events(monkeypatc if e.get("type") == "tool_start" and e.get("tool_name") == "code_execution" ] code_ends = [ - e for e in tool_events if e.get("type") == "tool_end" and "4" in str(e.get("result", "")) + e + for e in tool_events + if e.get("type") == "tool_end" and "4" in str(e.get("result", "")) ] assert len(code_starts) == 1, tool_events assert code_starts[0]["arguments"]["code"] == "print(2+2)" @@ -1645,7 +1667,9 @@ def test_safe_fetch_image_rejects_resolved_private_host(monkeypatch): monkeypatch.setattr(socket, "getaddrinfo", fake_getaddrinfo) res = asyncio.new_event_loop().run_until_complete( - ep_mod._safe_fetch_image_for_gemini("https://internal.example/x.png", "image/png") + ep_mod._safe_fetch_image_for_gemini( + "https://internal.example/x.png", "image/png" + ) ) assert res is None @@ -1719,7 +1743,9 @@ def test_youtube_and_files_api_uris_stay_as_file_data(monkeypatch): parts = captured["body"]["contents"][-1]["parts"] file_uris = [p["fileData"]["fileUri"] for p in parts if "fileData" in p] assert "https://www.youtube.com/watch?v=abc123" in file_uris, parts - assert "https://generativelanguage.googleapis.com/v1beta/files/abc" in file_uris, parts + assert ( + "https://generativelanguage.googleapis.com/v1beta/files/abc" in file_uris + ), parts def test_tool_use_prompt_tokens_added_to_input_tokens(monkeypatch): @@ -1912,7 +1938,9 @@ def test_inline_image_tool_end_carries_thought_signature(monkeypatch): ) chunks = _parse_chunks(lines) tool_events = [c["_toolEvent"] for c in chunks if "_toolEvent" in c] - image_ends = [e for e in tool_events if e.get("type") == "tool_end" and e.get("image_b64")] + image_ends = [ + e for e in tool_events if e.get("type") == "tool_end" and e.get("image_b64") + ] assert image_ends, tool_events assert image_ends[0]["google"]["thought_signature"] == "SIG-IMG" # Multi-turn image edit must replay the original inlineData part with its @@ -1978,7 +2006,9 @@ def test_code_execution_plot_attaches_inline_image_native_part(monkeypatch): chunks = _parse_chunks(lines) tool_events = [c["_toolEvent"] for c in chunks if "_toolEvent" in c] code_ends = [ - e for e in tool_events if e.get("type") == "tool_end" and e.get("tool_call_id") == "code_a" + e + for e in tool_events + if e.get("type") == "tool_end" and e.get("tool_call_id") == "code_a" ] # Two tool_end events on the same id: one for codeExecutionResult, one # merging in the inlineData plot. The plot one must carry the native @@ -2025,7 +2055,9 @@ def test_text_chunk_carries_thought_signature(monkeypatch): lines = _collect(monkeypatch, sse) chunks = _parse_chunks(lines) text_chunks = [ - c for c in chunks if c.get("choices") and c["choices"][0]["delta"].get("content") == "hello" + c + for c in chunks + if c.get("choices") and c["choices"][0]["delta"].get("content") == "hello" ] assert text_chunks, chunks extra = text_chunks[0]["choices"][0]["delta"].get("extra_content") @@ -2203,7 +2235,8 @@ def test_code_execution_tool_call_replays_native_executable_code(monkeypatch): assert "executableCode" in native_keys, parts assert "codeExecutionResult" in native_keys, parts assert not any( - "functionCall" in p and (p["functionCall"] or {}).get("name") == "code_execution" + "functionCall" in p + and (p["functionCall"] or {}).get("name") == "code_execution" for p in parts ), parts exec_part = next(p for p in parts if "executableCode" in p) @@ -2258,7 +2291,8 @@ def test_image_generation_tool_call_replays_native_inline_data(monkeypatch): assert inline_parts[0]["inlineData"]["data"] == pixel assert inline_parts[0].get("thoughtSignature") == "SIG-IMG", inline_parts assert not any( - "functionCall" in p and (p["functionCall"] or {}).get("name") == "image_generation" + "functionCall" in p + and (p["functionCall"] or {}).get("name") == "image_generation" for p in parts ), parts @@ -2325,7 +2359,11 @@ def test_function_declarations_strip_openai_only_schema_keys(monkeypatch): ) tools_arr = captured["body"].get("tools") or [] decls = next( - (t.get("functionDeclarations") for t in tools_arr if "functionDeclarations" in t), + ( + t.get("functionDeclarations") + for t in tools_arr + if "functionDeclarations" in t + ), None, ) assert decls is not None, captured["body"] @@ -2378,7 +2416,11 @@ def test_function_declarations_inline_local_refs_into_gemini_schema(monkeypatch) ) tools_arr = captured["body"].get("tools") or [] decls = next( - (t.get("functionDeclarations") for t in tools_arr if "functionDeclarations" in t), + ( + t.get("functionDeclarations") + for t in tools_arr + if "functionDeclarations" in t + ), None, ) assert decls is not None, captured["body"] @@ -2429,7 +2471,11 @@ def test_function_declarations_inline_local_refs_in_anyof_and_items(monkeypatch) ) tools_arr = captured["body"].get("tools") or [] decls = next( - (t.get("functionDeclarations") for t in tools_arr if "functionDeclarations" in t), + ( + t.get("functionDeclarations") + for t in tools_arr + if "functionDeclarations" in t + ), None, ) assert decls is not None @@ -2444,7 +2490,10 @@ def test_function_declarations_inline_local_refs_in_anyof_and_items(monkeypatch) extras = params["properties"]["extras"] assert extras.get("type") == "array" assert extras.get("items", {}).get("type") == "object" - assert extras.get("items", {}).get("properties", {}).get("zip", {}).get("type") == "string" + assert ( + extras.get("items", {}).get("properties", {}).get("zip", {}).get("type") + == "string" + ) def test_function_declarations_self_referential_schema_terminates(monkeypatch): @@ -2482,7 +2531,11 @@ def test_function_declarations_self_referential_schema_terminates(monkeypatch): ) tools_arr = captured["body"].get("tools") or [] decls = next( - (t.get("functionDeclarations") for t in tools_arr if "functionDeclarations" in t), + ( + t.get("functionDeclarations") + for t in tools_arr + if "functionDeclarations" in t + ), None, ) assert decls is not None @@ -2545,7 +2598,9 @@ def test_gemini_native_skips_orphan_function_response_for_dropped_builtin(monkey assert fr.get("name") != "web_search", contents -def test_gemini_native_skips_orphan_function_response_for_native_part_replay(monkeypatch): +def test_gemini_native_skips_orphan_function_response_for_native_part_replay( + monkeypatch, +): """Round 26: code_execution / image_generation tool_calls are replayed as Gemini-native executableCode / codeExecutionResult / inlineData parts. The matching role="tool" follow-up must NOT then be emitted as a @@ -2761,7 +2816,9 @@ def test_chat_message_extra_content_round_trips_through_validation(): base_url = "https://generativelanguage.googleapis.com/v1beta", ) assistant_out = built[1] - assert assistant_out["extra_content"] == {"google": {"thought_signature": "SIG-TEXT"}} + assert assistant_out["extra_content"] == { + "google": {"thought_signature": "SIG-TEXT"} + } # Non-Gemini providers must NOT receive extra_content; Google's # thought_signature is unknown to OpenAI / Mistral / etc. built_openai = _build_external_messages( @@ -2828,7 +2885,10 @@ def test_parallel_tool_results_group_into_one_user_block(monkeypatch): c for c in contents if c.get("role") == "user" - and all(isinstance(p, dict) and "functionResponse" in p for p in (c.get("parts") or [])) + and all( + isinstance(p, dict) and "functionResponse" in p + for p in (c.get("parts") or []) + ) ] assert len(tool_result_users) == 1, contents fr_parts = tool_result_users[0]["parts"] @@ -2886,7 +2946,9 @@ def test_image_picker_model_with_search_off_pill_strips_text_tools(monkeypatch): ) body = captured["body"] assert "tools" not in body, body.get("tools") - assert "thinkingConfig" not in body.get("generationConfig", {}), body["generationConfig"] + assert "thinkingConfig" not in body.get("generationConfig", {}), body[ + "generationConfig" + ] def test_image_models_drop_function_declarations(monkeypatch): @@ -2904,7 +2966,10 @@ def test_image_models_drop_function_declarations(monkeypatch): ], ) assert captured["body"].get("tools") is None - assert captured["body"]["generationConfig"]["responseModalities"] == ["TEXT", "IMAGE"] + assert captured["body"]["generationConfig"]["responseModalities"] == [ + "TEXT", + "IMAGE", + ] def test_safe_fetch_image_rejects_malformed_bracketed_url(): @@ -2970,14 +3035,22 @@ def test_safe_fetch_image_pins_validated_ip_no_hostname_in_request(monkeypatch): ) return _StubResp() - monkeypatch.setattr("urllib.request.build_opener", lambda *_args, **_kw: _StubOpener()) + monkeypatch.setattr( + "urllib.request.build_opener", lambda *_args, **_kw: _StubOpener() + ) - res = _drive(ep_mod._safe_fetch_image_for_gemini("https://cdn.example.com/x.png", "image/png")) + res = _drive( + ep_mod._safe_fetch_image_for_gemini( + "https://cdn.example.com/x.png", "image/png" + ) + ) assert res is not None assert res[0] == "image/png" # Outgoing URL must use the pinned IP literal, not the hostname. assert any("8.8.8.8" in r["url"] for r in captured["requests"]), captured - assert all("cdn.example.com" not in r["url"] for r in captured["requests"]), captured + assert all( + "cdn.example.com" not in r["url"] for r in captured["requests"] + ), captured # Host header still carries the original hostname for vhost/SNI. assert captured["requests"][0]["host_header"] == "cdn.example.com" @@ -3030,9 +3103,15 @@ def test_safe_fetch_image_redirect_to_private_host_rejected(monkeypatch): None, ) - monkeypatch.setattr("urllib.request.build_opener", lambda *_args, **_kw: _StubOpener()) + monkeypatch.setattr( + "urllib.request.build_opener", lambda *_args, **_kw: _StubOpener() + ) - res = _drive(ep_mod._safe_fetch_image_for_gemini("https://cdn.example.com/x.png", "image/png")) + res = _drive( + ep_mod._safe_fetch_image_for_gemini( + "https://cdn.example.com/x.png", "image/png" + ) + ) assert res is None @@ -3182,7 +3261,9 @@ def test_legacy_gemini3_pro_medium_coerced_to_high(monkeypatch): model = "gemini-3-pro-preview", reasoning_effort = "medium", ) - assert captured["body"]["generationConfig"]["thinkingConfig"] == {"thinkingLevel": "high"} + assert captured["body"]["generationConfig"]["thinkingConfig"] == { + "thinkingLevel": "high" + } def test_gemini_3_1_pro_medium_passes_through(monkeypatch): @@ -3193,7 +3274,9 @@ def test_gemini_3_1_pro_medium_passes_through(monkeypatch): model = "gemini-3.1-pro-preview", reasoning_effort = "medium", ) - assert captured["body"]["generationConfig"]["thinkingConfig"] == {"thinkingLevel": "medium"} + assert captured["body"]["generationConfig"]["thinkingConfig"] == { + "thinkingLevel": "medium" + } def test_tool_calls_extra_content_stripped_for_non_native_gemini(): @@ -3289,7 +3372,9 @@ def test_user_function_named_with_server_tool_arg_not_dropped(monkeypatch): "type": "function", "function": { "name": "user_function", - "arguments": json.dumps({"_server_tool": True, "q": "x"}), + "arguments": json.dumps( + {"_server_tool": True, "q": "x"} + ), }, } ], @@ -3354,7 +3439,9 @@ def test_builtin_named_with_server_tool_marker_dropped(monkeypatch): "type": "function", "function": { "name": "web_search", - "arguments": json.dumps({"_server_tool": True, "query": "x"}), + "arguments": json.dumps( + {"_server_tool": True, "query": "x"} + ), }, } ], @@ -3442,7 +3529,9 @@ def test_schema_anyof_multitype_with_null_keeps_anyof_and_nullable(monkeypatch): assert either.get("nullable") is True inner = either.get("anyOf") assert isinstance(inner, list) and len(inner) == 2, either - assert all(not (isinstance(b, dict) and b.get("type") == "null") for b in inner), inner + assert all( + not (isinstance(b, dict) and b.get("type") == "null") for b in inner + ), inner def test_safe_fetch_image_redirect_malformed_url_no_crash(monkeypatch): @@ -3483,16 +3572,26 @@ def test_safe_fetch_image_redirect_malformed_url_no_crash(monkeypatch): None, ) - monkeypatch.setattr("urllib.request.build_opener", lambda *_args, **_kw: _StubOpener()) + monkeypatch.setattr( + "urllib.request.build_opener", lambda *_args, **_kw: _StubOpener() + ) - res = _drive(ep_mod._safe_fetch_image_for_gemini("https://cdn.example.com/x.png", "image/png")) + res = _drive( + ep_mod._safe_fetch_image_for_gemini( + "https://cdn.example.com/x.png", "image/png" + ) + ) assert res is None def test_safe_fetch_image_malformed_port_no_crash(): """Round 18: a URL with a non-numeric port (`https://h:bad/x.png`) must not raise; urlparse's port property lazily ValueErrors.""" - res = _drive(ep_mod._safe_fetch_image_for_gemini("https://example.com:bad/x.png", "image/png")) + res = _drive( + ep_mod._safe_fetch_image_for_gemini( + "https://example.com:bad/x.png", "image/png" + ) + ) assert res is None @@ -3541,10 +3640,14 @@ def test_safe_fetch_image_missing_content_type_uses_fallback(monkeypatch): ): return _StubResp() - monkeypatch.setattr("urllib.request.build_opener", lambda *_args, **_kw: _StubOpener()) + monkeypatch.setattr( + "urllib.request.build_opener", lambda *_args, **_kw: _StubOpener() + ) res = _drive( - ep_mod._safe_fetch_image_for_gemini("https://cdn.example.com/cat.png", "image/png") + ep_mod._safe_fetch_image_for_gemini( + "https://cdn.example.com/cat.png", "image/png" + ) ) assert res is not None assert res[0] == "image/png" @@ -3625,7 +3728,9 @@ def test_anthropic_translates_openai_tool_calls_into_tool_use_blocks(monkeypatch tool_results: list[dict] = [] for m in msgs: if m.get("role") == "user" and isinstance(m.get("content"), list): - tool_results.extend(b for b in m["content"] if b.get("type") == "tool_result") + tool_results.extend( + b for b in m["content"] if b.get("type") == "tool_result" + ) assert any( tr.get("tool_use_id") == "call_a" and tr.get("content") == "result_text" for tr in tool_results @@ -4105,7 +4210,9 @@ def test_orphan_function_call_output_dropped_when_call_skipped(monkeypatch): "type": "function", "function": { "name": "web_search", - "arguments": json.dumps({"_server_tool": True, "query": "x"}), + "arguments": json.dumps( + {"_server_tool": True, "query": "x"} + ), }, } ], @@ -4166,7 +4273,9 @@ def test_schema_multitype_union_with_null_preserves_anyof(monkeypatch): assert either.get("nullable") is True inner = either.get("anyOf") assert isinstance(inner, list) and len(inner) == 2, either - types = sorted(b.get("type") for b in inner if isinstance(b, dict) and b.get("type")) + types = sorted( + b.get("type") for b in inner if isinstance(b, dict) and b.get("type") + ) assert types == ["integer", "string"], inner @@ -4368,7 +4477,9 @@ def test_openrouter_no_synthetic_web_search_event_on_tool_choice_none(monkeypatc _drive(run()) # No synthetic web_search tool_start / tool_end emitted. - assert all(e.get("tool_name") != "web_search" for e in captured_events), captured_events + assert all( + e.get("tool_name") != "web_search" for e in captured_events + ), captured_events def test_anthropic_role_tool_list_content_translates_to_tool_result(monkeypatch): @@ -4435,7 +4546,9 @@ def test_anthropic_role_tool_list_content_translates_to_tool_result(monkeypatch) tool_results: list[dict] = [] for m in msgs: if m.get("role") == "user" and isinstance(m.get("content"), list): - tool_results.extend(b for b in m["content"] if b.get("type") == "tool_result") + tool_results.extend( + b for b in m["content"] if b.get("type") == "tool_result" + ) assert any( tr.get("tool_use_id") == "call_a" and tr.get("content") == "result_text" for tr in tool_results @@ -4561,7 +4674,13 @@ def test_openai_responses_assistant_text_serialized_before_function_call(monkeyp # function_call (get_weather) # function_call_output (sunny) # user ("thanks") - assert types == ["user", "assistant", "function_call", "function_call_output", "user"], items + assert types == [ + "user", + "assistant", + "function_call", + "function_call_output", + "user", + ], items def test_gemini_tool_choice_none_disables_image_generation(monkeypatch): @@ -4631,7 +4750,9 @@ def test_gemini_forced_function_tool_choice_drops_image_generation(monkeypatch): assert body["generationConfig"].get("responseModalities") == ["TEXT"], body -def test_gemini_code_execution_native_part_list_replays_per_part_signatures(monkeypatch): +def test_gemini_code_execution_native_part_list_replays_per_part_signatures( + monkeypatch, +): """Round 21: merged code-execution history must replay per-part `thoughtSignature`s, not fan one top-level signature across every native subpart. Gemini 3 strict validators reject a signature on the wrong @@ -4847,7 +4968,9 @@ def test_safe_fetch_image_threads_per_request_byte_budget(monkeypatch): ): return _StubResp() - monkeypatch.setattr("urllib.request.build_opener", lambda *_args, **_kw: _StubOpener()) + monkeypatch.setattr( + "urllib.request.build_opener", lambda *_args, **_kw: _StubOpener() + ) res = _drive( ep_mod._safe_fetch_image_for_gemini( @@ -4870,7 +4993,9 @@ def test_openai_chat_delta_type_includes_tool_calls_and_extra_content(): import os here = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) - types_path = os.path.join(here, "frontend", "src", "features", "chat", "types", "api.ts") + types_path = os.path.join( + here, "frontend", "src", "features", "chat", "types", "api.ts" + ) with open(types_path, "r", encoding = "utf-8") as f: src = f.read() assert "tool_calls?: OpenAIToolCallPart[]" in src, src[:200] @@ -5076,7 +5201,9 @@ def test_openai_responses_forced_function_tool_choice_drops_hosted_tools(monkeyp assert not (hosted_seen & hosted_types), body # The user function declaration must still be present so the pin has a # target. - user_function_seen = any(isinstance(t, dict) and t.get("type") == "function" for t in tools) + user_function_seen = any( + isinstance(t, dict) and t.get("type") == "function" for t in tools + ) assert user_function_seen, body # And the forced-function tool_choice must be forwarded in Responses shape: # `{type:"function", name:"..."}`. @@ -5300,7 +5427,9 @@ def test_strip_provider_synthetic_tool_history_drops_empty_assistant(): assert roles == ["user", "user"], out -def test_openrouter_no_synthetic_web_search_event_on_forced_function_tool_choice(monkeypatch): +def test_openrouter_no_synthetic_web_search_event_on_forced_function_tool_choice( + monkeypatch, +): """Round 22 sibling of the round-20 `tool_choice='none'` test: when the caller forces a specific function via `tool_choice={"type":"function", ...}` AND passes `enabled_tools=["web_search"]`, the OpenRouter path must NOT @@ -5311,7 +5440,10 @@ def test_openrouter_no_synthetic_web_search_event_on_forced_function_tool_choice def handler(request: httpx.Request) -> httpx.Response: return httpx.Response( 200, - content = (b'data: {"choices":[{"delta":{"content":"ok"}}]}\n\n' b"data: [DONE]\n\n"), + content = ( + b'data: {"choices":[{"delta":{"content":"ok"}}]}\n\n' + b"data: [DONE]\n\n" + ), headers = {"content-type": "text/event-stream"}, ) diff --git a/studio/backend/tests/test_gemma4_chat_template_override.py b/studio/backend/tests/test_gemma4_chat_template_override.py index f726741aa5..e167ffcbc7 100644 --- a/studio/backend/tests/test_gemma4_chat_template_override.py +++ b/studio/backend/tests/test_gemma4_chat_template_override.py @@ -32,7 +32,9 @@ chat_templates = importlib.util.module_from_spec(_ct_spec) _ct_spec.loader.exec_module(chat_templates) is_unsloth_gemma4_gguf = chat_templates.is_unsloth_gemma4_gguf -resolve_effective_chat_template_override = chat_templates.resolve_effective_chat_template_override +resolve_effective_chat_template_override = ( + chat_templates.resolve_effective_chat_template_override +) load_bundled_chat_template = chat_templates.load_bundled_chat_template is_unsloth_gemma4_edge_gguf = chat_templates.is_unsloth_gemma4_edge_gguf @@ -137,7 +139,9 @@ def test_is_unsloth_gemma4_edge_gguf(model_id, expected_edge): def test_resolver_returns_edge_template_for_e2b_e4b(): for mid in ("unsloth/gemma-4-E2B-it-GGUF", "unsloth/gemma-4-E4B-it-GGUF"): - out = resolve_effective_chat_template_override(model_identifier = mid, user_override = None) + out = resolve_effective_chat_template_override( + model_identifier = mid, user_override = None + ) assert out == EDGE assert out != BUNDLED @@ -165,7 +169,9 @@ def test_resolver_returns_standard_template_for_larger_models(): "unsloth/gemma-4-26B-A4B-it-GGUF", "unsloth/gemma-4-31B-it-GGUF", ): - out = resolve_effective_chat_template_override(model_identifier = mid, user_override = None) + out = resolve_effective_chat_template_override( + model_identifier = mid, user_override = None + ) assert out == BUNDLED @@ -268,7 +274,9 @@ def _convo_with_prior_tool_reasoning(): { "role": "assistant", "reasoning_content": "SECRET_THOUGHT", - "tool_calls": [{"id": "c1", "function": {"name": "f", "arguments": {"x": 1}}}], + "tool_calls": [ + {"id": "c1", "function": {"name": "f", "arguments": {"x": 1}}} + ], }, {"role": "tool", "tool_call_id": "c1", "content": "42"}, {"role": "user", "content": "q2"}, @@ -281,12 +289,18 @@ def test_preserve_thinking_off_omits_prior_reasoning(): def test_preserve_thinking_on_keeps_prior_reasoning(): - assert "SECRET_THOUGHT" in _render(_convo_with_prior_tool_reasoning(), preserve_thinking = True) + assert "SECRET_THOUGHT" in _render( + _convo_with_prior_tool_reasoning(), preserve_thinking = True + ) def test_enable_thinking_gates_think_token(): - assert "<|think|>" in _render([{"role": "user", "content": "hi"}], enable_thinking = True) - assert "<|think|>" not in _render([{"role": "user", "content": "hi"}], enable_thinking = False) + assert "<|think|>" in _render( + [{"role": "user", "content": "hi"}], enable_thinking = True + ) + assert "<|think|>" not in _render( + [{"role": "user", "content": "hi"}], enable_thinking = False + ) # ── Reload dedup interaction (why the route resolves the effective override) ── @@ -334,9 +348,14 @@ def test_already_in_target_state_consistent_with_bundled_override(): is_vision = False, ) # Effective (resolved bundled) override -> already loaded, no reload. - assert backend._already_in_target_state(chat_template_override = BUNDLED, **common) is True + assert ( + backend._already_in_target_state(chat_template_override = BUNDLED, **common) + is True + ) # Raw None (unresolved) -> false match, would force a needless reload. - assert backend._already_in_target_state(chat_template_override = None, **common) is False + assert ( + backend._already_in_target_state(chat_template_override = None, **common) is False + ) def _import_backend(): diff --git a/studio/backend/tests/test_gguf_completion_usage.py b/studio/backend/tests/test_gguf_completion_usage.py index d1e05f3e0e..d7093979e9 100644 --- a/studio/backend/tests/test_gguf_completion_usage.py +++ b/studio/backend/tests/test_gguf_completion_usage.py @@ -30,8 +30,12 @@ class _GgufBackend: def _request_completion(monkeypatch, usage): - monkeypatch.setattr(inference_route, "get_llama_cpp_backend", lambda: _GgufBackend(usage)) - monkeypatch.setattr(inference_route, "_effective_enable_tools", lambda payload: False) + monkeypatch.setattr( + inference_route, "get_llama_cpp_backend", lambda: _GgufBackend(usage) + ) + monkeypatch.setattr( + inference_route, "_effective_enable_tools", lambda payload: False + ) app = FastAPI() app.include_router(inference_route.router) diff --git a/studio/backend/tests/test_gguf_metadata.py b/studio/backend/tests/test_gguf_metadata.py index d3d4387720..586597fb2b 100644 --- a/studio/backend/tests/test_gguf_metadata.py +++ b/studio/backend/tests/test_gguf_metadata.py @@ -35,11 +35,17 @@ def _enc_kv_string(key: str, value: str) -> bytes: def _enc_kv_uint32(key: str, value: int) -> bytes: - return _enc_string(key) + struct.pack(" bytes: - return _enc_string(key) + struct.pack(" bytes: @@ -64,7 +70,10 @@ def _write_synthetic_gguf( extra_string_arrays = extra_string_arrays or {} extra_bools = extra_bools or {} kv_count = ( - len(general_strings) + len(extra_uint32) + len(extra_string_arrays) + len(extra_bools) + len(general_strings) + + len(extra_uint32) + + len(extra_string_arrays) + + len(extra_bools) ) body = b"" for k, v in general_strings.items(): @@ -117,7 +126,10 @@ def test_extracts_general_string_fields(tmp_path: Path): assert meta is not None assert meta["general.architecture"] == "qwen2vl" assert meta["general.basename"] == "Qwen3.5" - assert meta["general.base_model.0.repo_url"] == "https://huggingface.co/Qwen/Qwen3.5-9B" + assert ( + meta["general.base_model.0.repo_url"] + == "https://huggingface.co/Qwen/Qwen3.5-9B" + ) def test_skips_unrelated_fields_without_breaking(tmp_path: Path): diff --git a/studio/backend/tests/test_gpu_selection.py b/studio/backend/tests/test_gpu_selection.py index cb26330ed0..590627e0ad 100644 --- a/studio/backend/tests/test_gpu_selection.py +++ b/studio/backend/tests/test_gpu_selection.py @@ -64,7 +64,9 @@ class TestResolveRequestedGpuIds(_GpuCacheResetMixin, unittest.TestCase): def test_parent_visibility_uses_empty_numeric_ids_for_uuid_masks(self): with ( - patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": "GPU-aaa,GPU-bbb"}, clear = True), + patch.dict( + os.environ, {"CUDA_VISIBLE_DEVICES": "GPU-aaa,GPU-bbb"}, clear = True + ), patch("utils.hardware.hardware.get_physical_gpu_count", return_value = 8), ): self.assertEqual(get_parent_visible_gpu_ids(), []) @@ -94,7 +96,9 @@ class TestResolveRequestedGpuIds(_GpuCacheResetMixin, unittest.TestCase): def test_explicit_ids_are_rejected_for_uuid_parent_visibility(self): with ( - patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": "GPU-aaa,GPU-bbb"}, clear = True), + patch.dict( + os.environ, {"CUDA_VISIBLE_DEVICES": "GPU-aaa,GPU-bbb"}, clear = True + ), patch("utils.hardware.hardware.get_physical_gpu_count", return_value = 8), ): with self.assertRaisesRegex( @@ -200,9 +204,13 @@ class TestVisibleGpuUtilization(_GpuCacheResetMixin, unittest.TestCase): }, ] with ( - patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": "GPU-aaa,GPU-bbb"}, clear = True), + patch.dict( + os.environ, {"CUDA_VISIBLE_DEVICES": "GPU-aaa,GPU-bbb"}, clear = True + ), patch("utils.hardware.hardware.get_device", return_value = DeviceType.CUDA), - patch("utils.hardware.hardware._torch_get_physical_gpu_count", return_value = 2), + patch( + "utils.hardware.hardware._torch_get_physical_gpu_count", return_value = 2 + ), patch( "utils.hardware.hardware._torch_get_per_device_info", return_value = fake_torch_devices, @@ -246,7 +254,9 @@ class TestGpuAutoSelection(_GpuCacheResetMixin, unittest.TestCase): def test_get_device_map_uses_all_inherited_visible_gpus_for_uuid_masks(self): with ( - patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": "GPU-aaa,GPU-bbb"}, clear = True), + patch.dict( + os.environ, {"CUDA_VISIBLE_DEVICES": "GPU-aaa,GPU-bbb"}, clear = True + ), patch("utils.hardware.hardware.get_device", return_value = DeviceType.CUDA), ): self.assertEqual(get_device_map(None), "balanced") @@ -366,7 +376,9 @@ class TestGpuAutoSelection(_GpuCacheResetMixin, unittest.TestCase): return_value = 1234, ), ): - model_size_bytes, source = _hw_module.estimate_fp16_model_size_bytes("unsloth/test") + model_size_bytes, source = _hw_module.estimate_fp16_model_size_bytes( + "unsloth/test" + ) self.assertEqual(model_size_bytes, 1234) self.assertEqual(source, "vllm_utils") @@ -463,7 +475,9 @@ class TestGpuAutoSelection(_GpuCacheResetMixin, unittest.TestCase): def test_prepare_gpu_selection_preserves_uuid_parent_visibility_in_auto_mode(self): with ( - patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": "GPU-aaa,GPU-bbb"}, clear = True), + patch.dict( + os.environ, {"CUDA_VISIBLE_DEVICES": "GPU-aaa,GPU-bbb"}, clear = True + ), patch( "utils.hardware.hardware.estimate_required_model_memory_gb", return_value = ( @@ -510,7 +524,9 @@ class TestPreSpawnGpuResolution(_GpuCacheResetMixin, unittest.TestCase): patch( "core.training.training._CTX.Process", return_value = DummyProcess() ) as mock_process, - patch("core.training.training.threading.Thread", return_value = DummyThread()), + patch( + "core.training.training.threading.Thread", return_value = DummyThread() + ), ): backend.start_training( job_id = "test-job-1", @@ -551,7 +567,9 @@ class TestPreSpawnGpuResolution(_GpuCacheResetMixin, unittest.TestCase): patch( "core.training.training._CTX.Process", return_value = DummyProcess() ) as mock_process, - patch("core.training.training.threading.Thread", return_value = DummyThread()), + patch( + "core.training.training.threading.Thread", return_value = DummyThread() + ), ): backend.start_training( job_id = "test-job-2", @@ -581,7 +599,9 @@ class TestPreSpawnGpuResolution(_GpuCacheResetMixin, unittest.TestCase): dummy_queue = object() with ( - patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": "GPU-aaa,GPU-bbb"}, clear = True), + patch.dict( + os.environ, {"CUDA_VISIBLE_DEVICES": "GPU-aaa,GPU-bbb"}, clear = True + ), patch( "core.training.training._CTX.Queue", side_effect = [dummy_queue, dummy_queue], @@ -589,7 +609,9 @@ class TestPreSpawnGpuResolution(_GpuCacheResetMixin, unittest.TestCase): patch( "core.training.training._CTX.Process", return_value = DummyProcess() ) as mock_process, - patch("core.training.training.threading.Thread", return_value = DummyThread()), + patch( + "core.training.training.threading.Thread", return_value = DummyThread() + ), patch( "utils.hardware.hardware.estimate_required_model_memory_gb", return_value = ( @@ -607,7 +629,9 @@ class TestPreSpawnGpuResolution(_GpuCacheResetMixin, unittest.TestCase): config = mock_process.call_args.kwargs["kwargs"]["config"] self.assertIsNone(config["resolved_gpu_ids"]) - self.assertEqual(config["gpu_selection"]["selection_mode"], "inherit_parent_visible") + self.assertEqual( + config["gpu_selection"]["selection_mode"], "inherit_parent_visible" + ) def test_inference_orchestrator_resolves_explicit_gpu_ids_before_spawn(self): class DummyThread: @@ -635,7 +659,9 @@ class TestPreSpawnGpuResolution(_GpuCacheResetMixin, unittest.TestCase): "_wait_response", return_value = {"success": True, "model_info": {}}, ), - patch("utils.transformers_version.needs_transformers_5", return_value = False), + patch( + "utils.transformers_version.needs_transformers_5", return_value = False + ), ): self.assertTrue(orchestrator.load_model(config = config, gpu_ids = [1])) @@ -670,7 +696,9 @@ class TestPreSpawnGpuResolution(_GpuCacheResetMixin, unittest.TestCase): "_wait_response", return_value = {"success": True, "model_info": {}}, ), - patch("utils.transformers_version.needs_transformers_5", return_value = False), + patch( + "utils.transformers_version.needs_transformers_5", return_value = False + ), ): self.assertTrue(orchestrator.load_model(config = config, gpu_ids = None)) @@ -752,7 +780,9 @@ class TestRouteErrors(unittest.TestCase): raise ValueError("Invalid gpu_ids [99]") with ( - patch.object(training_route, "get_training_backend", return_value = DummyBackend()), + patch.object( + training_route, "get_training_backend", return_value = DummyBackend() + ), patch( "core.inference.get_inference_backend", return_value = SimpleNamespace(active_model_name = None), @@ -763,7 +793,9 @@ class TestRouteErrors(unittest.TestCase): ), ): with self.assertRaises(HTTPException) as exc_info: - asyncio.run(training_route.start_training(request, current_subject = "test-user")) + asyncio.run( + training_route.start_training(request, current_subject = "test-user") + ) self.assertEqual(exc_info.exception.status_code, 400) self.assertIn("gpu_ids [99]", exc_info.exception.detail) @@ -792,7 +824,9 @@ class TestRouteErrors(unittest.TestCase): ) with ( - patch.object(training_route, "get_training_backend", return_value = DummyBackend()), + patch.object( + training_route, "get_training_backend", return_value = DummyBackend() + ), patch( "core.inference.get_inference_backend", return_value = SimpleNamespace(active_model_name = None), @@ -803,7 +837,9 @@ class TestRouteErrors(unittest.TestCase): ), ): with self.assertRaises(HTTPException) as exc_info: - asyncio.run(training_route.start_training(request, current_subject = "test-user")) + asyncio.run( + training_route.start_training(request, current_subject = "test-user") + ) self.assertEqual(exc_info.exception.status_code, 400) self.assertIn("UUID/MIG", exc_info.exception.detail) @@ -943,7 +979,9 @@ class TestRaiseIfOffloaded(unittest.TestCase): def test_cpu_offload_raises(self): from utils.hardware import raise_if_offloaded - model = SimpleNamespace(hf_device_map = {"model.layers.0": 0, "model.layers.1": "cpu"}) + model = SimpleNamespace( + hf_device_map = {"model.layers.0": 0, "model.layers.1": "cpu"} + ) with self.assertRaisesRegex(ValueError, "offloaded"): raise_if_offloaded(model, "balanced", "Test") diff --git a/studio/backend/tests/test_gpu_selection_sandbox.py b/studio/backend/tests/test_gpu_selection_sandbox.py index 733933271b..1bd01ef16e 100644 --- a/studio/backend/tests/test_gpu_selection_sandbox.py +++ b/studio/backend/tests/test_gpu_selection_sandbox.py @@ -382,7 +382,9 @@ class TestResolveRequestedGpuIds(unittest.TestCase): def test_uuid_env_var_rejects_explicit_ids(self): from utils.hardware.hardware import resolve_requested_gpu_ids with ( - patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": "GPU-abc,GPU-def"}, clear = False), + patch.dict( + os.environ, {"CUDA_VISIBLE_DEVICES": "GPU-abc,GPU-def"}, clear = False + ), patch("utils.hardware.hardware.get_physical_gpu_count", return_value = 8), ): with self.assertRaises(ValueError): diff --git a/studio/backend/tests/test_host_defaults.py b/studio/backend/tests/test_host_defaults.py index 5c7129bc65..fc9dfc5dbf 100644 --- a/studio/backend/tests/test_host_defaults.py +++ b/studio/backend/tests/test_host_defaults.py @@ -20,7 +20,10 @@ def _parse_function_param_defaults(source: str, func_name: str) -> dict: """ tree = ast.parse(source) for node in ast.walk(tree): - if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)) and node.name == func_name: + if ( + isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)) + and node.name == func_name + ): result = {} all_args = node.args.args defaults = node.args.defaults @@ -66,7 +69,9 @@ def test_run_server_default_host_is_loopback(): """ source = _RUN_PY.read_text() defaults = _parse_function_param_defaults(source, "run_server") - assert "host" in defaults, "run_server() must have a 'host' parameter with a default" + assert ( + "host" in defaults + ), "run_server() must have a 'host' parameter with a default" host_default = defaults["host"] assert host_default == "127.0.0.1", ( f"run_server() host default must be '127.0.0.1' (loopback) " @@ -83,7 +88,9 @@ def test_argparse_default_host_is_loopback(): """ source = _RUN_PY.read_text() host_default = _parse_argparse_add_argument_default(source, "--host") - assert host_default is not None, "Could not find add_argument('--host', ...) in run.py" + assert ( + host_default is not None + ), "Could not find add_argument('--host', ...) in run.py" assert ( host_default == "127.0.0.1" ), f"run.py argparse --host default must be '127.0.0.1', got '{host_default}'" diff --git a/studio/backend/tests/test_index_bootstrap_origin.py b/studio/backend/tests/test_index_bootstrap_origin.py index b9d7f58867..962d3f6cb0 100644 --- a/studio/backend/tests/test_index_bootstrap_origin.py +++ b/studio/backend/tests/test_index_bootstrap_origin.py @@ -65,7 +65,9 @@ def test_is_same_origin_request_https_default_port_stripped_on_origin(): """RFC 6454 strips default ports on Origin; canonicalise both sides so this stays same-origin.""" from main import _is_same_origin_request - req = _build_request("example.com:443", origin = "https://example.com", scheme = "https") + req = _build_request( + "example.com:443", origin = "https://example.com", scheme = "https" + ) assert _is_same_origin_request(req) is True @@ -79,7 +81,9 @@ def test_is_same_origin_request_default_port_present_on_origin(): """Mirror case: Origin carries the default port, netloc doesn't. Same-origin.""" from main import _is_same_origin_request - req = _build_request("example.com", origin = "https://example.com:443", scheme = "https") + req = _build_request( + "example.com", origin = "https://example.com:443", scheme = "https" + ) assert _is_same_origin_request(req) is True @@ -127,5 +131,7 @@ def test_is_same_origin_request_explicit_non_default_port_still_mismatch(): """Canonicalisation does NOT collapse non-default ports to default.""" from main import _is_same_origin_request - req = _build_request("example.com", origin = "https://example.com:9999", scheme = "https") + req = _build_request( + "example.com", origin = "https://example.com:9999", scheme = "https" + ) assert _is_same_origin_request(req) is False diff --git a/studio/backend/tests/test_index_bootstrap_origin_extra.py b/studio/backend/tests/test_index_bootstrap_origin_extra.py index feda88c14c..b00b51fb74 100644 --- a/studio/backend/tests/test_index_bootstrap_origin_extra.py +++ b/studio/backend/tests/test_index_bootstrap_origin_extra.py @@ -91,7 +91,9 @@ def test_is_same_origin_request_data_url_origin_is_cross_origin(): """``data:`` URLs are opaque origins (HTML living standard); no host, never same-origin.""" from main import _is_same_origin_request - req = _build_request("127.0.0.1:8902", origin = "data:text/html,") + req = _build_request( + "127.0.0.1:8902", origin = "data:text/html," + ) assert _is_same_origin_request(req) is False diff --git a/studio/backend/tests/test_inference_model_validation.py b/studio/backend/tests/test_inference_model_validation.py index 2427ad35fa..56cc6d215e 100644 --- a/studio/backend/tests/test_inference_model_validation.py +++ b/studio/backend/tests/test_inference_model_validation.py @@ -201,7 +201,10 @@ def test_walkback_skips_explicitly_consumed_tool_call_id(): {"role": "tool", "content": "second result"}, ] ) - assert [m.tool_call_id for m in req.messages if m.role == "tool"] == ["call_a", "call_b"] + assert [m.tool_call_id for m in req.messages if m.role == "tool"] == [ + "call_a", + "call_b", + ] def test_walkback_handles_malformed_function_string(): diff --git a/studio/backend/tests/test_inference_orchestrator_crash_message.py b/studio/backend/tests/test_inference_orchestrator_crash_message.py index be1a673e62..b75e049280 100644 --- a/studio/backend/tests/test_inference_orchestrator_crash_message.py +++ b/studio/backend/tests/test_inference_orchestrator_crash_message.py @@ -25,7 +25,9 @@ def test_subprocess_crash_message_includes_signal_and_oom_hint(): msg = orchestrator._subprocess_crash_message("wait") - assert msg.startswith("The inference worker stopped unexpectedly while loading the model.") + assert msg.startswith( + "The inference worker stopped unexpectedly while loading the model." + ) assert "memory pressure" in msg assert "smaller model" in msg assert "Details:" in msg diff --git a/studio/backend/tests/test_install_resolve_prebuilt.py b/studio/backend/tests/test_install_resolve_prebuilt.py index ede5629664..ca15b07391 100644 --- a/studio/backend/tests/test_install_resolve_prebuilt.py +++ b/studio/backend/tests/test_install_resolve_prebuilt.py @@ -60,24 +60,38 @@ def test_published_repo_for_host(): # CPU-only Linux (x64 and arm64) -> ggml-org upstream. assert ilp.published_repo_for_host(_host(is_linux = True, is_x86_64 = True)) == UPSTREAM assert ( - ilp.published_repo_for_host(_host(is_linux = True, is_arm64 = True, machine = "aarch64")) + ilp.published_repo_for_host( + _host(is_linux = True, is_arm64 = True, machine = "aarch64") + ) == UPSTREAM ) # GPU Linux -> fork. assert ( - ilp.published_repo_for_host(_host(is_linux = True, is_x86_64 = True, has_usable_nvidia = True)) + ilp.published_repo_for_host( + _host(is_linux = True, is_x86_64 = True, has_usable_nvidia = True) + ) + == FORK + ) + assert ( + ilp.published_repo_for_host(_host(is_linux = True, is_x86_64 = True, has_rocm = True)) == FORK ) - assert ilp.published_repo_for_host(_host(is_linux = True, is_x86_64 = True, has_rocm = True)) == FORK # CPU-only Windows -> ggml-org (setup.ps1: the fork ships no win-cpu bundle). assert ( - ilp.published_repo_for_host(_host(system = "Windows", is_windows = True, is_x86_64 = True)) + ilp.published_repo_for_host( + _host(system = "Windows", is_windows = True, is_x86_64 = True) + ) == UPSTREAM ) # GPU Windows -> fork. assert ( ilp.published_repo_for_host( - _host(system = "Windows", is_windows = True, is_x86_64 = True, has_usable_nvidia = True) + _host( + system = "Windows", + is_windows = True, + is_x86_64 = True, + has_usable_nvidia = True, + ) ) == FORK ) @@ -121,7 +135,13 @@ def _run_resolve(monkeypatch, capsys, plans_or_exc): monkeypatch.setattr( sys, "argv", - ["install_llama_prebuilt.py", "--resolve-prebuilt", "latest", "--output-format", "json"], + [ + "install_llama_prebuilt.py", + "--resolve-prebuilt", + "latest", + "--output-format", + "json", + ], ) rc = ilp.main() assert rc == ilp.EXIT_SUCCESS @@ -133,7 +153,9 @@ def test_resolve_prebuilt_available(monkeypatch, capsys): release_tag = "b9585", llama_tag = "b9585", attempts = [ - SimpleNamespace(name = "llama-b9585-bin-macos-arm64.tar.gz", install_kind = "macos-arm64") + SimpleNamespace( + name = "llama-b9585-bin-macos-arm64.tar.gz", install_kind = "macos-arm64" + ) ], ) out = _run_resolve(monkeypatch, capsys, [plan]) @@ -153,7 +175,9 @@ def test_resolve_prebuilt_unavailable(monkeypatch, capsys): def test_resolve_prebuilt_linux_amd_tooling_routes_to_fork(monkeypatch, capsys): # CPU-probed Linux host but rocminfo on PATH: the dispatch must route to the # fork so a HIP source build is not offered an upstream CPU prebuilt. - monkeypatch.setattr(ilp, "detect_host", lambda: _host(is_linux = True, is_x86_64 = True)) + monkeypatch.setattr( + ilp, "detect_host", lambda: _host(is_linux = True, is_x86_64 = True) + ) monkeypatch.setattr(ilp.shutil, "which", lambda tool: tool == "rocminfo") seen = {} @@ -165,7 +189,13 @@ def test_resolve_prebuilt_linux_amd_tooling_routes_to_fork(monkeypatch, capsys): monkeypatch.setattr( sys, "argv", - ["install_llama_prebuilt.py", "--resolve-prebuilt", "latest", "--output-format", "json"], + [ + "install_llama_prebuilt.py", + "--resolve-prebuilt", + "latest", + "--output-format", + "json", + ], ) assert ilp.main() == ilp.EXIT_SUCCESS out = json.loads(capsys.readouterr().out.strip().splitlines()[-1]) diff --git a/studio/backend/tests/test_kv_cache_estimation.py b/studio/backend/tests/test_kv_cache_estimation.py index cd834b345b..8b79495554 100644 --- a/studio/backend/tests/test_kv_cache_estimation.py +++ b/studio/backend/tests/test_kv_cache_estimation.py @@ -338,7 +338,8 @@ class TestArchSwaPatternDefaults: assert kv_default > 0 assert kv_legacy > 0 assert kv_default < kv_legacy, ( - f"arch fallback should under-shoot legacy estimate: " f"{kv_default} >= {kv_legacy}" + f"arch fallback should under-shoot legacy estimate: " + f"{kv_default} >= {kv_legacy}" ) def test_scalar_sliding_window_pattern_expanded(self): @@ -421,9 +422,21 @@ class TestDynamicSwaResolver: from core.inference.llama_cpp import _period_from_layer_types # gemma3 (1 global/6), gpt-oss (alternating), gemma3n (1/5). - assert _period_from_layer_types((["sliding_attention"] * 5 + ["full_attention"]) * 4) == 6 - assert _period_from_layer_types(["sliding_attention", "full_attention"] * 12) == 2 - assert _period_from_layer_types((["sliding_attention"] * 4 + ["full_attention"]) * 7) == 5 + assert ( + _period_from_layer_types( + (["sliding_attention"] * 5 + ["full_attention"]) * 4 + ) + == 6 + ) + assert ( + _period_from_layer_types(["sliding_attention", "full_attention"] * 12) == 2 + ) + assert ( + _period_from_layer_types( + (["sliding_attention"] * 4 + ["full_attention"]) * 7 + ) + == 5 + ) def test_period_from_layer_types_returns_none_for_aperiodic(self): from core.inference.llama_cpp import _period_from_layer_types @@ -447,7 +460,9 @@ class TestDynamicSwaResolver: == "google/gemma-3-1b-it" ) assert ( - _hf_repo_from_url("https://huggingface.co/google/gemma-3-1b-it/blob/main/config.json") + _hf_repo_from_url( + "https://huggingface.co/google/gemma-3-1b-it/blob/main/config.json" + ) == "google/gemma-3-1b-it" ) for bad in [ @@ -500,7 +515,9 @@ class TestDynamicSwaResolver: b = _backend_from_gguf( "newmodel", _SWA_FIELDS, - general = {"general.source.huggingface.repository": "vendor/newmodel-1b-instruct"}, + general = { + "general.source.huggingface.repository": "vendor/newmodel-1b-instruct" + }, ) assert b._sliding_window_pattern == [(i + 1) % 4 != 0 for i in range(12)] assert calls == ["vendor/newmodel-1b-instruct"] @@ -547,7 +564,9 @@ class TestDynamicSwaResolver: monkeypatch.setattr(lc, "_fetch_swa_entry_from_hf", lambda repo_id: None) # Force failure into Tier 3; bypass Tier 2.5. - monkeypatch.setattr(lc, "_resolve_swa_entry_from_transformers", lambda arch: None) + monkeypatch.setattr( + lc, "_resolve_swa_entry_from_transformers", lambda arch: None + ) b = _backend_from_gguf( "newmodel", _SWA_FIELDS, @@ -593,14 +612,18 @@ class TestTransformersIntrospection: class _FakeLazyMapping(dict): def __getitem__(self, k): - return _FakeBrokenConfig if k == "brokenarch" else super().__getitem__(k) + return ( + _FakeBrokenConfig if k == "brokenarch" else super().__getitem__(k) + ) import sys, types as _types fake_auto = _types.ModuleType("transformers.models.auto.configuration_auto") fake_auto.CONFIG_MAPPING_NAMES = {"brokenarch": "FakeBroken"} fake_auto.CONFIG_MAPPING = _FakeLazyMapping({"brokenarch": "FakeBroken"}) - monkeypatch.setitem(sys.modules, "transformers.models.auto.configuration_auto", fake_auto) + monkeypatch.setitem( + sys.modules, "transformers.models.auto.configuration_auto", fake_auto + ) assert lc._resolve_swa_entry_from_transformers("brokenarch") == 7 def test_returns_none_when_transformers_unavailable(self, monkeypatch): @@ -628,7 +651,9 @@ class TestTransformersIntrospection: from core.inference.llama_cpp import _resolve_swa_entry_from_transformers assert _resolve_swa_entry_from_transformers("totally-fake-arch-xyz") is None - def test_full_resolver_uses_transformers_before_hf_fetch(self, monkeypatch, tmp_path): + def test_full_resolver_uses_transformers_before_hf_fetch( + self, monkeypatch, tmp_path + ): # Bootstrap empty: Tier 2.5 must answer before Tier 3 fires. self._isolate_cache(monkeypatch, tmp_path) from core.inference import llama_cpp as lc @@ -1273,7 +1298,8 @@ class TestServerFlags: "_kv_key_length": 256, "_kv_value_length": 256, "_sliding_window": 512, - "_sliding_window_pattern": [True, True, True, True, True, False] * 4 + [True, True], + "_sliding_window_pattern": [True, True, True, True, True, False] * 4 + + [True, True], } defaults.update(overrides) b = LlamaCppBackend() @@ -1329,7 +1355,9 @@ class TestServerFlags: def test_swa_full_suppresses_checkpoint_term(self): b = self._swa_backend() with_cp = b._estimate_kv_cache_bytes(8192, "f16", ctx_checkpoints = 8) - with_cp_full = b._estimate_kv_cache_bytes(8192, "f16", ctx_checkpoints = 8, swa_full = True) + with_cp_full = b._estimate_kv_cache_bytes( + 8192, "f16", ctx_checkpoints = 8, swa_full = True + ) no_cp_full = b._estimate_kv_cache_bytes(8192, "f16", swa_full = True) # Checkpoints only matter when SWA layers don't already keep n_ctx. assert with_cp_full == no_cp_full @@ -1346,7 +1374,9 @@ class TestServerFlags: for slots in (1, 2, 4, 8): for unified in (True, False): assert ( - b._estimate_kv_cache_bytes(4096, "f16", n_parallel = slots, kv_unified = unified) + b._estimate_kv_cache_bytes( + 4096, "f16", n_parallel = slots, kv_unified = unified + ) == baseline ) @@ -1355,7 +1385,9 @@ class TestServerFlags: baseline = b._estimate_kv_cache_bytes(4096, "f16") for unified in (True, False): assert ( - b._estimate_kv_cache_bytes(4096, "f16", n_parallel = 0, kv_unified = unified) + b._estimate_kv_cache_bytes( + 4096, "f16", n_parallel = 0, kv_unified = unified + ) == baseline ) @@ -1369,7 +1401,9 @@ class TestServerFlags: per_token_swa = 4 * (256 + 256) * 2 # k_swa/val_swa fall back per_slot_swa_cells = min(ctx, 2 * swa) # not clamped at parallel=1 global_bytes = sum( - ctx * per_token_global for f in b._sliding_window_pattern[: b._n_layers] if not f + ctx * per_token_global + for f in b._sliding_window_pattern[: b._n_layers] + if not f ) swa_bytes_per_slot = sum( per_slot_swa_cells * per_token_swa @@ -1380,12 +1414,16 @@ class TestServerFlags: assert global_bytes + swa_bytes_per_slot == baseline # Only the SWA portion scales by parallel for slots in (1, 2, 3, 4): - scaled = b._estimate_kv_cache_bytes(ctx, "f16", n_parallel = slots, kv_unified = False) + scaled = b._estimate_kv_cache_bytes( + ctx, "f16", n_parallel = slots, kv_unified = False + ) # SWA cells clamp to per_slot_ctx when ctx/slots < 2*swa per_slot_ctx = max(1, ctx // slots) cells = min(ctx, 2 * swa, per_slot_ctx) swa_bps = sum( - cells * per_token_swa for f in b._sliding_window_pattern[: b._n_layers] if f + cells * per_token_swa + for f in b._sliding_window_pattern[: b._n_layers] + if f ) assert scaled == global_bytes + slots * swa_bps @@ -1400,7 +1438,9 @@ class TestServerFlags: for slots in (1, 2, 4, 8): for unified in (True, False): assert ( - b._estimate_kv_cache_bytes(8192, "f16", n_parallel = slots, kv_unified = unified) + b._estimate_kv_cache_bytes( + 8192, "f16", n_parallel = slots, kv_unified = unified + ) == baseline ) @@ -1422,7 +1462,9 @@ class TestServerFlags: baseline = b._estimate_kv_cache_bytes(ctx, "f16") flagged = b._estimate_kv_cache_bytes(ctx, "f16", ctx_checkpoints = 4) # 22 SWA layers * 4 cps * 512 cells * 4 heads * (256+256) * 2 bytes - n_swa_layers = sum(1 for f in [True, True, True, True, True, False] * 4 + [True, True] if f) + n_swa_layers = sum( + 1 for f in [True, True, True, True, True, False] * 4 + [True, True] if f + ) per_layer = 4 * 512 * 4 * (256 + 256) * 2 assert flagged == baseline + n_swa_layers * per_layer @@ -1456,7 +1498,9 @@ class TestServerFlags: flagged = b._estimate_kv_cache_bytes( ctx, "f16", ctx_checkpoints = 4, n_parallel = slots, kv_unified = False ) - assert flagged == global_bytes + slots * (swa_bytes_per_slot + cp_extra_per_slot) + assert flagged == global_bytes + slots * ( + swa_bytes_per_slot + cp_extra_per_slot + ) # ── --kv-offload (kv_on_gpu) ─────────────────────────────────── @@ -1578,7 +1622,9 @@ class TestParallelSWAScaling: "_kv_value_length": 256, "_sliding_window": 512, # 15 SWA + 3 global, mirrors gemma-3-270m - "_sliding_window_pattern": [t == "swa" for t in (["swa"] * 5 + ["global"]) * 3], + "_sliding_window_pattern": [ + t == "swa" for t in (["swa"] * 5 + ["global"]) * 3 + ], } defaults.update(overrides) b = LlamaCppBackend() @@ -1594,7 +1640,9 @@ class TestParallelSWAScaling: for slots in (1, 2, 4, 8): for unified in (True, False): assert ( - b._estimate_kv_cache_bytes(8192, "f16", n_parallel = slots, kv_unified = unified) + b._estimate_kv_cache_bytes( + 8192, "f16", n_parallel = slots, kv_unified = unified + ) == baseline ) @@ -1648,7 +1696,9 @@ class TestParallelSWAScaling: cells = min(ctx, 2 * swa, per_slot_ctx) swa_bps = n_swa * cells * per_token for unified in (True, False): - got = b._estimate_kv_cache_bytes(ctx, "f16", n_parallel = slots, kv_unified = unified) + got = b._estimate_kv_cache_bytes( + ctx, "f16", n_parallel = slots, kv_unified = unified + ) assert got == global_bytes + slots * swa_bps def test_swa_fallback_scales_only_swa_portion(self): @@ -1693,7 +1743,8 @@ class TestParallelSWAScaling: baseline = b._estimate_kv_cache_bytes(ctx, "f16", swa_full = True) for slots in (1, 2, 4, 8): assert ( - b._estimate_kv_cache_bytes(ctx, "f16", swa_full = True, n_parallel = slots) == baseline + b._estimate_kv_cache_bytes(ctx, "f16", swa_full = True, n_parallel = slots) + == baseline ) # ── kv_unified: no-op for memory math ────────────────────────── @@ -1707,8 +1758,12 @@ class TestParallelSWAScaling: ] for label, b in backends: for slots in (1, 2, 4, 8): - u = b._estimate_kv_cache_bytes(8192, "f16", n_parallel = slots, kv_unified = True) - nu = b._estimate_kv_cache_bytes(8192, "f16", n_parallel = slots, kv_unified = False) + u = b._estimate_kv_cache_bytes( + 8192, "f16", n_parallel = slots, kv_unified = True + ) + nu = b._estimate_kv_cache_bytes( + 8192, "f16", n_parallel = slots, kv_unified = False + ) assert u == nu, f"{label} parallel={slots} unified-mismatch" # ── Empirical Gemma-3 270m formula ───────────────────────────── @@ -1845,7 +1900,9 @@ class TestSharedKVLayers: assert full_in_unshared == 4 kv_per = 4 * (256 + 256) * 2 swa_cells = min(ctx, 2 * 1024) - expected = full_in_unshared * ctx * kv_per + sliding_in_unshared * swa_cells * kv_per + expected = ( + full_in_unshared * ctx * kv_per + sliding_in_unshared * swa_cells * kv_per + ) assert b._estimate_kv_cache_bytes(ctx, "f16") == expected def test_shared_layers_reduces_estimate(self): @@ -1901,7 +1958,9 @@ class TestSharedKVLayers: per_slot_ctx = max(1, ctx // slots) swa_cells = min(ctx, 2 * swa, per_slot_ctx) swa_bytes_per_slot = sliding_in_unshared * swa_cells * per_token - flagged = b._estimate_kv_cache_bytes(ctx, "f16", n_parallel = slots, kv_unified = False) + flagged = b._estimate_kv_cache_bytes( + ctx, "f16", n_parallel = slots, kv_unified = False + ) assert flagged == global_bytes + slots * swa_bytes_per_slot def test_composes_with_ctx_checkpoints(self): diff --git a/studio/backend/tests/test_llama_cpp_context_fit.py b/studio/backend/tests/test_llama_cpp_context_fit.py index 58226f938c..a792686748 100644 --- a/studio/backend/tests/test_llama_cpp_context_fit.py +++ b/studio/backend/tests/test_llama_cpp_context_fit.py @@ -218,7 +218,9 @@ def _drive( elif gpus: gpu_indices, use_fit = inst._select_gpus(model_size, gpus) if use_fit and not explicit_ctx: - effective_ctx = min(FALLBACK_CTX, effective_ctx) if effective_ctx > 0 else FALLBACK_CTX + effective_ctx = ( + min(FALLBACK_CTX, effective_ctx) if effective_ctx > 0 else FALLBACK_CTX + ) return { "c_arg": effective_ctx if effective_ctx > 0 else 0, diff --git a/studio/backend/tests/test_llama_cpp_freshness.py b/studio/backend/tests/test_llama_cpp_freshness.py index f90c4ba0e7..078c808f0f 100644 --- a/studio/backend/tests/test_llama_cpp_freshness.py +++ b/studio/backend/tests/test_llama_cpp_freshness.py @@ -189,7 +189,9 @@ def test_latest_published_release_returns_none_on_network_failure(monkeypatch): assert fr.latest_published_release("unslothai/llama.cpp") is None -def test_latest_published_release_keeps_old_cache_on_transient_failure(monkeypatch, tmp_path): +def test_latest_published_release_keeps_old_cache_on_transient_failure( + monkeypatch, tmp_path +): # Disk entry older than TTL + network fail -> return cached value. cache_dir = tmp_path / ".freshness" cache_dir.mkdir() @@ -203,7 +205,9 @@ def test_latest_published_release_keeps_old_cache_on_transient_failure(monkeypat # check_prebuilt_freshness end-to-end. -def test_check_prebuilt_freshness_reports_stale_when_old_and_behind(monkeypatch, tmp_path): +def test_check_prebuilt_freshness_reports_stale_when_old_and_behind( + monkeypatch, tmp_path +): install_dir = tmp_path / "llama.cpp" _write_marker( install_dir, @@ -213,7 +217,9 @@ def test_check_prebuilt_freshness_reports_stale_when_old_and_behind(monkeypatch, .replace("+00:00", "Z"), ) bin_path = _fake_binary(install_dir, layout = "root") - monkeypatch.setattr(fr, "_fetch_latest_release_tag", lambda repo, timeout = 5.0: "b9300") + monkeypatch.setattr( + fr, "_fetch_latest_release_tag", lambda repo, timeout = 5.0: "b9300" + ) info = fr.check_prebuilt_freshness(str(bin_path)) assert info["has_marker"] is True assert info["stale"] is True @@ -233,7 +239,9 @@ def test_check_prebuilt_freshness_not_stale_when_tag_matches(monkeypatch, tmp_pa .replace("+00:00", "Z"), ) bin_path = _fake_binary(install_dir, layout = "root") - monkeypatch.setattr(fr, "_fetch_latest_release_tag", lambda repo, timeout = 5.0: "b9300") + monkeypatch.setattr( + fr, "_fetch_latest_release_tag", lambda repo, timeout = 5.0: "b9300" + ) info = fr.check_prebuilt_freshness(str(bin_path)) assert info["stale"] is False assert info["installed_tag"] == "b9300" @@ -251,7 +259,9 @@ def test_check_prebuilt_freshness_not_stale_within_threshold(monkeypatch, tmp_pa .replace("+00:00", "Z"), ) bin_path = _fake_binary(install_dir, layout = "root") - monkeypatch.setattr(fr, "_fetch_latest_release_tag", lambda repo, timeout = 5.0: "b9300") + monkeypatch.setattr( + fr, "_fetch_latest_release_tag", lambda repo, timeout = 5.0: "b9300" + ) info = fr.check_prebuilt_freshness(str(bin_path)) assert info["stale"] is False assert info["age_days"] == 1 @@ -264,7 +274,9 @@ def test_check_prebuilt_freshness_fails_open_without_marker(tmp_path): assert info["stale"] is False -def test_check_prebuilt_freshness_fails_open_when_github_unreachable(monkeypatch, tmp_path): +def test_check_prebuilt_freshness_fails_open_when_github_unreachable( + monkeypatch, tmp_path +): install_dir = tmp_path / "llama.cpp" _write_marker( install_dir, @@ -281,11 +293,15 @@ def test_check_prebuilt_freshness_fails_open_when_github_unreachable(monkeypatch assert info["latest_tag"] is None -def test_check_prebuilt_freshness_handles_unparseable_install_timestamp(monkeypatch, tmp_path): +def test_check_prebuilt_freshness_handles_unparseable_install_timestamp( + monkeypatch, tmp_path +): install_dir = tmp_path / "llama.cpp" _write_marker(install_dir, tag = "b9190", installed_at_utc = "not-a-date") bin_path = _fake_binary(install_dir, layout = "root") - monkeypatch.setattr(fr, "_fetch_latest_release_tag", lambda repo, timeout = 5.0: "b9300") + monkeypatch.setattr( + fr, "_fetch_latest_release_tag", lambda repo, timeout = 5.0: "b9300" + ) info = fr.check_prebuilt_freshness(str(bin_path)) assert info["stale"] is False assert info["age_days"] is None @@ -301,7 +317,9 @@ def test_check_prebuilt_freshness_respects_custom_threshold(monkeypatch, tmp_pat .replace("+00:00", "Z"), ) bin_path = _fake_binary(install_dir, layout = "root") - monkeypatch.setattr(fr, "_fetch_latest_release_tag", lambda repo, timeout = 5.0: "b9300") + monkeypatch.setattr( + fr, "_fetch_latest_release_tag", lambda repo, timeout = 5.0: "b9300" + ) info = fr.check_prebuilt_freshness(str(bin_path), threshold_days = 1) assert info["stale"] is True @@ -310,7 +328,9 @@ def test_check_prebuilt_freshness_respects_custom_threshold(monkeypatch, tmp_pat def test_format_stale_warning_contains_actionable_command(): - msg = fr.format_stale_warning({"installed_tag": "b9190", "latest_tag": "b9300", "age_days": 5}) + msg = fr.format_stale_warning( + {"installed_tag": "b9190", "latest_tag": "b9300", "age_days": 5} + ) assert "b9190" in msg assert "b9300" in msg assert "5 days" in msg @@ -318,7 +338,9 @@ def test_format_stale_warning_contains_actionable_command(): def test_format_stale_warning_singular_day(): - msg = fr.format_stale_warning({"installed_tag": "b9190", "latest_tag": "b9300", "age_days": 1}) + msg = fr.format_stale_warning( + {"installed_tag": "b9190", "latest_tag": "b9300", "age_days": 1} + ) assert "1 day" in msg assert "1 days" not in msg @@ -329,7 +351,9 @@ def test_format_stale_warning_singular_day(): def test_parse_base_build(): assert fr.parse_base_build("b9596") == 9596 assert fr.parse_base_build(" b9596 ") == 9596 - assert fr.parse_base_build("b9596-mix-e6f2453") == 9596 # mix suffix doesn't defeat it + assert ( + fr.parse_base_build("b9596-mix-e6f2453") == 9596 + ) # mix suffix doesn't defeat it assert fr.parse_base_build("9596") is None assert fr.parse_base_build("master-abc") is None assert fr.parse_base_build("") is None @@ -387,7 +411,9 @@ def test_check_prebuilt_freshness_downgrade_guard(monkeypatch, tmp_path): .replace("+00:00", "Z"), ) bin_path = _fake_binary(install_dir, layout = "root") - monkeypatch.setattr(fr, "_fetch_latest_release_tag", lambda repo, timeout = 5.0: "b9518") + monkeypatch.setattr( + fr, "_fetch_latest_release_tag", lambda repo, timeout = 5.0: "b9518" + ) info = fr.check_prebuilt_freshness(str(bin_path)) assert info["behind"] is False assert info["stale"] is False @@ -431,7 +457,9 @@ def test_fetch_latest_release_tag_uses_publish_time(monkeypatch): "published_at": "2026-06-12T00:00:00Z", }, ] - monkeypatch.setattr(urllib.request, "urlopen", lambda req, timeout = 5.0: _Resp(payload)) + monkeypatch.setattr( + urllib.request, "urlopen", lambda req, timeout = 5.0: _Resp(payload) + ) assert fr._fetch_latest_release_tag("unslothai/llama.cpp") == "b9596-mix-e6f2453" @@ -443,7 +471,9 @@ def _seed_disk_cache(tmp_path: Path, latest_tag: str) -> Path: cache_dir = tmp_path / ".freshness" cache_dir.mkdir(exist_ok = True) cache_file = cache_dir / "unslothai__llama.cpp.json" - cache_file.write_text(json.dumps({"fetched_at": time.time(), "latest_tag": latest_tag})) + cache_file.write_text( + json.dumps({"fetched_at": time.time(), "latest_tag": latest_tag}) + ) return cache_file @@ -468,7 +498,9 @@ def test_reset_caches_drop_disk_on_missing_dir_is_noop(tmp_path): fr.reset_caches(drop_disk = True) # must not raise -def test_drop_disk_lets_banner_fail_open_after_same_base_mix_swap(monkeypatch, tmp_path): +def test_drop_disk_lets_banner_fail_open_after_same_base_mix_swap( + monkeypatch, tmp_path +): # P2 #2: the disk cache holds a still-fresh same-base mix (b9596-mix-aaa) # from before an update to a *different* same-base mix (b9596-mix-bbb). # The post-install path drops the disk cache; if the forced refresh is then diff --git a/studio/backend/tests/test_llama_cpp_mmproj_fallback.py b/studio/backend/tests/test_llama_cpp_mmproj_fallback.py index 6ef94545fe..b8a65e98bc 100644 --- a/studio/backend/tests/test_llama_cpp_mmproj_fallback.py +++ b/studio/backend/tests/test_llama_cpp_mmproj_fallback.py @@ -31,7 +31,9 @@ _loggers_stub = _types.ModuleType("loggers") _loggers_stub.get_logger = lambda name: __import__("logging").getLogger(name) sys.modules.setdefault("loggers", _loggers_stub) _structlog_stub = _types.ModuleType("structlog") -_structlog_stub.get_logger = lambda *a, **k: __import__("logging").getLogger("structlog") +_structlog_stub.get_logger = lambda *a, **k: __import__("logging").getLogger( + "structlog" +) sys.modules.setdefault("structlog", _structlog_stub) if not hasattr(sys.modules["structlog"], "get_logger"): sys.modules["structlog"].get_logger = _structlog_stub.get_logger @@ -54,7 +56,9 @@ _OOM_OUT = ( "ggml_backend_cuda_buffer_type_alloc_buffer: allocating 12000.00 MiB on " "device 0: cudaMalloc failed: out of memory" ) -_BAD_ARCH_OUT = "llama_model_load: error loading model: unknown model architecture: 'qwen_image'" +_BAD_ARCH_OUT = ( + "llama_model_load: error loading model: unknown model architecture: 'qwen_image'" +) _PORT_OUT = "srv start: failed to bind: address already in use" _MISSING_OUT = "error: failed to open GGUF file: no such file or directory" # A healthy startup log that merely mentions the projector must not match. diff --git a/studio/backend/tests/test_llama_cpp_mtp_detection.py b/studio/backend/tests/test_llama_cpp_mtp_detection.py index b00cd7169e..a86e9823e4 100644 --- a/studio/backend/tests/test_llama_cpp_mtp_detection.py +++ b/studio/backend/tests/test_llama_cpp_mtp_detection.py @@ -77,7 +77,9 @@ def _enc_kv_string(key: str, value: str) -> bytes: def _enc_kv_uint32(key: str, value: int) -> bytes: - return _enc_string(key) + struct.pack(" forward --has-rocm. - assert upd._rocm_install_args("llama-b9334-bin-ubuntu-rocm-6.4-x64.tar.gz") == ["--has-rocm"] + assert upd._rocm_install_args("llama-b9334-bin-ubuntu-rocm-6.4-x64.tar.gz") == [ + "--has-rocm" + ] def test_rocm_install_args_windows_hip(): - assert upd._rocm_install_args("llama-b9334-bin-win-hip-radeon-x64.zip") == ["--has-rocm"] + assert upd._rocm_install_args("llama-b9334-bin-win-hip-radeon-x64.zip") == [ + "--has-rocm" + ] def test_rocm_install_args_non_rocm_and_missing(): @@ -460,8 +488,12 @@ def _capture_install_cmd( install_dir = tmp_path / "llama.cpp" binary = _write_install(install_dir, tag, repo = repo, asset = asset) monkeypatch.setattr(upd, "_find_binary", lambda: binary) - monkeypatch.setattr(upd, "_installer_script", lambda: tmp_path / "install_llama_prebuilt.py") - monkeypatch.setattr(freshness, "_fetch_latest_release_tag", lambda repo, timeout = 5.0: latest) + monkeypatch.setattr( + upd, "_installer_script", lambda: tmp_path / "install_llama_prebuilt.py" + ) + monkeypatch.setattr( + freshness, "_fetch_latest_release_tag", lambda repo, timeout = 5.0: latest + ) captured = {} @@ -545,7 +577,9 @@ def test_install_cmd_cuda_marker_minimal_and_backward_compatible(monkeypatch, tm def test_start_update_already_running_refuses(monkeypatch, tmp_path): binary = _write_install(tmp_path / "llama.cpp", "b9493") monkeypatch.setattr(upd, "_find_binary", lambda: binary) - monkeypatch.setattr(upd, "_installer_script", lambda: tmp_path / "install_llama_prebuilt.py") + monkeypatch.setattr( + upd, "_installer_script", lambda: tmp_path / "install_llama_prebuilt.py" + ) with upd._job_lock: upd._job.update(state = upd._JOB_RUNNING) res = upd.start_update() @@ -590,8 +624,12 @@ def test_update_sets_maintenance_flag_and_unloads(monkeypatch, tmp_path): install_dir = tmp_path / "llama.cpp" binary = _write_install(install_dir, "b9493") monkeypatch.setattr(upd, "_find_binary", lambda: binary) - monkeypatch.setattr(upd, "_installer_script", lambda: tmp_path / "install_llama_prebuilt.py") - monkeypatch.setattr(freshness, "_fetch_latest_release_tag", lambda repo, timeout = 5.0: "b9518") + monkeypatch.setattr( + upd, "_installer_script", lambda: tmp_path / "install_llama_prebuilt.py" + ) + monkeypatch.setattr( + freshness, "_fetch_latest_release_tag", lambda repo, timeout = 5.0: "b9518" + ) backend = _FakeBackend() _inject_backend(monkeypatch, backend) @@ -623,8 +661,12 @@ def test_update_clears_maintenance_flag_on_installer_failure(monkeypatch, tmp_pa install_dir = tmp_path / "llama.cpp" binary = _write_install(install_dir, "b9493") monkeypatch.setattr(upd, "_find_binary", lambda: binary) - monkeypatch.setattr(upd, "_installer_script", lambda: tmp_path / "install_llama_prebuilt.py") - monkeypatch.setattr(freshness, "_fetch_latest_release_tag", lambda repo, timeout = 5.0: "b9518") + monkeypatch.setattr( + upd, "_installer_script", lambda: tmp_path / "install_llama_prebuilt.py" + ) + monkeypatch.setattr( + freshness, "_fetch_latest_release_tag", lambda repo, timeout = 5.0: "b9518" + ) backend = _FakeBackend() _inject_backend(monkeypatch, backend) @@ -646,8 +688,12 @@ def test_update_fails_open_when_backend_unavailable(monkeypatch, tmp_path): install_dir = tmp_path / "llama.cpp" binary = _write_install(install_dir, "b9493") monkeypatch.setattr(upd, "_find_binary", lambda: binary) - monkeypatch.setattr(upd, "_installer_script", lambda: tmp_path / "install_llama_prebuilt.py") - monkeypatch.setattr(freshness, "_fetch_latest_release_tag", lambda repo, timeout = 5.0: "b9518") + monkeypatch.setattr( + upd, "_installer_script", lambda: tmp_path / "install_llama_prebuilt.py" + ) + monkeypatch.setattr( + freshness, "_fetch_latest_release_tag", lambda repo, timeout = 5.0: "b9518" + ) def _raise(): raise RuntimeError("no backend") @@ -659,7 +705,9 @@ def test_update_fails_open_when_backend_unavailable(monkeypatch, tmp_path): monkeypatch.setitem(sys.modules, "routes", routes_pkg) monkeypatch.setitem(sys.modules, "routes.inference", inference_mod) - _patch_installer_popen(monkeypatch, on_start = lambda cmd: _write_install(install_dir, "b9518")) + _patch_installer_popen( + monkeypatch, on_start = lambda cmd: _write_install(install_dir, "b9518") + ) res = upd.start_update() assert res["started"] is True @@ -676,15 +724,15 @@ def test_update_fails_open_when_backend_unavailable(monkeypatch, tmp_path): def test_resolve_prebuilt_parses_and_caches(monkeypatch, tmp_path): - monkeypatch.setattr(upd, "_installer_script", lambda: tmp_path / "install_llama_prebuilt.py") + monkeypatch.setattr( + upd, "_installer_script", lambda: tmp_path / "install_llama_prebuilt.py" + ) calls = {"n": 0} class _Proc: returncode = 0 # stderr noise plus the JSON line on stdout (installer logs to stderr). - stdout = ( - '{"prebuilt_available": true, "repo": "unslothai/llama.cpp", "release_tag": "b9585"}' - ) + stdout = '{"prebuilt_available": true, "repo": "unslothai/llama.cpp", "release_tag": "b9585"}' stderr = "[llama-prebuilt] some log\n" def _fake_run(cmd, **kwargs): @@ -701,7 +749,9 @@ def test_resolve_prebuilt_parses_and_caches(monkeypatch, tmp_path): def test_resolve_prebuilt_fails_open(monkeypatch, tmp_path): - monkeypatch.setattr(upd, "_installer_script", lambda: tmp_path / "install_llama_prebuilt.py") + monkeypatch.setattr( + upd, "_installer_script", lambda: tmp_path / "install_llama_prebuilt.py" + ) def _boom(cmd, **kwargs): raise OSError("subprocess failed") @@ -777,7 +827,9 @@ def test_llama_install_root_ignores_inactive_env_root(monkeypatch, tmp_path): assert upd._llama_install_root(str(binary)) == active -def test_llama_install_root_refuses_pinned_checkout_under_llama_cpp(monkeypatch, tmp_path): +def test_llama_install_root_refuses_pinned_checkout_under_llama_cpp( + monkeypatch, tmp_path +): # The LLAMA_SERVER_PATH pin guard must run before the ancestor scan, or a # user's own llama.cpp checkout could be handed to the installer. root = tmp_path / "my-project" / "llama.cpp" @@ -797,7 +849,9 @@ def test_start_update_source_build_refuses_when_newer(monkeypatch, tmp_path): binary.parent.mkdir(parents = True) binary.write_text("stub") # no marker monkeypatch.setattr(upd, "_find_binary", lambda: str(binary)) - monkeypatch.setattr(upd, "_installer_script", lambda: tmp_path / "install_llama_prebuilt.py") + monkeypatch.setattr( + upd, "_installer_script", lambda: tmp_path / "install_llama_prebuilt.py" + ) _prebuilt(monkeypatch, release_tag = "b9518") monkeypatch.setattr(upd, "_installed_build_number", lambda b: 9600) res = upd.start_update() @@ -810,10 +864,14 @@ def test_start_update_source_build_refuses_when_newer(monkeypatch, tmp_path): def test_status_not_offered_on_mix_latest(monkeypatch, tmp_path): # Installed the mix latest; GitHub latest is that same full tag -> no banner. - binary = _write_install(tmp_path / "llama.cpp", "b9596", release_tag = "b9596-mix-e6f2453") + binary = _write_install( + tmp_path / "llama.cpp", "b9596", release_tag = "b9596-mix-e6f2453" + ) monkeypatch.setattr(upd, "_find_binary", lambda: binary) monkeypatch.setattr( - freshness, "_fetch_latest_release_tag", lambda repo, timeout = 5.0: "b9596-mix-e6f2453" + freshness, + "_fetch_latest_release_tag", + lambda repo, timeout = 5.0: "b9596-mix-e6f2453", ) st = upd.get_update_status() assert st["update_available"] is False @@ -825,18 +883,26 @@ def test_status_not_offered_when_latest_lags(monkeypatch, tmp_path): # A lagging latest (older build than installed) must never be offered. binary = _write_install(tmp_path / "llama.cpp", "b9585") monkeypatch.setattr(upd, "_find_binary", lambda: binary) - monkeypatch.setattr(freshness, "_fetch_latest_release_tag", lambda repo, timeout = 5.0: "b9518") + monkeypatch.setattr( + freshness, "_fetch_latest_release_tag", lambda repo, timeout = 5.0: "b9518" + ) st = upd.get_update_status() assert st["update_available"] is False def test_start_update_marked_refuses_when_not_behind(monkeypatch, tmp_path): # A direct POST / stale banner must not reinstall when already on the latest. - binary = _write_install(tmp_path / "llama.cpp", "b9596", release_tag = "b9596-mix-e6f2453") + binary = _write_install( + tmp_path / "llama.cpp", "b9596", release_tag = "b9596-mix-e6f2453" + ) monkeypatch.setattr(upd, "_find_binary", lambda: binary) - monkeypatch.setattr(upd, "_installer_script", lambda: tmp_path / "install_llama_prebuilt.py") monkeypatch.setattr( - freshness, "_fetch_latest_release_tag", lambda repo, timeout = 5.0: "b9596-mix-e6f2453" + upd, "_installer_script", lambda: tmp_path / "install_llama_prebuilt.py" + ) + monkeypatch.setattr( + freshness, + "_fetch_latest_release_tag", + lambda repo, timeout = 5.0: "b9596-mix-e6f2453", ) res = upd.start_update() assert res["started"] is False diff --git a/studio/backend/tests/test_llama_cpp_wait_for_health.py b/studio/backend/tests/test_llama_cpp_wait_for_health.py index 1ba6c9f7b5..277ce4b6d1 100644 --- a/studio/backend/tests/test_llama_cpp_wait_for_health.py +++ b/studio/backend/tests/test_llama_cpp_wait_for_health.py @@ -162,7 +162,9 @@ class TestCrashLogTail: records: list = [] fake_logger = mock.Mock() - fake_logger.error = mock.Mock(side_effect = lambda msg, *a, **k: records.append(msg)) + fake_logger.error = mock.Mock( + side_effect = lambda msg, *a, **k: records.append(msg) + ) monkeypatch.setattr(_llama_mod, "logger", fake_logger) return records @@ -207,7 +209,10 @@ class TestRetryLogFilenameUnique: def test_log_name_includes_attempt_index(self): src = ( - Path(__file__).resolve().parent.parent / "core" / "inference" / "llama_cpp.py" + Path(__file__).resolve().parent.parent + / "core" + / "inference" + / "llama_cpp.py" ).read_text(encoding = "utf-8") assert "-try{_spawn_attempt}.log" in src diff --git a/studio/backend/tests/test_llama_cpp_wait_for_vram_settle.py b/studio/backend/tests/test_llama_cpp_wait_for_vram_settle.py index 493bb93e8c..e603c742b4 100644 --- a/studio/backend/tests/test_llama_cpp_wait_for_vram_settle.py +++ b/studio/backend/tests/test_llama_cpp_wait_for_vram_settle.py @@ -127,7 +127,9 @@ def test_stale_kill_skips_wait(): LlamaCppBackend._wait_for_vram_settle( **_kw(since_kill = long_ago, max_wait = 2.0, interval = 0.25) ) - assert state["calls"] == 0, "kill older than _VRAM_SETTLE_WINDOW_S must skip the wait" + assert ( + state["calls"] == 0 + ), "kill older than _VRAM_SETTLE_WINDOW_S must skip the wait" def test_empty_first_sample_returns_immediately(): @@ -215,7 +217,9 @@ def test_max_wait_respected_when_probe_is_slow(): elapsed = time.monotonic() - start # First probe (0.30 s) + at most one clipped sleep + bail. # Hard cap well below the old 0.30 + 0.25 + 0.30 = 0.85. - assert elapsed < 0.85, f"helper exceeded the deadline due to slow probes: {elapsed:.3f}s" + assert ( + elapsed < 0.85 + ), f"helper exceeded the deadline due to slow probes: {elapsed:.3f}s" def test_gpu_index_set_change_returns(): diff --git a/studio/backend/tests/test_llama_cpp_windows_nvidia_path.py b/studio/backend/tests/test_llama_cpp_windows_nvidia_path.py index 957de4bad6..bdf6f8df75 100644 --- a/studio/backend/tests/test_llama_cpp_windows_nvidia_path.py +++ b/studio/backend/tests/test_llama_cpp_windows_nvidia_path.py @@ -177,13 +177,17 @@ class TestWindowsPipNvidiaDllDirs: def test_missing_prefix_does_not_raise(self): # Nonexistent sys.prefix: resolver must return [], not raise. - result = LlamaCppBackend._windows_pip_nvidia_dll_dirs("/this/path/does/not/exist/anywhere") + result = LlamaCppBackend._windows_pip_nvidia_dll_dirs( + "/this/path/does/not/exist/anywhere" + ) assert result == [] def test_picks_up_cu13_bin_x86_64_layout(self, tmp_path): # nvidia 13.x Windows wheels ship DLLs under nvidia/cu13/bin/x86_64/ # not nvidia//bin/; else the new CUDA 13 wheels hit #5106. - dll_dir = tmp_path / "Lib" / "site-packages" / "nvidia" / "cu13" / "bin" / "x86_64" + dll_dir = ( + tmp_path / "Lib" / "site-packages" / "nvidia" / "cu13" / "bin" / "x86_64" + ) dll_dir.mkdir(parents = True) for name in ("cudart64_13.dll", "cublas64_13.dll", "cublasLt64_13.dll"): (dll_dir / name).write_bytes(b"") diff --git a/studio/backend/tests/test_llama_route_timeouts.py b/studio/backend/tests/test_llama_route_timeouts.py index 5aee6198ba..cb3c132408 100644 --- a/studio/backend/tests/test_llama_route_timeouts.py +++ b/studio/backend/tests/test_llama_route_timeouts.py @@ -66,7 +66,9 @@ def test_preheader_send_cleanup_on_disconnect_and_cancel(): return state.disconnected task = asyncio.create_task( - inf_mod._send_stream_with_preheader_cancel(_Client(), object(), request = _Request()) + inf_mod._send_stream_with_preheader_cancel( + _Client(), object(), request = _Request() + ) ) await started.wait() if cancel_parent: diff --git a/studio/backend/tests/test_llama_server_args.py b/studio/backend/tests/test_llama_server_args.py index f3a3ea1ec4..bea76203d4 100644 --- a/studio/backend/tests/test_llama_server_args.py +++ b/studio/backend/tests/test_llama_server_args.py @@ -18,7 +18,12 @@ import pytest # Load llama_server_args.py directly to avoid dragging in the full backend # chain via core/inference/__init__.py. The validator is dependency-free. -_LSA_PATH = Path(__file__).resolve().parent.parent / "core" / "inference" / "llama_server_args.py" +_LSA_PATH = ( + Path(__file__).resolve().parent.parent + / "core" + / "inference" + / "llama_server_args.py" +) _spec = importlib.util.spec_from_file_location("_lsa_test_only", _LSA_PATH) _lsa = importlib.util.module_from_spec(_spec) _spec.loader.exec_module(_lsa) @@ -487,7 +492,9 @@ def test_strip_shadowing_flags_jinja_boolean_preserves_positional(): def test_strip_shadowing_flags_no_jinja_boolean_preserves_positional(): - out = strip_shadowing_flags(["--no-jinja", "trailing-positional"], strip_template = True) + out = strip_shadowing_flags( + ["--no-jinja", "trailing-positional"], strip_template = True + ) assert out == ["trailing-positional"] @@ -664,11 +671,15 @@ def test_strip_shadowing_flags_keeps_split_mode_when_not_requested(): def test_strip_shadowing_flags_drops_split_mode_short_alias_and_equals(): - assert strip_shadowing_flags(["-sm", "tensor", "--top-k", "20"], strip_split_mode = True) == [ + assert strip_shadowing_flags( + ["-sm", "tensor", "--top-k", "20"], strip_split_mode = True + ) == [ "--top-k", "20", ] - assert strip_shadowing_flags(["--split-mode=row", "--seed", "-1"], strip_split_mode = True) == [ + assert strip_shadowing_flags( + ["--split-mode=row", "--seed", "-1"], strip_split_mode = True + ) == [ "--seed", "-1", ] diff --git a/studio/backend/tests/test_llm_assist_startup_opt_in.py b/studio/backend/tests/test_llm_assist_startup_opt_in.py index e81b1d3775..8b72105f25 100644 --- a/studio/backend/tests/test_llm_assist_startup_opt_in.py +++ b/studio/backend/tests/test_llm_assist_startup_opt_in.py @@ -72,9 +72,13 @@ def test_settings_route_persists_helper_precache_toggle(monkeypatch): def test_main_startup_uses_helper_precache_gate_instead_of_unconditional_precache(): - source = (Path(__file__).resolve().parent.parent / "main.py").read_text(encoding = "utf-8") + source = (Path(__file__).resolve().parent.parent / "main.py").read_text( + encoding = "utf-8" + ) startup_section = source[ - source.index("cleanup_orphaned_runs") : source.index("# Initialize RSA key pair") + source.index("cleanup_orphaned_runs") : source.index( + "# Initialize RSA key pair" + ) ] assert "_start_helper_precache_if_enabled()" in startup_section diff --git a/studio/backend/tests/test_login_rate_limit.py b/studio/backend/tests/test_login_rate_limit.py index 14b10576da..1a58f4fa6c 100644 --- a/studio/backend/tests/test_login_rate_limit.py +++ b/studio/backend/tests/test_login_rate_limit.py @@ -101,22 +101,30 @@ class TestClientIp: def test_xff_strips_ipv4_port(self, env_trust_proxy): from routes.auth import _client_ip - req = _FakeRequest("127.0.0.1", {"x-forwarded-for": "198.51.100.7:50001, 10.0.0.1"}) + req = _FakeRequest( + "127.0.0.1", {"x-forwarded-for": "198.51.100.7:50001, 10.0.0.1"} + ) assert _client_ip(req) == "198.51.100.7" def test_xff_strips_bracketed_ipv6_port(self, env_trust_proxy): from routes.auth import _client_ip - req = _FakeRequest("127.0.0.1", {"x-forwarded-for": "[2001:db8::1]:50001, 10.0.0.1"}) + req = _FakeRequest( + "127.0.0.1", {"x-forwarded-for": "[2001:db8::1]:50001, 10.0.0.1"} + ) assert _client_ip(req) == "2001:db8::1" def test_forwarded_strips_ipv4_port(self, env_trust_proxy): from routes.auth import _client_ip - req = _FakeRequest("127.0.0.1", {"forwarded": 'for="198.51.100.7:50001";proto=https'}) + req = _FakeRequest( + "127.0.0.1", {"forwarded": 'for="198.51.100.7:50001";proto=https'} + ) assert _client_ip(req) == "198.51.100.7" def test_forwarded_strips_bracketed_ipv6_port(self, env_trust_proxy): from routes.auth import _client_ip - req = _FakeRequest("127.0.0.1", {"forwarded": 'for="[2001:db8::1]:50001";proto=https'}) + req = _FakeRequest( + "127.0.0.1", {"forwarded": 'for="[2001:db8::1]:50001";proto=https'} + ) assert _client_ip(req) == "2001:db8::1" def test_forwarded_isolates_first_element(self, env_trust_proxy): @@ -229,7 +237,9 @@ class TestLogin429Body: import secrets as _secrets monkeypatch.setattr(storage, "DB_PATH", tmp_path / "auth.db") - monkeypatch.setattr(storage, "_BOOTSTRAP_PW_PATH", tmp_path / ".bootstrap_password") + monkeypatch.setattr( + storage, "_BOOTSTRAP_PW_PATH", tmp_path / ".bootstrap_password" + ) monkeypatch.setattr(storage, "_bootstrap_password", None) storage.create_initial_user( username = storage.DEFAULT_ADMIN_USERNAME, diff --git a/studio/backend/tests/test_mcp_config_import.py b/studio/backend/tests/test_mcp_config_import.py index 4082733490..86e0a61f18 100644 --- a/studio/backend/tests/test_mcp_config_import.py +++ b/studio/backend/tests/test_mcp_config_import.py @@ -102,7 +102,11 @@ def test_parse_windows_apostrophes_as_literals(monkeypatch): "C:\\Users\\O'Reilly\\server.js", ] assert mcp_client.parse_stdio_command("node 'draft'") == ["node", "'draft'"] - assert mcp_client.parse_stdio_command("node 'open close'") == ["node", "'open", "close'"] + assert mcp_client.parse_stdio_command("node 'open close'") == [ + "node", + "'open", + "close'", + ] def test_parse_rejects_unterminated_windows_double_quote(monkeypatch): @@ -187,11 +191,18 @@ def test_parse_accepts_cline_streamable_http_alias(): [ {"command": "node", "args": ["server.js"], "cwd": "/tmp/server"}, {"command": "node", "args": ["server.js"], "envFile": ".env"}, - {"command": "node", "args": ["server.js"], "env": {"API_KEY": "${input:api-key}"}}, + { + "command": "node", + "args": ["server.js"], + "env": {"API_KEY": "${input:api-key}"}, + }, {"command": "node", "args": ["${workspaceFolder}/server.js"]}, {"command": "node", "args": ["server.js"], "env": {"HTTP_PROXY": None}}, {"command": "node", "args": ["server.js"], "sandboxEnabled": True}, - {"url": "https://example.com/mcp", "headers": {"Authorization": "Bearer ${input:token}"}}, + { + "url": "https://example.com/mcp", + "headers": {"Authorization": "Bearer ${input:token}"}, + }, {"url": "https://example.com/mcp", "headers": {"Authorization": None}}, {"type": "http", "url": "https://example.com/sse"}, {"type": "http", "url": "https://example.com/sse "}, @@ -217,7 +228,9 @@ def test_servers_alias_key(): def test_env_and_args_values_coerced_to_str(): - cfg = {"mcpServers": {"fs": {"command": "node", "args": [8080], "env": {"PORT": 8080}}}} + cfg = { + "mcpServers": {"fs": {"command": "node", "args": [8080], "env": {"PORT": 8080}}} + } entries, errors = parse_mcp_config(cfg) assert errors == [] assert entries[0].headers == {"PORT": "8080"} @@ -296,11 +309,18 @@ def test_import_route_creates_and_dedups(tmp_path, monkeypatch): } } res = asyncio.run( - routes_mcp.import_mcp_servers(McpServerImportRequest(config = cfg), current_subject = "u") + routes_mcp.import_mcp_servers( + McpServerImportRequest(config = cfg), current_subject = "u" + ) ) assert res.errors == [] assert res.skipped == [] - assert {c.display_name for c in res.created} == {"fs", "remote", "oauth", "disabled"} + assert {c.display_name for c in res.created} == { + "fs", + "remote", + "oauth", + "disabled", + } fs = next(c for c in res.created if c.display_name == "fs") assert fs.headers == {"API_KEY": "sk"} assert fs.use_oauth is False @@ -312,7 +332,9 @@ def test_import_route_creates_and_dedups(tmp_path, monkeypatch): # Re-importing the same config skips both by url. res2 = asyncio.run( - routes_mcp.import_mcp_servers(McpServerImportRequest(config = cfg), current_subject = "u") + routes_mcp.import_mcp_servers( + McpServerImportRequest(config = cfg), current_subject = "u" + ) ) assert res2.created == [] assert set(res2.skipped) == {"fs", "remote", "oauth", "disabled"} @@ -333,7 +355,9 @@ def test_import_route_gates_stdio_when_disabled(tmp_path, monkeypatch): } } res = asyncio.run( - routes_mcp.import_mcp_servers(McpServerImportRequest(config = cfg), current_subject = "u") + routes_mcp.import_mcp_servers( + McpServerImportRequest(config = cfg), current_subject = "u" + ) ) # Remote still imports; the stdio entry is rejected per-entry (gate off). assert {c.display_name for c in res.created} == {"remote"} diff --git a/studio/backend/tests/test_mcp_servers.py b/studio/backend/tests/test_mcp_servers.py index 90b1ade03c..b22eb74a8e 100644 --- a/studio/backend/tests/test_mcp_servers.py +++ b/studio/backend/tests/test_mcp_servers.py @@ -45,7 +45,9 @@ def test_list_servers_ordered_by_created_at(tmp_path, monkeypatch): def test_update_server_coerces_bools(tmp_path, monkeypatch): _reset_db(tmp_path, monkeypatch) mcp_servers_db.create_server(id = "srv1", display_name = "A", url = "https://a/m") - assert mcp_servers_db.update_server("srv1", {"is_enabled": False, "use_oauth": True}) + assert mcp_servers_db.update_server( + "srv1", {"is_enabled": False, "use_oauth": True} + ) row = mcp_servers_db.get_server("srv1") assert row["is_enabled"] == 0 assert row["use_oauth"] == 1 @@ -87,7 +89,9 @@ def test_validate_url_rejects_bad(bad): def test_normalize_headers(): from routes.mcp_servers import _normalize_headers - assert _normalize_headers({" Auth ": "Bearer x", "": "ignored"}) == {"Auth": "Bearer x"} + assert _normalize_headers({" Auth ": "Bearer x", "": "ignored"}) == { + "Auth": "Bearer x" + } assert _normalize_headers({"X": 42}) == {"X": "42"} assert _normalize_headers({}) is None assert _normalize_headers(None) is None @@ -99,12 +103,15 @@ def test_changes_from_payload_tristate_headers(): from models.mcp_servers import McpServerUpdate # omitted → key absent - assert "headers_json" not in _changes_from_payload(McpServerUpdate(display_name = "x")) + assert "headers_json" not in _changes_from_payload( + McpServerUpdate(display_name = "x") + ) # null → stored as None (clear all headers) assert _changes_from_payload(McpServerUpdate(headers = None))["headers_json"] is None # dict → serialised JSON assert ( - _changes_from_payload(McpServerUpdate(headers = {"a": "1"}))["headers_json"] == '{"a": "1"}' + _changes_from_payload(McpServerUpdate(headers = {"a": "1"}))["headers_json"] + == '{"a": "1"}' ) @@ -134,7 +141,10 @@ def test_execute_tool_malformed_mcp_name(): def test_execute_tool_unknown_server(tmp_path, monkeypatch): _reset_db(tmp_path, monkeypatch) from core.inference.tools import execute_tool - assert execute_tool("mcp__missing__do_thing", {}) == "Error: MCP server 'missing' not found" + assert ( + execute_tool("mcp__missing__do_thing", {}) + == "Error: MCP server 'missing' not found" + ) def test_execute_tool_disabled_server(tmp_path, monkeypatch): @@ -147,7 +157,10 @@ def test_execute_tool_disabled_server(tmp_path, monkeypatch): ) from core.inference.tools import execute_tool - assert execute_tool("mcp__srv1__do_thing", {}) == "Error: MCP server 'srv1' is disabled" + assert ( + execute_tool("mcp__srv1__do_thing", {}) + == "Error: MCP server 'srv1' is disabled" + ) def test_mcp_specs_skip_invalid_openai_function_names(): @@ -616,7 +629,9 @@ def test_get_enabled_mcp_tools_caches_discovery(tmp_path, monkeypatch): from core.inference import tools as tools_mod monkeypatch.setattr(mcp_client, "_tool_cache", {}) - mcp_servers_db.create_server(id = "s1", display_name = "A", url = "https://x/mcp", is_enabled = True) + mcp_servers_db.create_server( + id = "s1", display_name = "A", url = "https://x/mcp", is_enabled = True + ) calls: list[str] = [] @@ -649,7 +664,9 @@ def test_get_enabled_mcp_tools_does_not_cache_failures(tmp_path, monkeypatch): monkeypatch.setattr(mcp_client, "_tool_cache", {}) monkeypatch.setattr(mcp_client, "_probe_cooloff_until", {}) - mcp_servers_db.create_server(id = "s1", display_name = "A", url = "https://x/mcp", is_enabled = True) + mcp_servers_db.create_server( + id = "s1", display_name = "A", url = "https://x/mcp", is_enabled = True + ) attempts = {"n": 0} @@ -684,7 +701,9 @@ def test_refresh_warms_tool_cache(tmp_path, monkeypatch): import routes.mcp_servers as routes_mcp monkeypatch.setattr(mcp_client, "_tool_cache", {}) - mcp_servers_db.create_server(id = "s1", display_name = "A", url = "https://x/mcp", is_enabled = True) + mcp_servers_db.create_server( + id = "s1", display_name = "A", url = "https://x/mcp", is_enabled = True + ) async def fake_refresh( url, @@ -716,7 +735,9 @@ def test_update_url_evicts_tool_cache(tmp_path, monkeypatch): import routes.mcp_servers as routes_mcp monkeypatch.setattr(mcp_client, "_tool_cache", {"s1": _one_tool("stale")}) - mcp_servers_db.create_server(id = "s1", display_name = "A", url = "https://old/mcp", is_enabled = True) + mcp_servers_db.create_server( + id = "s1", display_name = "A", url = "https://old/mcp", is_enabled = True + ) asyncio.run( routes_mcp.update_mcp_server( @@ -737,10 +758,14 @@ def test_update_display_name_keeps_tool_cache(tmp_path, monkeypatch): cached = _one_tool() monkeypatch.setattr(mcp_client, "_tool_cache", {"s1": cached}) - mcp_servers_db.create_server(id = "s1", display_name = "A", url = "https://x/mcp", is_enabled = True) + mcp_servers_db.create_server( + id = "s1", display_name = "A", url = "https://x/mcp", is_enabled = True + ) asyncio.run( - routes_mcp.update_mcp_server("s1", McpServerUpdate(display_name = "B"), current_subject = "u") + routes_mcp.update_mcp_server( + "s1", McpServerUpdate(display_name = "B"), current_subject = "u" + ) ) assert mcp_client.get_cached_tools("s1") == cached @@ -755,10 +780,14 @@ def test_update_disable_evicts_tool_cache(tmp_path, monkeypatch): import routes.mcp_servers as routes_mcp monkeypatch.setattr(mcp_client, "_tool_cache", {"s1": _one_tool()}) - mcp_servers_db.create_server(id = "s1", display_name = "A", url = "https://x/mcp", is_enabled = True) + mcp_servers_db.create_server( + id = "s1", display_name = "A", url = "https://x/mcp", is_enabled = True + ) asyncio.run( - routes_mcp.update_mcp_server("s1", McpServerUpdate(is_enabled = False), current_subject = "u") + routes_mcp.update_mcp_server( + "s1", McpServerUpdate(is_enabled = False), current_subject = "u" + ) ) assert mcp_client.get_cached_tools("s1") is None @@ -772,7 +801,9 @@ def test_delete_evicts_tool_cache(tmp_path, monkeypatch): import routes.mcp_servers as routes_mcp monkeypatch.setattr(mcp_client, "_tool_cache", {"s1": _one_tool()}) - mcp_servers_db.create_server(id = "s1", display_name = "A", url = "https://x/mcp", is_enabled = True) + mcp_servers_db.create_server( + id = "s1", display_name = "A", url = "https://x/mcp", is_enabled = True + ) asyncio.run(routes_mcp.delete_mcp_server("s1", current_subject = "u")) assert mcp_client.get_cached_tools("s1") is None @@ -795,8 +826,12 @@ def test_get_enabled_mcp_tools_probes_only_uncached(tmp_path, monkeypatch): from core.inference import tools as tools_mod monkeypatch.setattr(mcp_client, "_tool_cache", {"s1": _one_tool("cached")}) - mcp_servers_db.create_server(id = "s1", display_name = "A", url = "https://a/mcp", is_enabled = True) - mcp_servers_db.create_server(id = "s2", display_name = "B", url = "https://b/mcp", is_enabled = True) + mcp_servers_db.create_server( + id = "s1", display_name = "A", url = "https://a/mcp", is_enabled = True + ) + mcp_servers_db.create_server( + id = "s2", display_name = "B", url = "https://b/mcp", is_enabled = True + ) probed: list[str] = [] @@ -813,7 +848,10 @@ def test_get_enabled_mcp_tools_probes_only_uncached(tmp_path, monkeypatch): specs = asyncio.run(tools_mod.get_enabled_mcp_tools()) assert probed == ["https://b/mcp"] # only the uncached server is probed - assert sorted(t["function"]["name"] for t in specs) == ["mcp__s1__cached", "mcp__s2__fresh"] + assert sorted(t["function"]["name"] for t in specs) == [ + "mcp__s1__cached", + "mcp__s2__fresh", + ] def test_get_enabled_mcp_tools_partial_failure_caches_healthy(tmp_path, monkeypatch): @@ -826,8 +864,12 @@ def test_get_enabled_mcp_tools_partial_failure_caches_healthy(tmp_path, monkeypa monkeypatch.setattr(mcp_client, "_tool_cache", {}) monkeypatch.setattr(mcp_client, "_probe_cooloff_until", {}) - mcp_servers_db.create_server(id = "s1", display_name = "A", url = "https://bad/mcp", is_enabled = True) - mcp_servers_db.create_server(id = "s2", display_name = "B", url = "https://good/mcp", is_enabled = True) + mcp_servers_db.create_server( + id = "s1", display_name = "A", url = "https://bad/mcp", is_enabled = True + ) + mcp_servers_db.create_server( + id = "s2", display_name = "B", url = "https://good/mcp", is_enabled = True + ) async def fake( url, @@ -856,7 +898,9 @@ def test_get_enabled_mcp_tools_caches_empty_tool_list(tmp_path, monkeypatch): from core.inference import tools as tools_mod monkeypatch.setattr(mcp_client, "_tool_cache", {}) - mcp_servers_db.create_server(id = "s1", display_name = "A", url = "https://x/mcp", is_enabled = True) + mcp_servers_db.create_server( + id = "s1", display_name = "A", url = "https://x/mcp", is_enabled = True + ) calls: list[str] = [] @@ -887,7 +931,9 @@ def test_update_headers_evicts_tool_cache(tmp_path, monkeypatch): import routes.mcp_servers as routes_mcp monkeypatch.setattr(mcp_client, "_tool_cache", {"s1": _one_tool()}) - mcp_servers_db.create_server(id = "s1", display_name = "A", url = "https://x/mcp", is_enabled = True) + mcp_servers_db.create_server( + id = "s1", display_name = "A", url = "https://x/mcp", is_enabled = True + ) asyncio.run( routes_mcp.update_mcp_server( @@ -899,7 +945,9 @@ def test_update_headers_evicts_tool_cache(tmp_path, monkeypatch): assert mcp_client.get_cached_tools("s1") is None -def test_get_enabled_mcp_tools_skips_cache_when_config_changes_mid_probe(tmp_path, monkeypatch): +def test_get_enabled_mcp_tools_skips_cache_when_config_changes_mid_probe( + tmp_path, monkeypatch +): """A config edit landing during an in-flight probe must not be clobbered by the now-stale probe result (TOCTOU on the cache write).""" import asyncio @@ -909,7 +957,9 @@ def test_get_enabled_mcp_tools_skips_cache_when_config_changes_mid_probe(tmp_pat from core.inference import tools as tools_mod monkeypatch.setattr(mcp_client, "_tool_cache", {}) - mcp_servers_db.create_server(id = "s1", display_name = "A", url = "https://old/mcp", is_enabled = True) + mcp_servers_db.create_server( + id = "s1", display_name = "A", url = "https://old/mcp", is_enabled = True + ) async def fake( url, @@ -942,7 +992,9 @@ def test_get_enabled_mcp_tools_no_cooloff_when_config_changes_mid_failed_probe( monkeypatch.setattr(mcp_client, "_tool_cache", {}) monkeypatch.setattr(mcp_client, "_probe_cooloff_until", {}) - mcp_servers_db.create_server(id = "s1", display_name = "A", url = "https://old/mcp", is_enabled = True) + mcp_servers_db.create_server( + id = "s1", display_name = "A", url = "https://old/mcp", is_enabled = True + ) async def fake( url, @@ -974,7 +1026,9 @@ def test_get_enabled_mcp_tools_no_cooloff_when_server_deleted_mid_failed_probe( monkeypatch.setattr(mcp_client, "_tool_cache", {}) monkeypatch.setattr(mcp_client, "_probe_cooloff_until", {}) - mcp_servers_db.create_server(id = "s1", display_name = "A", url = "https://x/mcp", is_enabled = True) + mcp_servers_db.create_server( + id = "s1", display_name = "A", url = "https://x/mcp", is_enabled = True + ) async def fake( url, @@ -991,7 +1045,9 @@ def test_get_enabled_mcp_tools_no_cooloff_when_server_deleted_mid_failed_probe( assert "s1" not in mcp_client._probe_cooloff_until # no orphan cool-off -def test_get_enabled_mcp_tools_skips_failed_server_during_cooloff(tmp_path, monkeypatch): +def test_get_enabled_mcp_tools_skips_failed_server_during_cooloff( + tmp_path, monkeypatch +): """A down server is probed once, then skipped during the cool-off instead of being re-probed (and re-hung) on every send.""" import asyncio @@ -1002,7 +1058,9 @@ def test_get_enabled_mcp_tools_skips_failed_server_during_cooloff(tmp_path, monk monkeypatch.setattr(mcp_client, "_tool_cache", {}) monkeypatch.setattr(mcp_client, "_probe_cooloff_until", {}) - mcp_servers_db.create_server(id = "s1", display_name = "A", url = "https://x/mcp", is_enabled = True) + mcp_servers_db.create_server( + id = "s1", display_name = "A", url = "https://x/mcp", is_enabled = True + ) attempts = {"n": 0} @@ -1043,7 +1101,10 @@ def test_oauth_failure_cools_off_longer_than_plain(monkeypatch): monkeypatch.setattr(mcp_client, "_probe_cooloff_until", {}) mcp_client.record_probe_failure("plain", use_oauth = False) mcp_client.record_probe_failure("oauth", use_oauth = True) - assert mcp_client._probe_cooloff_until["oauth"] > mcp_client._probe_cooloff_until["plain"] + assert ( + mcp_client._probe_cooloff_until["oauth"] + > mcp_client._probe_cooloff_until["plain"] + ) def test_invalidate_clears_failure_cooloff(monkeypatch): @@ -1070,7 +1131,9 @@ def test_refresh_failure_records_cooloff(tmp_path, monkeypatch): monkeypatch.setattr(mcp_client, "_tool_cache", {}) monkeypatch.setattr(mcp_client, "_probe_cooloff_until", {}) - mcp_servers_db.create_server(id = "s1", display_name = "A", url = "https://x/mcp", is_enabled = True) + mcp_servers_db.create_server( + id = "s1", display_name = "A", url = "https://x/mcp", is_enabled = True + ) async def boom( url, @@ -1096,7 +1159,9 @@ def test_refresh_drops_result_when_config_changes_mid_probe(tmp_path, monkeypatc import routes.mcp_servers as routes_mcp monkeypatch.setattr(mcp_client, "_tool_cache", {}) - mcp_servers_db.create_server(id = "s1", display_name = "A", url = "https://old/mcp", is_enabled = True) + mcp_servers_db.create_server( + id = "s1", display_name = "A", url = "https://old/mcp", is_enabled = True + ) async def fake_refresh( url, @@ -1113,7 +1178,9 @@ def test_refresh_drops_result_when_config_changes_mid_probe(tmp_path, monkeypatc assert mcp_client.get_cached_tools("s1") is None -def test_refresh_failure_no_cooloff_when_config_changes_mid_probe(tmp_path, monkeypatch): +def test_refresh_failure_no_cooloff_when_config_changes_mid_probe( + tmp_path, monkeypatch +): """A manual refresh failure for an old config must not cool off the freshly edited server.""" import asyncio @@ -1124,7 +1191,9 @@ def test_refresh_failure_no_cooloff_when_config_changes_mid_probe(tmp_path, monk monkeypatch.setattr(mcp_client, "_tool_cache", {}) monkeypatch.setattr(mcp_client, "_probe_cooloff_until", {}) - mcp_servers_db.create_server(id = "s1", display_name = "A", url = "https://old/mcp", is_enabled = True) + mcp_servers_db.create_server( + id = "s1", display_name = "A", url = "https://old/mcp", is_enabled = True + ) async def boom( url, @@ -1141,7 +1210,9 @@ def test_refresh_failure_no_cooloff_when_config_changes_mid_probe(tmp_path, monk assert not mcp_client.in_failure_cooloff("s1") -def test_get_enabled_mcp_tools_drops_result_when_server_deleted_mid_probe(tmp_path, monkeypatch): +def test_get_enabled_mcp_tools_drops_result_when_server_deleted_mid_probe( + tmp_path, monkeypatch +): """A delete landing while a probe is in flight must drop the now-orphan result -- the `fresh is None` arm of the mid-probe TOCTOU guard. The result is neither served nor cached under the since-removed id.""" @@ -1153,7 +1224,9 @@ def test_get_enabled_mcp_tools_drops_result_when_server_deleted_mid_probe(tmp_pa monkeypatch.setattr(mcp_client, "_tool_cache", {}) monkeypatch.setattr(mcp_client, "_probe_cooloff_until", {}) - mcp_servers_db.create_server(id = "s1", display_name = "A", url = "https://x/mcp", is_enabled = True) + mcp_servers_db.create_server( + id = "s1", display_name = "A", url = "https://x/mcp", is_enabled = True + ) async def fake( url, diff --git a/studio/backend/tests/test_mcp_stdio_improvements.py b/studio/backend/tests/test_mcp_stdio_improvements.py index b0bfd45135..b997067034 100644 --- a/studio/backend/tests/test_mcp_stdio_improvements.py +++ b/studio/backend/tests/test_mcp_stdio_improvements.py @@ -62,7 +62,9 @@ def test_create_forces_oauth_off_for_stdio(tmp_path, monkeypatch): _enable(monkeypatch) resp = asyncio.run( routes_mcp.create_mcp_server( - McpServerCreate(display_name = "FS", url = "npx -y server /tmp", use_oauth = True), + McpServerCreate( + display_name = "FS", url = "npx -y server /tmp", use_oauth = True + ), current_subject = "u", ) ) @@ -92,8 +94,12 @@ def test_update_url_to_stdio_clears_oauth(tmp_path, monkeypatch): _reset_db(tmp_path, monkeypatch) _enable(monkeypatch) monkeypatch.setattr(mcp_client, "_oauth_token_store", None) - monkeypatch.setattr(routes_mcp, "clear_oauth_tokens_async", lambda *a, **k: asyncio.sleep(0)) - mcp_servers_db.create_server(id = "s1", display_name = "A", url = "https://a/mcp", use_oauth = True) + monkeypatch.setattr( + routes_mcp, "clear_oauth_tokens_async", lambda *a, **k: asyncio.sleep(0) + ) + mcp_servers_db.create_server( + id = "s1", display_name = "A", url = "https://a/mcp", use_oauth = True + ) resp = asyncio.run( routes_mcp.update_mcp_server( "s1", McpServerUpdate(url = "npx -y server /tmp"), current_subject = "u" @@ -142,7 +148,9 @@ def test_switch_keeps_explicitly_supplied_headers(tmp_path, monkeypatch): resp = asyncio.run( routes_mcp.update_mcp_server( "s1", - McpServerUpdate(url = "https://remote/mcp", headers = {"Authorization": "Bearer new"}), + McpServerUpdate( + url = "https://remote/mcp", headers = {"Authorization": "Bearer new"} + ), current_subject = "u", ) ) @@ -163,7 +171,9 @@ def test_same_transport_edit_keeps_headers(tmp_path, monkeypatch): ) # editing only the display name (still stdio) must keep env vars resp = asyncio.run( - routes_mcp.update_mcp_server("s1", McpServerUpdate(display_name = "B"), current_subject = "u") + routes_mcp.update_mcp_server( + "s1", McpServerUpdate(display_name = "B"), current_subject = "u" + ) ) assert resp.headers == {"API_KEY": "secret"} @@ -184,7 +194,9 @@ def test_validate_url_allows_url_in_argument(monkeypatch): from routes.mcp_servers import _validate_url _enable(monkeypatch) # :// inside an ARGUMENT (not the first token) is a valid command - assert _validate_url("npx server --url https://x/mcp") == ("npx server --url https://x/mcp") + assert _validate_url("npx server --url https://x/mcp") == ( + "npx server --url https://x/mcp" + ) # ── P6: Data Recipe stdio path obeys the same host gate ───────────── diff --git a/studio/backend/tests/test_mcp_stdio_pr5863.py b/studio/backend/tests/test_mcp_stdio_pr5863.py index 9a3e8d6882..d821dde84b 100644 --- a/studio/backend/tests/test_mcp_stdio_pr5863.py +++ b/studio/backend/tests/test_mcp_stdio_pr5863.py @@ -84,7 +84,9 @@ def transport(monkeypatch): monkeypatch.setattr( mcp_client, "_client", - lambda url, headers, use_oauth = False: _RecordingClient(url, headers, use_oauth, recorder), + lambda url, headers, use_oauth = False: _RecordingClient( + url, headers, use_oauth, recorder + ), ) return recorder @@ -129,7 +131,9 @@ def test_parse_basic_argv(): def test_parse_keeps_url_argument_as_one_command(): # gemini "high": a :// inside an ARGUMENT must not break the command. - assert mcp_client.parse_stdio_command("npx server --endpoint https://example.com/mcp") == [ + assert mcp_client.parse_stdio_command( + "npx server --endpoint https://example.com/mcp" + ) == [ "npx", "server", "--endpoint", @@ -160,7 +164,9 @@ def test_parse_windows_strips_wrapping_quotes(monkeypatch): # gemini "medium": posix=False keeps backslash paths but also the # wrapping quotes; the PR strips a matched pair so argv[0] is clean. monkeypatch.setattr(sys, "platform", "win32") - parts = mcp_client.parse_stdio_command(r'"C:\Program Files\node\node.exe" server.js') + parts = mcp_client.parse_stdio_command( + r'"C:\Program Files\node\node.exe" server.js' + ) assert parts[0] == r"C:\Program Files\node\node.exe" assert parts[1] == "server.js" @@ -242,10 +248,14 @@ def test_validate_url_gate_on_accepts_stdio(monkeypatch): # http still works when stdio is on assert _validate_url("https://x/mcp") == "https://x/mcp" # url-bearing argument accepted as a command - assert _validate_url("npx server --url https://x/mcp") == ("npx server --url https://x/mcp") + assert _validate_url("npx server --url https://x/mcp") == ( + "npx server --url https://x/mcp" + ) # A lone token is ambiguous; accept it as a command rather than # guessing it's a URL (no regression for single binaries). - assert _validate_url("/usr/local/bin/my-mcp-server") == "/usr/local/bin/my-mcp-server" + assert ( + _validate_url("/usr/local/bin/my-mcp-server") == "/usr/local/bin/my-mcp-server" + ) assert _validate_url("mcp-server-sqlite") == "mcp-server-sqlite" # empty / unparseable still rejected for bad in [" ", '"unclosed']: @@ -332,7 +342,9 @@ def test_refresh_route_gate(tmp_path, monkeypatch, transport): assert transport == [] _enable(monkeypatch) - res = asyncio.run(routes_mcp.refresh_mcp_server_tools("stdio1", current_subject = "u")) + res = asyncio.run( + routes_mcp.refresh_mcp_server_tools("stdio1", current_subject = "u") + ) assert res.ok and res.tool_count == 2 assert len(transport) == 1 @@ -343,7 +355,9 @@ def test_discovery_gate(tmp_path, monkeypatch, transport): from core.inference.tools import get_enabled_mcp_tools _reset_db(tmp_path, monkeypatch) - mcp_servers_db.create_server(id = "stdio1", display_name = "FS", url = "npx server", is_enabled = True) + mcp_servers_db.create_server( + id = "stdio1", display_name = "FS", url = "npx server", is_enabled = True + ) _disable(monkeypatch) assert asyncio.run(get_enabled_mcp_tools()) == [] @@ -359,7 +373,9 @@ def test_execute_gate(tmp_path, monkeypatch, transport): from core.inference.tools import execute_tool _reset_db(tmp_path, monkeypatch) - mcp_servers_db.create_server(id = "stdio1", display_name = "FS", url = "npx server", is_enabled = True) + mcp_servers_db.create_server( + id = "stdio1", display_name = "FS", url = "npx server", is_enabled = True + ) _disable(monkeypatch) out = execute_tool("mcp__stdio1__list_directory", {"path": "/tmp"}) diff --git a/studio/backend/tests/test_middleware.py b/studio/backend/tests/test_middleware.py index 1005431926..56c5b6288f 100644 --- a/studio/backend/tests/test_middleware.py +++ b/studio/backend/tests/test_middleware.py @@ -165,7 +165,9 @@ class TestMaxBodyMiddleware: assert r.status_code == 200 assert r.json()["total"] == 512 - def test_upload_passthrough_rejects_declared_body_over_dedicated_cap(self, main_module): + def test_upload_passthrough_rejects_declared_body_over_dedicated_cap( + self, main_module + ): app = _make_protected_app( 128, main_module, @@ -267,7 +269,9 @@ class TestSecurityHeadersMiddleware: csp = r.headers["content-security-policy"] assert f"'nonce-{nonce}'" in csp # Internal handoff header must not leak to clients. - assert main_module._CSP_SCRIPT_NONCE_HEADER not in {k.lower() for k in r.headers.keys()} + assert main_module._CSP_SCRIPT_NONCE_HEADER not in { + k.lower() for k in r.headers.keys() + } def test_build_csp_helper_shape(self, main_module): plain = main_module._build_csp() diff --git a/studio/backend/tests/test_mlx_training_worker_config.py b/studio/backend/tests/test_mlx_training_worker_config.py index 4402031467..afd141563a 100644 --- a/studio/backend/tests/test_mlx_training_worker_config.py +++ b/studio/backend/tests/test_mlx_training_worker_config.py @@ -45,8 +45,12 @@ def _load_worker_module(): setattr(wheel_utils, name, lambda *_args, **_kwargs: None) sys.modules["utils.wheel_utils"] = wheel_utils - worker_path = Path(__file__).resolve().parents[1] / "core" / "training" / "worker.py" - spec = importlib.util.spec_from_file_location("mlx_training_worker_under_test", worker_path) + worker_path = ( + Path(__file__).resolve().parents[1] / "core" / "training" / "worker.py" + ) + spec = importlib.util.spec_from_file_location( + "mlx_training_worker_under_test", worker_path + ) module = importlib.util.module_from_spec(spec) assert spec.loader is not None spec.loader.exec_module(module) @@ -138,7 +142,9 @@ def test_mlx_vlm_resized_image_layout_probes_processor_contract(): == "chw" ) assert ( - _mlx_vlm_resized_image_layout(types.SimpleNamespace(image_processor = HwcImageProcessor())) + _mlx_vlm_resized_image_layout( + types.SimpleNamespace(image_processor = HwcImageProcessor()) + ) is None ) @@ -157,7 +163,9 @@ def test_mlx_vlm_layout_probe_copies_image_processor(): image_processor = StatefulImageProcessor() - layout = _mlx_vlm_resized_image_layout(types.SimpleNamespace(image_processor = image_processor)) + layout = _mlx_vlm_resized_image_layout( + types.SimpleNamespace(image_processor = image_processor) + ) assert layout == "chw" assert image_processor.calls == 0 diff --git a/studio/backend/tests/test_models_get_model_config_case_resolution.py b/studio/backend/tests/test_models_get_model_config_case_resolution.py index 417ec74d17..3481e29948 100644 --- a/studio/backend/tests/test_models_get_model_config_case_resolution.py +++ b/studio/backend/tests/test_models_get_model_config_case_resolution.py @@ -50,7 +50,9 @@ def test_get_model_config_resolves_cached_case_before_model_checks(monkeypatch): return _DummyModelConfig() monkeypatch.setattr(models_route, "is_local_path", lambda _: False) - monkeypatch.setattr(models_route, "resolve_cached_repo_id_case", lambda _: "Org/Model") + monkeypatch.setattr( + models_route, "resolve_cached_repo_id_case", lambda _: "Org/Model" + ) monkeypatch.setattr(models_route, "load_model_defaults", _record_load) monkeypatch.setattr(models_route, "is_vision_model", _record_vision) monkeypatch.setattr(models_route, "is_embedding_model", _record_embedding) diff --git a/studio/backend/tests/test_mtp_drafter_companion.py b/studio/backend/tests/test_mtp_drafter_companion.py index d5e8d13652..6451be2901 100644 --- a/studio/backend/tests/test_mtp_drafter_companion.py +++ b/studio/backend/tests/test_mtp_drafter_companion.py @@ -201,7 +201,9 @@ def test_detect_mtp_file_search_root(tmp_path): sub.mkdir() (sub / "gemma-4-12b-it-Q4_K_M.gguf").write_bytes(b"x") (tmp_path / "mtp-gemma-4-12b-it.gguf").write_bytes(b"x") - found = detect_mtp_file(str(sub / "gemma-4-12b-it-Q4_K_M.gguf"), search_root = str(tmp_path)) + found = detect_mtp_file( + str(sub / "gemma-4-12b-it-Q4_K_M.gguf"), search_root = str(tmp_path) + ) assert found is not None and found.endswith("mtp-gemma-4-12b-it.gguf") diff --git a/studio/backend/tests/test_multimodal_document.py b/studio/backend/tests/test_multimodal_document.py index 5cd7c876cc..3d38d40b41 100644 --- a/studio/backend/tests/test_multimodal_document.py +++ b/studio/backend/tests/test_multimodal_document.py @@ -309,7 +309,11 @@ def test_openai_base64_pdf_becomes_input_file(monkeypatch): user_msg = captured["body"]["input"][0] parts = user_msg["content"] fileblk = next(p for p in parts if p.get("type") == "input_file") - assert fileblk == {"type": "input_file", "file_data": _PDF_DATA_URI, "filename": "paper.pdf"} + assert fileblk == { + "type": "input_file", + "file_data": _PDF_DATA_URI, + "filename": "paper.pdf", + } def test_openai_url_pdf_becomes_input_file(monkeypatch): @@ -490,7 +494,9 @@ def test_build_external_messages_passes_input_document_for_anthropic_and_openai( ) ] for provider in ("anthropic", "openai"): - out = _build_external_messages(msgs, supports_vision = True, provider_type = provider) + out = _build_external_messages( + msgs, supports_vision = True, provider_type = provider + ) assert len(out) == 1, (provider, out) parts = out[0]["content"] assert parts[0] == {"type": "text", "text": "summarise"}, provider @@ -526,7 +532,9 @@ def test_build_external_messages_strips_input_document_for_unmapped_providers(): ) ] for provider in ("gemini", "mistral", "kimi", "openrouter", "deepseek", "qwen"): - out = _build_external_messages(msgs, supports_vision = True, provider_type = provider) + out = _build_external_messages( + msgs, supports_vision = True, provider_type = provider + ) assert len(out) == 1, (provider, out) parts = out[0]["content"] types = [p.get("type") for p in parts if isinstance(p, dict)] diff --git a/studio/backend/tests/test_native_context_length.py b/studio/backend/tests/test_native_context_length.py index de1ca0649e..ba2337755d 100644 --- a/studio/backend/tests/test_native_context_length.py +++ b/studio/backend/tests/test_native_context_length.py @@ -334,7 +334,9 @@ class TestPydanticModels: def test_status_response_chat_template_roundtrip(self): """chat_template serializes and validates as part of status.""" resp = InferenceStatusResponse(chat_template = "{{ messages }}") - roundtripped = InferenceStatusResponse.model_validate_json(resp.model_dump_json()) + roundtripped = InferenceStatusResponse.model_validate_json( + resp.model_dump_json() + ) assert roundtripped.chat_template == "{{ messages }}" def test_roundtrip_preserves_value(self): @@ -402,7 +404,9 @@ class TestRouteCompleteness: def test_gguf_load_responses_have_field(self): """Every GGUF LoadResponse (is_gguf = True) includes native_context_length.""" blocks = self._find_construction_blocks("LoadResponse") - gguf_blocks = [b for b in blocks if "is_gguf = True" in b or "is_gguf=True" in b] + gguf_blocks = [ + b for b in blocks if "is_gguf = True" in b or "is_gguf=True" in b + ] assert ( len(gguf_blocks) >= 2 ), f"Expected at least 2 GGUF LoadResponse blocks, found {len(gguf_blocks)}" @@ -414,7 +418,9 @@ class TestRouteCompleteness: def test_non_gguf_load_responses_omit_field(self): """Non-GGUF LoadResponse blocks do not set native_context_length (defaults to None).""" blocks = self._find_construction_blocks("LoadResponse") - non_gguf = [b for b in blocks if "is_gguf = True" not in b and "is_gguf=True" not in b] + non_gguf = [ + b for b in blocks if "is_gguf = True" not in b and "is_gguf=True" not in b + ] # Non-GGUF paths shouldn't reference native_context_length # (Pydantic defaults it to None, so omitting it is correct). for block in non_gguf: @@ -425,7 +431,9 @@ class TestRouteCompleteness: def test_non_gguf_load_responses_set_runtime_context_length(self): """Non-GGUF LoadResponse blocks report runtime context_length.""" blocks = self._find_construction_blocks("LoadResponse") - non_gguf = [b for b in blocks if "is_gguf = True" not in b and "is_gguf=True" not in b] + non_gguf = [ + b for b in blocks if "is_gguf = True" not in b and "is_gguf=True" not in b + ] assert non_gguf, "Expected at least one non-GGUF LoadResponse block" for block in non_gguf: assert ( @@ -440,9 +448,7 @@ class TestRouteCompleteness: if "llama_backend" in block and "native_context_length" in block: found = True break - assert ( - found - ), "No InferenceStatusResponse block with llama_backend has native_context_length" + assert found, "No InferenceStatusResponse block with llama_backend has native_context_length" def test_non_gguf_status_path_reports_runtime_context_length(self): """Non-GGUF InferenceStatusResponse reports context_length from model_info.""" diff --git a/studio/backend/tests/test_offline_gguf_cache_fallback.py b/studio/backend/tests/test_offline_gguf_cache_fallback.py index aab58adfff..17916155d2 100644 --- a/studio/backend/tests/test_offline_gguf_cache_fallback.py +++ b/studio/backend/tests/test_offline_gguf_cache_fallback.py @@ -134,7 +134,8 @@ def _siblings(items: dict[str, int]): """Mock ``hf_model_info(...).siblings`` payload.""" return _types.SimpleNamespace( siblings = [ - _types.SimpleNamespace(rfilename = name, size = size) for name, size in items.items() + _types.SimpleNamespace(rfilename = name, size = size) + for name, size in items.items() ], ) @@ -158,8 +159,12 @@ class TestIterHfCacheSnapshots: assert list(_iter_hf_cache_snapshots("unsloth/bare")) == [] def test_yields_newest_first(self, hf_cache): - old = _build_cache(hf_cache, "unsloth/multi", {"x.gguf": 1}, snapshot_sha = "a" * 40) - new = _build_cache(hf_cache, "unsloth/multi", {"y.gguf": 1}, snapshot_sha = "b" * 40) + old = _build_cache( + hf_cache, "unsloth/multi", {"x.gguf": 1}, snapshot_sha = "a" * 40 + ) + new = _build_cache( + hf_cache, "unsloth/multi", {"y.gguf": 1}, snapshot_sha = "b" * 40 + ) os.utime(old, (1000, 1000)) os.utime(new, (2000, 2000)) out = list(_iter_hf_cache_snapshots("unsloth/multi")) @@ -198,7 +203,9 @@ class TestListGgufVariantsFromCache: class TestListGgufVariantsOffline: - def test_offline_env_short_circuits_api(self, hf_cache, clean_offline_env, monkeypatch): + def test_offline_env_short_circuits_api( + self, hf_cache, clean_offline_env, monkeypatch + ): _build_cache(hf_cache, "unsloth/a", {"a-UD-Q4_K_XL.gguf": 1}) monkeypatch.setenv("HF_HUB_OFFLINE", "1") @@ -274,7 +281,9 @@ class TestDetectGgufFromCache: class TestDetectGgufModelRemoteOffline: - def test_offline_env_short_circuits_retries(self, hf_cache, clean_offline_env, monkeypatch): + def test_offline_env_short_circuits_retries( + self, hf_cache, clean_offline_env, monkeypatch + ): _build_cache(hf_cache, "unsloth/a", {"a-Q4_K_M.gguf": 1}) monkeypatch.setenv("HF_HUB_OFFLINE", "1") @@ -298,7 +307,9 @@ class TestDetectGgufModelRemoteOffline: out = detect_gguf_model_remote("unsloth/a") assert out == "a-Q4_K_M.gguf" - def test_repository_not_found_does_not_consult_cache(self, hf_cache, clean_offline_env): + def test_repository_not_found_does_not_consult_cache( + self, hf_cache, clean_offline_env + ): # Cache has a file but the API says the repo is gone. _build_cache(hf_cache, "unsloth/a", {"a-Q4_K_M.gguf": 1}) @@ -393,7 +404,9 @@ class TestHfOfflineIfDnsDead: assert did_set is False assert "HF_HUB_OFFLINE" not in os.environ - def test_user_set_hf_hub_offline_is_preserved(self, dns, clean_offline_env, monkeypatch): + def test_user_set_hf_hub_offline_is_preserved( + self, dns, clean_offline_env, monkeypatch + ): # User explicitly set offline before launching Studio. monkeypatch.setenv("HF_HUB_OFFLINE", "1") dns.fail() @@ -403,7 +416,9 @@ class TestHfOfflineIfDnsDead: # Helper must not pop a variable it did not set. assert os.environ.get("HF_HUB_OFFLINE") == "1" - def test_user_set_transformers_offline_is_preserved(self, dns, clean_offline_env, monkeypatch): + def test_user_set_transformers_offline_is_preserved( + self, dns, clean_offline_env, monkeypatch + ): monkeypatch.setenv("TRANSFORMERS_OFFLINE", "1") dns.fail() with _hf_offline_if_dns_dead(): @@ -447,7 +462,9 @@ class TestDownloadMmprojOfflineCacheFallback: """``_download_mmproj`` must resolve cached mmproj GGUFs offline, like ``_download_gguf``; else the offline vision load returns None despite a cache hit.""" - def test_cache_lookup_returns_cached_mmproj_when_list_repo_files_fails(self, hf_cache): + def test_cache_lookup_returns_cached_mmproj_when_list_repo_files_fails( + self, hf_cache + ): _build_cache( hf_cache, "unsloth/vision-GGUF", @@ -606,7 +623,9 @@ class TestListGgufVariantsPermanentErrors: list_gguf_variants("u/gated-gguf") assert type(exc_info.value).__name__ == "GatedRepoError" - def test_transient_error_still_falls_back_to_cache(self, hf_cache, clean_offline_env): + def test_transient_error_still_falls_back_to_cache( + self, hf_cache, clean_offline_env + ): from utils.models.model_config import list_gguf_variants _build_cache(hf_cache, "u/transient-gguf", {"foo-Q4_K_M.gguf": 1}) diff --git a/studio/backend/tests/test_offline_inference_parent.py b/studio/backend/tests/test_offline_inference_parent.py index 71331220d6..a116671ae1 100644 --- a/studio/backend/tests/test_offline_inference_parent.py +++ b/studio/backend/tests/test_offline_inference_parent.py @@ -115,7 +115,9 @@ class TestTransformersVersionOfflineShortCircuits: with patch("urllib.request.urlopen", boom): assert _check_tokenizer_config_needs_v5(unique) is False - def test_config_550_skips_urllib_when_offline(self, monkeypatch, clean_offline_env, tmp_path): + def test_config_550_skips_urllib_when_offline( + self, monkeypatch, clean_offline_env, tmp_path + ): monkeypatch.setenv("HF_HUB_OFFLINE", "1") unique = f"unsloth/never-cached-{tmp_path.name}-cfg" diff --git a/studio/backend/tests/test_openai_citation_markers_edge.py b/studio/backend/tests/test_openai_citation_markers_edge.py index f44975ce33..d8bf0ad5c1 100644 --- a/studio/backend/tests/test_openai_citation_markers_edge.py +++ b/studio/backend/tests/test_openai_citation_markers_edge.py @@ -319,7 +319,9 @@ def test_split_helper_buffers_only_after_last_open_byte(): assert head == f"pre {complete} mid " assert tail == partial # Head, once rewritten, drops every private-use byte. - rewritten = _replace_openai_citation_markers(head, [{"source_id": "done", "url": "https://d"}]) + rewritten = _replace_openai_citation_markers( + head, [{"source_id": "done", "url": "https://d"}] + ) assert rewritten == "pre [[1]](https://d) mid " diff --git a/studio/backend/tests/test_openai_code_execution.py b/studio/backend/tests/test_openai_code_execution.py index 63e94613ed..8488979de3 100644 --- a/studio/backend/tests/test_openai_code_execution.py +++ b/studio/backend/tests/test_openai_code_execution.py @@ -250,7 +250,11 @@ def test_shell_call_emits_tool_start_and_end(monkeypatch): assert starts[0]["tool_call_id"] == "scall_1" # `_server_tool: True` marks a synthetic builtin so the frontend can tell # hosted tools from user-declared functions on history replay. - assert starts[0]["arguments"] == {"kind": "bash", "command": "ls -la", "_server_tool": True} + assert starts[0]["arguments"] == { + "kind": "bash", + "command": "ls -la", + "_server_tool": True, + } assert ends[0]["tool_call_id"] == "scall_1" assert "total 24" in ends[0]["result"] @@ -509,5 +513,7 @@ def test_expired_container_retries_only_once(monkeypatch): # Exactly two calls (first + one retry); a third would be a loop. assert call_count["n"] == 2 # The second failure surfaces normally as an error SSE line. - error_lines = [line for line in lines if '"error"' in line and "_toolEvent" not in line] + error_lines = [ + line for line in lines if '"error"' in line and "_toolEvent" not in line + ] assert len(error_lines) >= 1 diff --git a/studio/backend/tests/test_openai_container_crud.py b/studio/backend/tests/test_openai_container_crud.py index e0604527fd..33280fd95a 100644 --- a/studio/backend/tests/test_openai_container_crud.py +++ b/studio/backend/tests/test_openai_container_crud.py @@ -79,7 +79,9 @@ def test_create_sends_openai_beta_header(monkeypatch): return httpx.Response(200, json = {"id": "cntr_new", "name": "analysis"}) _mock_http_client(monkeypatch, handler) - result = _drive(_make_client().create_openai_container(name = "analysis", ttl_minutes = 30)) + result = _drive( + _make_client().create_openai_container(name = "analysis", ttl_minutes = 30) + ) assert result == {"id": "cntr_new", "name": "analysis"} assert seen["headers"].get("openai-beta") == "containers=v1" diff --git a/studio/backend/tests/test_openai_responses_translation.py b/studio/backend/tests/test_openai_responses_translation.py index f7d7e83a43..eb09243a39 100644 --- a/studio/backend/tests/test_openai_responses_translation.py +++ b/studio/backend/tests/test_openai_responses_translation.py @@ -494,7 +494,8 @@ def test_responses_response_incomplete_maps_to_length_finish_reason(monkeypatch) finish_reasons = [ json.loads(line[len("data:") :].strip())["choices"][0]["finish_reason"] for line in lines - if line.startswith("data:") and line[len("data:") :].strip() not in ("", "[DONE]") + if line.startswith("data:") + and line[len("data:") :].strip() not in ("", "[DONE]") ] assert "length" in finish_reasons @@ -726,7 +727,8 @@ def test_responses_reasoning_summary_wrapped_in_think_tags(monkeypatch): data_lines = [ line[len("data:") :].strip() for line in lines - if line.startswith("data:") and line[len("data:") :].strip() not in ("", "[DONE]") + if line.startswith("data:") + and line[len("data:") :].strip() not in ("", "[DONE]") ] payloads = [json.loads(raw) for raw in data_lines] combined = "".join( diff --git a/studio/backend/tests/test_openai_tool_passthrough.py b/studio/backend/tests/test_openai_tool_passthrough.py index b54f4c130d..5c408f7e01 100644 --- a/studio/backend/tests/test_openai_tool_passthrough.py +++ b/studio/backend/tests/test_openai_tool_passthrough.py @@ -267,7 +267,10 @@ class TestChatCompletionRequestToolFields: assert self._make(stop = "\nUser:").stop == "\nUser:" def test_stop_list(self): - assert self._make(stop = ["\nUser:", "\nAssistant:"]).stop == ["\nUser:", "\nAssistant:"] + assert self._make(stop = ["\nUser:", "\nAssistant:"]).stop == [ + "\nUser:", + "\nAssistant:", + ] def test_tools_default_none(self): req = self._make() @@ -307,7 +310,9 @@ class TestChatCompletionRequestToolFields: req = self._make() assert req.stream is False - def test_post_without_stream_field_decodes_to_stream_false_over_http(self, monkeypatch): + def test_post_without_stream_field_decodes_to_stream_false_over_http( + self, monkeypatch + ): # Wire-level guard: a POST body omitting `stream` must deserialise to # stream=False and return application/json, never text/event-stream. # Mounts the real router to catch middleware/aliasing regressions; @@ -357,9 +362,13 @@ class TestChatCompletionRequestToolFields: from auth.authentication import get_current_subject from utils.api_errors import install_api_error_handlers - monkeypatch.setattr(inference_route, "get_llama_cpp_backend", lambda: llama_backend) + monkeypatch.setattr( + inference_route, "get_llama_cpp_backend", lambda: llama_backend + ) if inference_backend is not None: - monkeypatch.setattr(inference_route, "get_inference_backend", lambda: inference_backend) + monkeypatch.setattr( + inference_route, "get_inference_backend", lambda: inference_backend + ) app = FastAPI() app.include_router(inference_route.router, prefix = "/v1") @@ -514,7 +523,9 @@ class TestChatCompletionRequestToolFields: ) self._assert_unsupported_n(resp) - def test_confirm_tool_calls_requires_streaming_for_safetensors_tools(self, monkeypatch): + def test_confirm_tool_calls_requires_streaming_for_safetensors_tools( + self, monkeypatch + ): import routes.inference as inference_route class _NoGGUFBackend: @@ -611,7 +622,9 @@ class TestAnthropicToolChoiceToOpenAI: assert anthropic_tool_choice_to_openai({"type": "none"}) == "none" def test_tool_named(self): - result = anthropic_tool_choice_to_openai({"type": "tool", "name": "get_weather"}) + result = anthropic_tool_choice_to_openai( + {"type": "tool", "name": "get_weather"} + ) assert result == {"type": "function", "function": {"name": "get_weather"}} def test_tool_missing_name_returns_none(self): @@ -856,11 +869,16 @@ class TestOpenAICompatibilityHelpers: usage = {"prompt_tokens": 3, "completion_tokens": 2, "total_tokens": 5} payload = SimpleNamespace(stream_options = None) assert ( - _openai_stream_usage_chunk(payload, "chatcmpl-test", 123, "model", usage, None) is None + _openai_stream_usage_chunk( + payload, "chatcmpl-test", 123, "model", usage, None + ) + is None ) payload.stream_options = {"include_usage": True} - line = _openai_stream_usage_chunk(payload, "chatcmpl-test", 123, "model", usage, None) + line = _openai_stream_usage_chunk( + payload, "chatcmpl-test", 123, "model", usage, None + ) assert line is not None assert '"choices":[]' in line assert '"usage"' in line @@ -895,7 +913,9 @@ class TestOpenAICompatibilityHelpers: if message.role == "developer": message.role = "system" - system_prompt, chat_messages, image_b64 = _extract_content_parts(payload.messages) + system_prompt, chat_messages, image_b64 = _extract_content_parts( + payload.messages + ) assert system_prompt == "original system\n\ndeveloper rules" assert chat_messages == [{"role": "user", "content": "hi"}] @@ -930,11 +950,15 @@ class TestFriendlyErrorHttpx: def test_non_httpx_unchanged(self): # Non-httpx exceptions still fall through to the substring heuristics # — a context-size message must still produce "Message too long". - ctx_msg = "request (4096 tokens) exceeds the available context size (2048 tokens)" + ctx_msg = ( + "request (4096 tokens) exceeds the available context size (2048 tokens)" + ) assert "Message too long" in _friendly_error(ValueError(ctx_msg)) def test_generic_exception_returns_generic_message(self): - assert _friendly_error(RuntimeError("unrelated")) == "An internal error occurred" + assert ( + _friendly_error(RuntimeError("unrelated")) == "An internal error occurred" + ) from routes.inference import ( # noqa: E402 @@ -952,7 +976,10 @@ class TestDropEmptyAssistantSentinels: {"role": "user", "content": "again"}, ] out = _drop_empty_assistant_sentinels(msgs) - assert out == [{"role": "user", "content": "hi"}, {"role": "user", "content": "again"}] + assert out == [ + {"role": "user", "content": "hi"}, + {"role": "user", "content": "again"}, + ] def test_drops_assistant_with_no_content_key(self): # exclude_none=True strips the content key entirely; filter must catch it. @@ -962,7 +989,10 @@ class TestDropEmptyAssistantSentinels: {"role": "user", "content": "ok"}, ] out = _drop_empty_assistant_sentinels(msgs) - assert out == [{"role": "user", "content": "hi"}, {"role": "user", "content": "ok"}] + assert out == [ + {"role": "user", "content": "hi"}, + {"role": "user", "content": "ok"}, + ] def test_preserves_assistant_with_text(self): msgs = [ @@ -1062,10 +1092,16 @@ class TestGgufVisionMessages: messages, has_image = _openai_messages_for_gguf_chat(req, is_vision = True) assert has_image is True - assert messages[0]["content"][0] == {"type": "text", "text": "describe image one"} + assert messages[0]["content"][0] == { + "type": "text", + "text": "describe image one", + } assert messages[0]["content"][1]["type"] == "image_url" assert len(messages[0]["content"]) == 2 - assert messages[2]["content"][0] == {"type": "text", "text": "describe image two"} + assert messages[2]["content"][0] == { + "type": "text", + "text": "describe image two", + } assert messages[2]["content"][1]["type"] == "image_url" assert len(messages[2]["content"]) == 2 assert isinstance(messages[1]["content"], str) @@ -1088,9 +1124,14 @@ class TestGgufVisionMessages: messages, has_image = _openai_messages_for_gguf_chat(req, is_vision = True) assert has_image is True - assert messages[0]["content"][0] == {"type": "text", "text": "describe this image"} + assert messages[0]["content"][0] == { + "type": "text", + "text": "describe this image", + } assert messages[0]["content"][1]["type"] == "image_url" - assert messages[0]["content"][1]["image_url"]["url"].startswith("data:image/png;base64,") + assert messages[0]["content"][1]["image_url"]["url"].startswith( + "data:image/png;base64," + ) def test_rejects_image_parts_for_text_only_gguf(self): req = ChatCompletionRequest( @@ -1156,7 +1197,9 @@ class TestGgufVisionMessages: {"role": "user", "content": "now"}, ] - updated = _set_or_prepend_system_message(messages, "Mid instructions.\n\nUse tools.") + updated = _set_or_prepend_system_message( + messages, "Mid instructions.\n\nUse tools." + ) assert [m["role"] for m in updated] == ["system", "user", "user"] assert updated[0]["content"] == "Mid instructions.\n\nUse tools." @@ -1216,7 +1259,9 @@ class TestGgufVisionToolRouting: { "type": "image_url", "image_url": { - "url": (f"data:image/png;base64,{TestGgufVisionMessages._PNG_B64}"), + "url": ( + f"data:image/png;base64,{TestGgufVisionMessages._PNG_B64}" + ), }, }, ], @@ -1225,7 +1270,9 @@ class TestGgufVisionToolRouting: ) response = self._drive( - openai_chat_completions(payload, request = self._Request(), current_subject = "test") + openai_chat_completions( + payload, request = self._Request(), current_subject = "test" + ) ) self._consume_response(response) @@ -1268,7 +1315,9 @@ class TestGgufVisionToolRouting: ) response = self._drive( - openai_chat_completions(payload, request = self._Request(), current_subject = "test") + openai_chat_completions( + payload, request = self._Request(), current_subject = "test" + ) ) self._consume_response(response) @@ -1323,7 +1372,11 @@ class TestGgufVisionToolRouting: yield "done" yield { "type": "metadata", - "usage": {"prompt_tokens": 3, "completion_tokens": 1, "total_tokens": 4}, + "usage": { + "prompt_tokens": 3, + "completion_tokens": 1, + "total_tokens": 4, + }, "finish_reason": "stop", } @@ -1346,7 +1399,9 @@ class TestGgufVisionToolRouting: ) self._drive( - openai_chat_completions(payload, request = self._Request(), current_subject = "test") + openai_chat_completions( + payload, request = self._Request(), current_subject = "test" + ) ) assert captured["messages"] == [ @@ -1361,7 +1416,9 @@ class TestGgufVisionToolRouting: (-1, [-1, -1, -1]), ], ) - def test_gguf_n_choices_vary_explicit_non_negative_seed(self, monkeypatch, seed, expected): + def test_gguf_n_choices_vary_explicit_non_negative_seed( + self, monkeypatch, seed, expected + ): import routes.inference as inf_mod seen_seeds = [] @@ -1396,7 +1453,9 @@ class TestGgufVisionToolRouting: ) response = self._drive( - openai_chat_completions(payload, request = self._Request(), current_subject = "test") + openai_chat_completions( + payload, request = self._Request(), current_subject = "test" + ) ) body = json.loads(response.body) diff --git a/studio/backend/tests/test_pricing.py b/studio/backend/tests/test_pricing.py index 8cd7796f14..198b9d453f 100644 --- a/studio/backend/tests/test_pricing.py +++ b/studio/backend/tests/test_pricing.py @@ -245,7 +245,9 @@ def test_openai_cache_read_subtracted_from_input_at_discount(): ) # 20k charged at full price, 80k charged at 0.1x assert _isclose(out["input_usd"], 20_000 / 1_000_000.0 * base) - assert _isclose(out["cache_read_usd"], 80_000 / 1_000_000.0 * base * OPENAI_CACHE_READ_MULT) + assert _isclose( + out["cache_read_usd"], 80_000 / 1_000_000.0 * base * OPENAI_CACHE_READ_MULT + ) def test_openai_billable_input_tokens_does_not_double_count_cache_read(): @@ -392,7 +394,9 @@ def test_openai_web_search_charged_per_thousand(): "openai_tool_use": {"web_search_requests": 250}, }, ) - assert _isclose(out["server_tools_usd"], 250 / 1_000.0 * OPENAI_WEB_SEARCH_USD_PER_1K) + assert _isclose( + out["server_tools_usd"], 250 / 1_000.0 * OPENAI_WEB_SEARCH_USD_PER_1K + ) assert _isclose(out["total_usd"], 250 / 1_000.0 * OPENAI_WEB_SEARCH_USD_PER_1K) @@ -426,7 +430,8 @@ def test_openai_tool_surcharges_added_to_total(): expected_input = 100_000 / 1_000_000.0 * 5.0 expected_output = 5_000 / 1_000_000.0 * 30.0 expected_tools = ( - 3 / 1_000.0 * OPENAI_WEB_SEARCH_USD_PER_1K + 0.25 * OPENAI_CONTAINER_USD_PER_HOUR + 3 / 1_000.0 * OPENAI_WEB_SEARCH_USD_PER_1K + + 0.25 * OPENAI_CONTAINER_USD_PER_HOUR ) assert _isclose( out["total_usd"], @@ -599,7 +604,10 @@ def test_openai_chat_style_envelope_reads_cache_from_prompt_tokens_details(): ) # Both envelopes must price identically. assert _isclose(chat_style["input_usd"], raw["input_usd"]), (chat_style, raw) - assert _isclose(chat_style["cache_read_usd"], raw["cache_read_usd"]), (chat_style, raw) + assert _isclose(chat_style["cache_read_usd"], raw["cache_read_usd"]), ( + chat_style, + raw, + ) # 80k at 0.1x base, 20k at full. assert _isclose( chat_style["cache_read_usd"], diff --git a/studio/backend/tests/test_pricing_edge.py b/studio/backend/tests/test_pricing_edge.py index 1fcc428f90..6c7f4038f1 100644 --- a/studio/backend/tests/test_pricing_edge.py +++ b/studio/backend/tests/test_pricing_edge.py @@ -191,7 +191,9 @@ def test_anthropic_chat_cache_read_exceeds_prompt_no_negative_billable(): assert out["billable_input_tokens"] == 500 # 0 uncached + 500 cache_read # cache_read still priced at the discount rate. base = ANTHROPIC_PRICING["claude-opus-4-7"]["input_per_mtok"] - assert _isclose(out["cache_read_usd"], 500 / 1_000_000.0 * base * ANTHROPIC_CACHE_READ_MULT) + assert _isclose( + out["cache_read_usd"], 500 / 1_000_000.0 * base * ANTHROPIC_CACHE_READ_MULT + ) def test_openai_raw_cached_tokens_exceeds_input_clamp_non_cached(): @@ -208,7 +210,9 @@ def test_openai_raw_cached_tokens_exceeds_input_clamp_non_cached(): ) assert out["input_usd"] == 0.0 # Cache read still priced (the 0.1x bucket). - assert _isclose(out["cache_read_usd"], 500 / 1_000_000.0 * base * OPENAI_CACHE_READ_MULT) + assert _isclose( + out["cache_read_usd"], 500 / 1_000_000.0 * base * OPENAI_CACHE_READ_MULT + ) # ── long-context tier crosses on billable, including cache_creation ── diff --git a/studio/backend/tests/test_providers_api.py b/studio/backend/tests/test_providers_api.py index 5e24ed752d..dff539e8e6 100644 --- a/studio/backend/tests/test_providers_api.py +++ b/studio/backend/tests/test_providers_api.py @@ -211,7 +211,9 @@ class TestAuth: json = {"username": USERNAME, "password": PASSWORD}, timeout = 10, ) - assert resp.status_code == 200, f"Login failed ({resp.status_code}): {resp.text}" + assert ( + resp.status_code == 200 + ), f"Login failed ({resp.status_code}): {resp.text}" body = resp.json() assert body.get("access_token"), "access_token is missing or empty" assert body.get("token_type") == "bearer" @@ -221,7 +223,9 @@ class TestAuth: class TestPublicKey: - def test_public_key_is_valid_pem(self, auth_headers: dict[str, str], public_key_pem: str): + def test_public_key_is_valid_pem( + self, auth_headers: dict[str, str], public_key_pem: str + ): """GET /api/providers/public-key returns an importable RSA PEM key.""" pem_bytes = public_key_pem.encode("utf-8") key = serialization.load_pem_public_key(pem_bytes) @@ -243,7 +247,9 @@ class TestRegistry: ) assert resp.status_code == 200, f"Registry failed: {resp.text}" providers = resp.json() - assert len(providers) == 9, f"Expected 9 providers, got {len(providers)}: {providers}" + assert ( + len(providers) == 9 + ), f"Expected 9 providers, got {len(providers)}: {providers}" print(f"\n {'Provider':<12} {'Base URL'}") print(f" {'-'*12} {'-'*45}") for p in providers: @@ -263,7 +269,9 @@ class TestRegistry: def test_registry_entries_have_required_fields(self, auth_headers: dict[str, str]): """Each registry entry has provider_type, display_name, base_url, default_models.""" - resp = requests.get(_url("/api/providers/registry"), headers = auth_headers, timeout = 10) + resp = requests.get( + _url("/api/providers/registry"), headers = auth_headers, timeout = 10 + ) assert resp.status_code == 200 for entry in resp.json(): for field in ( @@ -298,7 +306,9 @@ class TestProviderCRUD: json = {"provider_type": "openai", "display_name": "Test OpenAI (pytest)"}, timeout = 10, ) - assert resp.status_code == 201, f"Create failed ({resp.status_code}): {resp.text}" + assert ( + resp.status_code == 201 + ), f"Create failed ({resp.status_code}): {resp.text}" body = resp.json() assert body.get("id"), "No id in response" assert body["provider_type"] == "openai" @@ -309,7 +319,9 @@ class TestProviderCRUD: def test_list_includes_created(self, auth_headers: dict[str, str]): """GET /api/providers/ includes the newly created config.""" - assert TestProviderCRUD._created_id, "No created_id (run test_create_provider first)" + assert ( + TestProviderCRUD._created_id + ), "No created_id (run test_create_provider first)" resp = requests.get(_url("/api/providers/"), headers = auth_headers, timeout = 10) assert resp.status_code == 200 ids = [p["id"] for p in resp.json()] @@ -328,7 +340,9 @@ class TestProviderCRUD: json = {"display_name": new_name}, timeout = 10, ) - assert resp.status_code == 200, f"Update failed ({resp.status_code}): {resp.text}" + assert ( + resp.status_code == 200 + ), f"Update failed ({resp.status_code}): {resp.text}" assert resp.json()["display_name"] == new_name print(f"\n updated display_name to '{new_name}'") @@ -340,10 +354,14 @@ class TestProviderCRUD: headers = auth_headers, timeout = 10, ) - assert resp.status_code == 204, f"Delete failed ({resp.status_code}): {resp.text}" + assert ( + resp.status_code == 204 + ), f"Delete failed ({resp.status_code}): {resp.text}" # Confirm gone - list_resp = requests.get(_url("/api/providers/"), headers = auth_headers, timeout = 10) + list_resp = requests.get( + _url("/api/providers/"), headers = auth_headers, timeout = 10 + ) ids = [p["id"] for p in list_resp.json()] assert TestProviderCRUD._created_id not in ids, "Deleted provider still in list" print(f"\n deleted id={TestProviderCRUD._created_id} confirmed gone") @@ -391,7 +409,9 @@ class TestProviderInference: json = {"provider_type": provider_type, "encrypted_api_key": encrypted}, timeout = 30, ) - assert resp.status_code == 200, f"Request failed ({resp.status_code}): {resp.text}" + assert ( + resp.status_code == 200 + ), f"Request failed ({resp.status_code}): {resp.text}" body = resp.json() assert ( body["success"] is True @@ -415,7 +435,9 @@ class TestProviderInference: json = {"provider_type": provider_type, "encrypted_api_key": encrypted}, timeout = 30, ) - assert resp.status_code == 200, f"Request failed ({resp.status_code}): {resp.text}" + assert ( + resp.status_code == 200 + ), f"Request failed ({resp.status_code}): {resp.text}" models = resp.json() assert isinstance(models, list), f"Expected list, got {type(models)}" assert len(models) > 0, f"No models returned for {provider_type}" @@ -462,7 +484,9 @@ class TestProviderInference: # ── TestVisionInference ───────────────────────────────────────────── # Sloth photo for testing vision routing across providers -_VISION_IMAGE_URL = "https://www.travelexcellence.com/images/where-to-see-sloths-in-costa-rica.jpg" +_VISION_IMAGE_URL = ( + "https://www.travelexcellence.com/images/where-to-see-sloths-in-costa-rica.jpg" +) _VISION_PARAMS = [ pytest.param( @@ -562,6 +586,8 @@ class TestLocalInferenceUnaffected: f"This likely means the provider fields broke the base request schema." ) status_label = ( - "local model responded" if resp.status_code == 200 else "no model loaded (expected)" + "local model responded" + if resp.status_code == 200 + else "no model loaded (expected)" ) print(f"\n status={resp.status_code} ({status_label}) — local path unaffected") diff --git a/studio/backend/tests/test_rag_captioning.py b/studio/backend/tests/test_rag_captioning.py index 5d83a7d38d..b124c691f1 100644 --- a/studio/backend/tests/test_rag_captioning.py +++ b/studio/backend/tests/test_rag_captioning.py @@ -21,17 +21,28 @@ def test_caption_images_disabled_by_default(monkeypatch): def test_caption_images_groups_by_page(monkeypatch): monkeypatch.setattr(captioner.config, "CAPTION_IMAGES", True) monkeypatch.setattr(captioner.config, "CAPTION_MAX_IMAGES", 8) - monkeypatch.setattr(captioner, "_caption_one", lambda base, model, b, t: "a chart of results") - out = captioner.caption_images([_img(1), _img(1), _img(3)], endpoint = ("http://x", "local")) - assert out == {1: ["a chart of results", "a chart of results"], 3: ["a chart of results"]} + monkeypatch.setattr( + captioner, "_caption_one", lambda base, model, b, t: "a chart of results" + ) + out = captioner.caption_images( + [_img(1), _img(1), _img(3)], endpoint = ("http://x", "local") + ) + assert out == { + 1: ["a chart of results", "a chart of results"], + 3: ["a chart of results"], + } def test_caption_images_respects_cap(monkeypatch): monkeypatch.setattr(captioner.config, "CAPTION_IMAGES", True) monkeypatch.setattr(captioner.config, "CAPTION_MAX_IMAGES", 2) calls = [] - monkeypatch.setattr(captioner, "_caption_one", lambda *a: (calls.append(1) or "cap")) - captioner.caption_images([_img(i) for i in range(5)], endpoint = ("http://x", "local")) + monkeypatch.setattr( + captioner, "_caption_one", lambda *a: (calls.append(1) or "cap") + ) + captioner.caption_images( + [_img(i) for i in range(5)], endpoint = ("http://x", "local") + ) assert len(calls) == 2 diff --git a/studio/backend/tests/test_rag_chunking.py b/studio/backend/tests/test_rag_chunking.py index 3d3c9eedd8..e562a9a59f 100644 --- a/studio/backend/tests/test_rag_chunking.py +++ b/studio/backend/tests/test_rag_chunking.py @@ -27,12 +27,16 @@ def test_chunk_never_exceeds_max_with_overlap_carry(): """Overlap carry is trimmed so no chunk exceeds max_tokens (else the embedder overflows).""" s1 = " ".join("a" for _ in range(10)) s2 = " ".join("b" for _ in range(95)) # near max - chunks = chunk_pages([_page(f"{s1}. {s2}")], max_tokens = 100, overlap = 24, count = WORDS) + chunks = chunk_pages( + [_page(f"{s1}. {s2}")], max_tokens = 100, overlap = 24, count = WORDS + ) assert all(c.token_count <= 100 for c in chunks), [c.token_count for c in chunks] def test_chunk_indices_are_sequential(): - chunks = chunk_pages([_page("alpha. " * 200)], max_tokens = 32, overlap = 0, count = WORDS) + chunks = chunk_pages( + [_page("alpha. " * 200)], max_tokens = 32, overlap = 0, count = WORDS + ) assert [c.chunk_index for c in chunks] == list(range(len(chunks))) diff --git a/studio/backend/tests/test_rag_embed_llama_server.py b/studio/backend/tests/test_rag_embed_llama_server.py index 8321068afd..6c63c0890f 100644 --- a/studio/backend/tests/test_rag_embed_llama_server.py +++ b/studio/backend/tests/test_rag_embed_llama_server.py @@ -52,8 +52,12 @@ def _mock_auto(monkeypatch, *, gpus, binary): from core.inference.llama_cpp import LlamaCppBackend monkeypatch.setattr(config, "EMBED_BACKEND", "auto") - monkeypatch.setattr(LlamaCppBackend, "_get_gpu_free_memory", staticmethod(lambda: gpus)) - monkeypatch.setattr(LlamaCppBackend, "_find_llama_server_binary", staticmethod(lambda: binary)) + monkeypatch.setattr( + LlamaCppBackend, "_get_gpu_free_memory", staticmethod(lambda: gpus) + ) + monkeypatch.setattr( + LlamaCppBackend, "_find_llama_server_binary", staticmethod(lambda: binary) + ) def _stub_st_load(monkeypatch): @@ -117,7 +121,9 @@ def test_llama_backend_imports_no_torch(): "RAG_EMBED_BACKEND": "llama-server", "PYTHONPATH": str(backend_dir), } - proc = subprocess.run([sys.executable, "-c", code], capture_output = True, text = True, env = env) + proc = subprocess.run( + [sys.executable, "-c", code], capture_output = True, text = True, env = env + ) assert proc.returncode == 0, proc.stderr assert "OK" in proc.stdout @@ -163,16 +169,22 @@ def test_use_gpu_explicit_modes(monkeypatch): def test_use_gpu_auto_follows_probe(monkeypatch): b = LlamaServerBackend() monkeypatch.setattr(config, "EMBED_DEVICE", "auto") - monkeypatch.setattr(LlamaServerBackend, "_gpu_available", staticmethod(lambda: True)) + monkeypatch.setattr( + LlamaServerBackend, "_gpu_available", staticmethod(lambda: True) + ) assert b._use_gpu() is True - monkeypatch.setattr(LlamaServerBackend, "_gpu_available", staticmethod(lambda: False)) + monkeypatch.setattr( + LlamaServerBackend, "_gpu_available", staticmethod(lambda: False) + ) assert b._use_gpu() is False def test_use_gpu_sticky_cpu_fallback(monkeypatch): b = LlamaServerBackend() monkeypatch.setattr(config, "EMBED_DEVICE", "auto") - monkeypatch.setattr(LlamaServerBackend, "_gpu_available", staticmethod(lambda: True)) + monkeypatch.setattr( + LlamaServerBackend, "_gpu_available", staticmethod(lambda: True) + ) b._force_cpu = True # a prior GPU start failed assert b._use_gpu() is False @@ -183,11 +195,17 @@ def test_gpu_available_reuses_studio_probe(monkeypatch): monkeypatch.setattr(uh, "is_apple_silicon", lambda: False) # Ample free VRAM -> GPU; nearly full -> CPU; none -> CPU. - monkeypatch.setattr(LlamaCppBackend, "_get_gpu_free_memory", staticmethod(lambda: [(0, 40000)])) + monkeypatch.setattr( + LlamaCppBackend, "_get_gpu_free_memory", staticmethod(lambda: [(0, 40000)]) + ) assert LlamaServerBackend._gpu_available() is True - monkeypatch.setattr(LlamaCppBackend, "_get_gpu_free_memory", staticmethod(lambda: [(0, 100)])) + monkeypatch.setattr( + LlamaCppBackend, "_get_gpu_free_memory", staticmethod(lambda: [(0, 100)]) + ) assert LlamaServerBackend._gpu_available() is False - monkeypatch.setattr(LlamaCppBackend, "_get_gpu_free_memory", staticmethod(lambda: [])) + monkeypatch.setattr( + LlamaCppBackend, "_get_gpu_free_memory", staticmethod(lambda: []) + ) assert LlamaServerBackend._gpu_available() is False @@ -205,9 +223,15 @@ def _patch_spawn_deps( ): # Force CPU so spawn never depends on a host GPU. monkeypatch.setattr(config, "EMBED_DEVICE", "cpu") - monkeypatch.setattr(LlamaServerBackend, "_resolve_binary", lambda self: "/bin/llama-server") - monkeypatch.setattr(LlamaServerBackend, "_resolve_model_path", lambda self: "/m/bge.gguf") - monkeypatch.setattr(LlamaServerBackend, "_find_free_port", staticmethod(lambda: free_port)) + monkeypatch.setattr( + LlamaServerBackend, "_resolve_binary", lambda self: "/bin/llama-server" + ) + monkeypatch.setattr( + LlamaServerBackend, "_resolve_model_path", lambda self: "/m/bge.gguf" + ) + monkeypatch.setattr( + LlamaServerBackend, "_find_free_port", staticmethod(lambda: free_port) + ) monkeypatch.setattr(mod.subprocess, "Popen", lambda *a, **k: proc) @@ -239,7 +263,9 @@ def test_spawn_fails_loud_on_early_exit(monkeypatch): def test_spawn_auto_falls_back_to_cpu_on_gpu_failure(monkeypatch): monkeypatch.setattr(config, "EMBED_DEVICE", "auto") - monkeypatch.setattr(LlamaServerBackend, "_gpu_available", staticmethod(lambda: True)) + monkeypatch.setattr( + LlamaServerBackend, "_gpu_available", staticmethod(lambda: True) + ) b = LlamaServerBackend() calls = [] @@ -311,7 +337,9 @@ def test_encode_empty_returns_zero_rows(monkeypatch): def test_encode_rejects_count_mismatch(monkeypatch): b = LlamaServerBackend() monkeypatch.setattr(b, "_ensure_ready", lambda: None) - monkeypatch.setattr(b, "_post", lambda p, pl: {"data": [{"index": 0, "embedding": [1.0]}]}) + monkeypatch.setattr( + b, "_post", lambda p, pl: {"data": [{"index": 0, "embedding": [1.0]}]} + ) with pytest.raises(RuntimeError, match = "vectors for"): b.encode(["a", "b"], normalize = False) @@ -390,7 +418,9 @@ def test_post_restarts_once_on_connect_error(monkeypatch): b._port = 9000 monkeypatch.setattr(b, "_ensure_ready", lambda: None) restarts = {"n": 0} - monkeypatch.setattr(b, "_restart", lambda: restarts.__setitem__("n", restarts["n"] + 1)) + monkeypatch.setattr( + b, "_restart", lambda: restarts.__setitem__("n", restarts["n"] + 1) + ) attempts = {"n": 0} @@ -423,7 +453,9 @@ def test_post_restarts_once_on_read_timeout(monkeypatch): b._port = 9000 monkeypatch.setattr(b, "_ensure_ready", lambda: None) restarts = {"n": 0} - monkeypatch.setattr(b, "_restart", lambda: restarts.__setitem__("n", restarts["n"] + 1)) + monkeypatch.setattr( + b, "_restart", lambda: restarts.__setitem__("n", restarts["n"] + 1) + ) attempts = {"n": 0} diff --git a/studio/backend/tests/test_rag_embeddings.py b/studio/backend/tests/test_rag_embeddings.py index 28a2f69426..0d22b6cd54 100644 --- a/studio/backend/tests/test_rag_embeddings.py +++ b/studio/backend/tests/test_rag_embeddings.py @@ -147,7 +147,9 @@ def _patch_llama_backend(monkeypatch, *, binary): from core.inference.llama_cpp import LlamaCppBackend from core.rag import embed_llama_server - monkeypatch.setattr(LlamaCppBackend, "_find_llama_server_binary", staticmethod(lambda: binary)) + monkeypatch.setattr( + LlamaCppBackend, "_find_llama_server_binary", staticmethod(lambda: binary) + ) monkeypatch.setattr(embed_llama_server, "LlamaServerBackend", _SentinelLlamaBackend) @@ -189,7 +191,9 @@ class _BoomOnEncodeModel: def test_st_encode_runtime_failure_switches_to_llama(monkeypatch): # encode() blows up mid-run -> switch to llama-server and stay switched. - monkeypatch.setattr(embeddings, "_get", lambda model_name = None: _BoomOnEncodeModel()) + monkeypatch.setattr( + embeddings, "_get", lambda model_name = None: _BoomOnEncodeModel() + ) _patch_llama_backend(monkeypatch, binary = "/fake/llama-server") calls = {} @@ -203,7 +207,9 @@ def test_st_encode_runtime_failure_switches_to_llama(monkeypatch): calls["used"] = True return np.zeros((len(texts), 4), dtype = np.float32) - monkeypatch.setattr(_SentinelLlamaBackend, "encode", _sentinel_encode, raising = False) + monkeypatch.setattr( + _SentinelLlamaBackend, "encode", _sentinel_encode, raising = False + ) embeddings._reset_backend() out = embeddings.encode(["alpha", "beta"]) @@ -215,7 +221,9 @@ def test_st_encode_runtime_failure_switches_to_llama(monkeypatch): def test_st_encode_failure_without_llama_binary_reraises(monkeypatch): # No llama-server binary -> surface the encode error. - monkeypatch.setattr(embeddings, "_get", lambda model_name = None: _BoomOnEncodeModel()) + monkeypatch.setattr( + embeddings, "_get", lambda model_name = None: _BoomOnEncodeModel() + ) _patch_llama_backend(monkeypatch, binary = None) embeddings._reset_backend() with pytest.raises(RuntimeError, match = "CUDA error during encode"): diff --git a/studio/backend/tests/test_rag_ingestion.py b/studio/backend/tests/test_rag_ingestion.py index f0b71bc23b..ce2e418eba 100644 --- a/studio/backend/tests/test_rag_ingestion.py +++ b/studio/backend/tests/test_rag_ingestion.py @@ -39,7 +39,11 @@ def test_ingestion_lifecycle_pending_to_completed(rag_home, stub_embeddings, tmp conn = rag_db.get_connection() try: - assert store.get_document(conn, doc_id)["status"] in {"pending", "running", "completed"} + assert store.get_document(conn, doc_id)["status"] in { + "pending", + "running", + "completed", + } finally: conn.close() @@ -233,7 +237,9 @@ def test_ingestion_rejects_unsupported_ext(rag_home, stub_embeddings, tmp_path): ingestion.start_ingestion(store.kb_scope("K1"), "K1", None, "doc.xyz", path) -def test_ingestion_empty_doc_completes_with_zero_chunks(rag_home, stub_embeddings, tmp_path): +def test_ingestion_empty_doc_completes_with_zero_chunks( + rag_home, stub_embeddings, tmp_path +): path = _write(tmp_path, "empty.txt", " \n ") scope = store.kb_scope("K1") doc_id, job_id = ingestion.start_ingestion(scope, "K1", None, "empty.txt", path) @@ -249,7 +255,9 @@ def test_ingestion_empty_doc_completes_with_zero_chunks(rag_home, stub_embedding reason = "set RAG_REAL_EMBEDDER=1 to run the real sentence-transformers test", ) def test_ingestion_with_real_embedder(rag_home, tmp_path): - path = _write(tmp_path, "doc.txt", "The Kestrel-9 turbine is rated at 9.5 megawatts.") + path = _write( + tmp_path, "doc.txt", "The Kestrel-9 turbine is rated at 9.5 megawatts." + ) scope = store.kb_scope("K1") doc_id, job_id = ingestion.start_ingestion(scope, "K1", None, "doc.txt", path) _drain(job_id) @@ -260,7 +268,9 @@ def test_ingestion_with_real_embedder(rag_home, tmp_path): conn = rag_db.get_connection() try: - hits = retrieval.retrieve_hybrid(conn, scope, "how much power does the turbine make?", k = 5) + hits = retrieval.retrieve_hybrid( + conn, scope, "how much power does the turbine make?", k = 5 + ) assert hits and hits[0].chunk_id == f"{doc_id}:0" finally: conn.close() diff --git a/studio/backend/tests/test_rag_preview.py b/studio/backend/tests/test_rag_preview.py index e7f2a39792..68832fd20e 100644 --- a/studio/backend/tests/test_rag_preview.py +++ b/studio/backend/tests/test_rag_preview.py @@ -112,7 +112,9 @@ def test_preview_routes_and_signed_file(rag_home, stub_embeddings): assert res chunk_id = res[0]["chunkId"] - pt = c.get(f"/api/rag/documents/{doc_id}/preview-target", params = {"chunk_id": chunk_id}).json() + pt = c.get( + f"/api/rag/documents/{doc_id}/preview-target", params = {"chunk_id": chunk_id} + ).json() assert pt["mediaKind"] == "pdf" assert pt["text"] @@ -148,7 +150,9 @@ def test_locator_handles_midword_anchor_and_locates_line(): doc = pymupdf.open() page = doc.new_page() - page.insert_text((72, 200), "alpha beta gamma delta epsilon zeta eta theta", fontsize = 12) + page.insert_text( + (72, 200), "alpha beta gamma delta epsilon zeta eta theta", fontsize = 12 + ) page_text = doc[0].get_text("text") # mirrors what the parser stores start = page_text.index("lpha") end = page_text.index("theta") + 3 @@ -170,5 +174,7 @@ def test_sign_verify_roundtrip(rag_home): tok = rag_routes._sign_document("doc-123") assert rag_routes._verify_document_token(tok) == "doc-123" - assert rag_routes._verify_document_token("doc-123.0.deadbeef") is None # expired/bad + assert ( + rag_routes._verify_document_token("doc-123.0.deadbeef") is None + ) # expired/bad assert rag_routes._verify_document_token("garbage") is None diff --git a/studio/backend/tests/test_rag_retrieval.py b/studio/backend/tests/test_rag_retrieval.py index 69d9e90871..6cad207bb5 100644 --- a/studio/backend/tests/test_rag_retrieval.py +++ b/studio/backend/tests/test_rag_retrieval.py @@ -57,7 +57,9 @@ def _add_doc( text, page = None, ): - store.create_document(conn, scope = scope, filename = filename, sha256 = sha, document_id = doc_id) + store.create_document( + conn, scope = scope, filename = filename, sha256 = sha, document_id = doc_id + ) store.add_chunks(conn, scope, doc_id, [_chunk(text, 0, page)], [_embed(text)]) @@ -150,7 +152,9 @@ def test_tool_formats_chunks_and_sources(rag_conn, bow_embeddings, monkeypatch): def test_tool_kb_scope_retrieves_from_db(rag_conn, bow_embeddings): # End-to-end (no retrieve stub): doc found via its scope_kb_id (#8). _add_doc(rag_conn, "kb_K", "d1", "kb.pdf", "h1", "alpha bravo charlie", page = 1) - text, sources = tool.search_knowledge_base_with_sources(query = "alpha bravo", scope_kb_id = "K") + text, sources = tool.search_knowledge_base_with_sources( + query = "alpha bravo", scope_kb_id = "K" + ) assert "No matching chunks" not in text assert sources and sources[0]["chunkId"] == "d1:0" assert sources[0]["filename"] == "kb.pdf" @@ -192,11 +196,15 @@ def test_dispatcher_no_sentinel_when_no_hits(rag_home, monkeypatch): assert tools.RAG_SOURCES_SENTINEL not in out -def test_search_for_autoinject_gates_on_dense_score(rag_conn, bow_embeddings, monkeypatch): +def test_search_for_autoinject_gates_on_dense_score( + rag_conn, bow_embeddings, monkeypatch +): _add_doc(rag_conn, "kb_a", "d1", "paper.pdf", "h1", "body text here", page = 3) def _hits(score, **kw): - return lambda conn, scope, q, **k: [retrieval.Hit("d1:0", 1.0, **{kw["key"]: score})] + return lambda conn, scope, q, **k: [ + retrieval.Hit("d1:0", 1.0, **{kw["key"]: score}) + ] # Strong dense hit -> injected. monkeypatch.setattr(retrieval, "retrieve_hybrid", _hits(0.8, key = "dense_score")) @@ -207,14 +215,22 @@ def test_search_for_autoinject_gates_on_dense_score(rag_conn, bow_embeddings, mo # Dense below floor -> nothing injected. monkeypatch.setattr(retrieval, "retrieve_hybrid", _hits(0.30, key = "dense_score")) - assert tool.search_for_autoinject(query = "q", scope_kb_id = "a", min_dense_score = 0.55) is None + assert ( + tool.search_for_autoinject(query = "q", scope_kb_id = "a", min_dense_score = 0.55) + is None + ) # Lexical-only hit (no dense score) does not auto-inject. monkeypatch.setattr(retrieval, "retrieve_hybrid", _hits(1.0, key = "lexical_score")) - assert tool.search_for_autoinject(query = "q", scope_kb_id = "a", min_dense_score = 0.55) is None + assert ( + tool.search_for_autoinject(query = "q", scope_kb_id = "a", min_dense_score = 0.55) + is None + ) -def test_search_for_autoinject_bm25_gates_on_dense_probe(rag_conn, bow_embeddings, monkeypatch): +def test_search_for_autoinject_bm25_gates_on_dense_probe( + rag_conn, bow_embeddings, monkeypatch +): # BM25 hits carry no cosine, so the gate uses a dense 1-NN probe (#5). _add_doc(rag_conn, "kb_a", "d1", "paper.pdf", "h1", "body text here", page = 3) monkeypatch.setattr( @@ -226,7 +242,9 @@ def test_search_for_autoinject_bm25_gates_on_dense_probe(rag_conn, bow_embedding monkeypatch.setattr( retrieval, "retrieve_dense", - lambda conn, scope, q, k = None, **kw: [retrieval.Hit("d1:0", 0.82, dense_score = 0.82)], + lambda conn, scope, q, k = None, **kw: [ + retrieval.Hit("d1:0", 0.82, dense_score = 0.82) + ], ) found = tool.search_for_autoinject( query = "q", scope_kb_id = "a", mode = "lexical", min_dense_score = 0.70 @@ -236,10 +254,14 @@ def test_search_for_autoinject_bm25_gates_on_dense_probe(rag_conn, bow_embedding monkeypatch.setattr( retrieval, "retrieve_dense", - lambda conn, scope, q, k = None, **kw: [retrieval.Hit("d1:0", 0.40, dense_score = 0.40)], + lambda conn, scope, q, k = None, **kw: [ + retrieval.Hit("d1:0", 0.40, dense_score = 0.40) + ], ) assert ( - tool.search_for_autoinject(query = "q", scope_kb_id = "a", mode = "lexical", min_dense_score = 0.70) + tool.search_for_autoinject( + query = "q", scope_kb_id = "a", mode = "lexical", min_dense_score = 0.70 + ) is None ) @@ -271,7 +293,10 @@ def test_build_rag_autoinject_emits_pipeline(monkeypatch): te = next(e for e in out["events"] if e["type"] == "tool_end") assert te["tool_name"] == "search_knowledge_base" assert tools.RAG_SOURCES_SENTINEL in te["result"] - assert out["messages"][0]["tool_calls"][0]["function"]["name"] == "search_knowledge_base" + assert ( + out["messages"][0]["tool_calls"][0]["function"]["name"] + == "search_knowledge_base" + ) assert "__RAG_SOURCES__" not in out["messages"][1]["content"] @@ -282,7 +307,10 @@ def test_build_rag_autoinject_skips_without_hit(monkeypatch): monkeypatch.setattr(rag_db, "RAG_AVAILABLE", True, raising = False) monkeypatch.setattr(tool, "search_for_autoinject", lambda **k: None) assert ( - tools.build_rag_autoinject([{"role": "user", "content": "hi"}], {"thread_id": "t1"}) is None + tools.build_rag_autoinject( + [{"role": "user", "content": "hi"}], {"thread_id": "t1"} + ) + is None ) @@ -300,7 +328,9 @@ def test_build_rag_autoinject_enabled_by_default(monkeypatch): return ("x", [{"citationId": 1}]) monkeypatch.setattr(tool, "search_for_autoinject", fake) - out = tools.build_rag_autoinject([{"role": "user", "content": "hi"}], {"thread_id": "t1"}) + out = tools.build_rag_autoinject( + [{"role": "user", "content": "hi"}], {"thread_id": "t1"} + ) assert out is not None assert seen["min_dense_score"] == 0.70 # high-precision floor by default @@ -331,7 +361,10 @@ def test_build_rag_autoinject_disabled_by_env(monkeypatch): monkeypatch.setenv("RAG_AUTOINJECT", "0") assert ( - tools.build_rag_autoinject([{"role": "user", "content": "hi"}], {"thread_id": "t1"}) is None + tools.build_rag_autoinject( + [{"role": "user", "content": "hi"}], {"thread_id": "t1"} + ) + is None ) # No scope -> also a no-op. monkeypatch.delenv("RAG_AUTOINJECT", raising = False) @@ -429,4 +462,7 @@ def test_build_rag_autoinject_scope_overrides_env(monkeypatch): # Explicit False disables even with the env default on. monkeypatch.setenv("RAG_AUTOINJECT", "1") - assert tools.build_rag_autoinject(conv, {"thread_id": "t1", "autoinject": False}) is None + assert ( + tools.build_rag_autoinject(conv, {"thread_id": "t1", "autoinject": False}) + is None + ) diff --git a/studio/backend/tests/test_rag_store.py b/studio/backend/tests/test_rag_store.py index 4c54b02ea7..5d216af228 100644 --- a/studio/backend/tests/test_rag_store.py +++ b/studio/backend/tests/test_rag_store.py @@ -36,7 +36,9 @@ def _chunk( def _add_doc(conn, scope, doc_id, filename, sha, texts): chunks = [_chunk(t, i) for i, t in enumerate(texts)] vectors = [embed(t) for t in texts] - store.create_document(conn, scope = scope, filename = filename, sha256 = sha, document_id = doc_id) + store.create_document( + conn, scope = scope, filename = filename, sha256 = sha, document_id = doc_id + ) store.add_chunks(conn, scope, doc_id, chunks, vectors) @@ -50,7 +52,9 @@ def test_lexical_returns_only_matching_docs(rag_conn): def test_scope_isolation(rag_conn): _add_doc(rag_conn, "kb_a", "d1", "f", "h1", ["alpha bravo"]) _add_doc(rag_conn, "kb_b", "d2", "f", "h2", ["alpha bravo"]) - assert [cid for cid, _ in store.search_lexical(rag_conn, "kb_b", "alpha", 10)] == ["d2:0"] + assert [cid for cid, _ in store.search_lexical(rag_conn, "kb_b", "alpha", 10)] == [ + "d2:0" + ] def test_match_query_sanitizes_special_chars(): @@ -100,7 +104,9 @@ def test_incremental_add_is_flat(rag_conn): after = rag_conn.execute( "SELECT rowid, chunk_id FROM chunks_fts WHERE scope='kb_a' AND chunk_id LIKE 'd1:%'" ).fetchall() - before_d1 = [(r["rowid"], r["chunk_id"]) for r in before if r["chunk_id"].startswith("d1:")] + before_d1 = [ + (r["rowid"], r["chunk_id"]) for r in before if r["chunk_id"].startswith("d1:") + ] after_d1 = [(r["rowid"], r["chunk_id"]) for r in after] assert before_d1 == after_d1 diff --git a/studio/backend/tests/test_responses_api.py b/studio/backend/tests/test_responses_api.py index 693e832113..2ad5aeab94 100644 --- a/studio/backend/tests/test_responses_api.py +++ b/studio/backend/tests/test_responses_api.py @@ -168,7 +168,9 @@ class TestResponsesResponse: resp = ResponsesResponse( model = "test-model", output = [ - ResponsesOutputMessage(content = [ResponsesOutputTextContent(text = "Hello!")]), + ResponsesOutputMessage( + content = [ResponsesOutputTextContent(text = "Hello!")] + ), ], usage = ResponsesUsage(input_tokens = 10, output_tokens = 5, total_tokens = 15), ) diff --git a/studio/backend/tests/test_responses_tool_passthrough.py b/studio/backend/tests/test_responses_tool_passthrough.py index ae7ff729bd..e2e71b1d30 100644 --- a/studio/backend/tests/test_responses_tool_passthrough.py +++ b/studio/backend/tests/test_responses_tool_passthrough.py @@ -169,7 +169,9 @@ class TestResponsesMultiTurnInput: def test_function_call_output_missing_call_id_rejected(self): with pytest.raises(ValidationError): - ResponsesFunctionCallOutputInputItem(type = "function_call_output", output = "x") + ResponsesFunctionCallOutputInputItem( + type = "function_call_output", output = "x" + ) def test_function_call_output_accepts_content_array(self): item = ResponsesFunctionCallOutputInputItem( @@ -229,7 +231,9 @@ class TestToolsTranslation: assert _translate_responses_tools_to_chat([]) is None def test_only_builtin_tools_returns_none(self): - assert _translate_responses_tools_to_chat([{"type": "web_search_preview"}]) is None + assert ( + _translate_responses_tools_to_chat([{"type": "web_search_preview"}]) is None + ) def test_description_optional(self): out = _translate_responses_tools_to_chat( @@ -261,7 +265,9 @@ class TestToolChoiceTranslation: """A client sending the Chat Completions nested shape isn't double-wrapped.""" already_nested = {"type": "function", "function": {"name": "get_weather"}} - assert _translate_responses_tool_choice_to_chat(already_nested) == already_nested + assert ( + _translate_responses_tool_choice_to_chat(already_nested) == already_nested + ) def test_unknown_shape_passes_through(self): obj = {"type": "allowed_tools", "tools": [{"type": "function", "name": "x"}]} @@ -650,14 +656,20 @@ class TestResponsesNonStreamingAdapter: ) assert [item["type"] for item in body["output"]] == ["reasoning", "message"] - assert body["output"][0]["content"] == [{"type": "reasoning_text", "text": "plan"}] + assert body["output"][0]["content"] == [ + {"type": "reasoning_text", "text": "plan"} + ] assert body["output"][0]["summary"] == [] assert body["output"][1]["content"][0]["text"] == "33" assert "" not in body["output"][1]["content"][0]["text"] assert "" not in body["output"][1]["content"][0]["text"] - def test_literal_think_tags_remain_visible_without_reasoning_request(self, monkeypatch): - body = self._run_with_message(monkeypatch, {"content": "show x tags"}) + def test_literal_think_tags_remain_visible_without_reasoning_request( + self, monkeypatch + ): + body = self._run_with_message( + monkeypatch, {"content": "show x tags"} + ) assert [item["type"] for item in body["output"]] == ["message"] assert body["output"][0]["content"][0]["text"] == "show x tags" @@ -691,7 +703,9 @@ class TestResponsesNonStreamingAdapter: ) assert [item["type"] for item in body["output"]] == ["reasoning", "message"] - assert body["output"][0]["content"] == [{"type": "reasoning_text", "text": "plan next"}] + assert body["output"][0]["content"] == [ + {"type": "reasoning_text", "text": "plan next"} + ] assert body["output"][1]["content"][0]["text"] == "33" def test_plain_content_remains_message_only(self, monkeypatch): @@ -779,12 +793,16 @@ class TestResponsesStreamAdapter: supports_reasoning = supports_reasoning, reasoning_always_on = reasoning_always_on, _request_reasoning_kwargs = ( - lambda enable_thinking = None, reasoning_effort = None, preserve_thinking = None: None + lambda enable_thinking = None, + reasoning_effort = None, + preserve_thinking = None: None ), ), ) - def test_split_think_markers_stream_as_reasoning_and_visible_text(self, monkeypatch): + def test_split_think_markers_stream_as_reasoning_and_visible_text( + self, monkeypatch + ): chunks = [ {"choices": [{"delta": {"content": "pla"}}]}, @@ -793,7 +811,9 @@ class TestResponsesStreamAdapter: {"choices": [], "usage": {"prompt_tokens": 2, "completion_tokens": 3}}, ] self._install_stream_mock(monkeypatch, chunks) - payload = ResponsesRequest(input = "hi", stream = True, reasoning = {"effort": "high"}) + payload = ResponsesRequest( + input = "hi", stream = True, reasoning = {"effort": "high"} + ) messages = [ChatMessage(role = "user", content = "hi")] async def run(): @@ -814,7 +834,9 @@ class TestResponsesStreamAdapter: assert completed["response"]["output"][0]["content"][0]["text"] == "plan" assert completed["response"]["output"][1]["content"][0]["text"] == "33" - def test_literal_think_tags_stream_as_visible_text_without_reasoning_request(self, monkeypatch): + def test_literal_think_tags_stream_as_visible_text_without_reasoning_request( + self, monkeypatch + ): chunks = [ {"choices": [{"delta": {"content": "show x tags"}}]}, @@ -833,21 +855,28 @@ class TestResponsesStreamAdapter: reasoning_deltas = self._payloads(lines, "response.reasoning_text.delta") text_deltas = self._payloads(lines, "response.output_text.delta") assert reasoning_deltas == [] - assert "".join(event["delta"] for event in text_deltas) == "show x tags" + assert ( + "".join(event["delta"] for event in text_deltas) + == "show x tags" + ) completed = self._payloads(lines, "response.completed")[0] assert [item["type"] for item in completed["response"]["output"]] == ["message"] assert completed["response"]["output"][0]["content"][0]["text"] == ( "show x tags" ) - def test_non_reasoning_gguf_stream_keeps_literal_think_tags_visible(self, monkeypatch): + def test_non_reasoning_gguf_stream_keeps_literal_think_tags_visible( + self, monkeypatch + ): chunks = [ {"choices": [{"delta": {"content": "show x tags"}}]}, {"choices": [], "usage": {"prompt_tokens": 2, "completion_tokens": 3}}, ] self._install_stream_mock(monkeypatch, chunks, supports_reasoning = False) - payload = ResponsesRequest(input = "hi", stream = True, reasoning = {"effort": "high"}) + payload = ResponsesRequest( + input = "hi", stream = True, reasoning = {"effort": "high"} + ) messages = [ChatMessage(role = "user", content = "hi")] async def run(): @@ -859,7 +888,10 @@ class TestResponsesStreamAdapter: reasoning_deltas = self._payloads(lines, "response.reasoning_text.delta") text_deltas = self._payloads(lines, "response.output_text.delta") assert reasoning_deltas == [] - assert "".join(event["delta"] for event in text_deltas) == "show x tags" + assert ( + "".join(event["delta"] for event in text_deltas) + == "show x tags" + ) completed = self._payloads(lines, "response.completed")[0] assert [item["type"] for item in completed["response"]["output"]] == ["message"] assert completed["response"]["output"][0]["content"][0]["text"] == ( @@ -872,7 +904,9 @@ class TestResponsesStreamAdapter: {"choices": [], "usage": {"prompt_tokens": 2, "completion_tokens": 3}}, ] self._install_stream_mock(monkeypatch, chunks) - payload = ResponsesRequest(input = "hi", stream = True, reasoning = {"effort": "high"}) + payload = ResponsesRequest( + input = "hi", stream = True, reasoning = {"effort": "high"} + ) messages = [ChatMessage(role = "user", content = "hi")] async def run(): @@ -950,7 +984,9 @@ class TestResponsesStreamAdapter: text_deltas = self._payloads(lines, "response.output_text.delta") assert "".join(event["delta"] for event in reasoning_deltas) == "plan next" assert "".join(event["delta"] for event in text_deltas) == "33" - assert "reasoning_text" not in "".join(event["delta"] for event in reasoning_deltas) + assert "reasoning_text" not in "".join( + event["delta"] for event in reasoning_deltas + ) completed = self._payloads(lines, "response.completed")[0] assert completed["response"]["output"][0]["content"][0]["text"] == "plan next" assert completed["response"]["output"][1]["content"][0]["text"] == "33" @@ -988,7 +1024,10 @@ class TestResponsesStreamAdapter: done_events = self._payloads(lines, "response.output_item.done") assert [event["output_index"] for event in done_events] == [0, 1] - assert [event["item"]["type"] for event in done_events] == ["function_call", "message"] + assert [event["item"]["type"] for event in done_events] == [ + "function_call", + "message", + ] completed = self._payloads(lines, "response.completed")[0] assert [item["type"] for item in completed["response"]["output"]] == [ "function_call", @@ -1012,13 +1051,19 @@ class TestResponsesStreamAdapter: "index": 0, "id": "call_0", "type": "function", - "function": {"name": "first", "arguments": "{}"}, + "function": { + "name": "first", + "arguments": "{}", + }, }, { "index": 1, "id": "call_1", "type": "function", - "function": {"name": "second", "arguments": "{}"}, + "function": { + "name": "second", + "arguments": "{}", + }, }, ] } @@ -1055,7 +1100,9 @@ class TestResponsesStreamAdapter: base_url = "http://llama.test", # Non-reasoning template: the real backend returns None here. _request_reasoning_kwargs = ( - lambda enable_thinking = None, reasoning_effort = None, preserve_thinking = None: None + lambda enable_thinking = None, + reasoning_effort = None, + preserve_thinking = None: None ), ), ) @@ -1104,7 +1151,9 @@ class TestResponsesStreamAdapter: class TestResponsesOutputFunctionCall: def test_reasoning_output_item_serialises_full_reasoning_content(self): - item = ResponsesOutputReasoning(content = [{"type": "reasoning_text", "text": "plan"}]) + item = ResponsesOutputReasoning( + content = [{"type": "reasoning_text", "text": "plan"}] + ) d = item.model_dump() assert d["type"] == "reasoning" assert d["id"].startswith("rs_") @@ -1230,7 +1279,9 @@ class TestCodexStyleRequestShapes: msgs = _normalise_responses_input(payload) assert [m.role for m in msgs] == ["user", "assistant", "user"] - assert all("plan" not in (m.content or "") for m in msgs if isinstance(m.content, str)) + assert all( + "plan" not in (m.content or "") for m in msgs if isinstance(m.content, str) + ) def test_unknown_content_part_type_accepted(self): """Unknown content-part types (e.g. future input_audio) validate as @@ -1321,7 +1372,9 @@ class TestCodexStyleRequestShapes: input = [ { "role": "assistant", - "content": [{"type": "output_text", "text": "ok", "annotations": []}], + "content": [ + {"type": "output_text", "text": "ok", "annotations": []} + ], }, {"role": "user", "content": "next"}, ], diff --git a/studio/backend/tests/test_rocm_oom_guard.py b/studio/backend/tests/test_rocm_oom_guard.py index 6e70c7cde4..767a3e4caf 100644 --- a/studio/backend/tests/test_rocm_oom_guard.py +++ b/studio/backend/tests/test_rocm_oom_guard.py @@ -172,7 +172,9 @@ class TestDeviceNameFallback: props = _props(name = device_name) gcn, is_unified = _rocm_classify_unified_memory(props) assert gcn == "", f"expected empty gcn_arch, got {gcn!r}" - assert is_unified is True, f"device {device_name!r} should be classified as unified-memory" + assert ( + is_unified is True + ), f"device {device_name!r} should be classified as unified-memory" # --- discrete devices that must NOT be mis-classified --- diff --git a/studio/backend/tests/test_s3_dataset.py b/studio/backend/tests/test_s3_dataset.py index f47db565ff..791954b8f4 100644 --- a/studio/backend/tests/test_s3_dataset.py +++ b/studio/backend/tests/test_s3_dataset.py @@ -37,7 +37,9 @@ class _FakePaginator: def paginate(self, **kwargs): prefix = kwargs.get("Prefix") - contents = [{"Key": k} for k in self._keys if prefix is None or k.startswith(prefix)] + contents = [ + {"Key": k} for k in self._keys if prefix is None or k.startswith(prefix) + ] # Emit in two pages to exercise pagination handling. mid = len(contents) // 2 yield {"Contents": contents[:mid]} diff --git a/studio/backend/tests/test_safetensors_capability_advertise.py b/studio/backend/tests/test_safetensors_capability_advertise.py index 1e8fb9e2b2..dd5a44c91f 100644 --- a/studio/backend/tests/test_safetensors_capability_advertise.py +++ b/studio/backend/tests/test_safetensors_capability_advertise.py @@ -363,7 +363,11 @@ def test_worker_load_reply_payload_includes_chat_template_info(): "is_gguf": False, } _bm = getattr(backend, "models", {}) or {} - _entry = _bm.get(mc.identifier) or _bm.get(getattr(backend, "active_model_name", None)) or {} + _entry = ( + _bm.get(mc.identifier) + or _bm.get(getattr(backend, "active_model_name", None)) + or {} + ) _tpl_info = _entry.get("chat_template_info") if isinstance(_tpl_info, dict): model_info["chat_template_info"] = { diff --git a/studio/backend/tests/test_safetensors_tool_loop.py b/studio/backend/tests/test_safetensors_tool_loop.py index 12731783a0..7b9e35bb2c 100644 --- a/studio/backend/tests/test_safetensors_tool_loop.py +++ b/studio/backend/tests/test_safetensors_tool_loop.py @@ -40,7 +40,9 @@ from utils.datasets import is_gpt_oss_model_name class TestParser: def test_json_tool_call(self): - text = '{"name":"web_search","arguments":{"query":"hello"}}' + text = ( + '{"name":"web_search","arguments":{"query":"hello"}}' + ) result = parse_tool_calls_from_text(text) assert len(result) == 1 tc = result[0] @@ -85,9 +87,7 @@ class TestParser: def test_code_with_embedded_xml(self): # A code parameter with a literal must not truncate: the # parser uses end-of-body as the only boundary for single-param calls. - text = ( - "html = ''\nprint('hi')" - ) + text = "html = ''\nprint('hi')" result = parse_tool_calls_from_text(text) assert len(result) == 1 assert "print('hi')" in result[0]["function"]["arguments"] @@ -293,7 +293,10 @@ def test_active_tools_are_passed_to_single_turn_after_render_html_success(): assert exec_fn.calls == [("render_html", {"code": "one"})] assert captured_tool_names == [["render_html", "web_search"], ["web_search"]] - assert any(event.get("type") == "content" and event.get("text") == "Done." for event in events) + assert any( + event.get("type") == "content" and event.get("text") == "Done." + for event in events + ) class TestLoopBasic: @@ -392,7 +395,9 @@ class TestLoopBasic: assert exec_fn.calls[0][0] == "render_html" assert "" in exec_fn.calls[0][1]["code"] - def test_python_tool_containing_render_html_signal_does_not_emit_provisional_start(self): + def test_python_tool_containing_render_html_signal_does_not_emit_provisional_start( + self, + ): loop, exec_fn = _make_loop( turns = [ [ @@ -409,7 +414,9 @@ class TestLoopBasic: assert len(tool_starts) == 1 assert tool_starts[0]["tool_name"] == "python" - assert exec_fn.calls == [("python", {"code": "print('')"})] + assert exec_fn.calls == [ + ("python", {"code": "print('')"}) + ] def test_render_html_success_blocks_second_artifact_call(self): exec_fn = FakeExecuteTool(["Rendered HTML artifact."]) @@ -444,7 +451,10 @@ class TestLoopBasic: tool_starts = [e for e in events if e["type"] == "tool_start"] assert exec_fn.calls == [("render_html", {"code": "one"})] - assert [e["arguments"] for e in tool_starts] == [{}, {"code": "one"}] + assert [e["arguments"] for e in tool_starts] == [ + {}, + {"code": "one"}, + ] def test_truncated_unclosed_tool_call(self): loop, exec_fn = _make_loop( @@ -464,7 +474,9 @@ class TestLoopBasic: loop, exec_fn = _make_loop( turns = [ # ``arguments`` is a string _coerce_arguments can't parse, so heal runs. - ['{"name":"web_search","arguments":"hello world"}'], + [ + '{"name":"web_search","arguments":"hello world"}' + ], ["ok"], ], exec_results = ["..."], @@ -479,8 +491,12 @@ class TestLoopBehaviour: captured_messages: list[list[dict]] = [] turns = iter( [ - ['{"name":"web_search","arguments":{"query":"x"}}'], - ['{"name":"web_search","arguments":{"query":"x"}}'], + [ + '{"name":"web_search","arguments":{"query":"x"}}' + ], + [ + '{"name":"web_search","arguments":{"query":"x"}}' + ], ["final"], ] ) @@ -505,11 +521,14 @@ class TestLoopBehaviour: ) assert exec_fn.calls == [("web_search", {"query": "x"})] - assert [e["tool_call_id"] for e in events if e["type"] == "tool_end"] == ["call_0"] + assert [e["tool_call_id"] for e in events if e["type"] == "tool_end"] == [ + "call_0" + ] assert not [ e for e in events - if e.get("tool_call_id") == "call_1" and e.get("type") in {"tool_start", "tool_end"} + if e.get("tool_call_id") == "call_1" + and e.get("type") in {"tool_start", "tool_end"} ] duplicate_nudges = [ message @@ -524,9 +543,15 @@ class TestLoopBehaviour: captured_tool_names: list[list[str]] = [] turns = iter( [ - ['{"name":"web_search","arguments":{"query":"x"}}'], - ['{"name":"web_search","arguments":{"query":"x"}}'], - ['{"name":"python","arguments":{"code":"print(1)"}}'], + [ + '{"name":"web_search","arguments":{"query":"x"}}' + ], + [ + '{"name":"web_search","arguments":{"query":"x"}}' + ], + [ + '{"name":"python","arguments":{"code":"print(1)"}}' + ], ["final"], ] ) @@ -571,7 +596,8 @@ class TestLoopBehaviour: assert not [ e for e in events - if e.get("tool_call_id") == "call_1" and e.get("type") in {"tool_start", "tool_end"} + if e.get("tool_call_id") == "call_1" + and e.get("type") in {"tool_start", "tool_end"} ] duplicate_nudges = [ message @@ -586,9 +612,15 @@ class TestLoopBehaviour: captured_tool_names: list[list[str]] = [] turns = iter( [ - ['{"name":"web_search","arguments":{"query":"x"}}'], - ['{"name":"web_search","arguments":{"query":"x"}}'], - ['{"name":"web_search","arguments":{"query":"x"}}'], + [ + '{"name":"web_search","arguments":{"query":"x"}}' + ], + [ + '{"name":"web_search","arguments":{"query":"x"}}' + ], + [ + '{"name":"web_search","arguments":{"query":"x"}}' + ], ["final from first result"], ] ) @@ -620,11 +652,14 @@ class TestLoopBehaviour: assert exec_fn.calls == [("web_search", {"query": "x"})] assert [ - event.get("tool_call_id") for event in events if event.get("type") == "tool_end" + event.get("tool_call_id") + for event in events + if event.get("type") == "tool_end" ] == ["call_0"] assert captured_tool_names[-1] == [] assert any( - event.get("type") == "content" and "final from first result" in event.get("text", "") + event.get("type") == "content" + and "final from first result" in event.get("text", "") for event in events ) @@ -671,7 +706,9 @@ class TestLoopBehaviour: # carries the raw result for the UI. loop, exec_fn = _make_loop( turns = [ - ['{"name":"python","arguments":{"code":"plot()"}}'], + [ + '{"name":"python","arguments":{"code":"plot()"}}' + ], ["see chart"], ], exec_results = ["chart\n__IMAGES__:/tmp/chart.png"], @@ -709,7 +746,9 @@ class TestLoopBehaviour: tool_msgs = [m for m in captured[1] if m.get("role") == "tool"] assert tool_msgs, "no tool message reached the model" for tm in tool_msgs: - assert "__IMAGES__" not in tm["content"], f"sentinel leaked to model: {tm['content']!r}" + assert ( + "__IMAGES__" not in tm["content"] + ), f"sentinel leaked to model: {tm['content']!r}" def test_image_sentinel_stripped_with_multiple_markers(self): # Consecutive sentinels: cut at the first, nothing leaks. @@ -739,13 +778,19 @@ class TestLoopBehaviour: tool_msgs = [m for m in captured[1] if m.get("role") == "tool"] assert tool_msgs for tm in tool_msgs: - assert "__IMAGES__" not in tm["content"], f"second sentinel leaked: {tm['content']!r}" - assert tm["content"] == "panel", f"expected payload-only 'panel', got {tm['content']!r}" + assert ( + "__IMAGES__" not in tm["content"] + ), f"second sentinel leaked: {tm['content']!r}" + assert ( + tm["content"] == "panel" + ), f"expected payload-only 'panel', got {tm['content']!r}" def test_tool_execution_error_is_emitted_but_loop_continues(self): loop, exec_fn = _make_loop( turns = [ - ['{"name":"web_search","arguments":{"query":"x"}}'], + [ + '{"name":"web_search","arguments":{"query":"x"}}' + ], ["sorry, that failed"], ], exec_results = ["Error: network unreachable"], @@ -760,7 +805,9 @@ class TestLoopBehaviour: def test_exception_in_executor_does_not_raise(self): loop, exec_fn = _make_loop( turns = [ - ['{"name":"web_search","arguments":{"query":"x"}}'], + [ + '{"name":"web_search","arguments":{"query":"x"}}' + ], ["recovered"], ], exec_results = [RuntimeError("boom")], @@ -796,7 +843,9 @@ class TestLoopControl: loop, exec_fn = _make_loop( turns = [ # Tool call (executes once). - ['{"name":"web_search","arguments":{"query":"a"}}'], + [ + '{"name":"web_search","arguments":{"query":"a"}}' + ], # Model gives a final answer when nudged. ["here is the final answer"], ], @@ -813,19 +862,24 @@ class TestStatusFormatting: def test_status_for_known_tools(self): # Call the private helper directly to verify status formatting. assert ( - safetensors_agentic._status_for_tool("web_search", {"query": "abc"}) == "Searching: abc" + safetensors_agentic._status_for_tool("web_search", {"query": "abc"}) + == "Searching: abc" ) assert ( - safetensors_agentic._status_for_tool("web_search", {"url": "https://www.example.com/x"}) + safetensors_agentic._status_for_tool( + "web_search", {"url": "https://www.example.com/x"} + ) == "Reading: example.com" ) - assert safetensors_agentic._status_for_tool("python", {"code": "x = 1"}).startswith( - "Running Python:" + assert safetensors_agentic._status_for_tool( + "python", {"code": "x = 1"} + ).startswith("Running Python:") + assert safetensors_agentic._status_for_tool( + "terminal", {"command": "ls"} + ).startswith("Running:") + assert safetensors_agentic._status_for_tool("unknown_tool", {}).startswith( + "Calling:" ) - assert safetensors_agentic._status_for_tool("terminal", {"command": "ls"}).startswith( - "Running:" - ) - assert safetensors_agentic._status_for_tool("unknown_tool", {}).startswith("Calling:") class TestProseMentioningToolCall: @@ -835,7 +889,9 @@ class TestProseMentioningToolCall: loop, exec_fn = _make_loop( turns = [ # A real tool call so the loop advances a turn. - ['{"name":"web_search","arguments":{"query":"x"}}'], + [ + '{"name":"web_search","arguments":{"query":"x"}}' + ], # Prose that mentions the literal text. ["the docs say means an LLM tool call wrapper"], ], @@ -854,7 +910,9 @@ class TestProseMentioningToolCall: # loop parses only model output, so exactly one call. loop, exec_fn = _make_loop( turns = [ - ['{"name":"web_search","arguments":{"query":"x"}}'], + [ + '{"name":"web_search","arguments":{"query":"x"}}' + ], ["the docs mention wrappers"], ], exec_results = ["Page text: appears here in the docs"], @@ -955,18 +1013,23 @@ class TestGuardrails: ) assert exec_fn.calls == [] - assert not [event for event in events if event.get("type") in {"tool_start", "tool_end"}] + assert not [ + event for event in events if event.get("type") in {"tool_start", "tool_end"} + ] disabled_nudges = [ message for message in captured_messages[-1] - if message.get("role") == "user" and "not enabled" in message.get("content", "") + if message.get("role") == "user" + and "not enabled" in message.get("content", "") ] assert len(disabled_nudges) == 1 def test_empty_tools_list_means_allow_all_in_core_loop(self): turns = iter( [ - ['{"name":"python","arguments":{"code":"print(1)"}}'], + [ + '{"name":"python","arguments":{"code":"print(1)"}}' + ], ["done"], ] ) @@ -993,7 +1056,11 @@ class TestGuardrails: def test_max_iterations_zero_executes_no_tools(self): loop, exec_fn = _make_loop( - turns = [['{"name":"web_search","arguments":{"query":"x"}}']], + turns = [ + [ + '{"name":"web_search","arguments":{"query":"x"}}' + ] + ], exec_results = ["OK"], max_tool_iterations = 0, ) @@ -1024,7 +1091,9 @@ class TestGuardrails: def test_auto_heal_disabled_still_parses_valid_tool_call(self): loop, exec_fn = _make_loop( turns = [ - ['{"name":"web_search","arguments":{"query":"x"}}'], + [ + '{"name":"web_search","arguments":{"query":"x"}}' + ], ["done"], ], exec_results = ["OK"], @@ -1039,7 +1108,11 @@ class TestGuardrails: monkeypatch.setattr(safetensors_agentic, "new_approval_id", lambda: approval_id) loop, exec_fn = _make_loop( - turns = [['{"name":"python","arguments":{"code":"print(1)"}}']], + turns = [ + [ + '{"name":"python","arguments":{"code":"print(1)"}}' + ] + ], exec_results = ["OK"], confirm_tool_calls = True, session_id = "sess", @@ -1068,21 +1141,30 @@ class TestGuardrails: def fail_autoinject(*_args, **_kwargs): raise AssertionError("RAG autoinject must not run before approval") - monkeypatch.setattr("core.inference.tools.build_rag_autoinject", fail_autoinject) + monkeypatch.setattr( + "core.inference.tools.build_rag_autoinject", fail_autoinject + ) loop, exec_fn = _make_loop( turns = [["plain answer"]], confirm_tool_calls = True, rag_scope = {"thread_id": "t1"}, ) events = _collect_events(loop) - assert any(e.get("type") == "content" and e.get("text") == "plain answer" for e in events) + assert any( + e.get("type") == "content" and e.get("text") == "plain answer" + for e in events + ) assert exec_fn.calls == [] def test_auto_heal_disabled_preserves_xml_on_final_no_tools_pass(self): turns = iter( [ - ['{"name":"web_search","arguments":{"query":"x"}}'], - ['{"name":"web_search","arguments":{"query":"literal"}}'], + [ + '{"name":"web_search","arguments":{"query":"x"}}' + ], + [ + '{"name":"web_search","arguments":{"query":"literal"}}' + ], ] ) @@ -1142,18 +1224,29 @@ class TestGuardrails: def test_non_consecutive_duplicate_is_short_circuited(self): loop, exec_fn = _make_loop( turns = [ - ['{"name":"web_search","arguments":{"query":"A"}}'], - ['{"name":"web_search","arguments":{"query":"B"}}'], - ['{"name":"web_search","arguments":{"query":"A"}}'], + [ + '{"name":"web_search","arguments":{"query":"A"}}' + ], + [ + '{"name":"web_search","arguments":{"query":"B"}}' + ], + [ + '{"name":"web_search","arguments":{"query":"A"}}' + ], ["final"], ], exec_results = ["res-A", "res-B"], max_tool_iterations = 4, ) events = _collect_events(loop) - assert exec_fn.calls == [("web_search", {"query": "A"}), ("web_search", {"query": "B"})] + assert exec_fn.calls == [ + ("web_search", {"query": "A"}), + ("web_search", {"query": "B"}), + ] assert [ - event.get("tool_call_id") for event in events if event.get("type") == "tool_end" + event.get("tool_call_id") + for event in events + if event.get("type") == "tool_end" ] == ["call_0", "call_1"] assert not [ event @@ -1177,7 +1270,9 @@ class TestGuardrails: events = _collect_events(loop) assert exec_fn.calls == [("web_search", {"query": "A"})] assert [ - event.get("tool_call_id") for event in events if event.get("type") == "tool_end" + event.get("tool_call_id") + for event in events + if event.get("type") == "tool_end" ] == ["call_0"] assert not [ event @@ -1187,16 +1282,24 @@ class TestGuardrails: ] def test_coerce_string_args_python_uses_code_key(self): - assert _coerce_arguments("print(1)", heal = True, tool_name = "python") == {"code": "print(1)"} + assert _coerce_arguments("print(1)", heal = True, tool_name = "python") == { + "code": "print(1)" + } def test_coerce_string_args_terminal_uses_command_key(self): - assert _coerce_arguments("ls -la", heal = True, tool_name = "terminal") == {"command": "ls -la"} + assert _coerce_arguments("ls -la", heal = True, tool_name = "terminal") == { + "command": "ls -la" + } def test_tool_call_ids_unique_across_loop_iterations(self): loop, _exec = _make_loop( turns = [ - ['{"name":"web_search","arguments":{"query":"A"}}'], - ['{"name":"web_search","arguments":{"query":"B"}}'], + [ + '{"name":"web_search","arguments":{"query":"A"}}' + ], + [ + '{"name":"web_search","arguments":{"query":"B"}}' + ], ["done"], ], exec_results = ["A", "B"], diff --git a/studio/backend/tests/test_sandbox_tools.py b/studio/backend/tests/test_sandbox_tools.py index 24b1da1772..5578e9dd97 100644 --- a/studio/backend/tests/test_sandbox_tools.py +++ b/studio/backend/tests/test_sandbox_tools.py @@ -88,7 +88,9 @@ class TestTrustedHostAllowlist: _ok(f"import requests; requests.get({url!r})") def test_wikipedia_subdomain_passes(self): - _ok('import urllib.request; urllib.request.urlopen("https://m.en.wikipedia.org/wiki/Foo")') + _ok( + 'import urllib.request; urllib.request.urlopen("https://m.en.wikipedia.org/wiki/Foo")' + ) def test_hf_co_short_form_passes(self): _ok('import requests; requests.get("https://hf.co/unsloth/Qwen3.5-4B-GGUF")') @@ -219,7 +221,10 @@ class TestUploadDenylist: ) def test_plain_post_json_not_blocked(self): - _ok("import requests\n" 'requests.post("https://api.weather.gov/lookup", json={"k": "v"})') + _ok( + "import requests\n" + 'requests.post("https://api.weather.gov/lookup", json={"k": "v"})' + ) class TestSandboxEnvIsolation: diff --git a/studio/backend/tests/test_server_disk_logging.py b/studio/backend/tests/test_server_disk_logging.py index 05d03d869c..d4f6f9f229 100644 --- a/studio/backend/tests/test_server_disk_logging.py +++ b/studio/backend/tests/test_server_disk_logging.py @@ -91,7 +91,9 @@ class TestSetupServerDiskLogging: def test_run_server_wires_logging_before_main_import(self): src = (Path(_BACKEND_DIR) / "run.py").read_text(encoding = "utf-8") - call_idx = src.index("_setup_server_disk_logging()", src.index("def run_server")) + call_idx = src.index( + "_setup_server_disk_logging()", src.index("def run_server") + ) main_import_idx = src.index("from main import app", src.index("def run_server")) assert call_idx < main_import_idx, ( "disk logging must be armed before importing main so import-time " diff --git a/studio/backend/tests/test_studio_api.py b/studio/backend/tests/test_studio_api.py index 928b636e3e..ea4c55fa9c 100644 --- a/studio/backend/tests/test_studio_api.py +++ b/studio/backend/tests/test_studio_api.py @@ -72,7 +72,11 @@ DEFAULT_VARIANT = "UD-Q4_K_XL" PORT = 18222 # high port unlikely to collide HOST = "127.0.0.1" STARTUP_TIMEOUT = 120 # seconds -LOG_FILE = Path(__file__).resolve().parent.parent.parent.parent / "temp" / "test_studio_api.log" +LOG_FILE = ( + Path(__file__).resolve().parent.parent.parent.parent + / "temp" + / "test_studio_api.log" +) # Helpers @@ -216,7 +220,9 @@ def test_openai_sdk(base_url: str, api_key: str): client = OpenAI(base_url = f"{base_url}/v1", api_key = api_key) response = client.chat.completions.create( model = "current", - messages = [{"role": "user", "content": "What is 2+2? Answer with just the number."}], + messages = [ + {"role": "user", "content": "What is 2+2? Answer with just the number."} + ], stream = True, ) content_parts = [] @@ -379,7 +385,9 @@ def test_openai_tools_nonstream(base_url: str, api_key: str): assert "city" in parsed, f"Tool call missing required 'city' arg: {parsed}" # Usage must be non-zero (was 0 before the fix) usage = data.get("usage") or {} - assert usage.get("prompt_tokens", 0) > 0, f"Expected non-zero prompt_tokens; got {usage}" + assert ( + usage.get("prompt_tokens", 0) > 0 + ), f"Expected non-zero prompt_tokens; got {usage}" assert data.get("id"), "Missing response id" print( f" PASS openai tools non-stream: " @@ -404,7 +412,8 @@ def test_openai_tools_stream(base_url: str, api_key: str): assert status == 200, f"Expected 200, got {status}" assert len(chunks) > 0, "No SSE chunks received" assert _final_finish_reason(chunks) == "tool_calls", ( - f"Expected final finish_reason='tool_calls', got " f"{_final_finish_reason(chunks)!r}" + f"Expected final finish_reason='tool_calls', got " + f"{_final_finish_reason(chunks)!r}" ) assembled = _collect_streamed_tool_calls(chunks) assert len(assembled) >= 1, "No tool_calls reassembled from stream" @@ -487,7 +496,8 @@ def test_openai_sdk_tool_calling(base_url: str, api_key: str): stream = False, ) assert resp.choices[0].finish_reason == "tool_calls", ( - f"Expected finish_reason='tool_calls', got " f"{resp.choices[0].finish_reason!r}" + f"Expected finish_reason='tool_calls', got " + f"{resp.choices[0].finish_reason!r}" ) tool_calls = resp.choices[0].message.tool_calls assert tool_calls and len(tool_calls) >= 1, "No tool_calls from SDK" @@ -495,7 +505,9 @@ def test_openai_sdk_tool_calling(base_url: str, api_key: str): assert tc.function.name == "get_weather" parsed = json.loads(tc.function.arguments) assert "city" in parsed - print(f" PASS openai SDK tool calling: " f"tool={tc.function.name}, args={parsed}") + print( + f" PASS openai SDK tool calling: " f"tool={tc.function.name}, args={parsed}" + ) def test_invalid_key_rejected(base_url: str): @@ -638,7 +650,9 @@ def test_anthropic_sdk(base_url: str, api_key: str): message = client.messages.create( model = "default", max_tokens = 100, - messages = [{"role": "user", "content": "What is 2+2? Answer with just the number."}], + messages = [ + {"role": "user", "content": "What is 2+2? Answer with just the number."} + ], ) assert message.role == "assistant" assert len(message.content) > 0, "Empty content" @@ -689,7 +703,9 @@ def test_anthropic_with_tools(base_url: str, api_key: str): assert "message_stop" in event_types, "Missing message_stop" full = _collect_anthropic_text(events) - print(f" PASS anthropic with tools: {len(events)} events, {len(full)} chars content") + print( + f" PASS anthropic with tools: {len(events)} events, {len(full)} chars content" + ) def test_anthropic_tool_choice_any(base_url: str, api_key: str): @@ -749,7 +765,8 @@ def test_anthropic_tool_choice_any(base_url: str, api_key: str): tool_use_starts = [ e for e in events - if e[0] == "content_block_start" and e[1].get("content_block", {}).get("type") == "tool_use" + if e[0] == "content_block_start" + and e[1].get("content_block", {}).get("type") == "tool_use" ] assert len(tool_use_starts) >= 1, "No tool_use content block emitted" print( @@ -799,7 +816,9 @@ def _start_server(model: str, variant: str | None) -> tuple[subprocess.Popen, st if proc.poll() is not None: log_fh.flush() log_text = LOG_FILE.read_text() - raise RuntimeError(f"Server exited early (code {proc.returncode}):\n{log_text[-2000:]}") + raise RuntimeError( + f"Server exited early (code {proc.returncode}):\n{log_text[-2000:]}" + ) log_text = LOG_FILE.read_text() m = re.search(r"API Key:\s+(sk-unsloth-[a-f0-9]+)", log_text) if m: @@ -809,7 +828,9 @@ def _start_server(model: str, variant: str | None) -> tuple[subprocess.Popen, st if not api_key: log_text = LOG_FILE.read_text() _kill_server(proc) - raise RuntimeError(f"Timed out waiting for API key in server output:\n{log_text[-2000:]}") + raise RuntimeError( + f"Timed out waiting for API key in server output:\n{log_text[-2000:]}" + ) # Wait a moment for the model to be fully loaded time.sleep(2) @@ -836,7 +857,9 @@ def _kill_server(proc: subprocess.Popen): def main(): - parser = argparse.ArgumentParser(description = "End-to-end tests for unsloth studio run") + parser = argparse.ArgumentParser( + description = "End-to-end tests for unsloth studio run" + ) parser.add_argument( "--model", default = DEFAULT_MODEL, @@ -870,7 +893,9 @@ def main(): run_test(test_help_output) # 2-16. Start server and run API tests - print(f"\nStarting server: {args.model} (variant={args.gguf_variant}) on port {PORT}...") + print( + f"\nStarting server: {args.model} (variant={args.gguf_variant}) on port {PORT}..." + ) proc = None try: proc, api_key = _start_server(args.model, args.gguf_variant) diff --git a/studio/backend/tests/test_tensor_parallel.py b/studio/backend/tests/test_tensor_parallel.py index 30bfb91a08..c2f5ccbc9c 100644 --- a/studio/backend/tests/test_tensor_parallel.py +++ b/studio/backend/tests/test_tensor_parallel.py @@ -87,7 +87,9 @@ def test_load_request_accepts_tensor_parallel(): def test_load_request_round_trips_json_key(): # The frontend sends the snake_case key verbatim. - req = LoadRequest.model_validate({"model_path": "owner/repo", "tensor_parallel": True}) + req = LoadRequest.model_validate( + {"model_path": "owner/repo", "tensor_parallel": True} + ) assert req.tensor_parallel is True assert req.model_dump()["tensor_parallel"] is True @@ -259,7 +261,9 @@ def test_proportional_tensor_split_is_emitted_in_tensor_mode(): gate = src.find("if tensor_parallel:") ts = src.find('"--tensor-split"') nxt_else = src.find("self._tensor_parallel = False") - assert 0 <= gate < ts < nxt_else, "--tensor-split must be emitted under `if tensor_parallel:`" + assert ( + 0 <= gate < ts < nxt_else + ), "--tensor-split must be emitted under `if tensor_parallel:`" # ── tensor-mode allocation: conservative VRAM budget ───────────────── @@ -282,11 +286,15 @@ def test_fit_context_budget_frac_override_is_tighter(): pool_mib = 24 * 1024 # tight enough that KV capping bites fit_default = backend._fit_context_to_vram(131072, pool_mib, model_size, "f16") - fit_tp = backend._fit_context_to_vram(131072, pool_mib, model_size, "f16", budget_frac = 0.80) + fit_tp = backend._fit_context_to_vram( + 131072, pool_mib, model_size, "f16", budget_frac = 0.80 + ) assert fit_tp < 131072, "expected the context to be capped at this VRAM tier" assert fit_tp <= fit_default, "a tighter budget must not allow MORE context" # Omitting the override must reproduce the default budget exactly. - assert backend._fit_context_to_vram(131072, pool_mib, model_size, "f16") == fit_default + assert ( + backend._fit_context_to_vram(131072, pool_mib, model_size, "f16") == fit_default + ) # ── unsupported-arch load failure -> clean message ─────────────────── @@ -326,12 +334,16 @@ def _plan( mtp = False, ): b = _kv_seeded_backend() - return b, b._plan_tensor_parallel(gpus, int(model_gb * _GB), target, mtp_engaged = mtp) + return b, b._plan_tensor_parallel( + gpus, int(model_gb * _GB), target, mtp_engaged = mtp + ) def _kv_budget_b(model_gb, gpus = _ASYM): reserve = LlamaCppBackend._TENSOR_PARALLEL_BUFFER_RESERVE_MIB - return (sum(f for _, f in gpus) - len(gpus) * reserve) * 1024 * 1024 - int(model_gb * _GB) + return (sum(f for _, f in gpus) - len(gpus) * reserve) * 1024 * 1024 - int( + model_gb * _GB + ) def test_tp_plan_weighted_split_on_asymmetric_big_model(): @@ -396,7 +408,9 @@ def test_tp_plan_max_available_ctx_reports_native_not_explicit_ctx(): # An explicit small ctx caps effective_ctx but the UI ceiling # (max_available_ctx) must reflect the native/hardware cap, not the request. b = _kv_seeded_backend() - ec, mac, _gi, _ts = b._plan_tensor_parallel(_ASYM, int(50 * _GB), 8192, max_target_ctx = 131072) + ec, mac, _gi, _ts = b._plan_tensor_parallel( + _ASYM, int(50 * _GB), 8192, max_target_ctx = 131072 + ) _, native_mac, *_ = b._plan_tensor_parallel(_ASYM, int(50 * _GB), 131072) assert ec == 8192 # explicit request honored for the load assert mac == native_mac > ec # ceiling reflects the hardware cap @@ -436,7 +450,9 @@ def test_tp_plan_drops_gpu_below_buffer_reserve(): # split (and gpu_indices reflects only the usable device). b = _kv_seeded_backend() reserve = LlamaCppBackend._TENSOR_PARALLEL_BUFFER_RESERVE_MIB - ec, mac, gi, ts = b._plan_tensor_parallel([(0, 48000), (1, reserve - 1)], int(8 * _GB), 8192) + ec, mac, gi, ts = b._plan_tensor_parallel( + [(0, 48000), (1, reserve - 1)], int(8 * _GB), 8192 + ) assert gi == [0] assert ts is None @@ -458,7 +474,9 @@ class _RecordingLoader: self.calls: list[tuple] = [] async def __call__(self, tensor_parallel, extra_args): - self.calls.append((tensor_parallel, list(extra_args) if extra_args else extra_args)) + self.calls.append( + (tensor_parallel, list(extra_args) if extra_args else extra_args) + ) if resolve_tensor_parallel(extra_args, tensor_parallel): raise RuntimeError("llama-server failed to start") return True @@ -467,7 +485,9 @@ class _RecordingLoader: def test_tensor_fallback_retries_layer_on_crash(): loader = _RecordingLoader() ok = asyncio.run( - load_with_tensor_fallback(loader, requested_tensor = True, extra_args = None, label = "m") + load_with_tensor_fallback( + loader, requested_tensor = True, extra_args = None, label = "m" + ) ) assert ok is True # tensor first (crashes), then layer split. @@ -482,7 +502,9 @@ def test_tensor_fallback_no_retry_on_success(): return True ok = asyncio.run( - load_with_tensor_fallback(_ok, requested_tensor = True, extra_args = None, label = "m") + load_with_tensor_fallback( + _ok, requested_tensor = True, extra_args = None, label = "m" + ) ) assert ok is True assert calls == [True] # no fallback when the tensor load succeeds @@ -516,7 +538,9 @@ def test_tensor_fallback_returns_false_when_both_attempts_fail(): return False ok = asyncio.run( - load_with_tensor_fallback(_always_false, requested_tensor = True, extra_args = None, label = "m") + load_with_tensor_fallback( + _always_false, requested_tensor = True, extra_args = None, label = "m" + ) ) assert ok is False assert calls == [True, False] # tried tensor, then layer split @@ -559,7 +583,9 @@ def test_tensor_fallback_strips_split_mode_from_extras_on_retry(extras): # else resolve_tensor_parallel re-enables tensor and relaunches the crash. loader = _RecordingLoader() ok = asyncio.run( - load_with_tensor_fallback(loader, requested_tensor = False, extra_args = extras, label = "m") + load_with_tensor_fallback( + loader, requested_tensor = False, extra_args = extras, label = "m" + ) ) assert ok is True assert len(loader.calls) == 2 diff --git a/studio/backend/tests/test_tool_approvals.py b/studio/backend/tests/test_tool_approvals.py index af792e652c..9a5af18893 100644 --- a/studio/backend/tests/test_tool_approvals.py +++ b/studio/backend/tests/test_tool_approvals.py @@ -246,7 +246,9 @@ def test_concurrent_distinct_calls_route_their_own_decisions(): for i in range(n): aid = new_approval_id() waiters[aid] = _Waiter(f"s{i}", aid).start() - expected = {aid: ("allow" if i % 2 == 0 else "deny") for i, aid in enumerate(waiters)} + expected = { + aid: ("allow" if i % 2 == 0 else "deny") for i, aid in enumerate(waiters) + } for aid, decision in expected.items(): assert resolve_tool_decision(aid, decision) is True for aid, w in waiters.items(): diff --git a/studio/backend/tests/test_tool_call_parser_strict.py b/studio/backend/tests/test_tool_call_parser_strict.py index 8ff41342d7..2f7e7eb8a7 100644 --- a/studio/backend/tests/test_tool_call_parser_strict.py +++ b/studio/backend/tests/test_tool_call_parser_strict.py @@ -40,7 +40,9 @@ class TestFunctionStyleTrailingText: assert call == {"name": "web_search", "arguments": {"query": "weather london"}} def test_closed_function_with_trailing_whitespace_is_accepted(self): - text = "cats \n\n" + text = ( + "cats \n\n" + ) call = _only(text) assert call == {"name": "web_search", "arguments": {"query": "cats"}} diff --git a/studio/backend/tests/test_tool_confirm_loop.py b/studio/backend/tests/test_tool_confirm_loop.py index ce7852c95f..9e3aa90e2e 100644 --- a/studio/backend/tests/test_tool_confirm_loop.py +++ b/studio/backend/tests/test_tool_confirm_loop.py @@ -99,7 +99,9 @@ def _drive( if ev["type"] == "tool_start" and ev.get("awaiting_confirmation"): # Slot is already registered (begin ran before this yield), so # the decision lands before the loop enters its blocking wait. - resolve_tool_decision(ev["approval_id"], next(decision_iter), session_id = _SESSION) + resolve_tool_decision( + ev["approval_id"], next(decision_iter), session_id = _SESSION + ) return events, exec_fn.calls diff --git a/studio/backend/tests/test_tool_confirm_stream.py b/studio/backend/tests/test_tool_confirm_stream.py index b8e0472e12..986f0cb2fd 100644 --- a/studio/backend/tests/test_tool_confirm_stream.py +++ b/studio/backend/tests/test_tool_confirm_stream.py @@ -69,7 +69,9 @@ def _build_app() -> FastAPI: "approval_id": approval_id, "awaiting_confirmation": True, } - denied = wait_tool_decision(slot, approval_id, cancel_event = cancel_event) == "deny" + denied = ( + wait_tool_decision(slot, approval_id, cancel_event = cancel_event) == "deny" + ) result = TOOL_REJECTED_MESSAGE if denied else _EXECUTED_RESULT yield {"type": "tool_end", "tool_name": "python", "result": result} @@ -116,7 +118,9 @@ class _Server: def __init__(self, app): self.port = _free_port() - config = uvicorn.Config(app, host = "127.0.0.1", port = self.port, log_level = "warning") + config = uvicorn.Config( + app, host = "127.0.0.1", port = self.port, log_level = "warning" + ) self.server = uvicorn.Server(config) self._thread = threading.Thread(target = self.server.run, daemon = True) @@ -159,7 +163,9 @@ async def _drive(base_url, session_id, decision): resolved = None timeout = httpx.Timeout(10.0) async with httpx.AsyncClient(base_url = base_url, timeout = timeout) as client: - async with client.stream("POST", "/stream", json = {"session_id": session_id}) as resp: + async with client.stream( + "POST", "/stream", json = {"session_id": session_id} + ) as resp: assert resp.status_code == 200 async for line in resp.aiter_lines(): if not line.startswith("data: "): diff --git a/studio/backend/tests/test_tool_loop_controller.py b/studio/backend/tests/test_tool_loop_controller.py index dea5de6d6e..6783a59af6 100644 --- a/studio/backend/tests/test_tool_loop_controller.py +++ b/studio/backend/tests/test_tool_loop_controller.py @@ -69,7 +69,10 @@ def test_status_and_provenance_match_local_event_conventions(): status_for_tool("web_search", {"url": "https://www.example.com/a"}) == "Reading: example.com" ) - assert status_for_tool("python", {"code": "print(1)\nprint(2)"}) == "Running Python: print(1)" + assert ( + status_for_tool("python", {"code": "print(1)\nprint(2)"}) + == "Running Python: print(1)" + ) assert tool_event_provenance(healed = True, forced = False, provisional = None) == { "source": "local", "healed": True, @@ -85,7 +88,10 @@ def test_prepare_execute_builds_visible_events_and_model_tool_message(): assert decision.status_text == "Searching: gpu prices" assert decision.tool_start_payload()["arguments"] == {"query": "gpu prices"} assert decision.tool_start_event()["type"] == "tool_start" - assert decision.as_assistant_tool_call()["function"]["arguments"] == '{"query":"gpu prices"}' + assert ( + decision.as_assistant_tool_call()["function"]["arguments"] + == '{"query":"gpu prices"}' + ) completion = controller.record_result(decision, "Search result\n__IMAGES__:{...}") @@ -101,10 +107,14 @@ def test_prepare_execute_builds_visible_events_and_model_tool_message(): def test_successful_duplicate_is_internal_noop_and_keeps_remaining_tools(): controller = ToolLoopController(tools = [_tool("web_search"), _tool("python")]) - first = controller.prepare_call(_call("web_search", {"query": "gpu prices"}, "call_a")) + first = controller.prepare_call( + _call("web_search", {"query": "gpu prices"}, "call_a") + ) controller.record_result(first, "ok") - duplicate = controller.prepare_call(_call("web_search", {"query": "gpu prices"}, "call_b")) + duplicate = controller.prepare_call( + _call("web_search", {"query": "gpu prices"}, "call_b") + ) completion = controller.record_noop(duplicate) assert duplicate.action == "duplicate" @@ -123,10 +133,14 @@ def test_successful_duplicate_is_internal_noop_and_keeps_remaining_tools(): def test_repeated_successful_duplicate_becomes_terminal_after_one_recovery_nudge(): controller = ToolLoopController(tools = [_tool("web_search"), _tool("python")]) - first = controller.prepare_call(_call("web_search", {"query": "gpu prices"}, "call_a")) + first = controller.prepare_call( + _call("web_search", {"query": "gpu prices"}, "call_a") + ) controller.record_result(first, "ok") - duplicate_one = controller.prepare_call(_call("web_search", {"query": "gpu prices"}, "call_b")) + duplicate_one = controller.prepare_call( + _call("web_search", {"query": "gpu prices"}, "call_b") + ) completion_one = controller.record_noop(duplicate_one) assert duplicate_one.action == "duplicate" @@ -137,7 +151,9 @@ def test_repeated_successful_duplicate_becomes_terminal_after_one_recovery_nudge "python", ] - duplicate_two = controller.prepare_call(_call("web_search", {"query": "gpu prices"}, "call_c")) + duplicate_two = controller.prepare_call( + _call("web_search", {"query": "gpu prices"}, "call_c") + ) completion_two = controller.record_noop(duplicate_two) assert duplicate_two.action == "duplicate" @@ -190,12 +206,16 @@ def test_render_html_success_filters_active_tools_and_repeat_is_internal(): "web_search", ] - first = controller.prepare_call(_call("render_html", {"code": ""}, "call_html_1")) + first = controller.prepare_call( + _call("render_html", {"code": ""}, "call_html_1") + ) controller.record_result(first, "Rendered HTML artifact: Demo") assert [t["function"]["name"] for t in controller.active_tools()] == ["web_search"] - repeat = controller.prepare_call(_call("render_html", {"code": ""}, "call_html_2")) + repeat = controller.prepare_call( + _call("render_html", {"code": ""}, "call_html_2") + ) completion = controller.record_noop(repeat) assert repeat.action == "render_html_repeat" diff --git a/studio/backend/tests/test_tool_xml_strip.py b/studio/backend/tests/test_tool_xml_strip.py index 2ba3310fbe..4e3350144f 100644 --- a/studio/backend/tests/test_tool_xml_strip.py +++ b/studio/backend/tests/test_tool_xml_strip.py @@ -43,7 +43,9 @@ _strip_tool_xml_for_display = _ns["_strip_tool_xml_for_display"] def test_route_display_strip_respects_disabled_auto_heal_contract(): text = 'literal {"name":"web_search"} survives' assert _strip_tool_xml_for_display(text, auto_heal_tool_calls = False) == text - assert "" not in _strip_tool_xml_for_display(text, auto_heal_tool_calls = True) + assert "" not in _strip_tool_xml_for_display( + text, auto_heal_tool_calls = True + ) def test_strips_well_formed_tool_call(): @@ -91,7 +93,9 @@ def test_strips_orphan_tool_call_no_close(): def test_strips_orphan_function_no_close(): - text = "I'll call python:\n\n\nprint(1)\n" + text = ( + "I'll call python:\n\n\nprint(1)\n" + ) cleaned = _TOOL_XML_RE.sub("", text) assert "") self.assertEqual(result.dataset[1]["text"], "world") self.assertTrue( - any("null or non-string 'text' values" in notice.message for notice in result.notices) + any( + "null or non-string 'text' values" in notice.message + for notice in result.notices + ) ) diff --git a/studio/backend/tests/test_training_resume.py b/studio/backend/tests/test_training_resume.py index 91fdac9961..40df3c8f4b 100644 --- a/studio/backend/tests/test_training_resume.py +++ b/studio/backend/tests/test_training_resume.py @@ -67,7 +67,9 @@ def test_can_resume_run_rejects_s3_dataset_source(monkeypatch): def test_can_resume_run_rejects_s3_metadata_marker(monkeypatch): monkeypatch.setattr(resume, "has_resume_state", lambda _path: True) - run = _stopped_run(config_json = json.dumps({"s3_dataset": {"bucket": "training-data"}})) + run = _stopped_run( + config_json = json.dumps({"s3_dataset": {"bucket": "training-data"}}) + ) assert resume.can_resume_run(run) is False @@ -77,7 +79,9 @@ def test_list_runs_includes_config_json_for_resume_policy(monkeypatch, tmp_path) monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path)) monkeypatch.setattr(studio_db, "_schema_ready", False) - config_json = json.dumps({"dataset_source": "s3", "s3_dataset": {"bucket": "training-data"}}) + config_json = json.dumps( + {"dataset_source": "s3", "s3_dataset": {"bucket": "training-data"}} + ) studio_db.create_run( id = "run-s3", diff --git a/studio/backend/tests/test_training_worker_flash_attn.py b/studio/backend/tests/test_training_worker_flash_attn.py index 3c5d6cd094..363bf133a4 100644 --- a/studio/backend/tests/test_training_worker_flash_attn.py +++ b/studio/backend/tests/test_training_worker_flash_attn.py @@ -49,7 +49,9 @@ def _missing_module_import(missing: str): def test_should_try_runtime_flash_attn_install_threshold_and_skip(monkeypatch): monkeypatch.delenv(worker._FLASH_ATTN_SKIP_ENV, raising = False) assert worker._should_try_runtime_flash_attn_install(32767) is False - assert worker._should_try_runtime_flash_attn_install(32768) is sys.platform.startswith("linux") + assert worker._should_try_runtime_flash_attn_install( + 32768 + ) is sys.platform.startswith("linux") monkeypatch.setenv(worker._FLASH_ATTN_SKIP_ENV, "1") assert worker._should_try_runtime_flash_attn_install(32768) is False @@ -146,7 +148,9 @@ def test_runtime_flash_attn_skips_on_blackwell(monkeypatch): install_mock = mock.Mock() monkeypatch.delenv(worker._FLASH_ATTN_SKIP_ENV, raising = False) - monkeypatch.setattr(worker, "_should_try_runtime_flash_attn_install", lambda max_seq: True) + monkeypatch.setattr( + worker, "_should_try_runtime_flash_attn_install", lambda max_seq: True + ) monkeypatch.setattr(worker, "has_blackwell_gpu", lambda: True) monkeypatch.setattr(worker, "_install_package_wheel_first", install_mock) monkeypatch.setattr( @@ -512,10 +516,14 @@ def test_tilelang_backend_reinstalls_when_tvm_ffi_is_broken(monkeypatch): # Repair: --force-reinstall --no-deps, apache-tvm-ffi ONLY. assert "--force-reinstall" in repair_args - assert "--no-deps" in repair_args, "Repair MUST use --no-deps to avoid replacing torch / CUDA" + assert ( + "--no-deps" in repair_args + ), "Repair MUST use --no-deps to avoid replacing torch / CUDA" assert "--only-binary=:all:" in repair_args assert f"apache-tvm-ffi=={worker._APACHE_TVM_FFI_PACKAGE_VERSION}" in repair_args - assert all("tilelang" not in a for a in repair_args), "Repair MUST only touch apache-tvm-ffi" + assert all( + "tilelang" not in a for a in repair_args + ), "Repair MUST only touch apache-tvm-ffi" # Install: regular dep-resolving install, no --force-reinstall. assert "--force-reinstall" not in install_args @@ -689,12 +697,16 @@ def test_hook_installs_when_gate_returns_false(monkeypatch): conv_install = mock.Mock(side_effect = _conv_install_side_effect) - monkeypatch.setattr(worker, "_ensure_flash_linear_attention_unconditional", fla_install) + monkeypatch.setattr( + worker, "_ensure_flash_linear_attention_unconditional", fla_install + ) monkeypatch.setattr(worker, "_ensure_tilelang_backend_unconditional", tile_install) monkeypatch.setattr(worker, "_install_package_wheel_first", conv_install) monkeypatch.delenv(worker._FAST_PATH_HOOKS_SKIP_ENV, raising = False) - worker._install_fast_path_hooks(event_queue = _FakeQueue(), model_name = "unsloth/Qwen3.5-2B") + worker._install_fast_path_hooks( + event_queue = _FakeQueue(), model_name = "unsloth/Qwen3.5-2B" + ) from transformers.utils import import_utils as _iu @@ -718,7 +730,9 @@ def test_hook_skips_install_when_gate_already_true(monkeypatch): fla_install = mock.Mock() tile_install = mock.Mock() conv_install = mock.Mock() - monkeypatch.setattr(worker, "_ensure_flash_linear_attention_unconditional", fla_install) + monkeypatch.setattr( + worker, "_ensure_flash_linear_attention_unconditional", fla_install + ) monkeypatch.setattr(worker, "_ensure_tilelang_backend_unconditional", tile_install) monkeypatch.setattr(worker, "_install_package_wheel_first", conv_install) # Tilelang healthy -> post_available path is a no-op (otherwise it @@ -727,7 +741,9 @@ def test_hook_skips_install_when_gate_already_true(monkeypatch): monkeypatch.setattr(worker, "_installed_tvm_ffi_version", lambda: "0.1.9") monkeypatch.delenv(worker._FAST_PATH_HOOKS_SKIP_ENV, raising = False) - worker._install_fast_path_hooks(event_queue = _FakeQueue(), model_name = "unsloth/Qwen3.5-2B") + worker._install_fast_path_hooks( + event_queue = _FakeQueue(), model_name = "unsloth/Qwen3.5-2B" + ) from transformers.utils import import_utils as _iu @@ -755,12 +771,16 @@ def test_hook_idempotent_on_repeat_call(monkeypatch): return True conv_install = mock.Mock(side_effect = _conv_install_side_effect) - monkeypatch.setattr(worker, "_ensure_flash_linear_attention_unconditional", fla_install) + monkeypatch.setattr( + worker, "_ensure_flash_linear_attention_unconditional", fla_install + ) monkeypatch.setattr(worker, "_ensure_tilelang_backend_unconditional", tile_install) monkeypatch.setattr(worker, "_install_package_wheel_first", conv_install) monkeypatch.delenv(worker._FAST_PATH_HOOKS_SKIP_ENV, raising = False) - worker._install_fast_path_hooks(event_queue = _FakeQueue(), model_name = "unsloth/Qwen3.5-2B") + worker._install_fast_path_hooks( + event_queue = _FakeQueue(), model_name = "unsloth/Qwen3.5-2B" + ) from transformers.utils import import_utils as _iu @@ -781,12 +801,18 @@ def test_hook_handles_install_failure_gracefully(monkeypatch): def raising_install(eq): raise RuntimeError("pip failed to fetch wheel") - monkeypatch.setattr(worker, "_ensure_flash_linear_attention_unconditional", raising_install) - monkeypatch.setattr(worker, "_ensure_tilelang_backend_unconditional", lambda eq: None) + monkeypatch.setattr( + worker, "_ensure_flash_linear_attention_unconditional", raising_install + ) + monkeypatch.setattr( + worker, "_ensure_tilelang_backend_unconditional", lambda eq: None + ) monkeypatch.setattr(worker, "_install_package_wheel_first", lambda **kw: None) monkeypatch.delenv(worker._FAST_PATH_HOOKS_SKIP_ENV, raising = False) - worker._install_fast_path_hooks(event_queue = _FakeQueue(), model_name = "unsloth/Qwen3.5-2B") + worker._install_fast_path_hooks( + event_queue = _FakeQueue(), model_name = "unsloth/Qwen3.5-2B" + ) from transformers.utils import import_utils as _iu @@ -800,10 +826,14 @@ def test_hook_can_be_disabled_via_env(monkeypatch): _patch_iu_gates(monkeypatch, fla_gate, conv_gate) fla_install = mock.Mock() - monkeypatch.setattr(worker, "_ensure_flash_linear_attention_unconditional", fla_install) + monkeypatch.setattr( + worker, "_ensure_flash_linear_attention_unconditional", fla_install + ) monkeypatch.setenv(worker._FAST_PATH_HOOKS_SKIP_ENV, "1") - worker._install_fast_path_hooks(event_queue = _FakeQueue(), model_name = "unsloth/Qwen3.5-2B") + worker._install_fast_path_hooks( + event_queue = _FakeQueue(), model_name = "unsloth/Qwen3.5-2B" + ) from transformers.utils import import_utils as _iu @@ -818,12 +848,18 @@ def test_hook_clears_lru_cache_before_first_check(monkeypatch): conv_gate = _make_fake_gate(initial_return = True) _patch_iu_gates(monkeypatch, fla_gate, conv_gate) - monkeypatch.setattr(worker, "_ensure_flash_linear_attention_unconditional", lambda eq: None) - monkeypatch.setattr(worker, "_ensure_tilelang_backend_unconditional", lambda eq: None) + monkeypatch.setattr( + worker, "_ensure_flash_linear_attention_unconditional", lambda eq: None + ) + monkeypatch.setattr( + worker, "_ensure_tilelang_backend_unconditional", lambda eq: None + ) monkeypatch.setattr(worker, "_install_package_wheel_first", lambda **kw: None) monkeypatch.delenv(worker._FAST_PATH_HOOKS_SKIP_ENV, raising = False) - worker._install_fast_path_hooks(event_queue = _FakeQueue(), model_name = "unsloth/Qwen3.5-2B") + worker._install_fast_path_hooks( + event_queue = _FakeQueue(), model_name = "unsloth/Qwen3.5-2B" + ) from transformers.utils import import_utils as _iu _iu.is_flash_linear_attention_available() @@ -850,12 +886,18 @@ def test_hook_rewrites_previously_imported_module_bindings(monkeypatch): fla_gate.next_return = True return True - monkeypatch.setattr(worker, "_ensure_flash_linear_attention_unconditional", fake_install) - monkeypatch.setattr(worker, "_ensure_tilelang_backend_unconditional", lambda eq: True) + monkeypatch.setattr( + worker, "_ensure_flash_linear_attention_unconditional", fake_install + ) + monkeypatch.setattr( + worker, "_ensure_tilelang_backend_unconditional", lambda eq: True + ) monkeypatch.setattr(worker, "_install_package_wheel_first", lambda **kw: True) monkeypatch.delenv(worker._FAST_PATH_HOOKS_SKIP_ENV, raising = False) - worker._install_fast_path_hooks(event_queue = _FakeQueue(), model_name = "unsloth/Qwen3.5-2B") + worker._install_fast_path_hooks( + event_queue = _FakeQueue(), model_name = "unsloth/Qwen3.5-2B" + ) # The fake module's local binding is rewritten to the wrapper. assert fake_mod.is_flash_linear_attention_available is not fla_gate @@ -879,20 +921,30 @@ def test_hook_skips_when_import_utils_unavailable(monkeypatch): monkeypatch.delenv(worker._FAST_PATH_HOOKS_SKIP_ENV, raising = False) # Should not raise. - worker._install_fast_path_hooks(event_queue = _FakeQueue(), model_name = "unsloth/Qwen3.5-2B") + worker._install_fast_path_hooks( + event_queue = _FakeQueue(), model_name = "unsloth/Qwen3.5-2B" + ) def test_substring_fallback_unchanged_when_hook_skipped(monkeypatch): """Hook disabled -> legacy gate falls back to auto-discovered types.""" install_mock = mock.Mock() - monkeypatch.setattr(worker, "_ensure_flash_linear_attention_unconditional", install_mock) - monkeypatch.setattr(worker, "_discover_fla_model_types", lambda: frozenset({"qwen3_5"})) + monkeypatch.setattr( + worker, "_ensure_flash_linear_attention_unconditional", install_mock + ) + monkeypatch.setattr( + worker, "_discover_fla_model_types", lambda: frozenset({"qwen3_5"}) + ) monkeypatch.setenv(worker._FAST_PATH_HOOKS_SKIP_ENV, "1") - worker._ensure_flash_linear_attention(event_queue = [], model_name = "unsloth/Qwen3.5-2B") + worker._ensure_flash_linear_attention( + event_queue = [], model_name = "unsloth/Qwen3.5-2B" + ) assert install_mock.call_count == 1 - worker._ensure_flash_linear_attention(event_queue = [], model_name = "meta-llama/Llama-3.1-8B") + worker._ensure_flash_linear_attention( + event_queue = [], model_name = "meta-llama/Llama-3.1-8B" + ) assert install_mock.call_count == 1 @@ -920,9 +972,13 @@ def test_hook_does_not_install_tilelang_for_model_outside_allowlist(monkeypatch) fla_install = mock.Mock(side_effect = _fla_install) tile_install = mock.Mock(return_value = True) - monkeypatch.setattr(worker, "_ensure_flash_linear_attention_unconditional", fla_install) + monkeypatch.setattr( + worker, "_ensure_flash_linear_attention_unconditional", fla_install + ) monkeypatch.setattr(worker, "_ensure_tilelang_backend_unconditional", tile_install) - monkeypatch.setattr(worker, "_install_package_wheel_first", mock.Mock(return_value = True)) + monkeypatch.setattr( + worker, "_install_package_wheel_first", mock.Mock(return_value = True) + ) monkeypatch.delenv(worker._FAST_PATH_HOOKS_SKIP_ENV, raising = False) # Hermetize the auto-discovered set so the test stays valid as new # transformers releases add FLA-using model_types (eg olmo_hybrid in @@ -957,12 +1013,18 @@ def test_hook_does_install_tilelang_for_qwen35(monkeypatch): fla_install = mock.Mock(side_effect = _fla_install) tile_install = mock.Mock(return_value = True) - monkeypatch.setattr(worker, "_ensure_flash_linear_attention_unconditional", fla_install) + monkeypatch.setattr( + worker, "_ensure_flash_linear_attention_unconditional", fla_install + ) monkeypatch.setattr(worker, "_ensure_tilelang_backend_unconditional", tile_install) - monkeypatch.setattr(worker, "_install_package_wheel_first", mock.Mock(return_value = True)) + monkeypatch.setattr( + worker, "_install_package_wheel_first", mock.Mock(return_value = True) + ) monkeypatch.delenv(worker._FAST_PATH_HOOKS_SKIP_ENV, raising = False) - worker._install_fast_path_hooks(event_queue = _FakeQueue(), model_name = "unsloth/Qwen3.5-2B") + worker._install_fast_path_hooks( + event_queue = _FakeQueue(), model_name = "unsloth/Qwen3.5-2B" + ) from transformers.utils import import_utils as _iu @@ -1011,14 +1073,20 @@ def test_hook_trusts_installer_bool_not_metadata(monkeypatch): return False # but deep import is broken fake_fla_install = mock.Mock(side_effect = _bad_install) - monkeypatch.setattr(worker, "_ensure_flash_linear_attention_unconditional", fake_fla_install) + monkeypatch.setattr( + worker, "_ensure_flash_linear_attention_unconditional", fake_fla_install + ) monkeypatch.setattr( worker, "_ensure_tilelang_backend_unconditional", mock.Mock(return_value = True) ) - monkeypatch.setattr(worker, "_install_package_wheel_first", mock.Mock(return_value = True)) + monkeypatch.setattr( + worker, "_install_package_wheel_first", mock.Mock(return_value = True) + ) monkeypatch.delenv(worker._FAST_PATH_HOOKS_SKIP_ENV, raising = False) - worker._install_fast_path_hooks(event_queue = _FakeQueue(), model_name = "unsloth/Qwen3.5-2B") + worker._install_fast_path_hooks( + event_queue = _FakeQueue(), model_name = "unsloth/Qwen3.5-2B" + ) from transformers.utils import import_utils as _iu @@ -1071,10 +1139,14 @@ def test_hook_skips_tilelang_when_fla_install_is_skipped(monkeypatch): monkeypatch.setenv(worker._FLA_SKIP_ENV, "1") tile_install = mock.Mock(return_value = True) monkeypatch.setattr(worker, "_ensure_tilelang_backend_unconditional", tile_install) - monkeypatch.setattr(worker, "_install_package_wheel_first", mock.Mock(return_value = True)) + monkeypatch.setattr( + worker, "_install_package_wheel_first", mock.Mock(return_value = True) + ) monkeypatch.delenv(worker._FAST_PATH_HOOKS_SKIP_ENV, raising = False) - worker._install_fast_path_hooks(event_queue = _FakeQueue(), model_name = "unsloth/Qwen3.5-2B") + worker._install_fast_path_hooks( + event_queue = _FakeQueue(), model_name = "unsloth/Qwen3.5-2B" + ) from transformers.utils import import_utils as _iu @@ -1094,15 +1166,21 @@ def test_hook_runs_tilelang_repair_when_fla_already_true(monkeypatch): fla_install = mock.Mock(return_value = True) tile_install = mock.Mock(return_value = True) - monkeypatch.setattr(worker, "_ensure_flash_linear_attention_unconditional", fla_install) + monkeypatch.setattr( + worker, "_ensure_flash_linear_attention_unconditional", fla_install + ) monkeypatch.setattr(worker, "_ensure_tilelang_backend_unconditional", tile_install) - monkeypatch.setattr(worker, "_install_package_wheel_first", mock.Mock(return_value = True)) + monkeypatch.setattr( + worker, "_install_package_wheel_first", mock.Mock(return_value = True) + ) # tilelang missing AND tvm-ffi on broken list — both trigger repair. monkeypatch.setattr(worker, "_tilelang_importable", lambda: False) monkeypatch.setattr(worker, "_installed_tvm_ffi_version", lambda: "0.1.11") monkeypatch.delenv(worker._FAST_PATH_HOOKS_SKIP_ENV, raising = False) - worker._install_fast_path_hooks(event_queue = _FakeQueue(), model_name = "unsloth/Qwen3.5-2B") + worker._install_fast_path_hooks( + event_queue = _FakeQueue(), model_name = "unsloth/Qwen3.5-2B" + ) from transformers.utils import import_utils as _iu @@ -1199,11 +1277,17 @@ def test_install_fast_path_hooks_sets_fla_tilelang_zero_on_hip(monkeypatch): monkeypatch.delenv("FLA_TILELANG", raising = False) monkeypatch.delenv(worker._FAST_PATH_HOOKS_SKIP_ENV, raising = False) monkeypatch.setattr(worker, "_torch_has_hip", lambda: True) - monkeypatch.setattr(worker, "_ensure_flash_linear_attention_unconditional", lambda eq: True) - monkeypatch.setattr(worker, "_ensure_tilelang_backend_unconditional", lambda eq: True) + monkeypatch.setattr( + worker, "_ensure_flash_linear_attention_unconditional", lambda eq: True + ) + monkeypatch.setattr( + worker, "_ensure_tilelang_backend_unconditional", lambda eq: True + ) monkeypatch.setattr(worker, "_install_package_wheel_first", lambda **kw: True) - worker._install_fast_path_hooks(event_queue = _FakeQueue(), model_name = "unsloth/Qwen3.5-2B") + worker._install_fast_path_hooks( + event_queue = _FakeQueue(), model_name = "unsloth/Qwen3.5-2B" + ) assert _os.environ.get("FLA_TILELANG") == "0" @@ -1217,11 +1301,17 @@ def test_install_fast_path_hooks_respects_user_fla_tilelang_override(monkeypatch monkeypatch.setenv("FLA_TILELANG", "1") monkeypatch.delenv(worker._FAST_PATH_HOOKS_SKIP_ENV, raising = False) monkeypatch.setattr(worker, "_torch_has_hip", lambda: True) - monkeypatch.setattr(worker, "_ensure_flash_linear_attention_unconditional", lambda eq: True) - monkeypatch.setattr(worker, "_ensure_tilelang_backend_unconditional", lambda eq: True) + monkeypatch.setattr( + worker, "_ensure_flash_linear_attention_unconditional", lambda eq: True + ) + monkeypatch.setattr( + worker, "_ensure_tilelang_backend_unconditional", lambda eq: True + ) monkeypatch.setattr(worker, "_install_package_wheel_first", lambda **kw: True) - worker._install_fast_path_hooks(event_queue = _FakeQueue(), model_name = "unsloth/Qwen3.5-2B") + worker._install_fast_path_hooks( + event_queue = _FakeQueue(), model_name = "unsloth/Qwen3.5-2B" + ) assert _os.environ["FLA_TILELANG"] == "1" @@ -1233,11 +1323,17 @@ def test_install_fast_path_hooks_does_not_set_fla_tilelang_on_cuda(monkeypatch): monkeypatch.delenv("FLA_TILELANG", raising = False) monkeypatch.delenv(worker._FAST_PATH_HOOKS_SKIP_ENV, raising = False) monkeypatch.setattr(worker, "_torch_has_hip", lambda: False) - monkeypatch.setattr(worker, "_ensure_flash_linear_attention_unconditional", lambda eq: True) - monkeypatch.setattr(worker, "_ensure_tilelang_backend_unconditional", lambda eq: True) + monkeypatch.setattr( + worker, "_ensure_flash_linear_attention_unconditional", lambda eq: True + ) + monkeypatch.setattr( + worker, "_ensure_tilelang_backend_unconditional", lambda eq: True + ) monkeypatch.setattr(worker, "_install_package_wheel_first", lambda **kw: True) - worker._install_fast_path_hooks(event_queue = _FakeQueue(), model_name = "unsloth/Qwen3.5-2B") + worker._install_fast_path_hooks( + event_queue = _FakeQueue(), model_name = "unsloth/Qwen3.5-2B" + ) assert _os.environ.get("FLA_TILELANG") is None @@ -1247,7 +1343,9 @@ def test_install_fast_path_hooks_does_not_set_fla_tilelang_on_cuda(monkeypatch): # ─────────────────────────────────────────────────────────────────── -def _make_fake_transformers_tree(tmp_path, fla_types: list[str], non_fla_types: list[str]): +def _make_fake_transformers_tree( + tmp_path, fla_types: list[str], non_fla_types: list[str] +): """Lay out tmp dir as `transformers/models/{type}/modeling_{type}.py`.""" pkg = tmp_path / "transformers" models = pkg / "models" @@ -1290,7 +1388,9 @@ def test_discover_fla_model_types_returns_only_fla_users(tmp_path, monkeypatch): def test_discover_fla_model_types_caches_across_calls(tmp_path, monkeypatch): - pkg = _make_fake_transformers_tree(tmp_path, fla_types = ["qwen3_5"], non_fla_types = []) + pkg = _make_fake_transformers_tree( + tmp_path, fla_types = ["qwen3_5"], non_fla_types = [] + ) fake = mock.MagicMock(__file__ = str(pkg / "__init__.py")) monkeypatch.setitem(sys.modules, "transformers", fake) _reset_fla_cache(monkeypatch) @@ -1336,7 +1436,9 @@ def test_discover_fla_model_types_handles_missing_transformers(monkeypatch): def test_discover_fla_model_types_handles_unreadable_file(tmp_path, monkeypatch): - pkg = _make_fake_transformers_tree(tmp_path, fla_types = ["qwen3_5"], non_fla_types = []) + pkg = _make_fake_transformers_tree( + tmp_path, fla_types = ["qwen3_5"], non_fla_types = [] + ) fake = mock.MagicMock(__file__ = str(pkg / "__init__.py")) monkeypatch.setitem(sys.modules, "transformers", fake) _reset_fla_cache(monkeypatch) @@ -1382,7 +1484,9 @@ def test_model_wants_tilelang_empty_when_transformers_has_no_fla(monkeypatch): def test_model_wants_tilelang_normalizes_separators(monkeypatch): - monkeypatch.setattr(worker, "_discover_fla_model_types", lambda: frozenset({"qwen3_next"})) + monkeypatch.setattr( + worker, "_discover_fla_model_types", lambda: frozenset({"qwen3_next"}) + ) for variant in ( "qwen3-next", "Qwen3.Next", diff --git a/studio/backend/tests/test_transformers_version.py b/studio/backend/tests/test_transformers_version.py index 7c497ba1b1..4089ea98d8 100644 --- a/studio/backend/tests/test_transformers_version.py +++ b/studio/backend/tests/test_transformers_version.py @@ -480,7 +480,9 @@ class TestGetTransformersTier: return_value = False, ), ): - assert get_transformers_tier("mistralai/Ministral-3-8B-Instruct-2512") == "530" + assert ( + get_transformers_tier("mistralai/Ministral-3-8B-Instruct-2512") == "530" + ) def test_llama_returns_default(self): with ( diff --git a/studio/backend/tests/test_utils.py b/studio/backend/tests/test_utils.py index c66d56528a..49c7660226 100644 --- a/studio/backend/tests/test_utils.py +++ b/studio/backend/tests/test_utils.py @@ -187,7 +187,9 @@ class TestGetGpuMemoryInfo: # --- When a GPU IS available --- - @pytest.mark.skipif(_actual_device() == "cpu", reason = "No GPU available on this machine") + @pytest.mark.skipif( + _actual_device() == "cpu", reason = "No GPU available on this machine" + ) def test_gpu_available_fields(self): result = get_gpu_memory_info() assert result["available"] is True @@ -285,7 +287,9 @@ class TestLogGpuMemory: "free_gb": 14.0, } - with patch("utils.hardware.hardware.get_gpu_memory_info", return_value = fake_info): + with patch( + "utils.hardware.hardware.get_gpu_memory_info", return_value = fake_info + ): log_gpu_memory("unit-test") captured = capfd.readouterr() @@ -296,7 +300,9 @@ class TestLogGpuMemory: def test_logs_cpu_fallback_when_no_gpu(self, capfd): fake_info = {"available": False, "backend": "cpu"} - with patch("utils.hardware.hardware.get_gpu_memory_info", return_value = fake_info): + with patch( + "utils.hardware.hardware.get_gpu_memory_info", return_value = fake_info + ): log_gpu_memory("cpu-test") captured = capfd.readouterr() diff --git a/studio/backend/tests/test_vision_cache.py b/studio/backend/tests/test_vision_cache.py index d1bdec8449..d16ab29bcb 100644 --- a/studio/backend/tests/test_vision_cache.py +++ b/studio/backend/tests/test_vision_cache.py @@ -131,7 +131,9 @@ class TestLocalGgufVisionDetection: "utils.models.model_config._is_vision_model_subprocess", side_effect = AssertionError("GGUF must not use Transformers vision detection"), ) - def test_qwen36_gguf_with_mmproj_skips_transformers(self, mock_subprocess, tmp_path): + def test_qwen36_gguf_with_mmproj_skips_transformers( + self, mock_subprocess, tmp_path + ): model = tmp_path / "Qwen3.6-27B-UD-Q4_K_XL-MTP.gguf" model.write_bytes(b"") (tmp_path / "mmproj-F32.gguf").write_bytes(b"") @@ -143,7 +145,9 @@ class TestLocalGgufVisionDetection: "utils.models.model_config._is_vision_model_subprocess", side_effect = AssertionError("GGUF must not use Transformers vision detection"), ) - def test_direct_gguf_in_variant_subdir_finds_snapshot_mmproj(self, mock_subprocess, tmp_path): + def test_direct_gguf_in_variant_subdir_finds_snapshot_mmproj( + self, mock_subprocess, tmp_path + ): variant_dir = tmp_path / "BF16" variant_dir.mkdir() model = variant_dir / "Qwen3.6-27B-UD-Q4_K_XL-MTP.gguf" @@ -157,7 +161,9 @@ class TestLocalGgufVisionDetection: "utils.models.model_config._is_vision_model_subprocess", side_effect = AssertionError("GGUF must not use Transformers vision detection"), ) - def test_qwen36_gguf_without_mmproj_skips_transformers(self, mock_subprocess, tmp_path): + def test_qwen36_gguf_without_mmproj_skips_transformers( + self, mock_subprocess, tmp_path + ): model = tmp_path / "Qwen3.6-27B-UD-Q4_K_XL-MTP.gguf" model.write_bytes(b"") @@ -289,7 +295,9 @@ class TestVisionCacheDirectPath: @patch("utils.transformers_version.needs_transformers_5", return_value = False) @patch("utils.models.model_config.load_model_config") - def test_vision_config_attr_detected_and_cached(self, mock_load_config, mock_needs_t5): + def test_vision_config_attr_detected_and_cached( + self, mock_load_config, mock_needs_t5 + ): """Models with vision_config (LLaVA, Qwen2-VL, etc.) should be cached as True.""" cfg = MagicMock(spec = []) # strict: only explicitly set attrs exist cfg.model_type = "qwen2_vl" @@ -303,7 +311,9 @@ class TestVisionCacheDirectPath: @patch("utils.transformers_version.needs_transformers_5", return_value = False) @patch("utils.models.model_config.load_model_config") - def test_gemma4_model_type_detected_and_cached(self, mock_load_config, mock_needs_t5): + def test_gemma4_model_type_detected_and_cached( + self, mock_load_config, mock_needs_t5 + ): cfg = MagicMock(spec = []) cfg.model_type = "gemma4" cfg.architectures = ["Gemma4ForConditionalGeneration"] @@ -315,7 +325,9 @@ class TestVisionCacheDirectPath: @patch("utils.transformers_version.needs_transformers_5", return_value = False) @patch("utils.models.model_config.load_model_config") - def test_gemma4_audio_subconfig_not_detected_as_vision(self, mock_load_config, mock_needs_t5): + def test_gemma4_audio_subconfig_not_detected_as_vision( + self, mock_load_config, mock_needs_t5 + ): cfg = MagicMock(spec = []) cfg.model_type = "gemma4_audio" cfg.architectures = ["Gemma4AudioModel"] @@ -327,7 +339,9 @@ class TestVisionCacheDirectPath: @patch("utils.transformers_version.needs_transformers_5", return_value = False) @patch("utils.models.model_config.load_model_config") - def test_gemma4_text_subconfig_not_detected_as_vision(self, mock_load_config, mock_needs_t5): + def test_gemma4_text_subconfig_not_detected_as_vision( + self, mock_load_config, mock_needs_t5 + ): cfg = MagicMock(spec = []) cfg.model_type = "gemma4_text" cfg.architectures = ["Gemma4ForCausalLM"] @@ -472,10 +486,15 @@ class TestSubprocessScript: is True ) assert ( - inline_is_vlm(_C(model_type = "gemma4_text", architectures = ["Gemma4ForCausalLM"])) + inline_is_vlm( + _C(model_type = "gemma4_text", architectures = ["Gemma4ForCausalLM"]) + ) + is False + ) + assert ( + inline_is_vlm(_C(model_type = "llama", architectures = ["LlamaForCausalLM"])) is False ) - assert inline_is_vlm(_C(model_type = "llama", architectures = ["LlamaForCausalLM"])) is False # --------------------------------------------------------------------------- diff --git a/studio/backend/tests/test_vram_estimation.py b/studio/backend/tests/test_vram_estimation.py index 2def8738e2..d0e2b03623 100644 --- a/studio/backend/tests/test_vram_estimation.py +++ b/studio/backend/tests/test_vram_estimation.py @@ -316,8 +316,12 @@ class TestLoraParams(unittest.TestCase): self.assertLess(qv_only, all_mods) def test_moe_mlp_modules_scale_with_experts(self): - dense_lora = compute_lora_params(LLAMA_8B, 16, ["gate_proj", "up_proj", "down_proj"]) - moe_lora = compute_lora_params(MOE_CONFIG, 16, ["gate_proj", "up_proj", "down_proj"]) + dense_lora = compute_lora_params( + LLAMA_8B, 16, ["gate_proj", "up_proj", "down_proj"] + ) + moe_lora = compute_lora_params( + MOE_CONFIG, 16, ["gate_proj", "up_proj", "down_proj"] + ) ratio = moe_lora / dense_lora self.assertAlmostEqual(ratio, 8.0, delta = 0.5) @@ -334,8 +338,12 @@ class TestLoraParams(unittest.TestCase): self.assertGreater(moe_lora, dense_lora * 20) def test_attention_modules_same_for_moe(self): - dense_attn = compute_lora_params(LLAMA_8B, 16, ["q_proj", "k_proj", "v_proj", "o_proj"]) - moe_attn = compute_lora_params(MOE_CONFIG, 16, ["q_proj", "k_proj", "v_proj", "o_proj"]) + dense_attn = compute_lora_params( + LLAMA_8B, 16, ["q_proj", "k_proj", "v_proj", "o_proj"] + ) + moe_attn = compute_lora_params( + MOE_CONFIG, 16, ["q_proj", "k_proj", "v_proj", "o_proj"] + ) self.assertEqual(dense_attn, moe_attn) def test_all_linear_uses_default_text_modules(self): @@ -458,7 +466,9 @@ class TestActivationBytes(unittest.TestCase): def test_non_flash_attention_uses_quadratic_path(self): seq_len = 4096 - expected_quadratic = 1 * STRUCTURED_MIXED.num_attention_heads * seq_len * seq_len * 2 * 12.0 + expected_quadratic = ( + 1 * STRUCTURED_MIXED.num_attention_heads * seq_len * seq_len * 2 * 12.0 + ) for attention_implementation in ("eager", "unknown_impl", None): with self.subTest(attention_implementation = attention_implementation): non_flash = compute_activation_bytes( @@ -473,7 +483,9 @@ class TestActivationBytes(unittest.TestCase): def test_non_flash_attention_without_gc_scales_quadratic_path_by_layers(self): seq_len = 4096 - one_layer = 1 * STRUCTURED_MIXED.num_attention_heads * seq_len * seq_len * 2 * 12.0 + one_layer = ( + 1 * STRUCTURED_MIXED.num_attention_heads * seq_len * seq_len * 2 * 12.0 + ) non_flash = compute_activation_bytes( STRUCTURED_MIXED, 1, @@ -705,7 +717,9 @@ class TestEstimateTrainingVram(unittest.TestCase): ) v8 = estimate_training_vram(LLAMA_8B, opt8) v32 = estimate_training_vram(LLAMA_8B, opt32) - self.assertAlmostEqual(v32.optimizer_states / v8.optimizer_states, 1.5, delta = 0.1) + self.assertAlmostEqual( + v32.optimizer_states / v8.optimizer_states, 1.5, delta = 0.1 + ) def test_min_gpu_vram_treats_activations_as_per_gpu_fixed(self): config = TrainingVramConfig(training_method = "qlora", load_in_4bit = True) @@ -755,7 +769,9 @@ class TestEstimateTrainingVram(unittest.TestCase): optimizer = "adamw_8bit", load_in_4bit = False, ) - expected_floor = int(compute_model_weights_bytes(LLAMA_8B, "full", False) * 0.15) + expected_floor = int( + compute_model_weights_bytes(LLAMA_8B, "full", False) * 0.15 + ) with patch( "utils.hardware.vram_estimation.compute_gradient_bytes", return_value = 1, @@ -1275,7 +1291,9 @@ class TestSharedExperts(unittest.TestCase): delta_per_layer = 4096 * 1407 * 3 * 2 expected_delta = delta_per_layer * 32 * 2 actual_delta = w_yes - w_no - self.assertAlmostEqual(actual_delta, expected_delta, delta = expected_delta * 0.01) + self.assertAlmostEqual( + actual_delta, expected_delta, delta = expected_delta * 0.01 + ) def test_deepseek_v3_params_in_range(self): total = compute_total_params(DEEPSEEK_V3) @@ -1391,7 +1409,9 @@ class TestDenseMoEMix(unittest.TestCase): moe_intermediate_size = 1024, num_dense_layers = 5, ) - lora_all = compute_lora_params(all_moe, 16, ["gate_proj", "up_proj", "down_proj"]) + lora_all = compute_lora_params( + all_moe, 16, ["gate_proj", "up_proj", "down_proj"] + ) lora_mix = compute_lora_params(mixed, 16, ["gate_proj", "up_proj", "down_proj"]) self.assertNotEqual(lora_all, lora_mix) @@ -1475,7 +1495,9 @@ class TestPerLayerInputSkipAlias(unittest.TestCase): delta = _compute_skipped_quantizable_elements(arch) self.assertEqual( delta, - arch.hidden_size * arch.num_hidden_layers * arch.hidden_size_per_layer_input, + arch.hidden_size + * arch.num_hidden_layers + * arch.hidden_size_per_layer_input, ) def test_layer_aggregate_skip_includes_per_layer_input_modules(self): @@ -1554,7 +1576,9 @@ class TestSharedExpertVariants(unittest.TestCase): def test_shared_expert_size_separate_from_routed_changes_weight_count(self): from utils.hardware.vram_estimation import _compute_moe_mlp_elements - arch_separate = extract_arch_config(self._hf(shared_expert_intermediate_size = 64)) + arch_separate = extract_arch_config( + self._hf(shared_expert_intermediate_size = 64) + ) arch_implicit = extract_arch_config(self._hf(n_shared_experts = 1)) # Different shared sizes (64 vs default moe_intermediate_size=128) must # give different MoE element counts. @@ -1598,7 +1622,9 @@ class TestSharedExpertActivation(unittest.TestCase): moe_intermediate_size = 64, **fields, ) - return extract_arch_config(SimpleNamespace(text_config = text_config, quantization_config = {})) + return extract_arch_config( + SimpleNamespace(text_config = text_config, quantization_config = {}) + ) def test_shared_expert_increases_activation_bytes(self): with_shared = self._make(shared_expert_intermediate_size = 64) @@ -1650,7 +1676,9 @@ class TestPerLayerInputActivation(unittest.TestCase): tie_word_embeddings = False, **fields, ) - return extract_arch_config(SimpleNamespace(text_config = text_config, quantization_config = {})) + return extract_arch_config( + SimpleNamespace(text_config = text_config, quantization_config = {}) + ) def test_ple_increases_activation_bytes(self): with_ple = self._make( @@ -1714,7 +1742,9 @@ class TestKvSharedActivation(unittest.TestCase): num_kv_shared_layers = kv_shared, layer_types = ["full_attention"] * 4, ) - return extract_arch_config(SimpleNamespace(text_config = text_config, quantization_config = {})) + return extract_arch_config( + SimpleNamespace(text_config = text_config, quantization_config = {}) + ) def test_kv_shared_layers_keep_activation_bytes(self): shared = self._make(kv_shared = 2) @@ -1760,7 +1790,9 @@ class TestSparseMoeSkipAliases(unittest.TestCase): def test_gemma4_layers_experts_alias_pulls_routed(self): from utils.hardware.vram_estimation import _compute_skipped_quantizable_elements - arch = extract_arch_config(self._hf(["model.layers.0.experts"], enable_moe_block = True)) + arch = extract_arch_config( + self._hf(["model.layers.0.experts"], enable_moe_block = True) + ) self.assertGreater(_compute_skipped_quantizable_elements(arch), 0) def test_qwen_shared_expert_skip_pulls_only_shared(self): @@ -1811,7 +1843,9 @@ class TestAllLinearMoELoraExclusion(unittest.TestCase): moe_intermediate_size = 64, **fields, ) - return extract_arch_config(SimpleNamespace(text_config = text_config, quantization_config = {})) + return extract_arch_config( + SimpleNamespace(text_config = text_config, quantization_config = {}) + ) def test_all_linear_drops_routed_moe_expert_lora(self): arch = self._arch() @@ -1829,7 +1863,9 @@ class TestAllLinearMoELoraExclusion(unittest.TestCase): def test_all_linear_includes_attention_lora(self): arch = self._arch() all_linear = compute_lora_params(arch, 8, "all-linear") - attn_only = compute_lora_params(arch, 8, ["q_proj", "k_proj", "v_proj", "o_proj"]) + attn_only = compute_lora_params( + arch, 8, ["q_proj", "k_proj", "v_proj", "o_proj"] + ) # all-linear still attaches to attention nn.Linear modules. self.assertGreaterEqual(all_linear, attn_only) @@ -1847,7 +1883,9 @@ class TestExplicitPerLayerInputLora(unittest.TestCase): hidden_size_per_layer_input = 32, vocab_size_per_layer_input = 128, ) - return extract_arch_config(SimpleNamespace(text_config = text_config, quantization_config = {})) + return extract_arch_config( + SimpleNamespace(text_config = text_config, quantization_config = {}) + ) def test_explicit_per_layer_input_gate_returns_nonzero(self): arch = self._arch() @@ -1886,7 +1924,9 @@ class TestTopKExpertActivation(unittest.TestCase): moe_intermediate_size = 64, **fields, ) - return extract_arch_config(SimpleNamespace(text_config = text_config, quantization_config = {})) + return extract_arch_config( + SimpleNamespace(text_config = text_config, quantization_config = {}) + ) def test_num_experts_per_tok_extracted(self): arch = self._make(num_experts_per_tok = 4) diff --git a/studio/backend/tests/test_windows_gpu_detection_mock.py b/studio/backend/tests/test_windows_gpu_detection_mock.py index 88a1a28d14..bc06b83b30 100644 --- a/studio/backend/tests/test_windows_gpu_detection_mock.py +++ b/studio/backend/tests/test_windows_gpu_detection_mock.py @@ -167,7 +167,9 @@ def _build_path_dirs_like_start_llama_server( cuda_path: str = "", ) -> list[str]: """Wrapper around the real _build_windows_path_dirs staticmethod.""" - return LlamaCppBackend._build_windows_path_dirs(str(binary_dir), str(prefix), cuda_path) + return LlamaCppBackend._build_windows_path_dirs( + str(binary_dir), str(prefix), cuda_path + ) def _mock_nvidia_smi_run(fake_output: str, returncode: int = 0) -> "mock._patch": @@ -201,7 +203,9 @@ class TestWindowsGpuDetectionAfter5106Fix: fake_csv = "0, 22805\n" with _mock_nvidia_smi_run(fake_csv): gpus = LlamaCppBackend._get_gpu_free_memory() - assert gpus == [(0, 22805)], f"GPU probe failed to parse mocked nvidia-smi output: {gpus}" + assert gpus == [ + (0, 22805) + ], f"GPU probe failed to parse mocked nvidia-smi output: {gpus}" def test_nvidia_smi_probe_respects_cuda_visible_devices(self, monkeypatch): """CUDA_VISIBLE_DEVICES=1 -> only GPU 1 visible.""" @@ -236,7 +240,9 @@ class TestWindowsGpuDetectionAfter5106Fix: site / "nvidia" / "cu13" / "bin" / "x86_64", site / "torch" / "lib", ): - assert str(expected) in out, f"resolver missed {expected.relative_to(prefix)}: {out}" + assert ( + str(expected) in out + ), f"resolver missed {expected.relative_to(prefix)}: {out}" def test_path_assembly_makes_cudart_reachable_without_toolkit(self, tmp_path): """The #5106 scenario: GPU detected, pip nvidia wheels present, @@ -247,7 +253,9 @@ class TestWindowsGpuDetectionAfter5106Fix: _populate_studio_venv(prefix) _populate_studio_install(install, runtime = "13.1") binary_dir = install / "build" / "bin" / "Release" - path_dirs = _build_path_dirs_like_start_llama_server(binary_dir, prefix, cuda_path = "") + path_dirs = _build_path_dirs_like_start_llama_server( + binary_dir, prefix, cuda_path = "" + ) # binary_dir first -- Windows DLL search step 1. assert path_dirs[0] == str( binary_dir @@ -263,7 +271,9 @@ class TestWindowsGpuDetectionAfter5106Fix: ) # Defence in depth: both fix paths contribute cudart. sources = {Path(e).relative_to(tmp_path).parts[0] for e, _ in cudart_locations} - assert "studio_install" in sources, f"#5322's cudart drop not reachable: {cudart_locations}" + assert ( + "studio_install" in sources + ), f"#5322's cudart drop not reachable: {cudart_locations}" assert ( "studio_venv" in sources ), f"#5324's pip nvidia dir not contributing cudart: {cudart_locations}" @@ -280,7 +290,8 @@ class TestWindowsGpuDetectionAfter5106Fix: for required in REAL_UPSTREAM_CUDART_BUNDLE["13.1"]: reachable = any((Path(d) / required).exists() for d in path_dirs) assert reachable, ( - f"{required} unreachable from PATH; #5106 not fixed.\n" f"PATH entries: {path_dirs}" + f"{required} unreachable from PATH; #5106 not fixed.\n" + f"PATH entries: {path_dirs}" ) def test_no_pip_nvidia_wheels_still_works_via_install_dir(self, tmp_path): @@ -292,7 +303,9 @@ class TestWindowsGpuDetectionAfter5106Fix: _populate_studio_install(install, runtime = "13.1") binary_dir = install / "build" / "bin" / "Release" path_dirs = _build_path_dirs_like_start_llama_server(binary_dir, prefix) - assert path_dirs == [str(binary_dir)], f"bare venv produced unexpected PATH: {path_dirs}" + assert path_dirs == [ + str(binary_dir) + ], f"bare venv produced unexpected PATH: {path_dirs}" for required in REAL_UPSTREAM_CUDART_BUNDLE["13.1"]: assert ( binary_dir / required @@ -316,7 +329,8 @@ class TestWindowsGpuDetectionAfter5106Fix: (rel / fn).write_bytes(b"PE-stub") path_dirs = _build_path_dirs_like_start_llama_server(rel, prefix) cudart_reachable = any( - (Path(d) / "cudart64_12.dll").exists() or (Path(d) / "cudart64_13.dll").exists() + (Path(d) / "cudart64_12.dll").exists() + or (Path(d) / "cudart64_13.dll").exists() for d in path_dirs ) assert cudart_reachable, ( @@ -324,7 +338,8 @@ class TestWindowsGpuDetectionAfter5106Fix: f"on cudart-less install. PATH entries: {path_dirs}" ) cublas_reachable = any( - (Path(d) / "cublas64_12.dll").exists() or (Path(d) / "cublas64_13.dll").exists() + (Path(d) / "cublas64_12.dll").exists() + or (Path(d) / "cublas64_13.dll").exists() for d in path_dirs ) assert cublas_reachable, "cublas unreachable on cudart-less install" @@ -343,7 +358,8 @@ class TestWindowsGpuDetectionAfter5106Fix: # Pre-PR PATH: binary_dir only, no pip nvidia dirs, no toolkit. pre_pr_path_dirs = [str(rel)] cudart_reachable_pre = any( - (Path(d) / "cudart64_12.dll").exists() or (Path(d) / "cudart64_13.dll").exists() + (Path(d) / "cudart64_12.dll").exists() + or (Path(d) / "cudart64_13.dll").exists() for d in pre_pr_path_dirs ) assert not cudart_reachable_pre, ( @@ -364,5 +380,7 @@ class TestWindowsSysPlatformMocked: out = LlamaCppBackend._windows_pip_nvidia_dll_dirs(str(prefix)) assert out, f"resolver returned empty under sys.platform=win32: {out}" # cu13 arch dir must be in the output. - cu13_arch = prefix / "Lib" / "site-packages" / "nvidia" / "cu13" / "bin" / "x86_64" + cu13_arch = ( + prefix / "Lib" / "site-packages" / "nvidia" / "cu13" / "bin" / "x86_64" + ) assert str(cu13_arch) in out diff --git a/studio/backend/utils/api_errors.py b/studio/backend/utils/api_errors.py index b1c55b61b9..0ad10a7ae1 100644 --- a/studio/backend/utils/api_errors.py +++ b/studio/backend/utils/api_errors.py @@ -142,7 +142,9 @@ def error_body_for_path( """ if is_anthropic_path(path): return anthropic_error_body(message, status = status, err_type = err_type) - return openai_error_body(message, status = status, err_type = err_type, code = code, param = param) + return openai_error_body( + message, status = status, err_type = err_type, code = code, param = param + ) def _summarize_validation_errors(errors) -> tuple: @@ -175,7 +177,11 @@ def _summarize_validation_errors(errors) -> tuple: param = part break - label = ".".join(str(p) for p in loc_parts) if loc_parts else ".".join(str(p) for p in loc) + label = ( + ".".join(str(p) for p in loc_parts) + if loc_parts + else ".".join(str(p) for p in loc) + ) summary = f"{label}: {msg}" if label else str(msg) return summary, param @@ -214,7 +220,9 @@ def install_api_error_handlers(app) -> None: if path.startswith("/v1/"): detail = exc.detail # Already a fully-formed envelope: pass through untouched. - if isinstance(detail, dict) and ("error" in detail or detail.get("type") == "error"): + if isinstance(detail, dict) and ( + "error" in detail or detail.get("type") == "error" + ): return JSONResponse( status_code = exc.status_code, content = detail, diff --git a/studio/backend/utils/cache_cleanup.py b/studio/backend/utils/cache_cleanup.py index 210735973d..ede6047fcc 100644 --- a/studio/backend/utils/cache_cleanup.py +++ b/studio/backend/utils/cache_cleanup.py @@ -72,7 +72,8 @@ def clear_unsloth_compiled_cache(preserve_patterns: Optional[List[str]] = None) if preserve_patterns: logger.info( - f"Cleaning unsloth compiled cache (preserving {preserve_patterns}): " f"{cache_dir}" + f"Cleaning unsloth compiled cache (preserving {preserve_patterns}): " + f"{cache_dir}" ) for item in cache_dir.iterdir(): diff --git a/studio/backend/utils/datasets/data_collators.py b/studio/backend/utils/datasets/data_collators.py index 9bfb60ba17..73d66a2ee4 100644 --- a/studio/backend/utils/datasets/data_collators.py +++ b/studio/backend/utils/datasets/data_collators.py @@ -23,13 +23,19 @@ class DataCollatorSpeechSeq2SeqWithPadding: processor: Any def __call__(self, features: List[dict]) -> dict: - input_features = [{"input_features": feature["input_features"]} for feature in features] - batch = self.processor.feature_extractor.pad(input_features, return_tensors = "pt") + input_features = [ + {"input_features": feature["input_features"]} for feature in features + ] + batch = self.processor.feature_extractor.pad( + input_features, return_tensors = "pt" + ) label_features = [{"input_ids": feature["labels"]} for feature in features] labels_batch = self.processor.tokenizer.pad(label_features, return_tensors = "pt") - labels = labels_batch["input_ids"].masked_fill(labels_batch.attention_mask.ne(1), -100) + labels = labels_batch["input_ids"].masked_fill( + labels_batch.attention_mask.ne(1), -100 + ) if (labels[:, 0] == self.processor.tokenizer.bos_token_id).all().cpu().item(): labels = labels[:, 1:] @@ -136,7 +142,9 @@ class VLMDataCollator: all_images.append(img) texts = [ - self.processor.apply_chat_template(msgs, tokenize = False, add_generation_prompt = False) + self.processor.apply_chat_template( + msgs, tokenize = False, add_generation_prompt = False + ) for msgs in all_messages ] diff --git a/studio/backend/utils/datasets/dataset_none_detect.py b/studio/backend/utils/datasets/dataset_none_detect.py index a2fd8ef667..7fb39c6526 100644 --- a/studio/backend/utils/datasets/dataset_none_detect.py +++ b/studio/backend/utils/datasets/dataset_none_detect.py @@ -71,7 +71,9 @@ def _probe_conversation(dataset: Dataset, candidates = None): # No usable dict turn in 100 rows. Record an all_corrupt fallback, # plausible only with turn-shaped data (None cell or list of dict/None # turns); a later plausible candidate upgrades a non-plausible one. - if all_corrupt_fallback is None or not all_corrupt_fallback.get("has_plausible_turns"): + if all_corrupt_fallback is None or not all_corrupt_fallback.get( + "has_plausible_turns" + ): has_plausible_turns = False for i in range(min(len(dataset), 100)): cell = dataset[i][col] @@ -116,7 +118,9 @@ def _probe_conversation(dataset: Dataset, candidates = None): _CONV_KEYS = {"role", "from", "content", "value"} if not any(keys <= turn_keys for keys in _CHAT_KEY_SETS): schema_less_plausible = bool(turn_keys & _CONV_KEYS) - if all_corrupt_fallback is None or not all_corrupt_fallback.get("has_plausible_turns"): + if all_corrupt_fallback is None or not all_corrupt_fallback.get( + "has_plausible_turns" + ): all_corrupt_fallback = { "column": col, "turn_keys": turn_keys, @@ -159,11 +163,14 @@ def is_none_or_empty(value) -> bool: non_text_blocks = [item for item in dict_blocks if item.get("type") != "text"] if non_text_blocks: return False - text_values = [item.get("text") for item in dict_blocks if item.get("type") == "text"] + text_values = [ + item.get("text") for item in dict_blocks if item.get("type") == "text" + ] if text_values and all( t is None or ( - isinstance(t, str) and not t.strip().strip("\ufeff\u200b\u200c\u200d\u2060").strip() + isinstance(t, str) + and not t.strip().strip("\ufeff\u200b\u200c\u200d\u2060").strip() ) for t in text_values ): @@ -274,7 +281,9 @@ def find_none_chatml(dataset: Dataset, col: str = None) -> dict: stats["rows_with_none_turns"] += 1 stats["total_none_turns"] += 1 stats["rows_all_none"] += 1 - stats["none_by_role"]["unknown"] = stats["none_by_role"].get("unknown", 0) + 1 + stats["none_by_role"]["unknown"] = ( + stats["none_by_role"].get("unknown", 0) + 1 + ) stats["none_by_type"][vtype] = stats["none_by_type"].get(vtype, 0) + 1 stats["findings"].append( { @@ -293,7 +302,9 @@ def find_none_chatml(dataset: Dataset, col: str = None) -> dict: stats["rows_with_none_turns"] += 1 stats["total_none_turns"] += 1 stats["rows_all_none"] += 1 - stats["none_by_role"]["unknown"] = stats["none_by_role"].get("unknown", 0) + 1 + stats["none_by_role"]["unknown"] = ( + stats["none_by_role"].get("unknown", 0) + 1 + ) stats["none_by_type"]["empty_conversation"] = ( stats["none_by_type"].get("empty_conversation", 0) + 1 ) @@ -321,7 +332,9 @@ def find_none_chatml(dataset: Dataset, col: str = None) -> dict: "raw_value": repr(turn), } ) - stats["none_by_role"]["unknown"] = stats["none_by_role"].get("unknown", 0) + 1 + stats["none_by_role"]["unknown"] = ( + stats["none_by_role"].get("unknown", 0) + 1 + ) vtype = "None" if turn is None else "invalid_type" stats["none_by_type"][vtype] = stats["none_by_type"].get(vtype, 0) + 1 continue @@ -342,14 +355,20 @@ def find_none_chatml(dataset: Dataset, col: str = None) -> dict: if "from" in turn and "value" in turn: content = turn.get("value") elif "role" in turn: - content = turn.get("content") if "content" in turn else turn.get("value") + content = ( + turn.get("content") if "content" in turn else turn.get("value") + ) elif "from" in turn: content = turn.get("value") else: - content = turn.get("content") if "content" in turn else turn.get("value") + content = ( + turn.get("content") if "content" in turn else turn.get("value") + ) # Assistant tool-call turns carry empty content + tool_calls and are # valid; the exemption is assistant-only. - if is_none_or_empty(content) and not (role == "assistant" and turn.get("tool_calls")): + if is_none_or_empty(content) and not ( + role == "assistant" and turn.get("tool_calls") + ): vtype = _classify_empty(content) row_findings.append( { @@ -446,7 +465,9 @@ FORMAT_REGISTRY = [ }, { "name": "sharegpt", - "match": lambda ds, conv: (conv is not None and {"from", "value"} <= conv["turn_keys"]), + "match": lambda ds, conv: ( + conv is not None and {"from", "value"} <= conv["turn_keys"] + ), "scan": find_none_sharegpt, }, { @@ -731,7 +752,9 @@ def show_row( # Mirror scanner: tool_calls exemption is assistant-only; # other roles with empty content + tool_calls are still bad. r = t.get("role") if t.get("role") is not None else t.get("from") - if is_none_or_empty(c) and not (str(r) == "assistant" and t.get("tool_calls")): + if is_none_or_empty(c) and not ( + str(r) == "assistant" and t.get("tool_calls") + ): return True return False @@ -752,11 +775,19 @@ def show_row( if "from" in turn and "value" in turn: content = turn.get("value") elif "role" in turn: - content = turn.get("content") if "content" in turn else turn.get("value") + content = ( + turn.get("content") + if "content" in turn + else turn.get("value") + ) elif "from" in turn: content = turn.get("value") else: - content = turn.get("content") if "content" in turn else turn.get("value") + content = ( + turn.get("content") + if "content" in turn + else turn.get("value") + ) if is_none_or_empty(content) and not ( role == "assistant" and turn.get("tool_calls") ): @@ -798,8 +829,12 @@ examples: python dataset_none_detect.py org/my-dataset --token hf_... """, ) - parser.add_argument("dataset", help = "HuggingFace dataset repo id (e.g. org/my-dataset)") - parser.add_argument("--split", default = "train", help = "Dataset split to load (default: train)") + parser.add_argument( + "dataset", help = "HuggingFace dataset repo id (e.g. org/my-dataset)" + ) + parser.add_argument( + "--split", default = "train", help = "Dataset split to load (default: train)" + ) parser.add_argument( "--format", default = "auto", diff --git a/studio/backend/utils/datasets/dataset_utils.py b/studio/backend/utils/datasets/dataset_utils.py index faa3deac70..cb943d7402 100644 --- a/studio/backend/utils/datasets/dataset_utils.py +++ b/studio/backend/utils/datasets/dataset_utils.py @@ -341,7 +341,9 @@ def _apply_template_mapping( user_parts = [] for col in role_groups["user"]: if col in examples: - user_parts.append(_extract_column_value(examples[col][i], col, label_mapping)) + user_parts.append( + _extract_column_value(examples[col][i], col, label_mapping) + ) if user_parts: convo.append({"role": "user", "content": "\n".join(user_parts)}) @@ -349,7 +351,9 @@ def _apply_template_mapping( asst_parts = [] for col in role_groups["assistant"]: if col in examples: - asst_parts.append(_extract_column_value(examples[col][i], col, label_mapping)) + asst_parts.append( + _extract_column_value(examples[col][i], col, label_mapping) + ) if asst_parts: convo.append({"role": "assistant", "content": "\n".join(asst_parts)}) @@ -399,7 +403,11 @@ def _apply_user_mapping_alpaca( ("output", outputs), ): col = col_for[field] - val = str(examples[col][i]) if col and col in examples and examples[col][i] else "" + val = ( + str(examples[col][i]) + if col and col in examples and examples[col][i] + else "" + ) dest.append(val) return {"instruction": instructions, "input": inputs, "output": outputs} @@ -477,7 +485,9 @@ def format_dataset( else: # auto / chatml / sharegpt / conversational all produce chatml # conversations (sharegpt standardized to role/content internally) - mapped_dataset = _apply_user_mapping(dataset, custom_format_mapping, batch_size) + mapped_dataset = _apply_user_mapping( + dataset, custom_format_mapping, batch_size + ) final_format = "chatml_conversations" chat_column = "conversations" @@ -574,7 +584,9 @@ def format_dataset( elif detected["format"] == "chatml" and detected.get("chat_column"): return { "dataset": dataset, - "detected_format": _chatml_detected_format_label(detected["chat_column"]), + "detected_format": _chatml_detected_format_label( + detected["chat_column"] + ), "final_format": _chatml_final_format(detected["chat_column"]), "chat_column": detected["chat_column"], "is_standardized": True, @@ -586,7 +598,9 @@ def format_dataset( # Unknown - try standardization, pass as-is on failure else: - warnings.append(f"Unknown format detected. Keys found: {detected['sample_keys']}") + warnings.append( + f"Unknown format detected. Keys found: {detected['sample_keys']}" + ) # Try heuristic detection if auto_detect_custom: @@ -612,7 +626,9 @@ def format_dataset( if role == target_role and col_name in examples: content = examples[col_name][i] if content and str(content).strip(): - convo.append({"role": role, "content": str(content)}) + convo.append( + {"role": role, "content": str(content)} + ) conversations.append(convo) return {"conversations": conversations, **preserved_columns} @@ -661,7 +677,9 @@ def format_dataset( "warnings": warnings, } except Exception as e: - warnings.append(f"Could not standardize: {e}. Passing dataset as-is.") + warnings.append( + f"Could not standardize: {e}. Passing dataset as-is." + ) # Return as-is with warnings return { @@ -691,7 +709,9 @@ def format_dataset( "warnings": [], } - elif detected["format"] in ["sharegpt", "chatml"] and detected.get("chat_column"): + elif detected["format"] in ["sharegpt", "chatml"] and detected.get( + "chat_column" + ): try: # First standardize if ShareGPT if detected["format"] == "sharegpt": @@ -808,7 +828,9 @@ def format_dataset( elif detected["format"] == "chatml" and detected.get("chat_column"): return { "dataset": dataset, - "detected_format": _chatml_detected_format_label(detected["chat_column"]), + "detected_format": _chatml_detected_format_label( + detected["chat_column"] + ), "final_format": _chatml_final_format(detected["chat_column"]), "chat_column": detected["chat_column"], "is_standardized": True, @@ -966,7 +988,9 @@ def format_and_template_dataset( f"text='{user_vlm_text_column}') failed: {e} — " f"falling back to auto-detection" ) - logger.info(f"⚠️ User VLM mapping failed, falling back to auto-detection...") + logger.info( + f"⚠️ User VLM mapping failed, falling back to auto-detection..." + ) custom_format_mapping = None # so auto-detection runs below else: errors.append( @@ -1020,7 +1044,9 @@ def format_and_template_dataset( dataset_name = dataset_name, progress_callback = progress_callback, ) - warnings.append("Converted from ShareGPT+image format to standard VLM format") + warnings.append( + "Converted from ShareGPT+image format to standard VLM format" + ) except Exception as e: errors.append(f"Failed to convert ShareGPT+image format: {e}") import traceback @@ -1088,9 +1114,13 @@ def format_and_template_dataset( ) if vlm_instruction: - warnings.append(f"Using user-provided instruction: '{vlm_instruction}'") + warnings.append( + f"Using user-provided instruction: '{vlm_instruction}'" + ) else: - warnings.append("Auto-generated instruction based on dataset analysis") + warnings.append( + "Auto-generated instruction based on dataset analysis" + ) except Exception as e: errors.append(f"Failed to convert to VLM format: {e}") @@ -1195,7 +1225,9 @@ def format_and_template_dataset( summary = get_dataset_info_summary(dataset_info) # Combine results - all_warnings = dataset_info.get("warnings", []) + template_result.get("warnings", []) + all_warnings = dataset_info.get("warnings", []) + template_result.get( + "warnings", [] + ) all_errors = template_result.get("errors", []) # If apply_chat_template rescued an "unknown" format, update final_format. diff --git a/studio/backend/utils/datasets/format_conversion.py b/studio/backend/utils/datasets/format_conversion.py index cb24bd96ba..096a33c18f 100644 --- a/studio/backend/utils/datasets/format_conversion.py +++ b/studio/backend/utils/datasets/format_conversion.py @@ -100,7 +100,9 @@ def standardize_chat_format( role_key = keys[1] content_key = keys[0] else: - raise ValueError(f"Could not infer role/content keys for chat column '{chat_column}'") + raise ValueError( + f"Could not infer role/content keys for chat column '{chat_column}'" + ) # Mapping for aliases aliases_mapping = {} @@ -131,7 +133,9 @@ def standardize_chat_format( if original_role is None: original_role = message.get("role") or message.get("from") or "" if original_content is None: - original_content = message.get("content") or message.get("value") or "" + original_content = ( + message.get("content") or message.get("value") or "" + ) standard_role = aliases_mapping.get(original_role, original_role) @@ -188,11 +192,15 @@ def convert_chatml_to_alpaca( chatml_data = examples.get(chat_column) if chat_column else None if chatml_data is None: chatml_data = ( - examples.get("messages") or examples.get("conversations") or examples.get("texts") + examples.get("messages") + or examples.get("conversations") + or examples.get("texts") ) if chatml_data is None: - raise ValueError("No 'messages' or 'conversations' or 'texts' column found.") + raise ValueError( + "No 'messages' or 'conversations' or 'texts' column found." + ) instructions = [] outputs = [] @@ -352,12 +360,16 @@ def convert_to_vlm_format( instruction_column = instruction_info.get("instruction_column") uses_dynamic = instruction_info["uses_dynamic_instruction"] - logger.info(f"📝 Auto-detected instruction type: {instruction_info['instruction_type']}") + logger.info( + f"📝 Auto-detected instruction type: {instruction_info['instruction_type']}" + ) logger.info(f"📝 Confidence: {instruction_info['confidence']:.2f}") if not uses_dynamic: logger.info(f"📝 Using instruction: '{instruction}'") else: - logger.info(f"📝 Using dynamic instructions from column: '{instruction_column}'") + logger.info( + f"📝 Using dynamic instructions from column: '{instruction_column}'" + ) else: instruction_column = None uses_dynamic = False @@ -412,7 +424,9 @@ def convert_to_vlm_format( total = len(dataset) first_image = next(iter(dataset))[image_column] - has_urls = isinstance(first_image, str) and first_image.startswith(("http://", "https://")) + has_urls = isinstance(first_image, str) and first_image.startswith( + ("http://", "https://") + ) # ── Bare-filename detection: build a basename→repo_path lookup so # filename-only images resolve via hf_hub_download during conversion. @@ -461,7 +475,9 @@ def convert_to_vlm_format( num_workers = safe_thread_num_proc() _notify(f"Probing {PROBE_SIZE} image URLs with {num_workers} workers...") - logger.info(f"🔍 Probing {PROBE_SIZE}/{total} image URLs with {num_workers} workers...") + logger.info( + f"🔍 Probing {PROBE_SIZE}/{total} image URLs with {num_workers} workers..." + ) probe_samples = [dataset[i] for i in range(PROBE_SIZE)] probe_ok = 0 @@ -469,7 +485,9 @@ def convert_to_vlm_format( probe_start = time.time() with ThreadPoolExecutor(max_workers = num_workers) as executor: - futures = {executor.submit(_convert_single_sample, s): s for s in probe_samples} + futures = { + executor.submit(_convert_single_sample, s): s for s in probe_samples + } for future in as_completed(futures): try: future.result() @@ -561,7 +579,9 @@ def convert_to_vlm_format( except Exception as e: failed_count += 1 if failed_count == 1: - logger.info(f"First VLM conversion failure: {type(e).__name__}: {e}") + logger.info( + f"First VLM conversion failure: {type(e).__name__}: {e}" + ) converted_list.extend(r for r in batch_results if r is not None) @@ -586,7 +606,9 @@ def convert_to_vlm_format( failed_count += 1 if failed_count == 1: # Log the first failure to aid debugging - logger.info(f"First VLM conversion failure: {type(e).__name__}: {e}") + logger.info( + f"First VLM conversion failure: {type(e).__name__}: {e}" + ) pbar.set_postfix(ok = len(converted_list), failed = failed_count, refresh = False) pbar.close() @@ -752,7 +774,9 @@ def convert_sharegpt_with_images_to_vlm_format( return Image.open(local_path).convert("RGB") else: return Image.open(image_data).convert("RGB") - if isinstance(image_data, dict) and ("bytes" in image_data or "path" in image_data): + if isinstance(image_data, dict) and ( + "bytes" in image_data or "path" in image_data + ): if image_data.get("bytes"): from io import BytesIO return Image.open(BytesIO(image_data["bytes"])).convert("RGB") @@ -808,7 +832,9 @@ def convert_sharegpt_with_images_to_vlm_format( pbar.close() if failed_count > 0: - logger.info(f"⚠️ Skipped {failed_count}/{total} ({failed_count*100//total}%) samples") + logger.info( + f"⚠️ Skipped {failed_count}/{total} ({failed_count*100//total}%) samples" + ) if len(converted_list) == 0: raise ValueError( @@ -834,7 +860,9 @@ def convert_llava_to_vlm_format(dataset): """ from PIL import Image - logger.info(f"🔄 Converting {len(dataset)} samples from Llava format to standard VLM format...") + logger.info( + f"🔄 Converting {len(dataset)} samples from Llava format to standard VLM format..." + ) def _convert_single_sample(sample): """Convert one llava sample to standard VLM format.""" diff --git a/studio/backend/utils/datasets/format_detection.py b/studio/backend/utils/datasets/format_detection.py index f5ea5ca138..cdc816235f 100644 --- a/studio/backend/utils/datasets/format_detection.py +++ b/studio/backend/utils/datasets/format_detection.py @@ -8,7 +8,10 @@ import re def _keyword_in_column(keyword: str, col_name: str) -> bool: """Word-boundary keyword match to avoid false positives like 'pic' in 'topic'.""" - return re.search(r"\b" + re.escape(keyword) + r"\b", col_name, re.IGNORECASE) is not None + return ( + re.search(r"\b" + re.escape(keyword) + r"\b", col_name, re.IGNORECASE) + is not None + ) CONVERSATION_COLUMNS = ("messages", "conversations", "texts") @@ -89,7 +92,9 @@ def _inspect_conversation_column(rows: list[dict], column_name: str) -> dict | N return None -def _detect_conversation_column(rows: list[dict], column_names: list[str]) -> dict | None: +def _detect_conversation_column( + rows: list[dict], column_names: list[str] +) -> dict | None: column_name_set = set(column_names) unknown_exact = None for column_name in CONVERSATION_COLUMNS: @@ -281,7 +286,10 @@ def detect_custom_format_heuristic(dataset): return True for pattern in metadata_prefix_patterns: - if col_lower.startswith(pattern.split("_")[0] + "_") and col_lower != pattern: + if ( + col_lower.startswith(pattern.split("_")[0] + "_") + and col_lower != pattern + ): if "_" in col_lower: prefix = col_lower.split("_")[0] if prefix in ["generation", "pass", "inference"]: @@ -324,7 +332,9 @@ def detect_custom_format_heuristic(dataset): # Penalize ambiguous "task" so other user columns win. if role_type == "user": col_lower = col_name.lower() - if "task" in col_lower and not any(kw in col_lower for kw in user_words_high_priority): + if "task" in col_lower and not any( + kw in col_lower for kw in user_words_high_priority + ): score -= 15 priority_bonus = get_priority_score(col_name) @@ -354,13 +364,17 @@ def detect_custom_format_heuristic(dataset): content_columns = [col for col in all_columns if not is_metadata(col)] - assistant_potential = [col for col in content_columns if has_keyword(col, assistant_words)] + assistant_potential = [ + col for col in content_columns if has_keyword(col, assistant_words) + ] user_potential = [col for col in content_columns if has_keyword(col, user_words)] # STEP 1: best ASSISTANT column assistant_candidates = [] for col in assistant_potential: - score = score_column(col, assistant_words, "assistant", len(assistant_potential)) + score = score_column( + col, assistant_words, "assistant", len(assistant_potential) + ) if score > 0: assistant_candidates.append((col, score)) @@ -664,7 +678,9 @@ def detect_vlm_dataset_structure(dataset): if isinstance(content[0], dict) and "type" in content[0]: # Llava format? has_index = any( - "index" in item for item in content if isinstance(item, dict) + "index" in item + for item in content + if isinstance(item, dict) ) has_images_column = "images" in column_names @@ -679,7 +695,9 @@ def detect_vlm_dataset_structure(dataset): # Standard VLM format has_image = any( - "image" in item for item in content if isinstance(item, dict) + "image" in item + for item in content + if isinstance(item, dict) ) if has_image: return { @@ -782,7 +800,9 @@ def detect_vlm_dataset_structure(dataset): if any(col_lower.endswith(suffix) for suffix in metadata_patterns["suffixes"]): return True - if any(col_lower.startswith(prefix) for prefix in metadata_patterns["prefixes"]): + if any( + col_lower.startswith(prefix) for prefix in metadata_patterns["prefixes"] + ): return True return False @@ -794,7 +814,9 @@ def detect_vlm_dataset_structure(dataset): return 100 # HF Image feature dict. - if isinstance(sample_value, dict) and ("bytes" in sample_value or "path" in sample_value): + if isinstance(sample_value, dict) and ( + "bytes" in sample_value or "path" in sample_value + ): return 75 if isinstance(sample_value, str): @@ -816,7 +838,9 @@ def detect_vlm_dataset_structure(dataset): # Local file — check it exists. if not sample_value.startswith(("http://", "https://")): - return os.path.exists(sample_value) # bare filenames return False, that's OK + return os.path.exists( + sample_value + ) # bare filenames return False, that's OK # URL — quick HEAD with short timeout. try: diff --git a/studio/backend/utils/datasets/llm_assist.py b/studio/backend/utils/datasets/llm_assist.py index f7b35e2869..85a67724b2 100644 --- a/studio/backend/utils/datasets/llm_assist.py +++ b/studio/backend/utils/datasets/llm_assist.py @@ -53,7 +53,9 @@ def precache_helper_gguf(): return repo = os.environ.get("UNSLOTH_HELPER_MODEL_REPO", DEFAULT_HELPER_MODEL_REPO) - variant = os.environ.get("UNSLOTH_HELPER_MODEL_VARIANT", DEFAULT_HELPER_MODEL_VARIANT) + variant = os.environ.get( + "UNSLOTH_HELPER_MODEL_VARIANT", DEFAULT_HELPER_MODEL_VARIANT + ) try: from huggingface_hub import HfApi, hf_hub_download @@ -68,7 +70,9 @@ def precache_helper_gguf(): # GGUF files matching the variant (may be split into shards). variant_lower = variant.lower().replace("-", "_") - matching = sorted(f for f in gguf_files if variant_lower in f.lower().replace("-", "_")) + matching = sorted( + f for f in gguf_files if variant_lower in f.lower().replace("-", "_") + ) if matching: logger.info( @@ -95,7 +99,9 @@ def _run_with_helper(prompt: str, max_tokens: int = 256) -> Optional[str]: return None repo = os.environ.get("UNSLOTH_HELPER_MODEL_REPO", DEFAULT_HELPER_MODEL_REPO) - variant = os.environ.get("UNSLOTH_HELPER_MODEL_VARIANT", DEFAULT_HELPER_MODEL_VARIANT) + variant = os.environ.get( + "UNSLOTH_HELPER_MODEL_VARIANT", DEFAULT_HELPER_MODEL_VARIANT + ) backend = None try: @@ -117,7 +123,9 @@ def _run_with_helper(prompt: str, max_tokens: int = 256) -> Optional[str]: return None messages = [{"role": "user", "content": prompt}] - logger.info("Helper model request: enable_thinking=False (per-request override)") + logger.info( + "Helper model request: enable_thinking=False (per-request override)" + ) cumulative = "" for chunk in backend.generate_chat_completion( messages = messages, @@ -200,7 +208,9 @@ def llm_generate_vlm_instruction( } -def llm_classify_columns(column_names: list[str], samples: list[dict]) -> Optional[dict[str, str]]: +def llm_classify_columns( + column_names: list[str], samples: list[dict] +) -> Optional[dict[str, str]]: """Ask a helper LLM to classify columns into roles (when heuristic detection fails). Returns {column_name: role} for roles user|assistant|system|metadata, or None. @@ -258,7 +268,11 @@ def llm_classify_columns(column_names: list[str], samples: list[dict]) -> Option valid_roles = {"user", "assistant", "system", "metadata"} cleaned = {} for col, role in mapping.items(): - if col in column_names and isinstance(role, str) and role.lower() in valid_roles: + if ( + col in column_names + and isinstance(role, str) + and role.lower() in valid_roles + ): cleaned[col] = role.lower() if not cleaned: @@ -409,7 +423,9 @@ def fetch_hf_dataset_card( if val is not None: metadata[key] = val - logger.info(f"Fetched dataset card: {len(readme)} chars, {len(metadata)} metadata fields") + logger.info( + f"Fetched dataset card: {len(readme)} chars, {len(metadata)} metadata fields" + ) return readme, metadata except Exception as e: @@ -435,7 +451,9 @@ def _run_multi_pass_advisor( return None repo = os.environ.get("UNSLOTH_HELPER_MODEL_REPO", DEFAULT_HELPER_MODEL_REPO) - variant = os.environ.get("UNSLOTH_HELPER_MODEL_VARIANT", DEFAULT_HELPER_MODEL_VARIANT) + variant = os.environ.get( + "UNSLOTH_HELPER_MODEL_VARIANT", DEFAULT_HELPER_MODEL_VARIANT + ) backend = None try: @@ -465,7 +483,9 @@ def _run_multi_pass_advisor( samples_text += f"Row {i}:\n" + "\n".join(parts) + "\n" metadata_str = ( - json.dumps(dataset_metadata, indent = 2, default = str)[:500] if dataset_metadata else "N/A" + json.dumps(dataset_metadata, indent = 2, default = str)[:500] + if dataset_metadata + else "N/A" ) card_excerpt = (dataset_card or "")[:1200] or "N/A" @@ -667,7 +687,9 @@ def _run_multi_pass_advisor( # Must have at least one user AND one assistant roles_present = set(column_roles.values()) if "user" not in roles_present or "assistant" not in roles_present: - logger.warning(f"Pass 2 sanity fail: missing user or assistant role: {column_roles}") + logger.warning( + f"Pass 2 sanity fail: missing user or assistant role: {column_roles}" + ) return None # falls back to simple classification # ── Pass 3: System prompt (non-conversational datasets only) ── diff --git a/studio/backend/utils/datasets/raw_text.py b/studio/backend/utils/datasets/raw_text.py index 03315fb287..82e5683eb5 100644 --- a/studio/backend/utils/datasets/raw_text.py +++ b/studio/backend/utils/datasets/raw_text.py @@ -100,7 +100,8 @@ def prepare_raw_text_dataset( notices.append( RawTextNotice( message = ( - f"{mode_title}: renaming column '{renamed_col}' -> 'text' " f"for {split_scope}" + f"{mode_title}: renaming column '{renamed_col}' -> 'text' " + f"for {split_scope}" ), level = "info", ) diff --git a/studio/backend/utils/datasets/vlm_processing.py b/studio/backend/utils/datasets/vlm_processing.py index f018913fa8..5c336ee2a3 100644 --- a/studio/backend/utils/datasets/vlm_processing.py +++ b/studio/backend/utils/datasets/vlm_processing.py @@ -65,7 +65,9 @@ def generate_smart_vlm_instruction( # OCR / Transcription "ocr": { "keywords": ["ocr", "transcribe", "transcript"], - "content_hints": [r"[A-Za-z\u0600-\u06FF]{10,}"], # Long Latin/Arabic passages + "content_hints": [ + r"[A-Za-z\u0600-\u06FF]{10,}" + ], # Long Latin/Arabic passages "instruction": "Transcribe all the text shown in this image.", "confidence": 0.9, }, diff --git a/studio/backend/utils/downsample.py b/studio/backend/utils/downsample.py index 2d340ca248..bccf6a23b7 100644 --- a/studio/backend/utils/downsample.py +++ b/studio/backend/utils/downsample.py @@ -12,5 +12,7 @@ def downsample(values: list[float], target_count: int) -> list[float]: return [] if target_count == 1: return [values[-1]] - indices = [round(i * (len(values) - 1) / (target_count - 1)) for i in range(target_count)] + indices = [ + round(i * (len(values) - 1) / (target_count - 1)) for i in range(target_count) + ] return [values[i] for i in indices] diff --git a/studio/backend/utils/hardware/amd.py b/studio/backend/utils/hardware/amd.py index 27dfe187cc..1764f16c4a 100644 --- a/studio/backend/utils/hardware/amd.py +++ b/studio/backend/utils/hardware/amd.py @@ -215,7 +215,9 @@ def _extract_gpu_metrics(gpu_data: dict) -> dict[str, Any]: # Output structure varies by version; try common paths usage = gpu_data.get("usage", gpu_data.get("gpu_activity", {})) if isinstance(usage, dict): - gpu_util = _parse_numeric(usage.get("gfx_activity", usage.get("gpu_use_percent"))) + gpu_util = _parse_numeric( + usage.get("gfx_activity", usage.get("gpu_use_percent")) + ) else: gpu_util = _parse_numeric(usage) @@ -240,7 +242,9 @@ def _extract_gpu_metrics(gpu_data: dict) -> dict[str, Any]: power_data.get("average_socket_power", power_data.get("socket_power")), ) ) - power_limit = _parse_numeric(power_data.get("power_cap", power_data.get("max_power_limit"))) + power_limit = _parse_numeric( + power_data.get("power_cap", power_data.get("max_power_limit")) + ) else: power_draw = None power_limit = None @@ -254,10 +258,14 @@ def _extract_gpu_metrics(gpu_data: dict) -> dict[str, Any]: ) if isinstance(vram_data, dict): vram_used_mb = _parse_memory_mb( - vram_data.get("used_vram", vram_data.get("vram_used", vram_data.get("used"))) + vram_data.get( + "used_vram", vram_data.get("vram_used", vram_data.get("used")) + ) ) vram_total_mb = _parse_memory_mb( - vram_data.get("total_vram", vram_data.get("vram_total", vram_data.get("total"))) + vram_data.get( + "total_vram", vram_data.get("vram_total", vram_data.get("total")) + ) ) else: vram_used_mb = None @@ -265,7 +273,9 @@ def _extract_gpu_metrics(gpu_data: dict) -> dict[str, Any]: # Build the standardized dict (same shape as nvidia._build_gpu_metrics) vram_used_gb = round(vram_used_mb / 1024, 2) if vram_used_mb is not None else None - vram_total_gb = round(vram_total_mb / 1024, 2) if vram_total_mb is not None else None + vram_total_gb = ( + round(vram_total_mb / 1024, 2) if vram_total_mb is not None else None + ) vram_util = ( round((vram_used_mb / vram_total_mb) * 100, 1) if vram_used_mb is not None and vram_total_mb is not None and vram_total_mb > 0 @@ -375,7 +385,8 @@ def get_primary_gpu_utilization() -> dict[str, Any]: def get_visible_gpu_utilization( - parent_visible_ids: Optional[list[int]], parent_cuda_visible_devices: Optional[str] = None + parent_visible_ids: Optional[list[int]], + parent_cuda_visible_devices: Optional[str] = None, ) -> dict[str, Any]: """Return utilization metrics for visible AMD GPUs.""" if parent_visible_ids is None: @@ -415,7 +426,9 @@ def get_visible_gpu_utilization( continue # Use the AMD-reported GPU ID, else the enumeration index. _parse_numeric # handles bare ints/floats/strings and the {"value", "unit"} dict shape. - raw_id = gpu_data.get("gpu", gpu_data.get("gpu_id", gpu_data.get("id", fallback_idx))) + raw_id = gpu_data.get( + "gpu", gpu_data.get("gpu_id", gpu_data.get("id", fallback_idx)) + ) parsed_id = _parse_numeric(raw_id) if parsed_id is None: logger.warning( diff --git a/studio/backend/utils/hardware/apple.py b/studio/backend/utils/hardware/apple.py index 62dbd10b8d..3252dd3fa6 100644 --- a/studio/backend/utils/hardware/apple.py +++ b/studio/backend/utils/hardware/apple.py @@ -154,7 +154,11 @@ def _load_iokit() -> ctypes.CDLL: def _load_cf() -> ctypes.CDLL: cf = ctypes.CDLL(_CF_PATH) cf.CFStringCreateWithCString.restype = ctypes.c_void_p - cf.CFStringCreateWithCString.argtypes = [ctypes.c_void_p, ctypes.c_char_p, ctypes.c_uint32] + cf.CFStringCreateWithCString.argtypes = [ + ctypes.c_void_p, + ctypes.c_char_p, + ctypes.c_uint32, + ] cf.CFStringGetCString.restype = ctypes.c_bool cf.CFStringGetCString.argtypes = [ ctypes.c_void_p, @@ -191,9 +195,17 @@ def _load_ioreport() -> ctypes.CDLL: ctypes.c_void_p, ] ior.IOReportCreateSamples.restype = ctypes.c_void_p - ior.IOReportCreateSamples.argtypes = [ctypes.c_void_p, ctypes.c_void_p, ctypes.c_void_p] + ior.IOReportCreateSamples.argtypes = [ + ctypes.c_void_p, + ctypes.c_void_p, + ctypes.c_void_p, + ] ior.IOReportCreateSamplesDelta.restype = ctypes.c_void_p - ior.IOReportCreateSamplesDelta.argtypes = [ctypes.c_void_p, ctypes.c_void_p, ctypes.c_void_p] + ior.IOReportCreateSamplesDelta.argtypes = [ + ctypes.c_void_p, + ctypes.c_void_p, + ctypes.c_void_p, + ] ior.IOReportChannelGetChannelName.restype = ctypes.c_void_p ior.IOReportChannelGetChannelName.argtypes = [ctypes.c_void_p] ior.IOReportChannelGetUnitLabel.restype = ctypes.c_void_p @@ -204,7 +216,9 @@ def _load_ioreport() -> ctypes.CDLL: def _cfstr(cf: ctypes.CDLL, text: str) -> int: - return cf.CFStringCreateWithCString(None, text.encode("utf-8"), _CF_STRING_ENCODING_UTF8) + return cf.CFStringCreateWithCString( + None, text.encode("utf-8"), _CF_STRING_ENCODING_UTF8 + ) def _from_cfstr(cf: ctypes.CDLL, ref: Optional[int]) -> str: @@ -235,7 +249,12 @@ class _SMCConnection: def _open(self) -> int: iterator = ctypes.c_uint32(0) matching = self._iokit.IOServiceMatching(b"AppleSMC") - if self._iokit.IOServiceGetMatchingServices(0, matching, ctypes.byref(iterator)) != 0: + if ( + self._iokit.IOServiceGetMatchingServices( + 0, matching, ctypes.byref(iterator) + ) + != 0 + ): raise OSError("AppleSMC service not found") try: conn = self._open_keys_endpoint(iterator.value) @@ -290,7 +309,9 @@ class _SMCConnection: try: key_id = _fourcc(key) info = self._read_key_info(key_id) - oval = self._call(_SMCKeyData(key = key_id, data8 = _SMC_CMD_READ_BYTES, key_info = info)) + oval = self._call( + _SMCKeyData(key = key_id, data8 = _SMC_CMD_READ_BYTES, key_info = info) + ) return bytes(oval.bytes[: info.data_size]) except OSError: return None @@ -386,7 +407,9 @@ class _IOReportEnergy: watts = _watts(energy, unit, elapsed_s) if watts is not None: total = (total or 0.0) + watts - if total is None or total < 0: # negative = counter reset; show -- not a bogus draw + if ( + total is None or total < 0 + ): # negative = counter reset; show -- not a bogus draw return None return round(total, 1) diff --git a/studio/backend/utils/hardware/hardware.py b/studio/backend/utils/hardware/hardware.py index 86dfa8a93f..a391a90edc 100644 --- a/studio/backend/utils/hardware/hardware.py +++ b/studio/backend/utils/hardware/hardware.py @@ -51,7 +51,9 @@ class DeviceType(str, Enum): DEVICE: Optional[DeviceType] = None CHAT_ONLY: bool = True # No CUDA GPU -> GGUF chat only (Mac, CPU-only, etc.) -IS_ROCM: bool = False # True when running on AMD ROCm (HIP) -- routes GPU monitoring to amd.py +IS_ROCM: bool = ( + False # True when running on AMD ROCm (HIP) -- routes GPU monitoring to amd.py +) def _backend_label(device: DeviceType) -> str: @@ -622,7 +624,9 @@ def get_gpu_utilization() -> Dict[str, Any]: result["backend"] = _backend_label(device) if IS_ROCM: # Fix unified-memory VRAM on AMD iGPUs (Strix Halo etc.). - _reconcile_primary_rocm_unified_memory(result, _get_parent_visible_gpu_spec()) + _reconcile_primary_rocm_unified_memory( + result, _get_parent_visible_gpu_spec() + ) return result # SMI unavailable. On Windows, use Performance Counters (Task Manager # source) for system-wide VRAM, covering cross-process usage torch can't see. @@ -681,7 +685,9 @@ def get_gpu_utilization() -> Dict[str, Any]: "temperature_c": None, "vram_used_gb": _used, "vram_total_gb": _total, - "vram_utilization_pct": round((_used / _total) * 100, 1) if _total > 0 else None, + "vram_utilization_pct": round((_used / _total) * 100, 1) + if _total > 0 + else None, "power_draw_w": None, "power_limit_w": None, "power_utilization_pct": None, @@ -765,7 +771,9 @@ def _apply_unified_memory_correction( device_metrics["vram_total_gb"] = torch_total_gb device_metrics["vram_used_gb"] = torch_used_gb device_metrics["vram_utilization_pct"] = ( - round((torch_used_gb / torch_total_gb) * 100, 1) if torch_total_gb > 0 else None + round((torch_used_gb / torch_total_gb) * 100, 1) + if torch_total_gb > 0 + else None ) logger.debug( "ROCm unified memory: replaced amd-smi VRAM (%.2f GB) with " @@ -776,7 +784,9 @@ def _apply_unified_memory_correction( ) -def _reconcile_rocm_unified_memory(utilization: Dict[str, Any], device_indices: list[int]) -> None: +def _reconcile_rocm_unified_memory( + utilization: Dict[str, Any], device_indices: list[int] +) -> None: """Fix amd-smi VRAM for ROCm unified-memory GPUs (e.g. Strix Halo). amd-smi reports only the dedicated slice; torch sees the full GTT pool. When @@ -928,7 +938,9 @@ def _get_parent_visible_gpu_spec() -> Dict[str, Any]: # stale HIP_VISIBLE_DEVICES on NVIDIA can't override CUDA_VISIBLE_DEVICES. _is_rocm_spec = IS_ROCM or ( "CUDA_VISIBLE_DEVICES" not in os.environ - and ("HIP_VISIBLE_DEVICES" in os.environ or "ROCR_VISIBLE_DEVICES" in os.environ) + and ( + "HIP_VISIBLE_DEVICES" in os.environ or "ROCR_VISIBLE_DEVICES" in os.environ + ) ) if _is_rocm_spec: hip_vis = os.environ.get("HIP_VISIBLE_DEVICES") @@ -1018,7 +1030,9 @@ def resolve_requested_gpu_ids(gpu_ids: Optional[list[int]]) -> list[int]: max_parent_id = max(parent_visible_ids) if physical_gpu_count > max_parent_id: # Count is plausibly physical, so enforce it. - out_of_range = [gpu_id for gpu_id in requested_ids if gpu_id >= physical_gpu_count] + out_of_range = [ + gpu_id for gpu_id in requested_ids if gpu_id >= physical_gpu_count + ] if out_of_range: raise ValueError( f"Invalid gpu_ids {requested_ids}: IDs must be physical GPU IDs " @@ -1026,7 +1040,9 @@ def resolve_requested_gpu_ids(gpu_ids: Optional[list[int]]) -> list[int]: f"Rejected IDs: {out_of_range}. Parent-visible GPUs: {parent_visible_ids}" ) - disallowed_ids = [gpu_id for gpu_id in requested_ids if gpu_id not in parent_visible_ids] + disallowed_ids = [ + gpu_id for gpu_id in requested_ids if gpu_id not in parent_visible_ids + ] if disallowed_ids: raise ValueError( f"Invalid gpu_ids {requested_ids}: requested GPUs {disallowed_ids} are " @@ -1047,7 +1063,9 @@ def _resolve_model_identifier_for_gpu_estimate( return config.base_model return config.identifier if config else model_name except Exception as e: - logger.debug("Could not resolve base model for GPU estimate '%s': %s", model_name, e) + logger.debug( + "Could not resolve base model for GPU estimate '%s': %s", model_name, e + ) return model_name @@ -1192,15 +1210,17 @@ def _estimate_fp16_model_size_bytes_from_vllm_utils(config) -> Optional[int]: synthetic_total_bytes, synthetic_total_bytes, ) - _, _, _, memory_left_for_kv_cache_gb = _vllm_utils.approximate_vllm_memory_usage( - config, - load_in_4bit = False, - load_in_8bit = False, - max_seq_length = 1, - gpu_memory_utilization = 1.0, - enable_lora = False, - account_for_gradients = False, - cuda_graph_overhead = False, + _, _, _, memory_left_for_kv_cache_gb = ( + _vllm_utils.approximate_vllm_memory_usage( + config, + load_in_4bit = False, + load_in_8bit = False, + max_seq_length = 1, + gpu_memory_utilization = 1.0, + enable_lora = False, + account_for_gradients = False, + cuda_graph_overhead = False, + ) ) finally: _vllm_utils.get_mem_info = original_get_mem_info @@ -1222,11 +1242,15 @@ def _estimate_fp16_model_size_bytes_from_vllm_utils(config) -> Optional[int]: def estimate_fp16_model_size_bytes( model_name: str, hf_token: Optional[str] = None ) -> tuple[Optional[int], str]: - estimate_model = _resolve_model_identifier_for_gpu_estimate(model_name, hf_token = hf_token) + estimate_model = _resolve_model_identifier_for_gpu_estimate( + model_name, hf_token = hf_token + ) total_params = None if "/" in estimate_model and not Path(estimate_model).exists(): - total_params = _get_hf_safetensors_total_params(estimate_model, hf_token = hf_token) + total_params = _get_hf_safetensors_total_params( + estimate_model, hf_token = hf_token + ) if total_params: return int(total_params * 2), "safetensors" @@ -1281,7 +1305,9 @@ def estimate_required_model_memory_gb( DEFAULT_TARGET_MODULES, ) - model_size_bytes, source = estimate_fp16_model_size_bytes(model_name, hf_token = hf_token) + model_size_bytes, source = estimate_fp16_model_size_bytes( + model_name, hf_token = hf_token + ) metadata: Dict[str, Any] = { "mode": "inference" if training_type is None else "training", "model_size_source": source, @@ -1304,7 +1330,9 @@ def estimate_required_model_memory_gb( return required_gb, metadata training_method = ( - "full" if training_type == "Full Finetuning" else ("qlora" if load_in_4bit else "lora") + "full" + if training_type == "Full Finetuning" + else ("qlora" if load_in_4bit else "lora") ) vram_config = TrainingVramConfig( training_method = training_method, @@ -1317,12 +1345,14 @@ def estimate_required_model_memory_gb( load_in_4bit = load_in_4bit, ) - estimate_model = _resolve_model_identifier_for_gpu_estimate(model_name, hf_token = hf_token) + estimate_model = _resolve_model_identifier_for_gpu_estimate( + model_name, hf_token = hf_token + ) config = _load_config_for_gpu_estimate(estimate_model, hf_token = hf_token) if config is not None: try: - vram_config.attention_implementation = _determine_attention_impl_for_gpu_estimate( - config + vram_config.attention_implementation = ( + _determine_attention_impl_for_gpu_estimate(config) ) except Exception as e: # Debug-level: fires every estimate on Windows ROCm (stub lacks Store); @@ -1508,7 +1538,9 @@ def auto_select_gpu_ids( return selected, metadata # Use only GPUs with verified VRAM data. - fallback_all = [c["index"] for c in gpu_candidates] if gpu_candidates else parent_ids + fallback_all = ( + [c["index"] for c in gpu_candidates] if gpu_candidates else parent_ids + ) metadata["selection_mode"] = "fallback_all" if ranked: fallback_usable = ranked[0]["free_gb"] + sum( @@ -1857,7 +1889,10 @@ def get_device_map(gpu_ids: Optional[list[int]] = None) -> str: # UUID/MIG masks can't be split into numeric IDs; >1 visible GPU # means multi-GPU sharding is intended. parent_visible_spec = _get_parent_visible_gpu_spec() - if parent_visible_spec["numeric_ids"] is None and get_visible_gpu_count() > 1: + if ( + parent_visible_spec["numeric_ids"] is None + and get_visible_gpu_count() > 1 + ): multi_gpu = True if multi_gpu: @@ -1886,7 +1921,9 @@ def raise_if_offloaded( offloaded = get_offloaded_device_map_entries(model) if not offloaded: return - example = ", ".join(f"{name}={placement}" for name, placement in list(offloaded.items())[:5]) + example = ", ".join( + f"{name}={placement}" for name, placement in list(offloaded.items())[:5] + ) raise ValueError( f"{context} does not support models loaded with CPU or disk offload. " f"device_map='{device_map}' produced offloaded modules: {example}" diff --git a/studio/backend/utils/hardware/nvidia.py b/studio/backend/utils/hardware/nvidia.py index f98ca4343e..eaabd9d55b 100644 --- a/studio/backend/utils/hardware/nvidia.py +++ b/studio/backend/utils/hardware/nvidia.py @@ -29,8 +29,12 @@ def _build_gpu_metrics( ) -> dict[str, Any]: return { **extra, - "vram_used_gb": round(vram_used_mb / 1024, 2) if vram_used_mb is not None else None, - "vram_total_gb": round(vram_total_mb / 1024, 2) if vram_total_mb is not None else None, + "vram_used_gb": round(vram_used_mb / 1024, 2) + if vram_used_mb is not None + else None, + "vram_total_gb": round(vram_total_mb / 1024, 2) + if vram_total_mb is not None + else None, "vram_utilization_pct": round((vram_used_mb / vram_total_mb) * 100, 1) if vram_used_mb is not None and vram_total_mb and vram_total_mb > 0 else None, @@ -42,7 +46,9 @@ def _build_gpu_metrics( } -def _visible_ordinal_map(parent_visible_ids: Optional[list[int]]) -> Optional[dict[int, int]]: +def _visible_ordinal_map( + parent_visible_ids: Optional[list[int]], +) -> Optional[dict[int, int]]: if parent_visible_ids is None: return None return {gpu_id: ordinal for ordinal, gpu_id in enumerate(parent_visible_ids)} @@ -108,7 +114,8 @@ def get_primary_gpu_utilization() -> dict[str, Any]: def get_visible_gpu_utilization( - parent_visible_ids: Optional[list[int]], parent_cuda_visible_devices: Optional[str] = None + parent_visible_ids: Optional[list[int]], + parent_cuda_visible_devices: Optional[str] = None, ) -> dict[str, Any]: # parent_visible_ids None (UUID/MIG mask): can't map nvidia-smi rows to # visible devices, so return empty rather than exposing all physical GPUs. @@ -176,7 +183,9 @@ def get_visible_gpu_utilization( index = idx, index_kind = "physical", visible_ordinal = ( - visible_ordinals[idx] if visible_ordinals is not None else len(devices) + visible_ordinals[idx] + if visible_ordinals is not None + else len(devices) ), gpu_utilization_pct = _parse_smi_value(parts[1]), temperature_c = _parse_smi_value(parts[2]), @@ -259,7 +268,9 @@ def get_backend_visible_gpu_info( "index": idx, "index_kind": "physical", "visible_ordinal": ( - visible_ordinals[idx] if visible_ordinals is not None else len(devices) + visible_ordinals[idx] + if visible_ordinals is not None + else len(devices) ), "name": name, "memory_total_gb": round(mem_total_mb / 1024, 2), diff --git a/studio/backend/utils/hardware/vram_estimation.py b/studio/backend/utils/hardware/vram_estimation.py index 86069ead3d..8dd99fe55e 100644 --- a/studio/backend/utils/hardware/vram_estimation.py +++ b/studio/backend/utils/hardware/vram_estimation.py @@ -16,7 +16,9 @@ from dataclasses import dataclass, field from typing import Dict, Optional QUANT_4BIT_FACTOR = 16 / 5 -DOUBLE_QUANT_4BIT_FACTOR = 3.6 # bnb_4bit_use_double_quant; see VRAM_ESTIMATION.md section 1 +DOUBLE_QUANT_4BIT_FACTOR = ( + 3.6 # bnb_4bit_use_double_quant; see VRAM_ESTIMATION.md section 1 +) CUDA_OVERHEAD_BYTES = int(1.4 * 1024**3) # calibrated on RTX 5070 Ti NON_FLASH_ATTENTION_FACTOR = ( 12.0 # eager attention score+workspace overhead; see VRAM_ESTIMATION.md section 5 @@ -144,7 +146,12 @@ class VramBreakdown: Weights/LoRA/optimizer/gradients shard across GPUs; activations do NOT (the GPU running a layer holds them). """ - shardable = self.model_weights + self.lora_adapters + self.optimizer_states + self.gradients + shardable = ( + self.model_weights + + self.lora_adapters + + self.optimizer_states + + self.gradients + ) per_gpu_fixed = self.activations + self.cuda_overhead return shardable // max(n_gpus, 1) + per_gpu_fixed @@ -184,7 +191,9 @@ def _compute_dense_layer_indices(text_config, total_layers: int) -> tuple: layer_types = getattr(text_config, "mlp_layer_types", None) if layer_types: return tuple( - i for i, t in enumerate(layer_types[:total_layers]) if str(t).lower() == "dense" + i + for i, t in enumerate(layer_types[:total_layers]) + if str(t).lower() == "dense" ) # Llama4TextConfig.__init__ auto-populates self.moe_layers from @@ -221,7 +230,9 @@ def _compute_dense_layer_indices(text_config, total_layers: int) -> tuple: if sparse_step is not None and sparse_step > 0: mlp_only_set = {int(i) for i in mlp_only} return tuple( - i for i in range(total_layers) if i in mlp_only_set or (i + 1) % sparse_step != 0 + i + for i in range(total_layers) + if i in mlp_only_set or (i + 1) % sparse_step != 0 ) return () @@ -249,7 +260,8 @@ def extract_arch_config(hf_config) -> Optional[ModelArchConfig]: intermediate_size = hidden_size * 4 if not all( - v is not None for v in (hidden_size, num_layers, num_heads, intermediate_size, vocab_size) + v is not None + for v in (hidden_size, num_layers, num_heads, intermediate_size, vocab_size) ): return None if num_heads <= 0: @@ -312,7 +324,9 @@ def extract_arch_config(hf_config) -> Optional[ModelArchConfig]: # intermediate_size. One shared_expert per MoE layer (modeling_llama4.py). intermediate_size_mlp_raw = _first_scalar(_moe_attr("intermediate_size_mlp")) dense_intermediate_size = ( - int(intermediate_size_mlp_raw) if intermediate_size_mlp_raw is not None else None + int(intermediate_size_mlp_raw) + if intermediate_size_mlp_raw is not None + else None ) if ( intermediate_size_mlp_raw is not None @@ -371,7 +385,9 @@ def extract_arch_config(hf_config) -> Optional[ModelArchConfig]: None, ) or 0, - quantization_skip_modules = list(quantization_config.get("llm_int8_skip_modules", []) or []), + quantization_skip_modules = list( + quantization_config.get("llm_int8_skip_modules", []) or [] + ), quant_4bit_factor = quant_4bit_factor, moe_has_dense_mlp = bool(getattr(text_config, "enable_moe_block", False)), dense_layer_indices = dense_layer_indices, @@ -458,7 +474,11 @@ def _per_layer_input_lora_params(arch: ModelArchConfig, r: int, target_modules) pli = arch.hidden_size_per_layer_input if pli <= 0: return 0 - targets = {target_modules} if isinstance(target_modules, str) else set(target_modules or []) + targets = ( + {target_modules} + if isinstance(target_modules, str) + else set(target_modules or []) + ) n_layers = arch.num_hidden_layers hd = arch.hidden_size total = 0 @@ -475,7 +495,11 @@ def _layer_attention_dims(arch: ModelArchConfig, layer_idx: int) -> tuple: layer_types = _layer_types(arch) layer_type = layer_types[layer_idx] is_sliding = layer_type == "sliding_attention" - head_dim = arch.global_head_dim if not is_sliding and arch.global_head_dim else _head_dim(arch) + head_dim = ( + arch.global_head_dim + if not is_sliding and arch.global_head_dim + else _head_dim(arch) + ) use_alt_attention = arch.attention_k_eq_v and not is_sliding num_kv_heads = ( arch.num_global_key_value_heads @@ -495,7 +519,9 @@ def _layer_mlp_size(arch: ModelArchConfig, layer_idx: int) -> int: return _dense_mlp_size(arch) -def _text_linear_dims(arch: ModelArchConfig, layer_idx: int) -> Dict[str, tuple[int, int]]: +def _text_linear_dims( + arch: ModelArchConfig, layer_idx: int +) -> Dict[str, tuple[int, int]]: hd = arch.hidden_size if _uses_structured_layer_shapes(arch): q_size, kv_size, has_k, has_v = _layer_attention_dims(arch, layer_idx) @@ -561,7 +587,9 @@ def _add_module_aliases(aliases: Dict[str, str], canonical: str, suffix: str) -> aliases[alias] = canonical -def _build_text_module_elements(arch: ModelArchConfig) -> tuple[Dict[str, int], Dict[str, str]]: +def _build_text_module_elements( + arch: ModelArchConfig, +) -> tuple[Dict[str, int], Dict[str, str]]: elements: Dict[str, int] = {} aliases: Dict[str, str] = {} @@ -572,8 +600,12 @@ def _build_text_module_elements(arch: ModelArchConfig) -> tuple[Dict[str, int], for layer_idx in range(arch.num_hidden_layers): layer_modules: Dict[str, int] = {} dims = _text_linear_dims(arch, layer_idx) - attn_dims = {name: dim for name, dim in dims.items() if name in ATTENTION_TARGET_MODULES} - mlp_dims = {name: dim for name, dim in dims.items() if name in MLP_TARGET_MODULES} + attn_dims = { + name: dim for name, dim in dims.items() if name in ATTENTION_TARGET_MODULES + } + mlp_dims = { + name: dim for name, dim in dims.items() if name in MLP_TARGET_MODULES + } if is_mla: # MLA splits q/o into q_a/q_b/kv_a/kv_b; emit a single self_attn @@ -620,7 +652,10 @@ def _build_text_module_elements(arch: ModelArchConfig) -> tuple[Dict[str, int], ) else: layer_modules.update( - {f"mlp.{name}": in_dim * out_dim for name, (in_dim, out_dim) in mlp_dims.items()} + { + f"mlp.{name}": in_dim * out_dim + for name, (in_dim, out_dim) in mlp_dims.items() + } ) if pli > 0: @@ -643,7 +678,10 @@ def _build_text_module_elements(arch: ModelArchConfig) -> tuple[Dict[str, int], for name, value in layer_modules.items() if ( name == "mlp" - or (name.startswith("mlp.") and not (is_sibling_experts and name == "mlp.experts")) + or ( + name.startswith("mlp.") + and not (is_sibling_experts and name == "mlp.experts") + ) ) ) experts_total = layer_modules.get("mlp.experts", 0) if is_sibling_experts else 0 @@ -698,7 +736,10 @@ def _compute_skipped_quantizable_elements(arch: ModelArchConfig) -> int: pruned = { canonical for canonical in matched - if not any(canonical != parent and canonical.startswith(f"{parent}.") for parent in matched) + if not any( + canonical != parent and canonical.startswith(f"{parent}.") + for parent in matched + ) } return sum(module_elements[canonical] for canonical in pruned) @@ -829,7 +870,9 @@ def _compute_layer_elements(arch: ModelArchConfig): mlp_total = _compute_dense_mlp_elements(arch) * n_layers layernorms = 2 * hd - per_layer_embed = arch.vocab_size_per_layer_input * arch.hidden_size_per_layer_input * n_layers + per_layer_embed = ( + arch.vocab_size_per_layer_input * arch.hidden_size_per_layer_input * n_layers + ) ple_text_linear = _per_layer_input_quantizable(arch) ple_norms = _per_layer_input_norm_elements(arch) embed_tokens = arch.vocab_size * hd + per_layer_embed + ple_norms @@ -851,7 +894,9 @@ def compute_model_weights_bytes( ) quantized = total_quantizable - skipped_quantizable return int( - quantized * 2 / arch.quant_4bit_factor + skipped_quantizable * 2 + non_quantizable * 2 + quantized * 2 / arch.quant_4bit_factor + + skipped_quantizable * 2 + + non_quantizable * 2 ) return int((total_quantizable + non_quantizable) * 2) @@ -907,7 +952,9 @@ def _lora_mlp_elements( return total -def compute_lora_params(arch: ModelArchConfig, lora_rank: int, target_modules: list) -> int: +def compute_lora_params( + arch: ModelArchConfig, lora_rank: int, target_modules: list +) -> int: all_linear = _targets_all_linear(target_modules) selected_modules = list(DEFAULT_TARGET_MODULES) if all_linear else target_modules hd = arch.hidden_size @@ -972,7 +1019,11 @@ def compute_lora_params(arch: ModelArchConfig, lora_rank: int, target_modules: l mlp_total = moe_mlp * n_moe + dense_only else: mlp_total = structured_dense_mlp - return attn_total + mlp_total + _per_layer_input_lora_params(arch, r, target_modules) + return ( + attn_total + + mlp_total + + _per_layer_input_lora_params(arch, r, target_modules) + ) elif n_experts > 1: attn_total = _lora_attn_elements(arch, r, selected_modules) * n_layers n_dense = arch.num_dense_layers @@ -1024,7 +1075,9 @@ def compute_lora_params(arch: ModelArchConfig, lora_rank: int, target_modules: l * n_layers ) - return attn_total + mlp_total + _per_layer_input_lora_params(arch, r, target_modules) + return ( + attn_total + mlp_total + _per_layer_input_lora_params(arch, r, target_modules) + ) def compute_lora_adapter_bytes(lora_params: int) -> int: @@ -1100,7 +1153,9 @@ def _per_layer_activation_bytes( # layer when hidden_size_per_layer_input is set (gemma4 modular:1141-1145). pli = arch.hidden_size_per_layer_input activation_ple = seq_len * batch_size * (arch.hidden_size + pli) if pli > 0 else 0 - return int((activation_qkv + residual_memory + activation_mlp + activation_ple) * 2 * 1.25) + return int( + (activation_qkv + residual_memory + activation_mlp + activation_ple) * 2 * 1.25 + ) def compute_activation_bytes( @@ -1121,12 +1176,14 @@ def compute_activation_bytes( if gc_multiplier is None: effective_layers = n_layers linear_bytes = sum( - _per_layer_activation_bytes(arch, i, batch_size, seq_len) for i in range(n_layers) + _per_layer_activation_bytes(arch, i, batch_size, seq_len) + for i in range(n_layers) ) else: effective_layers = gc_multiplier max_layer_bytes = max( - _per_layer_activation_bytes(arch, i, batch_size, seq_len) for i in range(n_layers) + _per_layer_activation_bytes(arch, i, batch_size, seq_len) + for i in range(n_layers) ) linear_bytes = int(max_layer_bytes * effective_layers) @@ -1149,7 +1206,9 @@ def compute_activation_bytes( ) -def estimate_training_vram(arch: ModelArchConfig, config: TrainingVramConfig) -> VramBreakdown: +def estimate_training_vram( + arch: ModelArchConfig, config: TrainingVramConfig +) -> VramBreakdown: method = config.training_method.lower() is_lora = method in ("qlora", "lora") load_in_4bit = config.load_in_4bit or method == "qlora" diff --git a/studio/backend/utils/inference/inference_config.py b/studio/backend/utils/inference/inference_config.py index 05eb08067c..bbc65c35fa 100644 --- a/studio/backend/utils/inference/inference_config.py +++ b/studio/backend/utils/inference/inference_config.py @@ -28,7 +28,10 @@ def _load_family_defaults(): return json_path = ( - Path(__file__).parent.parent.parent / "assets" / "configs" / "inference_defaults.json" + Path(__file__).parent.parent.parent + / "assets" + / "configs" + / "inference_defaults.json" ) try: with open(json_path, "r", encoding = "utf-8") as f: diff --git a/studio/backend/utils/llama_cpp_freshness.py b/studio/backend/utils/llama_cpp_freshness.py index 87d0d2ec01..905200a348 100644 --- a/studio/backend/utils/llama_cpp_freshness.py +++ b/studio/backend/utils/llama_cpp_freshness.py @@ -153,7 +153,9 @@ def _fetch_latest_release_tag(repo: str, timeout: float = 5.0) -> Optional[str]: return newest["tag_name"] -def latest_published_release(repo: str, *, force_refresh: bool = False) -> Optional[str]: +def latest_published_release( + repo: str, *, force_refresh: bool = False +) -> Optional[str]: """Latest release tag for `repo`. Memo + disk-cached (24h TTL). None when offline and never previously cached.""" if not repo: diff --git a/studio/backend/utils/llama_cpp_update.py b/studio/backend/utils/llama_cpp_update.py index 8b90d36bc5..fc7af35ff9 100644 --- a/studio/backend/utils/llama_cpp_update.py +++ b/studio/backend/utils/llama_cpp_update.py @@ -110,7 +110,10 @@ def _installer_script() -> Optional[Path]: return Path(env) here = Path(__file__).resolve() for up in here.parents: - for cand in (up / "install_llama_prebuilt.py", up / "studio" / "install_llama_prebuilt.py"): + for cand in ( + up / "install_llama_prebuilt.py", + up / "studio" / "install_llama_prebuilt.py", + ): if cand.is_file(): return cand return None @@ -169,7 +172,9 @@ def _installed_build_number(binary: Optional[str]) -> Optional[int]: if not binary: return None try: - proc = subprocess.run([binary, "--version"], capture_output = True, text = True, timeout = 20) + proc = subprocess.run( + [binary, "--version"], capture_output = True, text = True, timeout = 20 + ) except Exception: # pragma: no cover - defensive return None m = re.search(r"version:\s*(\d+)", (proc.stderr or "") + (proc.stdout or "")) @@ -235,7 +240,9 @@ def _source_build_status(binary: str, *, force_refresh: bool) -> Optional[dict]: # Suppress only when the source build is reliably newer/equal; unknown # version (the involuntary source-build case) is treated as behind. update_available = ( - installed_build is None or latest_build is None or installed_build < latest_build + installed_build is None + or latest_build is None + or installed_build < latest_build ) with _job_lock: job = dict(_job) @@ -325,7 +332,9 @@ def _rocm_install_args(asset: Optional[str]) -> list[str]: return ["--has-rocm"] -def _run_update(install_dir: Path, repo: str, asset: Optional[str], script: Path) -> None: +def _run_update( + install_dir: Path, repo: str, asset: Optional[str], script: Path +) -> None: """Worker: put the backend into a maintenance state, run the installer for the latest prebuilt, then refresh caches so the next load uses the new build.""" backend = None @@ -338,7 +347,8 @@ def _run_update(install_dir: Path, repo: str, asset: Optional[str], script: Path backend = get_llama_cpp_backend() except Exception as exc: logger.debug( - "llama update: backend unavailable, skipping load coordination", error = str(exc) + "llama update: backend unavailable, skipping load coordination", + error = str(exc), ) backend = None @@ -395,7 +405,9 @@ def _run_update(install_dir: Path, repo: str, asset: Optional[str], script: Path m = _PROGRESS_LINE_RE.search(line) if m is None: continue - fraction = min(float(m.group(1)) / 100.0, 1.0) * _DOWNLOAD_PROGRESS_CEILING + fraction = ( + min(float(m.group(1)) / 100.0, 1.0) * _DOWNLOAD_PROGRESS_CEILING + ) with _job_lock: _job["progress"] = max(_job.get("progress") or 0.0, fraction) returncode = proc.wait() @@ -417,7 +429,9 @@ def _run_update(install_dir: Path, repo: str, asset: Optional[str], script: Path try: latest_published_release(repo, force_refresh = True) except Exception as exc: # pragma: no cover - network defensive - logger.debug("llama update: post-install freshness refresh failed", error = str(exc)) + logger.debug( + "llama update: post-install freshness refresh failed", error = str(exc) + ) new_marker = read_install_marker(_find_binary()) new_tag = (new_marker or {}).get("tag") or (new_marker or {}).get("release_tag") diff --git a/studio/backend/utils/models/checkpoints.py b/studio/backend/utils/models/checkpoints.py index 5a992926ec..3ccfceac86 100644 --- a/studio/backend/utils/models/checkpoints.py +++ b/studio/backend/utils/models/checkpoints.py @@ -92,7 +92,9 @@ def scan_checkpoints( name_part = parts[0] idx = name_part.find("_") if idx > 0: - metadata["base_model"] = name_part[:idx] + "/" + name_part[idx + 1 :] + metadata["base_model"] = ( + name_part[:idx] + "/" + name_part[idx + 1 :] + ) else: metadata["base_model"] = name_part @@ -122,7 +124,9 @@ def scan_checkpoints( ) models.append((item.name, checkpoints, metadata)) - logger.debug(f"Found model: {item.name} with {len(checkpoints)} checkpoint(s)") + logger.debug( + f"Found model: {item.name} with {len(checkpoints)} checkpoint(s)" + ) # Sort by modification time (newest first) models.sort(key = lambda x: Path(x[1][0][1]).stat().st_mtime, reverse = True) diff --git a/studio/backend/utils/models/model_config.py b/studio/backend/utils/models/model_config.py index 7245343eeb..ba6cc613f4 100644 --- a/studio/backend/utils/models/model_config.py +++ b/studio/backend/utils/models/model_config.py @@ -55,13 +55,19 @@ def _env_offline() -> bool: # ── Model size extraction ──────────────────────────────────── import re as _re -_MODEL_SIZE_RE = _re.compile(r"(?:^|[-_/])(\d+\.?\d*)\s*([bm])(?:$|[-_/])", _re.IGNORECASE) +_MODEL_SIZE_RE = _re.compile( + r"(?:^|[-_/])(\d+\.?\d*)\s*([bm])(?:$|[-_/])", _re.IGNORECASE +) # MoE active-parameter pattern: "A3B", "A3.5B", etc. -_ACTIVE_SIZE_RE = _re.compile(r"(?:^|[-_/])a(\d+\.?\d*)\s*([bm])(?:$|[-_/])", _re.IGNORECASE) +_ACTIVE_SIZE_RE = _re.compile( + r"(?:^|[-_/])a(\d+\.?\d*)\s*([bm])(?:$|[-_/])", _re.IGNORECASE +) # Gemma 3n/4 effective-parameter pattern: "E2B", "E4B" -- the runtime # footprint (MatFormer + per-layer embeddings), which is the size that # matters for size-gated policies like sub-3B speculative-decoding fallback. -_EFFECTIVE_SIZE_RE = _re.compile(r"(?:^|[-_/])e(\d+\.?\d*)\s*([bm])(?:$|[-_/])", _re.IGNORECASE) +_EFFECTIVE_SIZE_RE = _re.compile( + r"(?:^|[-_/])e(\d+\.?\d*)\s*([bm])(?:$|[-_/])", _re.IGNORECASE +) def extract_model_size_b(model_id: str) -> float | None: @@ -555,7 +561,10 @@ def _raw_config_has_vision_config( if model_type in _AUDIO_ONLY_MODEL_TYPES: return False return ( - any(isinstance(x, str) and x.endswith(_VLM_ARCH_SUFFIXES) for x in architectures) + any( + isinstance(x, str) and x.endswith(_VLM_ARCH_SUFFIXES) + for x in architectures + ) or "vision_config" in config or "img_processor" in config or "image_token_index" in config @@ -627,7 +636,9 @@ except Exception as exc: ) -def _is_vision_model_subprocess(model_name: str, hf_token: Optional[str] = None) -> Optional[bool]: +def _is_vision_model_subprocess( + model_name: str, hf_token: Optional[str] = None +) -> Optional[bool]: """Run is_vision_model in a subprocess with transformers 5.x. Spawns a clean subprocess with .venv_t5/ on sys.path so AutoConfig @@ -773,7 +784,9 @@ def is_vision_model(model_name: str, hf_token: Optional[str] = None) -> bool: return False -def _is_vision_model_uncached(model_name: str, hf_token: Optional[str] = None) -> Optional[bool]: +def _is_vision_model_uncached( + model_name: str, hf_token: Optional[str] = None +) -> Optional[bool]: """Uncached vision detection; use is_vision_model() instead. Returns True/False for definitive results, or None on transient errors @@ -855,7 +868,9 @@ _AUDIO_TOKEN_PATTERNS = { and "<|text_start|>" in tokens and "<|text_end|>" in tokens ), - "snac": lambda tokens: (sum(1 for t in tokens if t.startswith(" 10000), + "snac": lambda tokens: ( + sum(1 for t in tokens if t.startswith(" 10000 + ), } @@ -877,7 +892,9 @@ def detect_audio_type(model_name: str, hf_token: Optional[str] = None) -> Option return result -def _detect_audio_from_tokenizer(model_name: str, hf_token: Optional[str] = None) -> Optional[str]: +def _detect_audio_from_tokenizer( + model_name: str, hf_token: Optional[str] = None +) -> Optional[str]: """Detect audio type from tokenizer special tokens. Checks local HF cache first, then fetches tokenizer_config.json from HF; @@ -938,7 +955,9 @@ def _detect_audio_from_tokenizer(model_name: str, hf_token: Optional[str] = None return None except Exception as e: - logger.debug(f"Could not detect audio type from tokenizer for {model_name}: {e}") + logger.debug( + f"Could not detect audio type from tokenizer for {model_name}: {e}" + ) return None @@ -1209,7 +1228,9 @@ def detect_mtp_file(path: str, search_root: Optional[str] = None) -> Optional[st if not (name.startswith("mtp-") and name.endswith(".gguf")): continue stem = name[len("mtp-") : -len(".gguf")] - if not stem or (weight_name is not None and not weight_name.startswith(stem)): + if not stem or ( + weight_name is not None and not weight_name.startswith(stem) + ): continue try: if f.is_file(): @@ -1253,7 +1274,8 @@ def detect_gguf_model(path: str) -> Optional[str]: ( f for f in _iter_gguf_files(p) - if not _is_mmproj(f.name) and not _is_mtp_drafter(f"{f.parent.name}/{f.name}") + if not _is_mmproj(f.name) + and not _is_mtp_drafter(f"{f.parent.name}/{f.name}") ), key = lambda f: f.stat().st_size, reverse = True, @@ -1450,7 +1472,9 @@ def _iter_hf_cache_snapshots(repo_id: str): yield from snap_dirs -def _list_gguf_variants_from_hf_cache(repo_id: str) -> Optional[tuple[list[GgufVariantInfo], bool]]: +def _list_gguf_variants_from_hf_cache( + repo_id: str, +) -> Optional[tuple[list[GgufVariantInfo], bool]]: """Variants from the local HF cache snapshot, or None if not cached.""" for snap in _iter_hf_cache_snapshots(repo_id): variants, has_vision = list_local_gguf_variants(str(snap)) @@ -1652,14 +1676,17 @@ def _detect_gguf_from_hf_cache(repo_id: str) -> Optional[str]: rel_files = [ rel for f in _iter_gguf_files(snap, recursive = True) - if not _is_mtp_drafter(rel := f.relative_to(snap).as_posix()) and not _is_mmproj(f.name) + if not _is_mtp_drafter(rel := f.relative_to(snap).as_posix()) + and not _is_mmproj(f.name) ] if rel_files: return _pick_best_gguf(rel_files) return None -def detect_gguf_model_remote(repo_id: str, hf_token: Optional[str] = None) -> Optional[str]: +def detect_gguf_model_remote( + repo_id: str, hf_token: Optional[str] = None +) -> Optional[str]: """Return the best GGUF filename in a HF repo, or None. Retries (3 attempts, 1s/2s/4s backoff) on transient HF Hub failures: a @@ -1705,7 +1732,9 @@ def detect_gguf_model_remote(repo_id: str, hf_token: Optional[str] = None) -> Op ) return cached - logger.warning(f"Could not check GGUF files for '{repo_id}' after 3 attempts: {last_err}") + logger.warning( + f"Could not check GGUF files for '{repo_id}' after 3 attempts: {last_err}" + ) return None @@ -1828,7 +1857,9 @@ def _looks_like_lora_adapter(model_dir: Path) -> bool: ) -def scan_trained_models(outputs_dir: str = str(outputs_root())) -> List[Tuple[str, str, str]]: +def scan_trained_models( + outputs_dir: str = str(outputs_root()), +) -> List[Tuple[str, str, str]]: """Scan outputs folder for trained Studio models. Returns: @@ -1894,7 +1925,9 @@ def scan_exported_models( # Flat GGUF export (e.g. exports/gemma-3-4b-it-finetune-gguf/). # Skip mmproj (vision projection) files — not loadable as main models. - gguf_files = [f for f in _iter_gguf_files(run_dir) if not _is_mmproj(f.name)] + gguf_files = [ + f for f in _iter_gguf_files(run_dir) if not _is_mmproj(f.name) + ] if gguf_files: base_model = None export_meta = run_dir / "export_metadata.json" @@ -1968,7 +2001,9 @@ def scan_exported_models( # Fallback: base model from ./outputs/{run_name}/adapter_config.json if not base_model: - outputs_adapter_cfg = resolve_output_dir(run_dir.name) / "adapter_config.json" + outputs_adapter_cfg = ( + resolve_output_dir(run_dir.name) / "adapter_config.json" + ) try: if outputs_adapter_cfg.exists(): cfg = json.loads(outputs_adapter_cfg.read_text()) @@ -2001,7 +2036,9 @@ def get_base_model_from_checkpoint(checkpoint_path: str) -> Optional[str]: config = json.load(f) base_model = config.get("base_model_name_or_path") if base_model: - logger.info("Detected base model from adapter_config.json: %s", base_model) + logger.info( + "Detected base model from adapter_config.json: %s", base_model + ) return base_model config_path = checkpoint_path_obj / "config.json" @@ -2066,7 +2103,9 @@ def get_base_model_from_lora(lora_path: str) -> Optional[str]: config = json.load(f) base_model = config.get("base_model_name_or_path") if base_model: - logger.info(f"Detected base model from adapter_config.json: {base_model}") + logger.info( + f"Detected base model from adapter_config.json: {base_model}" + ) return base_model # Fallback: try training_args.bin (requires torch) @@ -2126,7 +2165,9 @@ def load_model_defaults(model_name: str) -> Dict[str, Any]: if config_path.is_file(): with open(config_path, "r", encoding = "utf-8") as f: config = yaml.safe_load(f) or {} - logger.info(f"Loaded model defaults from {config_path} (via mapping)") + logger.info( + f"Loaded model defaults from {config_path} (via mapping)" + ) return config # For local paths (e.g. /home/.../Spark-TTS-0.5B/LLM from @@ -2195,11 +2236,17 @@ class ModelConfig: is_lora: bool # LoRA adapter? is_gguf: bool = False # GGUF model? is_audio: bool = False # TTS audio model? - audio_type: Optional[str] = None # Audio codec type: 'snac', 'csm', 'bicodec', 'dac' + audio_type: Optional[str] = ( + None # Audio codec type: 'snac', 'csm', 'bicodec', 'dac' + ) has_audio_input: bool = False # Accepts audio input (ASR/speech understanding) gguf_file: Optional[str] = None # Full path to the .gguf file (local mode) - gguf_mmproj_file: Optional[str] = None # Full path to the mmproj .gguf file (vision projection) - gguf_mtp_file: Optional[str] = None # Full path to the separate MTP drafter (local mode) + gguf_mmproj_file: Optional[str] = ( + None # Full path to the mmproj .gguf file (vision projection) + ) + gguf_mtp_file: Optional[str] = ( + None # Full path to the separate MTP drafter (local mode) + ) gguf_hf_repo: Optional[str] = ( None # HF repo ID for -hf mode (e.g. "unsloth/gemma-3-4b-it-GGUF") ) @@ -2337,7 +2384,9 @@ class ModelConfig: gguf_is_vision = True logger.info(f"Detected mmproj for vision: {mmproj_file}") elif base_is_vision: - logger.warning(f"Base model is vision but no mmproj file found in {gguf_dir}") + logger.warning( + f"Base model is vision but no mmproj file found in {gguf_dir}" + ) # Separate MTP drafter sibling (Gemma 4), mirroring mmproj. mtp_file = detect_mtp_file(gguf_file, search_root = companion_root) @@ -2406,11 +2455,15 @@ class ModelConfig: # Auto-detect LoRA for local paths (adapter_config.json on disk) if not is_lora and is_local: detected_base = ( - get_base_model_from_lora(path) if _looks_like_lora_adapter(Path(path)) else None + get_base_model_from_lora(path) + if _looks_like_lora_adapter(Path(path)) + else None ) if detected_base: is_lora = True - logger.info(f"Auto-detected local LoRA adapter at '{path}' (base: {detected_base})") + logger.info( + f"Auto-detected local LoRA adapter at '{path}' (base: {detected_base})" + ) # Auto-detect LoRA for remote HF models. When offline, huggingface_hub # raises OfflineModeIsEnabled in ~0ms; we fall through to the cache. @@ -2424,14 +2477,18 @@ class ModelConfig: is_lora = True logger.info(f"Auto-detected remote LoRA adapter: '{identifier}'") except Exception as e: - logger.debug(f"Could not check remote LoRA status for '{identifier}': {e}") + logger.debug( + f"Could not check remote LoRA status for '{identifier}': {e}" + ) # API may have failed; adapter_config.json could still be cached. if not is_lora: for snap in _iter_hf_cache_snapshots(identifier): if (snap / "adapter_config.json").is_file(): is_lora = True - logger.info(f"Auto-detected cached LoRA adapter: '{identifier}'") + logger.info( + f"Auto-detected cached LoRA adapter: '{identifier}'" + ) break # Handle LoRA adapters @@ -2445,7 +2502,9 @@ class ModelConfig: try: from huggingface_hub import hf_hub_download - config_path = hf_hub_download(identifier, "adapter_config.json", token = hf_token) + config_path = hf_hub_download( + identifier, "adapter_config.json", token = hf_token + ) with open(config_path, "r") as f: adapter_config = json.load(f) base_model = adapter_config.get("base_model_name_or_path") @@ -2506,7 +2565,9 @@ class ModelConfig: # Resolve display names via the 'local_models' parameter if " (Active)" in selected or " (Ready)" in selected: - clean_display_name = selected.replace(" (Active)", "").replace(" (Ready)", "") + clean_display_name = selected.replace(" (Active)", "").replace( + " (Ready)", "" + ) if local_models: for local_display, local_path in local_models: if local_display == clean_display_name: diff --git a/studio/backend/utils/native_path_leases.py b/studio/backend/utils/native_path_leases.py index 7d8514abc8..58b4a3c653 100644 --- a/studio/backend/utils/native_path_leases.py +++ b/studio/backend/utils/native_path_leases.py @@ -68,7 +68,9 @@ def native_path_leases_supported() -> bool: return True -def child_env_without_native_path_secret(env: Mapping[str, str] | None = None) -> dict[str, str]: +def child_env_without_native_path_secret( + env: Mapping[str, str] | None = None, +) -> dict[str, str]: """Return a child-process env with the native path lease secret removed.""" if env is None: @@ -80,7 +82,9 @@ def child_env_without_native_path_secret(env: Mapping[str, str] | None = None) - return cleaned -def run_without_native_path_secret(target: Callable[..., Any], *args: Any, **kwargs: Any) -> Any: +def run_without_native_path_secret( + target: Callable[..., Any], *args: Any, **kwargs: Any +) -> Any: """Run a multiprocessing child target without the native path lease secret.""" global _CACHED_LEASE_SECRET, _SCRUB_SAVED_SECRET @@ -147,7 +151,9 @@ def verify_native_path_lease( raise NativePathLeaseError("Native path is no longer accessible.") from exc _reject_network_or_device_path(resolved) if not _same_native_path(resolved, path): - raise NativePathLeaseError("Native path grant no longer resolves to the selected path.") + raise NativePathLeaseError( + "Native path grant no longer resolves to the selected path." + ) grant = NativePathGrant( operation = str(payload["operation"]), @@ -211,7 +217,9 @@ def _decode_secret() -> bytes: if encoded is None and _SCRUB_SAVED_SECRET is not None: encoded = _SCRUB_SAVED_SECRET if not encoded: - raise NativePathLeaseError("Native path grants require the managed desktop backend.") + raise NativePathLeaseError( + "Native path grants require the managed desktop backend." + ) try: secret = _b64decode(encoded) except Exception as exc: @@ -262,7 +270,9 @@ def _validate_payload( ) missing = [key for key in required if key not in payload] if missing: - raise NativePathLeaseError("Native path grant payload is missing required fields.") + raise NativePathLeaseError( + "Native path grant payload is missing required fields." + ) if _required_int(payload, "version") != 1: raise NativePathLeaseError("Native path grant version is unsupported.") if payload["operation"] != operation: @@ -341,13 +351,19 @@ def _reject_network_or_device_path(path: Path) -> None: rest = normalized[4:] is_local_drive = len(rest) >= 3 and rest[0].isalpha() and rest[1:3] == ":\\" if not is_local_drive: - raise NativePathLeaseError("Network paths are not supported for native grants.") + raise NativePathLeaseError( + "Network paths are not supported for native grants." + ) elif normalized.startswith("\\\\"): - raise NativePathLeaseError("Network paths are not supported for native grants.") + raise NativePathLeaseError( + "Network paths are not supported for native grants." + ) if os.name != "nt": for root in ("/dev", "/proc", "/sys"): if path.is_relative_to(root): - raise NativePathLeaseError("Device and virtual filesystem paths are not supported.") + raise NativePathLeaseError( + "Device and virtual filesystem paths are not supported." + ) if "\x00" in text: raise NativePathLeaseError("Native path contains invalid characters.") @@ -379,7 +395,9 @@ def _optional_int(value: Any) -> int | None: def _required_int(payload: dict[str, Any], key: str) -> int: raw = payload.get(key) if raw is None: - raise NativePathLeaseError("Native path grant payload is missing required fields.") + raise NativePathLeaseError( + "Native path grant payload is missing required fields." + ) try: return int(raw) except (TypeError, ValueError) as exc: diff --git a/studio/backend/utils/paths/storage_roots.py b/studio/backend/utils/paths/storage_roots.py index d336bc2e71..cd7f9e3313 100644 --- a/studio/backend/utils/paths/storage_roots.py +++ b/studio/backend/utils/paths/storage_roots.py @@ -276,7 +276,9 @@ def _setup_cache_env() -> None: user hasn't, so explicit overrides are honored. """ root = cache_root() - xdg_cache = Path(os.environ.get("XDG_CACHE_HOME", Path.home() / ".cache")).expanduser() + xdg_cache = Path( + os.environ.get("XDG_CACHE_HOME", Path.home() / ".cache") + ).expanduser() hf_default = xdg_cache / "huggingface" defaults: dict[str, str] = { "HF_HOME": str(hf_default), @@ -309,7 +311,9 @@ def ensure_studio_directories() -> None: _setup_cache_env() -def _clean_relative_path(path_value: str, *, strip_prefixes: tuple[str, ...] = ()) -> Path: +def _clean_relative_path( + path_value: str, *, strip_prefixes: tuple[str, ...] = () +) -> Path: path = Path(path_value).expanduser() parts = [part for part in path.parts if part not in ("", ".")] while parts and parts[0] in strip_prefixes: @@ -348,7 +352,8 @@ def _assert_contained(resolved: Path, root: Path) -> None: resolved_real.relative_to(root_real) except ValueError as exc: raise ValueError( - f"path escapes root: {resolved!s} -> {resolved_real!s} " f"is not under {root_real!s}" + f"path escapes root: {resolved!s} -> {resolved_real!s} " + f"is not under {root_real!s}" ) from exc @@ -452,7 +457,9 @@ def resolve_dataset_path(path_value: str) -> Path: return path except ValueError: continue - raise ValueError(f"dataset path must be relative or under a dataset root: {raw!r}") + raise ValueError( + f"dataset path must be relative or under a dataset root: {raw!r}" + ) parts = [part for part in Path(path_value).parts if part not in ("", ".")] if parts[:2] == ["assets", "datasets"]: diff --git a/studio/backend/utils/studio_version.py b/studio/backend/utils/studio_version.py index 9c18070fbb..946873a2a1 100644 --- a/studio/backend/utils/studio_version.py +++ b/studio/backend/utils/studio_version.py @@ -40,7 +40,9 @@ def _path_is_in_site_packages(path: Path) -> bool: def _is_source_checkout(repo_root: Path) -> bool: - return (repo_root / ".git").exists() and not _path_is_in_site_packages(Path(__file__).resolve()) + return (repo_root / ".git").exists() and not _path_is_in_site_packages( + Path(__file__).resolve() + ) def _exact_git_studio_tag(repo_root: Path) -> str | None: diff --git a/studio/backend/utils/transformers_version.py b/studio/backend/utils/transformers_version.py index 6e1571d5ce..cd3af40148 100644 --- a/studio/backend/utils/transformers_version.py +++ b/studio/backend/utils/transformers_version.py @@ -313,7 +313,9 @@ def _check_tokenizer_config_needs_v5(model_name: str) -> bool: _tokenizer_class_cache[model_name] = result return result except Exception as exc: - logger.debug("Could not fetch tokenizer_config.json for '%s': %s", model_name, exc) + logger.debug( + "Could not fetch tokenizer_config.json for '%s': %s", model_name, exc + ) _tokenizer_class_cache[model_name] = False return False @@ -354,7 +356,9 @@ def _load_config_json(model_name: str) -> dict | None: return None -def _config_matches_tier(cfg: dict, architectures: set[str], model_types: set[str]) -> bool: +def _config_matches_tier( + cfg: dict, architectures: set[str], model_types: set[str] +) -> bool: archs = cfg.get("architectures", []) if any(a in architectures for a in archs): return True @@ -574,7 +578,8 @@ def _venv_dir_is_valid(venv_dir: str, packages: tuple[str, ...]) -> bool: pkg_name_norm = pkg_name.replace("-", "_") # Directory must exist. if not any( - (Path(venv_dir) / d).is_dir() for d in (pkg_name_norm, pkg_name_norm.replace("_", "-")) + (Path(venv_dir) / d).is_dir() + for d in (pkg_name_norm, pkg_name_norm.replace("_", "-")) ): return False # Unpinned packages: existence is enough. @@ -669,7 +674,9 @@ def _ensure_venv_dir(venv_dir: str, packages: tuple[str, ...], label: str) -> bo if _venv_dir_is_valid(venv_dir, packages): return True - logger.warning("%s not found or incomplete at %s -- installing at runtime", label, venv_dir) + logger.warning( + "%s not found or incomplete at %s -- installing at runtime", label, venv_dir + ) shutil.rmtree(venv_dir, ignore_errors = True) os.makedirs(venv_dir, exist_ok = True) for pkg in packages: @@ -681,7 +688,9 @@ def _ensure_venv_dir(venv_dir: str, packages: tuple[str, ...], label: str) -> bo def _ensure_venv_t5_530_exists() -> bool: """Ensure .venv_t5_530/ exists with transformers 5.3.0.""" - return _ensure_venv_dir(_VENV_T5_530_DIR, _VENV_T5_530_PACKAGES, "transformers 5.3.0") + return _ensure_venv_dir( + _VENV_T5_530_DIR, _VENV_T5_530_PACKAGES, "transformers 5.3.0" + ) def _ensure_venv_t5_550_exists() -> bool: @@ -812,12 +821,15 @@ def ensure_transformers_version(model_name: str) -> None: _deactivate_5x() if not ensure_fn(): raise RuntimeError( - f"Cannot activate transformers {target_version}: " f"venv missing at {venv_dir}" + f"Cannot activate transformers {target_version}: " + f"venv missing at {venv_dir}" ) logger.info("Activating transformers %s…", target_version) _activate_venv(venv_dir, f"transformers {target_version}") else: - logger.info("Reverting to default transformers %s…", TRANSFORMERS_DEFAULT_VERSION) + logger.info( + "Reverting to default transformers %s…", TRANSFORMERS_DEFAULT_VERSION + ) _deactivate_5x() final = _get_in_memory_version() diff --git a/studio/backend/utils/update_status.py b/studio/backend/utils/update_status.py index ad9dabcf36..bf9afcd23a 100644 --- a/studio/backend/utils/update_status.py +++ b/studio/backend/utils/update_status.py @@ -72,7 +72,11 @@ def detect_install_source() -> str: try: dist = distribution(PACKAGE_NAME) except PackageNotFoundError: - return "local_repo" if _path_has_git_parent(_repo_root_from_this_file()) else "unknown" + return ( + "local_repo" + if _path_has_git_parent(_repo_root_from_this_file()) + else "unknown" + ) try: direct_url = dist.read_text("direct_url.json") @@ -141,7 +145,9 @@ def get_studio_update_status(current_version: str) -> dict[str, Any]: current_version = current_version, latest_version = None, install_source = install_source, - reason = "invalid_current_version" if current_version != "dev" else "dev_build", + reason = "invalid_current_version" + if current_version != "dev" + else "dev_build", ) latest_result = get_latest_pypi_version() if latest_result.latest_version is None: @@ -209,7 +215,9 @@ def get_latest_pypi_version() -> LatestVersionResult: error = "Could not check PyPI update metadata.", ) - ttl = PYPI_SUCCESS_TTL_SECONDS if result.latest_version else PYPI_FAILURE_TTL_SECONDS + ttl = ( + PYPI_SUCCESS_TTL_SECONDS if result.latest_version else PYPI_FAILURE_TTL_SECONDS + ) with _cache_condition: _latest_version_cache = _LatestVersionCacheEntry( result = result, @@ -253,7 +261,9 @@ def _fetch_latest_pypi_version() -> LatestVersionResult: error = "Could not reach PyPI for update metadata.", ) - latest = payload.get("info", {}).get("version") if isinstance(payload, dict) else None + latest = ( + payload.get("info", {}).get("version") if isinstance(payload, dict) else None + ) if not isinstance(latest, str) or not latest.strip(): return LatestVersionResult( latest_version = None, @@ -355,4 +365,9 @@ def _parse_current_version(current_version: str) -> Version | None: def _utc_now_iso() -> str: - return datetime.now(timezone.utc).replace(microsecond = 0).isoformat().replace("+00:00", "Z") + return ( + datetime.now(timezone.utc) + .replace(microsecond = 0) + .isoformat() + .replace("+00:00", "Z") + ) diff --git a/studio/backend/utils/utils.py b/studio/backend/utils/utils.py index 3818253ac9..1d94fc39aa 100644 --- a/studio/backend/utils/utils.py +++ b/studio/backend/utils/utils.py @@ -19,7 +19,9 @@ logger = get_logger(__name__) # Never return raw exception text to clients; log server-side, return generic. -def safe_error_detail(error: Exception, fallback: str = "An internal error occurred") -> str: +def safe_error_detail( + error: Exception, fallback: str = "An internal error occurred" +) -> str: """Map an exception to a generic, client-safe message (never raw ``str(error)``, which can leak paths). Log the real exception server-side. """ @@ -36,7 +38,9 @@ def safe_error_detail(error: Exception, fallback: str = "An internal error occur return fallback -def safe_curated_detail(error: Exception, fallback: str = "An internal error occurred") -> str: +def safe_curated_detail( + error: Exception, fallback: str = "An internal error occurred" +) -> str: """Client-safe text for curated domain/validation exceptions. Keeps the message (paths stripped) instead of a generic fallback; for known diff --git a/studio/backend/utils/wheel_utils.py b/studio/backend/utils/wheel_utils.py index 98697df83c..4e005f8b50 100644 --- a/studio/backend/utils/wheel_utils.py +++ b/studio/backend/utils/wheel_utils.py @@ -19,7 +19,9 @@ from utils.subprocess_compat import windows_hidden_subprocess_kwargs _logger = logging.getLogger(__name__) -FLASH_ATTN_RELEASE_BASE_URL = "https://github.com/Dao-AILab/flash-attention/releases/download" +FLASH_ATTN_RELEASE_BASE_URL = ( + "https://github.com/Dao-AILab/flash-attention/releases/download" +) @functools.lru_cache(maxsize = 1) diff --git a/studio/install_llama_prebuilt.py b/studio/install_llama_prebuilt.py index 46aa14b041..baf0f9b22e 100644 --- a/studio/install_llama_prebuilt.py +++ b/studio/install_llama_prebuilt.py @@ -140,7 +140,9 @@ UPSTREAM_REPO = "ggml-org/llama.cpp" UPSTREAM_RELEASES_API = f"https://api.github.com/repos/{UPSTREAM_REPO}/releases/latest" -TEST_MODEL_URL = "https://huggingface.co/ggml-org/models/resolve/main/tinyllamas/stories260K.gguf" +TEST_MODEL_URL = ( + "https://huggingface.co/ggml-org/models/resolve/main/tinyllamas/stories260K.gguf" +) TEST_MODEL_SHA256 = "270cba1bd5109f42d03350f60406024560464db173c0e387d91f0426d3bd256d" VALIDATION_MODEL_CACHE_DIRNAME = ".cache" VALIDATION_MODEL_CACHE_FILENAME = "stories260K.gguf" @@ -211,8 +213,12 @@ _BLACKWELL_MIN_SM = 120 # windows-cuda build at or above this already covers Blackwell and makes the # older pinned 13.1 fallback unnecessary (cuda-12.4 is below it). _BLACKWELL_MIN_TOOLKIT = (12, 8) -_PINNED_BLACKWELL_LLAMA_SHA256 = "31ddb8b42d7ab4a47cab8c48c397519f580ca502df7e73f3ab396eacc16c8e8d" -_PINNED_BLACKWELL_CUDART_SHA256 = "f96935e7e385e3b2d0189239077c10fe8fd7e95690fea4afec455b1b6c7e3f18" +_PINNED_BLACKWELL_LLAMA_SHA256 = ( + "31ddb8b42d7ab4a47cab8c48c397519f580ca502df7e73f3ab396eacc16c8e8d" +) +_PINNED_BLACKWELL_CUDART_SHA256 = ( + "f96935e7e385e3b2d0189239077c10fe8fd7e95690fea4afec455b1b6c7e3f18" +) def _cuda_runtime_lines_for_major(major: int) -> list[str]: @@ -443,7 +449,9 @@ _LOG_TO_STDOUT = False def log(message: str) -> None: - print(f"[llama-prebuilt] {message}", file = sys.stdout if _LOG_TO_STDOUT else sys.stderr) + print( + f"[llama-prebuilt] {message}", file = sys.stdout if _LOG_TO_STDOUT else sys.stderr + ) def log_lines(lines: Iterable[str]) -> None: @@ -498,7 +506,9 @@ class _CrossHostAuthStrippingRedirectHandler(urllib.request.HTTPRedirectHandler) def redirect_request(self, req, fp, code, msg, headers, newurl): new_request = super().redirect_request(req, fp, code, msg, headers, newurl) - if new_request is not None and parsed_hostname(newurl) != parsed_hostname(req.full_url): + if new_request is not None and parsed_hostname(newurl) != parsed_hostname( + req.full_url + ): new_request.headers.pop("Authorization", None) new_request.unredirected_hdrs.pop("Authorization", None) return new_request @@ -750,9 +760,9 @@ def refs_match(candidate_ref: str | None, requested_ref: str | None) -> bool: candidate_commit = normalize_source_commit(candidate_ref) requested_commit = normalize_source_commit(requested_ref) if candidate_commit and requested_commit: - return candidate_commit.startswith(requested_commit) or requested_commit.startswith( - candidate_commit - ) + return candidate_commit.startswith( + requested_commit + ) or requested_commit.startswith(candidate_commit) return False @@ -856,7 +866,11 @@ class DownloadProgress: self.last_emit = 0.0 term_ok = os.environ.get("TERM", "").lower() != "dumb" self.stream = ( - sys.stderr if sys.stderr.isatty() else sys.stdout if sys.stdout.isatty() else sys.stderr + sys.stderr + if sys.stderr.isatty() + else sys.stdout + if sys.stdout.isatty() + else sys.stderr ) self.is_tty = term_ok and self.stream.isatty() self.completed = False @@ -890,7 +904,10 @@ class DownloadProgress: if self.is_tty: elapsed = now - self.start_time if not self.has_rendered_tty_progress: - if self.total_bytes is not None and downloaded_bytes >= self.total_bytes: + if ( + self.total_bytes is not None + and downloaded_bytes >= self.total_bytes + ): return if elapsed < TTY_PROGRESS_START_DELAY_SECONDS: return @@ -913,7 +930,10 @@ class DownloadProgress: percent = int((downloaded_bytes * 100) / max(self.total_bytes, 1)) step = self.milestone_step milestone_percent = min((percent // step) * step, 100) - if milestone_percent > self.last_milestone_percent and milestone_percent < 100: + if ( + milestone_percent > self.last_milestone_percent + and milestone_percent < 100 + ): self.last_milestone_percent = milestone_percent should_emit = True else: @@ -966,7 +986,11 @@ def download_bytes( content_length = response.headers.get("Content-Length") if content_length and content_length.isdigit(): total_bytes = int(content_length) - progress = DownloadProgress(progress_label, total_bytes) if progress_label else None + progress = ( + DownloadProgress(progress_label, total_bytes) + if progress_label + else None + ) data = bytearray() while True: chunk = response.read(1024 * 1024) @@ -996,13 +1020,17 @@ def fetch_json(url: str) -> Any: data = download_bytes( url, timeout = 30, - headers = github_api_headers(url) if is_github_api_url(url) else auth_headers(url), + headers = github_api_headers(url) + if is_github_api_url(url) + else auth_headers(url), ) except urllib.error.HTTPError as exc: if exc.code == 403 and is_github_api_url(url): hint = "" if not (os.environ.get("GH_TOKEN") or os.environ.get("GITHUB_TOKEN")): - hint = "; set GH_TOKEN or GITHUB_TOKEN to avoid GitHub API rate limits" + hint = ( + "; set GH_TOKEN or GITHUB_TOKEN to avoid GitHub API rate limits" + ) raise RuntimeError(f"GitHub API returned 403 for {url}{hint}") from exc raise if not data: @@ -1011,7 +1039,9 @@ def fetch_json(url: str) -> Any: try: payload = json.loads(data.decode("utf-8")) except (UnicodeDecodeError, json.JSONDecodeError) as exc: - last_decode_exc = RuntimeError(f"downloaded invalid JSON from {url}: {exc}") + last_decode_exc = RuntimeError( + f"downloaded invalid JSON from {url}: {exc}" + ) else: if not isinstance(payload, dict) and not isinstance(payload, list): raise RuntimeError( @@ -1045,7 +1075,9 @@ def download_file(url: str, destination: Path) -> None: content_length = response.headers.get("Content-Length") if content_length and content_length.isdigit(): total_bytes = int(content_length) - progress = DownloadProgress(f"Downloading {destination.name}", total_bytes) + progress = DownloadProgress( + f"Downloading {destination.name}", total_bytes + ) downloaded_bytes = 0 while True: chunk = response.read(1024 * 1024) @@ -1070,7 +1102,9 @@ def download_file(url: str, destination: Path) -> None: pass if attempt >= HTTP_FETCH_ATTEMPTS or not is_retryable_url_error(exc): raise - log(f"download failed ({attempt}/{HTTP_FETCH_ATTEMPTS}) for {url}: {exc}; retrying") + log( + f"download failed ({attempt}/{HTTP_FETCH_ATTEMPTS}) for {url}: {exc}; retrying" + ) sleep_backoff(attempt, exc = exc) assert last_exc is not None raise last_exc @@ -1082,7 +1116,9 @@ def download_file_verified( normalized_expected = normalize_sha256_digest(expected_sha256) if not normalized_expected: download_file(url, destination) - log(f"downloaded {label} without a published sha256; relying on install validation") + log( + f"downloaded {label} without a published sha256; relying on install validation" + ) return for attempt in range(1, 3): @@ -1180,7 +1216,9 @@ def latest_upstream_release_tag() -> str: payload = fetch_json(UPSTREAM_RELEASES_API) tag = payload.get("tag_name") if not isinstance(tag, str) or not tag: - raise RuntimeError(f"latest release tag was missing from {UPSTREAM_RELEASES_API}") + raise RuntimeError( + f"latest release tag was missing from {UPSTREAM_RELEASES_API}" + ) return tag @@ -1215,13 +1253,19 @@ def iter_release_payloads_by_time( yield github_release(repo, published_release_tag) return - if requested_tag and requested_tag != "latest" and is_release_tag_like(requested_tag): + if ( + requested_tag + and requested_tag != "latest" + and is_release_tag_like(requested_tag) + ): try: yield github_release(repo, requested_tag) return except urllib.error.HTTPError as exc: if exc.code == 404: - log(f"release tag {requested_tag} not found in {repo}; scanning recent releases") + log( + f"release tag {requested_tag} not found in {repo}; scanning recent releases" + ) else: raise except Exception: @@ -1229,15 +1273,21 @@ def iter_release_payloads_by_time( releases = [ release - for release in github_releases(repo, max_pages = DEFAULT_GITHUB_RELEASE_SCAN_MAX_PAGES) - if isinstance(release, dict) and not release.get("draft") and not release.get("prerelease") + for release in github_releases( + repo, max_pages = DEFAULT_GITHUB_RELEASE_SCAN_MAX_PAGES + ) + if isinstance(release, dict) + and not release.get("draft") + and not release.get("prerelease") ] releases.sort(key = release_time_sort_key, reverse = True) for release in releases: yield release -def direct_release_matches_request(*, release_tag: str, llama_tag: str, requested_tag: str) -> bool: +def direct_release_matches_request( + *, release_tag: str, llama_tag: str, requested_tag: str +) -> bool: if requested_tag == "latest": return True for candidate in (release_tag, llama_tag): @@ -1519,7 +1569,12 @@ def direct_upstream_release_plan( install_kind = "macos-x64", ) ) - elif host.is_linux and host.is_x86_64 and not host.has_usable_nvidia and not host.has_rocm: + elif ( + host.is_linux + and host.is_x86_64 + and not host.has_usable_nvidia + and not host.has_rocm + ): # ROCm hosts are excluded: this ggml-org path ships no per-gfx ROCm # asset, so they fall through to the empty-attempts raise (HIP source # build) rather than silently getting a CPU binary on a GPU host. @@ -1609,7 +1664,9 @@ def resolve_simple_install_release_plans( max_release_fallbacks = max_release_fallbacks, ) requested_tag = normalized_requested_llama_tag(llama_tag) - allow_older_release_fallback = requested_tag == "latest" and not published_release_tag + allow_older_release_fallback = ( + requested_tag == "latest" and not published_release_tag + ) # macOS: pin the last upstream build that loads on a pre-26 host instead of # fetching the latest (macOS 26 only) build and walking back release by # release. No-op on macOS 26+, unknown version, non-macOS, and the fork. @@ -1623,7 +1680,9 @@ def resolve_simple_install_release_plans( last_error: PrebuiltFallback | None = None try: - releases = iter_release_payloads_by_time(repo, published_release_tag, requested_tag) + releases = iter_release_payloads_by_time( + repo, published_release_tag, requested_tag + ) for release in releases: try: plan = direct_upstream_release_plan(release, host, repo, requested_tag) @@ -1646,13 +1705,17 @@ def resolve_simple_install_release_plans( except PrebuiltFallback: raise except Exception as exc: - raise PrebuiltFallback(f"failed to inspect published releases in {repo}: {exc}") from exc + raise PrebuiltFallback( + f"failed to inspect published releases in {repo}: {exc}" + ) from exc if plans: return requested_tag, plans if last_error is not None: raise last_error - raise PrebuiltFallback(f"no installable published llama.cpp releases were found in {repo}") + raise PrebuiltFallback( + f"no installable published llama.cpp releases were found in {repo}" + ) def normalized_requested_llama_tag(requested_tag: str | None) -> str: @@ -1704,7 +1767,9 @@ def parse_cuda_visible_devices(value: str | None) -> list[str] | None: return [token.strip() for token in raw.split(",") if token.strip()] -def supports_explicit_visible_device_matching(visible_devices: list[str] | None) -> bool: +def supports_explicit_visible_device_matching( + visible_devices: list[str] | None, +) -> bool: if not visible_devices: return False for token in visible_devices: @@ -1754,7 +1819,9 @@ def dir_provides_exact_library(directory: str | Path, library: str) -> bool: return candidate.exists() and (candidate.is_file() or candidate.is_symlink()) -def linux_runtime_dirs_for_required_libraries(required_libraries: Iterable[str]) -> list[str]: +def linux_runtime_dirs_for_required_libraries( + required_libraries: Iterable[str], +) -> list[str]: required = [library for library in required_libraries if library] candidates: list[str | Path] = [] @@ -1770,7 +1837,9 @@ def linux_runtime_dirs_for_required_libraries(required_libraries: Iterable[str]) value = os.environ.get(name) if value: cuda_roots.append(Path(value)) - cuda_roots.extend(Path(path) for path in glob_paths("/usr/local/cuda", "/usr/local/cuda-*")) + cuda_roots.extend( + Path(path) for path in glob_paths("/usr/local/cuda", "/usr/local/cuda-*") + ) for root in cuda_roots: candidates.extend( @@ -1795,7 +1864,8 @@ def linux_runtime_dirs_for_required_libraries(required_libraries: Iterable[str]) ) ) candidates.extend( - Path(path) for path in glob_paths("/usr/local/lib/ollama/cuda_v*", "/usr/lib/wsl/lib") + Path(path) + for path in glob_paths("/usr/local/lib/ollama/cuda_v*", "/usr/lib/wsl/lib") ) candidates.extend(Path(path) for path in python_runtime_dirs()) candidates.extend(Path(path) for path in ldconfig_runtime_dirs(required)) @@ -1807,7 +1877,9 @@ def linux_runtime_dirs_for_required_libraries(required_libraries: Iterable[str]) matched: list[tuple[int, str]] = [] for directory in resolved: base = Path(directory) - provided = sum(1 for library in required if dir_provides_exact_library(directory, library)) + provided = sum( + 1 for library in required if dir_provides_exact_library(directory, library) + ) if provided: matched.append((provided, directory)) @@ -1828,7 +1900,9 @@ def detected_linux_runtime_lines() -> tuple[list[str], dict[str, list[str]]]: matching_dirs: list[str] = [] for library in required: matched_dirs = [ - directory for directory in dirs if any(Path(directory).glob(f"{library}*")) + directory + for directory in dirs + if any(Path(directory).glob(f"{library}*")) ] if not matched_dirs: library_matches = {} @@ -1865,13 +1939,17 @@ def parse_published_artifact(raw: Any) -> PublishedLlamaArtifact | None: if not isinstance(asset_name, str) or not asset_name: raise ValueError("artifact.asset_name was missing or not a string") if not isinstance(install_kind, str) or not install_kind: - raise ValueError(f"artifact {asset_name} install_kind was missing or not a string") + raise ValueError( + f"artifact {asset_name} install_kind was missing or not a string" + ) supported_sms_raw = raw.get("supported_sms", []) if not isinstance(supported_sms_raw, (list, tuple)): raise ValueError(f"artifact {asset_name} supported_sms must be a list or tuple") if any(not isinstance(value, (int, str)) for value in supported_sms_raw): - raise ValueError(f"artifact {asset_name} supported_sms entries must be ints or strings") + raise ValueError( + f"artifact {asset_name} supported_sms entries must be ints or strings" + ) supported_sms = normalize_compute_caps(supported_sms_raw) min_sm_raw = raw.get("min_sm") @@ -1880,7 +1958,9 @@ def parse_published_artifact(raw: Any) -> PublishedLlamaArtifact | None: min_sm = int(min_sm_raw) if min_sm_raw is not None else None max_sm = int(max_sm_raw) if max_sm_raw is not None else None except (TypeError, ValueError) as exc: - raise ValueError(f"artifact {asset_name} min_sm/max_sm were not integers") from exc + raise ValueError( + f"artifact {asset_name} min_sm/max_sm were not integers" + ) from exc runtime_line = raw.get("runtime_line") coverage_class = raw.get("coverage_class") bundle_profile = raw.get("bundle_profile") @@ -1903,14 +1983,20 @@ def parse_published_artifact(raw: Any) -> PublishedLlamaArtifact | None: ) mapped_raw = raw.get("mapped_targets", []) mapped_targets = ( - [value.strip() for value in mapped_raw if isinstance(value, str) and value.strip()] + [ + value.strip() + for value in mapped_raw + if isinstance(value, str) and value.strip() + ] if isinstance(mapped_raw, (list, tuple)) else [] ) return PublishedLlamaArtifact( asset_name = asset_name, install_kind = install_kind, - runtime_line = runtime_line if isinstance(runtime_line, str) and runtime_line else None, + runtime_line = runtime_line + if isinstance(runtime_line, str) and runtime_line + else None, coverage_class = coverage_class if isinstance(coverage_class, str) and coverage_class else None, @@ -1987,7 +2073,9 @@ def parse_published_release_bundle( try: artifact = parse_published_artifact(raw_artifact) except ValueError as exc: - log(f"published artifact ignored for {repo}@{release_tag} artifact[{index}]: {exc}") + log( + f"published artifact ignored for {repo}@{release_tag} artifact[{index}]: {exc}" + ) continue if artifact is not None: artifacts.append(artifact) @@ -2006,7 +2094,9 @@ def parse_published_release_bundle( release_tag = release_tag, upstream_tag = upstream_tag, manifest_sha256 = manifest_sha256, - source_repo = source_repo if isinstance(source_repo, str) and source_repo else None, + source_repo = source_repo + if isinstance(source_repo, str) and source_repo + else None, source_repo_url = source_repo_url if isinstance(source_repo_url, str) and source_repo_url else None, @@ -2067,12 +2157,18 @@ def parse_approved_release_checksums( artifacts: dict[str, ApprovedArtifactHash] = {} for asset_name, raw_entry in artifacts_payload.items(): if not isinstance(asset_name, str) or not asset_name: - raise RuntimeError("published checksum asset used a non-string artifact key") + raise RuntimeError( + "published checksum asset used a non-string artifact key" + ) if not isinstance(raw_entry, dict): - raise RuntimeError(f"published checksum entry for {asset_name} was not an object") + raise RuntimeError( + f"published checksum entry for {asset_name} was not an object" + ) digest = normalize_sha256_digest(raw_entry.get("sha256")) if not digest: - raise RuntimeError(f"published checksum entry for {asset_name} omitted a valid sha256") + raise RuntimeError( + f"published checksum entry for {asset_name} omitted a valid sha256" + ) repo_value = raw_entry.get("repo") kind_value = raw_entry.get("kind") artifacts[asset_name] = ApprovedArtifactHash( @@ -2093,7 +2189,9 @@ def parse_approved_release_checksums( repo = repo, release_tag = release_tag, upstream_tag = upstream_tag, - source_repo = source_repo if isinstance(source_repo, str) and source_repo else None, + source_repo = source_repo + if isinstance(source_repo, str) and source_repo + else None, source_repo_url = source_repo_url if isinstance(source_repo_url, str) and source_repo_url else None, @@ -2112,7 +2210,9 @@ def parse_approved_release_checksums( ) -def load_approved_release_checksums(repo: str, release_tag: str) -> ApprovedReleaseChecksums: +def load_approved_release_checksums( + repo: str, release_tag: str +) -> ApprovedReleaseChecksums: try: release = github_release(repo, release_tag) except Exception as exc: @@ -2146,7 +2246,9 @@ def iter_published_release_bundles( else github_releases(repo, max_pages = DEFAULT_GITHUB_RELEASE_SCAN_MAX_PAGES) ) for release in releases: - if not published_release_tag and (release.get("draft") or release.get("prerelease")): + if not published_release_tag and ( + release.get("draft") or release.get("prerelease") + ): continue try: bundle = parse_published_release_bundle(repo, release) @@ -2159,13 +2261,18 @@ def iter_published_release_bundles( yield bundle -def _artifact_covers_sms(artifact: PublishedLlamaArtifact, host_sms: Iterable[str]) -> bool: +def _artifact_covers_sms( + artifact: PublishedLlamaArtifact, host_sms: Iterable[str] +) -> bool: """True when every host SM is listed in the artifact's supported_sms and falls within its [min_sm, max_sm] range.""" if not artifact.supported_sms or artifact.min_sm is None or artifact.max_sm is None: return False supported = {str(value) for value in artifact.supported_sms} - return all(sm in supported and artifact.min_sm <= int(sm) <= artifact.max_sm for sm in host_sms) + return all( + sm in supported and artifact.min_sm <= int(sm) <= artifact.max_sm + for sm in host_sms + ) def _sm_range(artifact: PublishedLlamaArtifact) -> int: @@ -2222,9 +2329,13 @@ def linux_cuda_choice_from_release( # below is arch-agnostic and applies to both. cuda_install_kind = "linux-arm64-cuda" if host.is_arm64 else "linux-cuda" published_artifacts = [ - artifact for artifact in release.artifacts if artifact.install_kind == cuda_install_kind + artifact + for artifact in release.artifacts + if artifact.install_kind == cuda_install_kind ] - published_asset_names = sorted(artifact.asset_name for artifact in published_artifacts) + published_asset_names = sorted( + artifact.asset_name for artifact in published_artifacts + ) selection_log.append( "linux_cuda_selection: published_assets=" + (",".join(published_asset_names) if published_asset_names else "none") @@ -2260,7 +2371,9 @@ def linux_cuda_choice_from_release( attempts: list[AssetChoice] = [] seen_attempts: set[str] = set() - def add_attempt(artifact: PublishedLlamaArtifact, asset_url: str, reason: str) -> None: + def add_attempt( + artifact: PublishedLlamaArtifact, asset_url: str, reason: str + ) -> None: asset_name = artifact.asset_name if asset_name in seen_attempts: return @@ -2297,7 +2410,9 @@ def linux_cuda_choice_from_release( asset_name = artifact.asset_name asset_url = release.assets.get(asset_name) if not asset_url: - selection_log.append(f"linux_cuda_selection: reject {asset_name} missing asset") + selection_log.append( + f"linux_cuda_selection: reject {asset_name} missing asset" + ) continue if not host_sms and artifact.coverage_class != "portable": selection_log.append( @@ -2325,7 +2440,9 @@ def linux_cuda_choice_from_release( supported_sms = {str(value) for value in artifact.supported_sms} missing_sms = [sm for sm in host_sms if sm not in supported_sms] out_of_range_sms = [ - sm for sm in host_sms if not (artifact.min_sm <= int(sm) <= artifact.max_sm) + sm + for sm in host_sms + if not (artifact.min_sm <= int(sm) <= artifact.max_sm) ] reasons: list[str] = [] if missing_sms: @@ -2369,7 +2486,8 @@ def linux_cuda_choice_from_release( return None selection_log.append( - "linux_cuda_selection: attempt_order=" + ",".join(choice.name for choice in attempts) + "linux_cuda_selection: attempt_order=" + + ",".join(choice.name for choice in attempts) ) for attempt in attempts: attempt.selection_log = list(selection_log) + [ @@ -2387,7 +2505,9 @@ def latest_published_linux_cuda_tag(host: HostInfo, published_repo: str) -> str def iter_upstream_releases() -> Iterable[dict[str, Any]]: - for release in github_releases(UPSTREAM_REPO, max_pages = DEFAULT_GITHUB_RELEASE_SCAN_MAX_PAGES): + for release in github_releases( + UPSTREAM_REPO, max_pages = DEFAULT_GITHUB_RELEASE_SCAN_MAX_PAGES + ): if release.get("draft") or release.get("prerelease"): continue yield release @@ -2431,7 +2551,9 @@ def validated_checksums_for_bundle( return checksums -def published_release_matches_request(bundle: PublishedReleaseBundle, requested_ref: str) -> bool: +def published_release_matches_request( + bundle: PublishedReleaseBundle, requested_ref: str +) -> bool: if requested_ref == "latest": return True for candidate in ( @@ -2486,7 +2608,9 @@ def resolve_published_release( raise PrebuiltFallback( f"no usable published llama.cpp releases were available in {repo}" ) - raise PrebuiltFallback(f"no published llama.cpp releases were available in {repo}") + raise PrebuiltFallback( + f"no published llama.cpp releases were available in {repo}" + ) raise PrebuiltFallback( f"no published prebuilt release in {repo} matched upstream tag {normalized_requested}" @@ -2545,7 +2669,9 @@ def iter_resolved_published_releases( return if normalized_requested == "latest": - raise PrebuiltFallback(f"no published llama.cpp releases were available in {repo}") + raise PrebuiltFallback( + f"no published llama.cpp releases were available in {repo}" + ) raise PrebuiltFallback( f"no published prebuilt release in {repo} matched upstream tag {normalized_requested}" @@ -2603,10 +2729,14 @@ def resolve_requested_install_tag( ).bundle.upstream_tag -def exact_source_archive_hash(checksums: ApprovedReleaseChecksums) -> ApprovedArtifactHash | None: +def exact_source_archive_hash( + checksums: ApprovedReleaseChecksums, +) -> ApprovedArtifactHash | None: if not checksums.source_commit: return None - return checksums.artifacts.get(exact_source_archive_logical_name(checksums.source_commit)) + return checksums.artifacts.get( + exact_source_archive_logical_name(checksums.source_commit) + ) def source_clone_url_for_release( @@ -2626,8 +2756,12 @@ def source_build_plan_for_release(release: ResolvedPublishedRelease) -> SourceBu exact_source = exact_source_archive_hash(checksums) source_repo = checksums.source_repo or release.bundle.source_repo source_repo_url = checksums.source_repo_url or release.bundle.source_repo_url - requested_source_ref = checksums.requested_source_ref or release.bundle.requested_source_ref - resolved_source_ref = checksums.resolved_source_ref or release.bundle.resolved_source_ref + requested_source_ref = ( + checksums.requested_source_ref or release.bundle.requested_source_ref + ) + resolved_source_ref = ( + checksums.resolved_source_ref or release.bundle.resolved_source_ref + ) source_commit = checksums.source_commit or release.bundle.source_commit source_ref_kind = checksums.source_ref_kind or release.bundle.source_ref_kind source_url = source_clone_url_for_release(checksums, release.bundle) @@ -2643,8 +2777,14 @@ def source_build_plan_for_release(release: ResolvedPublishedRelease) -> SourceBu resolved_source_ref = resolved_source_ref, source_commit = source_commit, ) - source_ref = checkout_friendly_ref(source_ref_kind, resolved_source_ref or requested_source_ref) - if source_url and source_ref and source_ref_kind in {"tag", "branch", "pull", "commit"}: + source_ref = checkout_friendly_ref( + source_ref_kind, resolved_source_ref or requested_source_ref + ) + if ( + source_url + and source_ref + and source_ref_kind in {"tag", "branch", "pull", "commit"} + ): return SourceBuildPlan( source_url = source_url, source_ref = source_ref, @@ -2837,7 +2977,9 @@ def detect_host() -> HostInfo: # ROCm host as NVIDIA and short-circuit the ROCm path. try: listing = run_capture([nvidia_smi, "-L"], timeout = 20) - gpu_lines = [line for line in listing.stdout.splitlines() if line.startswith("GPU ")] + gpu_lines = [ + line for line in listing.stdout.splitlines() if line.startswith("GPU ") + ] if gpu_lines: has_physical_nvidia = True has_usable_nvidia = visible_device_tokens != [] @@ -2995,7 +3137,9 @@ def detect_host() -> HostInfo: # AMD torch wheels ship hipInfo.exe into the venv Scripts dir # (next to python.exe) -- resolvable on driver-only hosts where no # SDK dir exists, so a standalone rerun can still detect the GPU. - _venv_candidate = os.path.join(os.path.dirname(sys.executable), f"{name}.exe") + _venv_candidate = os.path.join( + os.path.dirname(sys.executable), f"{name}.exe" + ) if os.path.isfile(_venv_candidate): return _venv_candidate return None @@ -3082,7 +3226,9 @@ def _apply_host_overrides( return host -def published_repo_for_host(host: HostInfo, *, linux_amd_tooling_present: bool = False) -> str: +def published_repo_for_host( + host: HostInfo, *, linux_amd_tooling_present: bool = False +) -> str: """The release repo setup.sh / setup.ps1 pick for this host: macOS always the fork (ggml-org macOS bundles need too-new macOS); else CPU-only Linux/Windows -> ggml-org upstream (the fork ships no CPU bundle) and any usable GPU (NVIDIA @@ -3092,7 +3238,9 @@ def published_repo_for_host(host: HostInfo, *, linux_amd_tooling_present: bool = if host.is_macos: return DEFAULT_PUBLISHED_REPO has_gpu = ( - host.has_usable_nvidia or host.has_rocm or (host.is_linux and linux_amd_tooling_present) + host.has_usable_nvidia + or host.has_rocm + or (host.is_linux and linux_amd_tooling_present) ) return DEFAULT_PUBLISHED_REPO if has_gpu else UPSTREAM_REPO @@ -3194,7 +3342,9 @@ def detect_torch_cuda_runtime_preference(host: HostInfo) -> CudaRuntimePreferenc try: cuda_available = bool(torch.cuda.is_available()) except Exception as exc: - selection_log.append(f"torch_cuda_preference: torch.cuda.is_available() failed: {exc}") + selection_log.append( + f"torch_cuda_preference: torch.cuda.is_available() failed: {exc}" + ) return CudaRuntimePreference(runtime_line = None, selection_log = selection_log) if not cuda_available: @@ -3282,10 +3432,14 @@ def windows_cuda_attempts( f"{preferred_runtime_line} unavailable_or_incompatible" ) else: - selection_log.append("windows_cuda_selection: no Torch runtime preference available") + selection_log.append( + "windows_cuda_selection: no Torch runtime preference available" + ) runtime_order.extend( - runtime_line for runtime_line in normal_runtime_lines if runtime_line not in runtime_order + runtime_line + for runtime_line in normal_runtime_lines + if runtime_line not in runtime_order ) # Keep every driver-compatible line reachable as a fallback, so a line gated # out by the driver version still drops to an older major (cuda13 -> cuda12). @@ -3309,7 +3463,9 @@ def windows_cuda_attempts( # Track whatever minor llama.cpp actually ships for this major # (cuda13 -> 13.1, 13.3, ...). Skip the line when the release has no # matching asset instead of guessing a now-missing name. - runtime = _published_windows_cuda_runtime(upstream_assets, major, host.driver_cuda_version) + runtime = _published_windows_cuda_runtime( + upstream_assets, major, host.driver_cuda_version + ) if runtime is None: selection_log.append( f"windows_cuda_selection: no driver-supported asset for {runtime_line}" @@ -3405,7 +3561,8 @@ def _drop_blackwell_incapable_windows_cuda( return [ attempt for attempt in attempts - if attempt.install_kind != "windows-cuda" or _windows_cuda_attempt_covers_blackwell(attempt) + if attempt.install_kind != "windows-cuda" + or _windows_cuda_attempt_covers_blackwell(attempt) ] @@ -3430,7 +3587,10 @@ def _pinned_windows_cuda_fallback( caps = normalize_compute_caps(host.compute_caps) if not caps or int(caps[-1]) < _BLACKWELL_MIN_SM: return None - if any(_windows_cuda_attempt_covers_blackwell(attempt) for attempt in existing_cuda_attempts): + if any( + _windows_cuda_attempt_covers_blackwell(attempt) + for attempt in existing_cuda_attempts + ): return None tag = _PINNED_BLACKWELL_FALLBACK_TAG runtime = _PINNED_BLACKWELL_FALLBACK_RUNTIME @@ -3505,7 +3665,9 @@ def published_windows_cuda_attempts( ) -> list[AssetChoice]: selection_log = list(release.selection_log) + list(selection_preamble) published_artifacts = [ - artifact for artifact in release.artifacts if artifact.install_kind == "windows-cuda" + artifact + for artifact in release.artifacts + if artifact.install_kind == "windows-cuda" ] artifacts_by_runtime: dict[str, list[PublishedLlamaArtifact]] = {} for artifact in published_artifacts: @@ -3549,7 +3711,9 @@ def published_windows_cuda_attempts( # the driver major is the real constraint. Mirrors the legacy # windows_cuda_attempts fallback; without it a torch-only host gets no # fork attempt and silently drops to the upstream build. - ordered_lines = [line for line in compatible if line in detected] or list(compatible) + ordered_lines = [line for line in compatible if line in detected] or list( + compatible + ) if preferred_runtime_line and preferred_runtime_line in ordered_lines: ordered_lines = [preferred_runtime_line] + [ line for line in ordered_lines if line != preferred_runtime_line @@ -3594,7 +3758,11 @@ def published_windows_cuda_attempts( and artifact.min_sm is not None and artifact.max_sm is not None ) - if host_sms and has_sm_info and not _artifact_covers_sms(artifact, host_sms): + if ( + host_sms + and has_sm_info + and not _artifact_covers_sms(artifact, host_sms) + ): continue if not host_sms and has_sm_info and artifact.coverage_class != "portable": continue @@ -3690,7 +3858,11 @@ def published_asset_choice_for_kind( release: PublishedReleaseBundle, install_kind: str ) -> AssetChoice | None: candidates = sorted( - (artifact for artifact in release.artifacts if artifact.install_kind == install_kind), + ( + artifact + for artifact in release.artifacts + if artifact.install_kind == install_kind + ), key = lambda artifact: (artifact.rank, artifact.asset_name), ) for artifact in candidates: @@ -3706,7 +3878,9 @@ def published_asset_choice_for_kind( install_kind = install_kind, runtime_line = artifact.runtime_line, selection_log = list(release.selection_log) - + [f"published_selection: selected {artifact.asset_name} install_kind={install_kind}"], + + [ + f"published_selection: selected {artifact.asset_name} install_kind={install_kind}" + ], ) return None @@ -3759,7 +3933,11 @@ def _detect_host_rocm_version() -> tuple[int, int] | None: if result.returncode == 0: raw = (result.stdout or "").strip().split("\n")[0] parts = raw.split(".") - if len(parts) >= 2 and parts[0].isdigit() and parts[1].split("-")[0].isdigit(): + if ( + len(parts) >= 2 + and parts[0].isdigit() + and parts[1].split("-")[0].isdigit() + ): return int(parts[0]), int(parts[1].split("-")[0]) except Exception: pass @@ -3873,7 +4051,9 @@ def resolve_upstream_asset_choice(host: HostInfo, llama_tag: str) -> AssetChoice _compatible: list[tuple[tuple[int, ...], str]] = rocm_candidates if _host_rocm_version is not None: _compatible = [ - item for item in rocm_candidates if item[0][:2] <= _host_rocm_version + item + for item in rocm_candidates + if item[0][:2] <= _host_rocm_version ] if rocm_candidates and not _compatible: # Fall back to the newest candidate so a source build is @@ -3931,7 +4111,9 @@ def resolve_upstream_asset_choice(host: HostInfo, llama_tag: str) -> AssetChoice if host.has_rocm: hip_name = f"llama-{llama_tag}-bin-win-hip-radeon-x64.zip" if hip_name in upstream_assets: - log(f"AMD ROCm detected on Windows -- trying upstream HIP prebuilt {hip_name}") + log( + f"AMD ROCm detected on Windows -- trying upstream HIP prebuilt {hip_name}" + ) return AssetChoice( repo = UPSTREAM_REPO, tag = llama_tag, @@ -3940,7 +4122,9 @@ def resolve_upstream_asset_choice(host: HostInfo, llama_tag: str) -> AssetChoice source_label = "upstream", install_kind = "windows-hip", ) - log("AMD ROCm detected on Windows but no HIP prebuilt found -- falling back to CPU") + log( + "AMD ROCm detected on Windows but no HIP prebuilt found -- falling back to CPU" + ) upstream_name = f"llama-{llama_tag}-bin-win-cpu-x64.zip" if upstream_name not in upstream_assets: @@ -3980,7 +4164,9 @@ def resolve_upstream_asset_choice(host: HostInfo, llama_tag: str) -> AssetChoice install_kind = "macos-x64", ) - raise PrebuiltFallback(f"no prebuilt policy exists for {host.system} {host.machine}") + raise PrebuiltFallback( + f"no prebuilt policy exists for {host.system} {host.machine}" + ) def resolve_asset_choice(host: HostInfo, llama_tag: str) -> AssetChoice: @@ -4034,7 +4220,9 @@ def resolve_release_asset_choice( # prebuilt. We still avoid hard-pinning windows-cpu here so a CPU bundle # never shadows that ROCm path. if host.has_rocm: - published_choice = published_rocm_choice_for_host(release, host, "windows-rocm") + published_choice = published_rocm_choice_for_host( + release, host, "windows-rocm" + ) else: published_choice = published_asset_choice_for_kind(release, "windows-cpu") elif host.is_macos and host.is_arm64: @@ -4062,14 +4250,18 @@ def extract_archive(archive_path: Path, destination: Path) -> None: normalized = member_name.replace("\\", "/") member_path = Path(normalized) if member_path.is_absolute(): - raise PrebuiltFallback(f"archive member used an absolute path: {member_name}") + raise PrebuiltFallback( + f"archive member used an absolute path: {member_name}" + ) target = (base / member_path).resolve() base_resolved = base.resolve() try: target.relative_to(base_resolved) except ValueError as exc: - raise PrebuiltFallback(f"archive member escaped destination: {member_name}") from exc + raise PrebuiltFallback( + f"archive member escaped destination: {member_name}" + ) from exc return target def _try_repair_missing_slash( @@ -4115,7 +4307,11 @@ def extract_archive(archive_path: Path, destination: Path) -> None: return candidates[0][len(prefix) :] def safe_link_target( - base: Path, member_name: str, link_name: str, target: Path, archive_names: set[str] + base: Path, + member_name: str, + link_name: str, + target: Path, + archive_names: set[str], ) -> tuple[str, Path]: normalized = link_name.replace("\\", "/") repaired = _try_repair_missing_slash(member_name, normalized, archive_names) @@ -4175,7 +4371,9 @@ def extract_archive(archive_path: Path, destination: Path) -> None: target.parent.mkdir(parents = True, exist_ok = True) extracted = archive.extractfile(member) if extracted is None: - raise PrebuiltFallback(f"tar archive entry could not be read: {member.name}") + raise PrebuiltFallback( + f"tar archive entry could not be read: {member.name}" + ) with extracted, target.open("wb") as dst: shutil.copyfileobj(extracted, dst) @@ -4209,7 +4407,9 @@ def extract_archive(archive_path: Path, destination: Path) -> None: details = ", ".join( f"{member.name} -> {member.linkname}" for member, _ in next_round ) - raise PrebuiltFallback(f"tar archive contained unresolved link entries: {details}") + raise PrebuiltFallback( + f"tar archive contained unresolved link entries: {details}" + ) unresolved = next_round destination.mkdir(parents = True, exist_ok = True) @@ -4392,7 +4592,9 @@ def hydrate_source_tree( for index, source_url in enumerate(source_urls): try: if index > 0: - log(f"retrying source tree download from fallback URL: {source_url}") + log( + f"retrying source tree download from fallback URL: {source_url}" + ) download_file_verified( source_url, archive_path, @@ -4417,11 +4619,14 @@ def hydrate_source_tree( source_root / "gguf-py", ] missing = [ - str(path.relative_to(source_root)) for path in required_paths if not path.exists() + str(path.relative_to(source_root)) + for path in required_paths + if not path.exists() ] if missing: raise PrebuiltFallback( - "upstream source archive was missing required repo files: " + ", ".join(missing) + "upstream source archive was missing required repo files: " + + ", ".join(missing) ) copy_directory_contents(source_root, install_dir) except PrebuiltFallback: @@ -4448,7 +4653,9 @@ def discover_installed_executable(install_dir: Path, executable_name: str) -> Pa direct = install_dir / executable_name if direct.exists() and direct.is_file(): return direct - candidate = next((path for path in install_dir.rglob(executable_name) if path.is_file()), None) + candidate = next( + (path for path in install_dir.rglob(executable_name) if path.is_file()), None + ) if candidate is None: raise PrebuiltFallback(f"{executable_name} was not installed") return candidate @@ -4478,7 +4685,9 @@ def create_exec_entrypoint(entrypoint: Path, target: Path) -> None: write_exec_wrapper(entrypoint, target) -def overlay_directory_for_choice(install_dir: Path, choice: AssetChoice, host: HostInfo) -> Path: +def overlay_directory_for_choice( + install_dir: Path, choice: AssetChoice, host: HostInfo +) -> Path: if host.is_windows or choice.install_kind.startswith("windows"): path = install_dir / "build" / "bin" / "Release" else: @@ -4514,7 +4723,12 @@ def runtime_patterns_for_choice(choice: AssetChoice) -> list[str]: "linux-rocm", "linux-arm64", }: - return ["llama-server", "llama-quantize", "llama-diffusion-gemma-visual-server", "lib*.so*"] + return [ + "llama-server", + "llama-quantize", + "llama-diffusion-gemma-visual-server", + "lib*.so*", + ] if choice.install_kind in {"macos-arm64", "macos-x64"}: return [ "llama-server", @@ -4535,7 +4749,9 @@ def runtime_patterns_for_choice(choice: AssetChoice) -> list[str]: "llama-diffusion-gemma-visual-server.exe", "*.dll", ] - raise PrebuiltFallback(f"unsupported install kind for runtime overlay: {choice.install_kind}") + raise PrebuiltFallback( + f"unsupported install kind for runtime overlay: {choice.install_kind}" + ) def runtime_subdirs_for_choice(choice: AssetChoice) -> list[str]: @@ -4741,7 +4957,9 @@ def confirm_install_tree(install_dir: Path, host: HostInfo) -> None: expected.append(install_dir / "UNSLOTH_PREBUILT_INFO.json") missing = [str(path) for path in expected if not path.exists()] if missing: - raise RuntimeError("activated install was missing expected files: " + ", ".join(missing)) + raise RuntimeError( + "activated install was missing expected files: " + ", ".join(missing) + ) def activate_staged_dir(staging_dir: Path, dst: Path) -> None: @@ -4768,7 +4986,9 @@ def activate_staged_dir(staging_dir: Path, dst: Path) -> None: except OSError as exc: if not is_busy_lock_error(exc): raise - log(f"os.replace failed ({exc!r}); falling back to file-by-file copy of staging tree") + log( + f"os.replace failed ({exc!r}); falling back to file-by-file copy of staging tree" + ) shutil.copytree(staging_dir, dst, dirs_exist_ok = True) remove_tree(staging_dir) @@ -4893,11 +5113,15 @@ def install_from_archives( expected_sha256 = choice.runtime_sha256, label = f"prebuilt runtime archive {choice.runtime_name}", ) - runtime_extract_dir = Path(tempfile.mkdtemp(prefix = "extract-runtime-", dir = work_dir)) + runtime_extract_dir = Path( + tempfile.mkdtemp(prefix = "extract-runtime-", dir = work_dir) + ) extract_archive(runtime_archive, runtime_extract_dir) source_dir = extract_dir overlay_dir = overlay_directory_for_choice(install_dir, choice, host) - copy_globs(source_dir, overlay_dir, runtime_patterns_for_choice(choice), required = True) + copy_globs( + source_dir, overlay_dir, runtime_patterns_for_choice(choice), required = True + ) for _subdir in runtime_subdirs_for_choice(choice): _src_subdir = source_dir / _subdir if _src_subdir.is_dir(): @@ -4942,7 +5166,9 @@ def install_from_archives( source_server = build_bin / "llama-server" source_quantize = build_bin / "llama-quantize" if not source_server.exists() or not source_quantize.exists(): - raise PrebuiltFallback("unix executables were not installed correctly into build/bin") + raise PrebuiltFallback( + "unix executables were not installed correctly into build/bin" + ) os.chmod(source_server, 0o755) os.chmod(source_quantize, 0o755) @@ -4968,9 +5194,13 @@ def ensure_repo_shape(install_dir: Path) -> None: install_dir / "convert_hf_to_gguf.py", install_dir / "gguf-py", ] - missing = [str(path.relative_to(install_dir)) for path in required if not path.exists()] + missing = [ + str(path.relative_to(install_dir)) for path in required if not path.exists() + ] if missing: - raise PrebuiltFallback("hydrated llama.cpp source tree was missing: " + ", ".join(missing)) + raise PrebuiltFallback( + "hydrated llama.cpp source tree was missing: " + ", ".join(missing) + ) def validation_model_cache_path(install_dir: Path) -> Path: @@ -5015,7 +5245,9 @@ def _fetch_validation_model_bytes() -> bytes: repo_id, revision, filename = parts try: from huggingface_hub import hf_hub_download - local = hf_hub_download(repo_id = repo_id, filename = filename, revision = revision) + local = hf_hub_download( + repo_id = repo_id, filename = filename, revision = revision + ) return validated_validation_model_bytes(Path(local).read_bytes()) except Exception as exc: log( @@ -5038,7 +5270,9 @@ def download_validation_model(path: Path, cache_path: Path | None = None) -> Non data = validated_validation_model_bytes(cache_path.read_bytes()) log(f"using cached tiny GGUF validation model from {cache_path}") except Exception as exc: - log(f"cached tiny GGUF validation model was invalid; refreshing cache ({exc})") + log( + f"cached tiny GGUF validation model was invalid; refreshing cache ({exc})" + ) data = None if data is None: log("downloading tiny GGUF validation model") @@ -5122,7 +5356,9 @@ def dedupe_existing_dirs(paths: Iterable[str | Path]) -> list[str]: return unique -def linux_missing_libraries(binary_path: Path, *, env: dict[str, str] | None = None) -> list[str]: +def linux_missing_libraries( + binary_path: Path, *, env: dict[str, str] | None = None +) -> list[str]: try: result = run_capture(["ldd", str(binary_path)], timeout = 20, env = env) except Exception: @@ -5279,7 +5515,9 @@ def _macho_slice_minos(data: bytes, offset: int) -> tuple[int, int] | None: return None -def macho_minimum_macos(path: Path, host: HostInfo | None = None) -> tuple[int, int] | None: +def macho_minimum_macos( + path: Path, host: HostInfo | None = None +) -> tuple[int, int] | None: """Minimum macOS (major, minor) a Mach-O binary or dylib requires. Pure-Python so it works on consumer Macs without the Xcode command line @@ -5318,7 +5556,9 @@ def macho_minimum_macos(path: Path, host: HostInfo | None = None) -> tuple[int, return None if host is not None: want = ( - _CPU_TYPE_ARM64 if host.is_arm64 else (_CPU_TYPE_X86_64 if host.is_x86_64 else None) + _CPU_TYPE_ARM64 + if host.is_arm64 + else (_CPU_TYPE_X86_64 if host.is_x86_64 else None) ) for cputype, minos in slices: if cputype == want: @@ -5379,7 +5619,8 @@ def preflight_macos_installed_binaries( issues = macos_binary_minos_issues(binaries, install_dir, host) if issues: raise PrebuiltFallback( - "macos prebuilt requires a newer macOS than this host:\n" + "\n".join(issues) + "macos prebuilt requires a newer macOS than this host:\n" + + "\n".join(issues) ) @@ -5395,14 +5636,18 @@ def preflight_linux_installed_binaries( missing = linux_missing_libraries(binary_path, env = env) if not missing: continue - runtime_dirs = [part for part in env.get("LD_LIBRARY_PATH", "").split(os.pathsep) if part] + runtime_dirs = [ + part for part in env.get("LD_LIBRARY_PATH", "").split(os.pathsep) if part + ] issues.append( f"{binary_path.name}: missing={','.join(missing)} " f"ld_library_path={','.join(runtime_dirs) if runtime_dirs else 'none'}" ) if issues: - raise PrebuiltFallback("linux extracted binary preflight failed:\n" + "\n".join(issues)) + raise PrebuiltFallback( + "linux extracted binary preflight failed:\n" + "\n".join(issues) + ) def glob_paths(*patterns: str) -> list[str]: @@ -5448,7 +5693,9 @@ def windows_runtime_dirs() -> list[str]: def windows_runtime_dirs_for_patterns( required_patterns: Iterable[str], candidate_dirs: Iterable[str] | None = None ) -> list[str]: - directories = list(candidate_dirs) if candidate_dirs is not None else windows_runtime_dirs() + directories = ( + list(candidate_dirs) if candidate_dirs is not None else windows_runtime_dirs() + ) matching_dirs: list[str] = [] for pattern in required_patterns: matched_dirs = [ @@ -5524,12 +5771,20 @@ def binary_env( if _wsl_rocm: ld_dirs = [*_wsl_rocm, *ld_dirs] env.setdefault("HSA_ENABLE_DXG_DETECTION", "1") - existing = [part for part in env.get("LD_LIBRARY_PATH", "").split(os.pathsep) if part] - env["LD_LIBRARY_PATH"] = os.pathsep.join(dedupe_existing_dirs([*ld_dirs, *existing])) + existing = [ + part for part in env.get("LD_LIBRARY_PATH", "").split(os.pathsep) if part + ] + env["LD_LIBRARY_PATH"] = os.pathsep.join( + dedupe_existing_dirs([*ld_dirs, *existing]) + ) elif host.is_macos: dyld_dirs = [str(binary_path.parent), str(install_dir)] - existing = [part for part in env.get("DYLD_LIBRARY_PATH", "").split(os.pathsep) if part] - env["DYLD_LIBRARY_PATH"] = os.pathsep.join(dedupe_existing_dirs([*dyld_dirs, *existing])) + existing = [ + part for part in env.get("DYLD_LIBRARY_PATH", "").split(os.pathsep) if part + ] + env["DYLD_LIBRARY_PATH"] = os.pathsep.join( + dedupe_existing_dirs([*dyld_dirs, *existing]) + ) return env @@ -5551,7 +5806,11 @@ def validate_quantize( env = binary_env(quantize_path, install_dir, host, runtime_line = runtime_line), **windows_hidden_subprocess_kwargs(), ) - if result.returncode != 0 or not quantized_path.exists() or quantized_path.stat().st_size == 0: + if ( + result.returncode != 0 + or not quantized_path.exists() + or quantized_path.stat().st_size == 0 + ): combined = result.stdout + ("\n" + result.stderr if result.stderr else "") # Backstop for prebuilts the static minos scan could not read: a dyld # "built for macOS N" / missing Metal symbol failure means this binary @@ -5561,7 +5820,9 @@ def validate_quantize( if looks_like_macos_incompatibility(combined) else "" ) - raise PrebuiltFallback(prefix + "llama-quantize validation failed:\n" + combined) + raise PrebuiltFallback( + prefix + "llama-quantize validation failed:\n" + combined + ) def validate_server( @@ -5619,7 +5880,9 @@ def validate_server( # is exercised against the actual hardware rather than the # CPU fallback. NVIDIA and macOS-arm64 are already covered. _enable_gpu_layers = ( - host.has_usable_nvidia or host.has_rocm or (host.is_macos and host.is_arm64) + host.has_usable_nvidia + or host.has_rocm + or (host.is_macos and host.is_arm64) ) if _enable_gpu_layers: command.extend(["--n-gpu-layers", "1"]) @@ -5635,7 +5898,9 @@ def validate_server( stdout = log_handle, stderr = subprocess.STDOUT, text = True, - env = binary_env(server_path, install_dir, host, runtime_line = runtime_line), + env = binary_env( + server_path, install_dir, host, runtime_line = runtime_line + ), **windows_hidden_subprocess_kwargs(), ) deadline = time.time() + 60 @@ -5650,7 +5915,9 @@ def validate_server( exited_quickly = ( time.time() - startup_started ) <= SERVER_BIND_RETRY_WINDOW_SECONDS - failure = PrebuiltFallback("llama-server exited during startup:\n" + output) + failure = PrebuiltFallback( + "llama-server exited during startup:\n" + output + ) if ( port_attempt < SERVER_PORT_BIND_ATTEMPTS and is_retryable_server_bind_error( @@ -5667,7 +5934,9 @@ def validate_server( break raise failure - payload = json.dumps({"prompt": "a", "n_predict": 1}).encode("utf-8") + payload = json.dumps({"prompt": "a", "n_predict": 1}).encode( + "utf-8" + ) request = urllib.request.Request( f"http://127.0.0.1:{port}/completion", data = payload, @@ -5679,7 +5948,9 @@ def validate_server( response_body = response.read().decode("utf-8", "replace") if status_code == 200: return - last_error = RuntimeError(f"unexpected HTTP status {status_code}") + last_error = RuntimeError( + f"unexpected HTTP status {status_code}" + ) except urllib.error.HTTPError as exc: response_body = exc.read().decode("utf-8", "replace") last_error = exc @@ -5713,7 +5984,9 @@ def validate_server( raise PrebuiltFallback("llama-server validation failed unexpectedly") -def collect_system_report(host: HostInfo, choice: AssetChoice | None, install_dir: Path) -> str: +def collect_system_report( + host: HostInfo, choice: AssetChoice | None, install_dir: Path +) -> str: lines = [ f"platform={host.system} machine={host.machine}", f"driver_cuda_version={host.driver_cuda_version}", @@ -5727,7 +6000,8 @@ def collect_system_report(host: HostInfo, choice: AssetChoice | None, install_di if host.is_linux and host.has_physical_nvidia: runtime_lines, runtime_dirs = detected_linux_runtime_lines() lines.append( - "linux_runtime_lines=" + (",".join(runtime_lines) if runtime_lines else "none") + "linux_runtime_lines=" + + (",".join(runtime_lines) if runtime_lines else "none") ) for runtime_line in ("cuda13", "cuda12"): lines.append( @@ -5756,7 +6030,10 @@ def collect_system_report(host: HostInfo, choice: AssetChoice | None, install_di server_env = binary_env(server_binary, install_dir, host) lines.append( "linux_missing_libs=" - + (",".join(linux_missing_libraries(server_binary, env = server_env)) or "none") + + ( + ",".join(linux_missing_libraries(server_binary, env = server_env)) + or "none" + ) ) lines.append( "linux_runtime_dirs=" @@ -5764,7 +6041,9 @@ def collect_system_report(host: HostInfo, choice: AssetChoice | None, install_di ",".join( [ part - for part in server_env.get("LD_LIBRARY_PATH", "").split(os.pathsep) + for part in server_env.get("LD_LIBRARY_PATH", "").split( + os.pathsep + ) if part ] ) @@ -5772,16 +6051,21 @@ def collect_system_report(host: HostInfo, choice: AssetChoice | None, install_di ) ) try: - ldd = run_capture(["ldd", str(server_binary)], timeout = 20, env = server_env) + ldd = run_capture( + ["ldd", str(server_binary)], timeout = 20, env = server_env + ) lines.append("ldd llama-server:") lines.append((ldd.stdout + ldd.stderr).strip()) except Exception as exc: lines.append(f"ldd error: {exc}") elif host.is_windows: - lines.append("windows_runtime_dirs=" + (",".join(windows_runtime_dirs()) or "none")) + lines.append( + "windows_runtime_dirs=" + (",".join(windows_runtime_dirs()) or "none") + ) runtime_lines, runtime_dirs = detected_windows_runtime_lines() lines.append( - "windows_runtime_lines=" + (",".join(runtime_lines) if runtime_lines else "none") + "windows_runtime_lines=" + + (",".join(runtime_lines) if runtime_lines else "none") ) for runtime_line in ("cuda13", "cuda12"): lines.append( @@ -5930,7 +6214,9 @@ def resolve_install_attempts( return requested_tag, plan.llama_tag, plan.attempts, plan.approved_checksums -def _linux_published_attempts(host: HostInfo, bundle: PublishedReleaseBundle) -> list[AssetChoice]: +def _linux_published_attempts( + host: HostInfo, bundle: PublishedReleaseBundle +) -> list[AssetChoice]: """Build the install attempts for a fork Linux host from a manifest-described bundle: CUDA (with a CPU fallback), per-gfx ROCm, or CPU. Same selection the upstream filename path used, just sourced from the manifest instead of @@ -5978,11 +6264,17 @@ def _fork_manifest_release_plans( llama-prebuilt-manifest.json rather than in the filename: arm64 CUDA, Windows CUDA, per-gfx ROCm, and macOS. Linux x64 takes the faster filename path.""" requested_tag = normalized_requested_llama_tag(llama_tag) - allow_older_release_fallback = requested_tag == "latest" and not published_release_tag + allow_older_release_fallback = ( + requested_tag == "latest" and not published_release_tag + ) release_limit = max(1, max_release_fallbacks) # macOS may need to walk past a run of too-new prebuilts. Only when the host # version is known; otherwise keep the default (cannot tell up front). - if host.is_macos and allow_older_release_fallback and host.macos_version is not None: + if ( + host.is_macos + and allow_older_release_fallback + and host.macos_version is not None + ): release_limit = max(release_limit, DEFAULT_MAX_MACOS_RELEASE_FALLBACKS) plans: list[InstallReleasePlan] = [] last_error: PrebuiltFallback | None = None @@ -5999,10 +6291,14 @@ def _fork_manifest_release_plans( if host.is_linux: linux_attempts = _linux_published_attempts(host, bundle) if not linux_attempts: - raise PrebuiltFallback("no compatible Linux prebuilt asset was found") + raise PrebuiltFallback( + "no compatible Linux prebuilt asset was found" + ) attempts = apply_approved_hashes(linux_attempts, checksums) if not attempts: - raise PrebuiltFallback("no compatible Linux prebuilt asset was found") + raise PrebuiltFallback( + "no compatible Linux prebuilt asset was found" + ) if attempts[0].selection_log: log_lines(attempts[0].selection_log) else: @@ -6103,7 +6399,9 @@ def write_prebuilt_metadata( "prebuilt_fallback_used": prebuilt_fallback_used, "installed_at_utc": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()), } - (install_dir / "UNSLOTH_PREBUILT_INFO.json").write_text(json.dumps(metadata, indent = 2) + "\n") + (install_dir / "UNSLOTH_PREBUILT_INFO.json").write_text( + json.dumps(metadata, indent = 2) + "\n" + ) def expected_install_fingerprint( @@ -6220,7 +6518,9 @@ def install_runtime_dir(install_dir: Path, host: HostInfo) -> Path: return install_dir / "build" / "bin" -def runtime_payload_is_healthy(install_dir: Path, host: HostInfo, choice: AssetChoice) -> bool: +def runtime_payload_is_healthy( + install_dir: Path, host: HostInfo, choice: AssetChoice +) -> bool: runtime_dir = install_runtime_dir(install_dir, host) if not runtime_dir.exists(): return False @@ -6343,13 +6643,17 @@ def validate_prebuilt_choice( # not in any repo, so fetch the asset directly; codeload stays the fallback. asset_url = ( release_asset_download_url( - approved_checksums.repo, approved_checksums.release_tag, source_archive.asset_name + approved_checksums.repo, + approved_checksums.release_tag, + source_archive.asset_name, ) if exact_source and source_archive is not None else None ) if exact_source: - log(f"hydrating exact llama.cpp source for {source_repo}@{source_ref} into {install_dir}") + log( + f"hydrating exact llama.cpp source for {source_repo}@{source_ref} into {install_dir}" + ) else: log(f"hydrating upstream llama.cpp source for {llama_tag} into {install_dir}") hydrate_source_tree( @@ -6367,7 +6671,9 @@ def validate_prebuilt_choice( asset_url = asset_url, ) log(f"overlaying prebuilt bundle {choice.name} into {install_dir}") - server_path, quantize_path = install_from_archives(choice, host, install_dir, work_dir) + server_path, quantize_path = install_from_archives( + choice, host, install_dir, work_dir + ) preflight_linux_installed_binaries((server_path, quantize_path), install_dir, host) preflight_macos_installed_binaries((server_path, quantize_path), install_dir, host) ensure_repo_shape(install_dir) @@ -6551,9 +6857,13 @@ def install_prebuilt( published_repo, published_release_tag, ) - if release_plans and existing_install_matches_plan(install_dir, host, release_plans[0]): + if release_plans and existing_install_matches_plan( + install_dir, host, release_plans[0] + ): current = release_plans[0] - if diffusion_visual_server_backfill_needed(install_dir, host, current.attempts[0]): + if diffusion_visual_server_backfill_needed( + install_dir, host, current.attempts[0] + ): log( f"existing install matches {current.release_tag} but is missing the " "DiffusionGemma visual-server; re-extracting the bundle to backfill it" @@ -6567,11 +6877,15 @@ def install_prebuilt( with tempfile.TemporaryDirectory(prefix = "unsloth-llama-prebuilt-") as tmp: work_dir = Path(tmp) probe_path = work_dir / "stories260K.gguf" - download_validation_model(probe_path, validation_model_cache_path(install_dir)) + download_validation_model( + probe_path, validation_model_cache_path(install_dir) + ) release_count = len(release_plans) for release_index, plan in enumerate(release_plans): choice = plan.attempts[0] - backfill = diffusion_visual_server_backfill_needed(install_dir, host, choice) + backfill = diffusion_visual_server_backfill_needed( + install_dir, host, choice + ) if existing_install_matches_plan(install_dir, host, plan): if backfill: log( @@ -6627,7 +6941,9 @@ def install_prebuilt( f"({textwrap.shorten(str(exc), width = 200, placeholder = '...')})" ) try: - ensure_diffusion_visual_server(install_dir, host, plan.release_tag) + ensure_diffusion_visual_server( + install_dir, host, plan.release_tag + ) except Exception as exc: log( "diffusion visual server step skipped; install remains valid " @@ -6794,7 +7110,9 @@ def main() -> int: ) emit_resolver_output( { - "requested_tag": normalized_requested_llama_tag(args.resolve_install_tag), + "requested_tag": normalized_requested_llama_tag( + args.resolve_install_tag + ), "llama_tag": resolved, }, output_format = args.output_format, @@ -6809,7 +7127,9 @@ def main() -> int: ) emit_resolver_output( { - "requested_tag": normalized_requested_llama_tag(args.resolve_source_build), + "requested_tag": normalized_requested_llama_tag( + args.resolve_source_build + ), "source_url": plan.source_url, "source_ref_kind": plan.source_ref_kind, "source_ref": plan.source_ref, diff --git a/studio/install_python_stack.py b/studio/install_python_stack.py index d404041ca7..3bbc80492f 100644 --- a/studio/install_python_stack.py +++ b/studio/install_python_stack.py @@ -243,7 +243,11 @@ def _detect_rocm_version() -> tuple[int, int] | None: if result.returncode == 0: raw = result.stdout.decode().strip().split("\n")[0] parts = raw.split(".") - if len(parts) >= 2 and parts[0].isdigit() and parts[1].split("-")[0].isdigit(): + if ( + len(parts) >= 2 + and parts[0].isdigit() + and parts[1].split("-")[0].isdigit() + ): return int(parts[0]), int(parts[1].split("-")[0]) except Exception: pass @@ -359,7 +363,8 @@ def _detect_windows_gfx_arch() -> str | None: # findall gets every gcnArchName line so multi-GPU hosts are # enumerable and HIP_VISIBLE_DEVICES selects correctly. _tokens = [ - t.strip().lower() for t in re.findall(r"(?im)^\s*gcnArchName\s*:\s*(\S+)", text) + t.strip().lower() + for t in re.findall(r"(?im)^\s*gcnArchName\s*:\s*(\S+)", text) ] _pick = _dedup_pick(_tokens) if _pick: @@ -536,7 +541,9 @@ def _persist_bnb_rocm_version(version: str) -> bool: try: sitecustomize_path.parent.mkdir(parents = True, exist_ok = True) existing = ( - sitecustomize_path.read_text(encoding = "utf-8") if sitecustomize_path.exists() else "" + sitecustomize_path.read_text(encoding = "utf-8") + if sitecustomize_path.exists() + else "" ) # Strip all managed regions, including one whose END marker was lost to # an interrupted write, then append exactly one fresh block. @@ -779,7 +786,8 @@ def _install_bnb_windows_rocm() -> bool: # Fall back to "72" if detection fails (e.g. install was a no-op / dry-run). _env_ver = os.environ.get("BNB_ROCM_VERSION") _env_is_persisted_default = ( - os.environ.get(_BNB_ROCM_VERSION_SOURCE_ENV) == _BNB_ROCM_VERSION_SOURCE_SITECUSTOMIZE + os.environ.get(_BNB_ROCM_VERSION_SOURCE_ENV) + == _BNB_ROCM_VERSION_SOURCE_SITECUSTOMIZE ) _persist_detected_version = False if _env_ver and not _env_is_persisted_default: @@ -913,7 +921,9 @@ def _ensure_cuda_torch() -> None: # Take the last non-empty stdout line: stray output from sitecustomize or # an import hook must not mask the marker (fail-closed either way). _marker_lines = [ - line.strip() for line in probe.stdout.decode(errors = "replace").splitlines() if line.strip() + line.strip() + for line in probe.stdout.decode(errors = "replace").splitlines() + if line.strip() ] if not _marker_lines or _marker_lines[-1] != "hip": return # healthy CUDA torch, or a deliberate CPU wheel -- leave as-is @@ -1022,7 +1032,9 @@ def _ensure_rocm_torch() -> None: if not _torch_already_rocm: index_url = _windows_rocm_index_url(gfx_arch) if index_url is None: - print(f" No AMD Windows torch index for GPU arch {gfx_arch} -- skipping") + print( + f" No AMD Windows torch index for GPU arch {gfx_arch} -- skipping" + ) return print(f" {gfx_arch} (Windows) -- installing torch from {index_url}") pip_install( @@ -1095,7 +1107,9 @@ def _ensure_rocm_torch() -> None: except (OSError, subprocess.TimeoutExpired): probe = None has_hip_torch = ( - probe is not None and probe.returncode == 0 and probe.stdout.decode().strip() != "" + probe is not None + and probe.returncode == 0 + and probe.stdout.decode().strip() != "" ) rocm_torch_ready = has_hip_torch @@ -1117,11 +1131,14 @@ def _ensure_rocm_torch() -> None: # Pick the runtime-visible GPU: use the HIP_VISIBLE_DEVICES index # into gfx_codes, else default to the first GPU. Skip the override # unless the resolved GPU is Strix. - _runtime_gfx = gfx_codes[_pick_visible_index(len(gfx_codes))] if gfx_codes else None + _runtime_gfx = ( + gfx_codes[_pick_visible_index(len(gfx_codes))] if gfx_codes else None + ) if _runtime_gfx in _strix_gfx: _selected_gfx = _runtime_gfx _amd_mirror = ( - os.environ.get("UNSLOTH_AMD_ROCM_MIRROR") or "https://repo.amd.com/rocm/whl" + os.environ.get("UNSLOTH_AMD_ROCM_MIRROR") + or "https://repo.amd.com/rocm/whl" ).rstrip("/") _strix_override_url = f"{_amd_mirror}/{_selected_gfx}/" _strix_override_pkgs = ( @@ -1182,7 +1199,10 @@ def _ensure_rocm_torch() -> None: None, ) if tag is None: - print(f" No PyTorch wheel for ROCm {ver[0]}.{ver[1]} -- " f"skipping torch reinstall") + print( + f" No PyTorch wheel for ROCm {ver[0]}.{ver[1]} -- " + f"skipping torch reinstall" + ) else: index_url = f"{_PYTORCH_WHL_BASE}/{tag}" print(f" ROCm {ver[0]}.{ver[1]} -- installing torch from {index_url}") @@ -1340,7 +1360,9 @@ CONSTRAINTS = SINGLE_ENV / "constraints.txt" LOCAL_DD_UNSTRUCTURED_PLUGIN = ( SCRIPT_DIR / "backend" / "plugins" / "data-designer-unstructured-seed" ) -LOCAL_DD_GITHUB_PLUGIN = SCRIPT_DIR / "backend" / "plugins" / "data-designer-github-repo-seed" +LOCAL_DD_GITHUB_PLUGIN = ( + SCRIPT_DIR / "backend" / "plugins" / "data-designer-github-repo-seed" +) # Apple Silicon: override mlx-vlm/mlx-lm's transformers pin (see overrides). _MLX_OVERRIDES = SINGLE_ENV / "overrides-darwin-arm64.txt" @@ -1465,7 +1487,9 @@ def _progress(label: str) -> None: pad = " " * (_COL - len(_LABEL)) end = "\n" if _STEP >= _TOTAL else "" try: - sys.stdout.write(f"\r {_dim(_LABEL)}{pad}[{bar}] {_STEP:2}/{_TOTAL} {label:<20}{end}") + sys.stdout.write( + f"\r {_dim(_LABEL)}{pad}[{bar}] {_STEP:2}/{_TOTAL} {label:<20}{end}" + ) sys.stdout.flush() except OSError: pass @@ -1525,7 +1549,9 @@ def _build_flash_attn_wheel_url(env: dict[str, str]) -> str | None: return flash_attn_wheel_url(env) -def _print_optional_install_failure(label: str, result: subprocess.CompletedProcess[str]) -> None: +def _print_optional_install_failure( + label: str, result: subprocess.CompletedProcess[str] +) -> None: _step("warning", f"{label} failed (exit code {result.returncode})", _cyan) if result.stdout: print(result.stdout.strip()) @@ -1620,7 +1646,9 @@ def _filter_requirements(req: Path, skip: set[str]) -> Path: """Return a temp copy of a requirements file with certain packages removed.""" lines = req.read_text(encoding = "utf-8").splitlines(keepends = True) filtered = [ - line for line in lines if not any(line.strip().lower().startswith(pkg) for pkg in skip) + line + for line in lines + if not any(line.strip().lower().startswith(pkg) for pkg in skip) ] tmp = tempfile.NamedTemporaryFile( mode = "w", @@ -1830,7 +1858,9 @@ def install_python_stack() -> int: if not IS_MACOS and not NO_TORCH: base_total += 1 # ROCm torch check (line 1526) -- all non-macOS platforms if not IS_WINDOWS: - base_total += 2 # flash-attn (line 1620) + ROCm torch final (line 1705) -- Linux only + base_total += ( + 2 # flash-attn (line 1620) + ROCm torch final (line 1705) -- Linux only + ) _TOTAL = (base_total - 1) if skip_base else base_total # 1. Try uv for faster installs (before pip upgrade -- uv venvs don't diff --git a/tests/python/conftest.py b/tests/python/conftest.py index f7b125edf6..9129e384e5 100644 --- a/tests/python/conftest.py +++ b/tests/python/conftest.py @@ -2,5 +2,9 @@ def pytest_configure(config): - config.addinivalue_line("markers", "server: heavyweight tests requiring studio venv") - config.addinivalue_line("markers", "e2e: end-to-end tests requiring network and venv creation") + config.addinivalue_line( + "markers", "server: heavyweight tests requiring studio venv" + ) + config.addinivalue_line( + "markers", "e2e: end-to-end tests requiring network and venv creation" + ) diff --git a/tests/python/test_cross_platform_parity.py b/tests/python/test_cross_platform_parity.py index 34f984714e..8f7f09477f 100644 --- a/tests/python/test_cross_platform_parity.py +++ b/tests/python/test_cross_platform_parity.py @@ -27,11 +27,17 @@ class TestNoTorchBackendAutoInInstallSh: for i, line in enumerate(lines): if fallback_start is None and "GPU detection failed" in line: fallback_start = i - elif fallback_start is not None and fallback_end is None and line.strip() == "fi": + elif ( + fallback_start is not None + and fallback_end is None + and line.strip() == "fi" + ): fallback_end = i break fallback_range = ( - range(fallback_start or 0, (fallback_end or 0) + 1) if fallback_start else range(0) + range(fallback_start or 0, (fallback_end or 0) + 1) + if fallback_start + else range(0) ) matches = [ diff --git a/tests/python/test_e2e_no_torch_sandbox.py b/tests/python/test_e2e_no_torch_sandbox.py index 9b3bda133f..b4ff73882d 100644 --- a/tests/python/test_e2e_no_torch_sandbox.py +++ b/tests/python/test_e2e_no_torch_sandbox.py @@ -235,8 +235,12 @@ class TestBeforeAfterImportChain: exec(source) """) result = _run_in_sandbox(no_torch_venv, code) - assert result.returncode != 0, "BEFORE chat_templates.py should crash without torch" - assert b"ModuleNotFoundError" in result.stderr or b"ImportError" in result.stderr + assert ( + result.returncode != 0 + ), "BEFORE chat_templates.py should crash without torch" + assert ( + b"ModuleNotFoundError" in result.stderr or b"ImportError" in result.stderr + ) def test_before_data_collators_crashes(self, no_torch_venv, sandbox_dir): """BEFORE: data_collators.py with top-level 'import torch' crashes.""" @@ -254,8 +258,12 @@ class TestBeforeAfterImportChain: exec(open({str(before_file)!r}).read()) """) result = _run_in_sandbox(no_torch_venv, code) - assert result.returncode != 0, "BEFORE data_collators.py should crash without torch" - assert b"ModuleNotFoundError" in result.stderr or b"ImportError" in result.stderr + assert ( + result.returncode != 0 + ), "BEFORE data_collators.py should crash without torch" + assert ( + b"ModuleNotFoundError" in result.stderr or b"ImportError" in result.stderr + ) def test_before_full_import_chain_crashes(self, no_torch_venv, sandbox_dir): """BEFORE: full utils/datasets/ package with top-level torch imports crashes.""" @@ -300,8 +308,12 @@ class TestBeforeAfterImportChain: from utils.datasets import detect_dataset_format """) result = _run_in_sandbox(no_torch_venv, code) - assert result.returncode != 0, "BEFORE full import chain should crash without torch" - assert b"ModuleNotFoundError" in result.stderr or b"ImportError" in result.stderr + assert ( + result.returncode != 0 + ), "BEFORE full import chain should crash without torch" + assert ( + b"ModuleNotFoundError" in result.stderr or b"ImportError" in result.stderr + ) # -- AFTER: succeeds -- @@ -515,7 +527,9 @@ class TestEdgeCasesBrokenTorch: print("OK: data_collators works despite broken torch on sys.path") """) result = _run_in_sandbox(no_torch_venv, code) - assert result.returncode == 0, f"Should work with broken torch:\n{result.stderr.decode()}" + assert ( + result.returncode == 0 + ), f"Should work with broken torch:\n{result.stderr.decode()}" assert b"OK:" in result.stdout def test_torch_import_error_hardware_fallback(self, no_torch_venv, sandbox_dir): @@ -578,10 +592,14 @@ class TestEdgeCasesBrokenTorch: print("OK: detect_hardware returned CPU with fake torch (no CUDA)") """) result = _run_in_sandbox(no_torch_venv, code) - assert result.returncode == 0, f"Should fall back to CPU:\n{result.stderr.decode()}" + assert ( + result.returncode == 0 + ), f"Should fall back to CPU:\n{result.stderr.decode()}" assert b"OK:" in result.stdout - def test_lazy_torch_fails_at_call_time_not_import_time(self, no_torch_venv, sandbox_dir): + def test_lazy_torch_fails_at_call_time_not_import_time( + self, no_torch_venv, sandbox_dir + ): """apply_chat_template_to_dataset is importable without torch. Calling the alpaca branch triggers the lazy 'from torch.utils.data' inside @@ -627,7 +645,9 @@ class TestEdgeCasesBrokenTorch: print("OK: call succeeded (unexpected but not a crash)") """) result = _run_in_sandbox(no_torch_venv, code) - assert result.returncode == 0, f"Should not crash at import time:\n{result.stderr.decode()}" + assert ( + result.returncode == 0 + ), f"Should not crash at import time:\n{result.stderr.decode()}" assert b"OK: import succeeded" in result.stdout @@ -979,7 +999,9 @@ class TestInstallPythonStackFiltering: source = Path(ips.__file__).read_text(encoding = "utf-8") # NO_TORCH guard before overrides - assert "if NO_TORCH:" in source, "NO_TORCH guard not found in install_python_stack.py" + assert ( + "if NO_TORCH:" in source + ), "NO_TORCH guard not found in install_python_stack.py" # macOS guard for triton assert ( @@ -1082,7 +1104,9 @@ class TestLiveServerStartup: for _ in range(30): time.sleep(1) try: - resp = urllib.request.urlopen(f"http://127.0.0.1:{port}/api/health", timeout = 2) + resp = urllib.request.urlopen( + f"http://127.0.0.1:{port}/api/health", timeout = 2 + ) if resp.status == 200: ready = True break @@ -1106,8 +1130,12 @@ class TestLiveServerStartup: capture_output = True, timeout = 300, ) - server_output = stdout.decode(errors = "replace") + stderr.decode(errors = "replace") - pytest.skip(f"Server failed to start within 30 seconds. Output:\n{server_output}") + server_output = stdout.decode(errors = "replace") + stderr.decode( + errors = "replace" + ) + pytest.skip( + f"Server failed to start within 30 seconds. Output:\n{server_output}" + ) yield proc, port @@ -1151,7 +1179,9 @@ class TestLiveServerStartup: import urllib.request _, port = server_process - resp = urllib.request.urlopen(f"http://127.0.0.1:{port}/openapi.json", timeout = 5) + resp = urllib.request.urlopen( + f"http://127.0.0.1:{port}/openapi.json", timeout = 5 + ) spec = json.loads(resp.read()) assert ( len(spec.get("paths", {})) >= 20 diff --git a/tests/python/test_fast_language_model_text_only.py b/tests/python/test_fast_language_model_text_only.py index ce4dd74439..9061240c05 100644 --- a/tests/python/test_fast_language_model_text_only.py +++ b/tests/python/test_fast_language_model_text_only.py @@ -126,7 +126,9 @@ def test_fast_language_model_forwards_text_only_to_fast_model(): # text_only defaults False (opt-in, not forced True), and both FastModel # delegations forward it. text_only_default = _param_default(method, "text_only") - assert isinstance(text_only_default, ast.Constant) and text_only_default.value is False + assert ( + isinstance(text_only_default, ast.Constant) and text_only_default.value is False + ) fast_model_calls = [ node @@ -150,13 +152,16 @@ def test_fast_model_text_only_does_not_override_explicit_auto_model(): method = _class_method(ast.parse(source), "FastModel", "from_pretrained") text_only_default = _param_default(method, "text_only") - assert isinstance(text_only_default, ast.Constant) and text_only_default.value is False + assert ( + isinstance(text_only_default, ast.Constant) and text_only_default.value is False + ) # load_text_only is text_only AND a check that the caller did not pass auto_model. def _is_guarded_bool(value): names = _names_in(value) has_none_check = any( - isinstance(n, ast.Compare) and any(isinstance(op, (ast.Is, ast.IsNot)) for op in n.ops) + isinstance(n, ast.Compare) + and any(isinstance(op, (ast.Is, ast.IsNot)) for op in n.ops) for n in ast.walk(value) ) return "text_only" in names and "auto_model" in names and has_none_check @@ -192,7 +197,9 @@ def test_fast_base_model_text_only_bypasses_vision_auto_model(): method = _class_method(ast.parse(source), "FastBaseModel", "from_pretrained") text_only_default = _param_default(method, "text_only") - assert isinstance(text_only_default, ast.Constant) and text_only_default.value is False + assert ( + isinstance(text_only_default, ast.Constant) and text_only_default.value is False + ) assert _assigns_name( method, @@ -327,7 +334,9 @@ def test_text_only_key_mapping_targets_published_prefixes(): # transformers >=5 (on 4.x base_model_prefix handles it and a mapping hurts). transformers = pytest.importorskip("transformers") get_key_mapping = _load_util_func("_get_text_only_key_mapping") - mapping = get_key_mapping(transformers.Gemma3Config(), transformers.Gemma3TextConfig()) + mapping = get_key_mapping( + transformers.Gemma3Config(), transformers.Gemma3TextConfig() + ) if int(transformers.__version__.split(".")[0]) < 5: assert mapping is None else: diff --git a/tests/python/test_fast_model_config_passthrough.py b/tests/python/test_fast_model_config_passthrough.py index b2ba3d2eef..35bdb6b619 100644 --- a/tests/python/test_fast_model_config_passthrough.py +++ b/tests/python/test_fast_model_config_passthrough.py @@ -29,7 +29,8 @@ def _assigns_from_kwargs_pop(method, target_name, key_name): if not isinstance(node, ast.Assign): continue if not any( - isinstance(target, ast.Name) and target.id == target_name for target in node.targets + isinstance(target, ast.Name) and target.id == target_name + for target in node.targets ): continue value = node.value @@ -50,7 +51,9 @@ def _assigns_from_kwargs_pop(method, target_name, key_name): def _calls_name(method, name): return any( - isinstance(node, ast.Call) and isinstance(node.func, ast.Name) and node.func.id == name + isinstance(node, ast.Call) + and isinstance(node.func, ast.Name) + and node.func.id == name for node in ast.walk(method) ) @@ -130,7 +133,9 @@ def test_fast_model_uses_user_config_num_labels_for_task_model_selection(): def test_fast_model_captures_user_config_num_labels_before_text_only_switch(): source = _source(LOADER_PATH) - fallback = source.index("task_config_attrs = _get_user_task_config_attrs(user_config)") + fallback = source.index( + "task_config_attrs = _get_user_task_config_attrs(user_config)" + ) text_only_switch = source.index("model_config = text_config") assert fallback < text_only_switch diff --git a/tests/python/test_fast_sentence_transformer_redirect_lifecycle.py b/tests/python/test_fast_sentence_transformer_redirect_lifecycle.py index 6fa7d06ec3..82d5278075 100644 --- a/tests/python/test_fast_sentence_transformer_redirect_lifecycle.py +++ b/tests/python/test_fast_sentence_transformer_redirect_lifecycle.py @@ -127,10 +127,18 @@ def _build_driver(transformer_class): return model if is_requested_model_name(a, kw) else original_model(*a, **kw) def return_existing_tokenizer(*a, **kw): - return tokenizer if is_requested_model_name(a, kw) else original_tokenizer(*a, **kw) + return ( + tokenizer + if is_requested_model_name(a, kw) + else original_tokenizer(*a, **kw) + ) def return_existing_processor(*a, **kw): - return tokenizer if is_requested_model_name(a, kw) else original_processor(*a, **kw) + return ( + tokenizer + if is_requested_model_name(a, kw) + else original_processor(*a, **kw) + ) try: AutoModel.from_pretrained = return_existing_model diff --git a/tests/python/test_flash_attn_install_python_stack.py b/tests/python/test_flash_attn_install_python_stack.py index 26ff03505a..a09b8e468f 100644 --- a/tests/python/test_flash_attn_install_python_stack.py +++ b/tests/python/test_flash_attn_install_python_stack.py @@ -33,42 +33,64 @@ class TestHasBlackwellGpu: def test_returns_true_for_sm_100(self): with ( - mock.patch.object(wheel_utils.shutil, "which", return_value = "/usr/bin/nvidia-smi"), - mock.patch.object(wheel_utils.subprocess, "run", return_value = _smi_result("10.0\n")), + mock.patch.object( + wheel_utils.shutil, "which", return_value = "/usr/bin/nvidia-smi" + ), + mock.patch.object( + wheel_utils.subprocess, "run", return_value = _smi_result("10.0\n") + ), ): assert wheel_utils.has_blackwell_gpu() is True def test_returns_true_for_sm_120(self): with ( - mock.patch.object(wheel_utils.shutil, "which", return_value = "/usr/bin/nvidia-smi"), - mock.patch.object(wheel_utils.subprocess, "run", return_value = _smi_result("12.0\n")), + mock.patch.object( + wheel_utils.shutil, "which", return_value = "/usr/bin/nvidia-smi" + ), + mock.patch.object( + wheel_utils.subprocess, "run", return_value = _smi_result("12.0\n") + ), ): assert wheel_utils.has_blackwell_gpu() is True def test_returns_true_for_sm_121(self): with ( - mock.patch.object(wheel_utils.shutil, "which", return_value = "/usr/bin/nvidia-smi"), - mock.patch.object(wheel_utils.subprocess, "run", return_value = _smi_result("12.1\n")), + mock.patch.object( + wheel_utils.shutil, "which", return_value = "/usr/bin/nvidia-smi" + ), + mock.patch.object( + wheel_utils.subprocess, "run", return_value = _smi_result("12.1\n") + ), ): assert wheel_utils.has_blackwell_gpu() is True def test_returns_false_for_sm_90(self): with ( - mock.patch.object(wheel_utils.shutil, "which", return_value = "/usr/bin/nvidia-smi"), - mock.patch.object(wheel_utils.subprocess, "run", return_value = _smi_result("9.0\n")), + mock.patch.object( + wheel_utils.shutil, "which", return_value = "/usr/bin/nvidia-smi" + ), + mock.patch.object( + wheel_utils.subprocess, "run", return_value = _smi_result("9.0\n") + ), ): assert wheel_utils.has_blackwell_gpu() is False def test_returns_false_for_sm_89(self): with ( - mock.patch.object(wheel_utils.shutil, "which", return_value = "/usr/bin/nvidia-smi"), - mock.patch.object(wheel_utils.subprocess, "run", return_value = _smi_result("8.9\n")), + mock.patch.object( + wheel_utils.shutil, "which", return_value = "/usr/bin/nvidia-smi" + ), + mock.patch.object( + wheel_utils.subprocess, "run", return_value = _smi_result("8.9\n") + ), ): assert wheel_utils.has_blackwell_gpu() is False def test_mixed_gpus_with_one_blackwell_returns_true(self): with ( - mock.patch.object(wheel_utils.shutil, "which", return_value = "/usr/bin/nvidia-smi"), + mock.patch.object( + wheel_utils.shutil, "which", return_value = "/usr/bin/nvidia-smi" + ), mock.patch.object( wheel_utils.subprocess, "run", @@ -79,7 +101,9 @@ class TestHasBlackwellGpu: def test_returns_false_when_nvidia_smi_fails(self): with ( - mock.patch.object(wheel_utils.shutil, "which", return_value = "/usr/bin/nvidia-smi"), + mock.patch.object( + wheel_utils.shutil, "which", return_value = "/usr/bin/nvidia-smi" + ), mock.patch.object( wheel_utils.subprocess, "run", @@ -90,7 +114,9 @@ class TestHasBlackwellGpu: def test_returns_false_on_subprocess_timeout(self): with ( - mock.patch.object(wheel_utils.shutil, "which", return_value = "/usr/bin/nvidia-smi"), + mock.patch.object( + wheel_utils.shutil, "which", return_value = "/usr/bin/nvidia-smi" + ), mock.patch.object( wheel_utils.subprocess, "run", @@ -101,7 +127,9 @@ class TestHasBlackwellGpu: def test_returns_false_on_malformed_output(self): with ( - mock.patch.object(wheel_utils.shutil, "which", return_value = "/usr/bin/nvidia-smi"), + mock.patch.object( + wheel_utils.shutil, "which", return_value = "/usr/bin/nvidia-smi" + ), mock.patch.object( wheel_utils.subprocess, "run", @@ -133,7 +161,10 @@ class TestFlashAttnWheelSelection: ) assert url is not None assert "v2.8.1" in url - assert "flash_attn-2.8.1+cu12torch2.10cxx11abiTRUE-cp313-cp313-linux_x86_64.whl" in url + assert ( + "flash_attn-2.8.1+cu12torch2.10cxx11abiTRUE-cp313-cp313-linux_x86_64.whl" + in url + ) def test_missing_cuda_major_disables_wheel_lookup(self): assert ( @@ -316,7 +347,10 @@ class TestEnsureFlashAttn: ips._ensure_flash_attn() mock_install_wheel.assert_not_called() - assert ("warning", "No published flash-attn prebuilt wheel found") in step_messages + assert ( + "warning", + "No published flash-attn prebuilt wheel found", + ) in step_messages def test_skip_env_disables_setup_install(self): with ( @@ -357,7 +391,9 @@ class TestEnsureFlashAttn: mock_probe.assert_not_called() mock_install_wheel.assert_not_called() - assert any(label == "warning" and "Blackwell" in msg for label, msg in step_messages) + assert any( + label == "warning" and "Blackwell" in msg for label, msg in step_messages + ) def test_blackwell_gpu_on_windows_emits_blackwell_warning(self): step_messages: list[tuple[str, str]] = [] @@ -383,7 +419,9 @@ class TestEnsureFlashAttn: mock_probe.assert_not_called() mock_install_wheel.assert_not_called() - assert any(label == "warning" and "Blackwell" in msg for label, msg in step_messages) + assert any( + label == "warning" and "Blackwell" in msg for label, msg in step_messages + ) def test_non_blackwell_windows_does_not_emit_blackwell_warning(self): step_messages: list[tuple[str, str]] = [] @@ -435,7 +473,9 @@ class TestInstallPythonStackFlashAttnIntegration: mock.patch("subprocess.run", side_effect = fake_run), mock.patch.object(ips, "_has_usable_nvidia_gpu", return_value = False), mock.patch.object(ips, "_has_rocm_gpu", return_value = False), - mock.patch.object(ips, "LOCAL_DD_UNSTRUCTURED_PLUGIN", Path("/fake/plugin")), + mock.patch.object( + ips, "LOCAL_DD_UNSTRUCTURED_PLUGIN", Path("/fake/plugin") + ), mock.patch("pathlib.Path.is_dir", return_value = True), mock.patch("pathlib.Path.is_file", return_value = True), mock.patch.dict(os.environ, {"SKIP_STUDIO_BASE": "1"}, clear = False), diff --git a/tests/python/test_gpu_init_ldconfig_guard.py b/tests/python/test_gpu_init_ldconfig_guard.py index 248bb84faa..081a6132b4 100644 --- a/tests/python/test_gpu_init_ldconfig_guard.py +++ b/tests/python/test_gpu_init_ldconfig_guard.py @@ -19,7 +19,9 @@ def _find_geteuid_guard(tree: ast.AST): def test_gpu_init_has_geteuid_guard(): tree = ast.parse(GPU_INIT.read_text()) guard = _find_geteuid_guard(tree) - assert guard is not None, "_gpu_init.py must guard ldconfig recovery on os.geteuid()" + assert ( + guard is not None + ), "_gpu_init.py must guard ldconfig recovery on os.geteuid()" def test_ldconfig_calls_only_inside_geteuid_guard(): diff --git a/tests/python/test_no_torch_filtering.py b/tests/python/test_no_torch_filtering.py index f7a90fc6a8..6460274145 100644 --- a/tests/python/test_no_torch_filtering.py +++ b/tests/python/test_no_torch_filtering.py @@ -156,7 +156,9 @@ class TestFilterRequirements: ) # First filter Windows packages, then NO_TORCH packages intermediate = ips._filter_requirements(req, ips.WINDOWS_SKIP_PACKAGES) - result = ips._filter_requirements(Path(intermediate), ips.NO_TORCH_SKIP_PACKAGES) + result = ips._filter_requirements( + Path(intermediate), ips.NO_TORCH_SKIP_PACKAGES + ) lines = Path(result).read_text(encoding = "utf-8").splitlines() non_blank = [l.strip() for l in lines if l.strip()] assert non_blank == [ @@ -175,7 +177,9 @@ class TestFilterRequirements: result = ips._filter_requirements(req, ips.NO_TORCH_SKIP_PACKAGES) lines = Path(result).read_text(encoding = "utf-8").splitlines() non_blank = [l.strip() for l in lines if l.strip()] - assert non_blank == ["numpy"], f"VCS URL line should be filtered, got: {non_blank}" + assert non_blank == [ + "numpy" + ], f"VCS URL line should be filtered, got: {non_blank}" def test_env_marker_line_filtered(self, tmp_path): """Package lines with env markers are still filtered by prefix.""" @@ -189,7 +193,9 @@ class TestFilterRequirements: result = ips._filter_requirements(req, ips.NO_TORCH_SKIP_PACKAGES) lines = Path(result).read_text(encoding = "utf-8").splitlines() non_blank = [l.strip() for l in lines if l.strip()] - assert non_blank == ["numpy"], f"Env marker line should be filtered, got: {non_blank}" + assert non_blank == [ + "numpy" + ], f"Env marker line should be filtered, got: {non_blank}" def test_git_plus_url_not_over_matched(self, tmp_path): """A git+ URL whose path contains a skip package name but does NOT start with it.""" @@ -241,7 +247,9 @@ class TestRealRequirementsFiltering: expected = [ l for l in original - if not any(l.strip().lower().startswith(p) for p in ips.NO_TORCH_SKIP_PACKAGES) + if not any( + l.strip().lower().startswith(p) for p in ips.NO_TORCH_SKIP_PACKAGES + ) ] assert filtered == expected, ( f"Filtered extras.txt should match expected.\n" @@ -251,7 +259,9 @@ class TestRealRequirementsFiltering: def test_extras_no_deps_txt_torchcodec_and_dlpack_removed(self): """extras-no-deps.txt: torchcodec and torch-c-dlpack-ext must be removed.""" - result = ips._filter_requirements(EXTRAS_NO_DEPS_TXT, ips.NO_TORCH_SKIP_PACKAGES) + result = ips._filter_requirements( + EXTRAS_NO_DEPS_TXT, ips.NO_TORCH_SKIP_PACKAGES + ) filtered = self._non_blank_non_comment(Path(result)) original = self._non_blank_non_comment(EXTRAS_NO_DEPS_TXT) @@ -263,7 +273,9 @@ class TestRealRequirementsFiltering: expected = [ l for l in original - if not any(l.strip().lower().startswith(p) for p in ips.NO_TORCH_SKIP_PACKAGES) + if not any( + l.strip().lower().startswith(p) for p in ips.NO_TORCH_SKIP_PACKAGES + ) ] assert filtered == expected @@ -279,7 +291,9 @@ class TestRealRequirementsFiltering: def test_extras_no_deps_txt_trl_preserved(self): """trl should survive NO_TORCH filtering in extras-no-deps.txt.""" - result = ips._filter_requirements(EXTRAS_NO_DEPS_TXT, ips.NO_TORCH_SKIP_PACKAGES) + result = ips._filter_requirements( + EXTRAS_NO_DEPS_TXT, ips.NO_TORCH_SKIP_PACKAGES + ) filtered_text = Path(result).read_text(encoding = "utf-8").lower() assert "trl" in filtered_text, "trl should survive NO_TORCH filtering" @@ -390,7 +404,9 @@ class TestInstallPythonStackSubprocessMock: captured_cmds: list[list[str]] = [] def mock_run(cmd, **kw): - captured_cmds.append(list(cmd) if isinstance(cmd, (list, tuple)) else [str(cmd)]) + captured_cmds.append( + list(cmd) if isinstance(cmd, (list, tuple)) else [str(cmd)] + ) return subprocess.CompletedProcess(cmd, 0, b"", b"") env = {"SKIP_STUDIO_BASE": "1"} if skip_base else {} @@ -407,7 +423,9 @@ class TestInstallPythonStackSubprocessMock: mock.patch.object(ips, "_has_rocm_gpu", return_value = False), mock.patch("subprocess.run", side_effect = mock_run), mock.patch.object(ips, "_bootstrap_uv", return_value = True), - mock.patch.object(ips, "LOCAL_DD_UNSTRUCTURED_PLUGIN", Path("/fake/plugin")), + mock.patch.object( + ips, "LOCAL_DD_UNSTRUCTURED_PLUGIN", Path("/fake/plugin") + ), mock.patch("pathlib.Path.is_dir", return_value = True), mock.patch("pathlib.Path.is_file", return_value = True), ): @@ -450,7 +468,9 @@ class TestInstallPythonStackSubprocessMock: has_extras_nd = self._cmds_contain_file(cmds, "extras-no-deps.txt") or any( "-r" in cmd and "tmp" in cmd.lower() for cmd in cmds ) - assert has_extras_nd, "extras-no-deps.txt (or its filtered temp) should be called" + assert ( + has_extras_nd + ), "extras-no-deps.txt (or its filtered temp) should be called" # -- IS_WINDOWS=True + NO_TORCH=True (stacked) -- @@ -549,13 +569,17 @@ class TestOverridesSkip: def test_no_torch_guard_exists_in_source(self): """The install_python_stack source must contain a NO_TORCH guard around overrides.""" source = Path(ips.__file__).read_text(encoding = "utf-8") - assert "if NO_TORCH:" in source, "NO_TORCH guard not found in install_python_stack.py" + assert ( + "if NO_TORCH:" in source + ), "NO_TORCH guard not found in install_python_stack.py" def test_overrides_skipped_when_no_torch(self): """With NO_TORCH=True on the module, pip_install should NOT be called for overrides.""" source = Path(ips.__file__).read_text(encoding = "utf-8") overrides_match = re.search(r"if NO_TORCH:.*?overrides", source, re.DOTALL) - assert overrides_match is not None, "Expected NO_TORCH conditional before overrides install" + assert ( + overrides_match is not None + ), "Expected NO_TORCH conditional before overrides install" # ── install.sh --no-torch flag tests ────────────────────────────────── @@ -574,21 +598,33 @@ class TestInstallShNoTorchFlag: def test_no_torch_flag_in_case_statement(self): """--no-torch must appear in the flag parser case statement.""" - assert "--no-torch)" in self.source, "--no-torch not found in install.sh flag parser" + assert ( + "--no-torch)" in self.source + ), "--no-torch not found in install.sh flag parser" def test_no_torch_flag_variable_initialized(self): """_NO_TORCH_FLAG must be initialized to false.""" - assert "_NO_TORCH_FLAG=false" in self.source, "_NO_TORCH_FLAG=false not found in install.sh" + assert ( + "_NO_TORCH_FLAG=false" in self.source + ), "_NO_TORCH_FLAG=false not found in install.sh" def test_skip_torch_variable_exists(self): """SKIP_TORCH variable must be defined.""" - assert "SKIP_TORCH=false" in self.source, "SKIP_TORCH=false not found in install.sh" - assert "SKIP_TORCH=true" in self.source, "SKIP_TORCH=true not found in install.sh" + assert ( + "SKIP_TORCH=false" in self.source + ), "SKIP_TORCH=false not found in install.sh" + assert ( + "SKIP_TORCH=true" in self.source + ), "SKIP_TORCH=true not found in install.sh" def test_skip_torch_driven_by_flag_and_mac_intel(self): """SKIP_TORCH must check both _NO_TORCH_FLAG and MAC_INTEL.""" - assert "_NO_TORCH_FLAG" in self.source, "_NO_TORCH_FLAG not referenced in SKIP_TORCH logic" - assert "MAC_INTEL" in self.source, "MAC_INTEL not referenced in SKIP_TORCH logic" + assert ( + "_NO_TORCH_FLAG" in self.source + ), "_NO_TORCH_FLAG not referenced in SKIP_TORCH logic" + assert ( + "MAC_INTEL" in self.source + ), "MAC_INTEL not referenced in SKIP_TORCH logic" def test_unsloth_no_torch_uses_skip_torch(self): """UNSLOTH_NO_TORCH must reference $SKIP_TORCH, not $MAC_INTEL.""" @@ -596,12 +632,18 @@ class TestInstallShNoTorchFlag: matches = re.findall(r'UNSLOTH_NO_TORCH="\$(\w+)"', self.source) for var in matches: - assert var == "SKIP_TORCH", f"UNSLOTH_NO_TORCH references ${var} instead of $SKIP_TORCH" + assert ( + var == "SKIP_TORCH" + ), f"UNSLOTH_NO_TORCH references ${var} instead of $SKIP_TORCH" def test_cpu_hint_message_exists(self): """CPU hint message must exist in install.sh.""" - assert "No GPU detected" in self.source, "CPU hint message not found in install.sh" - assert "--no-torch" in self.source, "--no-torch suggestion not found in CPU hint" + assert ( + "No GPU detected" in self.source + ), "CPU hint message not found in install.sh" + assert ( + "--no-torch" in self.source + ), "--no-torch suggestion not found in CPU hint" def test_no_torch_flag_parsing_subprocess(self): """--no-torch flag sets _NO_TORCH_FLAG=true (subprocess test).""" diff --git a/tests/python/test_studio_import_no_torch.py b/tests/python/test_studio_import_no_torch.py index 7b7a7103d2..09e04d254a 100644 --- a/tests/python/test_studio_import_no_torch.py +++ b/tests/python/test_studio_import_no_torch.py @@ -23,9 +23,15 @@ from pathlib import Path import pytest REPO_ROOT = Path(__file__).resolve().parents[2] -DATA_COLLATORS = REPO_ROOT / "studio" / "backend" / "utils" / "datasets" / "data_collators.py" -CHAT_TEMPLATES = REPO_ROOT / "studio" / "backend" / "utils" / "datasets" / "chat_templates.py" -FORMAT_CONVERSION = REPO_ROOT / "studio" / "backend" / "utils" / "datasets" / "format_conversion.py" +DATA_COLLATORS = ( + REPO_ROOT / "studio" / "backend" / "utils" / "datasets" / "data_collators.py" +) +CHAT_TEMPLATES = ( + REPO_ROOT / "studio" / "backend" / "utils" / "datasets" / "chat_templates.py" +) +FORMAT_CONVERSION = ( + REPO_ROOT / "studio" / "backend" / "utils" / "datasets" / "format_conversion.py" +) def _has_uv() -> bool: @@ -65,7 +71,9 @@ def no_torch_venv(request, tmp_path_factory): [str(venv_python), "-c", "import torch"], capture_output = True, ) - assert check.returncode != 0, f"torch should NOT be importable in fresh {py_version} venv" + assert ( + check.returncode != 0 + ), f"torch should NOT be importable in fresh {py_version} venv" return str(venv_python) @@ -214,7 +222,9 @@ class TestDataCollatorsNoTorchVenv: capture_output = True, timeout = 30, ) - assert result.returncode == 0, f"DeepSeekOCRDataCollator failed:\n{result.stderr.decode()}" + assert ( + result.returncode == 0 + ), f"DeepSeekOCRDataCollator failed:\n{result.stderr.decode()}" assert b"OK: DeepSeekOCRDataCollator instantiated" in result.stdout def test_dataclass_vlm_collator_instantiable(self, no_torch_venv): @@ -235,7 +245,9 @@ class TestDataCollatorsNoTorchVenv: capture_output = True, timeout = 30, ) - assert result.returncode == 0, f"VLMDataCollator failed:\n{result.stderr.decode()}" + assert ( + result.returncode == 0 + ), f"VLMDataCollator failed:\n{result.stderr.decode()}" assert b"OK: VLMDataCollator instantiated" in result.stdout @@ -516,9 +528,12 @@ class TestNegativeControls: capture_output = True, timeout = 30, ) - assert result.returncode != 0, "Expected failure when 'import torch' is prepended" assert ( - b"ModuleNotFoundError" in result.stderr or b"ImportError" in result.stderr + result.returncode != 0 + ), "Expected failure when 'import torch' is prepended" + assert ( + b"ModuleNotFoundError" in result.stderr + or b"ImportError" in result.stderr ), f"Expected ImportError, got:\n{result.stderr.decode()}" finally: os.unlink(temp_file) @@ -561,4 +576,6 @@ class TestNegativeControls: timeout = 30, ) assert result.returncode != 0, "import torch should fail in no-torch venv" - assert b"ModuleNotFoundError" in result.stderr or b"ImportError" in result.stderr + assert ( + b"ModuleNotFoundError" in result.stderr or b"ImportError" in result.stderr + ) diff --git a/tests/python/test_tokenizers_and_torch_constraint.py b/tests/python/test_tokenizers_and_torch_constraint.py index 40924efca2..ccdc879de3 100644 --- a/tests/python/test_tokenizers_and_torch_constraint.py +++ b/tests/python/test_tokenizers_and_torch_constraint.py @@ -17,7 +17,9 @@ _TESTS_DIR = pathlib.Path(__file__).resolve().parent.parent # tests/ _REPO_ROOT = _TESTS_DIR.parent # unsloth/ _INSTALL_SH = _REPO_ROOT / "install.sh" _INSTALL_PS1 = _REPO_ROOT / "install.ps1" -_NO_TORCH_RT = _REPO_ROOT / "studio" / "backend" / "requirements" / "no-torch-runtime.txt" +_NO_TORCH_RT = ( + _REPO_ROOT / "studio" / "backend" / "requirements" / "no-torch-runtime.txt" +) def _read(path: pathlib.Path) -> str: @@ -42,23 +44,30 @@ class TestStructuralTokenizers: def test_tokenizers_present(self): """tokenizers must be a standalone package line.""" pkgs = _lines(_NO_TORCH_RT) - bare_names = [p.split(">")[0].split("<")[0].split("!")[0].split("=")[0] for p in pkgs] + bare_names = [ + p.split(">")[0].split("<")[0].split("!")[0].split("=")[0] for p in pkgs + ] assert "tokenizers" in bare_names def test_tokenizers_before_transformers(self): """tokenizers should appear before transformers (install order intent).""" pkgs = _lines(_NO_TORCH_RT) - bare_names = [p.split(">")[0].split("<")[0].split("!")[0].split("=")[0] for p in pkgs] + bare_names = [ + p.split(">")[0].split("<")[0].split("!")[0].split("=")[0] for p in pkgs + ] idx_tok = bare_names.index("tokenizers") idx_tf = bare_names.index("transformers") assert idx_tok < idx_tf, ( - f"tokenizers at index {idx_tok} should appear before " f"transformers at index {idx_tf}" + f"tokenizers at index {idx_tok} should appear before " + f"transformers at index {idx_tf}" ) def test_torch_not_in_no_torch_file(self): """torch itself must NOT be listed in the no-torch requirements.""" pkgs = _lines(_NO_TORCH_RT) - bare_names = [p.split(">")[0].split("<")[0].split("!")[0].split("=")[0] for p in pkgs] + bare_names = [ + p.split(">")[0].split("<")[0].split("!")[0].split("=")[0] for p in pkgs + ] assert "torch" not in bare_names @@ -399,7 +408,9 @@ class TestE2ETokenizersFix: r = self._pip_install(venv, "--no-deps", "-r", str(_NO_TORCH_RT)) assert r.returncode == 0, f"Install failed: {r.stderr}" - result = self._run_python(venv, "from transformers import AutoConfig; print('OK')") + result = self._run_python( + venv, "from transformers import AutoConfig; print('OK')" + ) assert ( result.returncode == 0 ), f"AutoConfig import failed:\nstdout: {result.stdout}\nstderr: {result.stderr}" @@ -429,15 +440,22 @@ class TestE2ETokenizersFix: req_no_tokenizers = tmp_path / "no-tokenizers.txt" req_no_tokenizers.write_text( "\n".join( - line for line in _read(_NO_TORCH_RT).splitlines() if line.strip() != "tokenizers" + line + for line in _read(_NO_TORCH_RT).splitlines() + if line.strip() != "tokenizers" ), encoding = "utf-8", ) r = self._pip_install(venv, "--no-deps", "-r", str(req_no_tokenizers)) assert r.returncode == 0, f"Install failed: {r.stderr}" result = self._run_python(venv, "from transformers import AutoConfig") - assert result.returncode != 0, "AutoConfig should fail without tokenizers installed" - assert "tokenizers" in result.stderr.lower() or "ModuleNotFoundError" in result.stderr + assert ( + result.returncode != 0 + ), "AutoConfig should fail without tokenizers installed" + assert ( + "tokenizers" in result.stderr.lower() + or "ModuleNotFoundError" in result.stderr + ) # ====================================================================== @@ -516,7 +534,9 @@ class TestE2EFullNoTorchSandbox: venv = self._create_venv(tmp_path, "full-no-torch") r = self._pip_install(venv, "--no-deps", "-r", str(_NO_TORCH_RT)) assert r.returncode == 0, f"Install failed: {r.stderr}" - result = self._run_python(venv, "from transformers import AutoConfig; print('OK')") + result = self._run_python( + venv, "from transformers import AutoConfig; print('OK')" + ) assert ( result.returncode == 0 ), f"AutoConfig failed:\nstdout: {result.stdout}\nstderr: {result.stderr}" diff --git a/tests/python/test_unsloth_run_tool_policy_resolver.py b/tests/python/test_unsloth_run_tool_policy_resolver.py index 6e3e3a722d..6aff02494b 100644 --- a/tests/python/test_unsloth_run_tool_policy_resolver.py +++ b/tests/python/test_unsloth_run_tool_policy_resolver.py @@ -141,11 +141,15 @@ class TestZeroHost: class TestIsExternalHost: - @pytest.mark.parametrize("host", ["127.0.0.1", "localhost", "::1", "LOCALHOST", "Localhost"]) + @pytest.mark.parametrize( + "host", ["127.0.0.1", "localhost", "::1", "LOCALHOST", "Localhost"] + ) def test_loopback_aliases_are_local(self, host): assert is_external_host(host) is False - @pytest.mark.parametrize("host", ["0.0.0.0", "::", "192.168.1.5", "10.0.0.1", "example.com"]) + @pytest.mark.parametrize( + "host", ["0.0.0.0", "::", "192.168.1.5", "10.0.0.1", "example.com"] + ) def test_non_loopback_is_external(self, host): assert is_external_host(host) is True diff --git a/tests/qlora/test_hf_qlora_train_and_merge.py b/tests/qlora/test_hf_qlora_train_and_merge.py index 0892627c46..ae975b0266 100644 --- a/tests/qlora/test_hf_qlora_train_and_merge.py +++ b/tests/qlora/test_hf_qlora_train_and_merge.py @@ -91,7 +91,9 @@ if __name__ == "__main__": print(training_args) print(peft_config) - trainer = setup_trainer(model, tokenizer, dataset, training_args, peft_config = peft_config) + trainer = setup_trainer( + model, tokenizer, dataset, training_args, peft_config = peft_config + ) with header_footer_context("Model"): print(type(model.model)) diff --git a/tests/saving/gpt-oss-merge/test_merged_model.py b/tests/saving/gpt-oss-merge/test_merged_model.py index 024d48f276..d965f44966 100644 --- a/tests/saving/gpt-oss-merge/test_merged_model.py +++ b/tests/saving/gpt-oss-merge/test_merged_model.py @@ -42,7 +42,9 @@ inputs = merged_tokenizer.apply_chat_template( reasoning_effort = "low", # low, medium or high ).to(merged_model.device) -_ = merged_model.generate(**inputs, max_new_tokens = 512, streamer = TextStreamer(merged_tokenizer)) +_ = merged_model.generate( + **inputs, max_new_tokens = 512, streamer = TextStreamer(merged_tokenizer) +) print("\n✅ Inference complete.") # --- Final Cleanup --- @@ -52,5 +54,7 @@ torch.cuda.empty_cache() gc.collect() safe_remove_directory("./gpt-oss-finetuned-merged") -safe_remove_directory("./unsloth_compiled_cache") # Clean up cache created by this process +safe_remove_directory( + "./unsloth_compiled_cache" +) # Clean up cache created by this process print("✅ Final cleanup complete. Exiting inference script.") diff --git a/tests/saving/gpt-oss-merge/train_and_merge.py b/tests/saving/gpt-oss-merge/train_and_merge.py index 1d35c6759d..0ca5cd0db5 100644 --- a/tests/saving/gpt-oss-merge/train_and_merge.py +++ b/tests/saving/gpt-oss-merge/train_and_merge.py @@ -27,7 +27,9 @@ tokenizer = None def formatting_prompts_func(examples): convos = examples["messages"] texts = [ - tokenizer.apply_chat_template(convo, tokenize = False, add_generation_prompt = False) + tokenizer.apply_chat_template( + convo, tokenize = False, add_generation_prompt = False + ) for convo in convos ] return {"text": texts} @@ -81,7 +83,9 @@ print("Fine-tuning complete.") # --- Merge and Save --- print("\n💾 Merging and saving the 16-bit model to './gpt-oss-finetuned-merged'...") -model.save_pretrained_merged(save_directory = "./gpt-oss-finetuned-merged", tokenizer = tokenizer) +model.save_pretrained_merged( + save_directory = "./gpt-oss-finetuned-merged", tokenizer = tokenizer +) print("✅ Model merged and saved.") # --- Cleanup --- @@ -91,5 +95,7 @@ torch.cuda.empty_cache() gc.collect() safe_remove_directory("./outputs") -safe_remove_directory("./unsloth_compiled_cache") # Clean up the cache created by this process +safe_remove_directory( + "./unsloth_compiled_cache" +) # Clean up the cache created by this process print("✅ Cleanup complete. Exiting training script.") diff --git a/tests/saving/language_models/test_merge_4bit_validation.py b/tests/saving/language_models/test_merge_4bit_validation.py index 2e2e823bd3..9990028624 100644 --- a/tests/saving/language_models/test_merge_4bit_validation.py +++ b/tests/saving/language_models/test_merge_4bit_validation.py @@ -16,7 +16,9 @@ from tests.utils.cleanup_utils import safe_remove_directory def formatting_prompts_func(examples): convos = examples["messages"] texts = [ - tokenizer.apply_chat_template(convo, tokenize = False, add_generation_prompt = False) + tokenizer.apply_chat_template( + convo, tokenize = False, add_generation_prompt = False + ) for convo in convos ] return {"text": texts} @@ -48,7 +50,9 @@ tokenizer = get_chat_template( chat_template = "llama-3.1", ) -dataset_train = load_dataset("allenai/openassistant-guanaco-reformatted", split = "train[:100]") +dataset_train = load_dataset( + "allenai/openassistant-guanaco-reformatted", split = "train[:100]" +) dataset_train = dataset_train.map(formatting_prompts_func, batched = True) print("✅ Base model loaded successfully!") diff --git a/tests/saving/language_models/test_merge_model_perplexity_llama-3.2.py b/tests/saving/language_models/test_merge_model_perplexity_llama-3.2.py index f8f81c550e..710d86ab6b 100644 --- a/tests/saving/language_models/test_merge_model_perplexity_llama-3.2.py +++ b/tests/saving/language_models/test_merge_model_perplexity_llama-3.2.py @@ -34,7 +34,9 @@ from tests.utils.perplexity_eval import ( def formatting_prompts_func(examples): convos = examples["messages"] texts = [ - tokenizer.apply_chat_template(convo, tokenize = False, add_generation_prompt = False) + tokenizer.apply_chat_template( + convo, tokenize = False, add_generation_prompt = False + ) for convo in convos ] return {"text": texts} @@ -62,12 +64,16 @@ def load_and_compute_8bit_ppl( ) # Load dataset fresh in subprocess - dataset_ppl = load_dataset("allenai/openassistant-guanaco-reformatted", split = "eval") + dataset_ppl = load_dataset( + "allenai/openassistant-guanaco-reformatted", split = "eval" + ) def formatting_prompts_func(examples): convos = examples["messages"] texts = [ - merged_tokenizer.apply_chat_template(convo, tokenize = False, add_generation_prompt = False) + merged_tokenizer.apply_chat_template( + convo, tokenize = False, add_generation_prompt = False + ) for convo in convos ] return {"text": texts} @@ -120,8 +126,12 @@ if __name__ == "__main__": from unsloth.chat_templates import standardize_sharegpt - dataset_train = load_dataset("allenai/openassistant-guanaco-reformatted", split = "train") - dataset_ppl = load_dataset("allenai/openassistant-guanaco-reformatted", split = "eval") + dataset_train = load_dataset( + "allenai/openassistant-guanaco-reformatted", split = "train" + ) + dataset_ppl = load_dataset( + "allenai/openassistant-guanaco-reformatted", split = "eval" + ) dataset_train = dataset_train.map(formatting_prompts_func, batched = True) dataset_ppl = dataset_ppl.map(formatting_prompts_func, batched = True) diff --git a/tests/saving/language_models/test_merge_model_perplexity_mistral.py b/tests/saving/language_models/test_merge_model_perplexity_mistral.py index d467089a47..46c4062908 100644 --- a/tests/saving/language_models/test_merge_model_perplexity_mistral.py +++ b/tests/saving/language_models/test_merge_model_perplexity_mistral.py @@ -51,7 +51,9 @@ def load_and_compute_8bit_ppl( # ) # Load dataset fresh in subprocess. - dataset_ppl = load_dataset("allenai/openassistant-guanaco-reformatted", split = "eval") + dataset_ppl = load_dataset( + "allenai/openassistant-guanaco-reformatted", split = "eval" + ) alpaca_prompt = """Below is an instruction that describes a task, paired with an input that provides further context. Write a response that appropriately completes the request. @@ -87,7 +89,10 @@ def load_and_compute_8bit_ppl( inputs.append(user_message) outputs.append(assistant_message) - text = alpaca_prompt.format(instruction, user_message, assistant_message) + EOS_TOKEN + text = ( + alpaca_prompt.format(instruction, user_message, assistant_message) + + EOS_TOKEN + ) texts.append(text) return { @@ -172,7 +177,10 @@ if __name__ == "__main__": inputs.append(user_message) outputs.append(assistant_message) - text = alpaca_prompt.format(instruction, user_message, assistant_message) + EOS_TOKEN + text = ( + alpaca_prompt.format(instruction, user_message, assistant_message) + + EOS_TOKEN + ) texts.append(text) return { @@ -182,8 +190,12 @@ if __name__ == "__main__": "text": texts, } - dataset_train = load_dataset("allenai/openassistant-guanaco-reformatted", split = "train") - dataset_ppl = load_dataset("allenai/openassistant-guanaco-reformatted", split = "eval") + dataset_train = load_dataset( + "allenai/openassistant-guanaco-reformatted", split = "train" + ) + dataset_ppl = load_dataset( + "allenai/openassistant-guanaco-reformatted", split = "eval" + ) dataset_train = dataset_train.map(formatting_prompts_func, batched = True) dataset_ppl = dataset_ppl.map(formatting_prompts_func, batched = True) diff --git a/tests/saving/language_models/test_merge_model_perplexity_phi_4.py b/tests/saving/language_models/test_merge_model_perplexity_phi_4.py index 391d1b90db..5fbf2bc026 100644 --- a/tests/saving/language_models/test_merge_model_perplexity_phi_4.py +++ b/tests/saving/language_models/test_merge_model_perplexity_phi_4.py @@ -34,7 +34,9 @@ from tests.utils.perplexity_eval import ( def formatting_prompts_func(examples): convos = examples["messages"] texts = [ - tokenizer.apply_chat_template(convo, tokenize = False, add_generation_prompt = False) + tokenizer.apply_chat_template( + convo, tokenize = False, add_generation_prompt = False + ) for convo in convos ] return { @@ -64,12 +66,16 @@ def load_and_compute_8bit_ppl( ) # Load dataset fresh in subprocess - dataset_ppl = load_dataset("allenai/openassistant-guanaco-reformatted", split = "eval") + dataset_ppl = load_dataset( + "allenai/openassistant-guanaco-reformatted", split = "eval" + ) def formatting_prompts_func(examples): convos = examples["messages"] texts = [ - merged_tokenizer.apply_chat_template(convo, tokenize = False, add_generation_prompt = False) + merged_tokenizer.apply_chat_template( + convo, tokenize = False, add_generation_prompt = False + ) for convo in convos ] return {"text": texts} @@ -120,8 +126,12 @@ if __name__ == "__main__": chat_template = "phi-4", ) - dataset_train = load_dataset("allenai/openassistant-guanaco-reformatted", split = "train") - dataset_ppl = load_dataset("allenai/openassistant-guanaco-reformatted", split = "eval") + dataset_train = load_dataset( + "allenai/openassistant-guanaco-reformatted", split = "train" + ) + dataset_ppl = load_dataset( + "allenai/openassistant-guanaco-reformatted", split = "eval" + ) dataset_train = dataset_train.map(formatting_prompts_func, batched = True) dataset_ppl = dataset_ppl.map(formatting_prompts_func, batched = True) diff --git a/tests/saving/language_models/test_merged_model_perplexity_llama-3.1-8b.py b/tests/saving/language_models/test_merged_model_perplexity_llama-3.1-8b.py index 9c2c1acb41..1285cbe69e 100644 --- a/tests/saving/language_models/test_merged_model_perplexity_llama-3.1-8b.py +++ b/tests/saving/language_models/test_merged_model_perplexity_llama-3.1-8b.py @@ -34,7 +34,9 @@ from tests.utils.perplexity_eval import ( def formatting_prompts_func(examples): convos = examples["messages"] texts = [ - tokenizer.apply_chat_template(convo, tokenize = False, add_generation_prompt = False) + tokenizer.apply_chat_template( + convo, tokenize = False, add_generation_prompt = False + ) for convo in convos ] return {"text": texts} @@ -62,12 +64,16 @@ def load_and_compute_8bit_ppl( ) # Load dataset fresh in subprocess - dataset_ppl = load_dataset("allenai/openassistant-guanaco-reformatted", split = "eval") + dataset_ppl = load_dataset( + "allenai/openassistant-guanaco-reformatted", split = "eval" + ) def formatting_prompts_func(examples): convos = examples["messages"] texts = [ - merged_tokenizer.apply_chat_template(convo, tokenize = False, add_generation_prompt = False) + merged_tokenizer.apply_chat_template( + convo, tokenize = False, add_generation_prompt = False + ) for convo in convos ] return {"text": texts} @@ -120,8 +126,12 @@ if __name__ == "__main__": from unsloth.chat_templates import standardize_sharegpt - dataset_train = load_dataset("allenai/openassistant-guanaco-reformatted", split = "train") - dataset_ppl = load_dataset("allenai/openassistant-guanaco-reformatted", split = "eval") + dataset_train = load_dataset( + "allenai/openassistant-guanaco-reformatted", split = "train" + ) + dataset_ppl = load_dataset( + "allenai/openassistant-guanaco-reformatted", split = "eval" + ) dataset_train = dataset_train.map(formatting_prompts_func, batched = True) dataset_ppl = dataset_ppl.map(formatting_prompts_func, batched = True) diff --git a/tests/saving/language_models/test_merged_model_perplexity_qwen_2.5.py b/tests/saving/language_models/test_merged_model_perplexity_qwen_2.5.py index cdbad174a4..701c5c3f6d 100644 --- a/tests/saving/language_models/test_merged_model_perplexity_qwen_2.5.py +++ b/tests/saving/language_models/test_merged_model_perplexity_qwen_2.5.py @@ -96,7 +96,9 @@ def load_and_compute_8bit_ppl( # ) # Load dataset fresh in subprocess - dataset_ppl = load_dataset("allenai/openassistant-guanaco-reformatted", split = "eval") + dataset_ppl = load_dataset( + "allenai/openassistant-guanaco-reformatted", split = "eval" + ) alpaca_prompt = """Below is an instruction that describes a task, paired with an input that provides further context. Write a response that appropriately completes the request. @@ -182,8 +184,12 @@ if __name__ == "__main__": attn_implementation = attn_implementation, ) - dataset_train = load_dataset("allenai/openassistant-guanaco-reformatted", split = "train") - dataset_ppl = load_dataset("allenai/openassistant-guanaco-reformatted", split = "eval") + dataset_train = load_dataset( + "allenai/openassistant-guanaco-reformatted", split = "train" + ) + dataset_ppl = load_dataset( + "allenai/openassistant-guanaco-reformatted", split = "eval" + ) dataset_train = dataset_train.map(formatting_prompts_func, batched = True) dataset_ppl = dataset_ppl.map(formatting_prompts_func, batched = True) diff --git a/tests/saving/language_models/test_push_to_hub_merged.py b/tests/saving/language_models/test_push_to_hub_merged.py index cacb1e2a10..857802b9f0 100644 --- a/tests/saving/language_models/test_push_to_hub_merged.py +++ b/tests/saving/language_models/test_push_to_hub_merged.py @@ -35,7 +35,9 @@ from tests.utils.perplexity_eval import ( def formatting_prompts_func(examples): convos = examples["messages"] texts = [ - tokenizer.apply_chat_template(convo, tokenize = False, add_generation_prompt = False) + tokenizer.apply_chat_template( + convo, tokenize = False, add_generation_prompt = False + ) for convo in convos ] return {"text": texts} @@ -173,7 +175,9 @@ try: print("=== TESTING MODEL DOWNLOAD ===".center(80)) print("=" * 80 + "\n") # Force download even if cached - model, tokenizer = FastLanguageModel.from_pretrained(f"{hf_username}/merged_llama_text_model") + model, tokenizer = FastLanguageModel.from_pretrained( + f"{hf_username}/merged_llama_text_model" + ) success["download"] = True print("✅ Model downloaded successfully!") except Exception as e: diff --git a/tests/saving/language_models/test_push_to_hub_merged_sharded_index_file.py b/tests/saving/language_models/test_push_to_hub_merged_sharded_index_file.py index 59a09f06ec..4a3726dba3 100644 --- a/tests/saving/language_models/test_push_to_hub_merged_sharded_index_file.py +++ b/tests/saving/language_models/test_push_to_hub_merged_sharded_index_file.py @@ -36,7 +36,9 @@ from tests.utils.perplexity_eval import ( def formatting_prompts_func(examples): convos = examples["messages"] texts = [ - tokenizer.apply_chat_template(convo, tokenize = False, add_generation_prompt = False) + tokenizer.apply_chat_template( + convo, tokenize = False, add_generation_prompt = False + ) for convo in convos ] return {"text": texts} @@ -192,7 +194,9 @@ try: print("=== TESTING MODEL DOWNLOAD ===".center(80)) print("=" * 80 + "\n") # Force download even if cached - model, tokenizer = FastLanguageModel.from_pretrained(f"{hf_username}/merged_llama_text_model") + model, tokenizer = FastLanguageModel.from_pretrained( + f"{hf_username}/merged_llama_text_model" + ) success["download"] = True print("✅ Model downloaded successfully!") except Exception as e: diff --git a/tests/saving/language_models/test_save_merged_grpo_model.py b/tests/saving/language_models/test_save_merged_grpo_model.py index 63e4b970a6..357cee537f 100644 --- a/tests/saving/language_models/test_save_merged_grpo_model.py +++ b/tests/saving/language_models/test_save_merged_grpo_model.py @@ -177,7 +177,9 @@ def training_run(result_queue): avg_length = sum(lengths) / len(lengths) min_length = min(lengths) - print(f"Prompt lengths - Min: {min_length}, Max: {max_length}, Avg: {avg_length:.1f}") + print( + f"Prompt lengths - Min: {min_length}, Max: {max_length}, Avg: {avg_length:.1f}" + ) return max_length, avg_length def extract_unsloth_answer( @@ -268,7 +270,9 @@ def training_run(result_queue): ground_truth_num = float(norm_ground_truth) if ground_truth_num != 0: - relative_error = abs(extracted_num - ground_truth_num) / abs(ground_truth_num) + relative_error = abs(extracted_num - ground_truth_num) / abs( + ground_truth_num + ) if relative_error < 0.01: return True, True, 0.9 @@ -303,7 +307,10 @@ def training_run(result_queue): ) responses = [completion[0]["content"] for completion in completions] - rewards = [3.0 if re.match(pattern, response, re.DOTALL) else 0.0 for response in responses] + rewards = [ + 3.0 if re.match(pattern, response, re.DOTALL) else 0.0 + for response in responses + ] return rewards def match_format_approximately(completions, **kwargs): @@ -401,7 +408,9 @@ def training_run(result_queue): format_improvement = ( result["correct_format_pct"] - base_result["correct_format_pct"] ) - exact_improvement = result["exact_match_pct"] - base_result["exact_match_pct"] + exact_improvement = ( + result["exact_match_pct"] - base_result["exact_match_pct"] + ) plausible_improvement = ( result["plausible_match_pct"] - base_result["plausible_match_pct"] ) @@ -433,7 +442,9 @@ def training_run(result_queue): if torch.cuda.is_available(): allocated = torch.cuda.memory_allocated() / 1024**3 reserved = torch.cuda.memory_reserved() / 1024**3 - print(f"GPU memory - Allocated: {allocated:.2f} GB, Reserved: {reserved:.2f} GB") + print( + f"GPU memory - Allocated: {allocated:.2f} GB, Reserved: {reserved:.2f} GB" + ) """#### Data Loading and Preparation""" @@ -472,7 +483,9 @@ def training_run(result_queue): def formatting_prompts_func(examples): convos = examples["prompt"] texts = [ - tokenizer.apply_chat_template(convo, tokenize = False, add_generation_prompt = False) + tokenizer.apply_chat_template( + convo, tokenize = False, add_generation_prompt = False + ) for convo in convos ] return { @@ -689,7 +702,9 @@ def training_run(result_queue): print(f"{'='*60}") try: - model.save_pretrained_merged("final_merged_model", tokenizer, save_method = "merged_16bit") + model.save_pretrained_merged( + "final_merged_model", tokenizer, save_method = "merged_16bit" + ) print("✅ Merged model saved to: final_merged_model/") except Exception as e: print(f"⚠️ Could not save merged model: {e}") diff --git a/tests/saving/test_fix_sentencepiece_gguf_robustness.py b/tests/saving/test_fix_sentencepiece_gguf_robustness.py index 9c61ca4067..49bd70fa2f 100644 --- a/tests/saving/test_fix_sentencepiece_gguf_robustness.py +++ b/tests/saving/test_fix_sentencepiece_gguf_robustness.py @@ -44,7 +44,9 @@ def test_user_defined_special_piece_is_not_retyped(tmp_path): ] (tmp_path / "tokenizer.model").write_bytes(_build(pieces)) (tmp_path / "tokenizer.json").write_text( - json.dumps({"added_tokens": [{"id": 2, "content": "", "special": True}]}) + json.dumps( + {"added_tokens": [{"id": 2, "content": "", "special": True}]} + ) ) fix_sentencepiece_gguf(str(tmp_path)) got = dict(_read(str(tmp_path / "tokenizer.model"))) @@ -85,7 +87,10 @@ def test_save_py_except_clause_is_broad_exception(): with open(_SAVE_PY) as f: tree = ast.parse(f.read()) for node in ast.walk(tree): - if isinstance(node, ast.FunctionDef) and node.name == "unsloth_save_pretrained_gguf": + if ( + isinstance(node, ast.FunctionDef) + and node.name == "unsloth_save_pretrained_gguf" + ): for subnode in ast.walk(node): if isinstance(subnode, ast.Try): body_src = "\n".join(ast.unparse(s) for s in subnode.body) diff --git a/tests/saving/test_preserve_tokenizer_eos_token.py b/tests/saving/test_preserve_tokenizer_eos_token.py index 2ea40ab778..6e2f8c7f9d 100644 --- a/tests/saving/test_preserve_tokenizer_eos_token.py +++ b/tests/saving/test_preserve_tokenizer_eos_token.py @@ -16,7 +16,8 @@ def _load_preserve_helper(): helper = next( node for node in tree.body - if isinstance(node, ast.FunctionDef) and node.name == "_preserve_tokenizer_eos_token" + if isinstance(node, ast.FunctionDef) + and node.name == "_preserve_tokenizer_eos_token" ) module = ast.Module(body = [helper], type_ignores = []) ast.fix_missing_locations(module) @@ -45,7 +46,9 @@ def test_preserve_tokenizer_eos_token_supports_processor_tokenizer(tmp_path): preserve = _load_preserve_helper() tokenizer_config = tmp_path / "tokenizer_config.json" tokenizer_config.write_text(json.dumps({"eos_token": ""}), encoding = "utf-8") - processor = types.SimpleNamespace(tokenizer = types.SimpleNamespace(eos_token = "")) + processor = types.SimpleNamespace( + tokenizer = types.SimpleNamespace(eos_token = "") + ) preserve(processor, tmp_path) diff --git a/tests/saving/test_save_shell_injection.py b/tests/saving/test_save_shell_injection.py index b02748c250..c6c2c8fe15 100644 --- a/tests/saving/test_save_shell_injection.py +++ b/tests/saving/test_save_shell_injection.py @@ -19,7 +19,10 @@ def _assert_safe_ggml_calls(calls: list[ast.Call]) -> None: popen_calls = [] for call in calls: if isinstance(call.func, ast.Attribute) and call.func.attr == "Popen": - if isinstance(call.func.value, ast.Name) and call.func.value.id == "subprocess": + if ( + isinstance(call.func.value, ast.Name) + and call.func.value.id == "subprocess" + ): popen_calls.append(call) assert popen_calls, "Expected at least one subprocess.Popen call" @@ -51,7 +54,9 @@ def _assert_safe_ggml_calls(calls: list[ast.Call]) -> None: assert call.args, "subprocess.Popen must receive argv as a positional argument" argv = call.args[0] - assert isinstance(argv, ast.List), "subprocess.Popen must be called with an argv list" + assert isinstance( + argv, ast.List + ), "subprocess.Popen must be called with an argv list" assert len(argv.elts) == 5, "GGML conversion argv should have five elements" second_arg = argv.elts[1] diff --git a/tests/saving/test_save_subprocess_utf8_encoding.py b/tests/saving/test_save_subprocess_utf8_encoding.py index 4a609cd7b7..c81282369c 100644 --- a/tests/saving/test_save_subprocess_utf8_encoding.py +++ b/tests/saving/test_save_subprocess_utf8_encoding.py @@ -79,7 +79,9 @@ def _collect_text_mode_subprocess_calls() -> list[ast.Call]: return [ node for node in ast.walk(tree) - if isinstance(node, ast.Call) and _is_subprocess_call(node) and _is_text_mode(node) + if isinstance(node, ast.Call) + and _is_subprocess_call(node) + and _is_text_mode(node) ] diff --git a/tests/saving/test_unsloth_save.py b/tests/saving/test_unsloth_save.py index a85c119e04..3ee2037f39 100644 --- a/tests/saving/test_unsloth_save.py +++ b/tests/saving/test_unsloth_save.py @@ -128,13 +128,19 @@ def test_save_merged_16bit(model, tokenizer, temp_save_dir: str): model.config._name_or_path.replace("/", "_"), ) - model.save_pretrained_merged(save_path, tokenizer = tokenizer, save_method = "merged_16bit") + model.save_pretrained_merged( + save_path, tokenizer = tokenizer, save_method = "merged_16bit" + ) assert os.path.isdir(save_path), f"Directory {save_path} does not exist." - assert os.path.isfile(os.path.join(save_path, "config.json")), "config.json not found." + assert os.path.isfile( + os.path.join(save_path, "config.json") + ), "config.json not found." weight_files = [ - f for f in os.listdir(save_path) if f.endswith(".bin") or f.endswith(".safetensors") + f + for f in os.listdir(save_path) + if f.endswith(".bin") or f.endswith(".safetensors") ] assert len(weight_files) > 0, "No weight files found in the save directory." @@ -148,7 +154,9 @@ def test_save_merged_16bit(model, tokenizer, temp_save_dir: str): with open(config_path, "r") as f: config = json.load(f) - assert "quantization_config" not in config, "Quantization config not found in the model config." + assert ( + "quantization_config" not in config + ), "Quantization config not found in the model config." total_size = sum(os.path.getsize(os.path.join(save_path, f)) for f in weight_files) save_file_sizes["merged_16bit"][model.config._name_or_path] = total_size @@ -170,13 +178,19 @@ def test_save_merged_4bit(model, tokenizer, temp_save_dir: str): model.config._name_or_path.replace("/", "_"), ) - model.save_pretrained_merged(save_path, tokenizer = tokenizer, save_method = "merged_4bit_forced") + model.save_pretrained_merged( + save_path, tokenizer = tokenizer, save_method = "merged_4bit_forced" + ) assert os.path.isdir(save_path), f"Directory {save_path} does not exist." - assert os.path.isfile(os.path.join(save_path, "config.json")), "config.json not found." + assert os.path.isfile( + os.path.join(save_path, "config.json") + ), "config.json not found." weight_files = [ - f for f in os.listdir(save_path) if f.endswith(".bin") or f.endswith(".safetensors") + f + for f in os.listdir(save_path) + if f.endswith(".bin") or f.endswith(".safetensors") ] assert len(weight_files) > 0, "No weight files found in the save directory." @@ -199,7 +213,9 @@ def test_save_merged_4bit(model, tokenizer, temp_save_dir: str): with open(config_path, "r") as f: config = json.load(f) - assert "quantization_config" in config, "Quantization config not found in the model config." + assert ( + "quantization_config" in config + ), "Quantization config not found in the model config." # Verify the saved model loads loaded_model, loaded_tokenizer = FastModel.from_pretrained( @@ -231,18 +247,28 @@ def test_save_torchao(fp16_model_tokenizer, temp_save_dir: str): ) weight_files_16bit = [ - f for f in os.listdir(save_path) if f.endswith(".bin") or f.endswith(".safetensors") + f + for f in os.listdir(save_path) + if f.endswith(".bin") or f.endswith(".safetensors") ] - total_16bit_size = sum(os.path.getsize(os.path.join(save_path, f)) for f in weight_files_16bit) + total_16bit_size = sum( + os.path.getsize(os.path.join(save_path, f)) for f in weight_files_16bit + ) save_file_sizes["merged_16bit"][model.config._name_or_path] = total_16bit_size torchao_save_path = save_path + "-torchao" - assert os.path.isdir(torchao_save_path), f"Directory {torchao_save_path} does not exist." - assert os.path.isfile(os.path.join(torchao_save_path, "config.json")), "config.json not found." + assert os.path.isdir( + torchao_save_path + ), f"Directory {torchao_save_path} does not exist." + assert os.path.isfile( + os.path.join(torchao_save_path, "config.json") + ), "config.json not found." weight_files = [ - f for f in os.listdir(torchao_save_path) if f.endswith(".bin") or f.endswith(".safetensors") + f + for f in os.listdir(torchao_save_path) + if f.endswith(".bin") or f.endswith(".safetensors") ] assert len(weight_files) > 0, "No weight files found in the save directory." @@ -251,7 +277,9 @@ def test_save_torchao(fp16_model_tokenizer, temp_save_dir: str): os.path.join(torchao_save_path, file) ), f"{file} not found in the save directory." - total_size = sum(os.path.getsize(os.path.join(torchao_save_path, f)) for f in weight_files) + total_size = sum( + os.path.getsize(os.path.join(torchao_save_path, f)) for f in weight_files + ) save_file_sizes["torchao"][model.config._name_or_path] = total_size assert ( @@ -262,7 +290,9 @@ def test_save_torchao(fp16_model_tokenizer, temp_save_dir: str): with open(config_path, "r") as f: config = json.load(f) - assert "quantization_config" in config, "Quantization config not found in the model config." + assert ( + "quantization_config" in config + ), "Quantization config not found in the model config." # load_in_4bit must stay False: a torchao-quantized model can't be # re-quantized with bitsandbytes. @@ -287,7 +317,9 @@ def test_save_and_inference_torchao(fp16_model_tokenizer, temp_save_dir: str): print(f"Testing TorchAO save and inference for: {model_name}") - save_path = os.path.join(temp_save_dir, "torchao_models", model_name.replace("/", "_")) + save_path = os.path.join( + temp_save_dir, "torchao_models", model_name.replace("/", "_") + ) from torchao.quantization import Int8DynamicActivationInt8WeightConfig diff --git a/tests/saving/text_to_speech_models/test_csm.py b/tests/saving/text_to_speech_models/test_csm.py index 3d04bc39cb..3e4b2d6430 100644 --- a/tests/saving/text_to_speech_models/test_csm.py +++ b/tests/saving/text_to_speech_models/test_csm.py @@ -134,7 +134,9 @@ import torch output_audio_path = "csm_audio.wav" try: - text = "We just finished fine tuning a text to speech model... and it's pretty good!" + text = ( + "We just finished fine tuning a text to speech model... and it's pretty good!" + ) speaker_id = 0 inputs = processor(f"[{speaker_id}]{text}", add_special_tokens = True).to("cuda") audio_values = model.generate( diff --git a/tests/saving/text_to_speech_models/test_lasa.py b/tests/saving/text_to_speech_models/test_lasa.py index c0c4f80e0e..804ff512f9 100644 --- a/tests/saving/text_to_speech_models/test_lasa.py +++ b/tests/saving/text_to_speech_models/test_lasa.py @@ -167,7 +167,9 @@ def extract_speech_ids(speech_tokens_str): # TTS start! with torch.inference_mode(): with torch.amp.autocast("cuda", dtype = model.dtype): - formatted_text = f"<|TEXT_UNDERSTANDING_START|>{input_text}<|TEXT_UNDERSTANDING_END|>" + formatted_text = ( + f"<|TEXT_UNDERSTANDING_START|>{input_text}<|TEXT_UNDERSTANDING_END|>" + ) # Tokenize the text chat = [ diff --git a/tests/saving/text_to_speech_models/test_orpheus.py b/tests/saving/text_to_speech_models/test_orpheus.py index 83adaf8dbc..c748f023a5 100644 --- a/tests/saving/text_to_speech_models/test_orpheus.py +++ b/tests/saving/text_to_speech_models/test_orpheus.py @@ -151,7 +151,9 @@ for prompt in prompts_: all_input_ids.append(input_ids) start_token = torch.tensor([[128259]], dtype = torch.int64) # Start of human -end_tokens = torch.tensor([[128009, 128260]], dtype = torch.int64) # End of text, End of human +end_tokens = torch.tensor( + [[128009, 128260]], dtype = torch.int64 +) # End of text, End of human all_modified_input_ids = [] for input_ids in all_input_ids: @@ -162,7 +164,9 @@ for input_ids in all_input_ids: all_padded_tensors = [] all_attention_masks = [] -max_length = max([modified_input_ids.shape[1] for modified_input_ids in all_modified_input_ids]) +max_length = max( + [modified_input_ids.shape[1] for modified_input_ids in all_modified_input_ids] +) for modified_input_ids in all_modified_input_ids: padding = max_length - modified_input_ids.shape[1] padded_tensor = torch.cat( diff --git a/tests/saving/text_to_speech_models/test_whisper.py b/tests/saving/text_to_speech_models/test_whisper.py index e0271d098e..3fa44927b8 100644 --- a/tests/saving/text_to_speech_models/test_whisper.py +++ b/tests/saving/text_to_speech_models/test_whisper.py @@ -179,9 +179,13 @@ expected_phrases = [ ] transcribed_lower = transcribed_text["text"].lower() -all_phrases_found = all(phrase.lower() in transcribed_lower for phrase in expected_phrases) +all_phrases_found = all( + phrase.lower() in transcribed_lower for phrase in expected_phrases +) -assert all_phrases_found, f"Expected phrases not found in transcription: {transcribed_text['text']}" +assert ( + all_phrases_found +), f"Expected phrases not found in transcription: {transcribed_text['text']}" print("✅ Transcription contains all expected phrases!") diff --git a/tests/saving/vision_models/test_index_file_sharded_model.py b/tests/saving/vision_models/test_index_file_sharded_model.py index 6f2d0ff782..048df45469 100644 --- a/tests/saving/vision_models/test_index_file_sharded_model.py +++ b/tests/saving/vision_models/test_index_file_sharded_model.py @@ -128,7 +128,9 @@ try: per_device_train_batch_size = 2, gradient_accumulation_steps = 4, gradient_checkpointing = True, - gradient_checkpointing_kwargs = {"use_reentrant": False}, # use reentrant checkpointing + gradient_checkpointing_kwargs = { + "use_reentrant": False + }, # use reentrant checkpointing max_grad_norm = 0.3, # max gradient norm based on QLoRA paper warmup_ratio = 0.03, # num_train_epochs = 2, # Set this instead of max_steps for full training runs diff --git a/tests/saving/vision_models/test_push_to_hub_merged.py b/tests/saving/vision_models/test_push_to_hub_merged.py index d83c66de4b..ea1ec0d191 100644 --- a/tests/saving/vision_models/test_push_to_hub_merged.py +++ b/tests/saving/vision_models/test_push_to_hub_merged.py @@ -138,7 +138,9 @@ try: per_device_train_batch_size = 2, gradient_accumulation_steps = 4, gradient_checkpointing = True, - gradient_checkpointing_kwargs = {"use_reentrant": False}, # use reentrant checkpointing + gradient_checkpointing_kwargs = { + "use_reentrant": False + }, # use reentrant checkpointing max_grad_norm = 0.3, # max gradient norm based on QLoRA paper warmup_ratio = 0.03, # num_train_epochs = 2, # Set this instead of max_steps for full training runs diff --git a/tests/saving/vision_models/test_save_merge_qwen2.5vl32B_model_ocr_benchmark.py b/tests/saving/vision_models/test_save_merge_qwen2.5vl32B_model_ocr_benchmark.py index 548e6bfc37..84e2d4ea2b 100644 --- a/tests/saving/vision_models/test_save_merge_qwen2.5vl32B_model_ocr_benchmark.py +++ b/tests/saving/vision_models/test_save_merge_qwen2.5vl32B_model_ocr_benchmark.py @@ -128,7 +128,9 @@ trainer = SFTTrainer( per_device_train_batch_size = 2, gradient_accumulation_steps = 4, gradient_checkpointing = True, - gradient_checkpointing_kwargs = {"use_reentrant": False}, # use reentrant checkpointing + gradient_checkpointing_kwargs = { + "use_reentrant": False + }, # use reentrant checkpointing max_grad_norm = 0.3, # max gradient norm based on QLoRA paper warmup_ratio = 0.03, # num_train_epochs = 2, # Set this instead of max_steps for full training runs diff --git a/tests/saving/vision_models/test_save_merge_vision_model_ocr_benchmark.py b/tests/saving/vision_models/test_save_merge_vision_model_ocr_benchmark.py index 66064600ea..8351cb8282 100644 --- a/tests/saving/vision_models/test_save_merge_vision_model_ocr_benchmark.py +++ b/tests/saving/vision_models/test_save_merge_vision_model_ocr_benchmark.py @@ -124,7 +124,9 @@ trainer = SFTTrainer( per_device_train_batch_size = 2, gradient_accumulation_steps = 4, gradient_checkpointing = True, - gradient_checkpointing_kwargs = {"use_reentrant": False}, # use reentrant checkpointing + gradient_checkpointing_kwargs = { + "use_reentrant": False + }, # use reentrant checkpointing max_grad_norm = 0.3, # max gradient norm based on QLoRA paper warmup_ratio = 0.03, # num_train_epochs = 2, # Set this instead of max_steps for full training runs diff --git a/tests/security/test_lockfile_supply_chain_audit.py b/tests/security/test_lockfile_supply_chain_audit.py index 905c4afcac..1e5d950a45 100644 --- a/tests/security/test_lockfile_supply_chain_audit.py +++ b/tests/security/test_lockfile_supply_chain_audit.py @@ -254,7 +254,9 @@ def test_advisory_finding_emitted_as_single_line_annotation(tmp_path): npm_lockfiles = [FIXTURES / "clean_lockfile.json"], cargo_lockfiles = [lockfile], ) - warning_lines = [line for line in proc.stderr.splitlines() if line.startswith("::warning::")] + warning_lines = [ + line for line in proc.stderr.splitlines() if line.startswith("::warning::") + ] assert warning_lines, ( "expected at least one ::warning:: annotation; " f"stderr was:\n{proc.stderr}" ) diff --git a/tests/security/test_new_install_scripts.py b/tests/security/test_new_install_scripts.py index e71659dd88..e27d778651 100644 --- a/tests/security/test_new_install_scripts.py +++ b/tests/security/test_new_install_scripts.py @@ -104,7 +104,9 @@ def test_new_dep_with_postinstall_exits_1(tmp_path: Path): head_pkgs = dict(base_pkgs) head_pkgs["node_modules/evil-postinstall"] = { "version": "1.0.0", - "resolved": ("https://registry.npmjs.org/evil-postinstall/-/evil-postinstall-1.0.0.tgz"), + "resolved": ( + "https://registry.npmjs.org/evil-postinstall/-/evil-postinstall-1.0.0.tgz" + ), "integrity": "sha512-fake", "hasInstallScript": True, } @@ -184,7 +186,8 @@ def test_v2_v3_lockfile_format_support(tmp_path: Path): head = _write(tmp_path / "head.json", _v2_lockfile(head_pkgs, head_deps)) result = _run(base, head) assert result.returncode == 1, ( - f"expected exit 1 for v2 lockfile, got {result.returncode}; " f"stderr:\n{result.stderr}" + f"expected exit 1 for v2 lockfile, got {result.returncode}; " + f"stderr:\n{result.stderr}" ) assert "v2-postinstall-dep" in result.stderr diff --git a/tests/security/test_scan_npm_packages.py b/tests/security/test_scan_npm_packages.py index 2c79b9c3ed..bf04847b13 100644 --- a/tests/security/test_scan_npm_packages.py +++ b/tests/security/test_scan_npm_packages.py @@ -101,10 +101,16 @@ def test_blocked_npm_versions_complete(): table = snp.BLOCKED_NPM_VERSIONS tanstack_keys = [k for k in table if k.startswith("@tanstack/")] assert len(tanstack_keys) == 42, ( - f"expected 42 @tanstack/* entries, got {len(tanstack_keys)}: " f"{sorted(tanstack_keys)}" + f"expected 42 @tanstack/* entries, got {len(tanstack_keys)}: " + f"{sorted(tanstack_keys)}" ) assert "@opensearch-project/opensearch" in table - assert table["@opensearch-project/opensearch"] == {"3.5.3", "3.6.2", "3.7.0", "3.8.0"} + assert table["@opensearch-project/opensearch"] == { + "3.5.3", + "3.6.2", + "3.7.0", + "3.8.0", + } squawk = [k for k in table if k.startswith("@squawk/")] assert len(squawk) >= 22, ( f"expected at least 22 @squawk/* entries (full safedep.io enumeration), " diff --git a/tests/security/test_scan_packages.py b/tests/security/test_scan_packages.py index 5b6e115213..a6301c749e 100644 --- a/tests/security/test_scan_packages.py +++ b/tests/security/test_scan_packages.py @@ -116,7 +116,9 @@ def test_clean_wheel_no_findings(): str(FIXTURES / "clean_wheel.whl"), "clean_fixture", ) - assert findings == [], f"unexpected findings on clean wheel: {[str(f) for f in findings]}" + assert ( + findings == [] + ), f"unexpected findings on clean wheel: {[str(f) for f in findings]}" # --------------------------------------------------------------------------- @@ -231,7 +233,8 @@ def test_archive_corruption_produces_critical_finding(tmp_path): assert findings, "scan_archive returned 0 findings on corrupt wheel" corrupted = [f for f in findings if f.check == "archive_corrupted"] assert corrupted, ( - "no archive_corrupted finding; got " f"{[(f.severity, f.check) for f in findings]}" + "no archive_corrupted finding; got " + f"{[(f.severity, f.check) for f in findings]}" ) assert all(f.severity == sp.CRITICAL for f in corrupted) diff --git a/tests/studio/_playwright_robust.py b/tests/studio/_playwright_robust.py index bc0e3f783b..7fe94df613 100644 --- a/tests/studio/_playwright_robust.py +++ b/tests/studio/_playwright_robust.py @@ -149,7 +149,9 @@ def wait_for_health( # but accept any 200 -- different Studio builds report differently. if status == 200: if info is not None: - info(f"health pre-flight OK: status=200, body keys={list((body or {}).keys())}") + info( + f"health pre-flight OK: status=200, body keys={list((body or {}).keys())}" + ) return True time.sleep(0.5) if info is not None: @@ -188,7 +190,9 @@ def recover_or_replace_page( info(f"recovery: page.is_closed() check failed: {exc!r}") if goto_url is not None: try: - page.goto(goto_url, wait_until = "domcontentloaded", timeout = default_timeout_ms) + page.goto( + goto_url, wait_until = "domcontentloaded", timeout = default_timeout_ms + ) if settle_networkidle: try: page.wait_for_load_state("networkidle", timeout = 30_000) diff --git a/tests/studio/install/smoke_test_llama_prebuilt.py b/tests/studio/install/smoke_test_llama_prebuilt.py index f7fd58aaa4..d87537dc94 100644 --- a/tests/studio/install/smoke_test_llama_prebuilt.py +++ b/tests/studio/install/smoke_test_llama_prebuilt.py @@ -15,7 +15,9 @@ INSTALLER_PATH = PACKAGE_ROOT / "studio" / "install_llama_prebuilt.py" def load_installer_module(): - spec = importlib.util.spec_from_file_location("studio_install_llama_prebuilt", INSTALLER_PATH) + spec = importlib.util.spec_from_file_location( + "studio_install_llama_prebuilt", INSTALLER_PATH + ) if spec is None or spec.loader is None: raise RuntimeError(f"unable to load installer module from {INSTALLER_PATH}") module = importlib.util.module_from_spec(spec) @@ -110,13 +112,17 @@ def main() -> int: published_release_tag = args.published_release_tag, ) print(f"[smoke] PASS install_dir={install_dir}") - print("[smoke] note=This was a real prebuilt install into an isolated temp directory.") + print( + "[smoke] note=This was a real prebuilt install into an isolated temp directory." + ) return installer.EXIT_SUCCESS except SystemExit as exc: code = int(exc.code) if isinstance(exc.code, int) else installer.EXIT_ERROR if code == installer.EXIT_FALLBACK: print(f"[smoke] FALLBACK install_dir={install_dir}") - print("[smoke] note=Prebuilt path failed and would fall back to source build in setup.") + print( + "[smoke] note=Prebuilt path failed and would fall back to source build in setup." + ) print(installer.collect_system_report(host, choice, install_dir)) else: print(f"[smoke] ERROR exit_code={code} install_dir={install_dir}") diff --git a/tests/studio/install/smoke_test_parallel_studio_home.py b/tests/studio/install/smoke_test_parallel_studio_home.py index be01147740..318a39dc74 100644 --- a/tests/studio/install/smoke_test_parallel_studio_home.py +++ b/tests/studio/install/smoke_test_parallel_studio_home.py @@ -65,7 +65,12 @@ def _free_port() -> int: def _run_one_install( - label: str, repo: Path, studio_home: Path, fake_home: Path, uv_cache: Path, log_path: Path + label: str, + repo: Path, + studio_home: Path, + fake_home: Path, + uv_cache: Path, + log_path: Path, ) -> tuple[str, int]: studio_home.mkdir(parents = True, exist_ok = True) fake_home.mkdir(parents = True, exist_ok = True) @@ -133,7 +138,9 @@ def _wait_for_health(port: int, timeout: float) -> dict: except (urllib.error.URLError, ConnectionError, OSError) as e: last_err = e time.sleep(HEALTH_POLL_INTERVAL_S) - raise TestFailure(f"port {port}: /api/health never returned 200 (last_err={last_err})") + raise TestFailure( + f"port {port}: /api/health never returned 200 (last_err={last_err})" + ) def _http_status( @@ -187,7 +194,9 @@ def _check_install_layout(label: str, studio_home: Path) -> dict: raise TestFailure(f"[{label}] launch-studio.sh kept @@DATA_DIR@@ placeholder") expected_data_dir_line = f"DATA_DIR='{studio_home}/share'" if expected_data_dir_line not in launcher: - raise TestFailure(f"[{label}] launch-studio.sh missing {expected_data_dir_line!r}") + raise TestFailure( + f"[{label}] launch-studio.sh missing {expected_data_dir_line!r}" + ) return {"label": label, "studio_home": str(studio_home), "install_id": install_id} @@ -204,7 +213,9 @@ def _check_fake_home_clean(fake_home: Path) -> None: ] leaked = [str(p) for p in forbidden if (fake_home / p).exists()] if leaked: - raise TestFailure(f"redirected HOME picked up persistent install pollution: {leaked}") + raise TestFailure( + f"redirected HOME picked up persistent install pollution: {leaked}" + ) def _backend_pid_python(pid: int) -> Path | None: @@ -228,7 +239,9 @@ def run(n_installs: int, keep: bool) -> int: repo = PACKAGE_ROOT if not (repo / "install.sh").is_file(): - raise TestFailure(f"install.sh not found at {repo}; run from a clone of unslothai/unsloth") + raise TestFailure( + f"install.sh not found at {repo}; run from a clone of unslothai/unsloth" + ) test_root = Path(tempfile.mkdtemp(prefix = "unsloth_studio_clash_")) _log(f"test root: {test_root}") @@ -316,7 +329,8 @@ def run(n_installs: int, keep: bool) -> int: raise TestFailure(f"[{label}] chat_only is not true under --no-torch") if health["studio_root_id"] in seen_root_ids: raise TestFailure( - f"[{label}] studio_root_id collision at runtime: " f"{health['studio_root_id']}" + f"[{label}] studio_root_id collision at runtime: " + f"{health['studio_root_id']}" ) seen_root_ids.add(health["studio_root_id"]) @@ -327,7 +341,9 @@ def run(n_installs: int, keep: bool) -> int: exe = _backend_pid_python(proc.pid) if exe is not None: - expected_python = (studio_home / "unsloth_studio" / "bin" / "python").resolve() + expected_python = ( + studio_home / "unsloth_studio" / "bin" / "python" + ).resolve() if exe != expected_python: raise TestFailure( f"[{label}] PID {proc.pid} exe={exe}, expected {expected_python}" @@ -337,7 +353,10 @@ def run(n_installs: int, keep: bool) -> int: if len(versions) != 1: raise TestFailure(f"version mismatch across installs: {versions}") - _log(f"PASS: all install + runtime invariants hold " f"(version={next(iter(versions))})") + _log( + f"PASS: all install + runtime invariants hold " + f"(version={next(iter(versions))})" + ) return 0 except TestFailure as e: diff --git a/tests/studio/install/test_cuda_repair.py b/tests/studio/install/test_cuda_repair.py index 83c2d962e8..ac3a5a7ff2 100644 --- a/tests/studio/install/test_cuda_repair.py +++ b/tests/studio/install/test_cuda_repair.py @@ -19,7 +19,9 @@ import pytest PACKAGE_ROOT = Path(__file__).resolve().parents[3] _STACK_PATH = PACKAGE_ROOT / "studio" / "install_python_stack.py" -_STACK_SPEC = importlib.util.spec_from_file_location("studio_install_python_stack", _STACK_PATH) +_STACK_SPEC = importlib.util.spec_from_file_location( + "studio_install_python_stack", _STACK_PATH +) assert _STACK_SPEC is not None and _STACK_SPEC.loader is not None stack_mod = importlib.util.module_from_spec(_STACK_SPEC) sys.modules[_STACK_SPEC.name] = stack_mod @@ -54,7 +56,9 @@ def _make_run( return result # nvidia-smi version probe (text = True) result.returncode = smi_rc - out = f"CUDA Version: {cuda_version}\n" if cuda_version else "No devices found\n" + out = ( + f"CUDA Version: {cuda_version}\n" if cuda_version else "No devices found\n" + ) result.stdout = out if kwargs.get("text") else out.encode() return result diff --git a/tests/studio/install/test_gpu_detection_followups.py b/tests/studio/install/test_gpu_detection_followups.py index 983969a4d3..3386dda3a6 100644 --- a/tests/studio/install/test_gpu_detection_followups.py +++ b/tests/studio/install/test_gpu_detection_followups.py @@ -91,9 +91,15 @@ def _run_detect_host( patches = [ patch.object(prebuilt_mod.platform, "system", return_value = system), patch.object(prebuilt_mod.platform, "machine", return_value = machine), - patch.object(prebuilt_mod.platform, "mac_ver", return_value = ("", ("", "", ""), "")), - patch.object(prebuilt_mod.shutil, "which", side_effect = lambda n: which_map.get(n)), - patch.object(prebuilt_mod, "run_capture", side_effect = _make_run_capture(rocminfo_stdout)), + patch.object( + prebuilt_mod.platform, "mac_ver", return_value = ("", ("", "", ""), "") + ), + patch.object( + prebuilt_mod.shutil, "which", side_effect = lambda n: which_map.get(n) + ), + patch.object( + prebuilt_mod, "run_capture", side_effect = _make_run_capture(rocminfo_stdout) + ), patch.object(prebuilt_mod.os.path, "isdir", side_effect = fake_isdir), patch.object(prebuilt_mod.os, "listdir", side_effect = fake_listdir), patch.object(prebuilt_mod.os, "access", return_value = False), @@ -326,7 +332,9 @@ class TestBackendExportLeafClassification: out = sp.run( ["sh", str(script), url], capture_output = True, text = True, timeout = 30 ).stdout.strip() - assert out == expected, f"{url} classified as {out!r}, expected {expected!r}" + assert ( + out == expected + ), f"{url} classified as {out!r}, expected {expected!r}" # TEST: CUDA_VISIBLE_DEVICES=""/-1 hides NVIDIA in every usable-GPU helper @@ -414,7 +422,9 @@ class TestHiddenCvdNotUsable: with ( patch.object(stack_mod.shutil, "which", side_effect = which_map.get), patch.object(stack_mod.subprocess, "run", side_effect = fake_run), - patch.dict(stack_mod.os.environ, {"CUDA_VISIBLE_DEVICES": "-1"}, clear = False), + patch.dict( + stack_mod.os.environ, {"CUDA_VISIBLE_DEVICES": "-1"}, clear = False + ), ): assert stack_mod._has_rocm_gpu() is True @@ -474,7 +484,11 @@ class TestHiddenCvdNotUsable: out = self._run_sh_helper( tmp_path, src, - ["_setup_run_smi", "_setup_cvd_hides_nvidia", "_setup_has_usable_nvidia_gpu"], + [ + "_setup_run_smi", + "_setup_cvd_hides_nvidia", + "_setup_has_usable_nvidia_gpu", + ], cvd, ) assert out == expected diff --git a/tests/studio/install/test_hf_auth.py b/tests/studio/install/test_hf_auth.py index 7c3296d4fd..b516b6641a 100644 --- a/tests/studio/install/test_hf_auth.py +++ b/tests/studio/install/test_hf_auth.py @@ -28,7 +28,9 @@ sys.modules[_SPEC.name] = mod _SPEC.loader.exec_module(mod) _TOKEN_VARS = ("GH_TOKEN", "GITHUB_TOKEN", "HF_TOKEN", "HUGGING_FACE_HUB_TOKEN") -HF_URL = "https://huggingface.co/ggml-org/models/resolve/main/tinyllamas/stories260K.gguf" +HF_URL = ( + "https://huggingface.co/ggml-org/models/resolve/main/tinyllamas/stories260K.gguf" +) GH_URL = "https://api.github.com/repos/unslothai/llama.cpp/releases" diff --git a/tests/studio/install/test_install_llama_prebuilt_logic.py b/tests/studio/install/test_install_llama_prebuilt_logic.py index 9ccfe15a88..6b619cb392 100644 --- a/tests/studio/install/test_install_llama_prebuilt_logic.py +++ b/tests/studio/install/test_install_llama_prebuilt_logic.py @@ -13,7 +13,9 @@ import pytest PACKAGE_ROOT = Path(__file__).resolve().parents[3] MODULE_PATH = PACKAGE_ROOT / "studio" / "install_llama_prebuilt.py" -SPEC = importlib.util.spec_from_file_location("studio_install_llama_prebuilt", MODULE_PATH) +SPEC = importlib.util.spec_from_file_location( + "studio_install_llama_prebuilt", MODULE_PATH +) assert SPEC is not None and SPEC.loader is not None INSTALL_LLAMA_PREBUILT = importlib.util.module_from_spec(SPEC) sys.modules[SPEC.name] = INSTALL_LLAMA_PREBUILT @@ -211,7 +213,9 @@ def test_hydrate_source_tree_extracts_upstream_archive_contents( def test_release_asset_download_url(): fn = INSTALL_LLAMA_PREBUILT.release_asset_download_url assert fn( - "unslothai/llama.cpp", "b9000-mix-abc1234", "llama.cpp-source-commit-deadbeef.tar.gz" + "unslothai/llama.cpp", + "b9000-mix-abc1234", + "llama.cpp-source-commit-deadbeef.tar.gz", ) == ( "https://github.com/unslothai/llama.cpp/releases/download/" "b9000-mix-abc1234/llama.cpp-source-commit-deadbeef.tar.gz" @@ -225,14 +229,18 @@ def test_release_asset_download_url(): def _mk_source_tarball(path: Path, tag: str) -> None: with tarfile.open(path, "w:gz") as archive: add_bytes_to_tar( - archive, f"llama.cpp-{tag}/CMakeLists.txt", b"cmake_minimum_required(VERSION 3.14)\n" + archive, + f"llama.cpp-{tag}/CMakeLists.txt", + b"cmake_minimum_required(VERSION 3.14)\n", ) add_bytes_to_tar( archive, f"llama.cpp-{tag}/convert_hf_to_gguf.py", b"#!/usr/bin/env python3\nimport gguf\n", ) - add_bytes_to_tar(archive, f"llama.cpp-{tag}/gguf-py/gguf/__init__.py", b"__all__ = []\n") + add_bytes_to_tar( + archive, f"llama.cpp-{tag}/gguf-py/gguf/__init__.py", b"__all__ = []\n" + ) def test_hydrate_source_tree_prefers_release_asset_for_mix( @@ -244,7 +252,9 @@ def test_hydrate_source_tree_prefers_release_asset_for_mix( archive_path = tmp_path / "merged-source.tar.gz" _mk_source_tarball(archive_path, f"b9000-mix-{commit[:7]}") asset_url = INSTALL_LLAMA_PREBUILT.release_asset_download_url( - "unslothai/llama.cpp", "b9000-mix-abc1234", f"llama.cpp-source-commit-{commit}.tar.gz" + "unslothai/llama.cpp", + "b9000-mix-abc1234", + f"llama.cpp-source-commit-{commit}.tar.gz", ) codeload_urls = set( INSTALL_LLAMA_PREBUILT.commit_source_archive_urls("unslothai/llama.cpp", commit) @@ -254,7 +264,9 @@ def test_hydrate_source_tree_prefers_release_asset_for_mix( def fake_download_file(url: str, destination: Path) -> None: seen.append(url) if url in codeload_urls: - raise AssertionError("codeload was hit even though the release asset was available") + raise AssertionError( + "codeload was hit even though the release asset was available" + ) assert url == asset_url destination.write_bytes(archive_path.read_bytes()) @@ -287,7 +299,9 @@ def test_hydrate_source_tree_falls_back_to_codeload_when_asset_missing( asset_url = INSTALL_LLAMA_PREBUILT.release_asset_download_url( "unslothai/llama.cpp", "b9000", f"llama.cpp-source-commit-{commit}.tar.gz" ) - codeload_urls = INSTALL_LLAMA_PREBUILT.commit_source_archive_urls("unslothai/llama.cpp", commit) + codeload_urls = INSTALL_LLAMA_PREBUILT.commit_source_archive_urls( + "unslothai/llama.cpp", commit + ) def fake_download_file(url: str, destination: Path) -> None: if url == asset_url: @@ -378,8 +392,12 @@ def test_validate_prebuilt_choice_creates_repo_shaped_linux_install( "preflight_linux_installed_binaries", lambda *args, **kwargs: None, ) - monkeypatch.setattr(INSTALL_LLAMA_PREBUILT, "validate_quantize", lambda *args, **kwargs: None) - monkeypatch.setattr(INSTALL_LLAMA_PREBUILT, "validate_server", lambda *args, **kwargs: None) + monkeypatch.setattr( + INSTALL_LLAMA_PREBUILT, "validate_quantize", lambda *args, **kwargs: None + ) + monkeypatch.setattr( + INSTALL_LLAMA_PREBUILT, "validate_server", lambda *args, **kwargs: None + ) host = HostInfo( system = "Linux", @@ -495,8 +513,12 @@ def test_validate_prebuilt_choice_creates_repo_shaped_windows_install( "preflight_linux_installed_binaries", lambda *args, **kwargs: None, ) - monkeypatch.setattr(INSTALL_LLAMA_PREBUILT, "validate_quantize", lambda *args, **kwargs: None) - monkeypatch.setattr(INSTALL_LLAMA_PREBUILT, "validate_server", lambda *args, **kwargs: None) + monkeypatch.setattr( + INSTALL_LLAMA_PREBUILT, "validate_quantize", lambda *args, **kwargs: None + ) + monkeypatch.setattr( + INSTALL_LLAMA_PREBUILT, "validate_server", lambda *args, **kwargs: None + ) host = HostInfo( system = "Windows", @@ -587,7 +609,9 @@ def test_activate_install_tree_restores_existing_install_after_activation_failur monkeypatch.setattr( INSTALL_LLAMA_PREBUILT, "confirm_install_tree", - lambda *_args, **_kwargs: (_ for _ in ()).throw(RuntimeError("activation confirm failed")), + lambda *_args, **_kwargs: (_ for _ in ()).throw( + RuntimeError("activation confirm failed") + ), ) with pytest.raises( @@ -636,7 +660,9 @@ def test_activate_install_tree_cleans_all_paths_when_rollback_restore_fails( monkeypatch.setattr( INSTALL_LLAMA_PREBUILT, "confirm_install_tree", - lambda *_args, **_kwargs: (_ for _ in ()).throw(RuntimeError("activation confirm failed")), + lambda *_args, **_kwargs: (_ for _ in ()).throw( + RuntimeError("activation confirm failed") + ), ) original_replace = INSTALL_LLAMA_PREBUILT.os.replace @@ -663,7 +689,10 @@ def test_activate_install_tree_cleans_all_paths_when_rollback_restore_fails( captured = capsys.readouterr() output = captured.out + captured.err assert "rollback after failed activation also failed: restore failed" in output - assert "cleaning staging, install, and rollback paths before source build fallback" in output + assert ( + "cleaning staging, install, and rollback paths before source build fallback" + in output + ) assert "removing failed install path" in output assert "removing rollback path" in output @@ -889,7 +918,9 @@ def write_linux_install_shape(install_dir: Path) -> None: (runtime_dir / "libggml-base.so.0").write_bytes(b"DLL") (runtime_dir / "libggml-cpu-x64.so.0").write_bytes(b"DLL") (runtime_dir / "libmtmd.so.0").write_bytes(b"DLL") - (install_dir / "convert_hf_to_gguf.py").write_text("#!/usr/bin/env python3\n", encoding = "utf-8") + (install_dir / "convert_hf_to_gguf.py").write_text( + "#!/usr/bin/env python3\n", encoding = "utf-8" + ) (install_dir / "gguf-py" / "gguf").mkdir(parents = True, exist_ok = True) @@ -913,7 +944,9 @@ def write_windows_install_shape( (runtime_dir / "cudart64_12.dll").write_bytes(b"DLL") (runtime_dir / "cublas64_12.dll").write_bytes(b"DLL") (runtime_dir / "cublasLt64_12.dll").write_bytes(b"DLL") - (install_dir / "convert_hf_to_gguf.py").write_text("#!/usr/bin/env python3\n", encoding = "utf-8") + (install_dir / "convert_hf_to_gguf.py").write_text( + "#!/usr/bin/env python3\n", encoding = "utf-8" + ) (install_dir / "gguf-py" / "gguf").mkdir(parents = True, exist_ok = True) @@ -936,7 +969,9 @@ def write_macos_install_shape( (runtime_dir / "libggml.0.dylib").write_bytes(b"DLL") if include_libmtmd: (runtime_dir / "libmtmd.0.dylib").write_bytes(b"DLL") - (install_dir / "convert_hf_to_gguf.py").write_text("#!/usr/bin/env python3\n", encoding = "utf-8") + (install_dir / "convert_hf_to_gguf.py").write_text( + "#!/usr/bin/env python3\n", encoding = "utf-8" + ) (install_dir / "gguf-py" / "gguf").mkdir(parents = True, exist_ok = True) @@ -1015,7 +1050,8 @@ def test_existing_install_matches_plan_false_without_fingerprint(tmp_path: Path) install_dir.mkdir() write_linux_install_shape(install_dir) (install_dir / "UNSLOTH_PREBUILT_INFO.json").write_text( - json.dumps({"tag": "b9001", "asset": "llama-b9001-bin-ubuntu-x64.tar.gz"}) + "\n", + json.dumps({"tag": "b9001", "asset": "llama-b9001-bin-ubuntu-x64.tar.gz"}) + + "\n", encoding = "utf-8", ) @@ -1078,7 +1114,9 @@ def test_existing_install_matches_plan_false_with_malformed_metadata(tmp_path: P install_dir = tmp_path / "llama.cpp" install_dir.mkdir() write_linux_install_shape(install_dir) - (install_dir / "UNSLOTH_PREBUILT_INFO.json").write_text("{not-json\n", encoding = "utf-8") + (install_dir / "UNSLOTH_PREBUILT_INFO.json").write_text( + "{not-json\n", encoding = "utf-8" + ) host = HostInfo( system = "Linux", @@ -1209,7 +1247,9 @@ def test_existing_install_matches_plan_windows_cpu_requires_llama_dll(tmp_path: def test_existing_install_matches_plan_windows_cuda_requires_cuda_dll(tmp_path: Path): install_dir = tmp_path / "llama.cpp" install_dir.mkdir() - write_windows_install_shape(install_dir, include_llama_dll = True, include_cuda_dll = True) + write_windows_install_shape( + install_dir, include_llama_dll = True, include_cuda_dll = True + ) host = HostInfo( system = "Windows", @@ -1278,7 +1318,9 @@ def test_existing_install_matches_plan_windows_cuda_requires_cuda_dll(tmp_path: assert existing_install_matches_plan(install_dir, host, plan) is False -def test_existing_install_matches_plan_windows_cuda_paired_requires_cudart(tmp_path: Path): +def test_existing_install_matches_plan_windows_cuda_paired_requires_cudart( + tmp_path: Path, +): """When the choice ships a paired cudart bundle (#5106), the install is considered stale unless cudart64_*.dll and cublas64_*.dll are actually on disk. Otherwise existing broken installs would keep @@ -1394,7 +1436,9 @@ def test_existing_install_matches_plan_windows_cuda_paired_requires_cudart(tmp_p assert existing_install_matches_plan(install_dir, host, plan) is False -def test_existing_install_matches_plan_windows_cuda_unpaired_skips_cudart_check(tmp_path: Path): +def test_existing_install_matches_plan_windows_cuda_unpaired_skips_cudart_check( + tmp_path: Path, +): """If the choice has no paired runtime archive (manifest dropped it, or upstream did not ship cudart), legacy installs without cudart on disk must still pass the health check -- otherwise the installer @@ -1749,7 +1793,9 @@ def test_install_prebuilt_skips_download_when_existing_install_matches( INSTALL_LLAMA_PREBUILT, "download_validation_model", lambda *args, **kwargs: (_ for _ in ()).throw( - AssertionError("matching install should skip before validation model download") + AssertionError( + "matching install should skip before validation model download" + ) ), ) @@ -2265,7 +2311,9 @@ def test_install_prebuilt_same_tag_upstream_failure_uses_older_unsloth_release_p (staging_dir / "marker.txt").write_text("ready\n") return attempts[0], staging_dir, initial_fallback_used - monkeypatch.setattr(INSTALL_LLAMA_PREBUILT, "validate_prebuilt_attempts", fake_validate) + monkeypatch.setattr( + INSTALL_LLAMA_PREBUILT, "validate_prebuilt_attempts", fake_validate + ) activated = {} monkeypatch.setattr( @@ -2283,7 +2331,10 @@ def test_install_prebuilt_same_tag_upstream_failure_uses_older_unsloth_release_p install_prebuilt(install_dir, "latest", "unslothai/llama.cpp", "") - assert attempted == [("b9002", "release-2", "upstream"), ("b9001", "release-1", "upstream")] + assert attempted == [ + ("b9002", "release-2", "upstream"), + ("b9001", "release-1", "upstream"), + ] assert activated["install_dir"] == install_dir @@ -2311,7 +2362,9 @@ def add_symlink_to_tar(archive: tarfile.TarFile, name: str, target: str) -> None archive.addfile(info) -def test_existing_install_matches_choice_fails_when_install_tree_incomplete(tmp_path: Path): +def test_existing_install_matches_choice_fails_when_install_tree_incomplete( + tmp_path: Path, +): """confirm_install_tree guard rejects installs missing critical files.""" install_dir = tmp_path / "llama.cpp" install_dir.mkdir() @@ -2400,7 +2453,9 @@ def test_existing_install_matches_choice_fails_when_install_tree_incomplete(tmp_ ) -def test_existing_install_matches_choice_fails_when_install_tree_incomplete_macos(tmp_path: Path): +def test_existing_install_matches_choice_fails_when_install_tree_incomplete_macos( + tmp_path: Path, +): """confirm_install_tree guard rejects macOS arm64 installs missing critical files.""" install_dir = tmp_path / "llama.cpp" install_dir.mkdir() @@ -2614,7 +2669,9 @@ def test_runtime_overlay_cannot_overwrite_main_archive_payload(tmp_path: Path) - if expected_sha256: actual = hashlib.sha256(Path(target_path).read_bytes()).hexdigest() if actual != expected_sha256: - raise INSTALL_LLAMA_PREBUILT.PrebuiltFallback(f"sha256 mismatch on {label}") + raise INSTALL_LLAMA_PREBUILT.PrebuiltFallback( + f"sha256 mismatch on {label}" + ) INSTALL_LLAMA_PREBUILT.download_file_verified = fake_download try: @@ -2626,7 +2683,8 @@ def test_runtime_overlay_cannot_overwrite_main_archive_payload(tmp_path: Path) - server = release_dir / "llama-server.exe" assert server.exists() assert server.read_bytes() == b"MAIN-SERVER", ( - "runtime archive overwrote main llama-server.exe; " f"got {server.read_bytes()!r}" + "runtime archive overwrote main llama-server.exe; " + f"got {server.read_bytes()!r}" ) for name in ("cudart64_12.dll", "cublas64_12.dll", "cublasLt64_12.dll"): assert (release_dir / name).exists(), f"missing {name}" @@ -2708,7 +2766,9 @@ def test_linux_runtime_overlay_copies_llama_tool_impl_libraries(tmp_path: Path) if expected_sha256: actual = hashlib.sha256(Path(target_path).read_bytes()).hexdigest() if actual != expected_sha256: - raise INSTALL_LLAMA_PREBUILT.PrebuiltFallback(f"sha256 mismatch on {label}") + raise INSTALL_LLAMA_PREBUILT.PrebuiltFallback( + f"sha256 mismatch on {label}" + ) INSTALL_LLAMA_PREBUILT.download_file_verified = fake_download try: @@ -2726,7 +2786,9 @@ def test_linux_runtime_overlay_copies_llama_tool_impl_libraries(tmp_path: Path) assert not (runtime_dir / "llama-cli").exists() -def test_python_runtime_dirs_covers_cu13_and_library_bin(monkeypatch, tmp_path: Path) -> None: +def test_python_runtime_dirs_covers_cu13_and_library_bin( + monkeypatch, tmp_path: Path +) -> None: """Installer-side runtime DLL discovery must scan the same path set as the backend ``_windows_pip_nvidia_dll_dirs``: legacy ``nvidia//bin``, current ``nvidia//bin/x86_64`` @@ -2793,7 +2855,9 @@ def _run_validate_prebuilt_choice(monkeypatch, tmp_path, *, expected_sha256): src, "preferred_source_archive", lambda *a, **k: ("repo", "ref", None, False) ) monkeypatch.setattr(src, "hydrate_source_tree", lambda *a, **k: None) - monkeypatch.setattr(src, "install_from_archives", lambda *a, **k: (server_path, quantize_path)) + monkeypatch.setattr( + src, "install_from_archives", lambda *a, **k: (server_path, quantize_path) + ) monkeypatch.setattr(src, "preflight_linux_installed_binaries", lambda *a, **k: None) monkeypatch.setattr(src, "preflight_macos_installed_binaries", lambda *a, **k: None) monkeypatch.setattr(src, "ensure_repo_shape", lambda *a, **k: None) @@ -2804,7 +2868,9 @@ def _run_validate_prebuilt_choice(monkeypatch, tmp_path, *, expected_sha256): lambda *a, **k: calls.__setitem__("quantize", calls["quantize"] + 1), ) monkeypatch.setattr( - src, "validate_server", lambda *a, **k: calls.__setitem__("server", calls["server"] + 1) + src, + "validate_server", + lambda *a, **k: calls.__setitem__("server", calls["server"] + 1), ) bundle_name = "app-b9998-linux-x64-cuda13-newer.tar.gz" @@ -2846,14 +2912,20 @@ def _run_validate_prebuilt_choice(monkeypatch, tmp_path, *, expected_sha256): return calls -def test_validate_prebuilt_choice_approved_validation_skipped_when_flag_off(tmp_path, monkeypatch): +def test_validate_prebuilt_choice_approved_validation_skipped_when_flag_off( + tmp_path, monkeypatch +): # An approved (sha256-verified) bundle skips the staged smoke test while the # flag is off: the manifest hash is its integrity gate. - calls = _run_validate_prebuilt_choice(monkeypatch, tmp_path, expected_sha256 = "ab" * 32) + calls = _run_validate_prebuilt_choice( + monkeypatch, tmp_path, expected_sha256 = "ab" * 32 + ) assert calls == {"quantize": 0, "server": 0} -def test_validate_prebuilt_choice_hashless_build_always_validated(tmp_path, monkeypatch): +def test_validate_prebuilt_choice_hashless_build_always_validated( + tmp_path, monkeypatch +): # A hashless external build has no approved sha256, so the # functional smoke test is its only integrity gate and must run even while the # flag is off -- otherwise a corrupted/replaced archive could be activated. @@ -2861,9 +2933,13 @@ def test_validate_prebuilt_choice_hashless_build_always_validated(tmp_path, monk assert calls == {"quantize": 1, "server": 1} -def test_validate_prebuilt_choice_approved_validation_runs_when_flag_enabled(tmp_path, monkeypatch): +def test_validate_prebuilt_choice_approved_validation_runs_when_flag_enabled( + tmp_path, monkeypatch +): # Flipping _RUN_STAGED_PREBUILT_VALIDATION back on restores the full smoke test # for approved bundles too, proving the check is kept intact, only gated off. monkeypatch.setattr(INSTALL_LLAMA_PREBUILT, "_RUN_STAGED_PREBUILT_VALIDATION", True) - calls = _run_validate_prebuilt_choice(monkeypatch, tmp_path, expected_sha256 = "ab" * 32) + calls = _run_validate_prebuilt_choice( + monkeypatch, tmp_path, expected_sha256 = "ab" * 32 + ) assert calls == {"quantize": 1, "server": 1} diff --git a/tests/studio/install/test_llama_pr_force_and_source.py b/tests/studio/install/test_llama_pr_force_and_source.py index 8643806127..10e20c3ed7 100644 --- a/tests/studio/install/test_llama_pr_force_and_source.py +++ b/tests/studio/install/test_llama_pr_force_and_source.py @@ -386,7 +386,10 @@ class TestSourcePatternsSh: assert '_DEFAULT_LLAMA_PR_FORCE=""' in self.content def test_has_default_source(self): - assert '_DEFAULT_LLAMA_SOURCE="https://github.com/ggml-org/llama.cpp"' in self.content + assert ( + '_DEFAULT_LLAMA_SOURCE="https://github.com/ggml-org/llama.cpp"' + in self.content + ) def test_has_pr_force_env_read(self): assert "UNSLOTH_LLAMA_PR_FORCE" in self.content @@ -436,7 +439,9 @@ class TestSourcePatternsSh: lines = self.content.splitlines() for i, line in enumerate(lines, 1): if "git clone" in line and "ggml-org/llama.cpp.git" in line: - pytest.fail(f"Line {i} has hardcoded ggml-org clone URL: {line.strip()}") + pytest.fail( + f"Line {i} has hardcoded ggml-org clone URL: {line.strip()}" + ) # ========================================================================= @@ -453,7 +458,10 @@ class TestSourcePatternsPs1: assert '$DefaultLlamaPrForce = ""' in self.content def test_has_default_source(self): - assert '$DefaultLlamaSource = "https://github.com/ggml-org/llama.cpp"' in self.content + assert ( + '$DefaultLlamaSource = "https://github.com/ggml-org/llama.cpp"' + in self.content + ) def test_has_pr_force_env_read(self): assert "$env:UNSLOTH_LLAMA_PR_FORCE" in self.content @@ -465,7 +473,10 @@ class TestSourcePatternsPs1: def test_release_repo_override_removed(self): # No env-based release-repo override; the repo is chosen by GPU detection # (GPU -> fork, CPU -> ggml-org), mirroring setup.sh. - assert "$HelperReleaseRepo = if ($env:UNSLOTH_LLAMA_RELEASE_REPO)" not in self.content + assert ( + "$HelperReleaseRepo = if ($env:UNSLOTH_LLAMA_RELEASE_REPO)" + not in self.content + ) assert ( "$HelperReleaseRepo = if ($HasNvidiaSmi -or $HasROCm -or $script:ROCmGfxArch) " '{ "unslothai/llama.cpp" } else { "ggml-org/llama.cpp" }' in self.content @@ -487,7 +498,9 @@ class TestSourcePatternsPs1: def test_clone_urls_parameterized_pr_path(self): """PR clone path uses $LlamaSource.git, not hardcoded URL.""" - pr_idx = self.content.index("if ($LlamaPr) {\n", self.content.index("Cloning llama.cpp")) + pr_idx = self.content.index( + "if ($LlamaPr) {\n", self.content.index("Cloning llama.cpp") + ) else_idx = self.content.index("} else {", pr_idx) pr_block = self.content[pr_idx:else_idx] assert '"$LlamaSource.git"' in pr_block @@ -505,7 +518,9 @@ class TestSourcePatternsPs1: lines = self.content.splitlines() for i, line in enumerate(lines, 1): if "git clone" in line and "ggml-org/llama.cpp.git" in line: - pytest.fail(f"Line {i} has hardcoded ggml-org clone URL: {line.strip()}") + pytest.fail( + f"Line {i} has hardcoded ggml-org clone URL: {line.strip()}" + ) # ========================================================================= diff --git a/tests/studio/install/test_macos_version_compat.py b/tests/studio/install/test_macos_version_compat.py index 7f93b295eb..2b96d6c5e4 100644 --- a/tests/studio/install/test_macos_version_compat.py +++ b/tests/studio/install/test_macos_version_compat.py @@ -20,7 +20,9 @@ import pytest PACKAGE_ROOT = Path(__file__).resolve().parents[3] MODULE_PATH = PACKAGE_ROOT / "studio" / "install_llama_prebuilt.py" -SPEC = importlib.util.spec_from_file_location("studio_install_llama_prebuilt_macos", MODULE_PATH) +SPEC = importlib.util.spec_from_file_location( + "studio_install_llama_prebuilt_macos", MODULE_PATH +) assert SPEC is not None and SPEC.loader is not None ILP = importlib.util.module_from_spec(SPEC) sys.modules[SPEC.name] = ILP @@ -142,7 +144,10 @@ class TestMachoMinimumMacos: ) ) assert ILP.macho_minimum_macos(path, make_macos_host((14, 0))) == (14, 0) - assert ILP.macho_minimum_macos(path, make_macos_host((26, 0), arm64 = False)) == (26, 0) + assert ILP.macho_minimum_macos(path, make_macos_host((26, 0), arm64 = False)) == ( + 26, + 0, + ) def test_non_macho_returns_none(self, tmp_path): path = tmp_path / "script.sh" @@ -185,17 +190,23 @@ class TestPreflightMacosInstalledBinaries: def test_rejects_too_new_dylib(self, tmp_path): install_dir, binaries = self._install_dir(tmp_path, (26, 0)) with pytest.raises(PrebuiltFallback, match = "newer macOS"): - ILP.preflight_macos_installed_binaries(binaries, install_dir, make_macos_host((14, 0))) + ILP.preflight_macos_installed_binaries( + binaries, install_dir, make_macos_host((14, 0)) + ) def test_accepts_compatible_prebuilt(self, tmp_path): install_dir, binaries = self._install_dir(tmp_path, (14, 0)) # Must not raise on a macOS 15 host. - ILP.preflight_macos_installed_binaries(binaries, install_dir, make_macos_host((15, 5))) + ILP.preflight_macos_installed_binaries( + binaries, install_dir, make_macos_host((15, 5)) + ) def test_skips_when_host_version_unknown(self, tmp_path): install_dir, binaries = self._install_dir(tmp_path, (26, 0)) # Unknown host version -> defer to runtime validation, do not raise. - ILP.preflight_macos_installed_binaries(binaries, install_dir, make_macos_host(None)) + ILP.preflight_macos_installed_binaries( + binaries, install_dir, make_macos_host(None) + ) def test_noop_on_non_macos_host(self, tmp_path): install_dir, binaries = self._install_dir(tmp_path, (26, 0)) diff --git a/tests/studio/install/test_pr4562_bugfixes.py b/tests/studio/install/test_pr4562_bugfixes.py index 9dfa4e0005..35e3151ff7 100644 --- a/tests/studio/install/test_pr4562_bugfixes.py +++ b/tests/studio/install/test_pr4562_bugfixes.py @@ -29,7 +29,9 @@ import pytest # --------------------------------------------------------------------------- PACKAGE_ROOT = Path(__file__).resolve().parents[3] MODULE_PATH = PACKAGE_ROOT / "studio" / "install_llama_prebuilt.py" -SPEC = importlib.util.spec_from_file_location("studio_install_llama_prebuilt", MODULE_PATH) +SPEC = importlib.util.spec_from_file_location( + "studio_install_llama_prebuilt", MODULE_PATH +) assert SPEC is not None and SPEC.loader is not None MOD = importlib.util.module_from_spec(SPEC) sys.modules[SPEC.name] = MOD @@ -116,7 +118,9 @@ class TestBinaryEnvCrossPlatform: env = binary_env(binary_path, install_dir, host) ld_dirs = env["LD_LIBRARY_PATH"].split(os.pathsep) assert str(bin_dir) in ld_dirs, f"build/bin not in LD_LIBRARY_PATH: {ld_dirs}" - assert str(install_dir) in ld_dirs, f"install_dir not in LD_LIBRARY_PATH: {ld_dirs}" + assert ( + str(install_dir) in ld_dirs + ), f"install_dir not in LD_LIBRARY_PATH: {ld_dirs}" def test_linux_binary_parent_comes_before_install_dir( self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch @@ -135,7 +139,9 @@ class TestBinaryEnvCrossPlatform: ld_dirs = env["LD_LIBRARY_PATH"].split(os.pathsep) bin_idx = ld_dirs.index(str(bin_dir)) install_idx = ld_dirs.index(str(install_dir)) - assert bin_idx < install_idx, "binary_path.parent should come before install_dir" + assert ( + bin_idx < install_idx + ), "binary_path.parent should come before install_dir" def test_linux_deduplicates_when_binary_parent_equals_install_dir( self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch @@ -194,13 +200,17 @@ class TestBinaryEnvCrossPlatform: binary_path.write_bytes(b"MZ") host = make_host(system = "Windows") - monkeypatch.setattr(MOD, "windows_runtime_dirs_for_runtime_line", lambda _rt: []) + monkeypatch.setattr( + MOD, "windows_runtime_dirs_for_runtime_line", lambda _rt: [] + ) env = binary_env(binary_path, install_dir, host) path_dirs = env["PATH"].split(os.pathsep) assert str(bin_dir) in path_dirs, f"build/bin/Release not in PATH: {path_dirs}" - def test_macos_sets_dyld_library_path(self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch): + def test_macos_sets_dyld_library_path( + self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch + ): install_dir = tmp_path / "llama.cpp" install_dir.mkdir(parents = True) bin_dir = install_dir / "build" / "bin" @@ -213,8 +223,12 @@ class TestBinaryEnvCrossPlatform: env = binary_env(binary_path, install_dir, host) dyld_parts = [p for p in env["DYLD_LIBRARY_PATH"].split(os.pathsep) if p] - assert str(bin_dir) in dyld_parts, f"build/bin not in DYLD_LIBRARY_PATH: {dyld_parts}" - assert str(install_dir) in dyld_parts, f"install_dir not in DYLD_LIBRARY_PATH: {dyld_parts}" + assert ( + str(bin_dir) in dyld_parts + ), f"build/bin not in DYLD_LIBRARY_PATH: {dyld_parts}" + assert ( + str(install_dir) in dyld_parts + ), f"install_dir not in DYLD_LIBRARY_PATH: {dyld_parts}" # binary_path.parent (build/bin) should come before install_dir assert dyld_parts.index(str(bin_dir)) < dyld_parts.index(str(install_dir)) @@ -345,7 +359,9 @@ class TestResolveRequestedLlamaTag: class TestFetchJsonRetries: - def test_fetch_json_retries_invalid_github_api_json(self, monkeypatch: pytest.MonkeyPatch): + def test_fetch_json_retries_invalid_github_api_json( + self, monkeypatch: pytest.MonkeyPatch + ): calls = {"count": 0} def fake_download_bytes(url, **kwargs): @@ -680,7 +696,10 @@ class TestSourceCodePatterns: content = SETUP_SH.read_text() assert "--resolve-source-build" not in content assert "--resolve-install-tag" not in content - assert '--resolve-llama-tag latest --published-repo "ggml-org/llama.cpp"' in content + assert ( + '--resolve-llama-tag latest --published-repo "ggml-org/llama.cpp"' + in content + ) assert "--output-format json" in content assert "_RESOLVED_SOURCE_URL" in content assert "_RESOLVED_SOURCE_REF_KIND" in content @@ -756,7 +775,9 @@ class TestSourceCodePatterns: # Delivered via NVCC_PREPEND_FLAGS (covers the configure-time compiler # probe too), not embedded in the word-split CMAKE_ARGS string. assert "export NVCC_PREPEND_FLAGS=" in content - cmake_args_lines = [line for line in content.splitlines() if "CMAKE_ARGS=" in line] + cmake_args_lines = [ + line for line in content.splitlines() if "CMAKE_ARGS=" in line + ] assert all( "-allow-unsupported-compiler" not in line for line in cmake_args_lines ), "flag must stay out of CMAKE_ARGS (bash word-splitting safety)" @@ -772,7 +793,9 @@ class TestSourceCodePatterns: # Delivered via the process environment, not the $CmakeArgs array, so it # reaches both the configure-time compiler probe and `cmake --build`. assert "$env:NVCC_PREPEND_FLAGS" in content - cmake_args_lines = [line for line in content.splitlines() if "$CmakeArgs +=" in line] + cmake_args_lines = [ + line for line in content.splitlines() if "$CmakeArgs +=" in line + ] assert all( "-allow-unsupported-compiler" not in line for line in cmake_args_lines ), "flag must not be pushed into the $CmakeArgs array" @@ -788,10 +811,15 @@ class TestSourceCodePatterns: def test_macos_arm64_cpu_fallback_args_exclude_rpath(self): """CPU fallback args must NOT contain Metal-only RPATH flags at runtime.""" - script = '_IS_MACOS_ARM64=true\nNVCC_PATH=""\nGPU_BACKEND=""\n' + _GPU_BACKEND_FRAGMENT + script = ( + '_IS_MACOS_ARM64=true\nNVCC_PATH=""\nGPU_BACKEND=""\n' + + _GPU_BACKEND_FRAGMENT + ) output = run_bash(script) fallback_line = next( - line for line in output.splitlines() if line.startswith("CPU_FALLBACK_CMAKE_ARGS=") + line + for line in output.splitlines() + if line.startswith("CPU_FALLBACK_CMAKE_ARGS=") ) assert "-DGGML_METAL=OFF" in fallback_line assert ( @@ -812,7 +840,8 @@ class TestSourceCodePatterns: assert ( "x86_64" not in content[ - content.find("-DGGML_METAL=ON") - 200 : content.find("-DGGML_METAL=ON") + 200 + content.find("-DGGML_METAL=ON") - 200 : content.find("-DGGML_METAL=ON") + + 200 ] ) @@ -842,7 +871,9 @@ class TestSourceCodePatterns: # Allow git pull in other contexts context = "\n".join(lines[max(0, i - 5) : i + 5]) if "LlamaCppDir" in context: - pytest.fail(f"Found 'git pull' in llama.cpp build section at line {i+1}") + pytest.fail( + f"Found 'git pull' in llama.cpp build section at line {i+1}" + ) def test_setup_ps1_prebuilt_install_entrypoint(self): """PS1 prebuilt path should call the helper install entrypoint, not the @@ -869,7 +900,8 @@ class TestSourceCodePatterns: assert "--resolve-source-build" not in content assert "--resolve-install-tag" not in content assert ( - '"--resolve-llama-tag", "latest", "--published-repo", "ggml-org/llama.cpp"' in content + '"--resolve-llama-tag", "latest", "--published-repo", "ggml-org/llama.cpp"' + in content ) assert '--output-format", "json"' in content assert "$ResolvedSourceUrl" in content @@ -883,7 +915,10 @@ class TestSourceCodePatterns: block = content[max(0, install_idx - 800) : install_idx + 800] assert "$PSNativeCommandUseErrorActionPreference = $false" in block assert "$restoreNativeErrorPreference = $true" in block - assert "$PSNativeCommandUseErrorActionPreference = $previousNativeErrorPreference" in block + assert ( + "$PSNativeCommandUseErrorActionPreference = $previousNativeErrorPreference" + in block + ) def test_setup_ps1_helper_disables_error_action_abort(self): """Helper resolution should suppress terminating NativeCommandError on PS 5.1.""" @@ -904,7 +939,9 @@ class TestSourceCodePatterns: """The unconstrained nvcc fallback should not sort toolkit dirs lexicographically.""" content = SETUP_PS1.read_text() assert "Sort-Object Name | Select-Object -Last 1" not in content - assert "Sort-Object { [version]($_.Name -replace '^v','') } -Descending" in content + assert ( + "Sort-Object { [version]($_.Name -replace '^v','') } -Descending" in content + ) def test_binary_env_linux_has_binary_parent(self): """The Linux branch of binary_env should include binary_path.parent.""" @@ -965,7 +1002,10 @@ class TestMacOSMetalBuildLogic: def test_macos_arm64_cmake_args_contain_metal_flags(self): """macOS arm64 should enable Metal, not CUDA.""" - script = '_IS_MACOS_ARM64=true\nNVCC_PATH=""\nGPU_BACKEND=""\n' + _GPU_BACKEND_FRAGMENT + script = ( + '_IS_MACOS_ARM64=true\nNVCC_PATH=""\nGPU_BACKEND=""\n' + + _GPU_BACKEND_FRAGMENT + ) output = run_bash(script) assert "-DGGML_METAL=ON" in output assert "-DGGML_CUDA=ON" not in output @@ -973,7 +1013,10 @@ class TestMacOSMetalBuildLogic: def test_intel_macos_no_metal_flags(self): """Intel macOS (not arm64) should not get Metal flags.""" - script = '_IS_MACOS_ARM64=false\nNVCC_PATH=""\nGPU_BACKEND=""\n' + _GPU_BACKEND_FRAGMENT + script = ( + '_IS_MACOS_ARM64=false\nNVCC_PATH=""\nGPU_BACKEND=""\n' + + _GPU_BACKEND_FRAGMENT + ) output = run_bash(script) assert "-DGGML_METAL=ON" not in output assert "BUILD_DESC=building (CPU)" in output @@ -1059,14 +1102,18 @@ class TestMacOSMetalBuildLogic: # Verify cmake args: first call has Metal ON, second has Metal OFF calls = calls_file.read_text().splitlines() assert len(calls) >= 2, f"Expected >= 2 cmake calls, got {len(calls)}" - assert "-DGGML_METAL=ON" in calls[0], f"First cmake call should have Metal ON: {calls[0]}" + assert ( + "-DGGML_METAL=ON" in calls[0] + ), f"First cmake call should have Metal ON: {calls[0]}" assert ( "-DGGML_METAL=OFF" in calls[1] ), f"Second cmake call should have Metal OFF: {calls[1]}" assert ( "-DGGML_METAL=ON" not in calls[1] ), f"Second cmake call should NOT have Metal ON: {calls[1]}" - assert "@loader_path" not in calls[1], f"CPU fallback should not have RPATH: {calls[1]}" + assert ( + "@loader_path" not in calls[1] + ), f"CPU fallback should not have RPATH: {calls[1]}" assert ( "-DCMAKE_BUILD_WITH_INSTALL_RPATH=ON" not in calls[1] ), f"CPU fallback should not have RPATH build flag: {calls[1]}" @@ -1174,7 +1221,9 @@ class TestMacOSMetalBuildLogic: # Third call: re-configure with Metal OFF and no RPATH flags assert "-DGGML_METAL=OFF" in calls[2] assert "-DGGML_METAL=ON" not in calls[2] - assert "@loader_path" not in calls[2], f"CPU fallback should not have RPATH: {calls[2]}" + assert ( + "@loader_path" not in calls[2] + ), f"CPU fallback should not have RPATH: {calls[2]}" assert ( "-DCMAKE_BUILD_WITH_INSTALL_RPATH=ON" not in calls[2] ), f"CPU fallback should not have RPATH build flag: {calls[2]}" diff --git a/tests/studio/install/test_pr5940_followups.py b/tests/studio/install/test_pr5940_followups.py index 583532eb74..02e49aae93 100644 --- a/tests/studio/install/test_pr5940_followups.py +++ b/tests/studio/install/test_pr5940_followups.py @@ -67,8 +67,12 @@ def test_fetch_validation_model_prefers_huggingface_hub(tmp_path): model.write_bytes(b"GGUF-via-hf") fake_hf = MagicMock(return_value = str(model)) with ( - patch.object(prebuilt, "validated_validation_model_bytes", side_effect = lambda b: b), - patch.dict(sys.modules, {"huggingface_hub": MagicMock(hf_hub_download = fake_hf)}), + patch.object( + prebuilt, "validated_validation_model_bytes", side_effect = lambda b: b + ), + patch.dict( + sys.modules, {"huggingface_hub": MagicMock(hf_hub_download = fake_hf)} + ), ): assert prebuilt._fetch_validation_model_bytes() == b"GGUF-via-hf" assert fake_hf.called # hf path was taken, urllib not needed @@ -77,8 +81,12 @@ def test_fetch_validation_model_prefers_huggingface_hub(tmp_path): def test_fetch_validation_model_falls_back_to_urllib_on_hf_failure(): fake_hf = MagicMock(side_effect = RuntimeError("hf unreachable")) with ( - patch.object(prebuilt, "validated_validation_model_bytes", side_effect = lambda b: b), - patch.dict(sys.modules, {"huggingface_hub": MagicMock(hf_hub_download = fake_hf)}), + patch.object( + prebuilt, "validated_validation_model_bytes", side_effect = lambda b: b + ), + patch.dict( + sys.modules, {"huggingface_hub": MagicMock(hf_hub_download = fake_hf)} + ), patch.object(prebuilt, "download_bytes", return_value = b"GGUF-via-urllib") as dl, ): assert prebuilt._fetch_validation_model_bytes() == b"GGUF-via-urllib" @@ -192,7 +200,9 @@ def test_install_sh_name_arch_agrees_with_ps_for_strix_and_non_amd(): assert sh == expect, f"install.sh: {name!r} -> {sh!r}, expected {expect!r}" if expect is not None: # cross-check bash agrees with the PowerShell table ps = next((a for p, a in ps_rows if re.search(p, name)), None) - assert sh == ps, f"install.sh/install.ps1 drift for {name!r}: {sh!r} vs {ps!r}" + assert ( + sh == ps + ), f"install.sh/install.ps1 drift for {name!r}: {sh!r} vs {ps!r}" def test_setup_sh_name_arch_table_in_sync_with_install_sh(): @@ -268,7 +278,9 @@ def test_amd_smi_opt_in_forces_on_windows_no_sdk(): def test_amd_smi_opt_out_overrides_hip_sdk(): assert ( - _amd_smi_allowed_under("Windows", hipinfo_present = True, env = {"UNSLOTH_ENABLE_AMD_SMI": "0"}) + _amd_smi_allowed_under( + "Windows", hipinfo_present = True, env = {"UNSLOTH_ENABLE_AMD_SMI": "0"} + ) is False ) @@ -278,7 +290,9 @@ def test_ps_installers_gate_amd_smi_on_windows(): # UNSLOTH_ENABLE_AMD_SMI opt-in, mirroring _amd_smi_allowed(). for ps in (_INSTALL_PS1, _SETUP_PS1): text = ps.read_text(encoding = "utf-8") - assert "UNSLOTH_ENABLE_AMD_SMI" in text, f"{ps.name} missing amd-smi opt-in gate" + assert ( + "UNSLOTH_ENABLE_AMD_SMI" in text + ), f"{ps.name} missing amd-smi opt-in gate" assert "amdSmiAllowed" in text, f"{ps.name} missing amd-smi gate variable" @@ -289,7 +303,9 @@ def test_install_python_stack_gates_every_amd_smi_spawn(): # list` ungated on Adrenalin-only hosts; not-spawning is the only fix. import ast - src = (PACKAGE_ROOT / "studio" / "install_python_stack.py").read_text(encoding = "utf-8") + src = (PACKAGE_ROOT / "studio" / "install_python_stack.py").read_text( + encoding = "utf-8" + ) tree = ast.parse(src) def _names_amd_smi_command(node): @@ -313,7 +329,10 @@ def test_install_python_stack_gates_every_amd_smi_spawn(): return False def _references_gate(node): - return any(isinstance(n, ast.Name) and n.id == "_amd_smi_allowed" for n in ast.walk(node)) + return any( + isinstance(n, ast.Name) and n.id == "_amd_smi_allowed" + for n in ast.walk(node) + ) offenders = [ node.name diff --git a/tests/studio/install/test_probe_timeouts.py b/tests/studio/install/test_probe_timeouts.py index acea0ed34d..e5d25af0f8 100644 --- a/tests/studio/install/test_probe_timeouts.py +++ b/tests/studio/install/test_probe_timeouts.py @@ -173,11 +173,15 @@ def test_has_usable_nvidia_gpu_returns_under_timeout(): fake_dir.mkdir() fake_smi = fake_dir / "nvidia-smi" fake_smi.write_text("#!/bin/sh\nsleep 30\n") - fake_smi.chmod(fake_smi.stat().st_mode | stat.S_IEXEC | stat.S_IXGRP | stat.S_IXOTH) + fake_smi.chmod( + fake_smi.stat().st_mode | stat.S_IEXEC | stat.S_IXGRP | stat.S_IXOTH + ) # Build a minimal PATH that includes the fake nvidia-smi plus the real # `timeout`/`awk`/`ls` it needs. Use the fake dir first so it wins. - real_bins = {Path(shutil.which(c)).parent for c in ("timeout", "awk", "ls", "sh")} + real_bins = { + Path(shutil.which(c)).parent for c in ("timeout", "awk", "ls", "sh") + } path_env = os.pathsep.join([str(fake_dir)] + [str(p) for p in real_bins]) # Force the /proc fallback off so the result depends only on the probe, diff --git a/tests/studio/install/test_rocm_support.py b/tests/studio/install/test_rocm_support.py index a3c9555d0b..4f748eaa99 100644 --- a/tests/studio/install/test_rocm_support.py +++ b/tests/studio/install/test_rocm_support.py @@ -40,7 +40,9 @@ _normalize_forwarded_gfx = prebuilt_mod._normalize_forwarded_gfx # install_python_stack.py _STACK_PATH = PACKAGE_ROOT / "studio" / "install_python_stack.py" -_STACK_SPEC = importlib.util.spec_from_file_location("studio_install_python_stack", _STACK_PATH) +_STACK_SPEC = importlib.util.spec_from_file_location( + "studio_install_python_stack", _STACK_PATH +) assert _STACK_SPEC is not None and _STACK_SPEC.loader is not None stack_mod = importlib.util.module_from_spec(_STACK_SPEC) sys.modules[_STACK_SPEC.name] = stack_mod @@ -300,7 +302,9 @@ class TestResolveUpstreamAssetChoice: def test_rocm_linux_no_prebuilt_falls_back(self, mock_assets): """AMD ROCm host should fall back to source build when no ROCm prebuilt exists.""" # Remove the ROCm asset from available assets - assets_without_rocm = {k: v for k, v in UPSTREAM_ASSETS.items() if "rocm" not in k} + assets_without_rocm = { + k: v for k, v in UPSTREAM_ASSETS.items() if "rocm" not in k + } mock_assets.return_value = assets_without_rocm host = rocm_host() with pytest.raises(PrebuiltFallback, match = "ROCm detected"): @@ -591,7 +595,9 @@ class TestEnsureRocmTorch: @patch.object(stack_mod, "_has_usable_nvidia_gpu", return_value = False) @patch.object(stack_mod, "_has_rocm_gpu", return_value = True) @patch.object(stack_mod, "_detect_rocm_version", return_value = (7, 1)) - def test_torch_already_has_cuda_skips(self, mock_ver, mock_gpu, mock_nvidia, mock_pip): + def test_torch_already_has_cuda_skips( + self, mock_ver, mock_gpu, mock_nvidia, mock_pip + ): """If torch already has CUDA, should skip ROCm reinstall.""" mock_probe = MagicMock() mock_probe.returncode = 0 @@ -605,7 +611,9 @@ class TestEnsureRocmTorch: @patch.object(stack_mod, "_has_usable_nvidia_gpu", return_value = False) @patch.object(stack_mod, "_has_rocm_gpu", return_value = True) @patch.object(stack_mod, "_detect_rocm_version", return_value = (7, 1)) - def test_torch_already_has_hip_skips(self, mock_ver, mock_gpu, mock_nvidia, mock_pip): + def test_torch_already_has_hip_skips( + self, mock_ver, mock_gpu, mock_nvidia, mock_pip + ): """If torch already has HIP, should skip ROCm reinstall.""" mock_probe = MagicMock() mock_probe.returncode = 0 @@ -643,7 +651,9 @@ class TestEnsureRocmTorch: @patch.object(stack_mod, "_has_usable_nvidia_gpu", return_value = False) @patch.object(stack_mod, "_has_rocm_gpu", return_value = True) @patch.object(stack_mod, "_detect_rocm_version", return_value = (6, 3)) - def test_rocm_63_selects_correct_tag(self, mock_ver, mock_gpu, mock_nvidia, mock_pip): + def test_rocm_63_selects_correct_tag( + self, mock_ver, mock_gpu, mock_nvidia, mock_pip + ): """ROCm 6.3 should select rocm6.3 tag.""" mock_probe = MagicMock() mock_probe.returncode = 0 @@ -710,7 +720,9 @@ class TestEnsureRocmTorch: ): """Probe subprocess timeout should not crash; should proceed to reinstall.""" with patch("os.path.isdir", return_value = True): - with patch("subprocess.run", side_effect = subprocess.TimeoutExpired("python", 30)): + with patch( + "subprocess.run", side_effect = subprocess.TimeoutExpired("python", 30) + ): _ensure_rocm_torch() # If probe times out, the function should treat torch as unusable and reinstall # both torch (via pip_install) and bitsandbytes (via pip_install_try). @@ -735,7 +747,9 @@ class TestEnsureRocmTorch: @patch.object(stack_mod, "pip_install") @patch.object(stack_mod, "_has_rocm_gpu", return_value = True) @patch.object(stack_mod, "_has_usable_nvidia_gpu", return_value = True) - def test_torch_backend_cuda_env_skips_entirely(self, mock_nvidia, mock_gpu, mock_pip): + def test_torch_backend_cuda_env_skips_entirely( + self, mock_nvidia, mock_gpu, mock_pip + ): """UNSLOTH_TORCH_BACKEND=cuda must short-circuit before any GPU probe.""" with patch.dict(os.environ, {"UNSLOTH_TORCH_BACKEND": "cuda"}): # Reload _TORCH_BACKEND from the patched environment. @@ -746,7 +760,9 @@ class TestEnsureRocmTorch: @patch.object(stack_mod, "pip_install") @patch.object(stack_mod, "_has_rocm_gpu", return_value = True) @patch.object(stack_mod, "_has_usable_nvidia_gpu", return_value = True) - def test_torch_backend_cpu_env_skips_entirely(self, mock_nvidia, mock_gpu, mock_pip): + def test_torch_backend_cpu_env_skips_entirely( + self, mock_nvidia, mock_gpu, mock_pip + ): """UNSLOTH_TORCH_BACKEND=cpu must short-circuit before any GPU probe.""" with patch.dict(os.environ, {"UNSLOTH_TORCH_BACKEND": "cpu"}): with patch.object(stack_mod, "_TORCH_BACKEND", "cpu"): @@ -796,7 +812,9 @@ class TestHasRocmGpuKfdVendorGuard: def test_sysfs_fallback_guarded_by_non_win32(self): """KFD sysfs fallback must be Linux-only (guarded by sys.platform != 'win32').""" src = self._src() - assert "win32" in src, "_has_rocm_gpu sysfs fallback must be guarded by sys.platform check" + assert ( + "win32" in src + ), "_has_rocm_gpu sysfs fallback must be guarded by sys.platform check" def test_cpu_node_excluded(self): """gpu_id == '0' must be excluded (CPU topology nodes).""" @@ -812,8 +830,12 @@ class TestHasRocmGpuKfdVendorGuard: func_start = source.find("_has_amd_rocm_gpu()") func_end = source.find("\n}", func_start) func_body = source[func_start:func_end] - assert "vendor_id" in func_body, "_has_amd_rocm_gpu sysfs fallback must check vendor_id" - assert "4098" in func_body, "_has_amd_rocm_gpu must require AMD vendor_id 4098 (0x1002)" + assert ( + "vendor_id" in func_body + ), "_has_amd_rocm_gpu sysfs fallback must check vendor_id" + assert ( + "4098" in func_body + ), "_has_amd_rocm_gpu must require AMD vendor_id 4098 (0x1002)" def test_has_rocm_gpu_returns_false_when_nvidia_present(self): """_has_rocm_gpu must return False immediately when _has_usable_nvidia_gpu is True. @@ -935,19 +957,25 @@ class TestHardwareRocmFlag: def test_hardware_py_has_is_rocm(self): """hardware.py should define IS_ROCM.""" - hw_path = PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py" + hw_path = ( + PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py" + ) source = hw_path.read_text(encoding = "utf-8") assert "IS_ROCM: bool" in source and "False" in source def test_hardware_py_sets_is_rocm_on_hip(self): """detect_hardware() should set IS_ROCM when torch.version.hip is set.""" - hw_path = PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py" + hw_path = ( + PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py" + ) source = hw_path.read_text(encoding = "utf-8") assert 'torch.version, "hip"' in source or "torch.version.hip" in source def test_hardware_py_still_returns_cuda_for_rocm(self): """DeviceType should remain CUDA even on ROCm -- no DeviceType.ROCM.""" - hw_path = PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py" + hw_path = ( + PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py" + ) source = hw_path.read_text(encoding = "utf-8") # Ensure ROCM is NOT a DeviceType member enum_section = source.split("class DeviceType")[1].split("\n\n")[0] @@ -955,13 +983,17 @@ class TestHardwareRocmFlag: def test_hardware_py_has_rocm_in_package_versions(self): """get_package_versions() should include 'rocm' key.""" - hw_path = PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py" + hw_path = ( + PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py" + ) source = hw_path.read_text(encoding = "utf-8") assert '"rocm"' in source def test_hardware_py_device_type_cuda_references_intact(self): """All existing DeviceType.CUDA references should still be present.""" - hw_path = PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py" + hw_path = ( + PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py" + ) source = hw_path.read_text(encoding = "utf-8") # Key functions that must still reference DeviceType.CUDA assert "DeviceType.CUDA" in source @@ -969,20 +1001,26 @@ class TestHardwareRocmFlag: def test_is_rocm_exported_from_init(self): """IS_ROCM should be exported from hardware __init__.py.""" - init_path = PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "__init__.py" + init_path = ( + PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "__init__.py" + ) source = init_path.read_text(encoding = "utf-8") assert "IS_ROCM" in source def test_is_rocm_in_all_list(self): """IS_ROCM should be in __all__ list in __init__.py.""" - init_path = PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "__init__.py" + init_path = ( + PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "__init__.py" + ) source = init_path.read_text(encoding = "utf-8") # Extract __all__ section assert '"IS_ROCM"' in source def test_get_package_versions_returns_rocm_key(self): """get_package_versions() source should return both 'cuda' and 'rocm' keys.""" - hw_path = PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py" + hw_path = ( + PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py" + ) source = hw_path.read_text(encoding = "utf-8") # Find the get_package_versions function body func_start = source.find("def get_package_versions") @@ -997,16 +1035,22 @@ class TestHardwareRocmFlag: Windows ROCm where torch.distributed ships without that helper, causing a warning: 'module torch.distributed has no attribute is_torchelastic_launched'. """ - hw_path = PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py" + hw_path = ( + PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py" + ) source = hw_path.read_text(encoding = "utf-8") assert "is_torchelastic_launched" in source def test_distributed_stubs_cover_core_helpers(self): """_determine_attention_impl_for_gpu_estimate must stub the four core distributed helpers.""" - hw_path = PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py" + hw_path = ( + PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py" + ) source = hw_path.read_text(encoding = "utf-8") for attr in ("is_initialized", "is_available", "get_rank", "get_world_size"): - assert attr in source, f"distributed stub for '{attr}' missing from hardware.py" + assert ( + attr in source + ), f"distributed stub for '{attr}' missing from hardware.py" # TEST: tokenizer_utils.py -- error message @@ -1054,7 +1098,9 @@ class TestInstallShStructure: # a genuine here-string operator lives outside any quotes. unquoted = re.sub(r"'[^']*'", "", line) unquoted = re.sub(r'"[^"]*"', "", unquoted) - assert "<<<" not in unquoted, f"install.sh:{i} uses non-POSIX <<< here-string" + assert ( + "<<<" not in unquoted + ), f"install.sh:{i} uses non-POSIX <<< here-string" def test_rocm_detection_present(self): """install.sh should have ROCm detection in get_torch_index_url.""" @@ -1084,8 +1130,12 @@ class TestInstallShStructure: if no_nvidia_branch < 0: no_nvidia_branch = body.find('if [ -z "$_smi" ]') rocm_call = body.find("_has_amd_rocm_gpu") - assert nvidia_call >= 0, "get_torch_index_url should call _has_usable_nvidia_gpu" - assert no_nvidia_branch >= 0, "get_torch_index_url should gate ROCm on no-nvidia branch" + assert ( + nvidia_call >= 0 + ), "get_torch_index_url should call _has_usable_nvidia_gpu" + assert ( + no_nvidia_branch >= 0 + ), "get_torch_index_url should gate ROCm on no-nvidia branch" assert ( rocm_call > no_nvidia_branch ), "ROCm detection should sit inside the 'no NVIDIA' branch" @@ -1146,7 +1196,9 @@ class TestInstallShStructure: continue # Remove POSIX character classes [[:foo:]] before checking for [[ ]] cleaned = re.sub(r"\[\[:[a-z]+:\]\]", "", line) - assert "[[" not in cleaned, f"get_torch_index_url line {i} uses non-POSIX [[" + assert ( + "[[" not in cleaned + ), f"get_torch_index_url line {i} uses non-POSIX [[" def test_no_arithmetic_expansion_in_rocm_block(self): """ROCm detection block should not use (( )) (bash-only).""" @@ -1307,7 +1359,9 @@ class TestLiveRegression: # Load worker.py module _WORKER_PATH = PACKAGE_ROOT / "studio" / "backend" / "core" / "training" / "worker.py" -_EXPORT_WORKER_PATH = PACKAGE_ROOT / "studio" / "backend" / "core" / "export" / "worker.py" +_EXPORT_WORKER_PATH = ( + PACKAGE_ROOT / "studio" / "backend" / "core" / "export" / "worker.py" +) # The torchao Windows-ROCm stub was de-duplicated out of the export/training # workers into a shared module; both workers now call into it. _TORCHAO_STUB_PATH = PACKAGE_ROOT / "studio" / "backend" / "core" / "_torchao_stub.py" @@ -1333,7 +1387,9 @@ class TestWorkerRocmMambaSsm: def test_direct_wheel_url_returns_none_without_cuda_major(self, monkeypatch): """_direct_wheel_url should return None when cuda_major is empty (ROCm).""" # Load module for function access - _worker_spec = importlib.util.spec_from_file_location("test_worker", _WORKER_PATH) + _worker_spec = importlib.util.spec_from_file_location( + "test_worker", _WORKER_PATH + ) assert _worker_spec is not None and _worker_spec.loader is not None worker_mod = importlib.util.module_from_spec(_worker_spec) @@ -1565,7 +1621,9 @@ class TestHardwareAmdBranching: def test_hardware_imports_amd_module(self): """hardware.py should import from amd module when IS_ROCM.""" - hw_path = PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py" + hw_path = ( + PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py" + ) source = hw_path.read_text(encoding = "utf-8") assert "from . import amd" in source @@ -1573,13 +1631,17 @@ class TestHardwareAmdBranching: """get_gpu_utilization should dispatch to amd.py via _smi_query when IS_ROCM, and the dispatcher itself must check IS_ROCM and import the amd backend.""" - hw_path = PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py" + hw_path = ( + PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py" + ) source = hw_path.read_text(encoding = "utf-8") func_start = source.find("def get_gpu_utilization") func_body = source[func_start : source.find("\ndef ", func_start + 1)] assert '_smi_query("get_primary_gpu_utilization"' in func_body smi = source[ - source.find("def _smi_query") : source.find("\ndef ", source.find("def _smi_query") + 1) + source.find("def _smi_query") : source.find( + "\ndef ", source.find("def _smi_query") + 1 + ) ] assert "IS_ROCM" in smi assert "from . import amd" in smi @@ -1587,7 +1649,9 @@ class TestHardwareAmdBranching: def test_hardware_branches_on_is_rocm_for_visible(self): """get_visible_gpu_utilization should dispatch to amd.py via _smi_query when IS_ROCM.""" - hw_path = PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py" + hw_path = ( + PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py" + ) source = hw_path.read_text(encoding = "utf-8") func_start = source.find("def get_visible_gpu_utilization") func_body = source[func_start : source.find("\ndef ", func_start + 1)] @@ -1597,14 +1661,18 @@ class TestHardwareAmdBranching: assert _re.search(r'_smi_query\(\s*"get_visible_gpu_utilization"', func_body) smi = source[ - source.find("def _smi_query") : source.find("\ndef ", source.find("def _smi_query") + 1) + source.find("def _smi_query") : source.find( + "\ndef ", source.find("def _smi_query") + 1 + ) ] assert "IS_ROCM" in smi assert "from . import amd" in smi def test_hardware_branches_on_is_rocm_for_physical_count(self): """get_physical_gpu_count should try amd.py when IS_ROCM.""" - hw_path = PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py" + hw_path = ( + PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py" + ) source = hw_path.read_text(encoding = "utf-8") func_start = source.find("def get_physical_gpu_count") func_body = source[func_start : source.find("\ndef ", func_start + 1)] @@ -1621,7 +1689,9 @@ class TestApplyGpuIdsRocmFallback: def test_apply_gpu_ids_falls_back_to_torch_version_hip(self): """apply_gpu_ids should probe torch.version.hip when IS_ROCM is False and no ROCm env vars are set.""" - hw_path = PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py" + hw_path = ( + PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py" + ) source = hw_path.read_text(encoding = "utf-8") func_start = source.find("def apply_gpu_ids") func_body = source[func_start : source.find("\ndef ", func_start + 1)] @@ -1633,7 +1703,9 @@ class TestApplyGpuIdsRocmFallback: ROCR_VISIBLE_DEVICES uses HSA agent-level indexing, not physical GPU indices. Overwriting it breaks multi-GPU ROCm systems (see issue #6118). """ - hw_path = PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py" + hw_path = ( + PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py" + ) source = hw_path.read_text(encoding = "utf-8") func_start = source.find("def apply_gpu_ids") func_body = source[func_start : source.find("\ndef ", func_start + 1)] @@ -1642,7 +1714,9 @@ class TestApplyGpuIdsRocmFallback: def test_apply_gpu_ids_rocm_fallback_is_guarded_by_try_except(self): """torch import in apply_gpu_ids must be wrapped in try/except so a missing torch never crashes.""" - hw_path = PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py" + hw_path = ( + PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py" + ) source = hw_path.read_text(encoding = "utf-8") func_start = source.find("def apply_gpu_ids") func_body = source[func_start : source.find("\ndef ", func_start + 1)] @@ -1787,7 +1861,9 @@ class TestWindowsRocmIndexUrl: assert "repo.amd.com" in url def test_mirror_env_var_overrides_base(self, monkeypatch): - monkeypatch.setenv("UNSLOTH_ROCM_WINDOWS_MIRROR", "https://my-mirror.example.com/rocm/whl") + monkeypatch.setenv( + "UNSLOTH_ROCM_WINDOWS_MIRROR", "https://my-mirror.example.com/rocm/whl" + ) # Reload module-level constant by calling helper directly url = stack_mod._windows_rocm_index_url("gfx1200") # The env var is read at module load time for _ROCM_WINDOWS_INDEX_BASE, @@ -1900,7 +1976,9 @@ class TestGfxArchNameFallback: """hipinfo absent everywhere + amd-smi absent -> WMI name fallback.""" ps_result = MagicMock() ps_result.returncode = 0 - ps_result.stdout = b"AMD Radeon(TM) 8060S Graphics\r\nMicrosoft Basic Display Adapter\r\n" + ps_result.stdout = ( + b"AMD Radeon(TM) 8060S Graphics\r\nMicrosoft Basic Display Adapter\r\n" + ) def _run(cmd, **kwargs): if cmd and "powershell.exe" in str(cmd[0]).lower(): @@ -2032,7 +2110,9 @@ class TestInstallBnbWindowsRocm: os.environ.pop("BNB_ROCM_VERSION", None) os.environ.pop(stack_mod._BNB_ROCM_VERSION_SOURCE_ENV, None) with patch.object(stack_mod, "pip_install_try", return_value = True): - with patch.object(stack_mod, "_detect_bnb_rocm_dll_ver", return_value = "72"): + with patch.object( + stack_mod, "_detect_bnb_rocm_dll_ver", return_value = "72" + ): stack_mod._install_bnb_windows_rocm() assert os.environ.get("BNB_ROCM_VERSION") == "72" @@ -2042,7 +2122,9 @@ class TestInstallBnbWindowsRocm: os.environ.pop("BNB_ROCM_VERSION", None) os.environ.pop(stack_mod._BNB_ROCM_VERSION_SOURCE_ENV, None) with patch.object(stack_mod, "pip_install_try", return_value = True): - with patch.object(stack_mod, "_detect_bnb_rocm_dll_ver", return_value = "713"): + with patch.object( + stack_mod, "_detect_bnb_rocm_dll_ver", return_value = "713" + ): stack_mod._install_bnb_windows_rocm() assert os.environ.get("BNB_ROCM_VERSION") == "713" @@ -2052,7 +2134,9 @@ class TestInstallBnbWindowsRocm: os.environ.pop("BNB_ROCM_VERSION", None) os.environ.pop(stack_mod._BNB_ROCM_VERSION_SOURCE_ENV, None) with patch.object(stack_mod, "pip_install_try", return_value = True): - with patch.object(stack_mod, "_detect_bnb_rocm_dll_ver", return_value = None): + with patch.object( + stack_mod, "_detect_bnb_rocm_dll_ver", return_value = None + ): stack_mod._install_bnb_windows_rocm() assert os.environ.get("BNB_ROCM_VERSION") == "72" @@ -2091,7 +2175,9 @@ class TestInstallBnbWindowsRocm: }, ): with patch.object(stack_mod, "pip_install_try", return_value = True): - with patch.object(stack_mod, "_detect_bnb_rocm_dll_ver", return_value = "713"): + with patch.object( + stack_mod, "_detect_bnb_rocm_dll_ver", return_value = "713" + ): with patch.object( stack_mod, "_persist_bnb_rocm_version", return_value = True ) as mock_persist: @@ -2112,7 +2198,9 @@ class TestInstallBnbWindowsRocm: os.environ.pop("BNB_ROCM_VERSION", None) os.environ.pop(stack_mod._BNB_ROCM_VERSION_SOURCE_ENV, None) with patch.object(stack_mod, "pip_install_try", return_value = True): - with patch.object(stack_mod, "_detect_bnb_rocm_dll_ver", return_value = "72"): + with patch.object( + stack_mod, "_detect_bnb_rocm_dll_ver", return_value = "72" + ): with patch.object( stack_mod.sysconfig, "get_path", return_value = str(site_packages) ): @@ -2160,7 +2248,9 @@ class TestInstallBnbWindowsRocm: encoding = "utf-8", ) - with patch.object(stack_mod.sysconfig, "get_path", return_value = str(site_packages)): + with patch.object( + stack_mod.sysconfig, "get_path", return_value = str(site_packages) + ): assert stack_mod._persist_bnb_rocm_version("713") is True source = sitecustomize.read_text(encoding = "utf-8") @@ -2176,7 +2266,9 @@ class TestInstallBnbWindowsRocm: sitecustomize = site_packages / "sitecustomize.py" sitecustomize.write_bytes(b"\xff\xfe\x00") - with patch.object(stack_mod.sysconfig, "get_path", return_value = str(site_packages)): + with patch.object( + stack_mod.sysconfig, "get_path", return_value = str(site_packages) + ): assert stack_mod._persist_bnb_rocm_version("72") is False def test_persist_bnb_rocm_version_repairs_truncated_block(self, tmp_path): @@ -2192,7 +2284,9 @@ class TestInstallBnbWindowsRocm: encoding = "utf-8", ) - with patch.object(stack_mod.sysconfig, "get_path", return_value = str(site_packages)): + with patch.object( + stack_mod.sysconfig, "get_path", return_value = str(site_packages) + ): assert stack_mod._persist_bnb_rocm_version("713") is True source = sitecustomize.read_text(encoding = "utf-8") @@ -2215,7 +2309,9 @@ class TestInstallBnbWindowsRocm: ) sitecustomize.write_text(block + "USER_MID = 1\n" + block, encoding = "utf-8") - with patch.object(stack_mod.sysconfig, "get_path", return_value = str(site_packages)): + with patch.object( + stack_mod.sysconfig, "get_path", return_value = str(site_packages) + ): assert stack_mod._persist_bnb_rocm_version("713") is True source = sitecustomize.read_text(encoding = "utf-8") @@ -2230,7 +2326,9 @@ class TestInstallBnbWindowsRocm: site_packages = tmp_path / "site-packages" site_packages.mkdir() - with patch.object(stack_mod.sysconfig, "get_path", return_value = str(site_packages)): + with patch.object( + stack_mod.sysconfig, "get_path", return_value = str(site_packages) + ): assert stack_mod._persist_bnb_rocm_version("72") is True leftovers = [p.name for p in site_packages.iterdir() if "unsloth-tmp" in p.name] @@ -2242,16 +2340,24 @@ class TestRuntimeBnbRocmSourceGuards: """Runtime entrypoints redetect managed defaults but keep caller overrides.""" _MAIN_PATH = PACKAGE_ROOT / "studio" / "backend" / "main.py" - _TRAINING_WORKER_PATH = PACKAGE_ROOT / "studio" / "backend" / "core" / "training" / "worker.py" + _TRAINING_WORKER_PATH = ( + PACKAGE_ROOT / "studio" / "backend" / "core" / "training" / "worker.py" + ) def test_main_gate_redetects_persisted_default(self): source = self._MAIN_PATH.read_text(encoding = "utf-8") - assert 'os.environ.get("UNSLOTH_BNB_ROCM_VERSION_SOURCE") == "sitecustomize"' in source + assert ( + 'os.environ.get("UNSLOTH_BNB_ROCM_VERSION_SOURCE") == "sitecustomize"' + in source + ) assert 'os.environ["UNSLOTH_BNB_ROCM_VERSION_SOURCE"] = "detected"' in source def test_worker_gate_redetects_persisted_default(self): source = self._TRAINING_WORKER_PATH.read_text(encoding = "utf-8") - assert 'os.environ.get("UNSLOTH_BNB_ROCM_VERSION_SOURCE") == "sitecustomize"' in source + assert ( + 'os.environ.get("UNSLOTH_BNB_ROCM_VERSION_SOURCE") == "sitecustomize"' + in source + ) assert 'os.environ["UNSLOTH_BNB_ROCM_VERSION_SOURCE"] = "detected"' in source def test_fallback_prefers_seeded_value_over_hardcoded_72(self): @@ -2487,7 +2593,9 @@ class TestWorkerWindowsRocmPatches: # entry-point function (not the trainer helper which has its own "# ── 2."). idx_sec2 = source.find("# ── 2. Now import ML libraries") assert idx_bnb != -1, "BNB_ROCM_VERSION not found in worker.py" - assert idx_sec2 != -1, "'# ── 2. Now import ML libraries' marker not found in worker.py" + assert ( + idx_sec2 != -1 + ), "'# ── 2. Now import ML libraries' marker not found in worker.py" assert idx_bnb < idx_sec2, ( "BNB_ROCM_VERSION must be set before section 2 ML imports " f"(found at {idx_bnb}, section 2 at {idx_sec2})" @@ -2775,12 +2883,16 @@ class TestHipSdkEnvPathResolution: """setup.ps1 must tell the user how to add the HIP bin dir to PATH.""" source = _SETUP_PS1_PATH.read_text(encoding = "utf-8") # Should mention adding to PATH or SetEnvironmentVariable - assert "PATH" in source and ("SetEnvironmentVariable" in source or "Add" in source) + assert "PATH" in source and ( + "SetEnvironmentVariable" in source or "Add" in source + ) def test_install_provides_path_fix_hint(self): """install.ps1 must tell the user how to add the HIP bin dir to PATH.""" source = _INSTALL_PS1_PATH.read_text(encoding = "utf-8") - assert "PATH" in source and ("SetEnvironmentVariable" in source or "Add" in source) + assert "PATH" in source and ( + "SetEnvironmentVariable" in source or "Add" in source + ) # TEST: HIP SDK detected substep -- path + hipconfig version shown in terminal @@ -2942,7 +3054,9 @@ class TestSetupShGccInstallDir: # TEST: main.py -- BNB_ROCM_VERSION server startup + distributed stubs _MAIN_PY_PATH = PACKAGE_ROOT / "studio" / "backend" / "main.py" -_HARDWARE_PY_PATH = PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py" +_HARDWARE_PY_PATH = ( + PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py" +) class TestServerStartupRocmFixes: @@ -3184,7 +3298,9 @@ class TestApplyHostOverrides: assert out.rocm_gfx_target is None def test_malformed_forwarded_gfx_falls_back_to_has_rocm(self): - out = _apply_host_overrides(cpu_host(), override_has_rocm = True, override_rocm_gfx = "junk") + out = _apply_host_overrides( + cpu_host(), override_has_rocm = True, override_rocm_gfx = "junk" + ) assert out.has_rocm is True assert out.rocm_gfx_target is None @@ -3285,14 +3401,22 @@ class TestRocmGfxForwarding: # arch -> the host must still be treated as a GPU host and routed to the # fork's per-gfx prebuilt, not ggml-org / a source build. Linux x64 and # arm64 both go through the same fork branch. - assert self._resolve_setup_sh_repo("x86_64", False, "gfx1100") == "unslothai/llama.cpp" - assert self._resolve_setup_sh_repo("aarch64", False, "gfx1100") == "unslothai/llama.cpp" + assert ( + self._resolve_setup_sh_repo("x86_64", False, "gfx1100") + == "unslothai/llama.cpp" + ) + assert ( + self._resolve_setup_sh_repo("aarch64", False, "gfx1100") + == "unslothai/llama.cpp" + ) def test_setup_sh_env_forwarded_gfx_resolves_to_fork(self): # UNSLOTH_ROCM_GFX_ARCH set on a host where no probe fired (_setup_gfx # empty, no usable NVIDIA, no ROCm tooling): setup.sh adopts the env arch # and routes to the fork, same as the name-inference path. - repo = self._resolve_setup_sh_repo("x86_64", False, "", rocm_gfx_arch_env = "gfx1100") + repo = self._resolve_setup_sh_repo( + "x86_64", False, "", rocm_gfx_arch_env = "gfx1100" + ) assert repo == "unslothai/llama.cpp" def test_setup_sh_cpu_host_still_resolves_to_ggml(self): @@ -3339,7 +3463,10 @@ class TestRocmGfxForwarding: def test_setup_ps1_inferred_gfx_resolves_to_fork(self): # Adrenalin-only Windows host: $HasROCm is false (no HIP runtime) but a # gfx arch was inferred -> route to the fork's windows-rocm bundle. - assert self._resolve_setup_ps1_repo(False, False, "gfx1100") == "unslothai/llama.cpp" + assert ( + self._resolve_setup_ps1_repo(False, False, "gfx1100") + == "unslothai/llama.cpp" + ) def test_setup_ps1_cpu_host_still_resolves_to_ggml(self): # No NVIDIA, no ROCm, no inferred gfx -> CPU host stays on ggml-org. @@ -3394,7 +3521,9 @@ def test_pick_rocm_gfx_target_same_arch_multi_gpu(monkeypatch): _INSTALL_SH_PATH = PACKAGE_ROOT / "install.sh" -_LLAMA_CPP_PATH = PACKAGE_ROOT / "studio" / "backend" / "core" / "inference" / "llama_cpp.py" +_LLAMA_CPP_PATH = ( + PACKAGE_ROOT / "studio" / "backend" / "core" / "inference" / "llama_cpp.py" +) class TestWslSystemRocmLibDirs: @@ -3471,7 +3600,9 @@ class TestBinaryEnvWslOrdering: # real dir to stand in for the system ROCm lib path. sys_rocm = tmp_path / "sysrocm" sys_rocm.mkdir() - with patch.object(prebuilt_mod, "_wsl_system_rocm_lib_dirs", return_value = [str(sys_rocm)]): + with patch.object( + prebuilt_mod, "_wsl_system_rocm_lib_dirs", return_value = [str(sys_rocm)] + ): with patch.dict(os.environ, {}, clear = True): env = prebuilt_mod.binary_env(binary, tmp_path, self._linux_host()) ld = env["LD_LIBRARY_PATH"].split(os.pathsep) diff --git a/tests/studio/install/test_selection_logic.py b/tests/studio/install/test_selection_logic.py index 4e8beeb888..48e39d3eb3 100644 --- a/tests/studio/install/test_selection_logic.py +++ b/tests/studio/install/test_selection_logic.py @@ -25,7 +25,9 @@ import pytest PACKAGE_ROOT = Path(__file__).resolve().parents[3] MODULE_PATH = PACKAGE_ROOT / "studio" / "install_llama_prebuilt.py" RUN_MODULE_PATH = PACKAGE_ROOT / "studio" / "backend" / "run.py" -SPEC = importlib.util.spec_from_file_location("studio_install_llama_prebuilt", MODULE_PATH) +SPEC = importlib.util.spec_from_file_location( + "studio_install_llama_prebuilt", MODULE_PATH +) assert SPEC is not None and SPEC.loader is not None INSTALL_LLAMA_PREBUILT = importlib.util.module_from_spec(SPEC) sys.modules[SPEC.name] = INSTALL_LLAMA_PREBUILT @@ -50,7 +52,9 @@ supports_explicit_visible_device_matching = ( select_visible_gpu_rows = INSTALL_LLAMA_PREBUILT.select_visible_gpu_rows compatible_linux_runtime_lines = INSTALL_LLAMA_PREBUILT.compatible_linux_runtime_lines pick_windows_cuda_runtime = INSTALL_LLAMA_PREBUILT.pick_windows_cuda_runtime -compatible_windows_runtime_lines = INSTALL_LLAMA_PREBUILT.compatible_windows_runtime_lines +compatible_windows_runtime_lines = ( + INSTALL_LLAMA_PREBUILT.compatible_windows_runtime_lines +) runtime_line_from_cuda_version = INSTALL_LLAMA_PREBUILT.runtime_line_from_cuda_version apply_approved_hashes = INSTALL_LLAMA_PREBUILT.apply_approved_hashes linux_cuda_choice_from_release = INSTALL_LLAMA_PREBUILT.linux_cuda_choice_from_release @@ -62,11 +66,19 @@ _fork_manifest_release_plans = INSTALL_LLAMA_PREBUILT._fork_manifest_release_pla resolve_published_release = INSTALL_LLAMA_PREBUILT.resolve_published_release resolve_source_build_plan = INSTALL_LLAMA_PREBUILT.resolve_source_build_plan validated_checksums_for_bundle = INSTALL_LLAMA_PREBUILT.validated_checksums_for_bundle -parse_approved_release_checksums = INSTALL_LLAMA_PREBUILT.parse_approved_release_checksums -published_release_matches_request = INSTALL_LLAMA_PREBUILT.published_release_matches_request -exact_source_archive_logical_name = INSTALL_LLAMA_PREBUILT.exact_source_archive_logical_name +parse_approved_release_checksums = ( + INSTALL_LLAMA_PREBUILT.parse_approved_release_checksums +) +published_release_matches_request = ( + INSTALL_LLAMA_PREBUILT.published_release_matches_request +) +exact_source_archive_logical_name = ( + INSTALL_LLAMA_PREBUILT.exact_source_archive_logical_name +) source_archive_logical_name = INSTALL_LLAMA_PREBUILT.source_archive_logical_name -windows_cuda_upstream_asset_names = INSTALL_LLAMA_PREBUILT.windows_cuda_upstream_asset_names +windows_cuda_upstream_asset_names = ( + INSTALL_LLAMA_PREBUILT.windows_cuda_upstream_asset_names +) env_int = INSTALL_LLAMA_PREBUILT.env_int direct_upstream_release_plan = INSTALL_LLAMA_PREBUILT.direct_upstream_release_plan _pinned_windows_cuda_fallback = INSTALL_LLAMA_PREBUILT._pinned_windows_cuda_fallback @@ -77,7 +89,9 @@ _windows_cuda_attempt_covers_blackwell = ( ) resolve_release_asset_choice = INSTALL_LLAMA_PREBUILT.resolve_release_asset_choice pinned_macos_release_tag = INSTALL_LLAMA_PREBUILT.pinned_macos_release_tag -resolve_simple_install_release_plans = INSTALL_LLAMA_PREBUILT.resolve_simple_install_release_plans +resolve_simple_install_release_plans = ( + INSTALL_LLAMA_PREBUILT.resolve_simple_install_release_plans +) def load_studio_run_module(monkeypatch): @@ -230,7 +244,9 @@ def make_checksums_with_source( kind = "upstream-source", ), } - normalized_source_commit = source_commit.lower() if isinstance(source_commit, str) else None + normalized_source_commit = ( + source_commit.lower() if isinstance(source_commit, str) else None + ) if normalized_source_commit: artifacts[exact_source_archive_logical_name(normalized_source_commit)] = ( ApprovedArtifactHash( @@ -250,7 +266,9 @@ def make_checksums_with_source( requested_source_ref = requested_source_ref, resolved_source_ref = resolved_source_ref, source_commit = normalized_source_commit, - source_commit_short = normalized_source_commit[:7] if normalized_source_commit else None, + source_commit_short = normalized_source_commit[:7] + if normalized_source_commit + else None, artifacts = artifacts, ) @@ -345,7 +363,10 @@ class TestStudioLocalhostIpv6Warning: lambda host, port, timeout = 1.0: True, ) - assert run_module._localhost_ipv6_mismatch_url("127.0.0.1", 8888) == "http://127.0.0.1:8888" + assert ( + run_module._localhost_ipv6_mismatch_url("127.0.0.1", 8888) + == "http://127.0.0.1:8888" + ) @pytest.mark.parametrize("host", ["0.0.0.0", "::"]) def test_network_bind_suppresses_warning(self, monkeypatch, host): @@ -410,7 +431,9 @@ class TestStudioLocalhostIpv6Warning: monkeypatch.setattr( run_module, "_verify_global_reachability", - lambda display_host, port: calls["reachability"].append((display_host, port)), + lambda display_host, port: calls["reachability"].append( + (display_host, port) + ), ) return calls @@ -433,7 +456,9 @@ class TestStudioLocalhostIpv6Warning: def test_emit_startup_output_plain_localhost(self, monkeypatch): run_module = load_studio_run_module(monkeypatch) calls = self._wire_recorders(run_module, monkeypatch) - monkeypatch.setattr(run_module, "_localhost_ipv6_mismatch_url", lambda host, port: None) + monkeypatch.setattr( + run_module, "_localhost_ipv6_mismatch_url", lambda host, port: None + ) run_module._emit_startup_output("127.0.0.1", 8888, "127.0.0.1") @@ -446,7 +471,9 @@ class TestStudioLocalhostIpv6Warning: def test_emit_startup_output_wildcard_runs_reachability(self, monkeypatch, host): run_module = load_studio_run_module(monkeypatch) calls = self._wire_recorders(run_module, monkeypatch) - monkeypatch.setattr(run_module, "_localhost_ipv6_mismatch_url", lambda h, port: None) + monkeypatch.setattr( + run_module, "_localhost_ipv6_mismatch_url", lambda h, port: None + ) run_module._emit_startup_output(host, 8888, "203.0.113.5") @@ -833,7 +860,9 @@ class TestPublishedReleaseResolution: def fake_load(repo, release_tag): if release_tag == "v2.0": raise PrebuiltFallback("checksum asset missing") - return make_checksums_with_source([], release_tag = "v1.0", upstream_tag = "b8999") + return make_checksums_with_source( + [], release_tag = "v1.0", upstream_tag = "b8999" + ) monkeypatch.setattr( INSTALL_LLAMA_PREBUILT, @@ -863,7 +892,9 @@ class TestPublishedReleaseResolution: ), ) - assert resolve_requested_install_tag("b8508", "", "unslothai/llama.cpp") == "b8508" + assert ( + resolve_requested_install_tag("b8508", "", "unslothai/llama.cpp") == "b8508" + ) def test_concrete_tag_without_matching_release_raises(self, monkeypatch): release = make_release([], release_tag = "release-b9000", upstream_tag = "b9000") @@ -877,7 +908,9 @@ class TestPublishedReleaseResolution: resolve_requested_install_tag("b8508", "", "unslothai/llama.cpp") def test_pinned_release_must_match_requested_upstream_tag(self, monkeypatch): - bundle = make_release([], release_tag = "llama-prebuilt-latest", upstream_tag = "b9000") + bundle = make_release( + [], release_tag = "llama-prebuilt-latest", upstream_tag = "b9000" + ) monkeypatch.setattr( INSTALL_LLAMA_PREBUILT, "pinned_published_release_bundle", @@ -1053,13 +1086,15 @@ class TestSourceBuildPlanResolution: assert plan.source_ref == "main" assert plan.compatibility_upstream_tag == "b9000" - def test_direct_main_request_without_published_release_uses_branch_kind(self, monkeypatch): + def test_direct_main_request_without_published_release_uses_branch_kind( + self, monkeypatch + ): monkeypatch.setattr( INSTALL_LLAMA_PREBUILT, "resolve_published_release", - lambda requested_tag, published_repo, published_release_tag = "": (_ for _ in ()).throw( - PrebuiltFallback("missing") - ), + lambda requested_tag, published_repo, published_release_tag = "": ( + _ for _ in () + ).throw(PrebuiltFallback("missing")), ) plan = resolve_source_build_plan("main", "unslothai/llama.cpp") @@ -1134,7 +1169,9 @@ class TestValidatedChecksumsForBundle: def test_rejects_manifest_checksum_mismatch(self, monkeypatch): bundle = make_release([], release_tag = "r1", upstream_tag = "b8508") bundle.manifest_sha256 = "a" * 64 - checksums = make_checksums_with_source([], release_tag = "r1", upstream_tag = "b8508") + checksums = make_checksums_with_source( + [], release_tag = "r1", upstream_tag = "b8508" + ) checksums.artifacts[bundle.manifest_asset_name] = ApprovedArtifactHash( asset_name = bundle.manifest_asset_name, sha256 = "b" * 64, @@ -1187,9 +1224,13 @@ class TestValidatedChecksumsForBundle: lambda repo, release_tag: checksums, ) - assert validated_checksums_for_bundle("unslothai/llama.cpp", bundle) is checksums + assert ( + validated_checksums_for_bundle("unslothai/llama.cpp", bundle) is checksums + ) plan = INSTALL_LLAMA_PREBUILT.source_build_plan_for_release( - INSTALL_LLAMA_PREBUILT.ResolvedPublishedRelease(bundle = bundle, checksums = checksums) + INSTALL_LLAMA_PREBUILT.ResolvedPublishedRelease( + bundle = bundle, checksums = checksums + ) ) assert plan.source_url == "https://github.com/ggml-org/llama.cpp" assert plan.source_ref_kind == "commit" @@ -1233,7 +1274,9 @@ class TestLinuxCudaChoiceFromRelease: art12 = make_artifact("bundle-cuda12.tar.gz", runtime_line = "cuda12") art13 = make_artifact("bundle-cuda13.tar.gz", runtime_line = "cuda13") release = make_release([art12, art13]) - result = linux_cuda_choice_from_release(host, release, preferred_runtime_line = "cuda12") + result = linux_cuda_choice_from_release( + host, release, preferred_runtime_line = "cuda12" + ) assert result is not None assert result.primary.runtime_line == "cuda12" @@ -1242,7 +1285,9 @@ class TestLinuxCudaChoiceFromRelease: host = make_host(driver_cuda_version = (12, 8)) art = make_artifact("bundle-cuda12.tar.gz", runtime_line = "cuda12") release = make_release([art]) - result = linux_cuda_choice_from_release(host, release, preferred_runtime_line = "cuda13") + result = linux_cuda_choice_from_release( + host, release, preferred_runtime_line = "cuda13" + ) assert result is not None assert result.primary.runtime_line == "cuda12" log_entries = result.selection_log @@ -1283,7 +1328,9 @@ class TestLinuxCudaChoiceFromRelease: def test_exact_sm_match(self, monkeypatch): mock_linux_runtime(monkeypatch, ["cuda12"]) host = make_host(compute_caps = ["86"]) - art = make_artifact("bundle.tar.gz", supported_sms = ["75", "86", "89"], min_sm = 75, max_sm = 89) + art = make_artifact( + "bundle.tar.gz", supported_sms = ["75", "86", "89"], min_sm = 75, max_sm = 89 + ) release = make_release([art]) result = linux_cuda_choice_from_release(host, release) assert result is not None @@ -1292,7 +1339,9 @@ class TestLinuxCudaChoiceFromRelease: def test_sm_not_in_supported_sms(self, monkeypatch): mock_linux_runtime(monkeypatch, ["cuda12"]) host = make_host(compute_caps = ["86"]) - art = make_artifact("bundle.tar.gz", supported_sms = ["75", "80", "89"], min_sm = 75, max_sm = 89) + art = make_artifact( + "bundle.tar.gz", supported_sms = ["75", "80", "89"], min_sm = 75, max_sm = 89 + ) release = make_release([art]) result = linux_cuda_choice_from_release(host, release) assert result is None @@ -1300,7 +1349,9 @@ class TestLinuxCudaChoiceFromRelease: def test_sm_outside_min_range(self, monkeypatch): mock_linux_runtime(monkeypatch, ["cuda12"]) host = make_host(compute_caps = ["50"]) - art = make_artifact("bundle.tar.gz", supported_sms = ["50", "75", "86"], min_sm = 75, max_sm = 90) + art = make_artifact( + "bundle.tar.gz", supported_sms = ["50", "75", "86"], min_sm = 75, max_sm = 90 + ) release = make_release([art]) result = linux_cuda_choice_from_release(host, release) assert result is None @@ -1369,7 +1420,9 @@ class TestLinuxCudaChoiceFromRelease: def test_multi_gpu_not_all_covered(self, monkeypatch): mock_linux_runtime(monkeypatch, ["cuda12"]) host = make_host(compute_caps = ["50", "89"]) - art = make_artifact("bundle.tar.gz", supported_sms = ["75", "89"], min_sm = 75, max_sm = 89) + art = make_artifact( + "bundle.tar.gz", supported_sms = ["75", "89"], min_sm = 75, max_sm = 89 + ) release = make_release([art]) result = linux_cuda_choice_from_release(host, release) assert result is None @@ -1488,7 +1541,9 @@ class TestLinuxCudaChoiceFromRelease: class TestResolveInstallAttempts: - def test_windows_cuda_prefers_published_asset_from_selected_release(self, monkeypatch): + def test_windows_cuda_prefers_published_asset_from_selected_release( + self, monkeypatch + ): host = make_host(system = "Windows", machine = "AMD64") host.driver_cuda_version = (12, 4) mock_windows_runtime(monkeypatch, ["cuda12"]) @@ -1532,7 +1587,9 @@ class TestResolveInstallAttempts: INSTALL_LLAMA_PREBUILT, "github_release_assets", lambda repo, tag: (_ for _ in ()).throw( - AssertionError("published Windows CUDA choice should not query upstream") + AssertionError( + "published Windows CUDA choice should not query upstream" + ) ), ) @@ -1554,7 +1611,9 @@ class TestResolveInstallAttempts: host = make_host(system = "Windows", machine = "AMD64") host.driver_cuda_version = (12, 4) mock_windows_runtime(monkeypatch, ["cuda12"]) - release = make_release([], release_tag = "llama-prebuilt-latest", upstream_tag = "b9000") + release = make_release( + [], release_tag = "llama-prebuilt-latest", upstream_tag = "b9000" + ) checksums = make_checksums_with_source( ["llama-b9000-bin-win-cuda-12.4-x64.zip"], release_tag = release.release_tag, @@ -1609,7 +1668,9 @@ class TestResolveInstallAttempts: assert attempts[0].expected_sha256 == "a" * 64 assert approved.release_tag == "llama-prebuilt-latest" - def test_linux_cpu_fork_without_bundle_raises_no_upstream_fallback(self, monkeypatch): + def test_linux_cpu_fork_without_bundle_raises_no_upstream_fallback( + self, monkeypatch + ): # A CPU-only Linux host on the fork no longer falls back to the ggml-org # CPU asset: production routes CPU-only Linux to ggml-org, never the fork. # With no fork CPU bundle in the manifest the resolver raises rather than @@ -1619,7 +1680,9 @@ class TestResolveInstallAttempts: has_physical_nvidia = False, nvidia_smi = None, ) - release = make_release([], release_tag = "llama-prebuilt-latest", upstream_tag = "b9000") + release = make_release( + [], release_tag = "llama-prebuilt-latest", upstream_tag = "b9000" + ) checksums = make_checksums_with_source( [], release_tag = release.release_tag, @@ -1646,12 +1709,16 @@ class TestResolveInstallAttempts: ), ) - with pytest.raises(PrebuiltFallback, match = "no compatible Linux prebuilt asset was found"): + with pytest.raises( + PrebuiltFallback, match = "no compatible Linux prebuilt asset was found" + ): resolve_install_attempts("latest", host, "unslothai/llama.cpp", "") def test_linux_cuda_does_not_fall_back_to_upstream_cpu(self, monkeypatch): host = make_host(system = "Linux", machine = "x86_64", compute_caps = ["86"]) - release = make_release([], release_tag = "llama-prebuilt-latest", upstream_tag = "b9000") + release = make_release( + [], release_tag = "llama-prebuilt-latest", upstream_tag = "b9000" + ) checksums = make_checksums_with_source( [], release_tag = release.release_tag, @@ -1672,7 +1739,9 @@ class TestResolveInstallAttempts: ) mock_linux_runtime(monkeypatch, ["cuda12"]) - with pytest.raises(PrebuiltFallback, match = "no compatible Linux prebuilt asset was found"): + with pytest.raises( + PrebuiltFallback, match = "no compatible Linux prebuilt asset was found" + ): resolve_install_attempts("latest", host, "unslothai/llama.cpp", "") def test_windows_cpu_prefers_published_asset(self, monkeypatch): @@ -1882,7 +1951,9 @@ class TestResolveInstallReleasePlans: max_sm = 90, ) return INSTALL_LLAMA_PREBUILT.ResolvedPublishedRelease( - bundle = make_release([art], release_tag = release_tag, upstream_tag = upstream_tag), + bundle = make_release( + [art], release_tag = release_tag, upstream_tag = upstream_tag + ), checksums = make_checksums_with_source( [asset_name], release_tag = release_tag, @@ -1890,7 +1961,9 @@ class TestResolveInstallReleasePlans: ), ) - def test_latest_collects_multiple_older_release_plans_up_to_limit(self, monkeypatch): + def test_latest_collects_multiple_older_release_plans_up_to_limit( + self, monkeypatch + ): mock_linux_runtime(monkeypatch, ["cuda12"]) host = make_host(system = "Linux", machine = "x86_64", compute_caps = ["86"]) releases = [ @@ -1902,7 +1975,9 @@ class TestResolveInstallReleasePlans: monkeypatch.setattr( INSTALL_LLAMA_PREBUILT, "iter_resolved_published_releases", - lambda requested_tag, published_repo, published_release_tag = "": iter(releases), + lambda requested_tag, published_repo, published_release_tag = "": iter( + releases + ), ) requested_tag, plans = _fork_manifest_release_plans( @@ -1917,7 +1992,9 @@ class TestResolveInstallReleasePlans: assert [plan.release_tag for plan in plans] == ["r3", "r2"] assert [plan.llama_tag for plan in plans] == ["b9003", "b9002"] - def test_latest_skips_non_installable_release_and_keeps_searching(self, monkeypatch): + def test_latest_skips_non_installable_release_and_keeps_searching( + self, monkeypatch + ): mock_linux_runtime(monkeypatch, ["cuda12"]) host = make_host(system = "Linux", machine = "x86_64", compute_caps = ["86"]) releases = [ @@ -1936,7 +2013,9 @@ class TestResolveInstallReleasePlans: monkeypatch.setattr( INSTALL_LLAMA_PREBUILT, "iter_resolved_published_releases", - lambda requested_tag, published_repo, published_release_tag = "": iter(releases), + lambda requested_tag, published_repo, published_release_tag = "": iter( + releases + ), ) _requested_tag, plans = _fork_manifest_release_plans( @@ -1953,9 +2032,13 @@ class TestResolveInstallReleasePlans: def test_malformed_release_fallback_env_uses_default(self, monkeypatch): monkeypatch.setenv("UNSLOTH_LLAMA_MAX_PREBUILT_RELEASE_FALLBACKS", "not-an-int") - assert env_int("UNSLOTH_LLAMA_MAX_PREBUILT_RELEASE_FALLBACKS", 3, minimum = 1) == 3 + assert ( + env_int("UNSLOTH_LLAMA_MAX_PREBUILT_RELEASE_FALLBACKS", 3, minimum = 1) == 3 + ) - def test_import_with_malformed_release_fallback_env_does_not_crash(self, monkeypatch): + def test_import_with_malformed_release_fallback_env_does_not_crash( + self, monkeypatch + ): monkeypatch.setenv("UNSLOTH_LLAMA_MAX_PREBUILT_RELEASE_FALLBACKS", "bad-value") spec = importlib.util.spec_from_file_location( "studio_install_llama_prebuilt_env_reload", @@ -2225,14 +2308,23 @@ class TestPinnedBlackwellCudaFallback: assert pin.runtime_sha256 and len(pin.runtime_sha256) == 64 def test_pin_offered_for_driver_13_2(self): - assert _pinned_windows_cuda_fallback(self._win_host((13, 2), ["120"]), []) is not None + assert ( + _pinned_windows_cuda_fallback(self._win_host((13, 2), ["120"]), []) + is not None + ) def test_pin_offered_for_sm121_variant(self): # sm_121 is Blackwell-family and also needs toolkit >= 12.8. - assert _pinned_windows_cuda_fallback(self._win_host((13, 1), ["121"]), []) is not None + assert ( + _pinned_windows_cuda_fallback(self._win_host((13, 1), ["121"]), []) + is not None + ) def test_pin_uses_max_of_multi_gpu_caps(self): - assert _pinned_windows_cuda_fallback(self._win_host((13, 1), ["86", "120"]), []) is not None + assert ( + _pinned_windows_cuda_fallback(self._win_host((13, 1), ["86", "120"]), []) + is not None + ) @pytest.mark.parametrize("sm", ["89", "90", "100"]) def test_pin_not_offered_to_non_blackwell(self, sm): @@ -2243,11 +2335,16 @@ class TestPinnedBlackwellCudaFallback: # b9360 is native sm_120a SASS (no JIT) and ships a cuda-13.1 cudart, # both of which run on a 13.0 r580+ driver via CUDA minor-version # compatibility. 13.0 is the mainstream Blackwell branch, so it must fire. - assert _pinned_windows_cuda_fallback(self._win_host((13, 0), ["120"]), []) is not None + assert ( + _pinned_windows_cuda_fallback(self._win_host((13, 0), ["120"]), []) + is not None + ) def test_pin_not_offered_below_floor(self): # 12.x predates Blackwell entirely; the pin stays dormant below 13.0. - assert _pinned_windows_cuda_fallback(self._win_host((12, 9), ["120"]), []) is None + assert ( + _pinned_windows_cuda_fallback(self._win_host((12, 9), ["120"]), []) is None + ) def test_pin_not_offered_without_driver(self): assert _pinned_windows_cuda_fallback(self._win_host(None, ["120"]), []) is None @@ -2321,7 +2418,10 @@ class TestPinnedBlackwellCudaFallback: ], ) def test_attempt_covers_blackwell(self, minor, covers): - assert _windows_cuda_attempt_covers_blackwell(self._win_cuda_attempt(minor)) is covers + assert ( + _windows_cuda_attempt_covers_blackwell(self._win_cuda_attempt(minor)) + is covers + ) def test_attempt_covers_blackwell_ignores_non_cuda_kind(self): cpu = AssetChoice( @@ -2359,7 +2459,9 @@ class TestPinnedBlackwellCudaFallback: ("older", "cuda12", 89, False), # 12.4 toolkit app bundle stops at Ada ], ) - def test_attempt_covers_blackwell_app_bundle(self, profile, runtime_line, max_sm, covers): + def test_attempt_covers_blackwell_app_bundle( + self, profile, runtime_line, max_sm, covers + ): # App-named bundles carry no toolkit minor; coverage is read from max_sm. attempt = self._app_attempt(profile, runtime_line, max_sm) assert _windows_cuda_attempt_covers_blackwell(attempt) is covers @@ -2397,7 +2499,8 @@ class TestDirectUpstreamBlackwellPin: return { "tag_name": self.TAG, "assets": [ - {"name": n, "browser_download_url": f"https://example.com/{n}"} for n in names + {"name": n, "browser_download_url": f"https://example.com/{n}"} + for n in names ], } @@ -2417,7 +2520,9 @@ class TestDirectUpstreamBlackwellPin: driver_cuda_version = (13, 1), compute_caps = ["120"], ) - plan = direct_upstream_release_plan(self._release(), host, UPSTREAM_REPO, "latest") + plan = direct_upstream_release_plan( + self._release(), host, UPSTREAM_REPO, "latest" + ) order = [(a.tag, a.runtime_line or a.install_kind) for a in plan.attempts] # cuda-12.4 (toolkit 12.4, no sm_120) is dropped entirely on Blackwell: # behind the pin it would still be attempted if the pin download failed, @@ -2436,7 +2541,9 @@ class TestDirectUpstreamBlackwellPin: driver_cuda_version = (13, 3), compute_caps = ["120"], ) - plan = direct_upstream_release_plan(self._release(), host, UPSTREAM_REPO, "latest") + plan = direct_upstream_release_plan( + self._release(), host, UPSTREAM_REPO, "latest" + ) assert "b9360" not in [a.tag for a in plan.attempts] assert plan.attempts[0].tag == self.TAG assert plan.attempts[0].runtime_line == "cuda13" @@ -2514,7 +2621,9 @@ class TestBlackwellCuda124Exclusion: max_sm = 80, ) assert ( - INSTALL_LLAMA_PREBUILT._drop_blackwell_incapable_windows_cuda(self._bw_host(), [bundle]) + INSTALL_LLAMA_PREBUILT._drop_blackwell_incapable_windows_cuda( + self._bw_host(), [bundle] + ) == [] ) @@ -2527,7 +2636,9 @@ class TestBlackwellCuda124Exclusion: ) attempts = [self._upstream_cuda("12.4")] assert ( - INSTALL_LLAMA_PREBUILT._drop_blackwell_incapable_windows_cuda(host, attempts) + INSTALL_LLAMA_PREBUILT._drop_blackwell_incapable_windows_cuda( + host, attempts + ) == attempts ) @@ -2641,7 +2752,9 @@ class TestPublishedWindowsCudaAttemptsDynamicMajor: # the old hardcoded cuda12/cuda13 seed would never order it (the cuda14 # line would be skipped for want of a 14.x asset in the seed). mock_windows_runtime(monkeypatch, ["cuda14", "cuda13", "cuda12"]) - release = self._release([("14.0", "cuda14"), ("13.3", "cuda13"), ("12.4", "cuda12")]) + release = self._release( + [("14.0", "cuda14"), ("13.3", "cuda13"), ("12.4", "cuda12")] + ) host = make_host( system = "Windows", machine = "AMD64", @@ -3127,14 +3240,18 @@ class TestResolveUpstreamAssetChoice: def test_linux_x86_64_cpu(self, monkeypatch): name = f"llama-{self.TAG}-bin-ubuntu-x64.tar.gz" self._mock_github_assets(monkeypatch, {name: f"https://x/{name}"}) - host = make_host(has_usable_nvidia = False, nvidia_smi = None, has_physical_nvidia = False) + host = make_host( + has_usable_nvidia = False, nvidia_smi = None, has_physical_nvidia = False + ) result = resolve_upstream_asset_choice(host, self.TAG) assert result.install_kind == "linux-cpu" assert result.name == name def test_linux_cpu_missing(self, monkeypatch): self._mock_github_assets(monkeypatch, {}) - host = make_host(has_usable_nvidia = False, nvidia_smi = None, has_physical_nvidia = False) + host = make_host( + has_usable_nvidia = False, nvidia_smi = None, has_physical_nvidia = False + ) with pytest.raises(PrebuiltFallback, match = "Linux CPU"): resolve_upstream_asset_choice(host, self.TAG) @@ -3221,7 +3338,9 @@ class TestResolveUpstreamAssetChoice: has_physical_nvidia = False, has_usable_nvidia = False, ) - with pytest.raises(PrebuiltFallback, match = "no prebuilt policy exists for Linux aarch64"): + with pytest.raises( + PrebuiltFallback, match = "no prebuilt policy exists for Linux aarch64" + ): resolve_upstream_asset_choice(host, self.TAG) def test_windows_usable_nvidia_delegates(self, monkeypatch): @@ -3347,7 +3466,9 @@ class TestResolveSimpleMacosPin: for tag in self.TAGS: yield _release(tag) - monkeypatch.setattr(INSTALL_LLAMA_PREBUILT, "iter_release_payloads_by_time", fake_iter) + monkeypatch.setattr( + INSTALL_LLAMA_PREBUILT, "iter_release_payloads_by_time", fake_iter + ) return calls def test_pre26_host_pins_b9415_without_walkback(self, monkeypatch): @@ -3399,9 +3520,13 @@ class TestLinuxArm64ForkFallsBackToSource: called["args"] = (host.machine, repo) return "b9457", ["plan"] - monkeypatch.setattr(INSTALL_LLAMA_PREBUILT, "_fork_manifest_release_plans", _full) + monkeypatch.setattr( + INSTALL_LLAMA_PREBUILT, "_fork_manifest_release_plans", _full + ) host = make_host(system = "Linux", machine = "aarch64") - tag, plans = resolve_simple_install_release_plans("latest", host, "unslothai/llama.cpp", "") + tag, plans = resolve_simple_install_release_plans( + "latest", host, "unslothai/llama.cpp", "" + ) assert called.get("args") == ("aarch64", "unslothai/llama.cpp") assert plans == ["plan"] @@ -3415,9 +3540,13 @@ class TestLinuxArm64ForkFallsBackToSource: called["args"] = (host.machine, repo) return "b9457", ["plan"] - monkeypatch.setattr(INSTALL_LLAMA_PREBUILT, "_fork_manifest_release_plans", _full) + monkeypatch.setattr( + INSTALL_LLAMA_PREBUILT, "_fork_manifest_release_plans", _full + ) host = make_host(system = "Linux", machine = "x86_64") - tag, plans = resolve_simple_install_release_plans("latest", host, "unslothai/llama.cpp", "") + tag, plans = resolve_simple_install_release_plans( + "latest", host, "unslothai/llama.cpp", "" + ) assert called.get("args") == ("x86_64", "unslothai/llama.cpp") assert plans == ["plan"] @@ -3439,7 +3568,9 @@ class TestLinuxArm64ForkFallsBackToSource: has_usable_nvidia = False, ) with pytest.raises(PrebuiltFallback) as exc: - resolve_simple_install_release_plans("latest", host, "ggml-org/llama.cpp", "") + resolve_simple_install_release_plans( + "latest", host, "ggml-org/llama.cpp", "" + ) assert "linux-x64 prebuilts" not in str(exc.value) @@ -3525,7 +3656,9 @@ class TestCpuFallback: has_physical_nvidia = False, has_usable_nvidia = False, ) - plan = direct_upstream_release_plan(release, cpu_host, "ggml-org/llama.cpp", "latest") + plan = direct_upstream_release_plan( + release, cpu_host, "ggml-org/llama.cpp", "latest" + ) assert plan.attempts[0].install_kind == "linux-arm64" assert plan.attempts[0].name == f"llama-{tag}-bin-ubuntu-arm64.tar.gz" @@ -3771,7 +3904,9 @@ class TestCudaDriverToolkitMismatchMessage: nvcc.chmod(0o755) return nvcc - def test_setup_sh_major_mismatch_uses_newest_compatible_detected_toolkit(self, tmp_path): + def test_setup_sh_major_mismatch_uses_newest_compatible_detected_toolkit( + self, tmp_path + ): blocked_nvcc = self._fake_nvcc(tmp_path, "13.3") older_nvcc = self._fake_nvcc(tmp_path, "12.6") compatible_nvcc = self._fake_nvcc(tmp_path, "12.8") @@ -3955,7 +4090,9 @@ class TestCudaDriverToolkitMismatchMessage: assert "GPU_BACKEND=cuda" in output assert "ALLOWED=true" in output - def test_setup_sh_compatible_finder_rejects_newer_major_only_candidate(self, tmp_path): + def test_setup_sh_compatible_finder_rejects_newer_major_only_candidate( + self, tmp_path + ): # Only alternative is still newer-major than the driver: finder must fail, not pick it. blocked_nvcc = self._fake_nvcc(tmp_path, "13.3") other_newer_nvcc = self._fake_nvcc(tmp_path, "13.1") diff --git a/tests/studio/load_freeze/llama_server_shim.py b/tests/studio/load_freeze/llama_server_shim.py index 4030fcd6c4..2c0b8ff1b1 100644 --- a/tests/studio/load_freeze/llama_server_shim.py +++ b/tests/studio/load_freeze/llama_server_shim.py @@ -123,7 +123,9 @@ class _Handler(BaseHTTPRequestHandler): self._send_raw(srv.config.detok_status, srv.config.detok_body) return tids = body.get("tokens") or [] - content = "".join(srv.config.detok_map.get(int(t), f"") for t in tids) + content = "".join( + srv.config.detok_map.get(int(t), f"") for t in tids + ) self._send_json(srv.config.detok_status, {"content": content}) return if path == "/completion": @@ -226,7 +228,9 @@ class FakeLlamaServer: def start(self) -> "FakeLlamaServer": # port=0 lets ThreadingHTTPServer pick a free port atomically (no # find-then-bind race); read back via server_address[1]. - self._server = FakeLlamaServer._Server((self.host, self._requested_port), _Handler) + self._server = FakeLlamaServer._Server( + (self.host, self._requested_port), _Handler + ) self._server.config = self.config bound_port = self._server.server_address[1] self._thread = threading.Thread( diff --git a/tests/studio/load_freeze/test_load_orchestrator.py b/tests/studio/load_freeze/test_load_orchestrator.py index d091e56e70..a0ef080bc4 100644 --- a/tests/studio/load_freeze/test_load_orchestrator.py +++ b/tests/studio/load_freeze/test_load_orchestrator.py @@ -110,7 +110,9 @@ class _UvicornServerThread: self.host = host self.port = port - cfg = uvicorn.Config(app, host = host, port = port, log_level = "warning", access_log = False) + cfg = uvicorn.Config( + app, host = host, port = port, log_level = "warning", access_log = False + ) self._server = uvicorn.Server(cfg) self._server.install_signal_handlers = lambda: None # type: ignore[assignment] self._thread: threading.Thread | None = None @@ -214,7 +216,9 @@ def test_buggy_route_blocks_event_loop(): app = _build_app(backend, wrap_in_thread = False) port = _free_port() with _UvicornServerThread(app, port = port) as uv: - max_lat, probe_t, _ = _drive_concurrent_probe_and_health(f"http://127.0.0.1:{uv.port}") + max_lat, probe_t, _ = _drive_concurrent_probe_and_health( + f"http://127.0.0.1:{uv.port}" + ) assert probe_t >= 0.5 assert max_lat >= 0.4, f"expected >=0.4s stall, got {max_lat:.3f}s" @@ -437,7 +441,9 @@ def test_50_concurrent_probes_complete_without_deadlock(): with ThreadPoolExecutor(max_workers = 50) as pool: futs = [ pool.submit( - lambda: httpx.get(f"http://127.0.0.1:{uv.port}/probe", timeout = 30.0) + lambda: httpx.get( + f"http://127.0.0.1:{uv.port}/probe", timeout = 30.0 + ) ) for _ in range(50) ] diff --git a/tests/studio/playwright_chat_ime_i18n.py b/tests/studio/playwright_chat_ime_i18n.py index f27266193d..65d80770b8 100644 --- a/tests/studio/playwright_chat_ime_i18n.py +++ b/tests/studio/playwright_chat_ime_i18n.py @@ -240,7 +240,8 @@ with sync_playwright() as p: dir_attr = composer.evaluate("(el) => el.getAttribute('dir')") if dir_attr != "auto": soft_fail( - f'composer is missing dir="auto" (got {dir_attr!r}); RTL ' "languages will render LTR." + f'composer is missing dir="auto" (got {dir_attr!r}); RTL ' + "languages will render LTR." ) else: info('composer dir="auto" present') @@ -251,7 +252,9 @@ with sync_playwright() as p: _thread_src = ( _repo_root / "studio/frontend/src/components/assistant-ui/thread.tsx" ).read_text() - _shared_src = (_repo_root / "studio/frontend/src/features/chat/shared-composer.tsx").read_text() + _shared_src = ( + _repo_root / "studio/frontend/src/features/chat/shared-composer.tsx" + ).read_text() _edit_idx = _thread_src.find("aui-edit-composer-input") if _edit_idx == -1 or 'dir="auto"' not in _thread_src[_edit_idx : _edit_idx + 600]: soft_fail('edit composer source is missing dir="auto"') @@ -260,7 +263,8 @@ with sync_playwright() as p: _compare_idx = _shared_src.find("Send to both models") if ( _compare_idx == -1 - or 'dir="auto"' not in _shared_src[max(_compare_idx - 400, 0) : _compare_idx + 400] + or 'dir="auto"' + not in _shared_src[max(_compare_idx - 400, 0) : _compare_idx + 400] ): soft_fail('compare composer source is missing dir="auto"') else: @@ -467,7 +471,9 @@ with sync_playwright() as p: # IME keydown (isComposing=true / keyCode 229) must not slip preedit text # through submit. The onKeyDown gate re-pins composingRef so handleSubmit # refuses at form.requestSubmit() time, not at the (enabled) button. - step("BUG REPRO: keydown re-pin after watchdog cleared composing (issue #5546 follow-up)") + step( + "BUG REPRO: keydown re-pin after watchdog cleared composing (issue #5546 follow-up)" + ) clear() composer.click() composer.evaluate( @@ -513,7 +519,9 @@ with sync_playwright() as p: "Form submitted after an IME keydown -- preedit text leaked " "through the watchdog gap (#5546 follow-up regression)." ) - info(f"Form submit refused after IME keydown; textarea retained {submit_probe.get('after')!r}") + info( + f"Form submit refused after IME keydown; textarea retained {submit_probe.get('after')!r}" + ) shoot("06c-keydown-repin") info("keydown re-pin gate PASS") clear() diff --git a/tests/studio/playwright_chat_ui.py b/tests/studio/playwright_chat_ui.py index 4739d519d9..0ab9194fa5 100644 --- a/tests/studio/playwright_chat_ui.py +++ b/tests/studio/playwright_chat_ui.py @@ -141,7 +141,10 @@ def expected_default_model(): for node in tree.body: if not isinstance(node, ast.Assign): continue - if not any(isinstance(t, ast.Name) and t.id == "DEFAULT_MODELS_GGUF" for t in node.targets): + if not any( + isinstance(t, ast.Name) and t.id == "DEFAULT_MODELS_GGUF" + for t in node.targets + ): continue try: models = ast.literal_eval(node.value) @@ -318,7 +321,9 @@ with sync_playwright() as p: form_err: Exception | None = None for _form_attempt in range(3): try: - page.goto(f"{BASE}/change-password", wait_until = "domcontentloaded", timeout = 60_000) + page.goto( + f"{BASE}/change-password", wait_until = "domcontentloaded", timeout = 60_000 + ) try: page.wait_for_load_state("networkidle", timeout = 30_000) except Exception: @@ -377,7 +382,9 @@ with sync_playwright() as p: flush = True, ) if page_errors: - print(f"[ui] first pageerror: {page_errors[0][:200]!r}", flush = True) + print( + f"[ui] first pageerror: {page_errors[0][:200]!r}", flush = True + ) try: shoot(f"01-change-password-attempt-{_form_attempt + 1}-fail") except Exception: @@ -451,7 +458,9 @@ with sync_playwright() as p: flush = True, ) if page_errors: - print(f"[ui] first pageerror: {page_errors[0][:200]!r}", flush = True) + print( + f"[ui] first pageerror: {page_errors[0][:200]!r}", flush = True + ) try: shoot(f"03-composer-wait-attempt-{_attempt + 1}-fail") except Exception: @@ -548,7 +557,9 @@ with sync_playwright() as p: try: sel_text = (selector_btn.text_content(timeout = 2_000) or "").strip() except Exception as _sel_err: - info(f"WARN: model-selector probe skipped: {type(_sel_err).__name__}: {_sel_err}") + info( + f"WARN: model-selector probe skipped: {type(_sel_err).__name__}: {_sel_err}" + ) if sel_text: info(f"model selector button text: {sel_text!r}") shoot("03b-default-model-button") @@ -584,7 +595,10 @@ with sync_playwright() as p: if load_resp.get("error"): fail(f"/api/inference/load wedged: {load_resp['error']!r}") if load_resp["status"] != 200: - fail(f"/api/inference/load returned {load_resp['status']}: " f"{load_resp.get('body')!r}") + fail( + f"/api/inference/load returned {load_resp['status']}: " + f"{load_resp.get('body')!r}" + ) info(f"loaded model: {(load_resp['body'] or {}).get('display_name')}") # Studio caches the per-context model state in zustand; reload @@ -831,7 +845,8 @@ with sync_playwright() as p: # Look for either "Disable X" or "Enable X" -- whichever # is currently rendered. toggle = page.locator( - f'button[aria-label="Disable {feature}"], ' f'button[aria-label="Enable {feature}"]' + f'button[aria-label="Disable {feature}"], ' + f'button[aria-label="Enable {feature}"]' ).first if toggle.count() == 0: info(f"toggle '{feature}' not present on this layout") @@ -847,7 +862,8 @@ with sync_playwright() as p: page.wait_for_timeout(200) after = ( page.locator( - f'button[aria-label="Disable {feature}"], ' f'button[aria-label="Enable {feature}"]' + f'button[aria-label="Disable {feature}"], ' + f'button[aria-label="Enable {feature}"]' ).first.get_attribute("aria-label") or "" ) @@ -858,7 +874,8 @@ with sync_playwright() as p: # Flip back so test state is unchanged. try: page.locator( - f'button[aria-label="Disable {feature}"], ' f'button[aria-label="Enable {feature}"]' + f'button[aria-label="Disable {feature}"], ' + f'button[aria-label="Enable {feature}"]' ).first.click() except Exception: pass @@ -951,7 +968,8 @@ with sync_playwright() as p: except Exception as exc: if attempt == 1: soft_fail( - f"theme cycle {cycle + 1}: account-menu click failed " f"({exc!r})" + f"theme cycle {cycle + 1}: account-menu click failed " + f"({exc!r})" ) continue try: @@ -1002,7 +1020,8 @@ with sync_playwright() as p: if click_err is not None: page.keyboard.press("Escape") soft_fail( - f"theme cycle {cycle + 1}: theme menuitem click failed " f"({click_err!r})" + f"theme cycle {cycle + 1}: theme menuitem click failed " + f"({click_err!r})" ) break # Settle. The ".dark" class on is the ground @@ -1059,7 +1078,9 @@ with sync_playwright() as p: # progressively more permissive locators so the test stays # green on both platforms. candidates = [ - page.get_by_role("button", name = re.compile(rf"^\s*{label}\s*$", re.I)).first, + page.get_by_role( + "button", name = re.compile(rf"^\s*{label}\s*$", re.I) + ).first, page.locator(f'button:has-text("{label}")').first, page.locator(f'a:has-text("{label}")').first, page.locator(f'[data-sidebar="menu-button"]:has-text("{label}")').first, @@ -1095,15 +1116,21 @@ with sync_playwright() as p: click_nav("New Chat", r"/chat") shoot("11-new-chat") # Compare moved into the composer + menu (Tools and attachments). - plus_btn = page.get_by_role("button", name = re.compile(r"Tools and attachments", re.I)).first + plus_btn = page.get_by_role( + "button", name = re.compile(r"Tools and attachments", re.I) + ).first if plus_btn.count() > 0: plus_btn.click(force = True) page.wait_for_timeout(400) - compare_item = page.get_by_role("menuitem", name = re.compile(r"Compare chat", re.I)).first + compare_item = page.get_by_role( + "menuitem", name = re.compile(r"Compare chat", re.I) + ).first if compare_item.count() == 0: # The plus menu was decluttered: Compare chat now lives in the # "More" submenu; hover (then click as fallback) to open it. - more_trigger = page.get_by_role("menuitem", name = re.compile(r"^More$", re.I)).first + more_trigger = page.get_by_role( + "menuitem", name = re.compile(r"^More$", re.I) + ).first if more_trigger.count() > 0: more_trigger.hover() page.wait_for_timeout(400) @@ -1150,7 +1177,9 @@ with sync_playwright() as p: step("Developer (API) tab via account menu") acct.click() page.wait_for_timeout(400) - dev = page.get_by_role("menuitem", name = re.compile(r"developer|api", re.I)).first + dev = page.get_by_role( + "menuitem", name = re.compile(r"developer|api", re.I) + ).first if dev.count() > 0: dev.click() page.wait_for_timeout(800) @@ -1167,7 +1196,9 @@ with sync_playwright() as p: re.compile(r"api keys|developer", re.I), ).first if keys_section.count() > 0: - info(f"OK API tab text: {(keys_section.text_content() or '').strip()[:80]!r}") + info( + f"OK API tab text: {(keys_section.text_content() or '').strip()[:80]!r}" + ) # Close dialog with Escape. page.keyboard.press("Escape") page.wait_for_timeout(300) @@ -1185,7 +1216,9 @@ with sync_playwright() as p: page.wait_for_timeout(1500) # Recipe cards are rendered as or button elements; count # all clickable headings under main + screenshot. - headings = page.locator("main h2, main h3, [data-recipe], a[href*='/data-recipes/']") + headings = page.locator( + "main h2, main h3, [data-recipe], a[href*='/data-recipes/']" + ) n_cards = headings.count() info(f"Recipes route headings/cards: {n_cards}") shoot("15b-recipes-cards") @@ -1274,7 +1307,10 @@ with sync_playwright() as p: info(f"recent-thread click {i} failed: {_click_err!s}") continue if not clicked_recent: - soft_fail(f"no Recents entry was clickable within 30s deadline " f"(n_threads={n_threads})") + soft_fail( + f"no Recents entry was clickable within 30s deadline " + f"(n_threads={n_threads})" + ) # Back to chat. page.goto(f"{BASE}/chat") composer = page.locator('textarea[aria-label="Message input"]') diff --git a/tests/studio/playwright_extra_ui.py b/tests/studio/playwright_extra_ui.py index 26c3c244ca..6434f64bd3 100644 --- a/tests/studio/playwright_extra_ui.py +++ b/tests/studio/playwright_extra_ui.py @@ -170,7 +170,9 @@ with sync_playwright() as p: form_err: Exception | None = None for _form_attempt in range(3): try: - page.goto(f"{BASE}/change-password", wait_until = "domcontentloaded", timeout = 60_000) + page.goto( + f"{BASE}/change-password", wait_until = "domcontentloaded", timeout = 60_000 + ) try: page.wait_for_load_state("networkidle", timeout = 30_000) except Exception: @@ -327,14 +329,20 @@ with sync_playwright() as p: step("Compare tab: send to two panes") # Compare moved into the composer + menu (Tools and attachments). compare_opened = False - plus_btn = page.get_by_role("button", name = re.compile(r"Tools and attachments", re.I)).first + plus_btn = page.get_by_role( + "button", name = re.compile(r"Tools and attachments", re.I) + ).first if plus_btn.count() > 0: plus_btn.click(force = True) page.wait_for_timeout(400) - compare_item = page.get_by_role("menuitem", name = re.compile(r"Compare chat", re.I)).first + compare_item = page.get_by_role( + "menuitem", name = re.compile(r"Compare chat", re.I) + ).first if compare_item.count() == 0: # Compare chat moved into the "More" submenu; hover, then click fallback. - more_trigger = page.get_by_role("menuitem", name = re.compile(r"^More$", re.I)).first + more_trigger = page.get_by_role( + "menuitem", name = re.compile(r"^More$", re.I) + ).first if more_trigger.count() > 0: more_trigger.hover() page.wait_for_timeout(400) @@ -425,7 +433,9 @@ with sync_playwright() as p: arg = ok_count_before + 4, timeout = 60_000, ) - info("OK Compare: 4 total new assistant bubbles after second prompt") + info( + "OK Compare: 4 total new assistant bubbles after second prompt" + ) except Exception as exc: runtime_warn( f"Compare: 4 bubbles didn't appear (panes likely " @@ -446,7 +456,9 @@ with sync_playwright() as p: page.wait_for_timeout(1500) shoot("05-recipes-list") # Template cards render as