<commit_before>#!/usr/bin/env python
from tempfile import TemporaryFile
import os, sys, re, socket, time, pickle, csv, uuid, subprocess, argparse, decimal, select

class Server:
  def __init__(self, arguments):
    self.arguments = arguments
    self.arguments.cwd = os.getcwd()
    # Test to see if we are starting as a server
    if self.arguments.pbs == True:
      if os.getenv('PBS_NODEFILE') != None:
        # Initialize an agent, strictly for holding our stdout logs. Give it the UUID of 'server'
        self.agent = Agent(self.arguments, 'server')
        if self.arguments.recover:
          self.logfile = WriteCSV(self.arguments.outfile[0], False)
        else:
          self.logfile = WriteCSV(self.arguments.outfile[0], True)
        self.client_connections = []
        self.startServer()
      else:
        print 'I could not find your PBS_NODEFILE. Is PBS loaded?'
        sys.exit(1)
    # If we are not a server, start the single client
    else:
      self.startClient()

  def startServer(self):
    # Setup the TCP socket
    self.server_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    self.server_socket.bind((socket.gethostname(), 0))
    self.server_socket.listen(5)
    (self.host, self.port) = self.server_socket.getsockname()

    # We will store all connections (sockets objects) made to the server in a list
    self.client_connections.append(self.server_socket)

    # Launch the actual binary we want to track
    self._launchJob()

    # Now launch all pbs agents
    self._launchClients()

    # This is a try so we can handle a keyboard ctrl-c
    try:

      # Continue to listen and accept active connections from agents
      # until all agents report a STOP command.
      AGENTS_ACTIVE = True
      while AGENTS_ACTIVE:
        read_sockets, write_sockets, error_sockets = select.select(self.client_connections,[],[])
        for sock in read_sockets:
          if sock == self.server_socket:
            # Accept an incomming connection
            self.client_connections.append(self.server_socket.accept()[0])
          else:
            # Deal with the data being sent to the server by its agents
            self.handleAgent()

            # Check to see if _all_ agents are telling the server to stop
            agent_count = len(self.agent.agent_data.keys())
            current_count = 0
            for agent in self.agent.agent_data.keys():
              if self.agent.agent_data[agent]['STOP']:
                current_count += 1

            # if All Agents have reported a STOP command, begin to exit
            if current_count == agent_count:
              AGENTS_ACTIVE = False
              # Gotta get out of the for loop somehow...
              break

          # Sleep a bit before reading additional data
          time.sleep(self.arguments.repeat_rate[-1])

      # Close the server socket
      self.server_socket.close()

      # Close the logfile as the server is about to exit
      self.logfile.close()

    # Cancel server operations if ctrl-c was pressed
    except KeyboardInterrupt:
      print 'Canceled by user. Wrote log:', self.arguments.outfile[0]
      sys.exit(0)

    # Normal exiting procedures
    print '\n\nAll agents have stopped. Log file saved to:', self.arguments.outfile[0]
    sys.exit(0)

  def startClient(self):
    Client(self.arguments)

  def _launchClients(self):
    # Read the environment PBS_NODEFILE
    self._PBS_NODEFILE = open(os.getenv('PBS_NODEFILE'), 'r')
    nodes = set(self._PBS_NODEFILE.read().split())

    # Print some useful information about our setup
    print 'Memory Logger running on Host:', self.host, 'Port:', self.port, '\nNodes:', ', '.join(nodes), '\nSample rate (including stdout):', self.arguments.repeat_rate[-1], 's (use --repeat-rate to adjust)\nRemote agents delaying', self.arguments.pbs_delay[-1], 'second/s before tracking. (use --pbs-delay to adjust)\n'

    # Build our command list based on the PBS_NODEFILE
    command = []
    for node in nodes:
      command.append([  'ssh', node,
                   'bash --login -c "source /etc/profile && ' \
                   + 'sleep ' + str(self.arguments.pbs_delay[-1]) + ' && ' \
                   + os.path.abspath(__file__) \
                   + ' --call-back-host ' \
                   + self.host + ' ' + str(self.port) \
                   + '"'])

    # remote into each node and execute another copy of memory_logger.py
    # with a call back argument to recieve further instructions
    for pbs_node in command:
      subprocess.Popen(pbs_node, stdout=subprocess.PIPE, stderr=subprocess.PIPE)

  # Launch the binary we intend to track
  def _launchJob(self):
    subprocess.Popen(self.arguments.run[-1].split(), stdout=self.agent.log, stderr=self.agent.log)

  # A connection has been made from client to server
  # Capture that data, and determin what to do with it
  def handleAgent(self):
    # Loop through all client connections, and receive data if any
    for agent_socket in self.client_connections:

      # Completely ignore the server_socket object
      if agent_socket == self.server_socket:
        continue

      # Assign an AgentConnector for the task of handling data between client and server
      reporting_agent = AgentConnector(self.arguments, agent_socket)

      # OK... get data from a client and begin
      new_data = reporting_agent.readData()
      if new_data != None:
        # There should be only one dictionary key (were reading data from just one client at a time)
        agent_uuid = new_data.keys()[0]

        # Update our dictionary of an agents data
        self.agent.agent_data[agent_uuid] = new_data[agent_uuid]

        # Modify incoming Agents timestamp to match Server's time (because every node is a little bit off)
        if self.arguments.recover:
          self.agent.agent_data[agent_uuid]['TIMESTAMP'] = GetTime().now - self.agent.delta
        else:
          self.agent.agent_data[agent_uuid]['TIMESTAMP'] = GetTime().now

        # update total usage for all known reporting agents
        total_usage = 0
        for one_agent in self.agent.agent_data.keys():
          total_usage += self.agent.agent_data[one_agent]['MEMORY']
        self.agent.agent_data[agent_uuid]['TOTAL'] = int(total_usage)

        # Get any stdout thats happened thus far and apply it to what ever agent just sent us data
        self.agent.agent_data[agent_uuid]['STDOUT'] = self.agent._getStdout()

        # Write to our logfile
        self.logfile.write(self.agent.agent_data[agent_uuid])

        # Check for any agents sending a stop command. If we find one,
        # set some zeroing values, and close that agent's socket.
        if self.agent.agent_data[agent_uuid]['STOP']:
          self.agent.agent_data[agent_uuid]['MEMORY'] = 0
          agent_socket.close()
          if agent_socket != self.server_socket:
            self.client_connections.remove(agent_socket)

          # Go ahead and set our server agent to STOP as well.
          # The server will continue recording samples from agents
          self.agent.agent_data['server']['STOP'] = True

        # If an Agent has made a request for instructions, handle it here
        update_client = False
        if new_data[agent_uuid]['REQUEST'] != None:
          for request in new_data[agent_uuid]['REQUEST'].iteritems():
            if new_data[agent_uuid]['REQUEST'][request[0]] == '':
              update_client = True
              # We only support sending any arguments supplied to ther server, back to the agent
              for request_type in dir(self.arguments):
                if request[0] == str(request_type):
                  self.agent.agent_data[agent_uuid]['REQUEST'][request[0]] = getattr(self.arguments, request[0])

          # If an Agent needed additional instructions, go ahead and re-send those instructions
          if update_client:
            reporting_agent.sendData(self.agent.agent_data[agent_uuid])

