stringjax-tools 0.1.0__tar.gz → 0.2.0__tar.gz

{stringjax_tools-0.1.0/src/stringjax_tools.egg-info → stringjax_tools-0.2.0}/PKG-INFO

@@ -1,22 +1,22 @@
 Metadata-Version: 2.4
 Name: stringjax-tools
-Version: 0.1.0
-Summary: Reusable JAX transform, vectorisation, cache, and pytree helpers for StringJAX packages.
+Version: 0.2.0
+Summary: Reusable JAX transform, vectorisation, cache and pytree helpers, and general-purpose infrastructure (child processes with enforced timeouts, provenance), for StringJAX packages.
 Author-email: Andreas Schachner <as3475@cornell.edu>
 License-Expression: GPL-3.0-or-later
-Project-URL: Homepage, https://github.com/StringJAX/stringjax-tools
-Project-URL: Issues, https://github.com/StringJAX/stringjax-tools/issues
-Keywords: jax,vmap,jit,pytrees,automatic-vectorisation,scientific-computing
+Project-URL: Homepage, https://github.com/StringJAX/stringjax_tools
+Project-URL: Documentation, https://stringjax-tools.readthedocs.io
+Project-URL: Changelog, https://github.com/StringJAX/stringjax_tools/blob/main/CHANGELOG.md
+Project-URL: Issues, https://github.com/StringJAX/stringjax_tools/issues
+Keywords: jax,stringjax,vmap,jit,pytrees,automatic-vectorisation,scientific-computing
 Classifier: Development Status :: 3 - Alpha
 Classifier: Intended Audience :: Developers
 Classifier: Intended Audience :: Science/Research
 Classifier: Programming Language :: Python :: 3
-Classifier: Programming Language :: Python :: 3.10
-Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
 Classifier: Programming Language :: Python :: 3.13
 Classifier: Topic :: Scientific/Engineering
-Requires-Python: >=3.10
+Requires-Python: >=3.12
 Description-Content-Type: text/markdown
 License-File: LICENSE
 Requires-Dist: jax>=0.4.26
@@ -44,13 +44,20 @@ Dynamic: license-file
 # StringJAX Tools
 
 `stringjax_tools` is a small utility package for reusable JAX transform
