nornir-automation / nornir

Pluggable multi-threaded framework with inventory management to help operate collections of devices
https://nornir.readthedocs.io/
Apache License 2.0
1.37k stars 233 forks source link

Processors: ability to know the run hosts #885

Open u1735067 opened 8 months ago

u1735067 commented 8 months ago

Hi,

processors are very powerful but one limitation is that is doesn't seem possible to know the (number of) hosts actually in the run, as the list according to on_good and on_failed is calculated (https://github.com/nornir-automation/nornir/blob/v3.4.1/nornir/core/__init__.py#L118) after task_started() is called, could it be possible to have access to this information (as well as on_good and on_failed and maybe raise_on_error)?

One hacky way to get this informations is:

import logging
from dataclasses import dataclass
from nornir.core.processor import Processor

logger = logging.getLogger(__name__)

@dataclass
class _NornirRunVars:
    on_good: bool
    on_failed: bool
    raise_on_error: bool | None

class ProgressProcessor(Processor):
    def task_started(self, task):
        _run_vars = self._get_nornir_run_vars()
        if _run_vars.on_good:
            self.run_hosts = len(task.nornir.inventory.hosts)
        else:
            self.run_hosts = len(task.nornir.data.failed_hosts)
        if not _run_vars.on_failed:
            self.run_hosts -= len(task.nornir.data.failed_hosts)

    # [..]

    def _get_nornir_run_vars(self) -> _NornirRunVars:
        """
        Get the run() `on_good`, `on_failed` and `raise_on_error` variables, to recalculate `run_on`
        Issue for a better way: https://github.com/nornir-automation/nornir/issues/885
        """
        try:
            # https://stackoverflow.com/questions/15608987/how-can-i-access-variables-from-the-caller-even-if-it-isnt-an-enclosing-scope
            import inspect
            # Find the frame containing these variables
            required_vars = ("task", "on_good", "on_failed", "raise_on_error")

            # Iterate on some frames, skipping the known not relevant ones
            frame = inspect.currentframe()
            for i in range(6):  # only allow 6 back
                if frame is None:
                    raise RuntimeError('cannot inspect stack frames')
                if i > 2:  # 0 is self, 1 is self.task_started(), 2 is Processors().task_started()
                    if all(var in frame.f_locals for var in required_vars):
                        return _NornirRunVars(
                            on_good=frame.f_locals["on_good"],
                            on_failed=frame.f_locals["on_failed"],
                            raise_on_error=frame.f_locals["raise_on_error"],
                        )
                frame = frame.f_back
            # Not found
            raise AttributeError('vars not found in outer scope')
        except Exception as e:
            logger.exception(e)
            # Assume default values
            return _NornirRunVars(
                on_good=True,
                on_failed=False,
                raise_on_error=None,
            )

Thanks

dbarrosop commented 8 months ago

The problem is that data resides in completely different parts of the stack. An easier solution would be to pass that information explicitly to your processor.

class ProgressProcessor(Processor):
    def __init__(self, on_good, on_failed, raise_on_error):
       ...

Alternatively we could move this:

https://github.com/nornir-automation/nornir/blob/v3.4.1/nornir/core/__init__.py#L116

to here:

https://github.com/nornir-automation/nornir/blob/v3.4.1/nornir/core/__init__.py#L138

and pass run_on to task_started but that would be a breaking change so we need to think carefully about it.

ubaumann commented 8 months ago

The progress bar from nornir_rich has the option to pass the number of hosts implemented. Maybe the code helps you.

https://github.com/InfrastructureAsCode-ch/nornir_rich/blob/b1661687005a15059acbbc6d0665a25550f17a45/nornir_rich/progress_bar.py#L22