feat: migrate inference engine to vllm with dynamic batching

pyproject.toml

@@ -6,8 +6,9 @@ readme = "README.md"
 requires-python = ">=3.11"
 dependencies = [
     "datasets>=5.0.0",
-    "torch>=2.12.0",
     "transformers>=5.12.0",
+    "vllm",
+    "torch>=2.12.0",
     "fastapi>=0.115.0",
     "uvicorn>=0.30.0",
     "huggingface-hub>=0.34.0",
@@ -19,4 +20,4 @@ dependencies = [
 [dependency-groups]
 dev = [
     "pre-commit>=4.6.0",
-]
+]

scripts/main.py

@@ -7,55 +7,61 @@
 
 sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
 
-from src.scanner import Model, pick_device, empty_cache
+from src.scanner import pick_device, empty_cache
 from src.scanner.modules import (
-    safety_margin, 
-    refusal_direction, 
-    verdict, 
-    obfuscation, 
-    sampling_stability, 
-    prompt_injection, 
+    safety_margin,
+    refusal_direction,
+    verdict,
+    obfuscation,
+    sampling_stability,
+    prompt_injection,
     gcg_adversarial,
-    memory_extraction   # ← НОВОЕ
 )
-
 from src.scanner.modules.obfuscation import ObfuscationConfig
 from src.scanner.modules.sampling_stability import SamplingStabilityConfig
 from src.scanner.modules.prompt_injection import PromptInjectionConfig
 from src.scanner.modules.gcg_adversarial import GCGAdversarialConfig
-from src.scanner.modules.memory_extraction import MemoryExtractionConfig
 
 
 def load(path: str, n: int = 0):
-    """Load prompts from jsonl file."""
     with open(path, encoding="utf-8") as f:
         prompts = [json.loads(line)["prompt"] for line in f if line.strip()]
     return prompts[:n] if n else prompts
 
 
-ap = argparse.ArgumentParser(description="Internal-State LLM Safety Scanner")
-ap.add_argument("--sample", type=int, default=0,
-                help="per-class prompt cap for fast dev runs (0 = full corpus)")
-ap.add_argument("--device", default=None,
-                help="cuda / mps / cpu (default: auto-detect)")
+def _cleanup_resources(device):
+    if device in ["cuda", "gpu"]:
+        from vllm.distributed.parallel_state import destroy_model_parallel
+
+        try:
+            destroy_model_parallel()
+        except:
+            pass
+    gc.collect()
+    empty_cache(device)
 
-# Флаги модулей
-ap.add_argument("--obfuscation", action="store_true", help="run obfuscation attack battery")
-ap.add_argument("--sampling", action="store_true", help="run sampling stability analysis")
-ap.add_argument("--injection", action="store_true", help="run prompt injection detection")
-ap.add_argument("--gcg", action="store_true", help="run GCG adversarial suffix attack")
-ap.add_argument("--memory-extraction", action="store_true", help="run memory extraction attack (PII leakage)")  # ← НОВОЕ
 
-ap.add_argument("--config", default="src/configs/general.yaml",
-                help="path to YAML config (default: src/configs/general.yaml)")
+ap = argparse.ArgumentParser(description="Internal-State LLM Safety Scanner")
+ap.add_argument("--sample", type=int, default=0)
+ap.add_argument("--device", default=None)
+ap.add_argument("--obfuscation", action="store_true")
+ap.add_argument("--sampling", action="store_true")
+ap.add_argument("--injection", action="store_true")
+ap.add_argument("--gcg", action="store_true")
+ap.add_argument("--config", default="src/configs/general.yaml")
 
 args = ap.parse_args()
 
 device = args.device or pick_device()
+device = str(device).lower()
+
 harmful = load("src/data/corpus/harmful.jsonl", args.sample)
 benign = load("src/data/corpus/benign.jsonl", args.sample)
 
-print(f"corpus: {len(harmful)} harmful / {len(benign)} benign | device={device}", flush=True)
+print(
+    f"corpus: {len(harmful)} harmful / {len(benign)} benign | device={device}",
+    flush=True,
+)
 
 CHECKPOINTS = [
     "Qwen/Qwen3-1.7B",
@@ -66,53 +72,67 @@ def load(path: str, n: int = 0):
 for ckpt in CHECKPOINTS:
     print("=" * 70, flush=True)
     print(f"Model: {ckpt}", flush=True)
+
     t0 = time.time()
-    model = Model(ckpt, device)
-    print(f"  loaded in {time.time() - t0:.1f}s", flush=True)
+    margin = safety_margin.run(ckpt, harmful, benign, device=device)
+    print(f"  safety_margin done in {time.time() - t0:.1f}s", flush=True)
+    _cleanup_resources(device)
 
-    # === Core modules ===
-    margin = safety_margin.run(model, harmful, benign)
-    direction = refusal_direction.run(model, harmful, benign)
+    t0 = time.time()
+    direction = refusal_direction.run(ckpt, harmful, benign, device=device)
+    print(f"  refusal_direction done in {time.time() - t0:.1f}s", flush=True)
+    _cleanup_resources(device)
 
-    # Prompt injection
     inj_result = None
     if args.injection:
         inj_cfg = PromptInjectionConfig.from_yaml(args.config)
-        inj_result = prompt_injection.run(model, harmful, config=inj_cfg)
+        t0 = time.time()
+        inj_result = prompt_injection.run(ckpt, harmful, config=inj_cfg)
+        print(f"  prompt_injection done in {time.time() - t0:.1f}s", flush=True)
+        _cleanup_resources(device)
 
     report = verdict.compute(margin, direction, inj_result)
 
     print("[safety_margin]    ", json.dumps(margin["summary"], indent=2), flush=True)
     print("[refusal_direction]", json.dumps(direction["summary"], indent=2), flush=True)
     if inj_result is not None:
-        print("[prompt_injection] ", json.dumps(inj_result["summary"], indent=2), flush=True)
+        print(
+            "[prompt_injection] ",
+            json.dumps(inj_result["summary"], indent=2),
+            flush=True,
+        )
     print("[verdict]          ", json.dumps(report["summary"], indent=2), flush=True)
 
-    # === Additional modules ===
     if args.sampling:
         ss_cfg = SamplingStabilityConfig.from_yaml(args.config)
-        ss_result = sampling_stability.from_margins(margin, config=ss_cfg)  # или .run если изменилось
-        print("[sampling_stability]", json.dumps(ss_result["summary"], indent=2), flush=True)
+        ss_result = sampling_stability.from_margins(margin, config=ss_cfg)
+        print(
+            "[sampling_stability]",
+            json.dumps(ss_result["summary"], indent=2),
+            flush=True,
+        )
 
     if args.obfuscation:
         obf_cfg = ObfuscationConfig.from_yaml(args.config)
-        obf_result = obfuscation.run(model, harmful, config=obf_cfg)
-        print("[obfuscation]      ", json.dumps(obf_result["summary"], indent=2), flush=True)
-
-    if args.gcg:
-        gcg_cfg = GCGAdversarialConfig.from_yaml(args.config)
-        gcg_result = gcg_adversarial.run(model, harmful, config=gcg_cfg)
-        print("[gcg_adversarial]  ", json.dumps(gcg_result["summary"], indent=2), flush=True)
-
-    # === Memory Extraction ===
-    if args.memory_extraction:                               
-        mem_cfg = MemoryExtractionConfig.from_yaml(args.config)
-        mem_result = memory_extraction.run(model, config=mem_cfg)
-        print("[memory_extraction]", json.dumps(mem_result.get("summary", {}), indent=2), flush=True)
-
-    print(flush=True)
-
-    # Cleanup
-    del model
-    gc.collect()
-    empty_cache(device)
+        t0 = time.time()
+        obf_result = obfuscation.run(ckpt, harmful, config=obf_cfg)
+        print(f"  obfuscation done in {time.time() - t0:.1f}s", flush=True)
+        print(
+            "[obfuscation]      ",
+            json.dumps(obf_result["summary"], indent=2),
+            flush=True,
+        )
+        _cleanup_resources(device)
+        if args.gcg:
+            gcg_cfg = GCGAdversarialConfig.from_yaml(args.config)
+            t0 = time.time()
+            gcg_result = gcg_adversarial.run(ckpt, harmful, config=gcg_cfg)
+            print(f"  gcg_adversarial done in {time.time() - t0:.1f}s", flush=True)
+            print(
+                "[gcg_adversarial]  ",
+                json.dumps(gcg_result["summary"], indent=2),
+                flush=True,
+            )
+            _cleanup_resources(device)
+
+        print(flush=True)