create a pipeline class with basic test coverage

pipeline_test.py → pipeline.py

@@ -1,11 +1,8 @@
 import asyncio
-import random
-import time
-
 
 class Job:
-    def __init__(self, id, data):
-        self.id = id
+    def __init__(self, data):
+        self._id = None
         self.data = data
 
 
@@ -31,6 +28,11 @@ class Node:
             self._jobs_dequeued += 1
             if self.sequential_node == False:
                 await self.process_job(job)
+                if self.output_queue is not None:
+                    await self.output_queue.put(job)
+                if self.job_sync is not None:
+                    self.job_sync.append(job)
+                self._jobs_processed += 1
             else:
                 # ensure that jobs are processed in order
                 self.buffer[job.id] = job
@@ -38,79 +40,65 @@ class Node:
                     job = self.buffer.pop(self.next_i)
                     await self.process_job(job)
                     self.next_i += 1
-            if self.output_queue is not None:
-                await self.output_queue.put(job)
-            if self.job_sync is not None:
-                self.job_sync.append(job)
-            self._jobs_processed += 1
-
-    async def process_job(self, job: Job):
-        raise NotImplementedError
+                    if self.output_queue is not None:
+                        await self.output_queue.put(job)
+                    if self.job_sync is not None:
+                        self.job_sync.append(job)
+                    self._jobs_processed += 1
 
-
-class Node1(Node):
     async def process_job(self, job: Job):
-        job.data += f' (processed by node 1, worker {self.worker_id})'
-
-
-class Node2(Node):
-    async def process_job(self, job: Job):
-        sleep_duration = 0.8 + 0.4 * random.random()
-        await asyncio.sleep(sleep_duration)
-        job.data += f' (processed by node 2, worker {self.worker_id})'
-
-
-class Node3(Node):
-    async def process_job(self, job: Job):
-        job.data += f' (processed by node 3, worker {self.worker_id})'
-        print(f'{job.id} - {job.data}')
-
-
-async def main():
-    node1_queue = asyncio.Queue()
-    node2_queue = asyncio.Queue()
-    node3_queue = asyncio.Queue()
-
-    num_jobs = 100
-    joe_source = [Job(i, "") for i in range(num_jobs)]
-    job_sync = []
-
-    # create the workers
-    num_workers = 5
-    node1_workers = [Node1(i + 1, node1_queue, node2_queue) for i in range(1)]
-    node2_workers = [Node2(i + 1, node2_queue, node3_queue) for i in range(num_workers)]
-    node3_workers = [Node3(i + 1, node3_queue, job_sync=job_sync, sequential_node=True) for i in range(1)]
-
-    # create tasks for the workers
-    tasks1 = [asyncio.create_task(worker.run()) for worker in node1_workers]
-    tasks2 = [asyncio.create_task(worker.run()) for worker in node2_workers]
-    tasks3 = [asyncio.create_task(worker.run()) for worker in node3_workers]
-
-    for job in joe_source:
-        await node1_queue.put(job)
-    # await input_queue.put(joe_source[0])
-
-    try:
-        while len(job_sync) < num_jobs:
-            # print(f"Waiting for jobs to finish... Job sync size: {len(job_sync)}, node1_queue size: {node1_queue.qsize()}, node2_queue size: {node2_queue.qsize()}, node3_queue size: {node3_queue.qsize()}")
-            await asyncio.sleep(0.1)
-    except asyncio.CancelledError:
-        print("Pipeline cancelled")
-    for task in tasks1:
-        task.cancel()
-    for task in tasks2:
-        task.cancel()
-    for task in tasks3:
-        task.cancel()
-    await asyncio.gather(*tasks1, *tasks2, *tasks3, return_exceptions=True)
-
-
-start_time = time.time()
-
-try:
-    asyncio.run(main())
-except KeyboardInterrupt:
-    print("Pipeline interrupted by user")
-
-end_time = time.time()
-print(f"Pipeline processed in {end_time - start_time} seconds.")
+        raise NotImplementedError()
+    
+class Pipeline:
+    def __init__(self):
+        self.input_queues = []
+        self.root_queue = None
+        # self.output_queues = []
+        # self.job_sysncs = []
+        self.nodes= []
+        self.node_workers = {}
+        self.tasks = []
+        self._job_id = 0
+
+    async def add_node(self, node: Node, num_workers=1, input_queue=None, output_queue=None, job_sync=None, sequential_node=False ):
+        # input_queue must not be None
+        if input_queue is None:
+            raise ValueError('input_queue is None')
+        # job_sync nodes must be sequential_nodes
+        if job_sync is not None and sequential_node == False:
+            raise ValueError('job_sync is not None and sequential_node is False')
+        # sequential_nodes should one have 1 worker
+        if sequential_node == True and num_workers != 1:
+            raise ValueError('sequentaial nodes can only have one node (sequential_node is True and num_workers is not 1)')
+        # output queue must not equal input_queue
+        if output_queue == input_queue:
+            raise ValueError('output_queue must not be the same as input_queue')        
+
+        node_name = node.__class__.__name__
+        if node_name not in self.nodes:
+            self.nodes.append(node_name)
+
+        # if input_queue is None then this is the root node
+        if len(self.input_queues) is 0:
+            self.root_queue = input_queue
+
+        self.input_queues.append(input_queue)
+
+        for i in range(num_workers):
+            worker_id = i
+            node_worker = node(worker_id, input_queue, output_queue, job_sync, sequential_node)
+            self.node_workers[node_name] = node_worker
+            task = asyncio.create_task(node_worker.run())
+            self.tasks.append(task)
+
+    async def enqueue_job(self, job: Job):
+        job.id = self._job_id
+        self._job_id += 1
+        await self.root_queue.put(job)
+
+    async def close(self):
+        for task in self.tasks:
+            task.cancel()
+        await asyncio.gather(*self.tasks, return_exceptions=True)
+
+        

