0
\$\begingroup\$

I have following code that provides an interface for executing commands in the localhost and the remotehost. Please help me in improving 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)
\$\endgroup\$

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Browse other questions tagged or ask your own question.