tensorlake 0.4.31__py3-none-win_amd64.whl

tensorlake/__init__.py

@@ -0,0 +1 @@
+from tensorlake.image import Image

tensorlake/applications/__init__.py

@@ -0,0 +1 @@
+from .interface import *

tensorlake/applications/algorithms/__init__.py

@@ -0,0 +1,6 @@
+from .dfs_bottom_up import dfs_bottom_up, dfs_bottom_up_unique_only
+
+__all__ = [
+    "dfs_bottom_up_unique_only",
+    "dfs_bottom_up",
+]

tensorlake/applications/algorithms/dfs_bottom_up.py

@@ -0,0 +1,94 @@
+from typing import Any, Generator
+
+from ..interface.exceptions import InternalError
+from ..interface.futures import (
+    FunctionCallFuture,
+    Future,
+    MapFuture,
+    ReduceOperationFuture,
+    _InitialMissingType,
+    _TensorlakeFutureWrapper,
+    _unwrap_future,
+)
+
+
+def dfs_bottom_up_unique_only(root: Future) -> Generator[Future, None, None]:
+    """Yields all unique Futures starting from leafs up to the root aka Post-Order DFS.
+
+    This traversal order is useful when at the moment of processing a node, all its children have
+    to be already processed.
+
+    The traversal order is deterministic (always the same for the same tree).
+    Doesn't yield the same Future referenced in the tree multiple times.
+    Doesn't yield user supplied values.
+
+    Raises InternalError if encounters an unexpected Future type.
+    """
+    seen_future_ids: set[str] = set()
+    for future in dfs_bottom_up(root):
+        if future._id not in seen_future_ids:
+            seen_future_ids.add(future._id)
+            yield future
+
+
+def dfs_bottom_up(root: Future) -> Generator[Future, None, None]:
+    """Yields all Futures starting from leafs up to the root aka Post-Order DFS.
+
+    This traversal order is useful when at the moment of processing a node, all its children have
+    to be already processed.
+
+    The traversal order is deterministic (always the same for the same tree).
+    Doesn't yield user supplied values.
+
+    Raises InternalError if encounters an unexpected Future type.
+    """
+    dfs_stack: list[Future] = [root]
+    yield_stack: list[Future] = []
+
+    while len(dfs_stack) > 0:
+        node: Future = dfs_stack.pop()
+        yield_stack.append(node)
+
+        if isinstance(node, MapFuture):
+            node: MapFuture
+            items: list[_TensorlakeFutureWrapper[Future] | Any] | Future = (
+                _unwrap_future(node._items)
+            )
+            if isinstance(items, Future):
+                dfs_stack.append(items)
+            else:
+                for item in node._items:
+                    item: Future | Any = _unwrap_future(item)
+                    if isinstance(item, Future):
+                        dfs_stack.append(item)
+        elif isinstance(node, ReduceOperationFuture):
+            node: ReduceOperationFuture
+            initial: Future | Any | _InitialMissingType = _unwrap_future(node._initial)
+            if isinstance(initial, Future):
+                dfs_stack.append(initial)
+
+            items: list[_TensorlakeFutureWrapper[Future] | Any] | Future = (
+                _unwrap_future(node._items)
+            )
+            if isinstance(items, Future):
+                dfs_stack.append(items)
+            else:
+                for item in node._items:
+                    item: Future | Any = _unwrap_future(item)
+                    if isinstance(item, Future):
+                        dfs_stack.append(item)
+        elif isinstance(node, FunctionCallFuture):
+            node: FunctionCallFuture
+            for arg in node._args:
+                arg: Future | Any = _unwrap_future(arg)
+                if isinstance(arg, Future):
+                    dfs_stack.append(arg)
+            # Sort dict keys to ensure deterministic traversal order.
+            for key in sorted(node._kwargs.keys()):
+                arg: Future | Any = _unwrap_future(node._kwargs[key])
+                if isinstance(arg, Future):
+                    dfs_stack.append(arg)
+        else:
+            raise InternalError(f"Unexpected type of Future tree node: {type(node)}")
+
+    yield from reversed(yield_stack)

