changes

Author: elazarg

checkpoint/persist.py

@@ -1,10 +1,10 @@
 import os
+import shutil
+import signal
 import subprocess
 import sys
 import tempfile
 import typing
-from typing import Any
-
 import pickle
 import pathlib
 import hashlib
@@ -46,9 +46,8 @@ def run_instrumented_file(
 
 class Loader:
     fuel: int
-    step: int
 
-    restored_state: tuple[Any, ...]
+    restored_state: tuple
     iterator: typing.Optional[typing.Iterable]
     i: int
     filename: pathlib.Path
@@ -72,7 +71,6 @@ def __init__(self, module_filename: str | pathlib.Path, env) -> None:
         self.fuel = read_fuel()
         self.i = -1
         self.printing_index = 1
-        self.step = read_step()
         self.iterator = None
         self.restored_state = ()
 
@@ -101,7 +99,7 @@ def iterate(self, iterable) -> typing.Iterable:
 
     def commit(self, *args) -> None:
         self.i += 1
-        if self.i % self.step in [0, 1]:
+        if self.i % STEP_VALUE in [0, 1]:
             self.fuel -= 1
             if self.fuel <= 0:
                 raise KeyboardInterrupt("Out of fuel")
@@ -161,13 +159,11 @@ def __init__(self, tag: str) -> None:
                 input()
             else:
                 break
-        self.step = read_step()
         self.i = -1
-        self.value = None
 
     def commit(self) -> None:
         self.i += 1
-        if self.i % self.step in [0, 1]:
+        if self.i % STEP_VALUE in [0, 1]:
             # send commend to save_snapshot server to take snapshot
             self.socket.send(struct.pack("Q", self.i))
             # wait for snapshot to start, and continue after it's done
@@ -180,5 +176,73 @@ def __exit__(self, exc_type, exc_val, exc_tb) -> None:
         self.socket.close()
 
     def iterate(self, iterable):
-        for self.value in iterable:
-            yield self.value
+        for value in iterable:
+            yield value
+
+
+def diff_vm_snapshot(folder: pathlib.Path, i: int) -> int:
+    folder = folder.as_posix()
+    result = subprocess.run(
+        [
+            "./diff_vm_snapshot",
+            f"{folder}/{i}.a.dump",
+            f"{folder}/{i}.b.dump",
+            "64",
+            "remove",
+        ],
+        capture_output=True,
+    )
+    if result.returncode != 0:
+        raise RuntimeError("Failed to run diff_vm_snapshot", result)
+    return int(result.stdout.strip())
+
+
+def diff_coredump(folder: pathlib.Path, i: int) -> int:
+    folder = folder.as_posix()
+    result = subprocess.run(
+        [
+            "./diff_coredump",
+            f"{folder}/{i}.a.dump",
+            f"{folder}/{i}.b.dump",
+        ],
+        capture_output=True,
+    )
+    if result.returncode != 0:
+        raise RuntimeError("Failed to run diff_coredump", result)
+    return int(result.stdout.strip())
+
+
+def coredump(tag: str, index: int) -> None:
+    pid = os.getpid()
+    result = subprocess.run(["gcore" "-o" f"{tag}" f"{pid}"], shell=True)
+    if result.returncode != 0:
+        raise RuntimeError("Failed to run diff_coredump", result)
+    os.rename(f"{tag}.{pid}", f"{tag}/{index}.core")
+
+
+def make_dumps_folder(tag: str) -> pathlib.Path:
+    cwd = pathlib.Path.cwd()
+    dumps_folder = cwd / "dumps" / tag
+    shutil.rmtree(dumps_folder, ignore_errors=True)
+    dumps_folder.mkdir(exist_ok=False, parents=True)
+    return dumps_folder
+
+
+def make_results_folder(tag: str) -> pathlib.Path:
+    cwd = pathlib.Path.cwd()
+    results_folder = cwd / "results" / tag
+    results_folder.mkdir(exist_ok=True, parents=True)
+    return results_folder
+
+
+PID = os.getpid()
+STEP_VALUE = read_step()
+
+
+def sigint() -> None:
+    os.kill(PID, signal.SIGINT)
+
+
+def self_coredump(i) -> None:
+    if i % STEP_VALUE in [0, 1]:
+        sigint()

experiment/k_means/proc.py

@@ -0,0 +1,63 @@
+from checkpoint.persist import self_coredump
+import numpy as np
+import sklearn.datasets
+
+np.random.seed(42)
+
+
+def run(X: np.ndarray, k: int, max_iterations: int) -> np.ndarray:
+    """A simple clustering method that forms k clusters by iteratively reassigning
+    samples to the closest centroids and after that moves the centroids to the center
+    of the new formed clusters. Do K-Means clustering and return cluster indices
+    @param X: np.ndarray
+        The dataset to cluster, where each row is a sample and each column is a feature.
+    @param k: int
+        The number of clusters the algorithm will form.
+    @param max_iterations: int
+        The number of iterations the algorithm will run for if it does
+        not converge before that.
+    """
+    nsamples, features = X.shape
+    centroids = X[np.random.choice(nsamples, k)]
+    clusters = list[list[int]]()
+    for i in range(max_iterations):  # type: int
+        self_coredump(i)
+        clusters = [list[int]() for _ in range(k)]
+        for sample_i in range(len(X)):
+            r = np.argmin(np.linalg.norm(X[sample_i] - centroids, axis=1))
+            clusters[r].append(sample_i)
+        prev_centroids = centroids
+        centroids = np.array([np.mean(X[cluster], axis=0) for cluster in clusters])
+        diff = centroids - prev_centroids
+        if not diff.any():
+            break
+    y_pred = np.zeros(nsamples)
+    for cluster_i in range(len(clusters)):
+        for sample_i in clusters[cluster_i]:
+            y_pred[sample_i] = cluster_i
+    return y_pred
+
+
+def compute_random(n_samples: int, k: int, plot: bool) -> None:
+    X, y = sklearn.datasets.make_blobs(
+        n_samples=n_samples, n_features=2, centers=k, cluster_std=1.8, shuffle=True
+    )
+    y_pred = run(X, k=k, max_iterations=1000)
+    if plot:
+        import matplotlib.pyplot as plt
+
+        plt.scatter(X[:, 0], X[:, 1], c=y_pred)
+        plt.show()
+    else:
+        print("Done.")
+
+
+if __name__ == "__main__":
+    import argparse
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument("samples", type=int)
+    parser.add_argument("k", type=int)
+    parser.add_argument("--plot", action="store_true")
+    args = parser.parse_args()
+    compute_random(args.samples, args.k, args.plot)

pythia/type_system.py

@@ -1804,7 +1804,10 @@ def visit_FunctionDef(self, fdef: ast.FunctionDef) -> FunctionType:
             if update is not None:
                 assert isinstance(update, ast.Call)
                 assert len(update.args) == 1
-                update_type = self.expr_to_type(update.args[0])
+                update_arg = update.args[0]
+                if isinstance(update_arg, ast.Constant) and isinstance(update_arg.value, str):
+                    update_arg = ast.parse(update_arg.s).body[0].value
+                update_type = self.expr_to_type(update_arg)
             else:
                 update_type = None
             # side_effect = parse_side_effect(fdef.body)

save_snapshot.py

@@ -14,16 +14,14 @@
 import asyncio
 import logging
 import os
-import pathlib
-import shutil
 import socket
 import struct
-import subprocess
 import argparse
 from concurrent.futures import ThreadPoolExecutor, Future
 from typing import Iterator
 
 from checkpoint.qmp_client import SimpleQmpClient
+from checkpoint import persist
 
 logging.basicConfig(level=logging.INFO)
 
@@ -96,30 +94,10 @@ def __next__(self) -> int:
         return self.i
 
 
-def count_diff(folder: pathlib.Path, i: int) -> int:
-    folder = folder.as_posix()
-    result = subprocess.run(
-        [
-            "./count_diff",
-            f"{folder}/{i}.a.dump",
-            f"{folder}/{i}.b.dump",
-            "64",
-            "remove",
-        ],
-        capture_output=True,
-    )
-    if result.returncode != 0:
-        raise RuntimeError("Failed to run count_diff", result)
-    return int(result.stdout.strip())
-
-
 async def relay_qmp_results(qmp_port: int, server: Server) -> None:
     tag = server.tag
     assert tag.replace("_", "").isalnum()
-    cwd = pathlib.Path.cwd()
-    dumps_folder = cwd / "dumps" / tag
-    shutil.rmtree(dumps_folder, ignore_errors=True)
-    os.makedirs(dumps_folder, exist_ok=False)
+    dumps_folder = persist.make_dumps_folder(tag)
     async with SimpleQmpClient(qmp_port) as vm:
         await vm.dump(f"{dumps_folder}/0.a.dump")
         with ThreadPoolExecutor() as executor:
@@ -130,15 +108,16 @@ async def relay_qmp_results(qmp_port: int, server: Server) -> None:
                     next_prev_file = dumps_folder / f"{i + 1}.a.dump"
                     await vm.dump(current_next_file)
                     os.link(current_next_file, next_prev_file)
-                    p: Future[int] = executor.submit(count_diff, dumps_folder, i)
+                    p: Future[int] = executor.submit(
+                        persist.diff_vm_snapshot, dumps_folder, i
+                    )
                     if p is not None:
                         ps.append(p)
             else:
                 server.finish()
                 os.unlink(next_prev_file)
 
-            results_folder = cwd / "results" / tag
-            results_folder.mkdir(exist_ok=True, parents=True)
+            results_folder = persist.make_results_folder(tag)
             version_number = len(
                 [x for x in results_folder.iterdir() if x.name.startswith("vm_")]
             )

scripts/Makefile

@@ -0,0 +1,16 @@
+
+
+.PHONY: all
+
+all: diff_coredump diff_vm_snapshot
+
+diff_coredump:
+	gcc -O2 -std=c23 diff_coredump.c -o diff_coredump
+	
+
+diff_vm_snapshot:
+	g++ -O2 -std=c++20 diff_vm_snapshot.cpp -o diff_vm_snapshot
+
+clean:
+	rm diff_coredump diff_vm_snapshot
+

scripts/all_diff.py

@@ -6,13 +6,13 @@
 CHUNK_SIZE = 64
 
 
-def run_count_diff(file1, file2):
+def run_count_diff(file1: Path, file2: Path) -> int:
     cmd = f"./count_diff {file1} {file2} {CHUNK_SIZE}"
     output = os.popen(cmd).read().split()[0]
     return int(output)
 
 
-def main(dumpdir, chunk_size, max_workers):
+def main(dumpdir: str, max_workers):
     dump_files = list(sorted(Path(dumpdir).iterdir(), key=os.path.getmtime))
     with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor:
         results = executor.map(run_count_diff, dump_files[:-1], dump_files[1:])
@@ -25,9 +25,8 @@ def main(dumpdir, chunk_size, max_workers):
     # Check if folder path is provided as a command-line argument
     parser = argparse.ArgumentParser(description="Compute count_diff for pairs of dump files")
     parser.add_argument("dump_dir", type=str, help="Directory containing the dump files")
-    parser.add_argument("-c", "--chunk_size", type=int, default=64, help="Chunk size for comparison")
     parser.add_argument("-w", "--max_workers", type=int, default=10, help="Maximum number of worker threads")
     args = parser.parse_args()
 
     CHUNK_SIZE = args.chunk_size
-    main(args.dump_dir, args.chunk_size, args.max_workers)
+    main(args.dump_dir, args.max_workers)