Multiprocessing with Netmiko for Plaid Speed

mp_netmiko.py

from __future__ import print_function

import logging
import time
from datetime import datetime
from multiprocessing import Pool, cpu_count, current_process

from netmiko import ConnectHandler
from netmiko.ssh_exception import NetMikoTimeoutException, NetMikoAuthenticationException

logger = logging.getLogger()


def getNetworkData(item_list, worker_function, procs=None):
    """ Takes list of something (devices, sites, etc) and a worker function then uses
    a multiprocess Pool to grab the data quickly. Returns a list."""

    if procs is None:
        procs = 20

    chunk_size = 5

    start_time = datetime.now()
    logger.info("Start time: " + str(start_time))

    proc_num = procs * int(cpu_count())
    logger.info("Number of Processes Used: " + str(proc_num))

    pool = Pool(processes=proc_num, initializer=start_process, maxtasksperchild=3)

    try:
        r = pool.map_async(worker_function, item_list, chunksize=chunk_size)
        logger.info("\nPool starting.")
        r.wait(timeout=900)
        logger.info("\n\nPool map finished.\n\n")
        time.sleep(5)
        pool.close()
        pool.join()
        logger.info("\nPool done.")

        results = r.get()

        end_time = datetime.now()
        total_time = str(end_time - start_time)

        logger.info("Elapsed time: " + total_time)
        logger.info("Number of Processes Used: " + str(proc_num))

        return results
    except:
        pool.terminate()
        raise
        
def start_process():
    """ Prints the worker's process name """
    print("Starting: " + str(current_process().name))
    
    
def getHpVlan90(site):
    """ Queries all HP switches in a site for ports active
    on VLAN 90. Returns a dictionary to the parent function. """

    p_name = str(current_process().name)
    pid_num = str(current_process().pid)

    gdf = 3
    username = "my_username"
    password = "my_password"

    switches = ["switch1", "switch2"]

    command = "sh vlan 90 | beg ----"
    vlan90_info = {}
    vlan90_info[site] = []

    for switch in switches:

        device = switch + "." + site

        logger.info("START," + device + "," + p_name + "," + pid_num)

        hp_swx = {'device_type': 'hp_procurve', 'ip': device,
                  'username': username, 'password': password, 'global_delay_factor': gdf}

        # Lets retry up to
        # three times if we get a failed authentication
        for attempt in range(3):
            try:
                try:
                    sw_con = ConnectHandler(**hp_swx)
                    swx_output = sw_con.send_command(command)
                    sw_con.disconnect()

                except IOError as e:
                    print(switch + "," + site + "," + e)
                    break

                for line in swx_output.splitlines():
                    if line == "" or line == " " or "----" in line:
                        continue

                    line = line.encode('ascii').split()

                    vlan90_info[site].append(line[0])

            except NetMikoTimeoutException as e:
                logger.warning(device + ": Connection timed out.")
                break

            except NetMikoAuthenticationException as e:
                logger.warning(device + ": Failed auth {} time(s).".format(int(attempt) + 1))
            else:
                logger.info("COMPLETED," + device)
                break

        else:
            logger.warning("FAILED," + switch + "." + site + ",all attempts.")

    return vlan90_info

site_list = ['office1', 'office2', 'office3', 'office4', 'remotesite1']

processes = 40

swx_audit = getNetworkData(site_list, getHpVlan90, processes)