tests/test_pipeline.py

@@ -0,0 +1,83 @@
+import asyncio
+import random
+import time
+import unittest
+import sys
+import os
+sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+
+from pipeline import Pipeline, Node, Job
+
+
+class Node1(Node):
+    async def process_job(self, job: Job):
+        job.data += f' (processed by node 1, worker {self.worker_id})'
+
+
+class Node2(Node):
+    async def process_job(self, job: Job):
+        sleep_duration = 0.08 + 0.04 * random.random()
+        await asyncio.sleep(sleep_duration)
+        job.data += f' (processed by node 2, worker {self.worker_id})'
+
+
+class Node3(Node):
+    async def process_job(self, job: Job):
+        job.data += f' (processed by node 3, worker {self.worker_id})'
+        print(f'{job.id} - {job.data}')
+
+
+class TestPipeline(unittest.TestCase):
+    def setUp(self):
+        pass
+
+    async def _test_pipeline_edge_cases(self):
+        # must have a input queue
+        with self.assertRaises(ValueError):
+            await self.pipeline.add_node(Node1, 1, None, None)
+        # too output queue must not equal from input queue
+        node1_queue = asyncio.Queue()
+        with self.assertRaises(ValueError):
+            await self.pipeline.add_node(Node1, 1, node1_queue, node1_queue)
+
+
+    async def _test_pipeline(self, num_jobs):
+        node1_queue = asyncio.Queue()
+        node2_queue = asyncio.Queue()
+        node3_queue = asyncio.Queue()
+        await self.pipeline.add_node(Node1, 1, node1_queue, node2_queue)
+        await self.pipeline.add_node(Node2, 5, node2_queue, node3_queue)
+        await self.pipeline.add_node(Node3, 1, node3_queue, job_sync=self.job_sync, sequential_node=True)
+        for i in range(num_jobs):
+            job = Job("")
+            await self.pipeline.enqueue_job(job)
+        while True:
+            if len(self.job_sync) == num_jobs:
+                break
+            await asyncio.sleep(0.1)
+        await self.pipeline.close()
+
+    def test_pipeline_edge_cases(self):
+        self.pipeline = Pipeline()
+        self.job_sync = []
+        asyncio.run(self._test_pipeline_edge_cases())
+
+
+    # def test_pipeline_keeps_order(self):
+    #     self.pipeline = Pipeline()
+    #     self.job_sync = []
+    #     num_jobs = 100
+    #     start_time = time.time()
+    #     asyncio.run(self._test_pipeline(num_jobs))
+    #     end_time = time.time()
+    #     print(f"Pipeline processed in {end_time - start_time} seconds.")
+    #     self.assertEqual(len(self.job_sync), num_jobs)
+    #     for i, job in enumerate(self.job_sync):
+    #         self.assertEqual(i, job.id)
+
+
+if __name__ == '__main__':
+    unittest.main()
+    # test = TestPipeline()
+    # test.setUp()
+    # test.test_pipeline()