<s> """<STR_LIT>""" <EOL> import functools <EOL> import logging <EOL> import os <EOL> import posixpath <EOL> import threading <EOL> import uuid <EOL> from perfkitbenchmarker import events <EOL> from perfkitbenchmarker import flags <EOL> from perfkitbenchmarker import vm_util <EOL> flags . DEFINE_boolean ( '<STR_LIT>' , False , <EOL> '<STR_LIT>' <EOL> '<STR_LIT>' <EOL> '<STR_LIT>' ) <EOL> flags . DEFINE_integer ( '<STR_LIT>' , None , <EOL> '<STR_LIT>' <EOL> '<STR_LIT>' ) <EOL> flags . DEFINE_string ( '<STR_LIT>' , None , <EOL> '<STR_LIT>' <EOL> '<STR_LIT>' <EOL> '<STR_LIT>' ) <EOL> class _DStatCollector ( object ) : <EOL> """<STR_LIT>""" <EOL> def __init__ ( self , interval = None , output_directory = None ) : <EOL> """<STR_LIT>""" <EOL> self . interval = interval <EOL> self . output_directory = output_directory or vm_util . GetTempDir ( ) <EOL> self . _lock = threading . Lock ( ) <EOL> self . _pids = { } <EOL> self . _file_names = { } <EOL> if not os . path . isdir ( self . output_directory ) : <EOL> raise IOError ( '<STR_LIT>' . format ( <EOL> self . output_directory ) ) <EOL> def _StartOnVm ( self , vm , suffix = '<STR_LIT>' ) : <EOL> vm . Install ( '<STR_LIT>' ) <EOL> num_cpus = vm . num_cpus <EOL> block_devices , _ = vm . RemoteCommand ( <EOL> '<STR_LIT>' ) <EOL> block_devices = block_devices . splitlines ( ) <EOL> dstat_file = posixpath . join ( <EOL> vm_util . VM_TMP_DIR , '<STR_LIT>' . format ( vm . name , suffix ) ) <EOL> cmd = ( '<STR_LIT>' <EOL> '<STR_LIT>' <EOL> '<STR_LIT>' <EOL> '<STR_LIT>' <EOL> '<STR_LIT>' <EOL> '<STR_LIT>' ) . format ( <EOL> max_cpu = num_cpus - <NUM_LIT:1> , <EOL> block_devices = '<STR_LIT:U+002C>' . join ( block_devices ) , <EOL> output = dstat_file , <EOL> dstat_interval = self . interval or '<STR_LIT>' ) <EOL> stdout , _ = vm . RemoteCommand ( cmd ) <EOL> with self . _lock : <EOL> self . _pids [ vm . name ] = stdout . strip ( ) <EOL> self . _file_names [ vm . name ] = dstat_file <EOL> def _StopOnVm ( self , vm ) : <EOL> """<STR_LIT>""" <EOL> if vm . name not in self . _pids : <EOL> logging . warn ( '<STR_LIT>' , vm . name ) <EOL> return <EOL> else : <EOL> with self . _lock : <EOL> pid = self . _pids . pop ( vm . name ) <EOL> file_name = self . _file_names . pop ( vm . name ) <EOL> cmd = '<STR_LIT>' . format ( pid ) <EOL> vm . RemoteCommand ( cmd ) <EOL> try : <EOL> vm . PullFile ( self . output_directory , file_name ) <EOL> except Exception : <EOL> logging . exception ( '<STR_LIT>' , vm . name ) <EOL> def Start ( self , sender , benchmark_spec ) : <EOL> """<STR_LIT>""" <EOL> suffix = '<STR_LIT>' . format ( benchmark_spec . uid , <EOL> str ( uuid . uuid4 ( ) ) [ : <NUM_LIT:8> ] ) <EOL> start_on_vm = functools . partial ( self . _StartOnVm , suffix = suffix ) <EOL> vm_util . RunThreaded ( start_on_vm , benchmark_spec . vms ) <EOL> def Stop ( self , sender , benchmark_spec ) : <EOL> """<STR_LIT>""" <EOL> vm_util . RunThreaded ( self . _StopOnVm , benchmark_spec . vms ) <EOL> def Register ( parsed_flags ) : <EOL> """<STR_LIT>""" <EOL> if not parsed_flags . dstat : <EOL> return <EOL> output_directory = ( parsed_flags . dstat_output <EOL> if parsed_flags [ '<STR_LIT>' ] . present <EOL> else vm_util . GetTempDir ( ) ) <EOL> logging . debug ( '<STR_LIT>' , <EOL> parsed_flags . dstat_interval , output_directory ) <EOL> if not os . path . isdir ( output_directory ) : <EOL> os . makedirs ( output_directory ) <EOL> collector = _DStatCollector ( interval = parsed_flags . dstat_interval , <EOL> output_directory = output_directory ) <EOL> events . before_phase . connect ( collector . Start , events . RUN_PHASE , weak = False ) <EOL> events . after_phase . connect ( collector . Stop , events . RUN_PHASE , weak = False ) </s>
