Decent starting point

.github/workflows/go.yml

@@ -0,0 +1,38 @@
+# This workflow will build a golang project
+# For more information see: https://docs.github.com/en/actions/automating-builds-and-tests/building-and-testing-go
+
+name: Go
+
+on:
+  push:
+    branches: [ "main" ]
+  pull_request:
+    branches: [ "main" ]
+
+jobs:
+
+  build:
+    runs-on: ubuntu-latest
+    steps:
+    - uses: actions/checkout@v4
+
+    - name: Set up Go
+      uses: actions/setup-go@v4
+      with:
+        go-version: '1.20'
+
+    - name: Install linux deps
+      run: |
+        sudo apt-get update
+        sudo apt-get -y install libvips-dev
+
+    - name: Build
+      run: cd src/cmd/main && go build -v main.go
+
+    - name: Test
+      env:
+        DATAROOM_API_KEY: ${{ secrets.DATAROOM_API_KEY }}
+        DATAROOM_TEST_SOURCE: ${{ secrets.DATAROOM_TEST_SOURCE }}
+        DATAROOM_API_URL: ${{ secrets.DATAROOM_API_URL }}
+
+      run: cd src/tests && go test -v .

MANIFEST.in

@@ -0,0 +1,5 @@
+include *.md
+include *.py
+include datago/*.h
+include datago/*.c
+include datago/*.so

README.md

@@ -1,2 +1,161 @@
-# datago
-A golang-based data loader which can be used from Python. Useful to interact with a typical VectorDB stack for machine learning purposes
+[![Build & Test](https://github.com/Photoroom/datago/actions/workflows/go.yml/badge.svg)](https://github.com/Photoroom/datago/actions/workflows/go.yml)
+
+datago
+======
+
+A golang-based data loader which can be used from Python. Compatible with a soon-to-be open sourced VectorDB-enabled data stack, which exposes HTTP requests.
+
+Datago will handle, outside of the Python GIL
+- per sample IO from object storage
+- deserialization
+- some optional vision processing (aligning different image payloads)
+- serialization
+
+Samples are then exposed in the Python scope and ready for consumption, typically using PIL and Numpy base types.
+Speed will be network dependent, but GB/s is relatively easily possible
+
+Datago can be rank and world-size aware, in which case the samples are dispatched depending on the samples hash.
+
+<img width="922" alt="Screenshot 2024-09-24 at 9 39 44 PM" src="https://github.com/user-attachments/assets/b58002ce-f961-438b-af72-9e1338527365">
+
+
+<details> <summary><strong>Use it</strong></summary>
+
+Use the package from Python
+---------------------------
+
+```python
+from datago import datago
+
+# source, has/lacks attributes, has/lacks masks, has/lacks latents, metadata prefetch, sample prefetch, concurrent download
+client = datago.GetClient(
+            source="SOURCE",
+            require_images=True,
+            has_attributes="",
+            lacks_attributes="",
+            has_masks="",
+            lacks_masks="",
+            has_latents="",
+            lacks_latents="",
+            crop_and_resize=True,
+            prefetch_buffer_size=64,
+            samples_buffer_size=64,
+            downloads_concurrency=64,
+        )
+
+client.Start()  # This can be done early for convenience, not mandatory (can fetch samples while models are instanciated for intance)
+
+for _ in range(10):
+    sample = client.GetSample() # This start the client if not previously done, in that case latency for the first sample is higher
+```
+
+Please note that the image buffers will be passed around as raw pointers, they can be re-interpreted in python with the attached helpers
+
+
+Match the raw exported buffers with typical python types
+--------------------------------------------------------
+
+See helper functions provided in `polyglot.py`, should be self explanatory
+
+</details><details> <summary><strong>Build it</strong></summary>
+
+Install deps
+------------
+
+```bash
+$ sudo apt install golang libjpeg-turbo8-dev libvips-dev
+$ sudo ldconfig
+```
+
+Build a benchmark CLI
+---------------------
+
+From the root of this project `datago_src`:
+
+```bash
+$ go build cmd/main/main.go
+```
+
+Running it:
+
+```bash
+$ ./main --help` will tell you all about it
+```
+
+Running it with additional sanity checks
+
+```bash
+$ go run -race cmd/main/main.go
+```
+
+Run the go test suite
+---------------------
+
+From the src folder
+
+```bash
+$ go test -v tests/client_test.go
+```
+
+Refresh the python package and its binaries
+-------------------------------------------
+
+- Install the dependencies as detailed in the next point
+- Run the `generate_python_package.sh` script
+
+Generate the python package binaries manually
+---------------------------------------------
+
+```bash
+$ python3 -m pip install pybindgen
+$ go install golang.org/x/tools/cmd/goimports@latest
+$ go install github.com/go-python/gopy@latest
+$ go install golang.org/x/image/draw
+```
+
+NOTE:
+- you may need to add `~/go/bin` to your PATH so that gopy is found.
+- - Either `export PATH=$PATH:~/go/bin` or add it to your .bashrc
+- you may need this to make sure that LDD looks at the current folder `export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:.`
+
+then from the /pkg/client folder:
+
+```bash
+$ gopy pkg -author="Photoroom" -email="[email protected]" -url="" -name="datago" -version="0.0.1" .
+```
+
+then you can `pip install -e .` from here.
+
+
+Update the pypi release (maintainers)
+-------------------------------------
+```
+python3 setup.py sdist
+python3 -m twine upload dist/* --verbose
+```
+</details>
+
+
+License
+=======
+    MIT License
+
+    Copyright (c) 2024 Photoroom
+
+    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.

generate_python_package.sh

@@ -0,0 +1,26 @@
+#!/usr/bin/zsh
+
+echo "Updating the datago binaries"
+
+# Get the current python version
+python_version=$(python3 --version 2>&1 | awk '{print $2}' | cut -d. -f1,2)
+echo "Building package for python" $python_version
+
+# Setup where the python package will be copied
+DESTINATION="../../../python_$python_version"
+rm -rf $DESTINATION
+
+# Build the python package via the gopy toolchain
+cd src/pkg/client
+gopy pkg -author="Photoroom" -email="[email protected]" -url="" -name="datago" -version="0.3" .
+mkdir -p $DESTINATION/datago
+mv datago/* $DESTINATION/datago/.
+mv setup.py $DESTINATION/.
+mv Makefile $DESTINATION/.
+mv README.md $DESTINATION/.
+rm LICENSE
+rm MANIFEST.in
+
+cd ../../..
+
+

pyproject.toml

@@ -0,0 +1,18 @@
+[project]
+name = "datago_blefaudeux"
+version = "0.0.1"
+authors = [
+  { name="Photoroom", email="[email protected]" },
+]
+description = "A high performance python module to access data ressources through HTTP, written in Golang"
+readme = "README.md"
+requires-python = "==3.11"
+classifiers = [
+    "Programming Language :: Python :: 3",
+    "License :: OSI Approved :: MIT License",
+    "Operating System :: OS Independent",
+]
+
+[project.urls]
+Homepage = "https://github.com/photoroom/datago"
+Issues = "https://github.com/photoroom/datago/issues"

src/benchmark.py

@@ -0,0 +1,80 @@
+from datago import datago  # type: ignore
+import time
+import typer
+from tqdm import tqdm
+import numpy as np
+from polyglot import go_array_to_pil_image, go_array_to_numpy
+
+
+def benchmark(
+    source: str = typer.Option("SOURCE", help="The source to test out"),
+    limit: int = typer.Option(2000, help="The number of samples to test on"),
+    crop_and_resize: bool = typer.Option(True, help="Crop and resize the images on the fly"),
+    require_images: bool = typer.Option(True, help="Request the original images"),
+    require_embeddings: bool = typer.Option(False, help="Request embeddings"),
+    test_masks: bool = typer.Option(True, help="Test masks"),
+    test_latents: bool = typer.Option(True, help="Test latents"),
+):
+    print(f"Running benchmark for {source} - {limit} samples")
+    client = datago.GetClient(
+        source=source,
+        require_images=require_images,
+        require_embeddings=require_embeddings,
+        has_attributes="",
+        lacks_attributes="",
+        has_masks="segmentation_mask" if test_masks else "",
+        lacks_masks="",
+        has_latents="masked_image,my_test_latents" if test_latents else "",
+        lacks_latents="",
+        crop_and_resize=crop_and_resize,
+        prefetch_buffer_size=256,
+        samples_buffer_size=256,
+        downloads_concurrency=64,
+    )
+    client.Start()
+    start = time.time()
+
+    # Make sure in the following that we compare apples to apples, meaning in that case
+    # that we materialize the payloads in the python scope in the expected format
+    # (PIL.Image for images and masks for instance, numpy arrays for latents)
+    img, mask, masked_image = None, None, None
+    for _ in tqdm(range(limit), dynamic_ncols=True):
+        sample = client.GetSample()
+        if sample.ID:
+            # Bring the masks and image to PIL
+            if hasattr(sample, "Image"):
+                img = go_array_to_pil_image(sample.Image)
+
+            if hasattr(sample, "Masks"):
+                for _, mask_buffer in sample.Masks.items():
+                    mask = go_array_to_pil_image(mask_buffer)
+
+            if hasattr(sample, "AdditionalImages") and "masked_image" in sample.AdditionalImages:
+                masked_image = go_array_to_pil_image(sample.AdditionalImages["masked_image"])
+
+            # Bring the latents to numpy
+            if hasattr(sample, "Latents"):
+                for _, latent_buffer in sample.Latents.items():
+                    _latents = go_array_to_numpy(latent_buffer)
+
+            # Bring the embeddings to numpy
+            if hasattr(sample, "CocaEmbedding"):
+                _embedding = np.array(sample.CocaEmbedding)
+
+    fps = limit / (time.time() - start)
+    print(f"FPS {fps:.2f}")
+    client.Stop()
+
+    # Save the last image as a test
+    if img is not None:
+        img.save("benchmark_last_image.png")
+
+    if mask is not None:
+        mask.save("benchmark_last_mask.png")
+
+    if masked_image is not None:
+        masked_image.save("benchmark_last_masked_image.png")
+
+
+if __name__ == "__main__":
+    typer.run(benchmark)