Repository: spotify/luigi

Issue:
Add Prometheus contrib for monitoring purpose
## Description
Prometheus is a tool that scrapes exposed metrics endpoints and allows for you to visualize and alert on the metrics. Adding this contrib will allow for users to monitor Luigi tasks via Prometheus.

## Motivation and Context
Based on the status change of a task, Prometheus metrics are updated with the parameters that were used to run that specific task.

The implementation idea was strongly based on the merged PR for Datadog #2434.

## Have you tested this? If so, how?
I have included unit tests.


Discussion:
@Tarrasch Yes, this has been tested in production as well.
Cool. I'm fine to merge this in 1-2 weeks then I think, so other's have a chance to try it out or give input.
@victoriaalee  Has this integration still be stable for your uses?
Hi all!

I've discovered some discrepancies between the exported Prometheus metric and the task history runtimes displayed in the `/tasklist/:name` endpoint.

I noticed that the `_update_task_history` method is called at different places than `set_status` (which also handles calling the functions to update metrics), and I'm wondering if this might be the source of the discrepancies. Is there any documentation or could I receive some guidance on why this is the case? Intuitively, it seems that updating task history and setting status of a task should go hand-in-hand.

Thanks in advance.
> or could I receive some guidance on why this is the case?

Unfortunately that's one part of the code bases I don't know well and frankly I think it could use some love. Also I think it's written before it was an obvious thing to submit your code with rigorous unit tests.

