Merge pull request #38 from Photoroom/ben/ml-2347-cleaner-way-to-pass…

…-bytes-buffers-around

[gopy] Cleaner bytes buffer transfer

cmd/main.go

@@ -14,7 +14,7 @@ func main() {
 	// Define flags
 	config := datago.GetDatagoConfig()
 
-	sourceConfig := datago.SourceFileSystemConfig{RootPath: os.Getenv("DATAROOM_TEST_FILESYSTEM")}
+	sourceConfig := datago.SourceFileSystemConfig{RootPath: os.Getenv("DATAGO_TEST_FILESYSTEM")}
 	sourceConfig.PageSize = 10
 	sourceConfig.Rank = 0
 	sourceConfig.WorldSize = 1

pkg/architecture.go

@@ -4,9 +4,8 @@ import "context"
 
 // --- Sample data structures - these will be exposed to the Python world ---------------------------------------------------------------------------------------------------------------------------------------------------------------
 type LatentPayload struct {
-	Data    []byte
-	Len     int
-	DataPtr uintptr
+	Data []byte
+	Len  int
 }
 
 type ImagePayload struct {
@@ -16,7 +15,6 @@ type ImagePayload struct {
 	Height         int // Useful to decode the current payload
 	Width          int
 	Channels       int
-	DataPtr        uintptr
 }
 
 type Sample struct {

pkg/client.go

@@ -345,7 +345,6 @@ func (c *DatagoClient) asyncDispatch() {
 			for _, item := range page.samplesDataPointers {
 				select {
 				case <-c.context.Done():
-					fmt.Println("Metadata fetch goroutine wrapping up")
 					close(c.chanSampleMetadata)
 					return
 				case c.chanSampleMetadata <- item:

pkg/generator_db.go

@@ -261,7 +261,6 @@ func (f datagoGeneratorDB) generatePages(ctx context.Context, chanPages chan Pag
 	for {
 		select {
 		case <-ctx.Done():
-			fmt.Println("Pages fetch goroutine wrapping up")
 			return
 
 		default:

pkg/generator_filesystem.go

@@ -28,7 +28,7 @@ func (c *SourceFileSystemConfig) setDefaults() {
 	c.Rank = 0
 	c.WorldSize = 1
 
-	c.RootPath = os.Getenv("DATAROOM_TEST_FILESYSTEM")
+	c.RootPath = os.Getenv("DATAGO_TEST_FILESYSTEM")
 }
 
 func GetSourceFileSystemConfig() SourceFileSystemConfig {

pkg/serdes.go

@@ -117,7 +117,6 @@ func imageFromBuffer(buffer []byte, transform *ARAwareTransform, aspect_ratio fl
 		Height:         height,
 		Width:          width,
 		Channels:       channels,
-		DataPtr:        dataPtrFromSlice(img_bytes),
 	}
 
 	return &img_payload, aspect_ratio, nil
@@ -243,7 +242,6 @@ func fetchSample(config *SourceDBConfig, http_client *http.Client, sample_result
 			latents[latent.LatentType] = LatentPayload{
 				latent_payload.content,
 				len(latent_payload.content),
-				dataPtrFromSlice(latent_payload.content),
 			}
 		}
 	}

pkg/utils.go

@@ -2,16 +2,8 @@ package datago
 
 import (
 	"time"
-	"unsafe"
 )
 
-func dataPtrFromSlice(a []uint8) uintptr {
-	if len(a) == 0 {
-		return 0
-	}
-	return uintptr(unsafe.Pointer(&a[0]))
-}
-
 func exponentialBackoffWait(retries int) {
 	baseDelay := time.Second
 	maxDelay := 64 * time.Second

python/benchmark_db.py

@@ -5,8 +5,9 @@
 from go_types import go_array_to_pil_image, go_array_to_numpy
 import typer
 
+
 def benchmark(
-    source: str = typer.Option("SOURCE", help="The source to test out"),
+    source: str = typer.Option("DATAGO_TEST_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"

python/benchmark_filesystem.py

@@ -9,7 +9,7 @@
 
 def benchmark(
     root_path: str = typer.Option(
-        os.getenv("DATAROOM_TEST_FILESYSTEM", ""), help="The source to test out"
+        os.getenv("DATAGO_TEST_FILESYSTEM", ""), 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(
@@ -35,8 +35,8 @@ def benchmark(
             "max_aspect_ratio": 2.0,
             "pre_encode_images": False,
         },
-        "prefetch_buffer_size": 256,
-        "samples_buffer_size": 128,
+        "prefetch_buffer_size": concurrency * 2,
+        "samples_buffer_size": concurrency * 2,
         "concurrency": concurrency,
         "limit": limit,
     }

python/go_types.py

@@ -3,13 +3,10 @@
 import io
 from typing import Optional
 import numpy as np
+from datago import go
 
 
 def uint8_array_to_numpy(go_array):
-    if go_array.DataPtr == 0:
-        print("Error: null pointer")
-        return np.array([], dtype=np.uint8)
-
     # By convention, arrays which are already serialized as jpg or png are not reshaped
     # We export them from Go with a Channels dimension of -1 to mark them as dimensionless.
     # Anything else is a valid number of channels and will thus lead to a reshape
@@ -27,25 +24,15 @@ def uint8_array_to_numpy(go_array):
 
     # Wrap the buffer around to create a numpy array. Strangely, shape needs to be passed twice
     # This is a zero-copy operation
-    return np.ctypeslib.as_array(
-        (ctypes.c_uint8 * length).from_address(go_array.DataPtr), shape=shape
-    ).reshape(shape)
+    bytes_buffer = bytes(go.Slice_byte(go_array.Data))
+    return np.frombuffer(bytes_buffer, dtype=np.uint8).reshape(shape)
 
 
 def go_array_to_numpy(go_array) -> Optional[np.ndarray]:
     # Generic numpy-serialized array
-
-    if go_array.DataPtr == 0:
-        # No data in the buffer
-        return None
-
-    np_bytes = np.ctypeslib.as_array(
-        (ctypes.c_uint8 * go_array.Len).from_address(go_array.DataPtr),
-        shape=(go_array.Len,),
-    )
-    bytes_io = io.BytesIO(np_bytes.tobytes())
+    bytes_buffer = bytes(go.Slice_byte(go_array.Data))
     try:
-        return np.load(bytes_io, allow_pickle=False)
+        return np.load(bytes_buffer, allow_pickle=False)
     except ValueError:
         # Do not try to handle these, return None and we'll handle it in the caller
         print("Could not deserialize numpy array")
@@ -54,11 +41,6 @@ def go_array_to_numpy(go_array) -> Optional[np.ndarray]:
 
 def go_array_to_pil_image(go_array):
     # Zero copy conversion of the image buffer from Go to PIL.Image
-
-    if go_array.DataPtr == 0:
-        # No data in the image buffer
-        return None
-
     np_array = uint8_array_to_numpy(go_array)
     if go_array.Channels <= 0:
         # Do not try to decode, we have a jpg or png buffer already

python/test_datago_db.py

@@ -6,7 +6,7 @@
 
 
 def get_test_source() -> str:
-    test_source = os.getenv("DATAROOM_TEST_SOURCE", "COYO")
+    test_source = os.getenv("DATAROOM_TEST_SOURCE")
     assert test_source is not None, "Please set DATAROOM_TEST_SOURCE"
     return test_source