evakit 0.0.0__tar.gz

evakit-0.0.0/LICENSE

@@ -0,0 +1,9 @@
+
+The MIT License (MIT)
+Copyright (c) 2026, garywei944
+
+Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.

evakit-0.0.0/PKG-INFO

@@ -0,0 +1,61 @@
+Metadata-Version: 2.4
+Name: evakit
+Version: 0.0.0
+Summary: A collection of utilities for the eva project.
+Author-email: Gary Wei <garywei944@gmail.com>
+License-Expression: MIT
+Classifier: Intended Audience :: Developers
+Classifier: Natural Language :: English
+Classifier: Operating System :: OS Independent
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Classifier: Programming Language :: Python :: 3.14
+Requires-Python: >=3.11
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: typed-argument-parser
+Requires-Dist: psutil
+Requires-Dist: typing-extensions
+Dynamic: license-file
+
+# evakit
+
+## Quick Access
+
+## Dataset
+
+## Notebooks
+
+- Notebooks for experiments are located under [`notebooks`](notebooks).
+- Notebooks for visualization and presentation are located
+  under [`references`](references).
+
+## Conda / Mamba Environment Setup
+
+```shell
+# To firstly install the environment
+conda env create -f environment.yml
+```
+
+To update the environment after updating environment.yml
+
+**_BEST PRACTICE, RECOMMENDED_** when updating a conda environment
+
+```shell
+conda env update -f environment.yml --prune
+```
+
+## Sandbox
+
+The `sandbox` is a practice in teamwork collaboration.
+Everyone works with the project should have a sub-folder named by their id
+under it, and any scripts, doc, or data that are temporary should be placed
+there.
+
+---
+
+Project based on
+the [cookiecutter machine learning template](https://github.com/garywei944/cookiecutter-machine-learning)
+.

evakit-0.0.0/README.md

@@ -0,0 +1,39 @@
+# evakit
+
+## Quick Access
+
+## Dataset
+
+## Notebooks
+
+- Notebooks for experiments are located under [`notebooks`](notebooks).
+- Notebooks for visualization and presentation are located
+  under [`references`](references).
+
+## Conda / Mamba Environment Setup
+
+```shell
+# To firstly install the environment
+conda env create -f environment.yml
+```
+
+To update the environment after updating environment.yml
+
+**_BEST PRACTICE, RECOMMENDED_** when updating a conda environment
+
+```shell
+conda env update -f environment.yml --prune
+```
+
+## Sandbox
+
+The `sandbox` is a practice in teamwork collaboration.
+Everyone works with the project should have a sub-folder named by their id
+under it, and any scripts, doc, or data that are temporary should be placed
+there.
+
+---
+
+Project based on
+the [cookiecutter machine learning template](https://github.com/garywei944/cookiecutter-machine-learning)
+.

evakit-0.0.0/evakit/__version__.py

@@ -0,0 +1 @@
+__version__ = "0.2.0"

evakit-0.0.0/evakit/args_base.py

@@ -0,0 +1,190 @@
+import functools
+import logging
+import os
+from abc import ABC
+from argparse import ArgumentParser, ArgumentTypeError, Namespace
+from typing import final, get_args, get_origin
+
+from tap import Tap
+from typing_extensions import override
+
+from evakit.python_tricks import freeze_dataclass
+from evakit.singleton import Singleton
+
+__all__ = [
+    "ArgsBase",
+    "str2bool",
+    "add_bool_arg",
+    "move_arg_from_to",
+    "env_to_bool",
+    "arg_env_consistent_bool",
+    "csv",
+    "tuple_parser",
+]
+
+logger = logging.getLogger(__name__)
+
+
+class ArgsBase(Tap, Singleton, ABC):
+    """Base class for argument parsing using Tap.
+
+    Supports
+    - Type hinting
+    - Automatic argument parsing
+    - Singleton pattern with Dependency Injection
+    """
+
+    @final
+    def __init__(self, args: list[str] | None = None, *, frozen: bool = True) -> None:
+        super().__init__(explicit_bool=False, allow_abbrev=False)
+
+        self.parse_args(args=args, known_only=True)
+
+        if frozen:
+            freeze_dataclass(self.__class__)
+
+    def _add_args(self) -> None:
+        """Add arguments to the parser.
+
+        Equivalent to self.configure(), only keep for backward compatibility.
+        """
+
+    @final
+    @override
+    def configure(self) -> None:
+
+        # The Tap accepts `--ports p1 p2 p3` as a list input, which is different from
+        # `--ports=p1,p2,p3`. To keep backward compatibility, we override them here.
+        # self._annotations is a dict of arg name to type hint Types.GenericAlias
+        for arg_name, type_hint in self._annotations.items():
+
+            arg_flag, default = f"--{arg_name}", getattr(self, arg_name)
+            type_origin, type_args = get_origin(type_hint), get_args(type_hint)
+
+            # 1. handler bool to add_bool_arg
+            if type_hint is bool:
+                add_bool_arg(self, arg_flag, default=getattr(self, arg_name))
+            # 2. handle list or set
+            elif type_origin in {list, set}:
+                elem_type = type_args[0]
+                if elem_type is int:
+                    self.add_argument(arg_flag, type=functools.partial(csv, int, type_origin))
+                elif elem_type is float:
+                    self.add_argument(arg_flag, type=functools.partial(csv, float, type_origin))
+                elif elem_type is bool:
+                    self.add_argument(arg_flag, type=functools.partial(csv, str2bool, type_origin))
+                elif elem_type is str:
+                    self.add_argument(arg_flag, type=functools.partial(csv, str, type_origin))
+            # 3. handle tuple
+            elif type_origin is tuple:
+
+                # homogeneous tuple: tuple[int, ...]
+                if len(type_args) == 2 and type_args[1] is ...:
+                    elem_type = type_args[0]
+                    self.add_argument(
+                        arg_flag,
+                        type=functools.partial(csv, elem_type, tuple),
+                        default=default,
+                    )
+                else:
+                    self.add_argument(
+                        arg_flag,
+                        type=functools.partial(tuple_parser, type_args),
+                        default=default,
+                    )
+
+        self._add_args()
+
+    @final
+    @override
+    def process_args(self) -> None:
+        self._merge_args()
+        self._process_args()
+        self._check_args()
+
+    def _merge_args(self) -> None:
+        """Merge some arguments for compatibility."""
+
+    def _process_args(self) -> None:
+        """Process the arguments after parsing."""
+
+    def _check_args(self) -> None:
+        """Check the arguments for validity."""
+
+
+def str2bool(v: str) -> bool:
+    if v.lower() in ("yes", "true", "t", "y", "1"):
+        return True
+    elif v.lower() in ("no", "false", "f", "n", "0"):
+        return False
+    else:
+        raise ArgumentTypeError("Boolean value expected.")
+
+
+def add_bool_arg(parser: ArgumentParser, *args, default=False, help=None, **kwargs):
+    # https://stackoverflow.com/a/43357954
+
+    parser.add_argument(
+        *args,
+        type=str2bool,
+        metavar="BOOLEAN",
+        nargs="?",
+        const=True,
+        default=default,
+        help=help,
+        **kwargs,
+    )
+
+
+def move_arg_from_to(
+    args: Namespace | ArgsBase,
+    old_arg: str,
+    new_arg: str,
+    final_value,
+    print_deprecated_msg: bool,
+):
+    setattr(args, new_arg, final_value)
+    setattr(args, old_arg, final_value)  # for backward compatibility
+    if print_deprecated_msg:
+        logger.warning(
+            '[ArgParser] --%s is deprecated, please use "--%s" instead!',
+            old_arg,
+            new_arg,
+        )
+
+
+def env_to_bool(env_var: str) -> bool:
+    return os.getenv(env_var, "0").lower() not in ["0", "false"]
+
+
+def arg_env_consistent_bool(args: Namespace | ArgsBase, arg_name: str, env_var: str):
+    """Make sure argument and environment variable are consistent.
+
+    Taking the logical OR of both values.
+    """
+    arg_value = bool(getattr(args, arg_name, False))
+    env_value = env_to_bool(env_var)
+
+    if arg_value or env_value:
+        setattr(args, arg_name, True)
+        os.environ[env_var] = "1"
+    else:
+        setattr(args, arg_name, False)
+        os.environ[env_var] = "0"
+
+
+def csv(elem_type, container_type: type = list, value: str = ""):
+    if value == "":
+        return container_type()
+    try:
+        return container_type(elem_type(v) for v in value.split(","))
+    except Exception as e:
+        raise ArgumentTypeError(f"Invalid {elem_type.__name__} list: {value}") from e
+
+
+def tuple_parser(type_args, value: str):
+    parts = value.split(",")
+    if len(parts) != len(type_args):
+        raise ArgumentTypeError(f"Expected {len(type_args)} values, got {len(parts)}")
+
+    return tuple(t(p) for t, p in zip(type_args, parts))

evakit-0.0.0/evakit/cronjob.py

@@ -0,0 +1,52 @@
+import logging
+import threading
+from typing import Any, Callable
+
+__all__ = ["CronJob"]
+
+logger = logging.getLogger(__name__)
+
+
+class CronJob(threading.Thread):
+    interval: float | int
+
+    _stop_event: threading.Event
+
+    _target: Callable[..., Any]
+    _args: tuple[Any, ...]
+    _kwargs: dict[str, Any]
+
+    def __init__(
+        self,
+        name: str,
+        task: Callable,
+        args: tuple[Any, ...] = (),
+        kwargs: dict[str, Any] | None = None,
+        *,
+        interval: float | int = 60,
+        daemon: bool | None = True,
+    ):
+        super().__init__(
+            target=task,
+            name=name,
+            args=args,
+            kwargs=kwargs,
+            daemon=daemon,
+        )
+        self.interval_sec = interval
+        self._stop_event = threading.Event()
+
+    def run(self):
+        logger.info("[CronJob %s] Started with interval %ss", self.name, self.interval_sec)
+        while not self._stop_event.is_set():
+            try:
+                self._target(*self._args, **self._kwargs)
+            except Exception as e:
+                logger.exception("[CronJob %s] Exception: %s", self.name, e)
+            # Wait with early exit if stopped
+            if self._stop_event.wait(self.interval_sec):
+                break
+
+    def stop(self):
+        self._stop_event.set()
+        logger.info("[CronJob %s] Stopped", self.name)

evakit-0.0.0/evakit/launcher_base.py

@@ -0,0 +1,286 @@
+import atexit
+import datetime
+import itertools
+import logging
+import os
+import signal
+import threading
+from abc import ABC, abstractmethod
+from contextlib import contextmanager, suppress
+from dataclasses import dataclass
+from typing import Callable, TextIO, cast
+
+import psutil
+
+__all__ = [
+    "LauncherBase",
+    "ProcessMeta",
+    "defer_termination_signals",
+    "restore_signal_mask",
+    "exit_handler",
+]
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class ProcessMeta:
+    proc: psutil.Popen
+    log_file: TextIO
+
+
+class LauncherBase(ABC):
+    """
+    Maintain the life cycle of the spawned processes.
+    """
+
+    procs: list[ProcessMeta | None]
+
+    _lock: threading.Lock
+    _log_counter: itertools.cycle
+    _logging_interval: int
+
+    def __init__(
+        self,
+        exit_timeout: float = 180,
+        logging_interval: int = 60,
+        *,
+        # For testability,
+        register_atexit: bool = True,
+    ):
+        self.procs = []
+        self._logging_interval = logging_interval
+        # log every 60 * 10s = 10 minutes
+        self._log_counter = itertools.cycle(range(logging_interval))
+        self._lock = threading.Lock()
+        if register_atexit:
+            atexit.register(exit_handler, self, exit_timeout, "Normal Exit signal received")
+
+    def is_healthy(self) -> bool:
+        return self._is_healthy()
+
+    def _is_healthy(self, log_name="Process", offset=None) -> bool:
+        """Check for all process, process is alive or process is None."""
+        count = next(self._log_counter)
+        with self._lock:
+            for i, proc_meta in enumerate(self.procs):
+                if proc_meta is None:
+                    continue
+
+                p = proc_meta.proc
+                idx = "" if offset is None else f" {offset + i}"
+                if p.poll() is not None:
+                    logger.warning(
+                        "[Launcher] %s%s pid=%d exited unexpectedly, exit code: %d",
+                        log_name,
+                        idx,
+                        p.pid,
+                        p.returncode,
+                    )
+                    date_str = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
+                    if proc_meta.log_file:
+                        proc_meta.log_file.write(
+                            f"{date_str} runner fail; only for report log to forge\n"
+                        )
+                    return False
+                if count == 0:
+                    logger.info(
+                        "[Launcher] %s%s is running. pid=%d pgid=%d sid=%d",
+                        log_name,
+                        idx,
+                        p.pid,
+                        os.getpgid(p.pid),
+                        os.getsid(p.pid),
+                    )
+        return True
+
+    def all_alive(self) -> bool:
+        with self._lock:
+            # return False if procs is empty
+            if not self.procs:
+                return False
+            # all([]) returns True
+            return all(p is not None and p.proc.poll() is None for p in self.procs)
+
+    def any_alive(self) -> bool:
+        with self._lock:
+            # any([]) returns False
+            return any(p is not None and p.proc.poll() is None for p in self.procs)
+
+    @abstractmethod
+    def launch(self, *args, **kwargs) -> None:
+        """Launch the processes and maintain their metadata in self.procs."""
+
+    def kill_and_wait(self, timeout: float = 5, *args, **kwargs):
+        with self._lock:
+            proc_metas = [proc for proc in self.procs if proc is not None]
+            kill_procs(proc_metas, timeout)
+
+            # Clean status
+            for idx in range(len(self.procs)):
+                self.procs[idx] = None
+            self._log_counter = itertools.cycle(range(self._logging_interval))
+
+    def wait(self) -> list[int | None]:
+        exit_codes: list[int | None] = [None] * len(self.procs)
+        for i, proc_meta in enumerate(self.procs):
+            if proc_meta is None:
+                continue
+            p = proc_meta.proc
+            p.wait()
+            exit_codes[i] = p.returncode
+            logger.info("[Launcher] Process pid=%d exited with code %d", p.pid, p.returncode)
+        # Clean status
+        for idx in range(len(self.procs)):
+            self.procs[idx] = None
+        self._log_counter = itertools.cycle(range(self._logging_interval))
+
+        return exit_codes
+
+
+def kill_procs(proc_metas: list[ProcessMeta], timeout: float = 5):
+    """Guarantee failed procs are killed and log files are closed."""
+
+    # ! Scan all process that has the same pgid with self.procs. This is necessary
+    # ! to avoid that when bash is killed by -9 SIGKILL, runner process become orphan
+    # ! and its ppid=1 and not listed in curr_process.children(recursive=True)
+
+    # Popen with start_new_session=True guarantees that p.pid == p.pgid
+    child_pgids = [p.proc.pid for p in proc_metas]
+    child_pgid_procs: list[psutil.Process] = []
+    for proc in psutil.process_iter(["pid"]):
+        try:
+            if os.getpgid(proc.pid) in child_pgids:
+                child_pgid_procs.append(proc)
+        except (psutil.NoSuchProcess, ProcessLookupError):
+            logger.info("[Launcher] Process %d already terminated, no need to wait", proc.pid)
+            continue
+
+    # ! a child process of nsys might spawn our base_runner process by setting new
+    # ! pgid, so we need to reap its children, trying our best.
+    # ! But if our base_runner process is already an orphan with new pgid,
+    # ! there's no way to figure out that it's formerly a child of this launcher.
+    all_children: list[psutil.Process] = [p.proc for p in proc_metas]
+    for proc in child_pgid_procs:
+        all_children.append(proc)
+        all_children.extend(proc.children(recursive=True))
+
+    # Send SIGTERM to all children
+    need_to_wait: list[psutil.Process] = []
+    for proc in set(all_children):
+        logger.info("[Launcher] Send SIGTERM to child process %d", proc.pid)
+        try:
+            proc.terminate()
+        except psutil.ZombieProcess:
+            logger.warning("[Launcher] Child process %d is a zombie, reaping it", proc.pid)
+            proc.wait()
+        except psutil.NoSuchProcess:
+            logger.info("[Launcher] Child process %d already terminated", proc.pid)
+            continue
+        except psutil.AccessDenied:
+            logger.warning("[Launcher] Access denied to terminate child process %d", proc.pid)
+            continue
+        else:
+            need_to_wait.append(proc)
+
+    # In some region(like GCP), there are very flaky kernel syscall which might cause
+    # process to be stuck in uninterruptible sleep state. If that happens, waiting for
+    # maximum 5 min then exit.
+    tolerance_extra_timeout = 300
+    retry_sigkill_every = 30
+
+    def kill_callback(p: psutil.Process):
+        with suppress(Exception):
+            p.kill()
+
+    for tries in range(tolerance_extra_timeout // retry_sigkill_every + 1):
+        # For the first try, we wait for the given timeout
+        _timeout = timeout if tries == 0 else retry_sigkill_every
+
+        need_to_wait = wait_for_procs(
+            need_to_wait, kill_callback, "killing it with SIGKILL", _timeout
+        )
+        if not need_to_wait:
+            break
+    else:
+        # After all retries, some processes are still alive, force exit
+        for p in need_to_wait:
+            logger.critical(
+                "[Launcher] Child process %d didn't respond to SIGKILL after %d seconds!",
+                p.pid,
+                tolerance_extra_timeout,
+            )
+        logger.critical("[Launcher] Exiting launcher process forcefully!")
+        os._exit(1)
+
+    for proc in proc_metas:
+        proc.log_file.close()
+
+    logger.info("[Launcher] All processes & log_files are closed.")
+
+
+def wait_for_procs(
+    procs: list[psutil.Process],
+    callback: Callable[[psutil.Process], None],
+    callback_msg: str,
+    timeout: float = 5,
+) -> list[psutil.Process]:
+    gone, alive = psutil.wait_procs(procs, timeout=timeout)
+    for p in gone:
+        p = cast(psutil.Popen, p)
+        logger.info(
+            "[Launcher] Child process %d terminated with exit code %d",
+            p.pid,
+            p.returncode,
+        )
+    for p in alive:
+        logger.warning(
+            "[Launcher] Child process %d did not terminate in time, %s",
+            p.pid,
+            callback_msg,
+        )
+        callback(p)
+
+    return alive
+
+
+ALL_TERM_SIGNALS = {
+    signal.SIGABRT,
+    signal.SIGBUS,
+    signal.SIGFPE,
+    signal.SIGILL,
+    signal.SIGINT,
+    signal.SIGKILL,
+    signal.SIGPIPE,
+    signal.SIGQUIT,
+    signal.SIGSEGV,
+    signal.SIGTERM,
+    signal.SIGSYS,
+}
+
+
+@contextmanager
+def defer_termination_signals():
+    # Block these signals and save the old signal mask.
+    old_mask = signal.pthread_sigmask(signal.SIG_BLOCK, ALL_TERM_SIGNALS)
+    try:
+        yield
+    finally:
+        # Restore the original signal mask.
+        signal.pthread_sigmask(signal.SIG_SETMASK, old_mask)
+
+
+def restore_signal_mask():
+    """restore the signal mask that might block all termination signals"""
+    signal.pthread_sigmask(signal.SIG_UNBLOCK, ALL_TERM_SIGNALS)
+
+
+def exit_handler(launcher: LauncherBase, timeout: float = 180, msg: str = ""):
+    with defer_termination_signals():
+        logger.info("=" * 80)
+        if msg:
+            logger.info(msg)
+        logger.info("[ExitHandler] Release resource, timeout=%d s", timeout)
+        lg_procs = [lg_proc for lg_proc in launcher.procs if lg_proc is not None]
+        kill_procs(lg_procs, timeout)
+        logger.info("[ExitHandler] Resource released")