wehub-resource-sync
476 lines
14 KiB
Python
476 lines
14 KiB
Python
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())
|