I have the following code that provides an interface to execute commands on the local host and the remote host. Please help me improve this code.
class CommandExecutor(object):
def __init__(self, command_spec):
"""
:param command_spec:
:type command_spec: C{domain.spec.command_spec.CommandSpec}
"""
self._command_spec = command_spec
self._command_agent = None
self._rusage_before_command_execution = None
self._job_service = JobServiceImpl()
@property
def _command_id(self):
""" Returns the command id """
checkpoint_spec = self._command_spec.parent
spec = checkpoint_spec.parent
job_id = JobId(name=spec._name,
group=spec.group,
job_key=JobKey.from__key(spec._key),
dr_mode=get_current_domain()
)
checkpoint_id = CheckpointId(checkpoint_name=checkpoint_spec.name,
job_id=job_id,
checkpoint_index=checkpoint_spec.index)
return CommandId(checkpoint_id=checkpoint_id,
command_index=self._command_spec.index)
def execute(self):
""" Executes the command """
try:
command_run_result = self.run_command()
except (PersistenceWriteError, HostIsNotReachableError) as e:
command_run_result = COMMAND_RUN_RESULT(state=CommandExecutionState.FAILED,
failure_reason=e.message)
if command_run_result.state == CommandExecutionState.SUCCESS:
return CommandExecutionResult(state=CommandExecutionState.SUCCESS)
return CommandExecutionResult(state=CommandExecutionState.FAILED,
failure_reason=command_run_result.failure_reason)
def run_command(self):
""" Runs command agent to launch the command """
# Do all pre run tasks required for running the command
command_pre_run_result = self._do_pre_run_tasks()
if command_pre_run_result.state == CommandExecutionState.FAILED_ON_DEPENDENCY:
return COMMAND_RUN_RESULT(state=command_pre_run_result.state,
failure_reason=command_pre_run_result.failure_reason)
# Try to run the command precisely at given time.
self.sleep_until(self._command_spec.try_to_run_at_exactly)
command_to_execute = self._get_command()
# Add command rerun args if the command is not running
# for the first time
if self._command_spec.cmd_rerun_args is not None and
self._get_command_number_of_execution() > 0:
command_to_execute = command_to_execute + self._command_spec.cmd_rerun_args
# Mark the command start to update all new command state
self._mark_command_start(command_to_execute)
# Launch the command agent as async
self.launch_command_agent(
self.__get_command_agent_launcher_info(command_to_execute)
)
# Wait for the command to complete or timeout
self._wait_till_command_agent_complete(
timeout=get_field_value(self._command_spec,
"cmd_timeout"),
timeout_action=self._command_timeout_action
)
# Do all post run tasks required after completion of the command
# like cleaning up tasks
self._do_post_run_tasks()
return COMMAND_RUN_RESULT(state=CommandExecutionState.SUCCESS,
failure_reason=None)
def _do_pre_run_tasks(self):
"""
List of tasks to be done before running the command
1. Checks whether the dependency is met
2. Creates the command output store in ds
3. Creates the command process info store in ds
"""
command_dependency_result = self.check_command_dependencies()
if not command_dependency_result.is_met:
return COMMAND_PRE_RUN_RESULT(state=CommandExecutionState.FAILED_ON_DEPENDENCY,
failure_reason=command_dependency_result.failure_reason)
self._command_output_store = self._create_command_output_store()
self._command_process_info_store = self._create_command_process_info_store()
self._rusage_before_command_execution = resource.getrusage(resource.RUSAGE_SELF)
return COMMAND_PRE_RUN_RESULT(state=CommandExecutionState.SUCCESS,
failure_reason=None)
def check_command_dependencies(self):
"""
Checks whether all command dependencies are met or not.
"""
checkpoint_id = self._command_id.checkpoint_id
dependency_manager = DependencyManager()
dependencies = self.get_all_command_dependencies()
logging.info("Checking the dependencies for the command " + self._get_command())
_dependency_result = dependency_manager.check_dependencies(dependencies, checkpoint_id)
return _dependency_result
def get_all_command_dependencies(self):
"""
Returns all command dependencies.
:rtype: C{list}
"""
# Get all dependencies defined in the config
dependencies = get_field_value(self._command_spec, 'dependency') or ()
return dependencies
def _create_command_output_store(self):
""" Creates required command output store """
command_output_store_info = CommandOutputStoreInfo(
stdout_content_type=get_field_value(self._command_spec, "output_type"),
stderr_content_type=DEFAULT_COMMAND_OUTPUT_CONTENT_TYPE,
command_id=self._command_id,
command_name=self._get_command_short_name(self._get_command())
)
try:
command_output_store =
self._job_service.create_command_output_store(
command_output_store_info
)
except PersistenceWriteError:
logging.ferror(
"Unable to create output store for the command: " +
self._get_command())
raise
return command_output_store
def _create_command_process_info_store(self):
""" Creates the command process info store """
try:
command_process_info_store =
self._job_service.create_command_process_info_store(
command_id=self._command_id,
command_name=self._get_command_short_name(self._get_command())
)
except PersistenceWriteError:
logging.ferror(
"Unable to create process info store for the command: " +
self._get_command())
raise
return command_process_info_store
def sleep_until(self, wait_until):
"""
Sleeps until the wait time.
:param wait_until: Time to wait until
:type wait_until: C{time}
:rtype: None
"""
if wait_until is None:
return
now = datetime.now()
# We just have time parameter, check whether time has past
wait_time_delta = datetime.combine(date.today(), wait_until) - now
total_seconds_to_wait = wait_time_delta.total_seconds()
if total_seconds_to_wait > 0:
logging.info(
"Sleeping for {total_seconds_to_wait} seconds till time %s" %
wait_until)
sleep(total_seconds_to_wait)
def _get__job_command_state(self):
""" Returns the command state. """
return self._job_service.get_command(self._command_id)
def _get_command_number_of_execution(self):
""" Returns the number of times the command is executed in the past. """
command_state = self._get__job_command_state()
if command_state:
return command_state.num_cmd_executions
# Command is getting executed for first time
return 0
def _get_command(self):
""" Returns the command to be executed """
return self._command_spec.exec_str
def _get_environments_for_command(self):
""" List of envs for the command executor agent """
env = dict()
# env.update(self._command_spec.parent.parent.get_user_command_env())
# set the INFO_date from job_key date this will be used by the user's command
job_key = self._command_spec.parent.parent._key
# job_key = self._command_spec.parent.parent._job_key
if job_key.is_set('date') and not self._command_spec.use_original_env:
env('INFO_date') = str(job_key.date)
return env
def _get_environment_for_command_agent(self):
""" Returns the list of env to be passed to the command agent """
env = dict()
for env_name in ENV_TO_EXPORT_TO_COMMAND_AGENT:
if os.environ.get(env_name) is not None:
env(env_name) = os.environ(env_name)
env('_DIR') = '/u/nagarajn/python-/'
return env
@property
def _command_resource_limits(self):
""" Returns the command resource limits for the command """
return CommandResourceLimits(
memory_size=get_field_value(self._command_spec, "cmd_memory_limit"),
core_file_size=get_field_value(self._command_spec, "max_core_dump_size"),
file_size=get_field_value(self._command_spec, "filesize_limit")
)
def _mark_command_start(self, command):
""" Marks the command start for execution """
logging.info("Marking the command start for the command %s" % self._command_id)
now = Timestamp.now()
self._update_command_state(dict(
actual_start_time=now,
updated_time=now,
num_cmd_executions=self._get_command_number_of_execution() + 1,
stderr_file=self._command_output_store.stderr,
stdout_file=self._command_output_store.stdout,
hostup_file=self._command_process_info_store.hostname,
cmd_user=self._command_spec.cmd_user,
actual_cmd=command,
cmd_host=get_field_value(self._command_spec, "cmd_host"),
cmd_timeout_time=get_field_value(self._command_spec, "cmd_timeout"),
filesize_limit=get_field_value(self._command_spec, "filesize_limit"),
short_cmd=self._get_command_short_name(self._get_command()),
output_type=get_field_value(self._command_spec, "output_type").value
))
def __get_command_agent_launcher_info(self, command):
"""
Returns the command agent launcher info to execute the command
"""
command_info = CommandInfo(
command=command,
resource_limits=self._command_resource_limits,
command_user=self._command_spec.cmd_user,
file_creation_mask=get_field_value(self._command_spec, "umask"),
environment=self._get_environments_for_command()
)
command_agent_info = CommandAgentInfo(
command_info=command_info,
process_info_store=self._command_process_info_store,
chdir=get_field_value(self._command_spec, "cmd_exec_dir")
)
return CommandAgentLauncherInfo(
command_agent_info=command_agent_info,
command_output_store=self._command_output_store,
command_id=self._command_id,
environment=self._get_environment_for_command_agent()
)
@abstractmethod
def launch_command_agent(self, command_agent_info):
pass
def _wait_till_command_agent_complete(self,
timeout=None,
timeout_action=None):
"""
Waits until the command executor agent completes
:param timeout:
Command timeout for the command
:type timeout:
C{float}
:param timeout_action:
Action callback to be executed on timeout
:type timeout_action:
C{function}
:rtype:
C{None}
"""
start = time.time()
timeout_action_completed = False
last_periodic_update_time = start
while not self._command_agent.poll():
current_time = time.time()
if not timeout_action_completed and current_time - start >= timeout:
# Do timeout action
timeout_action_completed = True
try:
logging.finfo(
self._get_command() +
" command is running for more than {timeout}s")
logging.info("Performing the timeout action")
timeout_action()
# Timeout call back can throw any exception,
# we catch them all and continue.
except Exception as err:
logging.fexception(err)
if current_time - last_periodic_update_time >= PERIODIC_UPDATE_TIME_INTERVAL:
logging.info(
"Performing periodic tasks for the command: " +
self._get_command())
last_periodic_update_time = current_time
self._do_periodic_tasks()
sleep(POLL_TIME_INTERVAL)
def is_command_success(self, returncode):
return False, "Completed"
def do_command_post_processing(self):
pass
@classmethod
@memoize_method
def _get_command_short_name(cls, command):
"""
Returns the program/short name of the command.
:rtype: C{str}
"""
if command is None or not isinstance(command, six.string_types):
raise ValueError("Command should be a valid string")
for pattern in SHORT_COMMAND_NAME_PATTERN:
command = pattern.sub(EMPTY_STRING, command)
command_name = os.path.basename(command.split()(0))
return command_name.split(".")(0)
def _update_command_state(self, new_command_state):
"""
Updates the current command state with new state
:param new_command_state:
New command state fields
:type new_command_state:
C{dict}
:rtype:
C{None}
"""
job_id = self._command_id.checkpoint_id.job_id
# Get all commands state for the job
commands_state = self._job_service.get_job_commands(job_id)
new_commands_state = {
command_state.command_id.command_index: command_state
for command_state in commands_state
}
# Update new state for the command
command_state = copy.copy(
new_commands_state(self._command_id.command_index)
)
command_state.update(new_command_state)
new_commands_state(self._command_id.command_index) = command_state
self._job_service.update_command_state(
new_commands_state.values()
)
def _update_command_last_active_time(self):
"""
Updates the command last active time (for heartbeat)
"""
try:
self._update_command_state({
'updated_time': Timestamp.now()
})
except PersistenceWriteError:
logging.fwarn("Failed to update the command active time")
def _do_periodic_tasks(self):
"""
Does all periodic tasks
Tasks:
(1) Sending the heartbeat by updating the command file
(2) Monitoring the command agent launched
"""
self._update_command_last_active_time()
def notify_command_timeout(self):
pass
def _command_timeout_action(self):
"""
Do tasks required after command times out
We notify the user if only_notify_on_cmd_timeout is
enabled. Otherwise we terminate the command and mark
the command state as 'failed'.
"""
if self._command_spec.only_notify_on_cmd_timeout:
logging.info(
"Sending notification for the command running beyond expected time")
# Notify the command is running for more than expected time.
self.notify_command_timeout()
return
# Kill the command and mark them as failed.
self._command_agent.terminate()
def _add_header_to_stdout_content(self):
"""
Appends header to the output content if
the command has produced any stdout content.
"""
if self._command_spec.add_header:
pass
def _do_post_run_tasks(self):
""" Tasks to be done after the command execution is done """
pass
def terminate(self):
""" Terminates the command agent launched """
if self._command_agent is not None:
self._command_agent.terminate()
class LocalCommandExecutor(CommandExecutor):
def launch_command_agent(self, command_agent_info):
self._command_agent = LocalCommandAgentLauncher(command_agent_info)
class RemoteCommandExecutor(CommandExecutor):
def launch_command_agent(self, command_agent_info):
cmd_host = get_field_value(self._command_spec, "cmd_host")
self._command_agent = RemoteCommandAgentLauncher(cmd_host,
command_agent_info)
