chore: import upstream snapshot with attribution

python/ray/experimental/raysort/tracing_utils.py

@@ -0,0 +1,120 @@
+import datetime
+import functools
+import logging
+import os
+import time
+from typing import List, Tuple
+
+import ray
+from ray.experimental.raysort import constants, logging_utils
+from ray.util.metrics import Gauge, Histogram
+
+HISTOGRAM_BOUNDARIES = list(range(50, 200, 50))
+
+
+def timeit(
+    event: str,
+    report_time=False,
+    report_in_progress=True,
+    report_completed=True,
+):
+    def decorator(f):
+        @functools.wraps(f)
+        def wrapped_f(*args, **kwargs):
+            progress_tracker = ray.get_actor(constants.PROGRESS_TRACKER_ACTOR)
+            progress_tracker.inc.remote(f"{event}_in_progress", echo=report_in_progress)
+            try:
+                start = time.time()
+                ret = f(*args, **kwargs)
+                end = time.time()
+                duration = end - start
+                progress_tracker.observe.remote(
+                    f"{event}_time",
+                    duration,
+                    echo=report_time,
+                )
+                progress_tracker.inc.remote(f"{event}_completed", echo=report_completed)
+                return ret
+            finally:
+                progress_tracker.dec.remote(f"{event}_in_progress")
+
+        return wrapped_f
+
+    return decorator
+
+
+def get_metrics(_args):
+    return {
+        "gauges": [
+            "map_in_progress",
+            "merge_in_progress",
+            "reduce_in_progress",
+            "sort_in_progress",
+            "map_completed",
+            "merge_completed",
+            "reduce_completed",
+            "sort_completed",
+        ],
+        "histograms": [
+            ("map_time", HISTOGRAM_BOUNDARIES),
+            ("merge_time", HISTOGRAM_BOUNDARIES),
+            ("reduce_time", HISTOGRAM_BOUNDARIES),
+            ("sort_time", HISTOGRAM_BOUNDARIES),
+        ],
+    }
+
+
+def create_progress_tracker(args):
+    return ProgressTracker.options(name=constants.PROGRESS_TRACKER_ACTOR).remote(
+        **get_metrics(args)
+    )
+
+
+@ray.remote
+class ProgressTracker:
+    def __init__(
+        self,
+        gauges: List[str],
+        histograms: List[Tuple[str, List[int]]],
+    ):
+        self.counts = {m: 0 for m in gauges}
+        self.gauges = {m: Gauge(m) for m in gauges}
+        self.reset_gauges()
+        self.histograms = {m: Histogram(m, boundaries=b) for m, b in histograms}
+        logging_utils.init()
+
+    def reset_gauges(self):
+        for g in self.gauges.values():
+            g.set(0)
+
+    def inc(self, metric_name, value=1, echo=False):
+        gauge = self.gauges.get(metric_name)
+        if gauge is None:
+            logging.warning(f"No such Gauge: {metric_name}")
+            return
+        self.counts[metric_name] += value
+        gauge.set(self.counts[metric_name])
+        if echo:
+            logging.info(f"{metric_name} {self.counts[metric_name]}")
+
+    def dec(self, metric_name, value=1, echo=False):
+        return self.inc(metric_name, -value, echo)
+
+    def observe(self, metric_name, value, echo=False):
+        histogram = self.histograms.get(metric_name)
+        if histogram is None:
+            logging.warning(f"No such Histogram: {metric_name}")
+            return
+        histogram.observe(value)
+        if echo:
+            logging.info(f"{metric_name} {value}")
+
+
+def export_timeline():
+    timestr = datetime.datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
+    filename = f"ray-timeline-{timestr}.json"
+
+    temp_dir = ray.get_runtime_context().get_temp_dir()
+    filepath = os.path.join(temp_dir, filename)
+    ray.timeline(filename=filepath)
+    logging.info(f"Exported Ray timeline to {filepath}")