code-loader 0.2.84.dev23__tar.gz → 0.2.84.dev31__tar.gz

{code_loader-0.2.84.dev23 → code_loader-0.2.84.dev31}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: code-loader
-Version: 0.2.84.dev23
+Version: 0.2.84.dev31
 Summary: 
 Home-page: https://github.com/tensorleap/code-loader
 License: MIT
@@ -12,7 +12,6 @@ Classifier: Programming Language :: Python :: 3
 Classifier: Programming Language :: Python :: 3.8
 Classifier: Programming Language :: Python :: 3.9
 Classifier: Programming Language :: Python :: 3.10
-Requires-Dist: multiprocess (>=0.70.15,<0.71.0)
 Requires-Dist: numpy (>=1.22.3,<2.0.0)
 Requires-Dist: psutil (>=5.9.5,<6.0.0)
 Requires-Dist: tensorflow (>=2.11.0,<3.0.0) ; platform_machine == "x86_64"

{code_loader-0.2.84.dev23 → code_loader-0.2.84.dev31}/code_loader/leap_loader_parallelized_base.py

@@ -11,11 +11,10 @@ import psutil
 
 class LeapLoaderParallelizedBase(ABC):
     def __init__(self, code_path: str, code_entry_name: str,
-                 n_workers: Optional[int] = 2, max_ready_results_in_queue: int = 128, bb = "spawn") -> None:
-        self.bb = bb
-        self.mp = multiprocessing
-        if self.bb is not None:
-            self.mp = multiprocessing.get_context(self.bb)
+                 n_workers: Optional[int] = 2, max_ready_results_in_queue: int = 128,
+                 multiprocessing_context: Optional[str] = None) -> None:
+        self.multiprocessing_context = multiprocessing.get_context(multiprocessing_context)
+
         self.code_entry_name = code_entry_name
         self.code_path = code_path
 
@@ -36,7 +35,7 @@ class LeapLoaderParallelizedBase(ABC):
         memory_usage_in_bytes = p.memory_info().rss
         total_memory_in_bytes = psutil.virtual_memory().total
 
