fix up waste script

old-commit-hash: 263d23ff3c

scripts/waste.py

@@ -1,13 +1,12 @@
 #!/usr/bin/env python3
-from multiprocessing import Process
-from setproctitle import setproctitle  # pylint: disable=no-name-in-module
 import os
 import numpy as np
 from common.realtime import sec_since_boot
+from multiprocessing import Process
+from setproctitle import setproctitle  # pylint: disable=no-name-in-module
 
-def waste(pid):
-  # set affinity
-  os.system("taskset -p %d %d" % (1 << pid, os.getpid()))
+def waste(core):
+  os.sched_setaffinity(0, [core,])
 
   m1 = np.zeros((200, 200)) + 0.8
   m2 = np.zeros((200, 200)) + 1.2
@@ -17,16 +16,16 @@ def waste(pid):
   j = 0
   while 1:
     if (i % 100) == 0:
-      setproctitle("%3d: %8d" % (pid, i))
+      setproctitle("%3d: %8d" % (core, i))
       lt = sec_since_boot()
-      print("%3d: %8d %f  %.2f" % (pid, i, lt-st, j))
+      print("%3d: %8d %f  %.2f" % (core, i, lt-st, j))
       st = lt
     i += 1
     j = np.sum(np.matmul(m1, m2))
 
 def main(gctx=None):
   print("1-2 seconds is baseline")
-  for i in range(4):
+  for i in range(os.cpu_count()):
     p = Process(target=waste, args=(i,))
     p.start()