Repository: OpenMined/PySyft

Issue:
Rank to Worker_ID translation should be linked to a specific CrypTen computation
## Description

Server workers are intended to run multiple queries from client in parallel. Starting multiple CrypTen computation on the same worker at a time should be possible, however, due to the current way of storing the rank-to-worker-id translation, starting a new computation will overwrite the old translation dictionary and lead to undefined behavior.

## How to Reproduce
1. Start a server worker supporting CrypTen messages
2. Start two consecutive and long CrypTen computations with different workers involved

## Expected Behavior
Running parallel CrypTen computation should work normally as if they were run sequentially.

## Additional Context
Also see #3516 


Discussion:

Resolving patch:
diff --git a/syft/frameworks/crypten/__init__.py b/syft/frameworks/crypten/__init__.py
--- a/syft/frameworks/crypten/__init__.py
+++ b/syft/frameworks/crypten/__init__.py
@@ -3,10 +3,47 @@
 import crypten.communicator as comm
 import crypten
 
+from syft.workers.base import BaseWorker
+
+
+RANK_TO_WORKER_ID = {
+    # Contains translation dictionaries for every computation.
+    # cid (computation id): {rank_to_worker_id dictionary for a specific computation}
+}
+CID = None
+
+
+def get_worker_from_rank(rank: int, cid: int = None) -> BaseWorker:
+    """Find the worker running CrypTen party with specific rank in a certain computation.
+
+    Args:
+        rank: rank of the CrypTen party.
+        cid: CrypTen computation id.
+
+    Returns:
+        BaseWorker corresponding to cid and rank.
+    """
+    if cid is None:
+        if CID is None:
+            # Neither CID have been set appropriately nor cid have been passed
+            raise ValueError("cid must be set.")
+        cid = CID
+
+    rank_to_worker_id = RANK_TO_WORKER_ID.get(cid, None)
+    if rank_to_worker_id is None:
+        raise RuntimeError(
+            "CrypTen computation not initiated properly, computation_id doesn't match any rank to"
+            "worker_id translation table"
+        )
+    return syft.local_worker._get_worker_based_on_id(rank_to_worker_id[rank])
+
 
 def load(tag: str, src: int, **kwargs):
     if src == comm.get().get_rank():
-        worker = syft.local_worker.get_worker_from_rank(src)
+        if CID is None:
+            raise RuntimeError("CrypTen computation id is not set.")
+
+        worker = get_worker_from_rank(src)
         results = worker.search(tag)
 
         # Make sure there is only one result
diff --git a/syft/frameworks/crypten/context.py b/syft/frameworks/crypten/context.py
--- a/syft/frameworks/crypten/context.py
+++ b/syft/frameworks/crypten/context.py
@@ -13,8 +13,12 @@
 
 
 def _launch(
-    func, rank, world_size, master_addr, master_port, queue, func_args, func_kwargs
+    cid, func, rank, world_size, master_addr, master_port, queue, func_args, func_kwargs
 ):  # pragma: no cover
