galley-jl-python 0.1.0__py3-none-any.whl

galley_jl_python/__init__.py

@@ -0,0 +1,361 @@
+from operator import (
+    abs as abs,
+)
+from operator import (
+    add as add,
+)
+from operator import (
+    and_ as bitwise_and,
+)
+from operator import (
+    eq as equal,
+)
+from operator import (
+    floordiv as floor_divide,
+)
+from operator import (
+    ge as greater_equal,
+)
+from operator import (
+    gt as greater,
+)
+from operator import (
+    invert as bitwise_invert,
+)
+from operator import (
+    le as less_equal,
+)
+from operator import (
+    lshift as bitwise_left_shift,
+)
+from operator import (
+    lt as less,
+)
+from operator import (
+    matmul as matmul,
+)
+from operator import (
+    mod as remainder,
+)
+from operator import (
+    mul as multiply,
+)
+from operator import (
+    ne as not_equal,
+)
+from operator import (
+    neg as negative,
+)
+from operator import (
+    or_ as bitwise_or,
+)
+from operator import (
+    pos as positive,
+)
+from operator import (
+    pow as pow,
+)
+from operator import (
+    rshift as bitwise_right_shift,
+)
+from operator import (
+    sub as subtract,
+)
+from operator import (
+    truediv as divide,
+)
+from operator import (
+    xor as bitwise_xor,
+)
+
+from numpy import (
+    e as e,
+)
+from numpy import (
+    inf as inf,
+)
+from numpy import (
+    nan as nan,
+)
+from numpy import (
+    newaxis as newaxis,
+)
+from numpy import (
+    pi as pi,
+)
+
+from . import linalg
+from ._array_api_info import __array_namespace_info__
+from .compiled import (
+    DefaultScheduler,
+    GalleyScheduler,
+    compiled,
+    compute,
+    lazy,
+    set_optimizer,
+)
+from .dtypes import (
+    bool,
+    can_cast,
+    complex64,
+    complex128,
+    finfo,
+    float16,
+    float32,
+    float64,
+    iinfo,
+    int8,
+    int16,
+    int32,
+    int64,
+    int_,
+    uint,
+    uint8,
+    uint16,
+    uint32,
+    uint64,
+)
+from .finch_fused import jit
+from .io import (
+    read,
+    write,
+)
+from .levels import (
+    Dense,
+    DenseStorage,
+    Element,
+    Pattern,
+    RepeatRLE,
+    SparseByteMap,
+    SparseCOO,
+    SparseHash,
+    SparseList,
+    SparseVBL,
+    Storage,
+)
+from .tensor import (
+    SparseArray,
+    Tensor,
+    acos,
+    acosh,
+    all,
+    any,
+    arange,
+    argmax,
+    argmin,
+    asarray,
+    asin,
+    asinh,
+    astype,
+    atan,
+    atan2,
+    atanh,
+    ceil,
+    conj,
+    cos,
+    cosh,
+    diagonal,
+    einop,
+    einsum,
+    empty,
+    empty_like,
+    exp,
+    expand_dims,
+    expm1,
+    eye,
+    floor,
+    full,
+    full_like,
+    imag,
+    isfinite,
+    isinf,
+    isnan,
+    linspace,
+    log,
+    log1p,
+    log2,
+    log10,
+    logaddexp,
+    logical_and,
+    logical_or,
+    logical_xor,
+    max,
+    mean,
+    min,
+    moveaxis,
+    nonzero,
+    ones,
+    ones_like,
+    permute_dims,
+    power,
+    prod,
+    random,
+    real,
+    reshape,
+    round,
+    sign,
+    sin,
+    sinh,
+    sqrt,
+    square,
+    squeeze,
+    std,
+    sum,
+    tan,
+    tanh,
+    tensordot,
+    trunc,
+    var,
+    where,
+    zeros,
+    zeros_like,
+)
+
+__all__ = [
+    "DefaultScheduler",
+    "Dense",
+    "DenseStorage",
+    "Element",
+    "GalleyScheduler",
+    "Pattern",
+    "RepeatRLE",
+    "SparseArray",
+    "SparseByteMap",
+    "SparseCOO",
+    "SparseHash",
+    "SparseList",
+    "SparseVBL",
+    "Storage",
+    "Tensor",
+    "__array_namespace_info__",
+    "abs",
+    "acos",
+    "acosh",
+    "add",
+    "all",
+    "any",
+    "arange",
+    "argmax",
+    "argmin",
+    "asarray",
+    "asin",
+    "asinh",
+    "astype",
+    "atan",
+    "atan2",
+    "atanh",
+    "bitwise_and",
+    "bitwise_invert",
+    "bitwise_left_shift",
+    "bitwise_or",
+    "bitwise_right_shift",
+    "bitwise_xor",
+    "bool",
+    "can_cast",
+    "ceil",
+    "compiled",
+    "complex64",
+    "complex128",
+    "compute",
+    "conj",
+    "cos",
+    "cosh",
+    "diagonal",
+    "divide",
+    "e",
+    "einop",
+    "einsum",
+    "empty",
+    "empty_like",
+    "equal",
+    "exp",
+    "expand_dims",
+    "expm1",
+    "eye",
+    "finfo",
+    "float16",
+    "float32",
+    "float64",
+    "floor",
+    "floor_divide",
+    "full",
+    "full_like",
+    "greater",
+    "greater_equal",
+    "iinfo",
+    "imag",
+    "inf",
+    "int8",
+    "int16",
+    "int32",
+    "int64",
+    "int_",
+    "isfinite",
+    "isinf",
+    "isnan",
+    "jit",
+    "lazy",
+    "less",
+    "less_equal",
+    "linalg",
+    "linspace",
+    "log",
+    "log1p",
+    "log2",
+    "log10",
+    "logaddexp",
+    "logical_and",
+    "logical_or",
+    "logical_xor",
+    "matmul",
+    "max",
+    "mean",
+    "min",
+    "moveaxis",
+    "multiply",
+    "nan",
+    "negative",
+    "newaxis",
+    "nonzero",
+    "not_equal",
+    "ones",
+    "ones_like",
+    "permute_dims",
+    "pi",
+    "positive",
+    "pow",
+    "power",
+    "prod",
+    "random",
+    "read",
+    "real",
+    "remainder",
+    "reshape",
+    "round",
+    "set_optimizer",
+    "sign",
+    "sin",
+    "sinh",
+    "sqrt",
+    "square",
+    "squeeze",
+    "std",
+    "subtract",
+    "sum",
+    "tan",
+    "tanh",
+    "tensordot",
+    "trunc",
+    "uint",
+    "uint8",
+    "uint16",
+    "uint32",
+    "uint64",
+    "var",
+    "where",
+    "write",
+    "zeros",
+    "zeros_like",
+]
+
+__array_api_version__: str = "2024.12"

