polymorph-ai 1.1.0__py3-none-any.whl

polymorph_ai/__init__.py

@@ -0,0 +1,24 @@
+"""polymorph_ai: pick sequential / threading / multiprocessing automatically.
+
+    from polymorph_ai import adaptive_exec
+
+    @adaptive_exec
+    def process(item):
+        ...
+
+    results = process.map(items)
+    print(process.last_run)
+"""
+
+from importlib.metadata import PackageNotFoundError, version
+
+from .decorator import Decision, adaptive_exec, run, shutdown, warm_up
+from .features import FEATURE_NAMES, extract_features
+
+try:
+    __version__ = version("polymorph-ai")
+except PackageNotFoundError:     # running from a source checkout that is not installed
+    __version__ = "unknown"
+
+__all__ = ["adaptive_exec", "run", "warm_up", "shutdown", "Decision",
+           "FEATURE_NAMES", "extract_features"]

polymorph_ai/decorator.py

@@ -0,0 +1,480 @@
+"""@adaptive_exec: run a per-item function over a list in the best mode.
+
+    from polymorph_ai import adaptive_exec
+
+    @adaptive_exec
+    def process(item):
+        ...
+
+    process(one_item)            # a normal call, nothing changes
+    results = process.map(items) # polymorph_ai picks sequential / threading / multiprocessing
+    print(process.last_run)      # what it picked and why
+
+How one .map() call works:
+  1. safety checks: empty / tiny list, one worker, nested call, can items and
+     the function be sent to other processes at all?
+  2. probe: run the first items for ~50 ms and measure them (features.py, the
+     exact code used to collect the training data). Their results are kept.
+  3. the MLP (model.py) turns the 13 features into a mode.
+  4. the remaining items run in that mode on a pool that stays alive for the
+     next call (the model was trained on warm-pool timings).
+"""
+
+import atexit
+import functools
+import importlib
+import logging
+import math
+import multiprocessing
+import os
+import pickle
+import sys
+import threading
+import time
+from dataclasses import dataclass, field
+from multiprocessing.pool import ThreadPool
+
+try:
+    import cloudpickle      # sends functions written in Jupyter/REPL to worker processes by value
+except ImportError:         # pragma: no cover
+    cloudpickle = None
+
+from .features import PROBE_MIN_ITEMS, ast_features, extract_features, thread_cpu_time, usable_cpu_count
+from .model import ModeModel
+
+log = logging.getLogger("polymorph_ai")
+
+MODES = ("sequential", "threading", "multiprocessing")
+# Below this estimated sequential time (whole job, or what is left after the probe), skip the model:
+# no pool can win back its own overhead (the training data starts at 10 ms).
+SMALL_JOB_S = 0.01
+# A cached decision is reused this many times, then the job is probed again.
+CACHE_USES = 50
+# Median time to start a process pool in our dataset (pool_startup_s):
+# Windows (spawn) 0.25 s, Linux (fork/forkserver) 0.07 s. Used only when the pool is not running yet.
+COLD_START_S = {"spawn": 0.25, "fork": 0.07, "forkserver": 0.07}
+
+
+def _start_method():
+    # Read the method without fixing it, so the user can still call set_start_method() later.
+    # get_all_start_methods() lists the platform default first.
+    return multiprocessing.get_start_method(allow_none=True) or multiprocessing.get_all_start_methods()[0]
+
+
+# ------------------------------------------------------------------- result
+
+@dataclass
+class Decision:
+    """What one .map() call did. `print(decision)` gives a one-line summary."""
+    function: str
+    mode: str
+    reason: str
+    n_items: int
+    n_workers: int
+    probabilities: dict = None      # model output, None if the model was not asked
+    features: dict = None           # the 13 features, None if not measured
+    probed_items: int = 0           # items already computed during the probe
+    overhead_s: float = 0.0         # time spent deciding that did no useful work
+    total_s: float = 0.0
+    notes: list = field(default_factory=list)
+
+    def __str__(self):
+        p = f" p={self.probabilities[self.mode]:.2f}" if self.probabilities else ""
+        return (f"[polymorph_ai] {self.function}: {self.n_items:,} items -> {self.mode}{p} "
+                f"({self.reason}) | overhead {self.overhead_s * 1e3:.1f} ms | total {self.total_s:.3f} s")
+
+
+# -------------------------------------------------------------------- pools
+
+_pools = {}
+_pools_lock = threading.Lock()
+_worker = threading.local()          # .active is True inside our pool workers
+_decisions = {}                      # cache key -> [mode, uses left]
+_model = None
+_model_lock = threading.Lock()
+
+
+def _mark_worker():
+    _worker.active = True
+
+
+class _ThreadPool(ThreadPool):
+    """ThreadPool that does not lock the program's default start method
+    (the stock one calls multiprocessing.get_context(), which does)."""
+
+    def __init__(self, processes, initializer):
+        multiprocessing.pool.Pool.__init__(self, processes, initializer,
+                                           context=multiprocessing.get_context(_start_method()))
+
+
+def _get_pool(mode, n_workers):
+    with _pools_lock:
+        pool = _pools.get((mode, n_workers))
+        if pool is None:
+            if mode == "threading":
+                pool = _ThreadPool(n_workers, _mark_worker)
+            else:   # explicit context: same reason as _ThreadPool
+                pool = multiprocessing.get_context(_start_method()).Pool(n_workers, initializer=_mark_worker)
+            _pools[(mode, n_workers)] = pool
+            _register_shutdown()
+        return pool
+
+
+_shutdown_registered = False
+
+
+def _register_shutdown():
+    # atexit runs the last registered function first. multiprocessing registers its own
+    # exit hook when the first pool starts; if that hook ran before ours, it would kill
+    # the pool's workers, the pool would start new ones, and those would be left running
+    # after the program (or Jupyter kernel) exits. Registering after the first pool
+    # exists makes our shutdown() run first.
+    global _shutdown_registered
+    if not _shutdown_registered:
+        atexit.register(shutdown)
+        _shutdown_registered = True
+
+
+def _pool_is_warm(mode, n_workers):
+    return (mode, n_workers) in _pools
+
+
+def warm_up(n_workers=None, modes=("threading", "multiprocessing")):
+    """Start the pools (and the Windows CPU-clock calibration) now, so the first
+    .map() call does not pay these one-off costs."""
+    thread_cpu_time()
+    _get_model()
+    n_workers = n_workers or usable_cpu_count()
+    for mode in modes:
+        _get_pool(mode, n_workers).map(int, range(n_workers), chunksize=1)
+
+
+def shutdown():
+    """Stop all pools (called automatically when the program exits)."""
+    with _pools_lock:
+        for pool in _pools.values():
+            # pool.terminate() can hang while the interpreter is shutting down (seen in
+            # Jupyter kernels on Windows), which left the worker processes running after
+            # the kernel was gone. Give it a few seconds, then kill the workers directly.
+            workers = list(getattr(pool, "_pool", []))
+            stopper = threading.Thread(target=pool.terminate, daemon=True)
+            stopper.start()
+            stopper.join(timeout=3)
+            for p in workers + list(getattr(pool, "_pool", [])):
+                if hasattr(p, "kill") and p.is_alive():
+                    p.kill()
+        _pools.clear()
+
+
+def _get_model():
+    global _model
+    with _model_lock:
+        if _model is None:
+            _model = ModeModel()
+        return _model
+
+
+# ------------------------------------------------- can we use processes?
+
+def _find(module, qualname):
+    """Unpickle helper: look a decorated function up by name in a worker process."""
+    obj = importlib.import_module(module)
+    for part in qualname.split("."):
+        obj = getattr(obj, part)
+    return obj
+
+
+def _defined_without_file(target):
+    """True for a function typed into Jupyter or a REPL: worker processes started
+    with spawn (Windows, macOS) cannot import it by name."""
+    module = sys.modules.get(getattr(target, "__module__", None))
+    return (module is not None and module.__name__ == "__main__" and not hasattr(module, "__file__")
+            and _start_method() != "fork")
+
+
+def _process_problem(target, sample_item):
+    """None if target(item) can run in a worker process, else the reason why not."""
+    if multiprocessing.parent_process() is not None or getattr(_worker, "active", False):
+        return "already inside a worker"
+    qualname = getattr(target, "__qualname__", "")
+    if "<locals>" in qualname or "<lambda>" in qualname:
+        return "function is not defined at module top level"
+    module = sys.modules.get(getattr(target, "__module__", None))
+    if module is None:
+        return "function's module cannot be found"
+    if _defined_without_file(target):
+        if cloudpickle is None or not isinstance(target, AdaptiveExec):
+            return "function defined in Jupyter/REPL cannot be sent to worker processes"
+        try:                                   # it travels by value: check that this really works
+            pickle.loads(pickle.dumps(target))
+        except Exception as e:
+            return f"function defined in Jupyter/REPL cannot be pickled ({type(e).__name__})"
+    else:
+        try:
+            found = _find(module.__name__, qualname)
+        except AttributeError:
+            found = None
+        if found is not target:
+            return "function cannot be found by its name"
+    try:
+        pickle.dumps(sample_item)
+    except Exception:
+        return "items cannot be pickled"
+    return None
+
+
+def _pickled_size(obj):
+    """Bytes after pickling, or None if obj cannot be pickled."""
+    try:
+        return len(pickle.dumps(obj))
+    except Exception:
+        return None
+
+
+def _rule_on_source(func):
+    """Fallback when the probe cannot run: threads if the code waits or calls GIL-free C code."""
+    static = ast_features(func)
+    return "threading" if static["has_io_call"] or static["has_gil_release_call"] else "sequential"
+
+
+class _Recorder:
+    """Wraps the user function during the probe.
+
+    - keeps the outputs, so if probing fails half-way (an output that cannot be
+      pickled) the items already computed are not lost or run twice;
+    - keeps exceptions raised in the probe's helper threads, which Python would
+      otherwise only print and swallow.
+    It hashes/compares equal to the wrapped function, so the cached AST
+    features of that function are reused.
+    """
+
+    def __init__(self, func):
+        self.func = func
+        self.__wrapped__ = func                   # inspect.getsource() follows this
+        self.caller = threading.get_ident()
+        self.outputs = []
+        self.errors = []
+
+    def __hash__(self):
+        return hash(self.func)
+
+    def __eq__(self, other):
+        return other is self.func or (isinstance(other, _Recorder) and other.func is self.func)
+
+    def __call__(self, item):
+        in_caller = threading.get_ident() == self.caller
+        try:
+            out = self.func(item)
+        except BaseException as e:
+            self.errors.append(e)
+            if in_caller:
+                raise
+            return None
+        if in_caller:
+            self.outputs.append(out)
+        return out
+
+
+# ----------------------------------------------------------------- the core
+
+def _run_mode(mode, func, target, items, n_workers):
+    if not items:
+        return []
+    if mode == "sequential":
+        return [func(item) for item in items]
+    pool = _get_pool(mode, n_workers)
+    # default chunksize, same as in the data collection
+    return pool.map(target if mode == "multiprocessing" else func, items)
+
+
+def _best_allowed(probabilities, allowed):
+    return max(allowed, key=lambda m: probabilities[m])
+
+
+def run(func, items, n_workers=None, mode=None, cache=True, *, _target=None, _name=None):
+    """Compute [func(item) for item in items] in the best mode.
+
+    Returns (results, Decision). `mode` forces one mode (also settable for the
+    whole program with the environment variable POLYMORPH_MODE). With `cache`,
+    a later call that looks the same (similar item count, first-item time and
+    item size) reuses the decision instead of probing again.
+    """
+    t_start = time.perf_counter()
+    target = _target or func                    # what multiprocessing pickles
+    if not isinstance(items, (list, tuple)):
+        items = list(items)
+    n_workers = max(1, int(n_workers or usable_cpu_count()))
+    d = Decision(function=_name or getattr(func, "__qualname__", repr(func)),
+                 mode="sequential", reason="", n_items=len(items), n_workers=n_workers)
+
+    def finish(results):
+        d.total_s = time.perf_counter() - t_start
+        log.debug("%s", d)
+        return results, d
+
+    forced = mode or os.environ.get("POLYMORPH_MODE")
+    if forced:
+        if forced not in MODES:
+            raise ValueError(f"mode must be one of {MODES}, got {forced!r}")
+        if forced == "multiprocessing" and items:
+            problem = _process_problem(target, items[0])
+            if problem:
+                raise ValueError(f"cannot use multiprocessing: {problem}")
+        d.mode, d.reason = forced, "forced"
+        return finish(list(_run_mode(forced, func, target, list(items), n_workers)))
+
+    # ---- 1. cases where measuring is pointless or impossible
+    if getattr(_worker, "active", False):
+        d.reason = "nested call inside a polymorph_ai worker"
+        return finish([func(item) for item in items])
+    if len(items) <= PROBE_MIN_ITEMS or n_workers == 1:
+        d.reason = "too few items" if n_workers > 1 else "only 1 worker"
+        return finish([func(item) for item in items])
+
+    allowed = list(MODES)
+    problem = _process_problem(target, items[0])
+    if problem:
+        allowed.remove("multiprocessing")
+        d.notes.append("no multiprocessing: " + problem)
+
+    item_bytes = _pickled_size(items[0])
+    if item_bytes is None:
+        # The probe measures pickling, so it cannot run. Fall back to the code itself.
+        d.mode, d.reason = _rule_on_source(func), "items cannot be pickled -> rule on the source code"
+        return finish(_run_mode(d.mode, func, target, list(items), n_workers))
+
+    # ---- 2. time the first item (useful work): is the whole job tiny?
+    t0 = time.perf_counter()
+    first = [func(items[0])]
+    first_s = time.perf_counter() - t0
+    if first_s * len(items) < SMALL_JOB_S:
+        d.reason, d.probed_items = "tiny job (from the first item's time)", 1
+        return finish(first + [func(item) for item in items[1:]])
+
+    # ---- 3. same kind of call as before? reuse that decision, skip the probe
+    key = (func, n_workers, round(math.log2(len(items))), round(math.log2(item_bytes + 1) / 2),
+           round(math.log2(max(first_s, 1e-7)))) if cache else None
+    cached = None
+    if key:   # one timing can land just across a bucket edge, so accept the neighbouring buckets too
+        cached = next((_decisions[k] for k in (key, key[:-1] + (key[-1] - 1,), key[:-1] + (key[-1] + 1,))
+                       if k in _decisions), None)
+    if cached and cached[1] > 0 and (cached[0] != "multiprocessing" or "multiprocessing" in allowed):
+        cached[1] -= 1
+        d.mode, d.reason, d.probed_items = cached[0], "same as an earlier similar call (cached)", 1
+        d.overhead_s = time.perf_counter() - t0 - first_s
+        return finish(first + _run_mode(d.mode, func, target, list(items[1:]), n_workers))
+
+    # ---- 4. probe the next items (they are really computed and kept)
+    ast_features(func)                          # cache under func itself, not the recorder
+    rec = _Recorder(func)
+    try:
+        features, done, _, extra = extract_features(rec, items[1:], n_workers)
+    except Exception:
+        if rec.errors:                          # the user's function raised: behave like a plain loop
+            raise
+        # an output could not be pickled: keep what was computed, finish without processes
+        done = first + rec.outputs
+        d.mode, d.reason = _rule_on_source(func), "results cannot be pickled -> rule on the source code"
+        d.probed_items = len(done)
+        return finish(done + _run_mode(d.mode, func, target, list(items[len(done):]), n_workers))
+    if rec.errors:                              # raised inside a thread-probe thread
+        raise rec.errors[0]
+
+    features["n_items"] = len(items)            # the model is about the whole job
+    done = first + done
+    d.features, d.probed_items = features, len(done)
+    rest = list(items[len(done):])
+    t_decide = time.perf_counter()
+
+    # ---- 5. decide
+    if not rest:
+        d.reason = "probe already finished every item"
+    elif features["time_per_item"] * len(rest) < SMALL_JOB_S:
+        d.reason = "remaining work too small for any pool"
+        if key:
+            _decisions[key] = ["sequential", CACHE_USES]
+    else:
+        best, d.probabilities = _get_model().predict(features)
+        d.mode, d.reason = _best_allowed(d.probabilities, allowed), "model"
+        if d.mode != best:
+            d.reason = f"model ({best} not possible)"
+        if key:     # cached before the start-up check: a repeated job is worth starting the pool for
+            _decisions[key] = [d.mode, CACHE_USES]
+        if d.mode == "multiprocessing" and not _pool_is_warm("multiprocessing", n_workers):
+            # The model learned on warm pools. If even a perfect speed-up would
+            # save less than starting the pool costs, do not start it for this call.
+            ideal_saving = features["time_per_item"] * len(rest) * (1 - 1 / min(n_workers, features["cpu_count"]))
+            if ideal_saving < COLD_START_S.get(_start_method(), 0.25):
+                d.mode = _best_allowed(d.probabilities, ["sequential", "threading"])
+                d.reason = "model said multiprocessing, but pool start-up would cost more than it saves"
+    d.overhead_s = extra + (time.perf_counter() - t_decide)
+
+    # ---- 6. run the rest
+    return finish(done + _run_mode(d.mode, func, target, rest, n_workers))
+
+
+# ---------------------------------------------------------------- decorator
+
+class AdaptiveExec:
+    """The object @adaptive_exec puts in place of your function."""
+
+    def __init__(self, func, n_workers=None, mode=None, cache=True, verbose=False):
+        functools.update_wrapper(self, func)
+        self.func = func
+        self.n_workers = n_workers
+        self.mode = mode
+        self.cache = cache
+        self.verbose = verbose
+        self.last_run = None
+
+    def __call__(self, item):
+        return self.func(item)
+
+    def map(self, items, n_workers=None, mode=None, cache=None):
+        """[self(item) for item in items], run in the mode polymorph_ai picks."""
+        results, self.last_run = run(self.func, items, n_workers or self.n_workers, mode or self.mode,
+                                     self.cache if cache is None else cache,
+                                     _target=self, _name=self.__qualname__)
+        if self.verbose:
+            print(self.last_run, file=sys.stderr)
+        return results
+
+    def __reduce__(self):
+        if _defined_without_file(self) and cloudpickle is not None:
+            # Written in Jupyter/REPL: there is no file to import it from, so send
+            # the code itself. Inside that cloudpickle call, a reference back to
+            # this wrapper (e.g. a recursive function calling itself by name) is
+            # rebuilt from the function, which cloudpickle has already memoised.
+            if getattr(_by_value, "active", False):
+                return _rewrap, (self.func, self.n_workers, self.mode, self.cache, self.verbose)
+            _by_value.active = True
+            try:
+                data = cloudpickle.dumps(self)
+            finally:
+                _by_value.active = False
+            return cloudpickle.loads, (data,)
+        # Worker processes import the module and look the function up by name,
+        # which finds this object again (pickling self.func by name would fail:
+        # its name now points at this wrapper).
+        return _find, (self.__module__, self.__qualname__)
+
+    def __repr__(self):
+        return f"<adaptive_exec {self.__module__}.{self.__qualname__}>"
+
+
+_by_value = threading.local()
+
+
+def _rewrap(func, n_workers, mode, cache, verbose):
+    return AdaptiveExec(func, n_workers, mode, cache, verbose)
+
+
+def adaptive_exec(func=None, *, n_workers=None, mode=None, cache=True, verbose=False):
+    """Decorator. Use as @adaptive_exec or @adaptive_exec(n_workers=4, verbose=True).
+
+    n_workers: pool size (default: usable CPUs)   mode: force one mode
+    cache: reuse decisions for similar calls       verbose: print every decision
+    """
+    if func is None:
+        return lambda f: AdaptiveExec(f, n_workers, mode, cache, verbose)
+    return AdaptiveExec(func, n_workers, mode, cache, verbose)

