Option to set processor affinity

Author: turboderp

example_chatbot.py

@@ -65,6 +65,10 @@
 
 model_init.print_options(args, print_opts)
 
+# Globals
+
+model_init.set_globals(args)
+
 # Load prompt file
 
 username = args.username

globals.py

@@ -0,0 +1,22 @@
+import os
+
+def set_affinity_mask(affinity_mask = None):
+
+    if affinity_mask is None:
+        cpu_count = os.cpu_count()
+        affinity_mask = set(range(cpu_count))
+
+    os.sched_setaffinity(0, affinity_mask)
+
+
+def set_affinity_list(affinity_list = None):
+
+    if affinity_list is None: set_affinity_mask(None)
+    else: set_affinity_mask(set(affinity_list))
+
+
+def set_affinity_str(affinity_str = None):
+
+    if affinity_str is None or affinity_str.isspace(): set_affinity_mask(None)
+    aff = [int(alloc) for alloc in affinity_str.split(",")]
+    set_affinity_list(aff)

model_init.py

@@ -2,6 +2,7 @@
 from tokenizer import ExLlamaTokenizer
 import argparse, sys, os, glob
 from torch import version as torch_version
+from globals import set_affinity_str
 
 def add_args(parser):
 
@@ -33,6 +34,8 @@ def add_args(parser):
     parser.add_argument("-fh2", "--force_half2", action = "store_true", help = "Force enable half2 even if unsupported")
     parser.add_argument("-cs", "--concurrent_streams", action = "store_true", help = "Use concurrent CUDA streams")
 
+    parser.add_argument("-aff", "--affinity", type = str, help = "Comma-separated list, sets processor core affinity. E.g.: -aff 0,1,2,3")
+
 
 def post_parse(args):
 
@@ -72,6 +75,7 @@ def print_options(args, extra_options = None):
     print_opts = []
     if args.gpu_split is not None: print_opts.append(f"gpu_split: {args.gpu_split}")
     if args.gpu_peer_fix: print_opts.append("gpu_peer_fix")
+    if args.affinity: print_opts.append(f" --affinity: {args.affinity}")
 
     if extra_options is not None: print_opts += extra_options
 
@@ -139,6 +143,13 @@ def make_config(args):
     return config
 
 
+# Global state
+
+def set_globals(args):
+
+    if args.affinity: set_affinity_str(args.affinity)
+
+
 # Print stats after loading model
 
 def print_stats(model):

test_benchmark_inference.py

@@ -118,6 +118,10 @@ def mem(name, total = False):
 
 model_init.print_options(args, print_opts)
 
+# Globals
+
+model_init.set_globals(args)
+
 # Instantiate model
 
 config = model_init.make_config(args)

webui/app.py

@@ -138,6 +138,8 @@ def api_append_block():
 model_init.print_options(args)
 config = model_init.make_config(args)
 
+model_init.set_globals(args)
+
 print(f" -- Loading model...")
 model = ExLlama(config)