Function attack_from_queue

textattack/attacker.py:558–597

(
    attack, attack_args, num_gpus, first_to_start, lock, in_queue, out_queue
)

def attack_from_queue(
    attack, attack_args, num_gpus, first_to_start, lock, in_queue, out_queue
):
    assert isinstance(
        attack, Attack
    ), f"`attack` must be of type `Attack`, but got type `{type(attack)}`."

    gpu_id = (torch.multiprocessing.current_process()._identity[0] - 1) % num_gpus
    set_env_variables(gpu_id)
    textattack.shared.utils.set_seed(attack_args.random_seed)
    if torch.multiprocessing.current_process()._identity[0] > 1:
        logging.disable()

    attack.cuda_()

    # Simple non-synchronized check to see if it's the first process to reach this point.
    # This let us avoid waiting for lock.
    if bool(first_to_start.value):
        # If it's first process to reach this step, we first try to acquire the lock to update the value.
        with lock:
            # Because another process could have changed `first_to_start=False` while we wait, we check again.
            if bool(first_to_start.value):
                first_to_start.value = 0
                if not attack_args.silent:
                    print(attack, "\n")

    while True:
        try:
            i, example, ground_truth_output = in_queue.get(timeout=5)
            if i == "END" and example == "END" and ground_truth_output == "END":
                # End process when sentinel value is received
                break
            else:
                result = attack.attack(example, ground_truth_output)
                out_queue.put((i, result))
        except Exception as e:
            if isinstance(e, queue.Empty):
                continue
            else:
                out_queue.put((i, (e, traceback.format_exc())))