galley_jl_python/_array_api_info.py

@@ -0,0 +1,94 @@
+from . import dtypes
+from .typing import DType
+
+
+class __array_namespace_info__:
+    def capabilities(self) -> dict[str, bool]:
+        return {
+            "boolean indexing": True,
+            "data-dependent shapes": True,
+        }
+
+    def default_device(self) -> str:
+        return "cpu"
+
+    def default_dtypes(self, *, device: str | None = None) -> dict[str, DType]:
+        if device not in ["cpu", None]:
+            raise ValueError(
+                f'Device not understood. Only "cpu" is allowed, but received: {device}'
+            )
+        return {
+            "real floating": dtypes.float64,
+            "complex floating": dtypes.complex128,
+            "integral": dtypes.int_,
+            "indexing": dtypes.int_,
+        }
+
+    _bool_dtypes = {"bool": dtypes.bool}
+    _signed_integer_dtypes = {
+        "int8": dtypes.int8,
+        "int16": dtypes.int16,
+        "int32": dtypes.int32,
+        "int64": dtypes.int64,
+    }
+    _unsigned_integer_dtypes = {
+        "uint8": dtypes.uint8,
+        "uint16": dtypes.uint16,
+        "uint32": dtypes.uint32,
+        "uint64": dtypes.uint64,
+    }
+    _real_floating_dtypes = {
+        "float32": dtypes.float32,
+        "float64": dtypes.float64,
+    }
+    _complex_floating_dtypes = {
+        "complex64": dtypes.complex64,
+        "complex128": dtypes.complex128,
+    }
+
+    def dtypes(
+        self,
+        *,
+        device: str | None = None,
+        kind: str | tuple[str, ...] | None = None,
+    ) -> dict[str, DType]:
+        if device not in ["cpu", None]:
+            raise ValueError(
+                f'Device not understood. Only "cpu" is allowed, but received: {device}'
+            )
+        if kind is None:
+            return (
+                self._bool_dtypes
+                | self._signed_integer_dtypes
+                | self._unsigned_integer_dtypes
+                | self._real_floating_dtypes
+                | self._complex_floating_dtypes
+            )
+        if kind == "bool":
+            return self._bool_dtypes
+        if kind == "signed integer":
+            return self._signed_integer_dtypes
+        if kind == "unsigned integer":
+            return self._unsigned_integer_dtypes
+        if kind == "integral":
+            return self._signed_integer_dtypes | self._unsigned_integer_dtypes
+        if kind == "real floating":
+            return self._real_floating_dtypes
+        if kind == "complex floating":
+            return self._complex_floating_dtypes
+        if kind == "numeric":
+            return (
+                self._signed_integer_dtypes
+                | self._unsigned_integer_dtypes
+                | self._real_floating_dtypes
+                | self._complex_floating_dtypes
+            )
+        if isinstance(kind, tuple):
+            res = {}
+            for k in kind:
+                res.update(self.dtypes(kind=k))
+            return res
+        raise ValueError(f"unsupported kind: {kind!r}")
+
+    def devices(self) -> list[str]:
+        return ["cpu"]

