tensorlake 0.1.0__tar.gz

tensorlake-0.1.0/PKG-INFO

@@ -0,0 +1,18 @@
+Metadata-Version: 2.3
+Name: tensorlake
+Version: 0.1.0
+Summary: Petabyte scale data framework for unstructured data of any modality
+Home-page: https://github.com/tensorlakeai/tensorlake
+Author: Tensorlake Inc.
+Author-email: support@tensorlake.ai
+Requires-Python: >=3.10,<4.0
+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
+Requires-Dist: cloudpickle (>=3.1.0,<4.0.0)
+Requires-Dist: docker (>=7.1.0,<8.0.0)
+Requires-Dist: httpx[http2] (>=0.28.1,<0.29.0)
+Requires-Dist: pydantic (==2.10.4)
+Project-URL: Repository, https://github.com/tensorlakeai/tensorlake

tensorlake-0.1.0/pyproject.toml

@@ -0,0 +1,25 @@
+[tool.poetry]
+name = "tensorlake"
+version = "0.1.0"
+description = "Petabyte scale data framework for unstructured data of any modality"
+authors = ["Tensorlake Inc. <support@tensorlake.ai>"]
+homepage = "https://github.com/tensorlakeai/tensorlake"
+repository = "https://github.com/tensorlakeai/tensorlake"
+
+[tool.poetry.dependencies]
+python = "^3.10"
+httpx = { version = "^0.28.1", extras = ["http2"] }
+pydantic = "2.10.4"
+cloudpickle = "^3.1.0"
+docker = "^7.1.0"
+
+[tool.poetry.scripts]
+tensorlake = "tensorlake.cli:tensorlake"
+
+[build-system]
+requires = ["poetry-core>=1.0.0"]
+build-backend = "poetry.core.masonry.api"
+
+[tool.poetry.group.dev.dependencies]
+black = "^24.10.0"
+pylint = "^3.3.0"

tensorlake-0.1.0/tensorlake/__init__.py

@@ -0,0 +1,31 @@
+from . import data_loaders
+from .functions_sdk.graph import Graph
+from .functions_sdk.image import Image
+from .functions_sdk.indexify_functions import (
+    IndexifyFunction,
+    IndexifyRouter,
+    get_ctx,
+    indexify_router,
+    tensorlake_function,
+)
+from .functions_sdk.pipeline import Pipeline
+from .http_client import TensorlakeClient
+from .remote_graph import RemoteGraph
+from .remote_pipeline import RemotePipeline
+from .settings import DEFAULT_SERVICE_URL
+
+__all__ = [
+    "data_loaders",
+    "Graph",
+    "RemoteGraph",
+    "Pipeline",
+    "RemotePipeline",
+    "Image",
+    "tensorlake_function",
+    "get_ctx",
+    "IndexifyFunction",
+    "IndexifyRouter",
+    "indexify_router",
+    "DEFAULT_SERVICE_URL",
+    "TensorlakeClient",
+]

tensorlake-0.1.0/tensorlake/builder/client.py

@@ -0,0 +1,73 @@
+import os
+from typing import Dict, Optional
+
+import httpx
+
+from tensorlake.functions_sdk.image import Build
+
+
+class ImageBuilderClient:
+    def __init__(self, build_service: str, api_key):
+        self.client = httpx
+        self.build_service = build_service
+        self.headers = {}
+        if api_key:
+            self.headers["Authorization"] = f"Bearer {api_key}"
+
+    @classmethod
+    def from_env(cls):
+        api_key = os.getenv("TENSORLAKE_API_KEY")
+        indexify_url = os.getenv("INDEXIFY_URL", "https://api.tensorlake.ai")
+        build_url = os.getenv(
+            "TENSORLAKE_BUILD_SERVICE", f"{indexify_url}/images"
+        )  # Mainly used for debugging/local testing
+        return cls(build_url, api_key)
+
+    def get(self, endpoint: str, params: Optional[Dict] = None):
+        res = self.client.get(
+            f"{self.build_service}{endpoint}", params=params, headers=self.headers
+        )
+        res.raise_for_status()
+        return res.json()
+
+    def get_build(self, build_id: int):
+        return Build.model_validate(self.get(f"/v1/builds/{build_id}"))
+
+    def get_build_logs(self, build_id: int):
+        log_response = self.client.get(f"{self.build_service}/builds/{build_id}/log")
+        log_response.raise_for_status()
+        return log_response.content.decode("utf-8")
+
+    def post(self, endpoint: str, data: Dict, files: Dict):
+        res = self.client.post(
+            f"{self.build_service}{endpoint}",
+            data=data,
+            files=files,
+            headers=self.headers,
+            timeout=60,
+        )
+        res.raise_for_status()
+        return res.json()
+
+    def build_exists(self, image_name: str, image_hash: str):
+        builds = self.find_build(image_name, image_hash)
+        return builds != []
+
+    def find_build(self, image_name: str, image_hash: str):
+        params = {"image_name": image_name, "image_hash": image_hash}
+        res = self.client.get(
+            f"{self.build_service}/v1/builds", headers=self.headers, params=params
+        )
+        res.raise_for_status()
+        return [Build.model_validate(b) for b in res.json()]
+
+    def get_latest_build(self, image_name: str) -> Build:
+        res = self.client.get(
+            f"{self.build_service}/v1/builds",
+            headers=self.headers,
+            params={"image_name": image_name},
+        )
+        res.raise_for_status()
+        builds = [Build.model_validate(b) for b in res.json()]
+        builds.sort(key=lambda b: b.created_at, reverse=True)
+        return builds[0]