class Client:
  def __init__(self, arguments):
    self.arguments = arguments

    # Initialize an Agent with a UUID based on our hostname
    self.my_agent = Agent(arguments, str(uuid.uuid3(uuid.NAMESPACE_DNS, socket.gethostname())))

    # Initialize an AgentConnector
    self.remote_server = AgentConnector(self.arguments)

    # If client will talk to a server (PBS)
    if self.arguments.call_back_host:
      # We know by initializing an agent, agent_data contains the necessary message asking for further instructions
      self.my_agent.agent_data[self.my_agent.my_uuid] = self.remote_server.sendData(self.my_agent.agent_data)

      # Apply new instructions received from server (this basically updates our arguments)
      for request in self.my_agent.agent_data[self.my_agent.my_uuid]['REQUEST'].iteritems():
        for request_type in dir(self.arguments):
          if request[0] == str(request_type):
            setattr(self.arguments, request[0], request[1])

      # Requests have been satisfied, set to None
      self.my_agent.agent_data[self.my_agent.my_uuid]['REQUEST'] = None

      # Change to the same directory as the server was when initiated (needed for PBS stuff)
      os.chdir(self.arguments.cwd)

    # Client will not be talking to a server, save data to a file instead
    else:
      # Deal with --recover
      if self.arguments.recover:
        # Do not overwrite the file
        self.logfile = WriteCSV(self.arguments.outfile[0], False)
      else:
        # Overwrite the file
        self.logfile = WriteCSV(self.arguments.outfile[0], True)

    # Lets begin!
    self.startProcess()

  # This function handles the starting and stoping of the sampler process.
  # We loop until an agent returns a stop command.
  def startProcess(self):
    AGENTS_ACTIVE = True

    # If we know we are the only client, go ahead and start the process we want to track.
    if self.arguments.call_back_host == None:
      subprocess.Popen(self.arguments.run[-1].split(), stdout=self.my_agent.log, stderr=self.my_agent.log)
      # Delay just a bit to keep from recording a possible zero memory usage as the binary starts up
      time.sleep(self.arguments.sample_delay[0])


    # This is a try so we can handle a keyboard ctrl-c
    try:

      # Continue to process data until an Agent reports a STOP command
      while AGENTS_ACTIVE:
        # Take a sample
        current_data = self.my_agent.takeSample()

        # Handle the data supplied by the Agent.
        self._handleData(current_data)

        # If an Agent reported a STOP command, go ahead and begin the shutdown phase
        if current_data[current_data.keys()[0]]['STOP']:
          AGENTS_ACTIVE = False

        # Sleep just a bit between samples, as to not saturate the machine
        time.sleep(self.arguments.repeat_rate[-1])

      # An agent reported a stop command... so let everyone know where the log was saved, and exit!
      if self.arguments.call_back_host == None:
        print 'Binary has exited. Wrote log:', self.arguments.outfile[0]

    # Cancel server operations if ctrl-c was pressed
    except KeyboardInterrupt:
      self.logfile.close()
      print 'Canceled by user. Wrote log:', self.arguments.outfile[0]
      sys.exit(0)

    # Everything went smooth.
    sys.exit(0)

  # Figure out what to do with the sampled data
  def _handleData(self, data):
    # Sending the sampled data to a server
    if self.arguments.call_back_host:
      self.remote_server.sendData(data)
    # Saving the sampled data to a file
    else:
      # Compute the TOTAL memory usage to be how much our one agent reported
      # Because were the only client doing any work
      data[self.my_agent.my_uuid]['TOTAL'] = data[self.my_agent.my_uuid]['MEMORY']
      self.logfile.write(data[self.my_agent.my_uuid])

      # If the agent has been told to stop, close the database file
      if self.my_agent.agent_data[self.my_agent.my_uuid]['STOP'] == True:
        self.logfile.close()

