code-loader 0.2.70.dev1__py3-none-any.whl → 0.2.71__py3-none-any.whl

code_loader/leaploader.py

@@ -98,13 +98,14 @@ class LeapLoader:
         }
 
     def get_sample(self, state: DataStateEnum, idx: int) -> DatasetSample:
-        self.exec_script()
-        sample = DatasetSample(inputs=self._get_inputs(state, idx),
-                               gt=None if state == DataStateEnum.unlabeled else self._get_gt(state, idx),
-                               metadata=self._get_metadata(state, idx),
-                               index=idx,
-                               state=state)
-        return sample
+        with tf.device('/cpu:0'):
+            self.exec_script()
+            sample = DatasetSample(inputs=self._get_inputs(state, idx),
+                                   gt=None if state == DataStateEnum.unlabeled else self._get_gt(state, idx),
+                                   metadata=self._get_metadata(state, idx),
+                                   index=idx,
+                                   state=state)
+            return sample
 
     def check_dataset(self) -> DatasetIntegParseResult:
         test_payloads: List[DatasetTestResultPayload] = []

code_loader/samples_generator_parallelized.py

@@ -0,0 +1,143 @@
+# mypy: ignore-errors
+
+import traceback
+from dataclasses import dataclass
+from functools import lru_cache
+from queue import Empty
+from threading import Thread
+from typing import List, Tuple, Optional
+from multiprocessing import Process, Queue
+
+from code_loader.leaploader import LeapLoader
+from code_loader.contract.enums import DataStateEnum
+
+
+@dataclass
+class GetSampleSerializableError:
+    state: DataStateEnum
+    index: int
+    leap_script_trace: str
+    exception_as_str: str
+
+
+class SamplesGeneratorParallelized:
+    def __init__(self, code_path: str, code_entry_name: str,
+                 n_workers: int = 2, max_samples_in_queue: int = 128) -> None:
+        self.code_entry_name = code_entry_name
+        self.code_path = code_path
+
+        if n_workers <= 0:
+            raise Exception("need at least one worker")
+        self.n_workers = n_workers
+        self.max_samples_in_queue = max_samples_in_queue
+
+        self._n_samples_to_process = 0
+        self._samples_to_process: Optional[Queue] = None
+        self._ready_samples: Optional[Queue] = None
+        self.processes: Optional[List[Process]] = None
+        self._generate_samples_thread: Optional[Thread] = None
+        self._should_stop_thread = False
+
+    @lru_cache()
+    def start(self) -> None:
+        self._samples_to_process = Queue(5000)
+        self._ready_samples = Queue(self.max_samples_in_queue)
+
+        self.processes = [
+            Process(target=SamplesGeneratorParallelized._process_func,
+                    args=(self.code_path, self.code_entry_name, self._samples_to_process, self._ready_samples))
+            for _ in range(self.n_workers)]
+
+        for proc in self.processes:
+            proc.daemon = True
+            proc.start()
+
+        # needed in order to make sure the preprocess func runs once in nonparallel
+        self._generate_samples([(DataStateEnum.training, 0)])
+        self._get_next_sample()
+
+    @staticmethod
+    def _process_func(code_path: str, code_entry_name: str,
+                      samples_to_process: Queue, ready_samples: Queue) -> None:
+        leap_loader = LeapLoader(code_path, code_entry_name)
+        leap_loader.exec_script()
+        while True:
+            state, idx = samples_to_process.get(block=True)
+            try:
+                sample = leap_loader.get_sample(state, idx)
+            except Exception as e:
+                leap_script_trace = traceback.format_exc().split('File "<string>"')[-1]
+                ready_samples.put(GetSampleSerializableError(state, idx, leap_script_trace, str(e)))
+                continue
+
+            ready_samples.put(sample)
+
+    def _generate_samples(self, sample_identities: List[Tuple[DataStateEnum, int]]):
+        assert self._samples_to_process is not None
+        assert self._ready_samples is not None
+
+        for sample in sample_identities:
+            if self._should_stop_thread:
+                break
+            self._n_samples_to_process += 1
+            self._samples_to_process.put(sample)
+
+    def _clear_queues(self):
+        if self._generate_samples_thread is not None:
+            self._should_stop_thread = True
+            try:
+                self._samples_to_process.get_nowait()
+                self._n_samples_to_process -= 1
+            except Empty:
+                pass
+            self._generate_samples_thread.join()
+        while not self._samples_to_process.empty():
+            try:
+                self._samples_to_process.get_nowait()
+                self._n_samples_to_process -= 1
+            except Empty:
+                pass
+
+        for _ in range(self._n_samples_to_process):
+            self._get_next_sample()
+
+        self._should_stop_thread = False
+
+    def _get_next_sample(self):
+        sample = self._ready_samples.get()
+        self._n_samples_to_process -= 1
+        return sample
+
+    def generate_samples(self, sample_identities: List[Tuple[DataStateEnum, int]]):
+        self.start()
+
+        self._clear_queues()
+
+        self._generate_samples_thread = Thread(target=self._generate_samples, args=(sample_identities,))
+        self._generate_samples_thread.start()
+        return self._get_next_sample
+
+    @staticmethod
+    def _release_queue(queue: Queue):
+        assert queue is not None
+        queue.close()
+        queue.join_thread()
+
+    def release(self) -> None:
+        if self.processes is None:
+            return
+        self._clear_queues()
+
+        self._release_queue(self._samples_to_process)
+        self._release_queue(self._ready_samples)
+
+        for process in self.processes:
+            process.terminate()
+            process.kill()
+            process.join()
+            process.close()
+
+        self.processes = None
+
+    def __del__(self) -> None:
+        self.release()