polymorph_ai/features.py

@@ -0,0 +1,289 @@
+"""Feature extraction shared by the data collector and the decorator.
+
+The same code must be used in both places, otherwise the model is trained on
+features that the decorator measures differently at run time.
+
+13 features in 4 groups (see CLAUDE.md):
+  call-site    : n_items, input_bytes, cpu_count, n_workers
+  ast          : max_loop_depth, arith_op_count, has_io_call, has_gil_release_call
+  probe        : time_per_item, cpu_ratio, time_cv, pickle_time_ratio
+  thread probe : thread_probe_speedup
+
+A workload is "apply func to every item in a list", like map(func, items).
+"""
+
+import ast
+import functools
+import inspect
+import os
+import pickle
+import statistics
+import textwrap
+import threading
+import time
+
+FEATURE_NAMES = [
+    # call-site
+    "n_items", "input_bytes", "cpu_count", "n_workers",
+    # ast
+    "max_loop_depth", "arith_op_count", "has_io_call", "has_gil_release_call",
+    # probe
+    "time_per_item", "cpu_ratio", "time_cv", "pickle_time_ratio",
+    # thread probe
+    "thread_probe_speedup",
+]
+
+# The probe runs real items one by one until it has used this much time
+# (at least PROBE_MIN_ITEMS, at most a quarter of all items or PROBE_MAX_ITEMS). The outputs are
+# kept, so the decorator does not waste this work.
+PROBE_TARGET_S = 0.05
+PROBE_MIN_ITEMS = 3
+PROBE_MAX_ITEMS = 2000      # tiny items: more adds bookkeeping, not information
+PICKLE_SAMPLE = 20         # items used to measure the pickling cost
+
+
+# ---------------------------------------------------------------- cpu clock
+# CPU time used by the current thread, in seconds.
+# Linux: time.thread_time() is precise. Windows: thread_time()/process_time()
+# only tick every 15.6 ms, useless for a probe of a few ms. Instead we read
+# the thread's CPU cycle counter (QueryThreadCycleTime) and convert cycles to
+# seconds with a rate measured once by a short busy loop.
+
+if os.name == "nt":
+    import ctypes
+    from ctypes import wintypes
+
+    _kernel32 = ctypes.WinDLL("kernel32")
+    _kernel32.GetCurrentThread.restype = wintypes.HANDLE
+    _kernel32.QueryThreadCycleTime.argtypes = [wintypes.HANDLE,
+                                               ctypes.POINTER(ctypes.c_ulonglong)]
+
+    def _thread_cycles():
+        cycles = ctypes.c_ulonglong()
+        _kernel32.QueryThreadCycleTime(_kernel32.GetCurrentThread(), ctypes.byref(cycles))
+        return cycles.value
+
+    @functools.lru_cache(maxsize=1)
+    def _cycles_per_second():
+        # Busy loop: the thread uses 100% CPU, so cycles / wall = rate.
+        c0, t0 = _thread_cycles(), time.perf_counter()
+        while time.perf_counter() - t0 < 0.02:
+            pass
+        return (_thread_cycles() - c0) / (time.perf_counter() - t0)
+
+    def thread_cpu_time():
+        rate = _cycles_per_second()       # calibrate first, outside the reading
+        return _thread_cycles() / rate
+else:
+    thread_cpu_time = time.thread_time
+
+
+# ---------------------------------------------------------------- call-site
+
+def usable_cpu_count():
+    """Number of CPUs this process may actually use (not just installed)."""
+    if hasattr(os, "process_cpu_count"):          # Python 3.13+
+        n = os.process_cpu_count()
+    elif hasattr(os, "sched_getaffinity"):        # Linux
+        n = len(os.sched_getaffinity(0))
+    else:                                         # Windows, older Python
+        n = os.cpu_count()
+    return n or 1
+
+
+# ---------------------------------------------------------------------- ast
+
+# Calls that mean "waits for I/O". Matched against the dotted call name,
+# e.g. `time.sleep(...)` -> "time.sleep".
+IO_CALLS = {"open", "print", "input", "time.sleep", "os.read", "os.write",
+            "os.fsync", "asyncio.sleep"}
+IO_PREFIXES = ("socket.", "requests.", "urllib.", "http.", "subprocess.",
+               "shutil.", "ftplib.", "smtplib.", "sqlite3.")
+# Method names that are almost always file/socket I/O: f.read(), s.recv() ...
+IO_METHODS = {"read", "write", "readline", "readlines", "recv", "send",
+              "sendall", "flush"}
+
+# C functions that do heavy work and release the GIL, so threads can run
+# them in parallel. Only functions we have checked; plain `math` is not here.
+GIL_RELEASE_CALLS = {"np.dot", "numpy.dot", "np.matmul", "numpy.matmul",
+                     "np.sort", "numpy.sort"}
+GIL_RELEASE_PREFIXES = ("zlib.", "bz2.", "lzma.", "hashlib.",
+                        "np.linalg.", "numpy.linalg.", "np.fft.", "numpy.fft.")
+
+LOOP_NODES = (ast.For, ast.AsyncFor, ast.While)
+COMPREHENSIONS = (ast.ListComp, ast.SetComp, ast.DictComp, ast.GeneratorExp)
+
+
+def _dotted_name(node):
+    """`zlib.compress` -> "zlib.compress", `f.read` -> "?.read" if f is complex."""
+    parts = []
+    while isinstance(node, ast.Attribute):
+        parts.append(node.attr)
+        node = node.value
+    parts.append(node.id if isinstance(node, ast.Name) else "?")
+    return ".".join(reversed(parts))
+
+
+def _loop_depth(node, depth=0):
+    if isinstance(node, LOOP_NODES):
+        depth += 1
+    elif isinstance(node, COMPREHENSIONS):
+        depth += len(node.generators)
+    child_depths = [_loop_depth(child, depth) for child in ast.iter_child_nodes(node)]
+    return max([depth] + child_depths)
+
+
+@functools.lru_cache(maxsize=256)
+def ast_features(func):
+    """Static features read from the function's source code (cached per function).
+
+    Only the function's own body is inspected, not the functions it calls.
+    """
+    try:
+        source = textwrap.dedent(inspect.getsource(func))
+        tree = ast.parse(source)
+    except (OSError, TypeError, SyntaxError):
+        # No source available (builtin, lambda in REPL, ...): neutral values.
+        return {"max_loop_depth": 0, "arith_op_count": 0,
+                "has_io_call": 0, "has_gil_release_call": 0}
+
+    arith = 0
+    has_io = False
+    has_gil_release = False
+    for node in ast.walk(tree):
+        if isinstance(node, (ast.BinOp, ast.AugAssign)):
+            arith += 1
+            if isinstance(node.op, ast.MatMult):      # a @ b on numpy arrays
+                has_gil_release = True
+        elif isinstance(node, ast.Call):
+            name = _dotted_name(node.func)
+            method = name.rsplit(".", 1)[-1]
+            if (name in IO_CALLS or name.startswith(IO_PREFIXES)
+                    or ("." in name and method in IO_METHODS)):
+                has_io = True
+            if name in GIL_RELEASE_CALLS or name.startswith(GIL_RELEASE_PREFIXES):
+                has_gil_release = True
+
+    return {"max_loop_depth": _loop_depth(tree), "arith_op_count": arith,
+            "has_io_call": int(has_io), "has_gil_release_call": int(has_gil_release)}
+
+
+# -------------------------------------------------------------------- probe
+
+def probe(func, items):
+    """Run the first few real items sequentially and measure them.
+
+    Returns (features, results) where results are the outputs of the probed
+    items, so the decorator can reuse them instead of computing them twice.
+    """
+    max_items = min(len(items), PROBE_MAX_ITEMS, max(PROBE_MIN_ITEMS, len(items) // 4))
+    item_times, results, pickle_times, input_sizes = [], [], [], []
+
+    cpu_start = thread_cpu_time()
+    wall_start = time.perf_counter()
+    for item in items[:max_items]:
+        t0 = time.perf_counter()
+        out = func(item)
+        item_times.append(time.perf_counter() - t0)
+        results.append(out)
+        if (len(item_times) >= PROBE_MIN_ITEMS
+                and time.perf_counter() - wall_start >= PROBE_TARGET_S):
+            break
+    wall = time.perf_counter() - wall_start
+    cpu = thread_cpu_time() - cpu_start
+
+    # What multiprocessing would pay per item: send input, get output back.
+    # Measured outside the timed loop above so it does not inflate cpu_ratio.
+    # A spread-out sample is enough for the mean; tiny items can mean
+    # tens of thousands of probed items.
+    step = max(1, len(results) // PICKLE_SAMPLE)
+    for item, out in list(zip(items, results))[::step][:PICKLE_SAMPLE]:
+        t0 = time.perf_counter()
+        data = pickle.dumps(item)
+        pickle.loads(data)
+        pickle.loads(pickle.dumps(out))
+        pickle_times.append(time.perf_counter() - t0)
+        input_sizes.append(len(data))
+
+    time_per_item = statistics.mean(item_times)
+    features = {
+        "input_bytes": statistics.mean(input_sizes),
+        "time_per_item": time_per_item,
+        "cpu_ratio": min(cpu / wall, 1.0) if wall > 0 else 0.0,
+        "time_cv": (statistics.pstdev(item_times) / time_per_item
+                    if time_per_item > 0 else 0.0),
+        "pickle_time_ratio": (statistics.mean(pickle_times) / time_per_item
+                              if time_per_item > 0 else 0.0),
+    }
+    return features, results, item_times
+
+
+def thread_probe(func, items, time_per_item):
+    """Run the NEXT items (not the probed ones again) on 2 threads.
+
+    Compares their wall time with what they would take sequentially
+    (len(items) * time_per_item from the probe):
+    ~2.0 means threads really run in parallel (I/O or GIL-releasing code),
+    ~1.0 or less means the GIL serialises them.
+    The outputs are kept, so this work is not wasted. Items are never run
+    twice, so functions with side effects (writing files, ...) stay correct.
+
+    Returns (speedup, results). With fewer than 2 items left: (1.0, []).
+    """
+    if len(items) < 2:
+        return 1.0, []
+    results = [None] * len(items)
+
+    def work(indexes):
+        for i in indexes:
+            results[i] = func(items[i])
+
+    halves = [range(0, len(items), 2), range(1, len(items), 2)]
+    threads = [threading.Thread(target=work, args=(h,)) for h in halves]
+    t0 = time.perf_counter()
+    for t in threads:
+        t.start()
+    for t in threads:
+        t.join()
+    wall = time.perf_counter() - t0
+    speedup = len(items) * time_per_item / wall if wall > 0 else 1.0
+    return speedup, results
+
+
+# ---------------------------------------------------------------------- all
+
+def extract_features(func, items, n_workers):
+    """All 13 features for running func over items with n_workers workers.
+
+    Returns (features, probe_results, overhead_s, extra_s).
+    probe_results are the outputs of the first len(probe_results) items.
+    overhead_s is the total time spent here; extra_s is the part that did
+    no useful work (pickle test, AST, bookkeeping), i.e. overhead_s minus
+    the time spent computing real items.
+    """
+    t0 = time.perf_counter()
+    probe_feats, results, item_times = probe(func, items)
+    t_probe = time.perf_counter()
+    # Skip the first probed item when estimating sequential speed: it may
+    # include one-off warm-up.
+    steady = item_times[1:] if len(item_times) >= 3 else item_times
+    next_items = items[len(results):len(results) + len(results)]
+    speedup, thread_results = thread_probe(func, next_items, statistics.mean(steady))
+    t_thread = time.perf_counter()
+    results = results + thread_results
+
+    features = {
+        "n_items": len(items),
+        "input_bytes": probe_feats["input_bytes"],
+        "cpu_count": usable_cpu_count(),
+        "n_workers": n_workers,
+        **ast_features(func),
+        "time_per_item": probe_feats["time_per_item"],
+        "cpu_ratio": probe_feats["cpu_ratio"],
+        "time_cv": probe_feats["time_cv"],
+        "pickle_time_ratio": probe_feats["pickle_time_ratio"],
+        "thread_probe_speedup": speedup,
+    }
+    overhead = time.perf_counter() - t0
+    useful = sum(item_times) + (t_thread - t_probe if thread_results else 0.0)
+    return features, results, overhead, max(overhead - useful, 0.0)

polymorph_ai/model.py

@@ -0,0 +1,48 @@
+"""The trained mode-selection MLP, run with plain numpy.
+
+The weights come from notebooks/04_export_model.ipynb (Keras, trained on the
+collected dataset). Re-implementing the forward pass here keeps the library's
+only dependency at numpy: no Keras/JAX import (seconds) at run time, and one
+decision takes microseconds.
+
+Steps, identical to training:
+  log10 of the heavy-tailed features -> standardise -> Dense+ReLU ... -> Dense -> softmax
+"""
+
+import json
+import os
+
+import numpy as np
+
+DEFAULT_PATH = os.path.join(os.path.dirname(os.path.abspath(__file__)), "model.npz")
+
+
+class ModeModel:
+    def __init__(self, path=DEFAULT_PATH):
+        data = np.load(path)
+        self.feature_names = [str(f) for f in data["feature_names"]]
+        self.modes = [str(m) for m in data["modes"]]
+        self.meta = json.loads(str(data["meta"]))
+        log_features = {str(f) for f in data["log_features"]}
+        self._is_log = np.array([f in log_features for f in self.feature_names])
+        self._mean = data["scaler_mean"]
+        self._scale = data["scaler_scale"]
+        n = int(data["n_layers"])
+        self._layers = [(data[f"w{2 * i}"], data[f"w{2 * i + 1}"]) for i in range(n)]
+
+    def predict_proba(self, features):
+        """features: dict with the 13 feature names -> probabilities in self.modes order."""
+        x = np.array([float(features[f]) for f in self.feature_names])
+        x[self._is_log] = np.log10(np.maximum(x[self._is_log], 1e-9))
+        h = (x - self._mean) / self._scale
+        for i, (w, b) in enumerate(self._layers):
+            h = h @ w + b
+            if i < len(self._layers) - 1:
+                h = np.maximum(h, 0.0)                     # ReLU (dropout is off at inference)
+        e = np.exp(h - h.max())
+        return e / e.sum()                                 # softmax
+
+    def predict(self, features):
+        """Most likely mode and the probability of every mode (dict)."""
+        p = self.predict_proba(features)
+        return self.modes[int(p.argmax())], dict(zip(self.modes, p.tolist()))

polymorph_ai-1.1.0.dist-info/METADATA

@@ -0,0 +1,128 @@
+Metadata-Version: 2.4
+Name: polymorph-ai
+Version: 1.1.0
+Summary: A decorator that uses a trained neural network to run your function sequentially, with threads, or with processes - whichever is fastest.
+Author: Worachat Songmuangnu
+License-Expression: MIT
+Project-URL: Homepage, https://github.com/Worachat-Songmuangnu/polymorph-ai
+Project-URL: Source, https://github.com/Worachat-Songmuangnu/polymorph-ai
+Project-URL: Issues, https://github.com/Worachat-Songmuangnu/polymorph-ai/issues
+Keywords: parallel,multiprocessing,threading,decorator,machine-learning,performance
+Classifier: Development Status :: 4 - Beta
+Classifier: Intended Audience :: Developers
+Classifier: Operating System :: Microsoft :: Windows
+Classifier: Operating System :: POSIX :: Linux
+Classifier: Programming Language :: Python :: 3
+Classifier: Topic :: Software Development :: Libraries
+Classifier: Topic :: System :: Distributed Computing
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: numpy
+Requires-Dist: cloudpickle
+Dynamic: license-file
+
+# polymorph-ai
+
+`@adaptive_exec` runs a function over a list of items **sequentially, with threads, or with processes**, and a small trained neural network picks whichever mode should be fastest for that job on that machine.
+
+```python
+from polymorph_ai import adaptive_exec
+
+@adaptive_exec
+def process(item):
+    ...
+
+if __name__ == "__main__":            # required on Windows/macOS for processes
+    results = process.map(items)      # same as [process(x) for x in items], in order
+    print(process.last_run)
+    # [polymorph_ai] process: 5,000 items -> multiprocessing p=0.97 (model) | overhead 0.7 ms | total 1.84 s
+```
+
+`process(x)` still works as a normal function call. Only `.map()` adds the automatic part.
+
+## Install
+
+```bash
+pip install polymorph-ai
+```
+
+Dependencies: numpy and cloudpickle. Tested on Linux with Python 3.10, 3.12 and 3.14, and on Windows with Python 3.14.
+
+Works in Jupyter too: a function written in a notebook cell is sent to the worker processes by value (with cloudpickle), so it can still use multiprocessing on Windows and macOS.
+
+## How it decides
+
+Every `.map(items)` call goes through these steps:
+
+1. **Safety checks.** Empty or very short lists, one worker, or a call nested inside another polymorph_ai worker all just run as a plain loop. If the function or the items cannot be sent to another process (a lambda, a function defined inside another function, unpicklable items), multiprocessing is ruled out.
+2. **First item.** The first item runs and is timed. If `time × number of items` is under 10 ms, the job is too small for any pool to pay off, so the rest runs as a plain loop.
+3. **Cache.** A call that looks like an earlier one (similar number of items, first-item time, and item size) reuses that decision and skips the probe.
+4. **Probe.** The next items run for about 50 ms, first one by one and then on 2 threads, to measure the 13 features: call site, static code analysis (AST), runtime probe, and thread probe. The probe uses the same code (`polymorph_ai/features.py`) that collected the training data. **Probed items are real work: their results are kept and no item ever runs twice.**
+5. **Model.** An MLP (13 → 225 ReLU → 3 softmax, 3,828 weights) turns the features into probabilities. It is trained in Keras and runs here in numpy, at about 30 µs per decision.
+6. **Run.** The remaining items run in the chosen mode on a pool that stays alive for the next call.
+
+## Options
+
+```python
+@adaptive_exec(n_workers=4, verbose=True, cache=True)   # verbose prints every decision
+def process(item): ...
+
+process.map(items, mode="threading")      # force one mode for this call
+# POLYMORPH_MODE=sequential python app.py  # force one mode for the whole program (A/B testing, debugging)
+
+from polymorph_ai import run, warm_up
+results, decision = run(some_function, items)   # for functions you cannot decorate
+warm_up()                                      # start the pools early (e.g. at server start-up)
+```
+
+`process.last_run` (a `Decision`) records the chosen mode, why it was chosen, the model's probabilities, the 13 features, how many items the probe computed, and the overhead.
+
+## Results
+
+**Model** (13,603 jobs collected on 16 machine settings, see [`dataset/`](https://github.com/Worachat-Songmuangnu/polymorph-ai/tree/main/dataset), GroupKFold by workload template, so every test job comes from code the model never trained on). Time lost compared with always picking the fastest mode:
+
+- Always multiprocessing: +22.1%
+- if-else rules: +13.0%
+- Random Forest: +5.8%
+- **MLP: +4.3%**
+
+**Library** ([`examples/demo.py`](https://github.com/Worachat-Songmuangnu/polymorph-ai/blob/main/examples/demo.py): 8 everyday jobs that are not in the training data, 8 CPUs). Total time compared with perfect picks:
+
+| | Windows | Linux (WSL) |
+|---|---|---|
+| always sequential | +349% | +426% |
+| always threading | +117% | +138% |
+| always multiprocessing | +19% | +30% |
+| polymorph_ai, first call | +52% | +66% |
+| polymorph_ai, repeated call | **+13%** | **+24%** |
+
+**When does it pay off?** The first call on a job pays about 0.2 s for the probe (`python examples/demo.py --overhead`).
+
+- Compared with a plain loop, polymorph_ai is already faster on jobs from about 0.3–0.6 s.
+- Compared with a perfect choice of mode, the first call is within 15% once the job takes 5 s or more.
+- Repeated calls skip the probe, so they cost almost nothing.
+
+## Limitations
+
+- The decorated function takes **one item** and is applied to every item (`map`). polymorph_ai cannot parallelise code that does not have this shape.
+- On Windows and macOS, the script that uses processes needs the `if __name__ == "__main__":` guard (a Python rule for processes, not specific to polymorph_ai).
+- The first item and the probe items run in the calling thread, so a job with a few very long items loses one item's worth of parallel time.
+- The model was trained on synthetic workloads from 13 families. Code that behaves unlike all of them can still be mispredicted. In the demo, sorting 20k-number lists stays sequential when processes would be twice as fast.
+- asyncio is not supported: a decorator cannot turn ordinary code into `async def` code.
+- In Jupyter on Windows, call `shutdown()` before you close or restart the kernel. If the kernel is stopped straight after the cell that started the process pool, the worker processes can be left running.
+
+## Tests
+
+```bash
+python -m pytest tests -v    # 28 tests, pass on Windows and Linux
+```
+
+## Project layout
+
+```
+polymorph_ai/   the library: decorator.py, features.py, model.py + model.npz (the trained network)
+tests/          pytest suite
+examples/       demo.py (benchmark vs fixed modes) and its results/
+dataset/        the training dataset (13,603 jobs) and the real-code results (92 jobs)
+```

polymorph_ai-1.1.0.dist-info/RECORD

@@ -0,0 +1,10 @@
+polymorph_ai/__init__.py,sha256=1aFA5RwshexIWECw1ihJu1b7A9FWZKvZCV06iF5q6lM,738
+polymorph_ai/decorator.py,sha256=zB5MnztiSfOOWXYmwXBbrAgWSlKD2g_ziR2gSngr_ig,19980
+polymorph_ai/features.py,sha256=JgSjN6FW6elrlYmwXNmvj_VavwxT6aty7MeBIFHOpFw,11818
+polymorph_ai/model.npz,sha256=PYXE0nZhIMtUiSw5JjXBhk-rPuWl2I_qpclF5ITfJRg,21120
+polymorph_ai/model.py,sha256=GrOTV-TpepgZlaDamEDiLpWEiUi0-iT0XdZTg1OEOKw,2089
+polymorph_ai-1.1.0.dist-info/licenses/LICENSE,sha256=H0aLa4UfHVK6fzZHZ_aEMtgywY18sh6wa71_q0H4thg,1081
+polymorph_ai-1.1.0.dist-info/METADATA,sha256=I1DTMUvZb-tR2eV-FvGCkJzuafExt_14JGIRDyP8fHg,6998
+polymorph_ai-1.1.0.dist-info/WHEEL,sha256=YVMoNqKzERt-wjUZwJ33xBGAwnFl-4cqbYkTtWa4itE,91
+polymorph_ai-1.1.0.dist-info/top_level.txt,sha256=5F0pgkilc2sBAikhdTIpRrxPIxDMS2D1Ae9qrLO6K5A,13
+polymorph_ai-1.1.0.dist-info/RECORD,,

polymorph_ai-1.1.0.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (84.0.0)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

polymorph_ai-1.1.0.dist-info/licenses/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 the polymorph-ai authors
+
+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.

polymorph_ai-1.1.0.dist-info/top_level.txt

@@ -0,0 +1 @@
+polymorph_ai