class AgentConnector:
  """
Functions used to communicate to and from Client and Server.
Both Client and Server classes use this object.

readData()
sendData('message', socket_connection=None)

if sendData's socket_connection is None, it will create a new connection to the server
based on supplied arguments
"""
  def __init__(self, arguments, connection=None):
    self.arguments = arguments
    self.connection = connection
    self.CREATED_CONNECTION = False

    # If the connection is None, meaning this object was instanced by a client,
    # we must create a connection to the server first
    if self.connection == None and self.arguments.call_back_host != None:
      self.CREATED_CONNECTION = True
      self.connection = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
      self.connection.settimeout(15)
      self.connection.connect((self.arguments.call_back_host[0], int(self.arguments.call_back_host[1])))

  # read all data sent by an agent
  def readData(self):
    # Get how much data there is to receive
    # The first eight bytes is our data length
    data_width = int(self.connection.recv(8))
    tmp_received = ''

    # We need to receive precisely the ammount of data the
    # client is trying to send us.
    while len(tmp_received) < data_width:
      if data_width - len(tmp_received) > 1024:
        tmp_received += self.connection.recv(1024)
      else:
        tmp_received += self.connection.recv(data_width - (len(tmp_received)))

    # unpickle the received message
    return self._unpickleMessage(tmp_received)

  # send data to an agent
  def sendData(self, message):
    # pickle the data up, and send the message
    self.connection.sendall(self._pickleMessage(message))

    # If we had to create the socket (connection was none), and this client/agent is requesting
    # instructions, go ahead and read the data that _better be there_ sent to us by the server.
    if self.CREATED_CONNECTION and message[message.keys()[0]]['REQUEST'] != None:
      return self.readData()

  # The following two functions pickle up the data for easy socket transport
  def _pickleMessage(self, message):
    t = TemporaryFile()
    pickle.dump(message, t)
    t.seek(0)
    str_msg = t.read()
    str_len = len(str_msg)
    message = "%-8d" % (str_len,) + str_msg
    return message

  def _unpickleMessage(self, message):
    t = TemporaryFile()
    t.write(message)
    t.seek(0)
    try:
      return pickle.load(t)
    except KeyError:
      print 'Socket data was not pickled data: ', message
    except:
      raise

class WriteCSV:
  def __init__(self, logfile, overwrite):
    if overwrite:
      self.file_object = open(logfile, 'w', 1)
    else:
      self.file_object = open(logfile, 'a', 1)
    self.log_file = csv.writer(self.file_object, delimiter=',', quotechar='|', escapechar='\\', quoting=csv.QUOTE_MINIMAL)

  # Close the logfile
  def close(self):
    self.file_object.close()

  # Write a CSV row
  def write(self, data):
    formatted_string = self._formatString(data)
    self.log_file.writerow(formatted_string)

  # Format the CSV output
  def _formatString(self, data):
    # We will be saving this data in CSV format. Before we do, lets format it a bit here
    format_order = ['TIMESTAMP', 'TOTAL', 'STDOUT', 'STACK', 'HOSTNAME', 'MEMORY']
    formatted_text = []
    for item in format_order:
      # We have to handle python's way of formatting floats to strings specially
      if item == 'TIMESTAMP':
        formatted_text.append('%.6f' % data[item])
      else:
        formatted_text.append(data[item])
    return formatted_text

