IQ.Pilot Release Commit @ 0798119

tinygrad_repo/extra/tinyfs/fetch_file.py

@@ -0,0 +1,36 @@
+from tinygrad.tensor import Tensor
+import argparse, math, hashlib
+
+def _python_hash_1mb(data:bytes|bytearray):
+  chunks = [data[i:i+4096] for i in range(0, len(data), 4096)]
+  chunk_hashes = [hashlib.shake_128(chunk).digest(16) for chunk in chunks]
+  return hashlib.shake_128(b''.join(chunk_hashes)).digest(16)
+
+def hash_file(data: bytes|bytearray):
+  if len(data) % Tensor.CHUNK_SIZE != 0: data += bytes(Tensor.CHUNK_SIZE - len(data) % Tensor.CHUNK_SIZE)
+  base_chunks = math.ceil(len(data) / Tensor.CHUNK_SIZE)
+  tree_depth = math.ceil(math.log(base_chunks, Tensor.CHUNK_SIZE // 16))
+
+  for _ in range(tree_depth + 1):
+    data_chunks = [data[i:i+Tensor.CHUNK_SIZE] for i in range(0, len(data), Tensor.CHUNK_SIZE)]
+    data_chunk_hashes = [_python_hash_1mb(chunk) for chunk in data_chunks]
+    data = b''.join(data_chunk_hashes)
+    if len(data) % Tensor.CHUNK_SIZE != 0: data += bytes(Tensor.CHUNK_SIZE - len(data) % Tensor.CHUNK_SIZE)
+
+  return data[:16]
+
+if __name__ == "__main__":
+  parser = argparse.ArgumentParser()
+  parser.add_argument("--hash", type=str, required=True, help="file hash to fetch")
+  parser.add_argument("--len", type=int, required=True, help="file length to fetch")
+  parser.add_argument("--dest", type=str, required=True, help="destination path to save the file")
+  parser.add_argument("--check", action="store_true", help="verify the file hash after fetching")
+  args = parser.parse_args()
+
+  Tensor(bytes.fromhex(args.hash), device="CPU").fs_load(args.len).to(f"disk:{args.dest}").realize()
+
+  if args.check:
+    with open(args.dest, "rb") as f:
+      data = f.read()
+      assert hash_file(data) == bytes.fromhex(args.hash), "Hash mismatch after fetching file"
+    print("File hash verified successfully!")

tinygrad_repo/extra/tinyfs/fetch_raid.py

@@ -0,0 +1,41 @@
+import json, multiprocessing, functools
+from pathlib import Path
+
+from tinygrad.tensor import Tensor
+from tinygrad.helpers import tqdm, getenv
+
+raid_root = Path(getenv("RAID_ROOT", "/raid"))
+
+def fetch_file(item):
+  path, info = item
+  h, size = info["hash"], info["size"]
+
+  path = raid_root / Path(path)
+  path.parent.mkdir(parents=True, exist_ok=True)
+
+  try:
+    pt = Tensor(bytes.fromhex(h), device="CPU").fs_load(size).to(f"disk:{path.as_posix()}").realize()
+  except Exception as e:
+    print(f"error fetching {path}, {h}, {size}: {e}")
+    raise
+
+  pt.uop.buffer.deallocate()
+
+def fetch_mapping(h, l):
+  mapping_tensor = Tensor(bytes.fromhex(h)).fs_load(l).realize()
+  mapping = mapping_tensor.data().tobytes().decode()
+  mapping = json.loads(mapping)
+  mapped_files = mapping.items()
+  return list(mapped_files)
+
+if __name__ == "__main__":
+  h, l = getenv("HASH", "d734f5e3be9f1e9d863bfaa4fc6c1ef2"), getenv("LENGTH", 175866113)
+
+  with multiprocessing.Pool(processes=1) as pool:
+    mapped_files = pool.apply(functools.partial(fetch_mapping, h, l))
+
+  print(f"fetched mapping for {len(mapped_files)} files")
+
+  with multiprocessing.Pool(processes=multiprocessing.cpu_count()) as pool:
+    for _ in tqdm(pool.imap_unordered(fetch_file, mapped_files), total=len(mapped_files)):
+      pass

tinygrad_repo/extra/tinyfs/upload_raid.py

@@ -0,0 +1,31 @@
+from pathlib import Path
+import multiprocessing, json
+
+from tinygrad.tensor import Tensor
+from tinygrad.helpers import tqdm
+
+raid_root = Path("/raid")
+
+def upload_file(path: Path):
+  pt = Tensor(path).realize()
+  h = pt.fs_store().realize()
+  pt.uop.realized.deallocate()
+  return h.data().hex(), path, pt.nbytes()
+
+if __name__ == "__main__":
+  raid_files = sorted([p for p in raid_root.rglob("*") if p.is_file()])
+  print(f"found {len(raid_files)} files in /raid")
+
+  mapping = {}
+  with multiprocessing.Pool(processes=multiprocessing.cpu_count()) as pool:
+    for h, p, s in tqdm(pool.imap_unordered(upload_file, raid_files), total=len(raid_files)):
+      mapping[p.relative_to(raid_root).as_posix()] = {"hash": h, "size": s}
+
+  # sort the mapping by key
+  mapping = dict(sorted(mapping.items()))
+
+  mapping = json.dumps(mapping).encode()
+  mapping_tensor = Tensor(mapping, device="CPU")
+  h = mapping_tensor.fs_store().realize()
+
+  print(f"final hash: {h.data().hex()}, size: {len(mapping)}")