tensorlake-0.1.0/tensorlake/cli.py

@@ -0,0 +1,212 @@
+import os
+import tempfile
+import time
+from typing import Dict, List
+
+import click
+
+from tensorlake import Graph, Image, RemoteGraph, TensorlakeClient
+from tensorlake.builder.client import ImageBuilderClient
+from tensorlake.functions_sdk.image import Build
+
+
+@click.group()
+def tensorlake():
+    pass
+
+
+@click.command()
+@click.argument("workflow_file", type=click.File("r"))
+def deploy(workflow_file: click.File):
+    """Deploy a workflow to tensorlake."""
+
+    click.echo(f"Preparing deployment for {workflow_file.name}")
+    builder = ImageBuilderClient.from_env()
+    seen_images: Dict[Image, str] = {}
+    deployed_graphs: List[Graph] = []
+
+    # Read the graph file and build the images
+    workflow_globals = {}
+    with open(workflow_file.name, "r") as f:
+        exec(f.read(), workflow_globals)
+
+    for name, obj in workflow_globals.items():
+        if isinstance(obj, Graph):
+            deployed_graphs.append(obj)
+            for node_name, node_obj in obj.nodes.items():
+                image = node_obj.image
+                if image in seen_images:
+                    continue
+                seen_images[image] = image.hash()
+
+    _prepare_images(builder, seen_images)
+
+    # If we are still here then our images should all have URIs
+
+    # TODO: Fold calls to the platform API into a client class.
+    indexify_addr = os.getenv("INDEXIFY_URL", "https://api.tensorlake.ai")
+    introspect_response = builder.client.post(
+        f"{indexify_addr}/platform/v1/keys/introspect", headers=builder.headers
+    )
+    introspect_response.raise_for_status()
+    project_id = introspect_response.json()["projectId"]
+
+    client = TensorlakeClient(namespace=project_id, service_url="http://localhost:8900")
+    click.secho("Everything looks good, deploying now", fg="green")
+    for graph in deployed_graphs:
+        # TODO: Every time we post we get a new version, is that expected or the client should do the checks?
+        remote = RemoteGraph.deploy(graph, client=client)
+
+
+def _wait_for_build(builder: ImageBuilderClient, build: Build):
+    click.echo(f"Waiting for {build.image_name} to build")
+    while build.status != "completed":
+        time.sleep(1)
+        build = builder.get_build(build.id)
+
+    if build.push_completed_at:
+        build_duration = build.build_completed_at - build.push_completed_at
+        click.echo(f"Building completed in {build.image_name} {build_duration.seconds}")
+    return build
+
+
+def _build_image(
+    builder: ImageBuilderClient, image: Image, image_hash: str = ""
+) -> Build:
+    click.echo(f"Building {image._image_name}")
+    fd, context_file = tempfile.mkstemp()
+    image.build_context(context_file)
+
+    click.echo(
+        f"{image._image_name}: Posting {os.path.getsize(context_file)} bytes of context to build service...."
+    )
+    files = {"context": open(context_file, "rb")}
+    data = {"name": image._image_name, "hash": image_hash}
+
+    res = builder.client.post(
+        f"{builder.build_service}/v1/builds",
+        data=data,
+        files=files,
+        headers=builder.headers,
+        timeout=60,
+    )
+    res.raise_for_status()
+    build = Build.model_validate(res.json())
+
+    return _wait_for_build(builder, build)
+
+
+def _show_failed_summary(builder: ImageBuilderClient, build: Build):
+    click.secho(
+        f"Building {build.image_name} failed with error message: {build.error_message}",
+        fg="red",
+    )
+
+    log_response = builder.client.get(
+        f"{builder.build_service}/v1/builds/{build.id}/log", headers=builder.headers
+    )
+    if log_response.status_code == 200:
+        log = log_response.content.decode("utf-8")
+        click.echo(log)
+    elif log_response.status_code == 404:
+        click.echo("Logs not found")
+    else:
+        log_response.raise_for_status()
+
+
+def _prepare_images(builder: ImageBuilderClient, images: Dict[Image, str]):
+    ready_builds: Dict[Image, Build] = {}
+    # Go through the images and build anything that hasn't been built
+    for image, image_hash in images.items():
+        builds = builder.find_build(image._image_name, image_hash)
+
+        if builds:
+            build = builds[0]
+            if build.status == "completed":
+                if build.result == "failed":
+                    _show_failed_summary(builder, build)
+                else:
+                    click.secho(f"Image '{build.image_name}' is built", fg="green")
+                    ready_builds[image] = build
+
+            elif build.status in ("ready", "building"):
+                build = _wait_for_build(builder, build)
+                if build.result != "failed":
+                    ready_builds[image] = build
+                else:
+                    _show_failed_summary(builder, build)
+
+        else:
+            build = _build_image(builder, image, image_hash=image_hash)
+
+    # Find any blockers and report them to the users
+    blockers = []
+    for image in images:
+        if image not in ready_builds:
+            blockers.append(image)
+            click.secho(
+                f"Image {image._image_name} could not be built, this is blocking deployment",
+                fg="red",
+            )
+        else:
+            build = ready_builds[image]
+            image.uri = build.uri
+    if blockers:
+        raise click.Abort
+
+
+@click.command()
+@click.argument("workflow_file", type=click.File("r"))
+def prepare(workflow_file: click.File):
+    """Prepare a workflow and it's artifacts for deployment."""
+
+    click.echo(f"Preparing deployment for {workflow_file.name}")
+    client = ImageBuilderClient.from_env()
+    seen_images: Dict[Image, str] = {}
+
+    # Read the graph file and build the images
+    workflow_globals = {}
+    with open(workflow_file.name, "r") as f:
+        exec(f.read(), workflow_globals)
+
+    for name, obj in workflow_globals.items():
+        if isinstance(obj, Graph):
+            click.echo(f"Found graph {name}")
+            for node_name, node_obj in obj.nodes.items():
+                image = node_obj.image
+                click.echo(
+                    f"graph function {node_name} uses image '{image._image_name}'"
+                )
+                if image in seen_images:
+                    continue
+                seen_images[image] = image.hash()
+
+    click.echo(f"Found {len(seen_images)} images in this workflow")
+    _prepare_images(client, seen_images)
+
+
+@click.command(help="Extract and display logs from tensorlake")
+@click.option("--image", "-i")
+def show_logs(image: str):
+    if image:
+        builder = ImageBuilderClient.from_env()
+        if ":" in image:
+            image, image_hash = image.split(":")
+            build = builder.find_build(image, image_hash)[0]
+        else:
+            build = builder.get_latest_build(image)
+
+        log_response = builder.client.get(
+            f"{builder.build_service}/v1/builds/{build.id}/log", headers=builder.headers
+        )
+        if log_response.status_code == 200:
+            log = log_response.content.decode("utf-8")
+            print(log)
+
+
+tensorlake.add_command(deploy)
+tensorlake.add_command(prepare)
+tensorlake.add_command(show_logs)
+
+if __name__ == "__main__":
+    tensorlake()