{code_loader-0.2.70.dev1.dist-info → code_loader-0.2.71.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: code-loader
-Version: 0.2.70.dev1
+Version: 0.2.71
 Summary: 
 Home-page: https://github.com/tensorleap/code-loader
 License: MIT
@@ -13,8 +13,8 @@ Classifier: Programming Language :: Python :: 3.8
 Classifier: Programming Language :: Python :: 3.9
 Classifier: Programming Language :: Python :: 3.10
 Requires-Dist: numpy (>=1.22.3,<2.0.0)
-Requires-Dist: tensorflow (==2.12.0); platform_machine == "x86_64"
-Requires-Dist: tensorflow-macos (==2.12.0); platform_machine == "arm64"
+Requires-Dist: tensorflow (>=2.11.0,<3.0.0) ; platform_machine == "x86_64"
+Requires-Dist: tensorflow-macos (>=2.11.0,<3.0.0) ; platform_machine == "arm64"
 Requires-Dist: typeguard (>=2.13.3,<3.0.0)
 Project-URL: Repository, https://github.com/tensorleap/code-loader
 Description-Content-Type: text/markdown

{code_loader-0.2.70.dev1.dist-info → code_loader-0.2.71.dist-info}/RECORD

@@ -17,13 +17,14 @@ code_loader/helpers/detection/yolo/pytorch_utils.py,sha256=RuB9tt13fIKuuGXoaTLOm
 code_loader/helpers/detection/yolo/utils.py,sha256=eV9-IZhglRNGud1UmSckIuA49WRKJitDQhPKK7iIaBc,6563
 code_loader/leap_binder/__init__.py,sha256=Oe7DNj6RmgX-tMruOrI7GYvf3SyAnzOxugF2e117_Z8,93
 code_loader/leap_binder/leapbinder.py,sha256=WjC-UA5xeY9UZtMqJQOxjidlFcWURC9mmy7OrQy1IY0,6692
-code_loader/leaploader.py,sha256=aM9u6GhhZhkN879OkDr-daLFS_Dd7IaELgA7joBBA_I,14982
+code_loader/leaploader.py,sha256=3vvzBzS0oev6kqXjtriz3G6z0tijGOJLIH8QoOvu1Ls,15044
 code_loader/metrics/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 code_loader/metrics/default_metrics.py,sha256=xtvTskR7LC5pyKSlFcvlyBpFY9TSXNGVAi5taEPSI60,6727
+code_loader/samples_generator_parallelized.py,sha256=3Q9EixQmZ7Pv1G0J3N_S0d3O2u1znVTpvxPi2iSp9gQ,4747
 code_loader/utils.py,sha256=WUcM97OuxrhfLCRPoH9EbXrxajNpYgX1CTMc3_PXtYU,1736
 code_loader/visualizers/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 code_loader/visualizers/default_visualizers.py,sha256=HqWx2qfTrroGl2n8Fpmr_4X-rk7tE2oGapjO3gzz4WY,2226
-code_loader-0.2.70.dev1.dist-info/LICENSE,sha256=qIwWjdspQeSMTtnFZBC8MuT-95L02FPvzRUdWFxrwJY,1067
-code_loader-0.2.70.dev1.dist-info/WHEEL,sha256=vxFmldFsRN_Hx10GDvsdv1wroKq8r5Lzvjp6GZ4OO8c,88
-code_loader-0.2.70.dev1.dist-info/METADATA,sha256=xBIskSlGlQ26Hr1xBfsqO4WGqLdsJ9g7mQvxE15SI4k,865
-code_loader-0.2.70.dev1.dist-info/RECORD,,
+code_loader-0.2.71.dist-info/LICENSE,sha256=qIwWjdspQeSMTtnFZBC8MuT-95L02FPvzRUdWFxrwJY,1067
+code_loader-0.2.71.dist-info/METADATA,sha256=JLp90ger6p7KRVnGmKZ1dUQht7euhaIQBNH9F5S50Z4,876
+code_loader-0.2.71.dist-info/WHEEL,sha256=7Z8_27uaHI_UZAc4Uox4PpBhQ9Y5_modZXWMxtUi4NU,88
+code_loader-0.2.71.dist-info/RECORD,,

{code_loader-0.2.70.dev1.dist-info → code_loader-0.2.71.dist-info}/WHEEL

@@ -1,4 +1,4 @@
 Wheel-Version: 1.0
-Generator: poetry-core 1.3.2
+Generator: poetry-core 1.5.2
 Root-Is-Purelib: true
 Tag: py3-none-any