disk cache: thread-local db conn (#17694)

tinygrad/helpers.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 import time
 START_TIME = time.perf_counter()
-import os, functools, re, contextlib, operator, hashlib, pickle, sqlite3, tempfile, pathlib, string, ctypes, sys, gzip, getpass, gc
+import os, functools, re, contextlib, operator, hashlib, pickle, sqlite3, tempfile, pathlib, string, ctypes, sys, gzip, getpass, gc, threading
 from collections import defaultdict
 import shutil, math, types, copyreg, inspect, importlib, decimal, itertools, difflib
 from dataclasses import dataclass, field, replace
@@ -398,18 +398,17 @@ cache_dir: str = os.path.join(getenv("XDG_CACHE_HOME", os.path.expanduser("~/Lib
 CACHEDB: str = getenv("CACHEDB", os.path.abspath(os.path.join(cache_dir, "cache.db")))
 
 VERSION = 22
-_db_connection = None
+_db_connection = threading.local()
 def db_connection():
-  global _db_connection
-  if _db_connection is None:
+  if (conn:=getattr(_db_connection, "conn", None)) is None:
     os.makedirs(CACHEDB.rsplit(os.sep, 1)[0], exist_ok=True)
-    _db_connection = sqlite3.connect(CACHEDB, timeout=60, isolation_level="IMMEDIATE")
+    conn = _db_connection.conn = sqlite3.connect(CACHEDB, timeout=60, isolation_level="IMMEDIATE")
     # another connection has set it already or is in the process of setting it
     # that connection will lock the database
-    with contextlib.suppress(sqlite3.OperationalError): _db_connection.execute("PRAGMA journal_mode=WAL").fetchone()
-    _db_connection.execute("PRAGMA synchronous=NORMAL")
-    if DEBUG >= 8: _db_connection.set_trace_callback(print)
-  return _db_connection
+    with contextlib.suppress(sqlite3.OperationalError): conn.execute("PRAGMA journal_mode=WAL").fetchone()
+    conn.execute("PRAGMA synchronous=NORMAL")
+    if DEBUG >= 8: conn.set_trace_callback(print)
+  return conn
 
 def diskcache_clear():
   cur = db_connection().cursor()