resolve the CUDA limit issue

Author: hadilaff

octopipes/benchmark.py

@@ -1,4 +1,4 @@
-from multiprocess import Pool
+from concurrent.futures import ThreadPoolExecutor, as_completed
 
 from tqdm import tqdm
 
@@ -9,7 +9,12 @@
 
 class Benchmark:
     def __init__(self, dataloader: Dataloader, workflows: list[Workflow],
-                 flows_factory: AggregateFlowsFactory | None = None) -> None:
+                 flows_factory: AggregateFlowsFactory | None = None, mode: str = "threaded") -> None:
+
+        """
+        :param str mode: "single" for sequential execution, "threaded" for multithreading.
+        """
+
         self.dataloader = dataloader
         self.workflows = workflows
         self.results: list[AggregateFlows] = []
@@ -28,16 +33,30 @@ def run_sample(factory: AggregateFlowsFactory, workflows, sample):
         return aggregate
 
     def run_tests(self):
-        for batch in tqdm(self.dataloader):
-            with Pool(processes=len(batch)) as pool:
-                for result in pool.map(func=Run(self.factory, self.workflows), iterable=batch):
-                    self.results.append(result)
-
-
-class Run:
-    def __init__(self, factory, workflows) -> None:
-        self.factory = factory
-        self.workflows = workflows
+        if self.mode == "single":
+            self._run_sequential()
+        elif self.mode == "threaded":
+            self._run_multithreaded()
+        else:
+            raise ValueError(f"Invalid model: {self.mode}. Choose 'single' or 'threaded' ")       
+         
+
+    def _run_sequential(self):
+        """Runs workflows sequentially (single-threaded execution)."""
+        for batch in tqdm(self.dataloader, desc= "processing batches"):
+            for sample in batch:
+                result=self.run_sample(self.factory, self.workflows, sample)
+                self.results.append(result)
+
+    def _run_multithreaded(self):
+        """Runs workflows using multithreading (threaded execution)."""
+        for batch in tqdm(self.dataloader,desc= "processing batches"):
+            with ThreadPoolExecutor(max_workers=len(batch)) as executor:
+                future_to_sample= {executor.submit(self.run_sample,self.factory, self.workflows, sample): sample for sample in batch}   
+                for future in as_completed(future_to_sample):
+                    try:
+                        result=future.result()
+                        self.results.append(result)
+                    except Exception as e:
+                        print(f"Error processing sample: {e}")    
 
-    def __call__(self, sample):
-        return Benchmark.run_sample(self.factory, self.workflows, sample)