-patterns used across StringJAX projects.  It currently focuses on:
+patterns used across [StringJAX](https://github.com/StringJAX) projects.  Use
+it when a package needs the same JAX mechanics in several places, but should
+not carry domain-specific names or conventions into a shared helper layer.
+
+It currently focuses on:
 
 - conservative rank-checked automatic vectorisation;
 - manual cached `jax.jit(jax.vmap(...))` wrappers;
 - small helpers for explicit static-argument JIT compilation;
 - explicit JAX persistent-compilation-cache setup;
-- configurable pytree registration for stateful model classes.
+- configurable pytree registration for stateful model classes;
+- running a call in a child process with an enforced timeout
+  (`run_in_subprocess`);
+- recording which code produced a result (`provenance`).
 
 The package keeps reusable transformation and pytree machinery separate from
 domain-specific StringJAX model code.
@@ -63,7 +70,7 @@ From this directory:
 pip install -e ".[dev,docs]"
 ```
 
-The intended published package name is:
+From PyPI:
 
 ```bash
 pip install stringjax-tools
@@ -92,6 +99,29 @@ make -C documentation html SPHINXOPTS="-W -q"
 The feature notebooks are included under `documentation/source/notebooks/` and
 are linked from the Sphinx tutorials page.
 
+Start with the compact feature tour:
+
+```text
+documentation/source/notebooks/00_crucial_features.ipynb
+```
+
+It shows the pieces most StringJAX packages usually need first:
+
+- `auto_vmap` for scalar-or-leading-batch public methods;
+- exact shape validation for package-local model conventions;
+- `vmapping_func_cached` for explicit `in_axes` kernels;
+- static-argument JIT helpers;
+- `PytreePolicy` for stateful model classes;
+- explicit persistent compilation-cache setup.
+
+## Package Boundary
+
+`stringjax_tools` is intentionally application-agnostic.  Consuming packages
+should define their own shape names, static keys, ignored cache attributes and
+state policies locally.  For example, a downstream package can wrap
+`stringjax_tools.auto_vmap` with package-specific defaults, while
+`stringjax_tools` itself remains free of domain language.
+
 ## Automatic Vectorisation
 
 Keep the single-sample function explicit, then wrap it at the boundary where
@@ -143,6 +173,12 @@ class Model:
 String dimensions are resolved from a bound object, usually `self`.  Derived
 dimensions should be callables rather than expression strings.
 
+The same public API can be imported from the package root when convenient:
+
+```python
+from stringjax_tools import auto_vmap
+```
+
 ## Manual Transform Helpers
 
 When a function already has a deliberate `in_axes` structure, use:

{stringjax_tools-0.1.0 → stringjax_tools-0.2.0}/README.md

@@ -1,13 +1,20 @@
 # StringJAX Tools
 
 `stringjax_tools` is a small utility package for reusable JAX transform
-patterns used across StringJAX projects.  It currently focuses on:
+patterns used across [StringJAX](https://github.com/StringJAX) projects.  Use
+it when a package needs the same JAX mechanics in several places, but should
+not carry domain-specific names or conventions into a shared helper layer.
+
+It currently focuses on:
 
 - conservative rank-checked automatic vectorisation;
 - manual cached `jax.jit(jax.vmap(...))` wrappers;
 - small helpers for explicit static-argument JIT compilation;
 - explicit JAX persistent-compilation-cache setup;
-- configurable pytree registration for stateful model classes.
+- configurable pytree registration for stateful model classes;
+- running a call in a child process with an enforced timeout
+  (`run_in_subprocess`);
+- recording which code produced a result (`provenance`).
 
 The package keeps reusable transformation and pytree machinery separate from
 domain-specific StringJAX model code.
@@ -20,7 +27,7 @@ From this directory:
 pip install -e ".[dev,docs]"
 ```
 
-The intended published package name is:
+From PyPI:
 
 ```bash
 pip install stringjax-tools
@@ -49,6 +56,29 @@ make -C documentation html SPHINXOPTS="-W -q"
 The feature notebooks are included under `documentation/source/notebooks/` and
 are linked from the Sphinx tutorials page.
 
+Start with the compact feature tour:
+
+```text
+documentation/source/notebooks/00_crucial_features.ipynb
+```
+
+It shows the pieces most StringJAX packages usually need first:
+
+- `auto_vmap` for scalar-or-leading-batch public methods;
+- exact shape validation for package-local model conventions;
+- `vmapping_func_cached` for explicit `in_axes` kernels;
+- static-argument JIT helpers;
+- `PytreePolicy` for stateful model classes;
+- explicit persistent compilation-cache setup.
+
+## Package Boundary
+
+`stringjax_tools` is intentionally application-agnostic.  Consuming packages
+should define their own shape names, static keys, ignored cache attributes and
+state policies locally.  For example, a downstream package can wrap
+`stringjax_tools.auto_vmap` with package-specific defaults, while
+`stringjax_tools` itself remains free of domain language.
+
 ## Automatic Vectorisation
 
 Keep the single-sample function explicit, then wrap it at the boundary where
@@ -100,6 +130,12 @@ class Model:
 String dimensions are resolved from a bound object, usually `self`.  Derived
 dimensions should be callables rather than expression strings.
 
+The same public API can be imported from the package root when convenient:
+
+```python
+from stringjax_tools import auto_vmap
+```
+
 ## Manual Transform Helpers
 
 When a function already has a deliberate `in_axes` structure, use:

{stringjax_tools-0.1.0 → stringjax_tools-0.2.0}/pyproject.toml

@@ -4,10 +4,10 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "stringjax-tools"
-version = "0.1.0"
-description = "Reusable JAX transform, vectorisation, cache, and pytree helpers for StringJAX packages."
+version = "0.2.0"
+description = "Reusable JAX transform, vectorisation, cache and pytree helpers, and general-purpose infrastructure (child processes with enforced timeouts, provenance), for StringJAX packages."
 readme = "README.md"
-requires-python = ">=3.10"
+requires-python = ">=3.12"
 license = "GPL-3.0-or-later"
 license-files = ["LICENSE"]
 authors = [
@@ -15,6 +15,7 @@ authors = [
 ]
 keywords = [
     "jax",
+    "stringjax",
     "vmap",
     "jit",
     "pytrees",
@@ -26,8 +27,6 @@ classifiers = [
     "Intended Audience :: Developers",
     "Intended Audience :: Science/Research",
     "Programming Language :: Python :: 3",
-    "Programming Language :: Python :: 3.10",
-    "Programming Language :: Python :: 3.11",
     "Programming Language :: Python :: 3.12",
     "Programming Language :: Python :: 3.13",
     "Topic :: Scientific/Engineering",
@@ -60,8 +59,10 @@ docs = [
 ]
 
 [project.urls]
-Homepage = "https://github.com/StringJAX/stringjax-tools"
-Issues = "https://github.com/StringJAX/stringjax-tools/issues"
+Homepage = "https://github.com/StringJAX/stringjax_tools"
+Documentation = "https://stringjax-tools.readthedocs.io"
+Changelog = "https://github.com/StringJAX/stringjax_tools/blob/main/CHANGELOG.md"
+Issues = "https://github.com/StringJAX/stringjax_tools/issues"
 
 [tool.setuptools.packages.find]
 where = ["src"]

{stringjax_tools-0.1.0 → stringjax_tools-0.2.0}/src/stringjax_tools/__init__.py

@@ -15,7 +15,7 @@
 # You should have received a copy of the GNU General Public License
 # along with StringJAX Tools. If not, see <https://www.gnu.org/licenses/>.
 
-"""Reusable JAX transform helpers for StringJAX packages."""
+"""Reusable JAX transform helpers and general-purpose infrastructure for StringJAX packages."""
 
 from .auto_vectorise import (
     ArgSpec,
@@ -32,18 +32,21 @@ from .auto_vectorise import (
 )
 from .cache import configure_compilation_cache
 from .jit import is_static, jit_with_dynamic_static_args, jit_with_static_args
+from .process import ChildProcessFailed, run_in_subprocess
 from .pytrees import (
     PytreePolicy,
     flatten_func,
     make_pytree_flatteners,
     unflatten_func_class,
 )
+from .reproducibility import git_state, package_info, provenance
 from .vmap import clear_vmap_caches, vmapping_func, vmapping_func_cached
 
-__version__ = "0.1.0"
+__version__ = "0.2.0"
 
 __all__ = [
     "ArgSpec",
+    "ChildProcessFailed",
     "PytreePolicy",
     "auto_vmap",
     "auto_vmap_default_shapes",
@@ -54,12 +57,16 @@ __all__ = [
     "flatten_func",
     "get_auto_vmap_default_shapes",
     "get_auto_vmap_defaults",
+    "git_state",
     "is_static",
     "jit_with_dynamic_static_args",
     "jit_with_static_args",
     "make_pytree_flatteners",
+    "package_info",
+    "provenance",
     "reset_auto_vmap_default_shapes",
     "reset_auto_vmap_defaults",
+    "run_in_subprocess",
     "set_auto_vmap_default_shapes",
     "set_auto_vmap_defaults",
     "unflatten_func_class",

stringjax_tools-0.2.0/src/stringjax_tools/_worker.py

@@ -0,0 +1,112 @@
+# Copyright 2023-2026 Andreas Schachner
+#
+# This file is part of StringJAX Tools.
+#
+# StringJAX Tools is free software: you can redistribute it and/or modify
+# it under the terms of the GNU General Public License as published by
+# the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# StringJAX Tools is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+# GNU General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License
+# along with StringJAX Tools. If not, see <https://www.gnu.org/licenses/>.
+
+"""Child process of :func:`stringjax_tools.process.run_in_subprocess`.
+
+Run as a script, ``python _worker.py <task> <result> <lifeline fd>``, so that it does not import
+the ``stringjax_tools`` package (nor JAX), only the module of the target. The task file holds two
+pickles: ``(target, sys.path)``, where ``target`` is ``"module:qualname"``, and then
+``(args, kwargs)``, unpickled after ``sys.path`` is set; the file is removed once read. The
+outcome is written atomically to ``<result>``: ``("ok", value, warnings)`` or ``("error",
+exception type name, message, traceback)``; then the worker kills its process group, so nothing
+the target started outlives the call. A daemon thread reads the lifeline pipe from the parent; at
+end of file (the parent is gone) it removes the temporary directory and kills the process group.
+The ``__main__`` guard is required: a target that starts workers with ``multiprocessing`` under
+``spawn`` (the macOS default) re-imports this file.
+"""
+
+import os
+import sys
+
+# Run as a script, Python puts this script's directory (inside the package) first on the path;
+# drop it before anything else is imported, so that no module of the package can shadow a
+# top-level one. Only then: a ``multiprocessing`` worker started under ``spawn`` re-runs this
+# file after restoring the parent's path, whose first entry must stay.
+_HERE = os.path.dirname(os.path.abspath(__file__))
+if sys.path and os.path.abspath(sys.path[0] or os.curdir) == _HERE:
+    del sys.path[0]
+
+import importlib  # noqa: E402
+import pickle  # noqa: E402
+import shutil  # noqa: E402
+import signal  # noqa: E402
+import threading  # noqa: E402
+import traceback  # noqa: E402
+import warnings  # noqa: E402
+
+
+def _watch_lifeline(fd, tmp_dir):
+    """Once the parent's end of the pipe closes: remove ``tmp_dir``, kill this process group."""
+    try:
+        while os.read(fd, 1):
+            pass
+    except OSError:
+        pass
+    shutil.rmtree(tmp_dir, ignore_errors=True)
+    os.killpg(0, signal.SIGKILL)
+
+
+def _resolve(target):
+    module_name, function_name = target.split(":", 1)
+    obj = importlib.import_module(module_name)
+    for part in function_name.split("."):
+        obj = getattr(obj, part)
+    return obj
+
+
+def _describe(exc):
+    try:
+        return str(exc)
+    except Exception:  # an exception whose __str__ fails
+        return f"<{type(exc).__name__}: str() failed>"
+
+
+def _flush():
+    for stream in (sys.stdout, sys.stderr):
+        try:
+            stream.flush()
+        except Exception:
+            pass
+
+
+def main(task_path, result_path, lifeline_fd):
+    threading.Thread(
+        target=_watch_lifeline, args=(lifeline_fd, os.path.dirname(task_path)), daemon=True
+    ).start()
+    try:
+        with open(task_path, "rb") as f:
+            target, path = pickle.load(f)
+            sys.path[:] = path
+            args, kwargs = pickle.load(f)
+        os.remove(task_path)
+        function = _resolve(target)
+        with warnings.catch_warnings(record=True) as caught:
+            warnings.simplefilter("always")
+            value = function(*args, **kwargs)
+        records = [(_describe(w.message), [c.__name__ for c in w.category.__mro__]) for w in caught]
+        payload = pickle.dumps(("ok", value, records))
+    except BaseException as e:  # report everything to the parent, including an interrupt here
+        payload = pickle.dumps(("error", type(e).__name__, _describe(e), traceback.format_exc()))
+    with open(result_path + ".tmp", "wb") as f:
+        f.write(payload)
+    os.replace(result_path + ".tmp", result_path)
+    _flush()
+    os.killpg(0, signal.SIGKILL)  # this process and anything the target left running
+
+
+if __name__ == "__main__":
+    main(sys.argv[1], sys.argv[2], int(sys.argv[3]))