chore: import upstream snapshot with attribution

python/ray/experimental/raysort/constants.py

@@ -0,0 +1,44 @@
+import os
+
+from ray.experimental.raysort.types import ByteCount, PartId, RecordCount
+
+__DIR__ = os.path.dirname(os.path.abspath(__file__))
+
+# Basics
+RECORD_SIZE = 100  # bytes
+
+# Progress Tracker Actor
+PROGRESS_TRACKER_ACTOR = "ProgressTrackerActor"
+
+# Executable locations
+GENSORT_PATH = os.path.join(__DIR__, "bin/gensort/64/gensort")
+VALSORT_PATH = os.path.join(__DIR__, "bin/gensort/64/valsort")
+
+# Filenames
+WORK_DIR = "/tmp/raysort"
+INPUT_MANIFEST_FILE = os.path.join(WORK_DIR, "input-manifest.csv")
+OUTPUT_MANIFEST_FILE = os.path.join(WORK_DIR, "output-manifest.csv")
+DATA_DIR_FMT = {
+    "input": "{mnt}/tmp/input/",
+    "output": "{mnt}/tmp/output/",
+    "temp": "{mnt}/tmp/temp/",
+}
+FILENAME_FMT = {
+    "input": "input-{part_id:08}",
+    "output": "output-{part_id:08}",
+    "temp": "temp-{part_id:08}",
+}
+
+# Prometheus config
+PROM_RAY_EXPORTER_PORT = 8090
+PROM_NODE_EXPORTER_PORT = 8091
+
+
+# Convenience functions
+def bytes_to_records(n_bytes: ByteCount) -> RecordCount:
+    assert n_bytes % RECORD_SIZE == 0
+    return int(n_bytes / RECORD_SIZE)
+
+
+def merge_part_ids(reducer_id: PartId, mapper_id: PartId) -> PartId:
+    return reducer_id * 1_000_000 + mapper_id

python/ray/experimental/raysort/logging_utils.py

@@ -0,0 +1,9 @@
+import logging
+
+
+def init():
+    fmt = "%(levelname)s %(asctime)s %(filename)s:%(lineno)s] %(message)s"
+    logging.basicConfig(
+        format=fmt,
+        level=logging.INFO,
+    )

python/ray/experimental/raysort/main.py