tensorlake/applications/applications.py

@@ -0,0 +1,49 @@
+from collections.abc import Generator, Iterator
+
+from .interface.function import Function, _is_application_function
+from .interface.request import Request
+from .interface.run import (
+    run_local_application,
+    run_remote_application,
+)
+
+# Internal utilities for working with applications.
+
+
+def filter_applications(
+    functions: Iterator[Function],
+) -> Generator[Function, None, None]:
+    """Yields all applications out of the supplied functions."""
+    for function in functions:
+        function: Function
+        if not _is_application_function(function):
+            continue
+        yield function
+
+
+def functions_for_application(
+    _application: Function,
+    functions: Iterator[Function],
+) -> list[Function]:
+    """Returns all loaded functions for the application build request.
+
+    We cannot reliably determine application membership statically, and one
+    application can call another application's function at runtime. For
+    compatibility with deploy/runtime packaging, each application build includes
+    all loaded Tensorlake functions.
+    """
+    return list(functions)
+
+
+def run_application(
+    application: Function | str, remote: bool, *args, **kwargs
+) -> Request:
+    """Runs the application remotely or locally depending on the `remote` parameter value.
+
+    This is a convenience wrapper around the `run_remote_application` and `run_local_application`.
+    It's not part of SDK interface, it's a helper function for writing tests.
+    """
+    if remote:
+        return run_remote_application(application, *args, **kwargs)
+    else:
+        return run_local_application(application, *args, **kwargs)

tensorlake/applications/blob_store/__init__.py

@@ -0,0 +1,4 @@
+from .blob import BLOB, BLOBChunk
+from .blob_store import BLOBStore
+
+__all__ = ["BLOB", "BLOBChunk", "BLOBStore"]

tensorlake/applications/blob_store/blob.py

@@ -0,0 +1,22 @@
+from pydantic import BaseModel
+
+
+class BLOBChunk(BaseModel):
+    # URI of the chunk.
+    # S3 URI if the data is stored in S3.
+    # Starts with "file://"" prefix if the data is stored on a local file system.
+    uri: str
+    # Actual size of chunk data if the BLOB has data.
+    # Requested chunk size if the BLOB has no data yet (can be used for data upload).
+    size: int
+    # None if the BLOB has data.
+    # ETag of the chunk data (generated by BLOB store) if the BLOB chunk was used to upload data.
+    etag: str | None
+
+
+class BLOB(BaseModel):
+    # ID of the BLOB
+    id: str
+    # Ordered chunks of the BLOB if the BLOB has data.
+    # Ordered chunks of the BLOB that can be used to upload data if the BLOB has no data yet.
+    chunks: list[BLOBChunk]

tensorlake/applications/blob_store/blob_store.py

