From 804c45b5fc17fc6b9032dbfde7ba12bcf141ba7b Mon Sep 17 00:00:00 2001
From: Yahya Lmallas <1588632+Maroc-OS@users.noreply.github.com>
Date: Wed, 14 Jun 2023 20:32:17 +0100
Subject: [PATCH] FIX: Can't pickle local object (#979)

_early_exec_process is a local function that is defined whiting the scope of another function, should be global
---
 extra/helpers.py | 17 +++++++++--------
 1 file changed, 9 insertions(+), 8 deletions(-)

diff --git a/extra/helpers.py b/extra/helpers.py
index 7a1daae15e..92376e0a6a 100644
--- a/extra/helpers.py
+++ b/extra/helpers.py
@@ -1,13 +1,15 @@
 from tinygrad.helpers import Timing
+import subprocess
+import multiprocessing
+
+def _early_exec_process(qin, qout):
+  while True:
+    path, inp = qin.get()
+    qout.put(subprocess.check_output(path, input=inp))
 
 def enable_early_exec():
-  import subprocess, multiprocessing
   qin: multiprocessing.Queue = multiprocessing.Queue()
   qout: multiprocessing.Queue = multiprocessing.Queue()
-  def _early_exec_process(qin, qout):
-    while 1:
-      path, inp = qin.get()
-      qout.put(subprocess.check_output(path, input=inp))
   p = multiprocessing.Process(target=_early_exec_process, args=(qin, qout))
   p.daemon = True
   p.start()
@@ -16,17 +18,16 @@ def enable_early_exec():
     return qout.get()
   return early_exec
 
-def proc(itermaker, q):
+def proc(itermaker, q) -> None:
   for x in itermaker(): q.put(x)
   q.close()
 
 def cross_process(itermaker, maxsize=8):
   # TODO: use cloudpickle for itermaker
-  import multiprocessing
   q: multiprocessing.Queue = multiprocessing.Queue(maxsize)
   p = multiprocessing.Process(target=proc, args=(itermaker, q))
   p.daemon = True
   p.start()
 
   # TODO: write tests and handle exit case
-  while 1: yield q.get()
+  while True: yield q.get()