Feature/multi processing (#134)

* ✨ Implement Scheduler for multi processing

- Add scheduler singleton to handle creating pools an closing pools. 
- Remove Pool references and replace scheduler
- Update tests and benchmarks

* Modify scheduler to include only one method

add tally methods to tally in batches

* fix naming

* fix unused inports

* return int

* increase the deadline on slow tests

* Address PR Feedback

Co-authored-by: Keith Fung <keith.fung@infernored.com>
Co-authored-by: Keith Fung <keithrfung@users.noreply.github.com>

Author: AddressXception

Makefile

@@ -1,4 +1,4 @@
-.PHONY: all environment install install-mac install-linux install-windows lint validate test test-example coverage coverage-html coverage-xml coverage-erase generate-sample-data
+.PHONY: all bench environment install install-mac install-linux install-windows lint validate test test-example coverage coverage-html coverage-xml coverage-erase generate-sample-data
 
 CODE_COVERAGE ?= 90
 WINDOWS_32BIT_GMPY2 ?= packages/gmpy2-2.0.8-cp38-cp38-win32.whl
@@ -8,6 +8,10 @@ IS_64_BIT ?= $(shell python -c 'from sys import maxsize; print(maxsize > 2**32)'
 
 all: environment install validate lint coverage
 
+bench:
+	@echo 📊 BENCHMARKS
+	pipenv run python -s bench/bench_chaum_pedersen.py
+
 environment:
 	@echo 🔧 PIPENV SETUP
 	pip install pipenv

Pipfile

@@ -22,6 +22,7 @@ typish = '*'
 [packages]
 numpy = '>=1.18.2'
 jsons = '>=1.1.2'
+psutil = '>=5.7.2'
 cryptography = "*"
 
 [requires]

Pipfile.lock

@@ -1,6 +1,5 @@
-from multiprocessing import Pool, cpu_count
 from timeit import default_timer as timer
-from typing import Tuple, NamedTuple, Dict
+from typing import Dict, List, NamedTuple, Tuple
 
 from numpy import average, std
 
@@ -13,7 +12,7 @@
 )
 from electionguard.group import ElementModQ, int_to_q_unchecked
 from electionguard.nonces import Nonces
-
+from electionguard.scheduler import Scheduler
 from electionguard.utils import get_optional
 
 
@@ -49,62 +48,64 @@ def identity(x: int) -> int:
     problem_sizes = (100, 500, 1000, 5000)
     rands = Nonces(int_to_q_unchecked(31337))
     speedup: Dict[int, float] = {}
-    print(f"CPUs detected: {cpu_count()}, launching thread pool")
-    pool = Pool(cpu_count())
 
     # warm up the pool to help get consistent measurements
-    results = pool.map(identity, range(1, 30000))
-    assert results == list(range(1, 30000))
-
-    bench_start = timer()
-
-    for size in problem_sizes:
-        print("Benchmarking on problem size: ", size)
-        seeds = rands[0:size]
-        inputs = [
-            BenchInput(
-                get_optional(elgamal_keypair_from_secret(a)),
-                rands[size],
-                rands[size + 1],
-            )
-            for a in seeds
-        ]
-        start_all_scalar = timer()
-        timing_data = [chaum_pedersen_bench(i) for i in inputs]
-        end_all_scalar = timer()
-
-        print(f"  Creating Chaum-Pedersen proofs ({size} iterations)")
-        avg_proof_scalar = average([t[0] for t in timing_data])
-        std_proof_scalar = std([t[0] for t in timing_data])
-        print(f"    Avg    = {avg_proof_scalar:.6f} sec")
-        print(f"    Stddev = {std_proof_scalar:.6f} sec")
-
-        print(f"  Validating Chaum-Pedersen proofs ({size} iterations)")
-        avg_verify_scalar = average([t[1] for t in timing_data])
-        std_verify_scalar = std([t[1] for t in timing_data])
-        print(f"    Avg    = {avg_verify_scalar:.6f} sec")
-        print(f"    Stddev = {std_verify_scalar:.6f} sec")
-
-        # Run in parallel
-        start_all_parallel = timer()
-        timing_data_parallel = pool.map(chaum_pedersen_bench, inputs)
-        end_all_parallel = timer()
-
-        speedup[size] = (end_all_scalar - start_all_scalar) / (
-            end_all_parallel - start_all_parallel
+    with Scheduler() as scheduler:
+        results: List[int] = scheduler.schedule(
+            identity, [list([x]) for x in range(1, 30000)]
         )
-        print(f"  Parallel speedup: {speedup[size]:.3f}x")
+        assert results == list(range(1, 30000))
+
+        bench_start = timer()
+
+        for size in problem_sizes:
+            print("Benchmarking on problem size: ", size)
+            seeds = rands[0:size]
+            inputs = [
+                BenchInput(
+                    get_optional(elgamal_keypair_from_secret(a)),
+                    rands[size],
+                    rands[size + 1],
+                )
+                for a in seeds
+            ]
+            start_all_scalar = timer()
+            timing_data = [chaum_pedersen_bench(i) for i in inputs]
+            end_all_scalar = timer()
+
+            print(f"  Creating Chaum-Pedersen proofs ({size} iterations)")
+            avg_proof_scalar = average([t[0] for t in timing_data])
+            std_proof_scalar = std([t[0] for t in timing_data])
+            print(f"    Avg    = {avg_proof_scalar:.6f} sec")
+            print(f"    Stddev = {std_proof_scalar:.6f} sec")
+
+            print(f"  Validating Chaum-Pedersen proofs ({size} iterations)")
+            avg_verify_scalar = average([t[1] for t in timing_data])
+            std_verify_scalar = std([t[1] for t in timing_data])
+            print(f"    Avg    = {avg_verify_scalar:.6f} sec")
+            print(f"    Stddev = {std_verify_scalar:.6f} sec")
+
+            # Run in parallel
+            start_all_parallel = timer()
+            timing_data_parallel: List[Tuple[float, float]] = scheduler.schedule(
+                chaum_pedersen_bench, [list([input]) for input in inputs]
+            )
+            end_all_parallel = timer()
 
-    print()
-    print("PARALLELISM SPEEDUPS")
-    print("Size / Speedup")
-    for size in problem_sizes:
-        print(f"{size:4d} / {speedup[size]:.3f}x")
-    pool.close()
+            speedup[size] = (end_all_scalar - start_all_scalar) / (
+                end_all_parallel - start_all_parallel
+            )
+            print(f"  Parallel speedup: {speedup[size]:.3f}x")
 
-    bench_end = timer()
-    print()
-    print(f"Total benchmark runtime: {bench_end - bench_start} sec")
+        print()
+        print("PARALLELISM SPEEDUPS")
+        print("Size / Speedup")
+        for size in problem_sizes:
+            print(f"{size:4d} / {speedup[size]:.3f}x")
+
+        bench_end = timer()
+        print()
+        print(f"Total benchmark runtime: {bench_end - bench_start} sec")
 
 ##############################################################################################################
 # Performance conclusions (Dan Wallach, 21 March 2020):
@@ -117,11 +118,3 @@ def identity(x: int) -> int:
 #  500 / 5.765x
 # 1000 / 5.507x
 # 5000 / 5.548x
-
-# I've never seen this break 6x, and tweaking various parameters (e.g., cpu_count() returns 12, so I've
-# tried 6) yielded no significant improvement.  One thing that seems to matter a lot: creating a Pool is
-# a heavyweight operation. Keeping it around and reusing it has a significant impact on performance.
-
-# Moral of the story: the Pool.map() method is very much "good enough" to squeeze useful parallelism out
-# of any machine where we'll be verifying a lot of ballots. If we need radically more throughput, we're
-# probably going to need to move to running on clusters.

bench/bench_chaum_pedersen.py

@@ -62,5 +62,11 @@
     classifiers=CLASSIFIERS,
     project_urls=PROJECT_URLS,
     python_requires="~=3.8",
-    install_requires=["gmpy2>=2.0.8", "numpy>=1.18.2", "jsons>=1.1.2", "cryptography",],
+    install_requires=[
+        "gmpy2>=2.0.8",
+        "numpy>=1.18.2",
+        "jsons>=1.1.2",
+        "cryptography",
+        "psutil>=5.7.2",
+    ],
 )

setup.py

@@ -1,6 +1,4 @@
-from multiprocessing import cpu_count
-from multiprocessing.dummy import Pool
-from typing import Dict, Optional
+from typing import Dict, List, Optional
 
 from .auxiliary import AuxiliaryDecrypt
 from .ballot import CiphertextAcceptedBallot, CiphertextSelection
@@ -21,6 +19,7 @@
 from .key_ceremony import ElectionPublicKey
 from .logs import log_warning
 from .rsa import rsa_decrypt
+from .scheduler import Scheduler
 from .tally import (
     CiphertextTally,
     CiphertextTallyContest,
@@ -111,17 +110,20 @@ def compute_decryption_share_for_cast_contests(
     """
     Compute the decryption for all of the cast contests in the Ciphertext Tally
     """
-    cpu_pool = Pool(cpu_count())
     contests: Dict[CONTEST_ID, CiphertextDecryptionContest] = {}
+    scheduler = Scheduler()
 
     for contest in tally.cast.values():
         selections: Dict[SELECTION_ID, CiphertextDecryptionSelection] = {}
-        selection_decryptions = cpu_pool.starmap(
+        selection_decryptions: List[
+            Optional[CiphertextDecryptionSelection]
+        ] = scheduler.schedule(
             compute_decryption_share_for_selection,
             [
                 (guardian, selection, context)
                 for (_, selection) in contest.tally_selections.items()
             ],
+            with_shared_resources=True,
         )
 
         # verify the decryptions are received and add them to the collection
@@ -136,7 +138,6 @@ def compute_decryption_share_for_cast_contests(
         contests[contest.object_id] = CiphertextDecryptionContest(
             contest.object_id, guardian.object_id, contest.description_hash, selections
         )
-    cpu_pool.close()
     return contests
 
 
@@ -150,17 +151,20 @@ def compute_compensated_decryption_share_for_cast_contests(
     """
     Compute the compensated decryption for all of the cast contests in the Ciphertext Tally
     """
-    cpu_pool = Pool(cpu_count())
+    scheduler = Scheduler()
     contests: Dict[CONTEST_ID, CiphertextCompensatedDecryptionContest] = {}
 
     for contest in tally.cast.values():
         selections: Dict[SELECTION_ID, CiphertextCompensatedDecryptionSelection] = {}
-        selection_decryptions = cpu_pool.starmap(
+        selection_decryptions: List[
+            Optional[CiphertextCompensatedDecryptionSelection]
+        ] = scheduler.schedule(
             compute_compensated_decryption_share_for_selection,
             [
                 (guardian, missing_guardian_id, selection, context, decrypt)
                 for (_, selection) in contest.tally_selections.items()
             ],
+            with_shared_resources=True,
         )
 
         # verify the decryptions are received and add them to the collection
@@ -179,7 +183,6 @@ def compute_compensated_decryption_share_for_cast_contests(
             contest.description_hash,
             selections,
         )
-    cpu_pool.close()
     return contests
 
 
@@ -189,19 +192,22 @@ def compute_decryption_share_for_spoiled_ballots(
     """
     Compute the decryption for all spoiled ballots in the Ciphertext Tally
     """
-    cpu_pool = Pool(cpu_count())
     spoiled_ballots: Dict[BALLOT_ID, BallotDecryptionShare] = {}
+    scheduler = Scheduler()
 
     for spoiled_ballot in tally.spoiled_ballots.values():
         contests: Dict[CONTEST_ID, CiphertextDecryptionContest] = {}
         for contest in spoiled_ballot.contests:
             selections: Dict[SELECTION_ID, CiphertextDecryptionSelection] = {}
-            selection_decryptions = cpu_pool.starmap(
+            selection_decryptions: List[
+                Optional[CiphertextDecryptionSelection]
+            ] = scheduler.schedule(
                 compute_decryption_share_for_selection,
                 [
                     (guardian, selection, context)
                     for selection in contest.ballot_selections
                 ],
+                with_shared_resources=True,
             )
             # verify the decryptions are received and add them to the collection
             for decryption in selection_decryptions:
@@ -225,7 +231,6 @@ def compute_decryption_share_for_spoiled_ballots(
             spoiled_ballot.object_id,
             contests,
         )
-    cpu_pool.close()
     return spoiled_ballots
 
 
@@ -239,21 +244,24 @@ def compute_compensated_decryption_share_for_spoiled_ballots(
     """
     Compute the decryption for all spoiled ballots in the Ciphertext Tally
     """
-    cpu_pool = Pool(cpu_count())
     spoiled_ballots: Dict[BALLOT_ID, CompensatedBallotDecryptionShare] = {}
+    scheduler = Scheduler()
 
     for spoiled_ballot in tally.spoiled_ballots.values():
         contests: Dict[CONTEST_ID, CiphertextCompensatedDecryptionContest] = {}
         for contest in spoiled_ballot.contests:
             selections: Dict[
                 SELECTION_ID, CiphertextCompensatedDecryptionSelection
             ] = {}
-            selection_decryptions = cpu_pool.starmap(
+            selection_decryptions: List[
+                Optional[CiphertextCompensatedDecryptionSelection]
+            ] = scheduler.schedule(
                 compute_compensated_decryption_share_for_selection,
                 [
                     (guardian, missing_guardian_id, selection, context, decrypt)
                     for selection in contest.ballot_selections
                 ],
+                with_shared_resources=True,
             )
             # verify the decryptions are received and add them to the collection
             for decryption in selection_decryptions:
@@ -279,7 +287,6 @@ def compute_compensated_decryption_share_for_spoiled_ballots(
             spoiled_ballot.object_id,
             contests,
         )
-    cpu_pool.close()
     return spoiled_ballots