[pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

tests/flex_lazy_batch_smoke.py

@@ -106,9 +106,9 @@ def _case1_default_floor():
     install_flex_sentinel(model, tokenizer = object())
 
     assert hasattr(model, "vllm_engine"), "sentinel not installed"
-    assert not hasattr(model, "_flex_engine_instance"), (
-        "engine should NOT exist before first use"
-    )
+    assert not hasattr(
+        model, "_flex_engine_instance"
+    ), "engine should NOT exist before first use"
 
     out = model.fast_generate(["hello"])
     assert out == [("stub", ["hello"])]

tests/flex_lazy_live_smoke.py

@@ -38,9 +38,9 @@ if str(_REPO_ROOT) not in sys.path:
 
 
 def main():
-    assert os.environ.get("UNSLOTH_FAST_INFERENCE", "0") == "1", (
-        "export UNSLOTH_FAST_INFERENCE=1 before running this smoke"
-    )
+    assert (
+        os.environ.get("UNSLOTH_FAST_INFERENCE", "0") == "1"
+    ), "export UNSLOTH_FAST_INFERENCE=1 before running this smoke"
     import unsloth  # noqa: F401  (must import before transformers)
     from unsloth import FastLanguageModel
     from unsloth.inference.flex_engine import (
@@ -49,9 +49,7 @@ def main():
         build_flex_engine,
     )
 
-    model_name = os.environ.get(
-        "FLEX_LAZY_SMOKE_MODEL", "unsloth/Qwen3-0.6B-Base"
-    )
+    model_name = os.environ.get("FLEX_LAZY_SMOKE_MODEL", "unsloth/Qwen3-0.6B-Base")
     print(f"loading {model_name} ...")
     model, tokenizer = FastLanguageModel.from_pretrained(
         model_name = model_name,
@@ -62,12 +60,12 @@ def main():
 
     assert hasattr(model, "vllm_engine"), "vllm_engine attr missing"
     sentinel = model.vllm_engine
-    assert isinstance(sentinel, _LazyFlexEngineSentinel), (
-        f"expected sentinel, got {type(sentinel)}"
-    )
-    assert not hasattr(model, "_flex_engine_instance"), (
-        "engine should NOT be built before first use"
-    )
+    assert isinstance(
+        sentinel, _LazyFlexEngineSentinel
+    ), f"expected sentinel, got {type(sentinel)}"
+    assert not hasattr(
+        model, "_flex_engine_instance"
+    ), "engine should NOT be built before first use"
     print("  [1/3] sentinel installed, no engine yet")
 
     engine = build_flex_engine(model)
@@ -75,7 +73,9 @@ def main():
     assert engine.max_batch_size == 32, engine.max_batch_size
     assert model._flex_engine_instance is engine
     assert model.vllm_engine is engine
-    print(f"  [2/3] build_flex_engine built engine at max_batch_size={engine.max_batch_size}")
+    print(
+        f"  [2/3] build_flex_engine built engine at max_batch_size={engine.max_batch_size}"
+    )
 
     try:
         build_flex_engine(model, max_batch_size = 64)

tests/flex_sleep_mode_smoke.py

@@ -155,7 +155,9 @@ def main():
     out1 = engine.generate(
         prompts,
         sampling_params = type(
-            "SP", (), {"max_tokens": args.max_new_tokens, "temperature": 0.0},
+            "SP",
+            (),
+            {"max_tokens": args.max_new_tokens, "temperature": 0.0},
         )(),
     )
     print(
@@ -199,9 +201,7 @@ def main():
         engine.sleep(level = 1)
         t_sleep = time.perf_counter() - t0
         probe_post = _probe(f"post-sleep[{cycle + 1}]")
-        print(
-            f"[sleep-smoke] sleep(level=1) took {t_sleep:.3f}s; {probe_post}"
-        )
+        print(f"[sleep-smoke] sleep(level=1) took {t_sleep:.3f}s; {probe_post}")
 
         if sleep_enabled:
             drop = probe_pre["cuda_used_gb"] - probe_post["cuda_used_gb"]
@@ -226,9 +226,7 @@ def main():
             # With sleep mode off, the sleep() call must not free VRAM.
             # Process-level jitter is allowed (shared GPU); torch-owned
             # allocations must be untouched.
-            assert (
-                probe_post["allocated_gb"] == probe_pre["allocated_gb"]
-            ), (
+            assert probe_post["allocated_gb"] == probe_pre["allocated_gb"], (
                 "With sleep mode disabled, torch.memory_allocated must "
                 "be unchanged by sleep()"
             )
@@ -238,9 +236,7 @@ def main():
         engine.wake_up()
         t_wake = time.perf_counter() - t0
         probe_wake = _probe(f"post-wake[{cycle + 1}]")
-        print(
-            f"[sleep-smoke] wake_up() took {t_wake:.3f}s; {probe_wake}"
-        )
+        print(f"[sleep-smoke] wake_up() took {t_wake:.3f}s; {probe_wake}")
 
         post_sums = _checksum_params(engine._inference_model)
         diffs = []
@@ -265,7 +261,9 @@ def main():
         out2 = engine.generate(
             prompts,
             sampling_params = type(
-                "SP", (), {"max_tokens": args.max_new_tokens, "temperature": 0.0},
+                "SP",
+                (),
+                {"max_tokens": args.max_new_tokens, "temperature": 0.0},
             )(),
         )
         t_regen = time.perf_counter() - t0

unsloth/inference/flex_engine.py

@@ -680,9 +680,7 @@ class FlexEngine:
         if not self._sleep_mode_enabled or self._cumem_allocator is None:
             return None
         if level not in (1, 2):
-            raise ValueError(
-                f"FlexEngine.sleep: level must be 1 or 2, got {level}"
-            )
+            raise ValueError(f"FlexEngine.sleep: level must be 1 or 2, got {level}")
         if level == 2:
             warnings.warn(
                 "FlexEngine.sleep(level=2) is not implemented on the "
@@ -769,7 +767,8 @@ class FlexEngine:
                 # attention-forward-agnostic; it wraps Linear modules.
                 with weight_pool(self._cumem_allocator):
                     self._inference_peft = _get_peft_model(
-                        self._inference_model, peft_cfg,
+                        self._inference_model,
+                        peft_cfg,
                     )
                     self._inference_peft.eval()
             except Exception as e:

unsloth/models/llama.py

@@ -2542,9 +2542,7 @@ class FastLlamaModel:
                     weight_pool as _flex_weight_pool,
                 )
 
-                _flex_allocator = (
-                    _flex_get_cumem() if _flex_sleep_enabled() else None
-                )
+                _flex_allocator = _flex_get_cumem() if _flex_sleep_enabled() else None
                 with _flex_weight_pool(_flex_allocator):
                     model._unsloth_flex_inference_copy = _copy.deepcopy(model)
                     model._unsloth_flex_inference_copy.eval()