@@ -0,0 +1,475 @@
+import argparse
+import contextlib
+import csv
+import logging
+import os
+import random
+import subprocess
+import tempfile
+from typing import Callable, Dict, Iterable, List
+
+import numpy as np
+
+import ray
+from ray.experimental.raysort import constants, logging_utils, sortlib, tracing_utils
+from ray.experimental.raysort.types import (
+    BlockInfo,
+    ByteCount,
+    PartId,
+    PartInfo,
+    Path,
+    RecordCount,
+)
+from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy
+
+Args = argparse.Namespace
+
+# ------------------------------------------------------------
+#     Parse Arguments
+# ------------------------------------------------------------
+
+
+def get_args(*args, **kwargs):
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--ray_address",
+        default="auto",
+        type=str,
+        help="if set to None, will launch a local Ray cluster",
+    )
+    parser.add_argument(
+        "--total_data_size",
+        default=1 * 1000 * 1024 * 1024 * 1024,
+        type=ByteCount,
+        help="total data size in bytes",
+    )
+    parser.add_argument(
+        "--num_mappers",
+        default=256,
+        type=int,
+        help="number of map tasks",
+    )
+    parser.add_argument(
+        "--num_mappers_per_round",
+        default=16,
+        type=int,
+        help="number of map tasks per first-stage merge tasks",
+    )
+    parser.add_argument(
+        "--num_reducers",
+        default=16,
+        type=int,
+        help="number of second-stage reduce tasks",
+    )
+    parser.add_argument(
+        "--num_concurrent_rounds",
+        default=4,
+        type=int,
+        help="max number of rounds of map/merge tasks in flight",
+    )
+    parser.add_argument(
+        "--reducer_input_chunk",
+        default=100 * 1024 * 1024,
+        type=ByteCount,
+        help="bytes to read from each file in reduce tasks",
+    )
+    parser.add_argument(
+        "--skip_sorting",
+        default=False,
+        action="store_true",
+        help="if set, no sorting is actually performed",
+    )
+    parser.add_argument(
+        "--skip_input",
+        default=False,
+        action="store_true",
+        help="if set, mappers will not read data from disk",
+    )
+    parser.add_argument(
+        "--skip_output",
+        default=False,
+        action="store_true",
+        help="if set, reducers will not write out results to disk",
+    )
+    # Which tasks to run?
+    tasks_group = parser.add_argument_group(
+        "tasks to run", "if no task is specified, will run all tasks"
+    )
+    tasks = ["generate_input", "sort", "validate_output"]
+    for task in tasks:
+        tasks_group.add_argument(f"--{task}", action="store_true")
+
+    args = parser.parse_args(*args, **kwargs)
+    # Derive additional arguments.
+    args.input_part_size = ByteCount(args.total_data_size / args.num_mappers)
+    assert args.num_mappers % args.num_mappers_per_round == 0
+    args.num_rounds = int(args.num_mappers / args.num_mappers_per_round)
+    args.mount_points = _get_mount_points()
+    # If no tasks are specified, run all tasks.
+    args_dict = vars(args)
+    if not any(args_dict[task] for task in tasks):
+        for task in tasks:
+            args_dict[task] = True
+    return args
+
+
+def _get_mount_points():
+    default_ret = [tempfile.gettempdir()]
+    mnt = "/mnt"
+    if os.path.exists(mnt):
+        ret = [os.path.join(mnt, d) for d in os.listdir(mnt)]
+        if len(ret) > 0:
+            return ret
+    return default_ret
+
+
+# ------------------------------------------------------------
+#     Generate Input
+# ------------------------------------------------------------
+
+
+def _part_info(args: Args, part_id: PartId, kind="input") -> PartInfo:
+    node = ray._private.worker.global_worker.node_ip_address
+    mnt = random.choice(args.mount_points)
+    filepath = _get_part_path(mnt, part_id, kind)
+    return PartInfo(part_id, node, filepath)
+
+
+def _get_part_path(mnt: Path, part_id: PartId, kind="input") -> Path:
+    assert kind in {"input", "output", "temp"}
+    dir_fmt = constants.DATA_DIR_FMT[kind]
+    dirpath = dir_fmt.format(mnt=mnt)
+    os.makedirs(dirpath, exist_ok=True)
+    filename_fmt = constants.FILENAME_FMT[kind]
+    filename = filename_fmt.format(part_id=part_id)
+    filepath = os.path.join(dirpath, filename)
+    return filepath
+
+
+@ray.remote
+def generate_part(
+    args: Args, part_id: PartId, size: RecordCount, offset: RecordCount
+) -> PartInfo:
+    logging_utils.init()
+    pinfo = _part_info(args, part_id)
+    subprocess.run(
+        [constants.GENSORT_PATH, f"-b{offset}", f"{size}", pinfo.path], check=True
+    )
+    logging.info(f"Generated input {pinfo}")
+    return pinfo
+
+
+def generate_input(args: Args):
+    if args.skip_input:
+        return
+    size = constants.bytes_to_records(args.input_part_size)
+    offset = 0
+    tasks = []
+    for part_id in range(args.num_mappers):
+        tasks.append(generate_part.remote(args, part_id, size, offset))
+        offset += size
+    assert offset == constants.bytes_to_records(args.total_data_size), args
+    logging.info(f"Generating {len(tasks)} partitions")
+    parts = ray.get(tasks)
+    with open(constants.INPUT_MANIFEST_FILE, "w") as fout:
+        writer = csv.writer(fout)
+        writer.writerows(parts)
+
+
+# ------------------------------------------------------------
+#     Sort
+# ------------------------------------------------------------
+
+
+def _load_manifest(args: Args, path: Path) -> List[PartInfo]:
+    if args.skip_input:
+        return [PartInfo(i, None, None) for i in range(args.num_mappers)]
+    with open(path) as fin:
+        reader = csv.reader(fin)
+        return [PartInfo(int(part_id), node, path) for part_id, node, path in reader]
+
+
+def _load_partition(args: Args, path: Path) -> np.ndarray:
+    if args.skip_input:
+        return np.frombuffer(
+            np.random.bytes(args.input_part_size), dtype=np.uint8
+        ).copy()
+    return np.fromfile(path, dtype=np.uint8)
+
+
+def _dummy_sort_and_partition(
+    part: np.ndarray, boundaries: List[int]
+) -> List[BlockInfo]:
+    N = len(boundaries)
+    offset = 0
+    size = int(np.ceil(part.size / N))
+    blocks = []
+    for _ in range(N):
+        blocks.append((offset, size))
+        offset += size
+    return blocks
+
+
+@ray.remote
+@tracing_utils.timeit("map")
+def mapper(
+    args: Args, mapper_id: PartId, boundaries: List[int], path: Path
+) -> List[np.ndarray]:
+    logging_utils.init()
+    part = _load_partition(args, path)
+    sort_fn = (
+        _dummy_sort_and_partition if args.skip_sorting else sortlib.sort_and_partition
+    )
+    blocks = sort_fn(part, boundaries)
+    return [part[offset : offset + size] for offset, size in blocks]
+
+
+def _dummy_merge(
+    num_blocks: int, _n: int, get_block: Callable[[int, int], np.ndarray]
+) -> Iterable[np.ndarray]:
+    blocks = [((i, 0), get_block(i, 0)) for i in range(num_blocks)]
+    while len(blocks) > 0:
+        (m, d), block = blocks.pop(random.randrange(len(blocks)))
+        yield block
+        d_ = d + 1
+        block = get_block(m, d_)
+        if block is None:
+            continue
+        blocks.append(((m, d_), block))
+
+
+def _merge_impl(
+    args: Args,
+    M: int,
+    pinfo: PartInfo,
+    get_block: Callable[[int, int], np.ndarray],
+    skip_output=False,
+):
+    merge_fn = _dummy_merge if args.skip_sorting else sortlib.merge_partitions
+    merger = merge_fn(M, get_block)
+
+    if skip_output:
+        for datachunk in merger:
+            del datachunk
+    else:
+        with open(pinfo.path, "wb") as fout:
+            for datachunk in merger:
+                fout.write(datachunk)
+    return pinfo
+
+
+# See worker_placement_groups() for why `num_cpus=0`.
+@ray.remote(num_cpus=0, resources={"worker": 1})
+@tracing_utils.timeit("merge")
+def merge_mapper_blocks(
+    args: Args, reducer_id: PartId, mapper_id: PartId, *blocks: List[np.ndarray]
+) -> PartInfo:
+    part_id = constants.merge_part_ids(reducer_id, mapper_id)
+    pinfo = _part_info(args, part_id, kind="temp")
+    M = len(blocks)
+
+    def get_block(i, d):
+        if i >= M or d > 0:
+            return None
+        return blocks[i]
+
+    return _merge_impl(args, M, pinfo, get_block)
+
+
+# See worker_placement_groups() for why `num_cpus=0`.
+@ray.remote(num_cpus=0, resources={"worker": 1})
+@tracing_utils.timeit("reduce")
+def final_merge(
+    args: Args, reducer_id: PartId, *merged_parts: List[PartInfo]
+) -> PartInfo:
+    M = len(merged_parts)
+
+    def _load_block_chunk(pinfo: PartInfo, d: int) -> np.ndarray:
+        return np.fromfile(
+            pinfo.path,
+            dtype=np.uint8,
+            count=args.reducer_input_chunk,
+            offset=d * args.reducer_input_chunk,
+        )
+
+    def get_block(i, d):
+        ret = _load_block_chunk(merged_parts[i], d)
+        if ret.size == 0:
+            return None
+        return ret
+
+    pinfo = _part_info(args, reducer_id, "output")
+    return _merge_impl(args, M, pinfo, get_block, args.skip_output)
+
+
+def _node_res(node: str) -> Dict[str, float]:
+    return {"resources": {f"node:{node}": 1e-3}}
+
+
+@contextlib.contextmanager
+def worker_placement_groups(args: Args) -> List[ray.PlacementGroupID]:
+    """
+    Returns one placement group per node with a `worker` resource. To run
+    tasks in the placement group, use
+    `@ray.remote(num_cpus=0, resources={"worker": 1})`. Ray does not
+    automatically reserve CPU resources, so tasks must specify `num_cpus=0`
+    in order to run in a placement group.
+    """
+    pgs = [ray.util.placement_group([{"worker": 1}]) for _ in range(args.num_reducers)]
+    ray.get([pg.ready() for pg in pgs])
+    try:
+        yield pgs
+    finally:
+        for pg in pgs:
+            ray.util.remove_placement_group(pg)
+
+
+@tracing_utils.timeit("sort", report_time=True)
+def sort_main(args: Args):
+    parts = _load_manifest(args, constants.INPUT_MANIFEST_FILE)
+    assert len(parts) == args.num_mappers
+    boundaries = sortlib.get_boundaries(args.num_reducers)
+
+    # The exception of 'ValueError("Resource quantities >1 must be whole numbers.")'
+    # will be raised if the `num_cpus` > 1 and not an integer.
+    num_cpus = os.cpu_count() / args.num_concurrent_rounds
+    if num_cpus > 1.0:
+        num_cpus = int(num_cpus)
+    mapper_opt = {
+        "num_returns": args.num_reducers,
+        "num_cpus": num_cpus,
+    }  # Load balance across worker nodes by setting `num_cpus`.
+    merge_results = np.empty((args.num_rounds, args.num_reducers), dtype=object)
+
+    part_id = 0
+    with worker_placement_groups(args) as pgs:
+        for round in range(args.num_rounds):
+            # Limit the number of in-flight rounds.
+            num_extra_rounds = round - args.num_concurrent_rounds + 1
+            if num_extra_rounds > 0:
+                ray.wait(
+                    [f for f in merge_results.flatten() if f is not None],
+                    num_returns=num_extra_rounds * args.num_reducers,
+                )
+
+            # Submit map tasks.
+            mapper_results = np.empty(
+                (args.num_mappers_per_round, args.num_reducers), dtype=object
+            )
+            for _ in range(args.num_mappers_per_round):
+                _, node, path = parts[part_id]
+                m = part_id % args.num_mappers_per_round
+                mapper_results[m, :] = mapper.options(**mapper_opt).remote(
+                    args, part_id, boundaries, path
+                )
+                part_id += 1
+
+            # Submit merge tasks.
+            merge_results[round, :] = [
+                merge_mapper_blocks.options(
+                    scheduling_strategy=PlacementGroupSchedulingStrategy(
+                        placement_group=pgs[r]
+                    )
+                ).remote(args, r, round, *mapper_results[:, r].tolist())
+                for r in range(args.num_reducers)
+            ]
+
+            # Delete local references to mapper results.
+            mapper_results = None
+
+        # Submit second-stage reduce tasks.
+        reducer_results = [
+            final_merge.options(
+                scheduling_strategy=PlacementGroupSchedulingStrategy(
+                    placement_group=pgs[r]
+                )
+            ).remote(args, r, *merge_results[:, r].tolist())
+            for r in range(args.num_reducers)
+        ]
+        reducer_results = ray.get(reducer_results)
+
+    if not args.skip_output:
+        with open(constants.OUTPUT_MANIFEST_FILE, "w") as fout:
+            writer = csv.writer(fout)
+            writer.writerows(reducer_results)
+
+    logging.info(ray._private.internal_api.memory_summary(stats_only=True))
+
+
+# ------------------------------------------------------------
+#     Validate Output
+# ------------------------------------------------------------
+
+
+def _run_valsort(args: List[str]):
+    proc = subprocess.run([constants.VALSORT_PATH] + args, capture_output=True)
+    if proc.returncode != 0:
+        logging.critical("\n" + proc.stderr.decode("ascii"))
+        raise RuntimeError(f"Validation failed: {args}")
+
+
+@ray.remote
+def validate_part(path: Path):
+    logging_utils.init()
+    sum_path = path + ".sum"
+    _run_valsort(["-o", sum_path, path])
+    logging.info(f"Validated output {path}")
+    with open(sum_path, "rb") as fin:
+        return os.path.getsize(path), fin.read()
+
+
+def validate_output(args: Args):
+    if args.skip_sorting or args.skip_output:
+        return
+    partitions = _load_manifest(args, constants.OUTPUT_MANIFEST_FILE)
+    results = []
+    for _, node, path in partitions:
+        results.append(validate_part.options(**_node_res(node)).remote(path))
+    logging.info(f"Validating {len(results)} partitions")
+    results = ray.get(results)
+    total = sum(s for s, _ in results)
+    assert total == args.total_data_size, total - args.total_data_size
+    all_checksum = b"".join(c for _, c in results)
+    with tempfile.NamedTemporaryFile() as fout:
+        fout.write(all_checksum)
+        fout.flush()
+        _run_valsort(["-s", fout.name])
+    logging.info("All OK!")
+
+
+# ------------------------------------------------------------
+#     Main
+# ------------------------------------------------------------
+
+
+def init(args: Args):
+    if not args.ray_address:
+        ray.init(resources={"worker": os.cpu_count()})
+    else:
+        ray.init(address=args.ray_address)
+    logging_utils.init()
+    logging.info(args)
+    os.makedirs(constants.WORK_DIR, exist_ok=True)
+    resources = ray.cluster_resources()
+    logging.info(resources)
+    args.num_workers = resources["worker"]
+    progress_tracker = tracing_utils.create_progress_tracker(args)
+    return progress_tracker
+
+
+def main(args: Args):
+    # Keep the actor handle in scope for the duration of the program.
+    _progress_tracker = init(args)  # noqa F841
+
+    if args.generate_input:
+        generate_input(args)
+
+    if args.sort:
+        sort_main(args)
+
+    if args.validate_output:
+        validate_output(args)
+
+
+if __name__ == "__main__":
+    main(get_args())

