code: import syft as sy
  import torch
  from typing import Dict
  from typing import Any
  import logging
  
  logger = logging.getLogger(__name__)
  
  
  def extract_batches_per_worker(federated_train_loader: sy.FederatedDataLoader):
      """Extracts the batches from the federated_train_loader and stores them
         in a dictionary (keys = data.location).
  
         Args:
         federated_train_loader: the connection object we use to send responses.
                      back to the client.
  
      """
      logging_interval = 100
      batches = {}
      for worker_id in federated_train_loader.workers:
          worker = federated_train_loader.federated_dataset.datasets[worker_id].location
          batches[worker] = []
  
      for batch_idx, (data, target) in enumerate(federated_train_loader):
          if batch_idx % logging_interval == 0:
              logger.debug("Extracted %s batches from federated_train_loader", batch_idx)
          batches[data.location].append((data, target))
  
      return batches
  
  
  def add_model(dst_model, src_model):
      """Add the parameters of two models.
  
          Args:
              dst_model (torch.nn.Module): the model to which the src_model will be added.
              src_model (torch.nn.Module): the model to be added to dst_model.
          Returns:
              torch.nn.Module: the resulting model of the addition.
  
          """
  
      params1 = src_model.named_parameters()
      params2 = dst_model.named_parameters()
      dict_params2 = dict(params2)
      with torch.no_grad():
          for name1, param1 in params1:
              if name1 in dict_params2:
                  dict_params2[name1].set_(param1.data + dict_params2[name1].data)
      return dst_model
  
  
  def scale_model(model, scale):
      """Scale the parameters of a model.
  
      Args:
          model (torch.nn.Module): the models whose parameters will be scaled.
          scale (float): the scaling factor.
      Returns:
          torch.nn.Module: the module with scaled parameters.
  
      """
      params = model.named_parameters()
      dict_params = dict(params)
      with torch.no_grad():
          for name, param in dict_params.items():
              dict_params[name].set_(dict_params[name].data * scale)
      return model
  
  
  def federated_avg(models: Dict[Any, torch.nn.Module]) -> torch.nn.Module:
      """Calculate the federated average of a dictionary containing models.
         The models are extracted from the dictionary
         via the models.values() command.
  
      Args:
          models (Dict[Any, torch.nn.Module]): a dictionary of models
          for which the federated average is calculated.
  
      Returns:
          torch.nn.Module: the module with averaged parameters.
      """
      nr_models = len(models)
      model_list = list(models.values())
      model = model_list[0]
      for i in range(1, nr_models):
          model = add_model(model, model_list[i])
      model = scale_model(model, 1.0 / nr_models)
      return model
  
  
  def accuracy(pred_softmax, target):
      """Calculate the accuray of a given prediction.
  
      This functions assumes pred_softmax to be converted into the final prediction by taking the argmax.
  
      Args:
          pred_softmax: array type(float), providing nr_classes values per element in target.
          target: array type(int), correct classes, taking values in range [0, nr_classes).
  
      Returns:
          accuracy: float, fraction of correct predictions.
  
      """
      nr_elems = len(target)
      pred = pred_softmax.argmax(dim=1)
      return (pred.float() == target.view(pred.shape).float()).sum().numpy() / float(nr_elems)
  
  
  def create_gaussian_mixture_toy_data(nr_samples: int):  # pragma: no cover
      """ Create a simple toy data for binary classification
  
      The data is drawn from two normal distributions
      target = 1: mu = 2, sigma = 1
      target = 0: mu = 0, sigma = 1
      The dataset is balanced with an equal number of positive and negative samples
  
      Args:
          nr_samples: number of samples to generate
  
      Returns:
          data, targets
  
  
      """
      sample_dim = 2
      one_half = int(nr_samples / 2)
      X1 = torch.randn(one_half, sample_dim, requires_grad=True) - 5
      X2 = torch.randn(one_half, sample_dim, requires_grad=True) + 5
      X = torch.cat([X1, X2], dim=0)
      Y1 = torch.zeros(one_half, requires_grad=False).long()
      Y2 = torch.ones(one_half, requires_grad=False).long()
      Y = torch.cat([Y1, Y2], dim=0)
      return X, Y
  
  
  def iris_data_partial():
      """
  
      Returns: 30 samples from the iris data set: https://archive.ics.uci.edu/ml/datasets/iris
  
      """
      data = [
          [5.1, 3.5, 1.4, 0.2],
          [4.9, 3.0, 1.4, 0.2],
          [4.7, 3.2, 1.3, 0.2],
          [4.6, 3.1, 1.5, 0.2],
          [5.0, 3.6, 1.4, 0.2],
          [5.4, 3.9, 1.7, 0.4],
          [4.6, 3.4, 1.4, 0.3],
          [5.0, 3.4, 1.5, 0.2],
          [4.4, 2.9, 1.4, 0.2],
          [4.9, 3.1, 1.5, 0.1],
      ]
  
      target_to_string = {0: "Iris-setosa", 1: "Iris-versicolor", 2: "Iris-virginica"}
      targets = [0, 0, 0, 0, 0, 0, 0, 0, 0, 0]
  
      data += [
          [7.0, 3.2, 4.7, 1.4],
          [6.4, 3.2, 4.5, 1.5],
          [6.9, 3.1, 4.9, 1.5],
          [5.5, 2.3, 4.0, 1.3],
          [6.5, 2.8, 4.6, 1.5],
          [5.7, 2.8, 4.5, 1.3],
          [6.3, 3.3, 4.7, 1.6],
          [4.9, 2.4, 3.3, 1.0],
          [6.6, 2.9, 4.6, 1.3],
          [5.2, 2.7, 3.9, 1.4],
      ]
  
      targets += [1, 1, 1, 1, 1, 1, 1, 1, 1, 1]
  
      data += [
          [6.3, 3.3, 6.0, 2.5],
          [5.8, 2.7, 5.1, 1.9],
          [7.1, 3.0, 5.9, 2.1],
          [6.3, 2.9, 5.6, 1.8],
          [6.5, 3.0, 5.8, 2.2],
          [7.6, 3.0, 6.6, 2.1],
          [4.9, 2.5, 4.5, 1.7],
          [7.3, 2.9, 6.3, 1.8],
          [6.7, 2.5, 5.8, 1.8],
          [7.2, 3.6, 6.1, 2.5],
      ]
  
      targets += [2, 2, 2, 2, 2, 2, 2, 2, 2, 2]
  
      return torch.tensor(data), torch.tensor(targets)
   import io
  from os.path import exists, join
  import time
  from socket import gethostname
  from OpenSSL import crypto, SSL
  import pytest
  import torch
  import syft as sy
  from syft.generic.frameworks.hook import hook_args
  from syft.frameworks.torch.fl import utils
  
  from syft.workers.websocket_client import WebsocketClientWorker
  from syft.workers.websocket_server import WebsocketServerWorker
  
  from test.conftest import instantiate_websocket_client_worker
  
  
  PRINT_IN_UNITTESTS = False
  
  
  @pytest.mark.parametrize("secure", [True, False])
  def test_websocket_worker_basic(hook, start_proc, secure, tmpdir):
      """Evaluates that you can do basic tensor operations using
      WebsocketServerWorker in insecure and secure mode."""
  
      def create_self_signed_cert(cert_path, key_path):
          # create a key pair
          k = crypto.PKey()
          k.generate_key(crypto.TYPE_RSA, 1024)
  
          # create a self-signed cert
          cert = crypto.X509()
          cert.gmtime_adj_notBefore(0)
          cert.gmtime_adj_notAfter(1000)
          cert.set_pubkey(k)
          cert.sign(k, "sha1")
  
          # store keys and cert
          open(cert_path, "wb").write(crypto.dump_certificate(crypto.FILETYPE_PEM, cert))
          open(key_path, "wb").write(crypto.dump_privatekey(crypto.FILETYPE_PEM, k))
  
      kwargs = {
          "id": "secure_fed" if secure else "not_secure_fed",
          "host": "localhost",
          "port": 8766,
          "hook": hook,
      }
  
      if secure:
          # Create cert and keys
          cert_path = tmpdir.join("test.crt")
          key_path = tmpdir.join("test.key")
          create_self_signed_cert(cert_path, key_path)
          kwargs["cert_path"] = cert_path
          kwargs["key_path"] = key_path
  
      process_remote_worker = start_proc(WebsocketServerWorker, **kwargs)
  
      time.sleep(0.1)
      x = torch.ones(5)
  
      if secure:
          # unused args
          del kwargs["cert_path"]
          del kwargs["key_path"]
  
      kwargs["secure"] = secure
      remote_proxy = instantiate_websocket_client_worker(**kwargs)
  
      x = x.send(remote_proxy)
      y = x + x
      y = y.get()
  
      assert (y == torch.ones(5) * 2).all()
  
      del x
  
      remote_proxy.close()
      time.sleep(0.1)
      remote_proxy.remove_worker_from_local_worker_registry()
      process_remote_worker.terminate()
  
  
  def test_websocket_workers_search(hook, start_remote_worker):
      """Evaluates that a client can search and find tensors that belong
      to another party"""
      # Args for initializing the websocket server and client
      server, remote_proxy = start_remote_worker(id="fed2", hook=hook, port=8767)
  
      # Sample tensor to store on the server
      sample_data = torch.tensor([1, 2, 3, 4]).tag("#sample_data", "#another_tag")
      _ = sample_data.send(remote_proxy)
  
      # Search for the tensor located on the server by using its tag
      results = remote_proxy.search(["#sample_data", "#another_tag"])
  
      assert results
      assert results[0].owner.id == "me"
      assert results[0].location.id == "fed2"
  
      # Search multiple times should still work
      results = remote_proxy.search(["#sample_data", "#another_tag"])
  
      assert results
      assert results[0].owner.id == "me"
      assert results[0].location.id == "fed2"
  
      remote_proxy.close()
      time.sleep(0.1)
      remote_proxy.remove_worker_from_local_worker_registry()
      server.terminate()
  
  
  def test_list_objects_remote(hook, start_remote_worker):
      server, remote_proxy = start_remote_worker(id="fed-list-objects", hook=hook, port=8765)
      remote_proxy.clear_objects()
  
      x = torch.tensor([1, 2, 3]).send(remote_proxy)
  
      res = remote_proxy.list_tensors_remote()
  
      res_dict = eval(res.replace("tensor", "torch.tensor"))
      assert len(res_dict) == 1
  
      y = torch.tensor([4, 5, 6]).send(remote_proxy)
      res = remote_proxy.list_tensors_remote()
      res_dict = eval(res.replace("tensor", "torch.tensor"))
      assert len(res_dict) == 2
  
      # delete x before terminating the websocket connection
      del x
      del y
      time.sleep(0.1)
      remote_proxy.close()
      time.sleep(0.1)
      remote_proxy.remove_worker_from_local_worker_registry()
      server.terminate()
  
  
  def test_objects_count_remote(hook, start_remote_worker):
      server, remote_proxy = start_remote_worker(id="fed-count-objects", hook=hook, port=8764)
      remote_proxy.clear_objects()
  
      x = torch.tensor([1, 2, 3]).send(remote_proxy)
  
      nr_objects = remote_proxy.tensors_count_remote()
      assert nr_objects == 1
  
      y = torch.tensor([4, 5, 6]).send(remote_proxy)
      nr_objects = remote_proxy.tensors_count_remote()
      assert nr_objects == 2
  
      x.get()
      nr_objects = remote_proxy.tensors_count_remote()
      assert nr_objects == 1
  
      # delete remote object before terminating the websocket connection
      del y
      time.sleep(0.1)
      remote_proxy.close()
      time.sleep(0.1)
      remote_proxy.remove_worker_from_local_worker_registry()
      server.terminate()
  
  
  def test_clear_objects_remote(hook, start_remote_worker):
      server, remote_proxy = start_remote_worker(id="fed-clear-objects", hook=hook, port=8769)
  
      x = torch.tensor([1, 2, 3]).send(remote_proxy, garbage_collect_data=False)
      y = torch.tensor(4).send(remote_proxy, garbage_collect_data=False)
  
      nr_objects = remote_proxy.tensors_count_remote()
      assert nr_objects == 2
  
      remote_proxy.clear_objects_remote()
      nr_objects = remote_proxy.objects_count_remote()
      assert nr_objects == 0
  
      remote_proxy.close()
      remote_proxy.remove_worker_from_local_worker_registry()
      server.terminate()
  
  
  def test_connect_close(hook, start_remote_worker):
      server, remote_proxy = start_remote_worker(id="fed-connect-close", hook=hook, port=8770)
  
      x = torch.tensor([1, 2, 3])
      x_ptr = x.send(remote_proxy)
  
      assert remote_proxy.tensors_count_remote() == 1
  
      remote_proxy.close()
  
      time.sleep(0.1)
  
      remote_proxy.connect()
  
      assert remote_proxy.tensors_count_remote() == 1
  
      x_val = x_ptr.get()
      assert (x_val == x).all()
  
      remote_proxy.close()
      remote_proxy.remove_worker_from_local_worker_registry()
  
      time.sleep(0.1)
  
      server.terminate()
  
  
  def test_websocket_worker_multiple_output_response(hook, start_remote_worker):
      """Evaluates that you can do basic tensor operations using
      WebsocketServerWorker."""
      server, remote_proxy = start_remote_worker(id="socket_multiple_output", hook=hook, port=8771)
  
      x = torch.tensor([1.0, 3, 2])
      x = x.send(remote_proxy)
  
      p1, p2 = torch.sort(x)
      x1, x2 = p1.get(), p2.get()
  
      assert (x1 == torch.tensor([1.0, 2, 3])).all()
      assert (x2 == torch.tensor([0, 2, 1])).all()
  
      x.get()  # retrieve remote object before closing the websocket connection
  
      remote_proxy.close()
      server.terminate()
  
  
  def test_send_command_whitelist(hook, start_remote_worker):
      server, remote_proxy = start_remote_worker(
          id="worker_call_api_good_methods", hook=hook, port=8772
      )
      whitelisted_methods = {
          "torch": {"tensor": [1, 2, 3], "rand": (2, 3), "randn": (2, 3), "zeros": (2, 3)}
      }
  
      for framework, methods in whitelisted_methods.items():
          attr = getattr(remote_proxy.remote, framework)
  
          for method, inp in methods.items():
              x = getattr(attr, method)(inp)
  
              if "rand" not in method:
                  assert (x.get() == getattr(torch, method)(inp)).all()
  
      remote_proxy.close()
      server.terminate()
  
  
  def test_send_command_not_whitelisted(hook, start_remote_worker):
      server, remote_proxy = start_remote_worker(
          id="worker_call_api_bad_method", hook=hook, port=8773
      )
  
      method_not_exist = "openmind"
  
      for framework in remote_proxy.remote.frameworks:
          if framework in dir(remote_proxy.remote):
              attr = getattr(remote_proxy.remote, framework)
              with pytest.raises(AttributeError):
                  getattr(attr, method_not_exist)
  
      remote_proxy.close()
      server.terminate()
  
  
  @pytest.mark.skip
  def test_evaluate(hook, start_proc):  # pragma: no cover
  
      sy.local_worker.clear_objects()
      sy.generic.frameworks.hook.hook_args.hook_method_args_functions = {}
      sy.generic.frameworks.hook.hook_args.hook_method_response_functions = {}
      sy.generic.frameworks.hook.hook_args.get_tensor_type_functions = {}
      sy.generic.frameworks.hook.hook_args.register_response_functions = {}
  
      data, target = utils.iris_data_partial()
  
      dataset = sy.BaseDataset(data=data, targets=target)
  
      kwargs = {"id": "evaluate_remote", "host": "localhost", "port": 8790, "hook": hook}
      dataset_key = "iris"
      # TODO: check why unit test sometimes fails when WebsocketServerWorker is started from the unit test. Fails when run after test_federated_client.py
      # process_remote_worker = start_proc(WebsocketServerWorker, dataset=(dataset, dataset_key), verbose=True, **kwargs)
  
      remote_proxy = instantiate_websocket_client_worker(**kwargs)
  
      def loss_fn(pred, target):
          return torch.nn.functional.cross_entropy(input=pred, target=target)
  
      class Net(torch.nn.Module):
          def __init__(self):
              super(Net, self).__init__()
              self.fc1 = torch.nn.Linear(4, 3)
  
              torch.nn.init.xavier_normal_(self.fc1.weight)
  
          def forward(self, x):
              x = torch.nn.functional.relu(self.fc1(x))
              return x
  
      model_untraced = Net()
      model = torch.jit.trace(model_untraced, data)
      loss_traced = torch.jit.trace(loss_fn, (torch.tensor([[0.3, 0.5, 0.2]]), torch.tensor([1])))
  
      pred = model(data)
      loss_before = loss_fn(target=target, pred=pred)
      if PRINT_IN_UNITTESTS:  # pragma: no cover
          print(f"Loss: {loss_before}")
  
      # Create and send train config
      train_config = sy.TrainConfig(
          batch_size=4,
          model=model,
          loss_fn=loss_traced,
          model_id=None,
          loss_fn_id=None,
          optimizer_args=None,
          epochs=1,
      )
      train_config.send(remote_proxy)
  
      result = remote_proxy.evaluate(
          dataset_key=dataset_key, return_histograms=True, nr_bins=3, return_loss=True
      )
  
      len_dataset = result["nr_predictions"]
      hist_target = result["histogram_target"]
  
      if PRINT_IN_UNITTESTS:  # pragma: no cover
          print(f"Evaluation result before training: {result}")
  
      assert len_dataset == 30
      assert (hist_target == [10, 10, 10]).all()
  
      remote_proxy.close()
      remote_proxy.remove_worker_from_local_worker_registry()
      # process_remote_worker.terminate()
  
apis: torch.ones torch.randn torch.no_grad torch.tensor torch.zeros torch.cat
apis: torch.ones torch.nn.Linear torch.nn.init.xavier_normal_ torch.tensor torch.nn.functional.cross_entropy torch.sort torch.jit.trace