-        n_workers = min(int(multiprocessing.cpu_count()),
+        n_workers = min(int(self.multiprocessing_context.cpu_count()),
                         int(total_memory_in_bytes * 0.7 / memory_usage_in_bytes))
         n_workers = max(n_workers, 1)
         return n_workers
@@ -50,9 +49,8 @@ class LeapLoaderParallelizedBase(ABC):
 
     @lru_cache()
     def start(self) -> None:
-            # multiprocessing.set_start_method(self.bb, force=True)
-        self._inputs_waiting_to_be_process = self.mp.Queue(5000)
-        self._ready_processed_results = self.mp.Queue(self.max_ready_results_in_queue)
+        self._inputs_waiting_to_be_process = self.multiprocessing_context.Queue(5000)
+        self._ready_processed_results = self.multiprocessing_context.Queue(self.max_ready_results_in_queue)
 
         self._run_and_warm_first_process()
         n_workers = self.n_workers

code_loader-0.2.84.dev31/code_loader/metric_calculator_parallelized.py

@@ -0,0 +1,50 @@
+# mypy: ignore-errors
+from typing import Optional, List, Tuple, Dict
+from multiprocessing import Process, Queue
+from code_loader.leap_loader_parallelized_base import LeapLoaderParallelizedBase
+import traceback
+from dataclasses import dataclass
+import tensorflow as tf
+from code_loader.leaploader import LeapLoader
+
+
+@dataclass
+class MetricSerializableError:
+    metric_id: str
+    metric_name: str
+    leap_script_trace: str
+    exception_as_str: str
+
+
+class MetricCalculatorParallelized(LeapLoaderParallelizedBase):
+    def __init__(self, code_path: str, code_entry_name: str, n_workers: Optional[int] = 2,
+                 max_samples_in_queue: int = 128) -> None:
+        super().__init__(code_path, code_entry_name, n_workers, max_samples_in_queue, "spawn")
+
+    @staticmethod
+    def _process_func(code_path: str, code_entry_name: str,
+                      metrics_to_process: Queue, ready_samples: Queue) -> None:
+        leap_loader = LeapLoader(code_path, code_entry_name)
+        while True:
+            metric_id, metric_name, input_arg_name_to_tensor = metrics_to_process.get(block=True)
+            try:
+                with tf.device('/cpu:0'):
+                    metric_result = leap_loader.metric_by_name()[metric_name].function(**input_arg_name_to_tensor)
+            except Exception as e:
+                leap_script_trace = traceback.format_exc().split('File "<string>"')[-1]
+                ready_samples.put(MetricSerializableError(metric_id, metric_name, leap_script_trace, str(e)))
+                continue
+
+            ready_samples.put((metric_id, metric_result))
+
+    def _create_and_start_process(self) -> Process:
+        process = self.multiprocessing_context.Process(
+            target=MetricCalculatorParallelized._process_func,
+            args=(self.code_path, self.code_entry_name, self._inputs_waiting_to_be_process,
+                  self._ready_processed_results))
+        process.daemon = True
+        process.start()
+        return process
+
+    def calculate_metrics(self, input_arg_name_to_tensor_list: List[Tuple[str, str, Dict[str, tf.Tensor]]]):
+        return self.start_process_inputs(input_arg_name_to_tensor_list)

{code_loader-0.2.84.dev23 → code_loader-0.2.84.dev31}/code_loader/samples_generator_parallelized.py

@@ -2,9 +2,8 @@
 import traceback
 from dataclasses import dataclass
 
-from typing import List, Tuple, Optional, Dict
+from typing import List, Tuple, Optional
 
-import multiprocessing
 from multiprocessing import Process, Queue
 
 from code_loader.leap_loader_parallelized_base import LeapLoaderParallelizedBase
@@ -23,10 +22,10 @@ class SampleSerializableError:
 class SamplesGeneratorParallelized(LeapLoaderParallelizedBase):
     def __init__(self, code_path: str, code_entry_name: str, n_workers: Optional[int] = 2,
                  max_samples_in_queue: int = 128) -> None:
-        super().__init__(code_path, code_entry_name, n_workers, max_samples_in_queue, None)
+        super().__init__(code_path, code_entry_name, n_workers, max_samples_in_queue)
 
     def _create_and_start_process(self) -> Process:
-        process = self.mp.Process(
+        process = self.multiprocessing_context.Process(
             target=SamplesGeneratorParallelized._process_func,
             args=(self.code_path, self.code_entry_name, self._inputs_waiting_to_be_process,
                   self._ready_processed_results))

{code_loader-0.2.84.dev23 → code_loader-0.2.84.dev31}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "code-loader"
-version = "0.2.84.dev23"
+version = "0.2.84.dev31"
 description = ""
 authors = ["dorhar <doron.harnoy@tensorleap.ai>"]
 license = "MIT"
@@ -19,7 +19,6 @@ tensorflow-macos = {version = "^2.11.0", markers = "platform_machine  == 'arm64'
 typeguard = "^2.13.3"
 psutil = "^5.9.5"
 torch = "1.12.1"
-multiprocess = "^0.70.15"
 
 [tool.poetry.dev-dependencies]
 pytest = "^7.1.1"

code_loader-0.2.84.dev23/code_loader/bla.py

@@ -1,33 +0,0 @@
-# mypy: ignore-errors
-import traceback
-from dataclasses import dataclass
-from multiprocess import Queue
-import tensorflow as tf
-
-from code_loader.leaploader import LeapLoader
-
-
-@dataclass
-class MetricSerializableError:
-    metric_id: str
-    metric_name: str
-    leap_script_trace: str
-    exception_as_str: str
-
-
-def _process_func(code_path: str, code_entry_name: str,
-                  metrics_to_process: Queue, ready_samples: Queue) -> None:
-    import os
-    os.environ["CUDA_VISIBLE_DEVICES"] = "-1"
-    leap_loader = LeapLoader(code_path, code_entry_name)
-    while True:
-        metric_id, metric_name, input_arg_name_to_tensor = metrics_to_process.get(block=True)
-        try:
-            with tf.device('/cpu:0'):
-                metric_result = leap_loader.metric_by_name()[metric_name].function(**input_arg_name_to_tensor)
-        except Exception as e:
-            leap_script_trace = traceback.format_exc().split('File "<string>"')[-1]
-            ready_samples.put(MetricSerializableError(metric_id, metric_name, leap_script_trace, str(e)))
-            continue
-
-        ready_samples.put((metric_id, metric_result))

code_loader-0.2.84.dev23/code_loader/metric_calculator_parallelized.py

@@ -1,26 +0,0 @@
-# mypy: ignore-errors
-import multiprocessing
-from typing import Optional, List, Tuple, Dict
-from multiprocessing import Process
-
-import tensorflow as tf
-
-from code_loader.bla import _process_func
-from code_loader.leap_loader_parallelized_base import LeapLoaderParallelizedBase
-
-class MetricCalculatorParallelized(LeapLoaderParallelizedBase):
-    def __init__(self, code_path: str, code_entry_name: str, n_workers: Optional[int] = 2,
-                 max_samples_in_queue: int = 128) -> None:
-        super().__init__(code_path, code_entry_name, n_workers, max_samples_in_queue)
-
-    def _create_and_start_process(self) -> Process:
-        process = self.mp.Process(
-            target=_process_func,
-            args=(self.code_path, self.code_entry_name, self._inputs_waiting_to_be_process,
-                  self._ready_processed_results))
-        process.daemon = True
-        process.start()
-        return process
-
-    def calculate_metrics(self, input_arg_name_to_tensor_list: List[Tuple[str, str, Dict[str, tf.Tensor]]]):
-        return self.start_process_inputs(input_arg_name_to_tensor_list)