@@ -0,0 +1,276 @@
+from concurrent.futures import FIRST_EXCEPTION, Future, ThreadPoolExecutor, wait
+from dataclasses import dataclass
+from typing import Any
+
+from tensorlake.applications.interface.exceptions import InternalError
+from tensorlake.applications.internal_logger import InternalLogger
+
+from .blob import BLOB, BLOBChunk
+from .local_fs_blob_store import LocalFSBLOBStore
+from .s3_blob_store import S3BLOBStore
+
+# S3 multipart uploads from EC2 instances gradually speed up until we reach 10 parallel chunk uploads.
+# S3 downloads from EC2 instances gradually speed up until we reach 5 parallel chunk downloads.
+# Then they slow down gradually. Using a common middleground value of 7 to be able to use a single
+# thread pool with a static max workers value.
+_MAX_WORKER_THREADS: int = 7
+# Upper limit in case a function has low CPU limit.
+_IO_WORKER_THREADS_PER_AVAILABLE_CPU: int = 3
+
+
+@dataclass
+class _ChunkInfo:
+    index: int
+    offset: int
+
+
+class BLOBStore:
+    """Dispatches generic BLOB store calls to their real backends.
+
+    Implements chunking. Thread-safe. Picklable.
+    """
+
+    def __init__(self, available_cpu_count: int):
+        """Creates a BLOB store that uses the supplied BLOB stores."""
+        self._available_cpu_count: int = available_cpu_count
+
+        max_io_workers: int = min(
+            available_cpu_count * _IO_WORKER_THREADS_PER_AVAILABLE_CPU,
+            _MAX_WORKER_THREADS,
+        )
+        self._io_workers_pool = ThreadPoolExecutor(
+            max_workers=max_io_workers, thread_name_prefix="BLOBStoreWorker"
+        )
+        self._local: LocalFSBLOBStore = LocalFSBLOBStore()
+        self._s3: S3BLOBStore = S3BLOBStore(io_workers_count=max_io_workers)
+
+    def close(self):
+        """Closes the BLOB store and its resources."""
+        self._io_workers_pool.shutdown(wait=True, cancel_futures=True)
+
+    def __getstate__(self):
+        """Get the state for pickling."""
+        return {
+            "available_cpu_count": self._available_cpu_count,
+        }
+
+    def __setstate__(self, state: dict[str, Any]):
+        """Set the state for unpickling."""
+        self.__init__(
+            available_cpu_count=state["available_cpu_count"],
+        )
+
+    def get(
+        self, blob: BLOB, offset: int, size: int, logger: InternalLogger
+    ) -> bytearray:
+        """Returns binary data stored in BLOB with the supplied URI at the supplied offset.
+
+        Raises InternalError on error.
+        """
+        if offset + size > _blob_size(blob):
+            raise InternalError(
+                f"Offset {offset} + size {size} is out of bounds for BLOB chunks of size {_blob_size(blob)}."
+            )
+
+        # Read data from BLOB chunks in parallel until all data is read.
+        # Minimize data copying by not creating any intermediate bytes/bytearray objects.
+        read_chunk_futures: list[Future] = []
+        destination: bytearray = bytearray(size)
+        destination_view: memoryview = memoryview(destination)
+        read_offset: int = offset
+
+        # Allow reads of size 0. This happens when an application function call with no arguments is made.
+        if size == 0:
+            return destination
+
+        first_chunk_info: _ChunkInfo = _find_chunk(blob, offset)
+        chunk_ix: int = first_chunk_info.index
+        offset_inside_chunk: int = offset - first_chunk_info.offset
+        while read_offset != (offset + size):
+            chunk: BLOBChunk = blob.chunks[chunk_ix]
+            chunk_read_size: int = min(
+                (offset + size) - read_offset, chunk.size - offset_inside_chunk
+            )
+            destination_offset: int = read_offset - offset
+            chunk_in_destination: memoryview = destination_view[
+                destination_offset : destination_offset + chunk_read_size
+            ]
+            read_chunk_futures.append(
+                self._io_workers_pool.submit(
+                    self._read_into,
+                    # Local file chunk URI points at the beginning of the file (not the chunk).
+                    # S3 chunk URI points at the beginning of the S3 object (not the chunk). This is performance optimization so we don't
+                    # need to presign a ranged S3 URI per chunk. We use a single presigned S3 URI for all BLOB chunks instead.
+                    blob_uri=chunk.uri,
+                    blob_read_offset=read_offset,
+                    destination=chunk_in_destination,
+                    logger=logger,
+                )
+            )
+
+            read_offset += chunk_read_size
+            offset_inside_chunk = (
+                0  # only read of first chunk can be not aligned at chunk boundary
+            )
+            chunk_ix += 1
+
+        wait(read_chunk_futures, return_when=FIRST_EXCEPTION)
+        for future in read_chunk_futures:
+            if future.exception() is not None:
+                raise InternalError(
+                    "Failed reading BLOB store chunk"
+                ) from future.exception()
+
+        return destination
+
+    def _read_into(
+        self,
+        blob_uri: str,
+        blob_read_offset: int,
+        destination: memoryview,
+        logger: InternalLogger,
+    ) -> bytes:
+        if _is_file_uri(blob_uri):
+            self._local.get(
+                uri=blob_uri,
+                offset=blob_read_offset,
+                destination=destination,
+                logger=logger,
+            )
+        else:
+            self._s3.get(
+                uri=blob_uri,
+                offset=blob_read_offset,
+                destination=destination,
+                logger=logger,
+            )
+
+    def put(self, blob: BLOB, data: list[bytes], logger: InternalLogger) -> BLOB:
+        """Stores the supplied binary data into the supplied BLOB starting from its very beginning.
+
+        Overwrites BLOB. Raises Exception on error.
+        Data can be smaller than the BLOB size, but not larger.
+        Returns the updated BLOB with chunks that were used for storing the data starting from the first chunk
+        in the original BLOB. The original order of the chunks is preserved. Chunks that were not used for
+        storing the data are not added to the returned BLOB. Each chunk in the returned BLOB has its size set to
+        the actual size of the data that was written to it and its etag returned by the storage backend.
+
+        Raises InternalError on error.
+        """
+        blob_size: int = _blob_size(blob)
+        data_size: int = sum(len(chunk) for chunk in data)
+
+        if data_size > blob_size:
+            raise InternalError(f"Data size {data_size} exceeds BLOB size {blob_size}.")
+
+        # Write data to BLOB chunks in parallel until all data is written.
+        # Minimize data copying by not creating any intermediate bytes/bytearray objects.
+        data_read_offset: int = 0
+        write_chunk_futures: list[Future] = []
+        uploaded_chunk_sizes: list[int] = []
+
+        data_ix: int = 0
+        read_offset_inside_data: int = 0
+        for chunk in blob.chunks:
+            chunk: BLOBChunk
+            chunk_data: list[memoryview] = []
+            chunk_data_size: int = 0
+            chunk_offset: int = data_read_offset
+            if data_ix == len(data):
+                break
+
+            # Fill the chunk with data until it is full.
+            while chunk_data_size != chunk.size and data_ix != len(data):
+                read_size: int = min(
+                    chunk.size - chunk_data_size,
+                    len(data[data_ix]) - read_offset_inside_data,
+                )
+                chunk_data.append(
+                    memoryview(data[data_ix])[
+                        read_offset_inside_data : read_offset_inside_data + read_size
+                    ]
+                )
+                chunk_data_size += read_size
+                read_offset_inside_data += read_size
+                data_read_offset += read_size
+                if read_offset_inside_data == len(data[data_ix]):
+                    data_ix += 1
+                    read_offset_inside_data = 0
+
+            # Write the chunk (should be full except the last one).
+            # Local file chunk URI points at the beginning of the file (not the chunk).
+            # S3 chunk URI contains chunk's index (part number).
+            write_chunk_futures.append(
+                self._io_workers_pool.submit(
+                    self._write_chunk,
+                    chunk_uri=chunk.uri,
+                    chunk_offset=chunk_offset,
+                    source=chunk_data,
+                    logger=logger,
+                )
+            )
+            uploaded_chunk_sizes.append(chunk_data_size)
+
+        wait(write_chunk_futures, return_when=FIRST_EXCEPTION)
+        uploaded_blob: BLOB = BLOB(
+            id=blob.id,
+            chunks=[],
+        )
+        for ix, future in enumerate(write_chunk_futures):
+            if future.exception() is not None:
+                raise InternalError(
+                    "Failed writing BLOB store chunk"
+                ) from future.exception()
+            # The futures list is ordered by the chunk index, so appending here preserves
+            # the original chunks order.
+            uploaded_chunk: BLOBChunk = blob.chunks[ix].model_copy()
+            uploaded_chunk.size = uploaded_chunk_sizes[ix]
+            uploaded_chunk.etag = future.result()
+            uploaded_blob.chunks.append(uploaded_chunk)
+
+        return uploaded_blob
+
+    def _write_chunk(
+        self,
+        chunk_uri: str,
+        chunk_offset: int,
+        source: list[memoryview],
+        logger: InternalLogger,
+    ) -> str:
+        if _is_file_uri(chunk_uri):
+            return self._local.put(
+                uri=chunk_uri,
+                offset=chunk_offset,
+                source=source,
+                logger=logger,
+            )
+        else:
+            return self._s3.put(
+                uri=chunk_uri,
+                source=source,
+                logger=logger,
+            )
+
+
+def _find_chunk(blob: BLOB, offset: int) -> _ChunkInfo:
+    """Returns info of the chunk where the supplied offset starts.
+
+    Raises IndexError if the offset is outside of the BLOB."""
+    current_offset: int = 0
+    for ix, chunk in enumerate(blob.chunks):
+        if current_offset + chunk.size > offset:
+            return _ChunkInfo(index=ix, offset=current_offset)
+        current_offset += chunk.size
+
+    raise InternalError(
+        f"Offset {offset} is out of bounds for BLOB chunks of size {current_offset}."
+    )
+
+
+def _blob_size(blob: BLOB) -> int:
+    """Returns the total size of the BLOB."""
+    return sum(chunk.size for chunk in blob.chunks)
+
+
+def _is_file_uri(uri: str) -> bool:
+    return uri.startswith("file://")