+
+    # set CrypTen computation id
+    sy.frameworks.crypten.CID = cid
+
     communicator_args = {
         "RANK": rank,
         "WORLD_SIZE": world_size,
@@ -34,19 +38,20 @@ def _launch(
     queue.put(return_value)
 
 
-def _new_party(func, rank, world_size, master_addr, master_port, func_args, func_kwargs):
+def _new_party(cid, func, rank, world_size, master_addr, master_port, func_args, func_kwargs):
     queue = multiprocessing.Queue()
     process = multiprocessing.Process(
         target=_launch,
-        args=(func, rank, world_size, master_addr, master_port, queue, func_args, func_kwargs),
+        args=(cid, func, rank, world_size, master_addr, master_port, queue, func_args, func_kwargs),
     )
     return process, queue
 
 
-def run_party(func, rank, world_size, master_addr, master_port, func_args, func_kwargs):
+def run_party(cid, func, rank, world_size, master_addr, master_port, func_args, func_kwargs):
     """Start crypten party localy and run computation.
 
     Args:
+        cid (int): CrypTen computation id.
         func (function): computation to be done.
         rank (int): rank of the crypten party.
         world_size (int): number of crypten parties involved in the computation.
@@ -60,7 +65,7 @@ def run_party(func, rank, world_size, master_addr, master_port, func_args, func_
     """
 
     process, queue = _new_party(
-        func, rank, world_size, master_addr, master_port, func_args, func_kwargs
+        cid, func, rank, world_size, master_addr, master_port, func_args, func_kwargs
     )
     was_initialized = DistributedCommunicator.is_initialized()
     if was_initialized:
@@ -130,8 +135,6 @@ def wrapper(*args, **kwargs):
 
             rank_to_worker_id = dict(zip(range(0, len(workers)), [worker.id for worker in workers]))
 
-            sy.local_worker.rank_to_worker_id = rank_to_worker_id
-
             # TODO: run ttp in a specified worker
             # if crypten.mpc.ttp_required():
             #     ttp_process, _ = _new_party(
@@ -189,8 +192,6 @@ def wrapper(*args, **kwargs):
             for thread in threads:
                 thread.join()
 
-            del sy.local_worker.rank_to_worker_id
-
             return return_values
 
         return wrapper
diff --git a/syft/frameworks/crypten/message_handler.py b/syft/frameworks/crypten/message_handler.py
--- a/syft/frameworks/crypten/message_handler.py
+++ b/syft/frameworks/crypten/message_handler.py
@@ -1,28 +1,21 @@
-import types
+import syft
 
 from syft.messaging.message import CryptenInitPlan
 from syft.messaging.message import CryptenInitJail
 from syft.messaging.message import ObjectMessage
 
+from syft.frameworks import crypten as syft_crypten
 from syft.frameworks.crypten.context import run_party
-
 from syft.frameworks.crypten.jail import JailRunner
 from syft.frameworks.crypten import utils
-from syft.workers.base import BaseWorker
 
 from syft.generic.abstract.message_handler import AbstractMessageHandler
 
 
-def get_worker_from_rank(worker: BaseWorker, rank: int) -> BaseWorker:
-    assert hasattr(worker, "rank_to_worker_id"), "First need to call run_crypten_party"
-    return worker._get_worker_based_on_id(worker.rank_to_worker_id[rank])
-
-
 class CryptenMessageHandler(AbstractMessageHandler):
     def __init__(self, object_store, worker):
         super().__init__(object_store)
         self.worker = worker
-        setattr(worker, "get_worker_from_rank", types.MethodType(get_worker_from_rank, worker))
 
     def init_routing_table(self):
         return {
@@ -41,7 +34,10 @@ def run_crypten_party_plan(self, msg: CryptenInitPlan) -> ObjectMessage:
             An ObjectMessage containing the return value of the crypten function computed.
         """
 
-        self.worker.rank_to_worker_id, world_size, master_addr, master_port = msg.crypten_context
+        rank_to_worker_id, world_size, master_addr, master_port = msg.crypten_context
+
+        cid = syft.ID_PROVIDER.pop()
+        syft_crypten.RANK_TO_WORKER_ID[cid] = rank_to_worker_id
 
         # TODO Change this, we need a way to handle multiple plan definitions
         plans = self.worker.search("crypten_plan")
@@ -49,10 +45,13 @@ def run_crypten_party_plan(self, msg: CryptenInitPlan) -> ObjectMessage:
 
         plan = plans[0].get()
 
-        rank = self._current_rank()
+        rank = self._current_rank(rank_to_worker_id)
         assert rank is not None
 
-        return_value = run_party(plan, rank, world_size, master_addr, master_port, (), {})
+        return_value = run_party(cid, plan, rank, world_size, master_addr, master_port, (), {})
+        # remove rank to id transaltion dict
+        del syft_crypten.RANK_TO_WORKER_ID[cid]
+
         return ObjectMessage(return_value)
 
     def run_crypten_party_jail(self, msg: CryptenInitJail):
@@ -66,23 +65,31 @@ def run_crypten_party_jail(self, msg: CryptenInitJail):
             An ObjectMessage containing the return value of the crypten function computed.
         """
 
-        self.worker.rank_to_worker_id, world_size, master_addr, master_port = msg.crypten_context
+        rank_to_worker_id, world_size, master_addr, master_port = msg.crypten_context
+
+        cid = syft.ID_PROVIDER.pop()
+        syft_crypten.RANK_TO_WORKER_ID[cid] = rank_to_worker_id
 
         ser_func = msg.jail_runner
         onnx_model = msg.model
         crypten_model = None if onnx_model is None else utils.onnx_to_crypten(onnx_model)
         jail_runner = JailRunner.detail(ser_func, model=crypten_model)
 
-        rank = self._current_rank()
+        rank = self._current_rank(rank_to_worker_id)
         assert rank is not None
 
-        return_value = run_party(jail_runner, rank, world_size, master_addr, master_port, (), {})
+        return_value = run_party(
+            cid, jail_runner, rank, world_size, master_addr, master_port, (), {}
+        )
+        # remove rank to id transaltion dict
+        del syft_crypten.RANK_TO_WORKER_ID[cid]
+
         return ObjectMessage(return_value)
 
-    def _current_rank(self):
+    def _current_rank(self, rank_to_worker_id):
         """Returns current rank based on worker_id."""
         rank = None
-        for r, worker_id in self.worker.rank_to_worker_id.items():
+        for r, worker_id in rank_to_worker_id.items():
             if worker_id == self.worker.id:
                 rank = r
                 break


Test patch:
diff --git a/test/crypten/test_context.py b/test/crypten/test_context.py
--- a/test/crypten/test_context.py
+++ b/test/crypten/test_context.py
@@ -186,7 +186,7 @@ def party():  # pragma: no cover
         t = crypten.cryptensor(expected)
         return t.get_plain_text()
 
-    t = run_party(party, 0, 1, "127.0.0.1", 15463, (), {})
+    t = run_party(None, party, 0, 1, "127.0.0.1", 15463, (), {})
     result = utils.unpack_values(t)
     assert result == expected