class Agent:
  """
Each agent object contains its own sampled log data. The Agent class is responsible for
collecting and storing data. machine_id is used to identify the agent.

machine_id is supplied by the client class. This allows for multiple agents if desired
"""
  def __init__(self, arguments, machine_id):
    self.arguments = arguments
    self.my_uuid = machine_id
    self.track_process = ''

    # This log object is for stdout purposes
    self.log = TemporaryFile()
    self.log_position = 0

    # Discover if --recover is being used. If so, we need to obtain the
    # timestamp of the last entry in the outfile log... a little bulky
    # to do... and not a very good place to do it.
    if self.arguments.recover:
      if os.path.exists(self.arguments.outfile[-1]):
        memory_list = []
        history_file = open(self.arguments.outfile[-1], 'r')
        reader = csv.reader(history_file, delimiter=',', quotechar='|', escapechar='\\', quoting=csv.QUOTE_MINIMAL)

        # Get last item in list. Unfortunately, no way to do this until
        # we have read the entire file...? Lucky for us, most memory log
        # files are in the single digit megabytes
        for row in reader:
          memory_list.append(row)
        history_file.close()
        last_entry = float(memory_list[-1][0]) + self.arguments.repeat_rate[-1]
        self.delta = (GetTime().now - last_entry)
      else:
        print 'Recovery options detected, but I could not find your previous memory log file.'
        sys.exit(1)
    else:
      self.delta = 0

    # Create the dictionary to which all sampled data will be stored
    # NOTE: REQUEST dictionary items are instructions (arguments) we will
    # ask the server to provide (if we are running with --pbs)
    # Simply add them here. We _can not_ make the arguments match the
    # server exactly, this would cause every agent launched to perform
    # like a server... bad stuff

    # Example: We added repeat_rate (see dictionary below). Now every
    # agent would update their repeat_rate according to what the user
    # supplied as an argument (--repeat_rate 0.02)
    self.agent_data = { self.my_uuid :
                        { 'HOSTNAME'  : socket.gethostname(),
                          'STDOUT'    : '',
                          'STACK'     : '',
                          'MEMORY'    : 0,
                          'TIMESTAMP' : GetTime().now - self.delta,
                          'REQUEST'   : { 'run'          : '',
                                          'pstack'       : '',
                                          'repeat_rate'  : '',
                                          'cwd'          : ''},
                          'STOP'      : False,
                          'TOTAL'     : 0,
                          'DEBUG_LOG' : ''
                        }
                      }

  # NOTE: This is the only function that should be called in this class
  def takeSample(self):
    if self.arguments.pstack:
      self.agent_data[self.my_uuid]['STACK'] = self._getStack()

    # Always do the following
    self.agent_data[self.my_uuid]['MEMORY'] = self._getMemory()
    self.agent_data[self.my_uuid]['STDOUT'] = self._getStdout()
    if self.arguments.recover:
      self.agent_data[self.my_uuid]['TIMESTAMP'] = GetTime().now - self.delta
    else:
      self.agent_data[self.my_uuid]['TIMESTAMP'] = GetTime().now

    # Return the data to whom ever asked for it
    return self.agent_data

  def _getStdout(self):
    self.log.seek(self.log_position)
    output = self.log.read()
    self.log_position = self.log.tell()
    sys.stdout.write(output)
    return output

  def _getMemory(self):
    tmp_pids = self._getPIDs()
    memory_usage = 0
    if tmp_pids != {}:
      for single_pid in tmp_pids.iteritems():
        memory_usage += int(single_pid[1][0])
      if memory_usage == 0:
        # Memory usage hit zero? Then assume the binary being tracked has exited. So lets begin doing the same.
        self.agent_data[self.my_uuid]['DEBUG_LOG'] = 'I found the total memory usage of all my processes hit 0. Stoping'
        self.agent_data[self.my_uuid]['STOP'] = True
        return 0
      return int(memory_usage)
    # No binay even detected? Lets assume it exited, so we should begin doing the same.
    self.agent_data[self.my_uuid]['STOP'] = True
    self.agent_data[self.my_uuid]['DEBUG_LOG'] = 'I found no processes running. Stopping'
    return 0

  def _getStack(self):
    # A quick way to safely check for the avilability of needed tools
    self._verifyCommand(['pstack'])

    tmp_pids = self._getPIDs()
    if self._darwin == True:
      return ''
    else:
      # We are Linux
      if tmp_pids != {}:
        lowest_pid = sorted([x for x in tmp_pids.keys()])[0]
        # This is very expensive on some systems (~1 second). Hence we only do this once with the first item found
        tmp_proc = subprocess.Popen([which('pstack'), str(lowest_pid)], stdout=subprocess.PIPE, stderr=subprocess.PIPE)
        return tmp_proc.communicate()[0]
      else:
        return ''

  def _getPIDs(self):
    pid_list = {}

    # Determin the binary to sample and store it. Doing the findCommand is a little expensive.
    if self.track_process == '':
      self.track_process = self._findCommand(''.join(self.arguments.run))

    # A quick way to safely check for the avilability of needed tools
    self._verifyCommand(['ps'])

    # If we are tracking a binary
    if self.arguments.run:
      command = [which('ps'), '-e', '-o', 'pid,rss,user,args']
      tmp_proc = subprocess.Popen(command, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
      all_pids = tmp_proc.communicate()[0].split('\n')
      # Figure out what we are allowed to track (strip away mpiexec, processes not owned by us, etc)
      for single_pid in all_pids:
        if single_pid.find(self.track_process) != -1 and \
           single_pid.find(__file__) == -1 and \
           single_pid.find('mpirun') == -1 and \
           single_pid.find(os.getenv('USER')) != -1 and \
           single_pid.find('mpiexec') == -1:
          pid_list[int(single_pid.split()[0])] = []
          pid_list[int(single_pid.split()[0])].extend([single_pid.split()[1], single_pid.split()[3]])<commit_msg>added LLDB and GDB class handlers to obtain stack traces. LLDB for Darwin, GDB for Linux. refs #3616<commit_after>#!/usr/bin/env python
from tempfile import TemporaryFile, SpooledTemporaryFile
import os, sys, re, socket, time, pickle, csv, uuid, subprocess, argparse, decimal, select, platform

class LLDB:
  def _parseStackTrace(self, jibberish):
    not_jibberish = re.findall(r'\(lldb\) bt(.*)\(lldb\)', jibberish, re.DOTALL)
    if len(not_jibberish) != 0:
      return not_jibberish[0].replace(' frame ', '')
    else:
      return 'Stack Trace failed:', jibberish

  def _waitForResponse(self, wait=True):
    while wait:
      self.lldb_stdout.seek(self.last_position)
      for line in self.lldb_stdout:
        if line == '(lldb) ':
          self.last_position = self.lldb_stdout.tell()
          return True
      time.sleep(0.05)
    time.sleep(0.05)
    return True

  def getStackTrace(self, pid):
    lldb_commands = [ 'attach -p ' + pid + '\n', 'bt\n', 'quit\n', 'Y\n' ]
    self.lldb_stdout = SpooledTemporaryFile()
    self.last_position = 0
    lldb_process = subprocess.Popen(['lldb', '-x'], stdin=subprocess.PIPE, stdout=self.lldb_stdout, stderr=self.lldb_stdout)
    while lldb_process.poll() == None:
      for command in lldb_commands:
        if command == lldb_commands[-1]:
          lldb_commands = []
          if self._waitForResponse(False):
            lldb_process.stdin.write(command)
        elif self._waitForResponse():
          lldb_process.stdin.write(command)
    self.lldb_stdout.seek(0)
    stack_trace = self._parseStackTrace(self.lldb_stdout.read())
    self.lldb_stdout.close()
    return stack_trace

class GDB:
  def _parseStackTrace(self, jibberish):
    not_jibberish = re.findall(r'\(gdb\) (#.*)\(gdb\)', jibberish, re.DOTALL)
    if len(not_jibberish) != 0:
      return not_jibberish[0]
    else:
      return 'Stack Trace failed:', jibberish

  def _waitForResponse(self, wait=True):
    while wait:
      self.gdb_stdout.seek(self.last_position)
      for line in self.gdb_stdout:
        if line == '(gdb) ':
          self.last_position = self.gdb_stdout.tell()
          return True
      time.sleep(0.05)
    time.sleep(0.05)
    return True

  def getStackTrace(self, pid):
    gdb_commands = [ 'attach ' + pid + '\n', 'set verbose off\n', 'thread\n', 'apply\n', 'all\n', 'bt\n', 'quit\n', 'y\n' ]
    self.gdb_stdout = SpooledTemporaryFile()
    self.last_position = 0
    gdb_process = subprocess.Popen(['gdb', '-nx'], stdin=subprocess.PIPE, stdout=self.gdb_stdout, stderr=self.gdb_stdout)
    while gdb_process.poll() == None:
      for command in gdb_commands:
        if command == gdb_commands[-1]:
          gdb_commands = []
        elif self._waitForResponse():
          gdb_process.stdin.write(command)
    self.gdb_stdout.seek(0)
    stack_trace = self._parseStackTrace(self.gdb_stdout.read())
    self.gdb_stdout.close()
    return stack_trace

class Server:
  def __init__(self, arguments):
    self.arguments = arguments
    self.arguments.cwd = os.getcwd()
    # Test to see if we are starting as a server
    if self.arguments.pbs == True:
      if os.getenv('PBS_NODEFILE') != None:
        # Initialize an agent, strictly for holding our stdout logs. Give it the UUID of 'server'
        self.agent = Agent(self.arguments, 'server')
        if self.arguments.recover:
          self.logfile = WriteCSV(self.arguments.outfile[0], False)
        else:
          self.logfile = WriteCSV(self.arguments.outfile[0], True)
        self.client_connections = []
        self.startServer()
      else:
        print 'I could not find your PBS_NODEFILE. Is PBS loaded?'
        sys.exit(1)
    # If we are not a server, start the single client
    else:
      self.startClient()

  def startServer(self):
    # Setup the TCP socket
    self.server_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    self.server_socket.bind((socket.gethostname(), 0))
    self.server_socket.listen(5)
    (self.host, self.port) = self.server_socket.getsockname()

    # We will store all connections (sockets objects) made to the server in a list
    self.client_connections.append(self.server_socket)

    # Launch the actual binary we want to track
    self._launchJob()

    # Now launch all pbs agents
    self._launchClients()

    # This is a try so we can handle a keyboard ctrl-c
    try:

      # Continue to listen and accept active connections from agents
      # until all agents report a STOP command.
      AGENTS_ACTIVE = True
      while AGENTS_ACTIVE:
        read_sockets, write_sockets, error_sockets = select.select(self.client_connections,[],[])
        for sock in read_sockets:
          if sock == self.server_socket:
            # Accept an incomming connection
            self.client_connections.append(self.server_socket.accept()[0])
          else:
            # Deal with the data being sent to the server by its agents
            self.handleAgent()

            # Check to see if _all_ agents are telling the server to stop
            agent_count = len(self.agent.agent_data.keys())
            current_count = 0
            for agent in self.agent.agent_data.keys():
              if self.agent.agent_data[agent]['STOP']:
                current_count += 1

            # if All Agents have reported a STOP command, begin to exit
            if current_count == agent_count:
              AGENTS_ACTIVE = False
              # Gotta get out of the for loop somehow...
              break

          # Sleep a bit before reading additional data
          time.sleep(self.arguments.repeat_rate[-1])

      # Close the server socket
      self.server_socket.close()

      # Close the logfile as the server is about to exit
      self.logfile.close()

    # Cancel server operations if ctrl-c was pressed
    except KeyboardInterrupt:
      print 'Canceled by user. Wrote log:', self.arguments.outfile[0]
      sys.exit(0)

    # Normal exiting procedures
    print '\n\nAll agents have stopped. Log file saved to:', self.arguments.outfile[0]
    sys.exit(0)

  def startClient(self):
    Client(self.arguments)

  def _launchClients(self):
    # Read the environment PBS_NODEFILE
    self._PBS_NODEFILE = open(os.getenv('PBS_NODEFILE'), 'r')
    nodes = set(self._PBS_NODEFILE.read().split())

    # Print some useful information about our setup
    print 'Memory Logger running on Host:', self.host, 'Port:', self.port, '\nNodes:', ', '.join(nodes), '\nSample rate (including stdout):', self.arguments.repeat_rate[-1], 's (use --repeat-rate to adjust)\nRemote agents delaying', self.arguments.pbs_delay[-1], 'second/s before tracking. (use --pbs-delay to adjust)\n'

    # Build our command list based on the PBS_NODEFILE
    command = []
    for node in nodes:
      command.append([  'ssh', node,
                   'bash --login -c "source /etc/profile && ' \
                   + 'sleep ' + str(self.arguments.pbs_delay[-1]) + ' && ' \
                   + os.path.abspath(__file__) \
                   + ' --call-back-host ' \
                   + self.host + ' ' + str(self.port) \
                   + '"'])

    # remote into each node and execute another copy of memory_logger.py
    # with a call back argument to recieve further instructions
    for pbs_node in command:
      subprocess.Popen(pbs_node, stdout=subprocess.PIPE, stderr=subprocess.PIPE)

  # Launch the binary we intend to track
  def _launchJob(self):
    subprocess.Popen(self.arguments.run[-1].split(), stdout=self.agent.log, stderr=self.agent.log)

  # A connection has been made from client to server
  # Capture that data, and determin what to do with it
  def handleAgent(self):
    # Loop through all client connections, and receive data if any
    for agent_socket in self.client_connections:

      # Completely ignore the server_socket object
      if agent_socket == self.server_socket:
        continue

      # Assign an AgentConnector for the task of handling data between client and server
      reporting_agent = AgentConnector(self.arguments, agent_socket)

      # OK... get data from a client and begin
      new_data = reporting_agent.readData()
      if new_data != None:
        # There should be only one dictionary key (were reading data from just one client at a time)
        agent_uuid = new_data.keys()[0]

        # Update our dictionary of an agents data
        self.agent.agent_data[agent_uuid] = new_data[agent_uuid]

        # Modify incoming Agents timestamp to match Server's time (because every node is a little bit off)
        if self.arguments.recover:
          self.agent.agent_data[agent_uuid]['TIMESTAMP'] = GetTime().now - self.agent.delta
        else:
          self.agent.agent_data[agent_uuid]['TIMESTAMP'] = GetTime().now

        # update total usage for all known reporting agents
        total_usage = 0
        for one_agent in self.agent.agent_data.keys():
          total_usage += self.agent.agent_data[one_agent]['MEMORY']
        self.agent.agent_data[agent_uuid]['TOTAL'] = int(total_usage)

        # Get any stdout thats happened thus far and apply it to what ever agent just sent us data
        self.agent.agent_data[agent_uuid]['STDOUT'] = self.agent._getStdout()

        # Write to our logfile
        self.logfile.write(self.agent.agent_data[agent_uuid])

        # Check for any agents sending a stop command. If we find one,
        # set some zeroing values, and close that agent's socket.
        if self.agent.agent_data[agent_uuid]['STOP']:
          self.agent.agent_data[agent_uuid]['MEMORY'] = 0
          agent_socket.close()
          if agent_socket != self.server_socket:
            self.client_connections.remove(agent_socket)

          # Go ahead and set our server agent to STOP as well.
          # The server will continue recording samples from agents
          self.agent.agent_data['server']['STOP'] = True

        # If an Agent has made a request for instructions, handle it here
        update_client = False
        if new_data[agent_uuid]['REQUEST'] != None:
          for request in new_data[agent_uuid]['REQUEST'].iteritems():
            if new_data[agent_uuid]['REQUEST'][request[0]] == '':
              update_client = True
              # We only support sending any arguments supplied to ther server, back to the agent
              for request_type in dir(self.arguments):
                if request[0] == str(request_type):
                  self.agent.agent_data[agent_uuid]['REQUEST'][request[0]] = getattr(self.arguments, request[0])

          # If an Agent needed additional instructions, go ahead and re-send those instructions
          if update_client:
            reporting_agent.sendData(self.agent.agent_data[agent_uuid])

class Client:
  def __init__(self, arguments):
    self.arguments = arguments

    # Initialize an Agent with a UUID based on our hostname
    self.my_agent = Agent(arguments, str(uuid.uuid3(uuid.NAMESPACE_DNS, socket.gethostname())))

    # Initialize an AgentConnector
    self.remote_server = AgentConnector(self.arguments)

    # If client will talk to a server (PBS)
    if self.arguments.call_back_host:
      # We know by initializing an agent, agent_data contains the necessary message asking for further instructions
      self.my_agent.agent_data[self.my_agent.my_uuid] = self.remote_server.sendData(self.my_agent.agent_data)

      # Apply new instructions received from server (this basically updates our arguments)
      for request in self.my_agent.agent_data[self.my_agent.my_uuid]['REQUEST'].iteritems():
        for request_type in dir(self.arguments):
          if request[0] == str(request_type):
            setattr(self.arguments, request[0], request[1])

      # Requests have been satisfied, set to None
      self.my_agent.agent_data[self.my_agent.my_uuid]['REQUEST'] = None

      # Change to the same directory as the server was when initiated (needed for PBS stuff)
      os.chdir(self.arguments.cwd)

    # Client will not be talking to a server, save data to a file instead
    else:
      # Deal with --recover
      if self.arguments.recover:
        # Do not overwrite the file
        self.logfile = WriteCSV(self.arguments.outfile[0], False)
      else:
        # Overwrite the file
        self.logfile = WriteCSV(self.arguments.outfile[0], True)

    # Lets begin!
    self.startProcess()

  # This function handles the starting and stoping of the sampler process.
  # We loop until an agent returns a stop command.
  def startProcess(self):
    AGENTS_ACTIVE = True

    # If we know we are the only client, go ahead and start the process we want to track.
    if self.arguments.call_back_host == None:
      subprocess.Popen(self.arguments.run[-1].split(), stdout=self.my_agent.log, stderr=self.my_agent.log)
      # Delay just a bit to keep from recording a possible zero memory usage as the binary starts up
      time.sleep(self.arguments.sample_delay[0])


    # This is a try so we can handle a keyboard ctrl-c
    try:

      # Continue to process data until an Agent reports a STOP command
      while AGENTS_ACTIVE:
        # Take a sample
        current_data = self.my_agent.takeSample()

        # Handle the data supplied by the Agent.
        self._handleData(current_data)

        # If an Agent reported a STOP command, go ahead and begin the shutdown phase
        if current_data[current_data.keys()[0]]['STOP']:
          AGENTS_ACTIVE = False

        # Sleep just a bit between samples, as to not saturate the machine
        time.sleep(self.arguments.repeat_rate[-1])

      # An agent reported a stop command... so let everyone know where the log was saved, and exit!
      if self.arguments.call_back_host == None:
        print 'Binary has exited. Wrote log:', self.arguments.outfile[0]

    # Cancel server operations if ctrl-c was pressed
    except KeyboardInterrupt:
      self.logfile.close()
      print 'Canceled by user. Wrote log:', self.arguments.outfile[0]
      sys.exit(0)

    # Everything went smooth.
    sys.exit(0)

  # Figure out what to do with the sampled data
  def _handleData(self, data):
    # Sending the sampled data to a server
    if self.arguments.call_back_host:
      self.remote_server.sendData(data)
    # Saving the sampled data to a file
    else:
      # Compute the TOTAL memory usage to be how much our one agent reported
      # Because were the only client doing any work
      data[self.my_agent.my_uuid]['TOTAL'] = data[self.my_agent.my_uuid]['MEMORY']
      self.logfile.write(data[self.my_agent.my_uuid])

      # If the agent has been told to stop, close the database file
      if self.my_agent.agent_data[self.my_agent.my_uuid]['STOP'] == True:
        self.logfile.close()

class AgentConnector:
  """
Functions used to communicate to and from Client and Server.
Both Client and Server classes use this object.

readData()
sendData('message', socket_connection=None)

if sendData's socket_connection is None, it will create a new connection to the server
based on supplied arguments
"""
  def __init__(self, arguments, connection=None):
    self.arguments = arguments
    self.connection = connection
    self.CREATED_CONNECTION = False

    # If the connection is None, meaning this object was instanced by a client,
    # we must create a connection to the server first
    if self.connection == None and self.arguments.call_back_host != None:
      self.CREATED_CONNECTION = True
      self.connection = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
      self.connection.settimeout(15)
      self.connection.connect((self.arguments.call_back_host[0], int(self.arguments.call_back_host[1])))

  # read all data sent by an agent
  def readData(self):
    # Get how much data there is to receive
    # The first eight bytes is our data length
    data_width = int(self.connection.recv(8))
    tmp_received = ''

    # We need to receive precisely the ammount of data the
    # client is trying to send us.
    while len(tmp_received) < data_width:
      if data_width - len(tmp_received) > 1024:
        tmp_received += self.connection.recv(1024)
      else:
        tmp_received += self.connection.recv(data_width - (len(tmp_received)))

    # unpickle the received message
    return self._unpickleMessage(tmp_received)

  # send data to an agent
  def sendData(self, message):
    # pickle the data up, and send the message
    self.connection.sendall(self._pickleMessage(message))

    # If we had to create the socket (connection was none), and this client/agent is requesting
    # instructions, go ahead and read the data that _better be there_ sent to us by the server.
    if self.CREATED_CONNECTION and message[message.keys()[0]]['REQUEST'] != None:
      return self.readData()

  # The following two functions pickle up the data for easy socket transport
  def _pickleMessage(self, message):
    t = TemporaryFile()
    pickle.dump(message, t)
    t.seek(0)
    str_msg = t.read()
    str_len = len(str_msg)
    message = "%-8d" % (str_len,) + str_msg
    return message

  def _unpickleMessage(self, message):
    t = TemporaryFile()
    t.write(message)
    t.seek(0)
    try:
      return pickle.load(t)
    except KeyError:
      print 'Socket data was not pickled data: ', message
    except:
      raise

class WriteCSV:
  def __init__(self, logfile, overwrite):
    if overwrite:
      self.file_object = open(logfile, 'w', 1)
    else:
      self.file_object = open(logfile, 'a', 1)
    self.log_file = csv.writer(self.file_object, delimiter=',', quotechar='|', escapechar='\\', quoting=csv.QUOTE_MINIMAL)

  # Close the logfile
  def close(self):
    self.file_object.close()

  # Write a CSV row
  def write(self, data):
    formatted_string = self._formatString(data)
    self.log_file.writerow(formatted_string)

  # Format the CSV output
  def _formatString(self, data):
    # We will be saving this data in CSV format. Before we do, lets format it a bit here
    format_order = ['TIMESTAMP', 'TOTAL', 'STDOUT', 'STACK', 'HOSTNAME', 'MEMORY']
    formatted_text = []
    for item in format_order:
      # We have to handle python's way of formatting floats to strings specially
      if item == 'TIMESTAMP':
        formatted_text.append('%.6f' % data[item])
      else:
        formatted_text.append(data[item])
    return formatted_text

class Agent:
  """
Each agent object contains its own sampled log data. The Agent class is responsible for
collecting and storing data. machine_id is used to identify the agent.

machine_id is supplied by the client class. This allows for multiple agents if desired
"""
  def __init__(self, arguments, machine_id):
    self.arguments = arguments
    self.my_uuid = machine_id
    self.track_process = ''

    # This log object is for stdout purposes
    self.log = TemporaryFile()
    self.log_position = 0

    # Discover if --recover is being used. If so, we need to obtain the
    # timestamp of the last entry in the outfile log... a little bulky
    # to do... and not a very good place to do it.
    if self.arguments.recover:
      if os.path.exists(self.arguments.outfile[-1]):
        memory_list = []
        history_file = open(self.arguments.outfile[-1], 'r')
        reader = csv.reader(history_file, delimiter=',', quotechar='|', escapechar='\\', quoting=csv.QUOTE_MINIMAL)

        # Get last item in list. Unfortunately, no way to do this until
        # we have read the entire file...? Lucky for us, most memory log
        # files are in the single digit megabytes
        for row in reader:
          memory_list.append(row)
        history_file.close()
        last_entry = float(memory_list[-1][0]) + self.arguments.repeat_rate[-1]
        self.delta = (GetTime().now - last_entry)
      else:
        print 'Recovery options detected, but I could not find your previous memory log file.'
        sys.exit(1)
    else:
      self.delta = 0

    # Create the dictionary to which all sampled data will be stored
    # NOTE: REQUEST dictionary items are instructions (arguments) we will
    # ask the server to provide (if we are running with --pbs)
    # Simply add them here. We _can not_ make the arguments match the
    # server exactly, this would cause every agent launched to perform
    # like a server... bad stuff

    # Example: We added repeat_rate (see dictionary below). Now every
    # agent would update their repeat_rate according to what the user
    # supplied as an argument (--repeat_rate 0.02)
    self.agent_data = { self.my_uuid :
                        { 'HOSTNAME'  : socket.gethostname(),
                          'STDOUT'    : '',
                          'STACK'     : '',
                          'MEMORY'    : 0,
                          'TIMESTAMP' : GetTime().now - self.delta,
                          'REQUEST'   : { 'run'          : '',
                                          'pstack'       : '',
                                          'repeat_rate'  : '',
                                          'cwd'          : ''},
                          'STOP'      : False,
                          'TOTAL'     : 0,
                          'DEBUG_LOG' : ''
                        }
                      }

  # NOTE: This is the only function that should be called in this class
  def takeSample(self):
    if self.arguments.pstack:
      self.agent_data[self.my_uuid]['STACK'] = self._getStack()

    # Always do the following
    self.agent_data[self.my_uuid]['MEMORY'] = self._getMemory()
    self.agent_data[self.my_uuid]['STDOUT'] = self._getStdout()
    if self.arguments.recover:
      self.agent_data[self.my_uuid]['TIMESTAMP'] = GetTime().now - self.delta
    else:
      self.agent_data[self.my_uuid]['TIMESTAMP'] = GetTime().now

    # Return the data to whom ever asked for it
    return self.agent_data

  def _getStdout(self):
    self.log.seek(self.log_position)
    output = self.log.read()
    self.log_position = self.log.tell()
    sys.stdout.write(output)
    return output

  def _getMemory(self):
    tmp_pids = self._getPIDs()
    memory_usage = 0
    if tmp_pids != {}:
      for single_pid in tmp_pids.iteritems():
        memory_usage += int(single_pid[1][0])
      if memory_usage == 0:
        # Memory usage hit zero? Then assume the binary being tracked has exited. So lets begin doing the same.
        self.agent_data[self.my_uuid]['DEBUG_LOG'] = 'I found the total memory usage of all my processes hit 0. Stoping'
        self.agent_data[self.my_uuid]['STOP'] = True
        return 0
      return int(memory_usage)
    # No binay even detected? Lets assume it exited, so we should begin doing the same.
    self.agent_data[self.my_uuid]['STOP'] = True
    self.agent_data[self.my_uuid]['DEBUG_LOG'] = 'I found no processes running. Stopping'
    return 0

  def _getStack(self):
    if self._darwin() == True:
      stack_trace = LLDB()
    else:
      stack_trace = GDB()
    tmp_pids = self._getPIDs()
    if tmp_pids != {}:
      lowest_pid = sorted([x for x in tmp_pids.keys()])[0]
      return stack_trace.getStackTrace(str(lowest_pid))
    else:
      return ''

  def _getPIDs(self):
    pid_list = {}

    # Determin the binary to sample and store it. Doing the findCommand is a little expensive.
    if self.track_process == '':
      self.track_process = self._findCommand(''.join(self.arguments.run))

    # A quick way to safely check for the avilability of needed tools
    self._verifyCommand(['ps'])

    # If we are tracking a binary
    if self.arguments.run:
      command = [which('ps'), '-e', '-o', 'pid,rss,user,args']
      tmp_proc = subprocess.Popen(command, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
      all_pids = tmp_proc.communicate()[0].split('\n')
      # Figure out what we are allowed to track (strip away mpiexec, processes not owned by us, etc)
      for single_pid in all_pids:
        if single_pid.find(self.track_process) != -1 and \
           single_pid.find(__file__) == -1 and \
           single_pid.find('mpirun') == -1 and \
           single_pid.find(os.getenv('USER')) != -1 and \
           single_pid.find('mpiexec') == -1:
          pid_list[int(single_pid.split()[0])] = []
          pid_list[int(single_pid.split()[0])].extend([single_pid.split()[1], single_pid.split()[3]])
{'fraction_non_alphanumeric': 0.0988289991092572, 'fraction_numerical': 0.0044754393969019535, 'mean_word_length': 2.6785742827459442, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 12, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4534436', 'n_tokens_mistral': 13753, 'n_tokens_neox': 13333, 'n_words': 4215}