Repository: spotify/luigi
File: luigi/contrib/prometheus_metric.py (Python)

Code before:
from prometheus_client import CollectorRegistry, Counter, Gauge, generate_latest, CONTENT_TYPE_LATEST
from luigi import parameter
from luigi.metrics import MetricsCollector
from luigi.task import Config


class prometheus(Config):
    use_task_family = parameter.BoolParameter(default=True, parsing=parameter.BoolParameter.EXPLICIT_PARSING)
    task_parameters = parameter.ListParameter(default=[])


class PrometheusMetricsCollector(MetricsCollector):

    def _generate_task_labels(self, task):
        return {
            label: task.family if label == "family" else task.params.get(label)
            for label in self.labels
        }

    def __init__(self):
        super(PrometheusMetricsCollector, self).__init__()
        self.registry = CollectorRegistry()
        config = prometheus()
        self.labels = list(config.task_parameters)
        if config.use_task_family:
            self.labels += ["family"]
        if not self.labels:
            raise ValueError("Prometheus labels cannot be empty (see prometheus configuration)")
        self.task_started_counter = Counter(
            'luigi_task_started_total',
            'number of started luigi tasks',
            self.labels,
            registry=self.registry
        )
        self.task_failed_counter = Counter(
            'luigi_task_failed_total',
            'number of failed luigi tasks',
            self.labels,
            registry=self.registry
        )
        self.task_disabled_counter = Counter(
            'luigi_task_disabled_total',
            'number of disabled luigi tasks',
            self.labels,
            registry=self.registry
        )
        self.task_done_counter = Counter(
            'luigi_task_done_total',

Review comment (suggestion):
I'm not sure why this is the case (probably a miss during PR review), but the initial implementation of this `PrometheusMetricsCollector` forgot to allow for `*args, **kwargs` to its `__init__()`.

In most other classes in Luigi, we allow for args/kwargs and pass `**kwargs` to the config class initializations to facilitate config overrides. e.g. https://github.com/spotify/luigi/blob/master/luigi/contrib/datadog_metric.py#L27

Code after:
from prometheus_client import CollectorRegistry, Counter, Gauge, generate_latest, CONTENT_TYPE_LATEST
from luigi import parameter
from luigi.metrics import MetricsCollector
from luigi.task import Config


class prometheus(Config):
    use_task_family_in_labels = parameter.BoolParameter(
        default=True, parsing=parameter.BoolParameter.EXPLICIT_PARSING
    )
    task_parameters_to_use_in_labels = parameter.ListParameter(default=[])


class PrometheusMetricsCollector(MetricsCollector):

    def _generate_task_labels(self, task):
        return {
            label: task.family if label == "family" else task.params.get(label)
            for label in self.labels
        }

    def __init__(self, *args, **kwargs):
        super(PrometheusMetricsCollector, self).__init__()
        self.registry = CollectorRegistry()
        config = prometheus(**kwargs)
        self.labels = list(config.task_parameters_to_use_in_labels)
        if config.use_task_family_in_labels:
            self.labels += ["family"]
        if not self.labels:
            raise ValueError("Prometheus labels cannot be empty (see prometheus configuration)")
        self.task_started_counter = Counter(
            'luigi_task_started_total',
            'number of started luigi tasks',
            self.labels,
            registry=self.registry
        )
        self.task_failed_counter = Counter(
            'luigi_task_failed_total',
            'number of failed luigi tasks',
            self.labels,
            registry=self.registry
        )
        self.task_disabled_counter = Counter(
            'luigi_task_disabled_total',
            'number of disabled luigi tasks',
            self.labels,
            registry=self.registry
        )