proc

Signed-off-by: Elazar Gershuni <elazarg@gmail.com>

Author: elazarg

checkpoint/criu_binding.py

@@ -0,0 +1,113 @@
+from os import PathLike
+
+from checkpoint.ctypes_modern import Clibrary
+import os
+
+criu = Clibrary("criu")
+
+
+@criu.function
+def criu_init_opts() -> int: ...
+
+
+@criu.function
+def criu_check() -> int: ...
+
+
+@criu.function
+def criu_dump() -> int: ...
+
+
+@criu.function
+def criu_restore() -> int: ...
+
+
+@criu.function
+def criu_set_images_dir_fd(fd: int) -> None: ...
+
+
+@criu.function
+def criu_set_log_file(log_file: bytes) -> None: ...
+
+
+@criu.function
+def criu_set_log_level(log_level: int) -> None: ...
+
+
+@criu.function
+def criu_set_pid(pid: int) -> None: ...
+
+
+@criu.function
+def criu_set_leave_running(leave_running: bool) -> None: ...
+
+
+@criu.function
+def criu_set_service_address(address: bytes) -> None: ...
+
+
+@criu.function
+def criu_set_track_mem(track_mem: bool) -> None: ...
+
+
+@criu.function
+def criu_set_ext_unix_sk(ext_unix_sk: bool) -> None: ...
+
+
+@criu.function
+def criu_set_tcp_established(tcp_established: bool) -> None: ...
+
+
+@criu.function
+def criu_set_evasive_devices(evasive_devices: bool) -> None: ...
+
+
+@criu.function
+def criu_set_shell_job(shell_job: bool) -> None: ...
+
+
+@criu.function
+def criu_set_file_locks(file_locks: bool) -> None: ...
+
+
+@criu.function
+def criu_set_log_level(log_level: int) -> None: ...
+
+
+@criu.function
+def criu_set_log_file(log_file: bytes) -> None: ...
+
+
+def set_criu(folder: str | bytes | PathLike[str] | PathLike[bytes]) -> None:
+    if criu_init_opts() < 0:
+        raise OSError("CRIU init failed")
+
+    try:
+        fd = os.open(folder, os.O_DIRECTORY)
+    except OSError as e:
+        raise OSError(f"Failed to open criu_images directory: {e}")
+    criu_set_images_dir_fd(fd)
+
+    criu_set_log_file(b"criu.log")
+    criu_set_log_level(4)
+    criu_set_pid(os.getpid())
+    criu_set_leave_running(True)
+    criu_set_service_address(b"/tmp/criu_service.socket")
+    criu_set_track_mem(False)
+
+
+if __name__ == "__main__":
+    set_criu("../scripts/criu_images")
+    c = criu_check()
+    if c < 0:
+        print(f"Failed to check CRIU: {os.strerror(c)}")
+    elif c == 0:
+        print("CRIU is available")
+    else:
+        print("CRIU is not available")
+
+    # Perform a checkpoint
+    criu_dump()
+    # move criu_images/pages-1.img to criu_images/pages_old.img
+    os.rename("../scripts/criu_images/pages-1.img", "criu_images/pages_old.img")
+    criu_dump()

checkpoint/ctypes_modern.py

@@ -0,0 +1,72 @@
+import ctypes
+from ctypes.util import find_library
+from typing import Any, Callable, Type, get_type_hints
+
+
+def _get_ctype(type_hint: Type) -> Any:
+    if type_hint is int:
+        return ctypes.c_int
+    elif type_hint is float:
+        return ctypes.c_float
+    elif type_hint is bool:
+        return ctypes.c_bool
+    elif type_hint is bytes:
+        return ctypes.c_char_p
+    elif type_hint is None:
+        return None
+    else:
+        raise TypeError(f"Unsupported type hint: {type_hint}")
+
+
+def structure[T](cls: Type[T]) -> Type:
+    """
+    A decorator to create a ctypes Structure class using type hints.
+    """
+    return type(
+        cls.__name__,
+        (ctypes.Structure,),
+        {
+            "_fields_": [
+                (name, _get_ctype(type_hint))
+                for name, type_hint in get_type_hints(cls).items()
+            ]
+        },
+    )
+
+
+class Clibrary:
+    """
+    A class for wrapping C libraries and their functions.
+    """
+
+    def __init__(self, library_name: str):
+        self.lib = ctypes.CDLL(find_library(library_name))
+
+    def function[
+        *Args, Res
+    ](self, func: Callable[[*Args], Res]) -> Callable[[*Args], Res]:
+        """
+        A decorator method for wrapping individual C functions.
+        """
+        func_name = func.__name__
+        annotations = func.__annotations__
+
+        c_func = getattr(self.lib, func_name)
+        c_func.restype = _get_ctype(annotations.get("return"))
+        c_func.argtypes = [
+            _get_ctype(arg_type)
+            for arg_name, arg_type in annotations.items()
+            if arg_name != "return"
+        ]
+
+        return c_func
+
+
+criu = Clibrary("criu")
+
+
+@criu.function
+def criu_init_opts() -> int: ...
+
+
+criu_init_opts()

checkpoint/persist.py

@@ -10,7 +10,7 @@
 import hashlib
 import socket
 import struct
-
+from checkpoint.criu_binding import set_criu, criu_dump
 
 FUEL = "FUEL"
 STEP = "STEP"
@@ -59,10 +59,10 @@ def __init__(self, module_filename: str | pathlib.Path, env) -> None:
 
         # make sure the cache is invalidated when the module changes
         h = compute_hash(module_filename, env)
