splime 0.1.2__py3-none-any.whl

spl/core/__init__.py

@@ -0,0 +1,23 @@
+import spl.core.entities.adapter
+import spl.core.entities.artifact
+import spl.core.entities.distribution
+import spl.core.entities.function
+# Registered before `module` so local user functions are inlined instead of
+# being captured as a bare `from local_module import ...` (see local_function).
+import spl.core.entities.local_function
+import spl.core.entities.misc
+import spl.core.entities.module
+import spl.core.entities.node
+import spl.core.entities.node_function
+import spl.core.entities.node_remote
+import spl.core.entities.pipeline
+import spl.core.entities.scalar  # noqa: F401
+from spl.core.entities.node_remote import NodeRemote
+from spl.core.ir.utils import spl_export_to_dir, spl_export_to_file, spl_import_from_file
+
+__all__ = [
+    'NodeRemote',
+    'spl_export_to_dir',
+    'spl_export_to_file',
+    'spl_import_from_file'
+]

spl/core/common.py

@@ -0,0 +1,350 @@
+import os
+import shutil
+import tempfile
+import weakref
+from collections.abc import Callable
+from dataclasses import dataclass, replace
+from functools import reduce
+from itertools import groupby
+from operator import itemgetter
+from pathlib import Path
+from types import FunctionType
+from typing import Any
+
+from spl.core.entities.adapter import Adapter
+from spl.core.entities.artifact import ArtifactRef, compute_sha256
+from spl.core.entities.node import (
+    DEFAULT_PORT,
+    FormattedOutputRef,
+    Node,
+    NodeInputRef,
+    NodeOutputRef,
+)
+from spl.core.entities.node_function import NodeFunction
+from spl.core.entities.node_remote import NodeRemote
+from spl.core.entities.pipeline import Pipeline
+from spl.core.entities.scalar import Scalar
+
+_JSON_NATIVE_TYPES = {str, int, float, bool, dict, list}
+
+
+@dataclass(frozen = True)
+class PipelineBuilder:
+    pipeline: Pipeline
+    root: Node
+    format: str | None = None
+
+    @staticmethod
+    def lift(x):
+
+        match x:
+            case PipelineBuilder():
+                return x
+
+            case FunctionType():
+                root = NodeFunction(x)
+                return PipelineBuilder(
+                    Pipeline(nodes = {root}, links = set()),
+                    root)
+
+            case NodeFunction():
+                root = x
+                return PipelineBuilder(
+                    Pipeline(nodes = {root}, links = set()),
+                    root)
+
+            case NodeRemote():
+                root = x
+                return PipelineBuilder(
+                    Pipeline(nodes = {root}, links = set()),
+                    root)
+
+            case _:
+                raise ValueError(x)
+
+    def get_input_node_refs(self, port_name: str, is_free: bool):
+        node_refs = [
+            NodeInputRef(node, port)
+            for node in self.pipeline.nodes
+            for port in node.inputs
+            if port.name == port_name]
+
+        if is_free:
+            bound_refs = set(map(itemgetter(0), self.pipeline.links))
+            node_refs = [x for x in node_refs if x not in bound_refs]
+
+        return node_refs
+
+
+    def bind(self, **kwargs):
+        return self._bind(kwargs, is_strict = True, is_free = False)
+
+    def bind_all(self, **kwargs):
+        return self._bind(kwargs, is_strict = False, is_free = True)
+
+    def _bind(self, kwargs, is_strict: bool, is_free: bool):
+        pipeline = self.pipeline
+        for port_name, v in kwargs.items():
+            match self.get_input_node_refs(port_name, is_free):
+                case []:
+                    raise ValueError('node(s) for port `{}` is not found'.format(port_name))
+
+                case [ref]:
+                    pipeline = self._update_pipeline(pipeline, ref, v)
+
+                case refs:
+                    if is_strict:
+                        raise ValueError('ambigious node for port `{}`'.format(port_name))
+
+                    pipeline = reduce(
+                        lambda acc, ref: self._update_pipeline(acc, ref, v),
+                        refs,
+                        pipeline)
+        return PipelineBuilder(
+            pipeline = pipeline,
+            root = self.root,
+            format = self.format)
+
+    def alias(self, name):
+        return replace(self, pipeline = self.pipeline.add_alias(self.root, name))
+
+    def as_format(self, format: str) -> 'PipelineBuilder':
+        """Return a builder whose output edge uses an artifact format."""
+
+        if not isinstance(format, str):
+            raise TypeError('pipeline builder format must be a string')
+        if not format:
+            raise ValueError('pipeline builder format must be a non-empty string')
+        return replace(self, format = format)
+
+    @staticmethod
+    def _update_pipeline(pipeline, ref, v):
+        match v:
+            case PipelineBuilder():
+                output_ref = NodeOutputRef(
+                    v.root,
+                    v.root.get_output_port(DEFAULT_PORT))
+                link_value = (
+                    output_ref
+                    if v.format is None
+                    else FormattedOutputRef(output_ref, v.format))
+                return (pipeline | v.pipeline).add_link(
+                    ref,
+                    link_value)
+            case _:
+                return pipeline.add_link(
+                    ref,
+                    Scalar(v))
+
+    def render(self, name: str | None = None):
+        return replace(self.pipeline, name = name)
+
+
+lift = PipelineBuilder.lift
+
+
+def encode(value: Any, adapter: Adapter, artifacts_dir: Path) -> ArtifactRef:
+    """Materialize a value with an adapter and return its artifact reference."""
+
+    fd, artifact_path_value = tempfile.mkstemp(
+        prefix = 'artifact-',
+        dir = artifacts_dir)
+    os.close(fd)
+    artifact_path = Path(artifact_path_value)
+
+    try:
+        adapter.save(str(artifact_path), value)
+    except BaseException:
+        artifact_path.unlink(missing_ok = True)
+        raise
+
+    size = artifact_path.stat().st_size
+    sha256 = compute_sha256(artifact_path)
+    return ArtifactRef(
+        key = adapter.key,
+        uri = str(artifact_path),
+        sha256 = sha256,
+        size = size)
+
+
+def decode(ref: ArtifactRef, adapter: Adapter) -> Any:
+    """Load an artifact reference with an adapter after validating its digest."""
+
+    if ref.key != adapter.key:
+        raise ValueError('artifact ref key does not match adapter')
+
+    artifact_path = Path(ref.uri)
+    if artifact_path.stat().st_size != ref.size:
+        raise ValueError('artifact ref size does not match file')
+    if compute_sha256(artifact_path) != ref.sha256:
+        raise ValueError('artifact ref sha256 does not match file')
+    return adapter.load(str(artifact_path))
+
+
+class Deployment:
+    def __init__(self, client=None, pipeline=None):
+        if pipeline is None:
+            pipeline = client
+            client = None
+        self._client = client
+        self._pipeline = pipeline
+
+    def setup(self):
+        pass
+
+    def teardown(self):
+        pass
+
+    def run(self, **kwargs):
+        return Run(self._callback, self._pipeline, **kwargs)
+
+    def _callback(self, node, kwargs):
+        final_kwargs = {port.name: v for port, v in kwargs.items()}
+        output_port = self._single_output_port(node)
+        match node:
+            case NodeFunction():
+                return {output_port.name: node.func(**final_kwargs)}
+
+            case NodeRemote():
+                if self._client is None:
+                    raise RuntimeError('remote node execution requires a client')
+                return {output_port.name: self._client.run_node(node, final_kwargs)}
+
+            case _:
+                raise ValueError(node)
+
+    @staticmethod
+    def _single_output_port(node):
+        outputs = node.outputs or []
+        if len(outputs) != 1:
+            raise RuntimeError(
+                'node {} has {} outputs; local Deployment currently supports '
+                'exactly one output and requires an explicit daemon/server '
+                'output selector for multi-output pipelines'.format(
+                    node,
+                    len(outputs)))
+        return outputs[0]
+
+
+class Run:
+    def __init__(
+            self,
+            callback: Callable[..., dict[str, Any]],
+            pipeline: Pipeline,
+            **kwargs: Any) -> None:
+        self._callback = callback
+        self._pipeline = pipeline
+        self._kwargs = kwargs
+        self._deps: dict[Node, dict[Any, Any]] = {
+            k: dict(map(itemgetter(slice(1, None)), vs))
+            for k, vs in groupby(
+                sorted(
+                    [(x.node, x.port, y) for (x, y) in pipeline.links],
+                    key = lambda x: hash(x[0])),
+                itemgetter(0))}
+        self._results: dict[Node, dict[str, Any]] = dict()
+        self._artifact_refs: dict[tuple[Node, str, str], ArtifactRef] = dict()
+        self._artifacts_dir: Path | None = None
+        self._artifacts_finalizer: Any = None
+        self._closed = False
+
+    def __enter__(self) -> 'Run':
+        self._ensure_open()
+        return self
+
+    def __exit__(self, exc_type: Any, exc: Any, traceback: Any) -> None:
+        self.close()
+
+    def close(self) -> None:
+        if not self._closed:
+            if self._artifacts_finalizer is not None:
+                self._artifacts_finalizer()
+            self._closed = True
+
+    def _ensure_open(self) -> None:
+        if self._closed:
+            raise RuntimeError('run is closed')
+
+    def _get_artifacts_dir(self) -> Path:
+        self._ensure_open()
+        if self._artifacts_dir is None:
+            artifacts_dir = Path(tempfile.mkdtemp(prefix = 'spl-run-'))
+            self._artifacts_dir = artifacts_dir
+            self._artifacts_finalizer = weakref.finalize(
+                self,
+                shutil.rmtree,
+                artifacts_dir,
+                ignore_errors = True)
+        return self._artifacts_dir
+
+    def _round_trip_artifact(
+            self,
+            value: Any,
+            source_ref: NodeOutputRef | None = None,
+            adapter_format: str | None = None) -> Any:
+        if adapter_format is None and type(value) in _JSON_NATIVE_TYPES:
+            return value
+
+        adapter = self._pipeline.resolve_adapter(
+            py_type = type(value),
+            format = adapter_format)
+        if adapter is None:
+            if adapter_format is not None:
+                raise ValueError(
+                    'pipeline adapter is not found for python type ({}) '
+                    'and format `{}`'.format(type(value), adapter_format))
+            return value
+
+        self._ensure_open()
+        if source_ref is None:
+            ref = encode(value, adapter, self._get_artifacts_dir())
+        else:
+            cache_key = (source_ref.node, source_ref.port.name, adapter.key)
+            if cache_key not in self._artifact_refs:
+                self._artifact_refs[cache_key] = encode(
+                    value,
+                    adapter,
+                    self._get_artifacts_dir())
+            ref = self._artifact_refs[cache_key]
+        return decode(ref, adapter)
+
+    def _get_input(self, x: Any) -> Any:
+        match x:
+            case Scalar():
+                return self._round_trip_artifact(x.value)
+
+            case NodeOutputRef():
+                return self._round_trip_artifact(
+                    (self._get_result(x.node))[x.port.name],
+                    x)
+
+            case FormattedOutputRef():
+                return self._round_trip_artifact(
+                    (self._get_result(x.out_ref.node))[x.out_ref.port.name],
+                    source_ref = x.out_ref,
+                    adapter_format = x.format)
+
+            case _: raise ValueError(x)
+
+    def _get_result(self, node: Node) -> dict[str, Any]:
+        if node not in self._results:
+            self._ensure_open()
+            kwargs = {
+                port: self._round_trip_artifact(self._kwargs[port.name])
+                for port in node.inputs
+                if port.name in self._kwargs}
+
+            if node in self._deps:
+                kwargs = kwargs | {
+                    port: self._get_input(v)
+                    for port, v in self._deps[node].items()}
+
+            self._results[node] = self._callback(node, kwargs)
+        return self._results[node]
+
+    def __getitem__(self, node: Node) -> dict[str, Any]:
+        try:
+            return self._get_result(node)
+        except BaseException:
+            self.close()
+            raise