Were you to refactor/rewrite any part of it (doesn't need to be perfect, just better) I'm happy to merge.
Gotcha. I'll be on vacation for the next two weeks, but I'll take some time afterward to delve into the logic behind task status transitions and see if I can resolve the discrepancies and get this PR merged!
Thank you all for your patience. I did some investigation with the jobs being run at my company and their expected runtimes and found that this Prometheus contrib produces accurate results whereas the `/tasklist` endpoint does not.

Please let me know what else needs to be done to get this merged. Thanks!
Can you find anyone to review this (just for style, as it seems to work already)? Maybe even at your own company?

Otherwise I'm happy to merge.
@Tarrasch The code has been reviewed by a colleague at my company.
We have an internal “thing” that puts our Luigi task history into Prometheus, mind if I take a look at this next week before merge? @jessicaaustin
We are looking forward it as well.
Looks good, :+1: for a merge.

What do you see as the best path to support a run-time defined `MetricCollector`?  I was thinking about allowing the path of a `MetricCollector` subclass as input to the config parameter (like `my.own.MetricsCollector`) and in the `else` of the `MetricsCollectors.get` we could try to load the passed in config option as a class using `importlib`. That would open up the option for custom metrics classes external from the luigi codebase.
@victoriaalee, @kwilcox Hello, which plans do you have related to the PR?
Can I help you somehow?

I believe it would be awesome to get prometheus support in next release!

Thank you for your contributions and make life easier. 
@kwilcox I love your idea to allow the use of custom metrics classes external to the Luigi codebase. I know at my company we've been running off of a fork with this contrib for a while, and this would've would've been a more elegant solution.

@dbolshak As far as I'm aware this PR is good to merge. @Tarrasch?
Thanks for everyone's input and this contribution @victoriaalee! Perhaps we (I) should finally make a new release of luigi after Easter. 

Resolving patch:
diff --git a/luigi/contrib/prometheus_metric.py b/luigi/contrib/prometheus_metric.py
--- a/luigi/contrib/prometheus_metric.py
+++ b/luigi/contrib/prometheus_metric.py
@@ -1,4 +1,4 @@
-from prometheus_client import CollectorRegistry, Counter, Gauge, generate_latest
+from prometheus_client import CollectorRegistry, Counter, Gauge, generate_latest, CONTENT_TYPE_LATEST
 from luigi.metrics import MetricsCollector
 
 
@@ -58,3 +58,6 @@ def handle_task_done(self, task):
         # time_running can be `None` if task was already complete
         if task.time_running is not None:
             self.task_execution_time.labels(family=task.family).set(task.updated - task.time_running)
+
+    def configure_http_handler(self, http_handler):
+        http_handler.set_header('Content-Type', CONTENT_TYPE_LATEST)
diff --git a/luigi/metrics.py b/luigi/metrics.py
--- a/luigi/metrics.py
+++ b/luigi/metrics.py
@@ -53,6 +53,9 @@ def handle_task_done(self, task):
     def generate_latest(self):
         return
 
+    def configure_http_handler(self, http_handler):
+        pass
+
 
 class NoMetricsCollector(MetricsCollector):
     """Empty MetricsCollector when no collector is being used
diff --git a/luigi/server.py b/luigi/server.py
--- a/luigi/server.py
+++ b/luigi/server.py
@@ -45,7 +45,6 @@
 import time
 
 import pkg_resources
-from prometheus_client import CONTENT_TYPE_LATEST
 import tornado.httpserver
 import tornado.ioloop
 import tornado.netutil
@@ -282,8 +281,8 @@ def initialize(self, scheduler):
     def get(self):
         metrics = self._scheduler._state._metrics_collector.generate_latest()
         if metrics:
+            metrics.configure_http_handler(self)
             self.write(metrics)
-            self.set_header('Content-Type', CONTENT_TYPE_LATEST)
 
 
 def app(scheduler):


Test patch:
diff --git a/test/contrib/prometheus_metric_test.py b/test/contrib/prometheus_metric_test.py
--- a/test/contrib/prometheus_metric_test.py
+++ b/test/contrib/prometheus_metric_test.py
@@ -1,10 +1,17 @@
 from helpers import unittest
 from nose.plugins.attrib import attr
+from prometheus_client import CONTENT_TYPE_LATEST
 
 from luigi.contrib.prometheus_metric import PrometheusMetricsCollector
 from luigi.metrics import MetricsCollectors
 from luigi.scheduler import Scheduler
 
+try:
+    from unittest import mock
+except ImportError:
+    import mock
+
+
 WORKER = 'myworker'
 TASK_ID = 'TaskID'
 TASK_FAMILY = 'TaskFamily'
@@ -68,3 +75,8 @@ def test_handle_task_done(self):
 
         assert self.collector.registry.get_sample_value(counter_name, labels=labels) == 1
         assert self.collector.registry.get_sample_value(gauge_name, labels=labels) == task.updated - task.time_running
+
+    def test_configure_http_handler(self):
+        mock_http_handler = mock.MagicMock()
+        self.collector.configure_http_handler(mock_http_handler)
+        mock_http_handler.set_header.assert_called_once_with('Content-Type', CONTENT_TYPE_LATEST)
diff --git a/test/server_test.py b/test/server_test.py
--- a/test/server_test.py
+++ b/test/server_test.py
@@ -448,3 +448,27 @@ def tearDown(self):
         shutil.rmtree(self.server_client.tempdir)
 
     server_client_class = ServerClient
+
+
+class MetricsHandlerTest(unittest.TestCase):
+    def setUp(self):
+        self.mock_scheduler = mock.MagicMock()
+        self.handler = luigi.server.MetricsHandler(tornado.web.Application(), mock.MagicMock(),
+                                                   scheduler=self.mock_scheduler)
+
+    def test_initialize(self):
+        self.assertIs(self.handler._scheduler, self.mock_scheduler)
+
+    def test_get(self):
+        mock_metrics = mock.MagicMock()
+        self.mock_scheduler._state._metrics_collector.generate_latest.return_value = mock_metrics
+        with mock.patch.object(self.handler, 'write') as patched_write:
+            self.handler.get()
+            patched_write.assert_called_once_with(mock_metrics)
+            mock_metrics.configure_http_handler.assert_called_once_with(self.handler)
+
+    def test_get_no_metrics(self):
+        self.mock_scheduler._state._metrics_collector.generate_latest.return_value = None
+        with mock.patch.object(self.handler, 'write') as patched_write:
+            self.handler.get()
+            patched_write.assert_not_called()