tensorlake-0.1.0/tensorlake/data_loaders/__init__.py

@@ -0,0 +1,58 @@
+import hashlib
+import mimetypes
+import os
+from abc import ABC, abstractmethod
+from typing import List
+
+from pydantic import BaseModel
+
+
+class FileMetadata(BaseModel):
+    path: str
+    file_size: int
+    mime_type: str
+    md5_hash: str
+    created_at: int
+    updated_at: int
+
+    @classmethod
+    def from_path(cls, path: str):
+        file_size = os.path.getsize(path)
+        mime_type = mimetypes.guess_type(path)[0]
+
+        # Compute MD5 hash
+        hash_md5 = hashlib.md5()
+        with open(path, "rb") as f:
+            for chunk in iter(lambda: f.read(4096), b""):
+                hash_md5.update(chunk)
+        md5_hash = hash_md5.hexdigest()
+
+        created_at = int(os.path.getctime(path))
+        updated_at = int(os.path.getmtime(path))
+
+        return cls(
+            path=path,
+            file_size=file_size,
+            mime_type=str(mime_type),
+            md5_hash=md5_hash,
+            created_at=created_at,
+            updated_at=updated_at,
+        )
+
+
+class DataLoader(ABC):
+    @abstractmethod
+    def load(self) -> List[FileMetadata]:
+        pass
+
+    @abstractmethod
+    def read_all_bytes(self, file_metadata: FileMetadata) -> bytes:
+        pass
+
+    @abstractmethod
+    def state(self) -> dict:
+        pass
+
+
+from .local_directory_loader import LocalDirectoryLoader
+from .url_loader import UrlLoader

