Improve serialization

src/great_ai/great_ai/models/load_model.py

@@ -1,6 +1,6 @@
 from typing import Any, Optional, Tuple
 
-from joblib import load
+from dill import load
 
 from ..context import get_context
 

src/great_ai/great_ai/models/save_model.py

@@ -1,7 +1,7 @@
 from pathlib import Path
 from typing import Optional, Union
 
-from joblib import dump
+from dill import dump
 
 from ..context import get_context
 

src/great_ai/utilities/parallel_map/get_config.py

@@ -2,8 +2,6 @@ import os
 from math import ceil
 from typing import Callable, Iterable, Optional, Sequence, Union
 
-import dill
-
 from ..logger import get_logger
 from .parallel_map_configuration import ParallelMapConfiguration
 
@@ -52,7 +50,6 @@ def get_config(
         chunk_count=chunk_count,
         chunk_size=chunk_size,
         input_length=input_length,
-        serialized_map_function=dill.dumps(function, byref=True, recurse=True),
         function_name=function.__name__,
     )
 

src/great_ai/utilities/parallel_map/mapper_function.py

@@ -3,7 +3,7 @@ import queue
 import threading
 import traceback
 from time import sleep
-from typing import Union
+from typing import Callable, Union
 
 import dill
 
@@ -12,10 +12,11 @@ def mapper_function(
     input_queue: Union[mp.Queue, queue.Queue],
     output_queue: Union[mp.Queue, queue.Queue],
     should_stop: Union[mp.Event, threading.Event],
-    serialized_map_function: bytes,
+    map_function: Union[bytes, Callable],
 ) -> None:
     try:
-        map_function = dill.loads(serialized_map_function)
+        if isinstance(map_function, bytes):
+            map_function = dill.loads(map_function)
 
         last_chunk = None
         while not should_stop.is_set():

src/great_ai/utilities/parallel_map/parallel_map.py

@@ -1,6 +1,8 @@
 import multiprocessing as mp
 from typing import Callable, Iterable, Literal, Optional, Sequence, TypeVar, overload
 
+import dill
+
 from .get_config import get_config
 from .manage_communication import manage_communication
 from .manage_serial import manage_serial
@@ -92,6 +94,8 @@ def parallel_map(
     output_queue = ctx.Queue(config.concurrency * 2)
     should_stop = ctx.Event()
 
+    serialized_map_function = dill.dumps(function, byref=True, recurse=True)
+
     processes = [
         ctx.Process(
             name=f"parallel_map_{config.function_name}_{i}",
@@ -100,7 +104,7 @@ def parallel_map(
                 input_queue=input_queue,
                 output_queue=output_queue,
                 should_stop=should_stop,
-                serialized_map_function=config.serialized_map_function,
+                map_function=serialized_map_function,
             ),
         )
         for i in range(config.concurrency)

src/great_ai/utilities/parallel_map/parallel_map_configuration.py

@@ -12,5 +12,4 @@ class ParallelMapConfiguration(BaseModel):
     chunk_count: Optional[int]
     chunk_size: int
     input_length: Optional[int]
-    serialized_map_function: bytes
     function_name: str

src/great_ai/utilities/parallel_map/threaded_parallel_map.py

@@ -99,7 +99,7 @@ def threaded_parallel_map(
                 input_queue=input_queue,
                 output_queue=output_queue,
                 should_stop=should_stop,
-                serialized_map_function=config.serialized_map_function,
+                map_function=function,
             ),
         )
         for i in range(config.concurrency)

src/great_ai/utilities/unchunk.py

@@ -1,8 +1,9 @@
-from typing import Iterable, TypeVar
+from typing import Iterable, Optional, TypeVar
 
 T = TypeVar("T")
 
 
-def unchunk(chunks: Iterable[Iterable[T]]) -> Iterable[T]:
+def unchunk(chunks: Iterable[Optional[Iterable[T]]]) -> Iterable[T]:
     for chunk in chunks:
-        yield from chunk
+        if chunk is not None:
+            yield from chunk