python/ray/experimental/raysort/sortlib.py

@@ -0,0 +1,28 @@
+from typing import Callable, Iterable, List
+
+import numpy as np
+
+from ray.experimental.raysort.types import BlockInfo
+
+
+def get_boundaries(num_parts: int) -> List[int]:
+    return [0] * num_parts
+
+
+def sort_and_partition(part: np.ndarray, boundaries: List[int]) -> List[BlockInfo]:
+    N = len(boundaries)
+    offset = 0
+    size = int(np.ceil(part.size / N))
+    blocks = []
+    for _ in range(N):
+        blocks.append((offset, size))
+        offset += size
+    return blocks
+
+
+def merge_partitions(
+    num_blocks: int, get_block: Callable[[int, int], np.ndarray]
+) -> Iterable[memoryview]:
+    blocks = [get_block(i, 0) for i in range(num_blocks)]
+    for block in blocks:
+        yield block

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}")

python/ray/experimental/raysort/types.py

@@ -0,0 +1,18 @@
+from typing import NamedTuple, Tuple
+
+ByteCount = int
+NodeAddress = str
+PartId = int
+Path = str
+RecordCount = int
+
+BlockInfo = Tuple[int, int]
+
+
+class PartInfo(NamedTuple):
+    part_id: PartId
+    node: NodeAddress
+    path: Path
+
+    def __repr__(self):
+        return f"Part({self.node}:{self.path})"