tensorlake-0.1.0/tensorlake/data_loaders/local_directory_loader.py

@@ -0,0 +1,37 @@
+import os
+from typing import List, Optional
+
+from . import DataLoader, FileMetadata
+
+
+class LocalDirectoryLoader(DataLoader):
+    def __init__(
+        self,
+        directory: str,
+        file_extensions: Optional[List[str]] = None,
+        state: dict = {},
+    ):
+        self.directory = directory
+        self.file_extensions = file_extensions
+        self.processed_files = set(state.get("processed_files", []))
+
+    def load(self) -> List[FileMetadata]:
+        file_metadata_list = []
+        for root, _, files in os.walk(self.directory):
+            for file in files:
+                if self.file_extensions is None or any(
+                    file.endswith(ext) for ext in self.file_extensions
+                ):
+                    file_path = os.path.join(root, file)
+                    if file_path not in self.processed_files:
+                        file_metadata_list.append(FileMetadata.from_path(file_path))
+                        self.processed_files.add(file_path)
+
+        return file_metadata_list
+
+    def read_all_bytes(self, file: FileMetadata) -> bytes:
+        with open(file.path, "rb") as f:
+            return f.read()
+
+    def state(self) -> dict:
+        return {"processed_files": list(self.processed_files)}

tensorlake-0.1.0/tensorlake/data_loaders/url_loader.py

@@ -0,0 +1,52 @@
+import email.utils
+from typing import List
+
+import httpx
+
+from . import DataLoader, FileMetadata
+
+
+def convert_date_to_epoch(date_str: str) -> int:
+    """
+    Convert a date string from URL header to Unix epoch time.
+
+    Args:
+        date_str (str): The date string from the URL header.
+
+    Returns:
+        int: The Unix epoch time.
+    """
+    if not date_str:
+        return 0
+    parsed_date = email.utils.parsedate_to_datetime(date_str)
+    return int(parsed_date.timestamp())
+
+
+class UrlLoader(DataLoader):
+    def __init__(self, urls: List[str], state: dict = {}):
+        self.urls = urls
+
+    def load(self) -> List[FileMetadata]:
+        file_metadata_list = []
+        for url in self.urls:
+            response = httpx.head(url, follow_redirects=True)
+            file_metadata_list.append(
+                FileMetadata(
+                    path=url,
+                    file_size=response.headers.get("content-length", 0),
+                    mime_type=response.headers.get("content-type"),
+                    md5_hash="",
+                    created_at=convert_date_to_epoch(response.headers.get("date")),
+                    updated_at=convert_date_to_epoch(
+                        response.headers.get("last-modified")
+                    ),
+                )
+            )
+        return file_metadata_list
+
+    def read_all_bytes(self, file: FileMetadata) -> bytes:
+        response = httpx.get(file.path, follow_redirects=True)
+        return response.content
+
+    def state(self) -> dict:
+        return {}

tensorlake-0.1.0/tensorlake/error.py

@@ -0,0 +1,8 @@
+class ApiException(Exception):
+    def __init__(self, message: str) -> None:
+        super().__init__(message)
+
+
+class GraphStillProcessing(Exception):
+    def __init__(self) -> None:
+        super().__init__("graph is still processing")

tensorlake-0.1.0/tensorlake/functions_sdk/data_objects.py

@@ -0,0 +1,27 @@
+from typing import Any, Dict, List, Literal, Optional, Union
+
+from pydantic import BaseModel, Json
+
+
+class FileInput(BaseModel):
+    url: str
+    mime_type: Optional[str] = None
+    metadata: Optional[Dict[str, Json]] = None
+    sha_256: Optional[str] = None
+
+
+class RouterOutput(BaseModel):
+    edges: List[str]
+
+
+class IndexifyData(BaseModel):
+    id: Optional[str] = None
+    payload: Union[bytes, str]
+    encoder: Literal["cloudpickle", "json"] = "cloudpickle"
+
+
+class File(BaseModel):
+    data: bytes
+    mime_type: Optional[str] = None
+    metadata: Optional[Dict[str, Any]] = None
+    sha_256: Optional[str] = None