galley_jl_python/_sysimage.py

@@ -0,0 +1,222 @@
+"""Locate, and if needed download, a prebuilt Julia sysimage for juliacall.
+
+This module must not import juliacall: it runs before Julia starts.
+
+A sysimage only works with the exact Julia version and package versions it was
+built from, so each image is named after a hash of both (see ``env_key``) and
+lives in a per-user cache directory. An image is used only when its name
+matches the current environment; otherwise Julia starts normally.
+
+On import, galley downloads the image CI published for this platform and
+environment if it isn't cached yet, and uses it right away. Set
+``GALLEY_JL_PYTHON_SYSIMAGE=0`` to never download or load one, or
+``PYTHON_JULIACALL_SYSIMAGE`` to use a specific image. In a galley checkout,
+``pixi run build-sysimage`` builds one locally.
+"""
+
+import hashlib
+import json
+import os
+import platform
+import subprocess
+import sys
+import tempfile
+import time
+import tomllib
+import urllib.error
+import urllib.request
+from pathlib import Path
+
+RELEASES_URL = "https://github.com/finch-tensor/galley-jl-python/releases/download"
+# Bump when images built from the same environment become incompatible, so the
+# new images get new names instead of reusing the published ones.
+# 2: PythonCall finds juliacall from Python (v1 only worked on the build host)
+IMAGE_FORMAT = 2
+_EXTENSIONS = {"linux": "so", "darwin": "dylib", "win32": "dll"}
+# after a 404, don't ask again for this long (images are published after CI)
+_MISSING_TTL = 24 * 60 * 60
+# a lock older than this is left over from a download that was killed
+_STALE_LOCK = 2 * 60 * 60
+
+
+def cache_dir() -> Path:
+    if custom := os.environ.get("GALLEY_JL_PYTHON_CACHE"):
+        return Path(custom)
+    if sys.platform == "win32":
+        base = Path(os.environ.get("LOCALAPPDATA", Path.home() / "AppData" / "Local"))
+    elif sys.platform == "darwin":
+        base = Path.home() / "Library" / "Caches"
+    else:
+        base = Path(os.environ.get("XDG_CACHE_HOME", Path.home() / ".cache"))
+    return base / "galley-jl-python"
+
+
+def env_key(project: str) -> str:
+    """Hash the Julia version and resolved package set of a juliapkg project."""
+    manifest = tomllib.loads((Path(project) / "Manifest.toml").read_text())
+    resolved = sorted(
+        (name, entry.get("uuid"), entry.get("version"), entry.get("git-tree-sha1"))
+        for name, entries in manifest["deps"].items()
+        for entry in entries
+    )
+    payload = json.dumps([manifest["julia_version"], resolved])
+    return hashlib.sha256(payload.encode()).hexdigest()[:16]
+
+
+def release_tag(key: str) -> str:
+    return f"sysimage-{key}"
+
+
+def image_name(key: str) -> str:
+    ext = _EXTENSIONS.get(sys.platform, "so")
+    arch = platform.machine().lower()
+    return f"galley-v{IMAGE_FORMAT}-{sys.platform}-{arch}-{key}.{ext}"
+
+
+def image_path(key: str) -> Path:
+    return cache_dir() / image_name(key)
+
+
+def current_key() -> str:
+    import juliapkg
+
+    return env_key(juliapkg.project())
+
+
+def _log(message: str) -> None:
+    print(f"[galley-jl-python] {message}", file=sys.stderr, flush=True)
+
+
+def _acquire_lock(lock: Path) -> bool:
+    for _ in range(2):
+        try:
+            os.close(os.open(lock, os.O_CREAT | os.O_EXCL | os.O_WRONLY))
+            return True
+        except FileExistsError:
+            try:
+                if time.time() - lock.stat().st_mtime < _STALE_LOCK:
+                    return False
+                lock.unlink()
+            except FileNotFoundError:
+                pass
+    return False
+
+
+def _download(url: str, dest: Path, attempts: int = 4) -> None:
+    """Download ``url`` to ``dest``, resuming after a dropped connection."""
+    size = done = 0
+    next_report = 0.25
+    for attempt in range(attempts):
+        headers = {"Range": f"bytes={done}-"} if done else {}
+        try:
+            req = urllib.request.Request(url, headers=headers)
+            with (
+                urllib.request.urlopen(req, timeout=60) as resp,
+                dest.open("ab") as out,
+            ):
+                if done and resp.status != 206:  # server ignored the range
+                    out.truncate(0)
+                    done = 0
+                if not size:
+                    size = done + int(resp.headers.get("Content-Length") or 0)
+                    _log(
+                        f"downloading the prebuilt Julia sysimage ({size >> 20} MB), "
+                        "once per environment; set GALLEY_JL_PYTHON_SYSIMAGE=0 to skip"
+                    )
+                while chunk := resp.read(1 << 20):
+                    out.write(chunk)
+                    done += len(chunk)
+                    if size and done / size >= next_report:
+                        _log(f"  {done * 100 // size}%")
+                        next_report += 0.25
+            if not size or done >= size:
+                return
+        except urllib.error.HTTPError:
+            raise
+        except (urllib.error.URLError, OSError):
+            if attempt == attempts - 1:
+                raise
+        _log(
+            f"  connection dropped at {done >> 20} MB; resuming"
+            if done
+            else "  couldn't connect; retrying"
+        )
+
+
+def _works(image: Path) -> bool:
+    """Check that juliacall starts with ``image``, in a throwaway process."""
+    env = {**os.environ, "PYTHON_JULIACALL_SYSIMAGE": str(image)}
+    env["GALLEY_JL_PYTHON_SYSIMAGE"] = "0"
+    try:
+        result = subprocess.run(
+            [sys.executable, "-c", "from juliacall import Main; Main.seval('1 + 1')"],
+            env=env,
+            capture_output=True,
+            timeout=600,
+        )
+    except (OSError, subprocess.TimeoutExpired):
+        return False
+    return result.returncode == 0
+
+
+def fetch(key: str) -> Path | None:
+    """Download the published image for ``key`` into the cache.
+
+    Returns the image path, or None if there is no image or it couldn't be
+    downloaded. Never raises: galley works without an image, just with a slow
+    first call. If another process is already downloading, returns None
+    immediately instead of waiting.
+    """
+    image = image_path(key)
+    if image.is_file():
+        return image
+    missing = image.with_name(image.name + ".missing")
+    lock = image.with_name(image.name + ".lock")
+    try:
+        if missing.is_file() and time.time() - missing.stat().st_mtime < _MISSING_TTL:
+            return None
+        image.parent.mkdir(parents=True, exist_ok=True)
+        if not _acquire_lock(lock):
+            return None
+    except OSError:
+        return None
+
+    url = f"{RELEASES_URL}/{release_tag(key)}/{image.name}"
+    partial = None
+    try:
+        fd, partial = tempfile.mkstemp(dir=image.parent, suffix=".part")
+        os.close(fd)
+        _download(url, Path(partial))
+        if not _works(Path(partial)):
+            _log("the downloaded sysimage doesn't load here; starting without it")
+            missing.touch()
+            return None
+        os.chmod(partial, 0o644)
+        os.replace(partial, image)
+        partial = None
+        _log(f"saved {image}")
+        return image
+    except urllib.error.HTTPError as e:
+        if e.code == 404:
+            missing.touch()
+        else:
+            _log(f"sysimage download failed ({e}); starting without it")
+        return None
+    except (urllib.error.URLError, OSError) as e:
+        _log(f"sysimage download failed ({e}); starting without it")
+        return None
+    finally:
+        if partial is not None and os.path.exists(partial):
+            os.remove(partial)
+        lock.unlink(missing_ok=True)
+
+
+def configure() -> None:
+    """Point juliacall at this environment's sysimage, downloading it if needed."""
+    if os.environ.get("GALLEY_JL_PYTHON_SYSIMAGE") == "0":
+        return
+    if "PYTHON_JULIACALL_SYSIMAGE" in os.environ:
+        return
+    image = fetch(current_key())
+    if image is not None:
+        os.environ["PYTHON_JULIACALL_SYSIMAGE"] = str(image)