spl/core/entities/adapter.py

@@ -0,0 +1,210 @@
+import ast
+from collections.abc import Callable
+from dataclasses import dataclass
+from pathlib import Path
+from types import FunctionType
+from typing import Any, Generator, cast
+
+import yaml
+
+from spl.core.entities.distribution import DDistribution
+from spl.core.entities.function import get_function_metadata
+from spl.core.ir.common import DBase
+from spl.core.ir.parse import _branch, ir_parse
+from spl.core.ir.unparse import ir_unparse
+
+
+def _validate_non_empty_string(name: str, value: str) -> None:
+    if not isinstance(value, str):
+        raise TypeError('adapter {} must be a string'.format(name))
+    if not value:
+        raise ValueError('adapter {} must be a non-empty string'.format(name))
+
+
+def _format_from_key(key: str) -> str:
+    _validate_non_empty_string('key', key)
+    key_head, separator, key_format = key.rpartition('@')
+    if not key_head or not separator or not key_format:
+        raise ValueError('adapter key must be `<python_type>@<format>`')
+    return key_format
+
+
+def make_key(py_type: type[Any], format: str) -> str:
+    """Return the stable adapter key for a Python type and storage format."""
+
+    if not isinstance(py_type, type):
+        raise TypeError('adapter python type must be a type')
+    _validate_non_empty_string('format', format)
+    return '{}.{}@{}'.format(py_type.__module__, py_type.__qualname__, format)
+
+
+def _validate_function(name: str, value: Callable[..., Any]) -> None:
+    if not isinstance(value, FunctionType):
+        raise TypeError('adapter {} must be a function'.format(name))
+
+
+def _normalize_distributions(value: Any) -> tuple[DDistribution, ...]:
+    if not isinstance(value, tuple | list):
+        raise TypeError('adapter distributions must be a tuple')
+
+    distributions = tuple(value)
+    if any(not isinstance(x, DDistribution) for x in distributions):
+        raise TypeError('adapter distributions must contain DDistribution values')
+    return tuple(sorted(distributions))
+
+
+def _validate_py_type(value: type[Any] | None) -> None:
+    if value is not None and not isinstance(value, type):
+        raise TypeError('adapter py_type must be a type or None')
+
+
+def _validate_key_format(key: str, py_type: type[Any] | None, format: str) -> None:
+    _validate_non_empty_string('format', format)
+    if _format_from_key(key) != format:
+        raise ValueError('adapter key format does not match format')
+    if py_type is not None and key != make_key(py_type, format):
+        raise ValueError('adapter key does not match python type and format')
+
+
+def _function_name(func: Callable[..., Any]) -> str:
+    metadata = get_function_metadata(cast(FunctionType, func))
+    return cast(str, metadata.name)
+
+
+@dataclass(frozen = True)
+class Adapter:
+    """Versioned save/load pair for materializing values as artifacts."""
+
+    key: str
+    save: Callable[..., Any]
+    load: Callable[..., Any]
+    py_type: type[Any] | None
+    format: str
+    distributions: tuple[DDistribution, ...] = ()
+
+    def __post_init__(self) -> None:
+        _validate_function('save', self.save)
+        _validate_function('load', self.load)
+        _validate_py_type(self.py_type)
+        _validate_key_format(
+            key = self.key,
+            py_type = self.py_type,
+            format = self.format)
+        object.__setattr__(
+            self,
+            'distributions',
+            _normalize_distributions(self.distributions))
+
+    def __hash__(self) -> int:
+        return hash((
+            self.key,
+            _function_name(self.save),
+            _function_name(self.load),
+            self.distributions))
+
+
+@dataclass(frozen = True)
+class DAdapter(DBase):
+    """Serialized Adapter value for pipeline YAML."""
+
+    key: str
+    save: str
+    load: str
+    distributions: tuple[DDistribution, ...] = ()
+
+    def __post_init__(self) -> None:
+        _format_from_key(self.key)
+        _validate_non_empty_string('save', self.save)
+        _validate_non_empty_string('load', self.load)
+        object.__setattr__(
+            self,
+            'distributions',
+            _normalize_distributions(self.distributions))
+
+    def __hash__(self) -> int:
+        return hash((
+            self.key,
+            self.save,
+            self.load,
+            self.distributions))
+
+
+yaml.add_representer(
+    DAdapter,
+    lambda dumper, data: dumper.represent_mapping('!DAdapter', {
+        'key': data.key,
+        'save': data.save,
+        'load': data.load,
+        'distributions': list(data.distributions)}))
+
+
+def _construct_dadapter(loader: Any, node: Any) -> DAdapter:
+    return DAdapter(**loader.construct_mapping(node, deep = True))
+
+
+yaml.add_constructor(
+    '!DAdapter',
+    _construct_dadapter)
+
+
+@ir_parse.register(
+    lambda x: isinstance(x, Adapter))
+def _ir_parse__adapter(
+        x: Adapter,
+        name: str | None = None) -> _branch:
+    def mk_dependencies(frame_offset: int) -> Generator[Any]:
+        yield ir_parse(x.save)
+        yield ir_parse(x.load)
+
+    return _branch(
+        x,
+        lambda: DAdapter(
+            key = x.key,
+            save = _function_name(x.save),
+            load = _function_name(x.load),
+            distributions = x.distributions),
+        mk_dependencies)
+
+
+def _unparse_distributions(distributions: tuple[DDistribution, ...]) -> ast.Tuple:
+    return ast.Tuple(
+        elts = [
+            ast.Call(
+                func = ast.Name(id = 'DDistribution', ctx = ast.Load()),
+                keywords = [
+                    ast.keyword(
+                        arg = 'package',
+                        value = ast.Constant(value = x.package)),
+                    ast.keyword(
+                        arg = 'version',
+                        value = ast.Constant(value = x.version))])
+            for x in distributions],
+        ctx = ast.Load())
+
+
+@ir_unparse.register(
+    lambda x: isinstance(x, DAdapter))
+def _ir_unparse__adapter(x: DAdapter, source: Path) -> Generator[ast.stmt]:
+    yield ast.Assign(
+        targets = [ast.Name(id = '_adapter', ctx = ast.Store())],
+        value = ast.Call(
+            func = ast.Name(id = 'Adapter', ctx = ast.Load()),
+            keywords = [
+                ast.keyword(
+                    arg = 'key',
+                    value = ast.Constant(value = x.key)),
+                ast.keyword(
+                    arg = 'save',
+                    value = ast.Name(id = x.save, ctx = ast.Load())),
+                ast.keyword(
+                    arg = 'load',
+                    value = ast.Name(id = x.load, ctx = ast.Load())),
+                ast.keyword(
+                    arg = 'py_type',
+                    value = ast.Constant(value = None)),
+                ast.keyword(
+                    arg = 'format',
+                    value = ast.Constant(value = _format_from_key(x.key))),
+                ast.keyword(
+                    arg = 'distributions',
+                    value = _unparse_distributions(x.distributions))]))