-        cachedir = pathlib.Path(tempfile.gettempdir()) / f"pythia-{h}"
-        print(f"Using cache directory {cachedir}", file=sys.stderr)
-        cachedir.mkdir(parents=False, exist_ok=True)
-        self.filename = cachedir / "store.pickle"
+        cache_dir = pathlib.Path(tempfile.gettempdir()) / f"pythia-{h}"
+        print(f"Using cache directory {cache_dir}", file=sys.stderr)
+        cache_dir.mkdir(parents=False, exist_ok=True)
+        self.filename = cache_dir / "store.pickle"
         (results_path / tag).mkdir(parents=True, exist_ok=True)
         self.tsv_filename = (results_path / tag / module_filename.stem).with_suffix(
             ".tsv"
@@ -243,6 +243,34 @@ def sigint() -> None:
     os.kill(PID, signal.SIGINT)
 
 
-def self_coredump(i) -> None:
-    if i % STEP_VALUE in [0, 1]:
-        sigint()
+SET_CRIU = False
+CRIU_FOLDER = pathlib.Path("criu_images")
+CRIU_DUMPS = CRIU_FOLDER / "dumps"
+
+coredump_iterations = 0
+coredump_steps = 0
+
+
+def self_coredump() -> None:
+    global coredump_iterations, coredump_steps
+    if not coredump_iterations:
+        CRIU_FOLDER.mkdir(exist_ok=True)
+        shutil.rmtree(CRIU_DUMPS, ignore_errors=True)
+        CRIU_DUMPS.mkdir(exist_ok=False)
+        set_criu(CRIU_FOLDER)
+    coredump_iterations += 1
+
+    if coredump_iterations % STEP_VALUE in [0, 1]:
+        criu_dump()
+        image_file = CRIU_FOLDER / "pages-1.img"
+        if not image_file.exists():
+            raise RuntimeError(
+                "CRIU image was not created. Make sure to run the CRIU service:\n"
+                "sudo criu service --shell-job --address /tmp/criu_service.socket"
+            )
+        target_image = CRIU_DUMPS / f"{coredump_steps:05d}.a.img"
+        os.rename(CRIU_FOLDER / "pages-1.img", target_image)
+        if coredump_steps > 0:
+            source_image = CRIU_DUMPS / f"{coredump_steps-1:05d}.b.img"
+            source_image.hardlink_to(target_image)
+        coredump_steps += 1

experiment/k_means/proc.py

@@ -1,4 +1,4 @@
-from checkpoint.persist import self_coredump
+from checkpoint import persist
 import numpy as np
 import sklearn.datasets
 
@@ -21,7 +21,7 @@ def run(X: np.ndarray, k: int, max_iterations: int) -> np.ndarray:
     centroids = X[np.random.choice(nsamples, k)]
     clusters = list[list[int]]()
     for i in range(max_iterations):  # type: int
-        self_coredump(i)
+        persist.self_coredump()
         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))

experiment/omp/proc.py

@@ -0,0 +1,87 @@
+from checkpoint import persist
+import sys
+import numpy as np
+import argparse
+
+
+def append_int(a: np.ndarray, n: int) -> np.ndarray:
+    return np.append(a, n)
+
+
+def get_float(array: np.ndarray, idx: int) -> float:
+    res = array[idx]
+    return res
+
+
+def log(idx: int, k: int) -> None:
+    print(f"{idx} / {k}", end="\r", flush=True, file=sys.stderr)
+
+
+def run(features: np.ndarray, target: np.ndarray, k: int) -> np.ndarray:
+    """select k features from features using target as the target variable"""
+    S = np.array([], "int")
+    for idx in range(k):  # type: int
+        persist.self_coredump()
+        log(idx, k)
+        dims = np.unique(S[S >= 0])
+        target = np.array(target).reshape(target.shape[0], -1)
+        X = features[:, dims]
+        if X.size == 0:
+            prediction = np.zeros(features.shape[0]).reshape(features.shape[0], -1)
+        else:
+            if X.ndim == 1:
+                X = X.reshape(X.shape[0], 1)
+            y = np.concatenate(target)
+            X = (X - X.mean()) / X.std()
+            X = np.c_[np.ones(X.shape[0]), X]
+            theta = np.zeros(X.shape[1])
+            for _ in range(10000):
+                error = np.dot(X, theta.T) - y
+                theta -= 0.1 * (1 / y.size) * np.dot(X.T, error)
+            prediction = np.zeros((len(X), 1))
+            for j in range(len(X)):
+                total = 0.0
+                xj = X[j, :]
+                for i in range(len(xj)):
+                    x = get_float(xj, i)
+                    t = get_float(theta, i)
+                    total += x * t
+                prediction[j] = total
+        grad = np.dot(features.T, target - prediction)
+        points = np.setdiff1d(np.array(range(len(grad))), S).astype("int")
+        if len(points) == 0:
+            break
+        a = points[0]
+        m = get_float(grad, a)
+        for i in range(len(points)):
+            p = points[i]
+            n = get_float(grad, p)
+            if n > m:
+                a = p
+                m = n
+        if m >= 0:
+            S = np.unique(append_int(S, a))
+        else:
+            break
+    return S
+
+
+def main(dataset: str, k: int) -> None:
+    features = np.load(f"experiment/omp/{dataset}_features.npy")
+    target = np.load(f"experiment/omp/{dataset}_target.npy")
+    S = run(features, target, k)
+    print(S)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "dataset",
+        choices=["dataset_20KB", "dataset_large", "healthstudy"],
+        help="dataset to use",
+    )
+    parser.add_argument(
+        "--k", type=int, default=100000, help="number of features to select"
+    )
+    args = parser.parse_args()
+    main(args.dataset, args.k)