github/tinygrad
1
1
Fork 0
mirror of https://github.com/tinygrad/tinygrad.git synced 2026-01-08 22:48:25 -05:00
Code Issues Packages Projects Releases Wiki Activity

hotfix: extract_dataset.py (#7029)

Browse Source
This commit is contained in:
qazal
2024-10-13 11:18:23 +03:00
committed by GitHub
parent 942a17109a
commit 13846930cd
2 changed files with 10 additions and 14 deletions
Download Patch File Download Diff File Expand all files Collapse all files

18
extra/optimization/extract_dataset.py
View File

@@ -1,22 +1,18 @@
#!/usr/bin/env python3
# extract asts from process replay artifacts
import os, pickle
from tinygrad.helpers import db_connection, getenv, VERSION
import os
from tinygrad.helpers import db_connection, VERSION
from test.external.process_replay.process_replay import _pmap
PAGE_SIZE = 100
RUN_ID = os.getenv("GITHUB_RUN_ID", "HEAD")
TABLE_NAME = f"process_replay_{RUN_ID}_{getenv('GITHUB_RUN_ATTEMPT')}_{VERSION}"
TABLE_NAME = f"kernel_process_replay_{VERSION}"
LOGOPS = os.getenv("LOGOPS", "/tmp/sops")
def extract_ast(offset:int) -> bool:
logops = open(LOGOPS, "a")
conn = db_connection()
for row in conn.execute(f"SELECT val FROM '{TABLE_NAME}' LIMIT ? OFFSET ?", (PAGE_SIZE, offset)).fetchall():
logops.write(str(pickle.loads(row[0])[0]).replace("\n", "").replace(" ", "")+"\n")
return False
def extract_ast(*args) -> bool:
open(LOGOPS, "a").write(str(args[0]).replace("\n", "").replace(" ", "")+"\n")
return args[-1]
if __name__ == "__main__":
conn = db_connection()
row_count = conn.execute(f"SELECT COUNT(*) FROM '{TABLE_NAME}'").fetchone()[0]
_pmap(row_count, extract_ast)
_pmap("kernel", extract_ast)