graphbook 0.4.0__tar.gz

graphbook-0.4.0/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2024 Graphbook AI
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

graphbook-0.4.0/PKG-INFO

@@ -0,0 +1,99 @@
+Metadata-Version: 2.1
+Name: graphbook
+Version: 0.4.0
+Summary: An extensible ML workflow framework built for data scientists and ML engineers.
+Home-page: https://graphbook.ai
+License: MIT
+Keywords: ml,workflow,framework,pytorch,data science,machine learning,ai
+Author: Richard Franklin
+Author-email: rsamf@graphbook.ai
+Requires-Python: >=3.11,<4.0
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Operating System :: OS Independent
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Requires-Dist: aiohttp (>=3.9.4,<4.0.0)
+Requires-Dist: dill (>=0.3.8,<0.4.0)
+Requires-Dist: gputil (>=1.4.0,<2.0.0)
+Requires-Dist: pillow (>=10.3.0,<11.0.0)
+Requires-Dist: psutil (>=6.0.0,<7.0.0)
+Requires-Dist: python-magic (>=0.4.27,<0.5.0)
+Requires-Dist: torch (>=2.3.1,<3.0.0)
+Requires-Dist: torchvision (>=0.18.1,<0.19.0)
+Requires-Dist: watchdog (>=4.0.0,<5.0.0)
+Project-URL: Documentation, https://docs.graphbook.ai
+Project-URL: Repository, https://github.com/graphbookai/graphbook
+Description-Content-Type: text/markdown
+
+<p align="center">
+  <a href="https://graphbook.ai">
+    <img src="docs/_static/graphbook.png" alt="Logo" width=256>
+  </a>
+
+  <h1 align="center">Graphbook</h1>
+
+  <p align="center">
+    The ML workflow framework
+    <br>
+    <a href="https://github.com/graphbookai/graphbook/issues/new?template=bug.md">Report bug</a>
+    ·
+    <a href="https://github.com/graphbookai/graphbook/issues/new?template=feature.md&labels=feature">Request feature</a>
+  </p>
+
+  <p align="center">
+    <a href="#overview">Overview</a> •
+    <a href="#current-features">Current Features</a> •
+    <a href="#getting-started">Getting Started</a> •
+    <a href="#collaboration">Collaboration</a>
+    </p>
+</p>
+
+## Overview
+Graphbook is a framework for building efficient, visual DAG-structured ML workflows composed of nodes written in Python. Graphbook provides common ML processing features such as multiprocessing IO and automatic batching, and it features a web-based UI to assemble, monitor, and execute data processing workflows. It can be used to prepare training data for custom ML models, experiment with custom trained or off-the-shelf models, and to build ML-based ETL applications. Custom nodes can be built in Python, and Graphbook will behave like a framework and call lifecycle methods on those nodes.
+
+## Current Features
+- ​​Graph-based visual editor to experiment and create complex ML workflows
+- Caches outputs and only re-executes parts of the workflow that changes between executions
+- UI monitoring components for logs and outputs per node
+- Custom buildable nodes with Python
+- Automatic batching for Pytorch tensors
+- Multiprocessing I/O to and from disk and network
+- Customizable multiprocessing functions
+- Ability to execute entire graphs, or individual subgraphs/nodes
+- Ability to execute singular batches of data
+- Ability to pause graph execution
+- Basic nodes for filtering, loading, and saving outputs
+- Node grouping and subflows
+- Autosaving and shareable serialized workflow files
+- Registers node code changes without needing a restart
+- Monitorable CPU and GPU resource usage
+
+## Getting Started
+### Install from PyPI
+1. `pip install graphbook`
+1. `graphbook`
+1. Visit http://localhost:8007
+
+### Install with Docker
+1. Pull and run the downloaded image
+    ```bash
+    docker run --rm -p 8005:8005 -p 8006:8006 -p 8007:8007 -v $PWD/workflows:/app/workflows rsamf/graphbook:latest
+    ```
+1. Visit http://localhost:8007
+
+Visit the [docs](https://docs.graphbook.ai) to learn more on how to create custom nodes and workflows with Graphbook.
+
+## Collaboration
+This is a guide on how to get started developing Graphbook. If you are simply using Graphbook, view the [Getting Started](#getting-started) section.
+
+### Run Graphbook in Development Mode
+You can use any other virtual environment solution, but `poetry` is used in the steps below.
+1. Clone the repo and `cd graphbook`
+1. `poetry install --with dev`
+1. `poetry shell`
+1. `python graphbook/server.py`
+1. `cd web`
+1. `npm install`
+1. `npm run dev`
+

graphbook-0.4.0/README.md

@@ -0,0 +1,70 @@
+<p align="center">
+  <a href="https://graphbook.ai">
+    <img src="docs/_static/graphbook.png" alt="Logo" width=256>
+  </a>
+
+  <h1 align="center">Graphbook</h1>
+
+  <p align="center">
+    The ML workflow framework
+    <br>
+    <a href="https://github.com/graphbookai/graphbook/issues/new?template=bug.md">Report bug</a>
+    ·
+    <a href="https://github.com/graphbookai/graphbook/issues/new?template=feature.md&labels=feature">Request feature</a>
+  </p>
+
+  <p align="center">
+    <a href="#overview">Overview</a> •
+    <a href="#current-features">Current Features</a> •
+    <a href="#getting-started">Getting Started</a> •
+    <a href="#collaboration">Collaboration</a>
+    </p>
+</p>
+
+## Overview
+Graphbook is a framework for building efficient, visual DAG-structured ML workflows composed of nodes written in Python. Graphbook provides common ML processing features such as multiprocessing IO and automatic batching, and it features a web-based UI to assemble, monitor, and execute data processing workflows. It can be used to prepare training data for custom ML models, experiment with custom trained or off-the-shelf models, and to build ML-based ETL applications. Custom nodes can be built in Python, and Graphbook will behave like a framework and call lifecycle methods on those nodes.
+
+## Current Features
+- ​​Graph-based visual editor to experiment and create complex ML workflows
+- Caches outputs and only re-executes parts of the workflow that changes between executions
+- UI monitoring components for logs and outputs per node
+- Custom buildable nodes with Python
+- Automatic batching for Pytorch tensors
+- Multiprocessing I/O to and from disk and network
+- Customizable multiprocessing functions
+- Ability to execute entire graphs, or individual subgraphs/nodes
+- Ability to execute singular batches of data
+- Ability to pause graph execution
+- Basic nodes for filtering, loading, and saving outputs
+- Node grouping and subflows
+- Autosaving and shareable serialized workflow files
+- Registers node code changes without needing a restart
+- Monitorable CPU and GPU resource usage
+
+## Getting Started
+### Install from PyPI
+1. `pip install graphbook`
+1. `graphbook`
+1. Visit http://localhost:8007
+
+### Install with Docker
+1. Pull and run the downloaded image
+    ```bash
+    docker run --rm -p 8005:8005 -p 8006:8006 -p 8007:8007 -v $PWD/workflows:/app/workflows rsamf/graphbook:latest
+    ```
+1. Visit http://localhost:8007
+
+Visit the [docs](https://docs.graphbook.ai) to learn more on how to create custom nodes and workflows with Graphbook.
+
+## Collaboration
+This is a guide on how to get started developing Graphbook. If you are simply using Graphbook, view the [Getting Started](#getting-started) section.
+
+### Run Graphbook in Development Mode
+You can use any other virtual environment solution, but `poetry` is used in the steps below.
+1. Clone the repo and `cd graphbook`
+1. `poetry install --with dev`
+1. `poetry shell`
+1. `python graphbook/server.py`
+1. `cd web`
+1. `npm install`
+1. `npm run dev`

graphbook-0.4.0/graphbook/__init__.py

@@ -0,0 +1 @@
+from .note import Note

graphbook-0.4.0/graphbook/custom_nodes.py

@@ -0,0 +1,122 @@
+import asyncio
+from watchdog.events import FileSystemEvent, FileSystemEventHandler
+from watchdog.observers import Observer
+import importlib
+import importlib.util
+import hashlib
+import sys
+import os
+import os.path as osp
+import inspect
+from graphbook.steps import Step, BatchStep, SourceStep, AsyncStep, Split, SplitNotesByItems, SplitItemField
+from graphbook.resources import Resource, FunctionResource
+
+BUILT_IN_STEPS = [Step, BatchStep, SourceStep, AsyncStep, Split, SplitNotesByItems, SplitItemField]
+BUILT_IN_RESOURCES = [Resource, FunctionResource]
+
+class CustomModuleEventHandler(FileSystemEventHandler):
+    def __init__(self, root_path, handler):
+        super().__init__()
+        self.root_path = osp.abspath(root_path)
+        self.handler = handler
+        self.ha = {}
+
+    def on_created(self, event):
+        if event.is_directory:
+            return
+        self.handle_new_file_sync(event.src_path)
+
+    def on_modified(self, event):
+        if event.is_directory:
+            return
+        self.handle_new_file_sync(event.src_path)
+
+    def on_deleted(self, event):
+        if event.is_directory:
+            return
+
+    def on_moved(self, event: FileSystemEvent) -> None:
+        if event.is_directory:
+            return
+        self.handle_new_file_sync(event.dest_path)
+
+    async def handle_new_file(self, filename: str):
+        filename = osp.abspath(filename)
+        assert filename.startswith(
+            self.root_path
+        ), f"Received extraneous file {filename} during tracking of {self.root_path}"
+        if not filename.endswith(".py"):
+            return
+
+        with open(filename, "r") as f:
+            contents = f.read()
+
+        hash_code = hashlib.md5(contents.encode()).hexdigest()
+        og_hash_code = self.ha.get(filename, None)
+        if hash_code == og_hash_code:
+            return
+
+        self.ha[filename] = hash_code
+        filename = filename[len(self.root_path) + 1 :]
+        components = filename[: filename.index(".py")].split("/")
+        module_name = ".".join(components)
+
+        try:
+            if og_hash_code is None:
+                importlib.import_module(module_name)
+                print("Loaded", module_name)
+            else:
+                module = importlib.import_module(module_name)
+                importlib.reload(module)
+                print("Reloaded", module_name)
+        except Exception as e:
+            print(f"Error loading {module_name}: {e}")
+            return
+
+        module = sys.modules[module_name]
+        await self.handler(filename, module)
+
+    def handle_new_file_sync(self, filename: str):
+        asyncio.run(self.handle_new_file(filename))
+
+    async def init_custom_nodes(self):
+        for root, dirs, files in os.walk(self.root_path):
+            for file in files:
+                await self.handle_new_file(osp.join(root, file))
+
+    def init_custom_nodes_sync(self):
+        asyncio.run(self.init_custom_nodes())
+
+
+class CustomNodeImporter:
+    def __init__(self, path, step_handler, resource_handler):
+        self.websocket = None
+        self.path = path
+        self.step_handler = step_handler
+        self.resource_handler = resource_handler
+        sys.path.append(path)
+        self.observer = Observer()
+        self.event_handler = CustomModuleEventHandler(path, self.on_module)
+        self.event_handler.init_custom_nodes_sync()
+
+    def set_websocket(self, websocket):
+        self.websocket = websocket
+
+    async def on_module(self, filename, mod):
+        for name, obj in inspect.getmembers(mod):
+            if inspect.isclass(obj):
+                if issubclass(obj, Step) and not obj in BUILT_IN_STEPS:
+                    await self.step_handler(filename, name, obj)
+                if issubclass(obj, Resource) and not obj in BUILT_IN_RESOURCES:
+                    await self.resource_handler(filename, name, obj)
+
+        if self.websocket is not None and not self.websocket.closed:
+            await self.websocket.send_json({"event": "node_updated"})
+
+    def start_observer(self):
+        self.observer.schedule(self.event_handler, self.path, recursive=True)
+        self.observer.start()
+
+    def stop_observer(self):
+        self.observer.stop()
+        self.observer.join()

graphbook-0.4.0/graphbook/dataloading.py

@@ -0,0 +1,250 @@
+from typing import List
+import queue
+import torch.multiprocessing as mp
+
+MP_WORKER_TIMEOUT = 5.0
+MAX_RESULT_QUEUE_SIZE = 32
+
+
+def do_load(work_queue: mp.Queue, result_queue: mp.Queue):
+    if result_queue.full():
+        return
+
+    try:
+        item, index, record_id, load_fn, consumer_id = work_queue.get(False)
+    except queue.Empty:
+        return
+
+    try:
+        result_tensor = load_fn(item)
+        result = (result_tensor, index)
+        to_return = ((result, record_id), consumer_id)
+    except Exception as e:
+        to_return = ((None, record_id), consumer_id)
+        print(
+            f"Worker Error: Could not process input {item}. The following exception was raised: {e}. Check your load_fn."
+        )
+
+    result_queue.put(to_return, block=False)
+
+
+def do_dump(
+    work_queue: mp.Queue, result_queue: mp.Queue, output_dir: str, uid: int
+) -> bool:
+    if result_queue.full():
+        return False
+
+    try:
+        data, item_key, record_id, dump_fn, consumer_id = work_queue.get(False)
+    except queue.Empty:
+        return False
+
+    try:
+        output_fn = dump_fn(data, output_dir, uid)
+        result = (item_key, output_fn)
+        to_return = ((result, record_id), consumer_id)
+    except Exception as e:
+        print(f"Could not dump input {data}. The following exception was raised: {e}")
+        to_return = ((None, record_id), consumer_id)
+
+    result_queue.put(to_return, block=False)
+    return True
+
+
+def worker_loop(
+    rank: int,
+    num_processes: int,
+    load_queue: mp.Queue,
+    dump_queue: mp.Queue,
+    load_result_queue: mp.Queue,
+    dump_result_queue: mp.Queue,
+    dump_dir: str,
+    close_event: mp.Event,
+):
+    try:
+        dump_ctr = rank
+        while not close_event.is_set():
+            do_load(load_queue, load_result_queue)
+            did_receive_work = do_dump(
+                dump_queue, dump_result_queue, dump_dir, dump_ctr
+            )
+            if did_receive_work:
+                dump_ctr += num_processes
+    except KeyboardInterrupt:
+        pass
+
+
+class Dataloader:
+    def __init__(self, dump_dir: str, num_workers: int = 1):
+        self.dump_dir = dump_dir
+        self.num_workers = num_workers
+        self.context = mp.get_context("spawn")
+        self.consumer_load_queues = {}
+        self.consumer_dump_queues = {}
+        self.total_consumer_size = 0
+        self._start_workers()
+
+    def _start_workers(self):
+        self._workers: List[mp.Process] = []
+        self._worker_queue_cycle = 0
+        self._load_queues: List[mp.Queue] = []
+        self._dump_queues: List[mp.Queue] = []
+        self._load_result_queues: List[mp.Queue] = []
+        self._dump_result_queues: List[mp.Queue] = []
+        self._close_event: mp.Event = self.context.Event()
+        for i in range(self.num_workers):
+            load_queue = self.context.Queue()
+            dump_queue = self.context.Queue()
+            load_result_queue = self.context.Queue(maxsize=MAX_RESULT_QUEUE_SIZE)
+            dump_result_queue = self.context.Queue(maxsize=MAX_RESULT_QUEUE_SIZE)
+            load_queue.cancel_join_thread()
+            dump_queue.cancel_join_thread()
+            w = self.context.Process(
+                target=worker_loop,
+                args=(
+                    i,
+                    self.num_workers,
+                    load_queue,
+                    dump_queue,
+                    load_result_queue,
+                    dump_result_queue,
+                    self.dump_dir,
+                    self._close_event,
+                ),
+            )
+            w.daemon = True
+            w.start()
+            self._load_queues.append(load_queue)
+            self._dump_queues.append(dump_queue)
+            self._load_result_queues.append(load_result_queue)
+            self._dump_result_queues.append(dump_result_queue)
+            self._workers.append(w)
+
+    def setup(self, consumer_ids: List[int]):
+        for c in consumer_ids:
+            if c not in self.consumer_load_queues:
+                self.consumer_load_queues[c] = queue.Queue()
+            if c not in self.consumer_dump_queues:
+                self.consumer_dump_queues[c] = queue.Queue()
+        unused_ids = set(self.consumer_load_queues.keys()) - set(consumer_ids)
+        for c in unused_ids:
+            self.total_consumer_size -= self.consumer_load_queues[c].qsize()
+            del self.consumer_load_queues[c]
+        unused_ids = set(self.consumer_dump_queues.keys()) - set(consumer_ids)
+        for c in unused_ids:
+            self.total_consumer_size -= self.consumer_dump_queues[c].qsize()
+            del self.consumer_dump_queues[c]
+
+    def shutdown(self):
+        if len(self._workers) == 0:
+            return
+        try:
+            self._close_event.set()
+            for q in self._load_queues:
+                q.cancel_join_thread()
+                q.close()
+            for q in self._dump_queues:
+                q.cancel_join_thread()
+                q.close()
+            for w in self._workers:
+                w.join(timeout=MP_WORKER_TIMEOUT)
+        finally:
+            for w in self._workers:
+                if w.is_alive():
+                    w.terminate()
+
+    def _handle_queues(self):
+        for queues, consumers in zip(
+            [self._load_result_queues, self._dump_result_queues],
+            [self.consumer_load_queues, self.consumer_dump_queues],
+        ):
+            for q in queues:
+                while (
+                    not q.empty() and self.total_consumer_size < MAX_RESULT_QUEUE_SIZE
+                ):
+                    result, consumer_id = q.get(False)
+                    if consumer_id not in consumers:
+                        continue
+                    consumers[consumer_id].put(result, block=False)
+                    self.total_consumer_size += 1
+
+    def get_all_sizes(self):
+        return {
+            "load": [q.qsize() for q in self._load_queues],
+            "dump": [q.qsize() for q in self._dump_queues],
+            "load_result": [q.qsize() for q in self._load_result_queues],
+            "dump_result": [q.qsize() for q in self._dump_result_queues],
+            "total_consumer_size": self.total_consumer_size,
+        }
+
+    def clear(self):
+        def clear_queue(q):
+            while not q.empty():
+                try:
+                    q.get(False)
+                except queue.Empty:
+                    print("Emptying an empty queue. Is the graph still executing?")
+                    break
+
+        for q in self._load_queues:
+            clear_queue(q)
+        for q in self._dump_queues:
+            clear_queue(q)
+        for q in self._load_result_queues:
+            clear_queue(q)
+        for q in self._dump_result_queues:
+            clear_queue(q)
+        for q in self.consumer_load_queues:
+            clear_queue(q)
+        for q in self.consumer_dump_queues:
+            clear_queue(q)
+        self.total_consumer_size = 0
+
+    def put_load(
+        self, items: list, record_id: int, load_fn: callable, consumer_id: int
+    ):
+        for i, item in enumerate(items):
+            self._load_queues[self._worker_queue_cycle].put(
+                (item, i, record_id, load_fn, consumer_id), block=False
+            )
+            self._worker_queue_cycle = (self._worker_queue_cycle + 1) % self.num_workers
+
+    def get_load(self, consumer_id):
+        self._handle_queues()
+        if consumer_id not in self.consumer_load_queues:
+            return None
+        try:
+            result, record_id = self.consumer_load_queues[consumer_id].get(False)
+            self.total_consumer_size -= 1
+            if result is None:
+                return None, record_id
+            t, index = result
+            t_clone = t.clone()
+            del t
+            return (t_clone, index), record_id
+        except queue.Empty:
+            return None
+
+    def put_dump(
+        self,
+        data: any,
+        item_key: str,
+        record_id: int,
+        dump_fn: callable,
+        consumer_id: int,
+    ):
+        self._dump_queues[self._worker_queue_cycle].put(
+            (data, item_key, record_id, dump_fn, consumer_id), block=False
+        )
+        self._worker_queue_cycle = (self._worker_queue_cycle + 1) % self.num_workers
+
+    def get_dump(self, consumer_id):
+        self._handle_queues()
+        if consumer_id not in self.consumer_dump_queues:
+            return None
+        try:
+            record_id = self.consumer_dump_queues[consumer_id].get(False)
+            self.total_consumer_size -= 1
+            return record_id
+        except queue.Empty:
+            return None

graphbook-0.4.0/graphbook/exports.py

@@ -0,0 +1,98 @@
+import graphbook.steps as steps
+import graphbook.resources.base as rbase
+import graphbook.custom_nodes as custom_nodes
+from aiohttp import web
+
+default_exported_steps = {
+    "Split": steps.Split,
+    "SplitNotesByItems": steps.SplitNotesByItems,
+    "SplitItemField": steps.SplitItemField,
+    "DumpJSONL": steps.DumpJSONL,
+    "LoadJSONL": steps.LoadJSONL,
+}
+
+default_exported_resources = {
+    "Text": rbase.Resource,
+    "Number": rbase.NumberResource,
+    "Function": rbase.FunctionResource,
+}
+
+
+class NodeHub:
+    def __init__(self, path):
+        self.exported_steps = default_exported_steps
+        self.exported_resources = default_exported_resources
+        self.custom_node_importer = custom_nodes.CustomNodeImporter(
+            path, self.handle_step, self.handle_resource
+        )
+
+    def start(self):
+        self.custom_node_importer.start_observer()
+
+    def stop(self):
+        self.custom_node_importer.stop_observer()
+
+    async def handle_step(self, filename, name, step):
+        if name in self.exported_steps:
+            print(f"Reloading custom step node {name} from {filename}")
+        else:
+            print(f"Loading custom step node {name} from {filename}")
+        self.exported_steps[name] = step
+
+    async def handle_resource(self, filename, name, resource):
+        if name in self.exported_resources:
+            print(f"Reloading custom resource node {name} from {filename}")
+        else:
+            print(f"Loading custom resource node {name} from {filename}")
+        self.exported_resources[name] = resource
+
+    def get_steps(self):
+        return self.exported_steps
+
+    def get_resources(self):
+        return self.exported_resources
+
+    def get_all(self):
+        return {"steps": self.get_steps(), "resources": self.get_resources()}
+
+    def get_exported_nodes(self):
+        # Create directory structure for nodes based on their category
+        def create_dir_structure(nodes):
+            node_tree = {}
+            for node_name in nodes:
+                node = nodes[node_name]
+                category_tree = node["category"].split("/")
+                curr_category = node_tree
+                for category in category_tree:
+                    if curr_category.get(category) is None:
+                        curr_category[category] = {"children": {}}
+                    curr_category = curr_category[category]["children"]
+                curr_category[node_name] = node
+            return node_tree
+
+        steps = {
+            k: {
+                "name": k,
+                "parameters": v.Parameters,
+                "inputs": ["in"] if v.RequiresInput else [],
+                "outputs": v.Outputs,
+                "category": v.Category,
+            }
+            for k, v in self.get_steps().items()
+        }
+        resources = {
+            k: {
+                "name": k,
+                "parameters": v.Parameters,
+                "category": v.Category,
+            }
+            for k, v in self.get_resources().items()
+        }
+
+        return {
+            "steps": create_dir_structure(steps),
+            "resources": create_dir_structure(resources),
+        }
+
+    def set_websocket(self, websocket: web.WebSocketResponse):
+        self.custom_node_importer.set_websocket(websocket)