tensorlake/applications/blob_store/local_fs_blob_store.py

@@ -0,0 +1,77 @@
+import hashlib
+import os
+import os.path
+from typing import List
+
+from tensorlake.applications.interface.exceptions import InternalError
+from tensorlake.applications.internal_logger import InternalLogger
+
+
+class LocalFSBLOBStore:
+    """BLOB store that stores BLOBs in local file system."""
+
+    def get(
+        self,
+        uri: str,
+        offset: int,
+        destination: memoryview,
+        logger: InternalLogger,
+    ) -> None:
+        """Reads binary data stored in file at the supplied URI and offset into the destination memoryview.
+
+        The URI must be a file URI (starts with "file://"). The path must be absolute.
+        Raises InternalError on error.
+        """
+        blob_path: str = _blob_path_from_uri(uri)
+        if not os.path.isabs(blob_path):
+            raise InternalError(f"BLOB file path {blob_path} must be absolute")
+
+        if os.path.exists(blob_path):
+            with open(blob_path, mode="rb") as blob_file:
+                blob_file.seek(offset)
+                # memoryview ensures that the slice we pass points at destination.
+                blob_file.readinto(destination)
+        else:
+            raise InternalError(f"BLOB file at {blob_path} does not exist")
+
+    def put(
+        self,
+        uri: str,
+        offset: int,
+        source: List[memoryview],
+        logger: InternalLogger,
+    ) -> str:
+        """Stores the supplied memoryviews of binary data in a file at the supplied URI and offset.
+
+        The URI must be a file URI (starts with "file://"). The path must be absolute.
+        Overwrites existing file. Raises Exception on error.
+        Returns the ETag of the stored data.
+        """
+        blob_path: str = _blob_path_from_uri(uri)
+        if not os.path.isabs(blob_path):
+            raise InternalError(f"BLOB file path {blob_path} must be absolute")
+
+        os.makedirs(os.path.dirname(blob_path), exist_ok=True)
+        _create_file_if_doesnt_exist(blob_path)
+
+        hasher: hashlib.md5 = hashlib.md5()
+        with open(blob_path, mode="rb+") as blob_file:
+            blob_file.seek(offset)  # Adds zeroes if the file is smaller than offset.
+            for source_data in source:
+                blob_file.write(source_data)
+                hasher.update(source_data)
+
+        return hasher.hexdigest()
+
+
+def _blob_path_from_uri(uri: str) -> str:
+    return uri[7:]  # strip "file://" prefix
+
+
+def _create_file_if_doesnt_exist(path: str) -> None:
+    """Creates an empty file at the specified path if it doesn't exist."""
+    try:
+        with open(path, "x"):
+            pass  # File was created
+    except FileExistsError:
+        pass  # File already exists, no action needed