wehub-resource-sync
110 lines
3.9 KiB
Python
110 lines
3.9 KiB
Python
from typing import Any, Dict, Iterable, Tuple
|
|
|
|
import ray
|
|
from ray.data.block import Block
|
|
|
|
|
|
def _iter_sliced_blocks(
|
|
blocks: Iterable[Block], per_task_row_limit: int
|
|
) -> Iterable[Block]:
|
|
"""Iterate over blocks, accumulating rows up to the per-task row limit."""
|
|
rows_read = 0
|
|
for block in blocks:
|
|
if rows_read >= per_task_row_limit:
|
|
break
|
|
|
|
from ray.data.block import BlockAccessor
|
|
|
|
accessor = BlockAccessor.for_block(block)
|
|
block_rows = accessor.num_rows()
|
|
|
|
if rows_read + block_rows <= per_task_row_limit:
|
|
yield block
|
|
rows_read += block_rows
|
|
else:
|
|
# Slice the block to meet the limit exactly
|
|
remaining_rows = per_task_row_limit - rows_read
|
|
sliced_block = accessor.slice(0, remaining_rows, copy=True)
|
|
yield sliced_block
|
|
break
|
|
|
|
|
|
def _validate_head_node_resources_for_local_scheduling(
|
|
ray_remote_args: Dict[str, Any],
|
|
*,
|
|
op_description: str,
|
|
default_num_cpus: int = 1,
|
|
default_num_gpus: int = 0,
|
|
default_memory: int = 0,
|
|
) -> None:
|
|
"""Ensure the head node has enough resources before pinning work there.
|
|
|
|
Local paths (``local://``) and other driver-local I/O schedule tasks on the
|
|
head node via ``NodeAffinitySchedulingStrategy``. If the head node was
|
|
intentionally started with zero logical resources (a common practice to
|
|
avoid OOMs), those tasks become unschedulable. Detect this upfront and
|
|
raise a clear error with remediation steps.
|
|
"""
|
|
|
|
# Ray defaults to reserving 1 CPU per task when num_cpus isn't provided.
|
|
num_cpus = ray_remote_args.get("num_cpus", default_num_cpus)
|
|
num_gpus = ray_remote_args.get("num_gpus", default_num_gpus)
|
|
memory = ray_remote_args.get("memory", default_memory)
|
|
|
|
# Resource keys follow the Resources map of ray.nodes() (e.g., CPU, GPU, memory).
|
|
required_resources: Dict[str, float] = {}
|
|
required_resources["CPU"] = float(num_cpus)
|
|
required_resources["GPU"] = float(num_gpus)
|
|
required_resources["memory"] = float(memory)
|
|
|
|
# Include any additional custom resources requested.
|
|
custom_resources = ray_remote_args.get("resources", {})
|
|
for name, amount in custom_resources.items():
|
|
if amount is None:
|
|
continue
|
|
try:
|
|
amount = float(amount)
|
|
except (TypeError, ValueError) as err:
|
|
raise ValueError(f"Invalid resource amount for '{name}': {amount}") from err
|
|
required_resources[name] = amount
|
|
|
|
head_node = next(
|
|
(
|
|
node
|
|
for node in ray.nodes()
|
|
if node.get("Alive")
|
|
and "node:__internal_head__" in node.get("Resources", {})
|
|
),
|
|
None,
|
|
)
|
|
if not head_node:
|
|
# The head node metadata is unavailable (e.g., during shutdown). Fall back
|
|
# to the default behavior and let Ray surface its own error.
|
|
return
|
|
|
|
# Build a map of required vs available resources on the head node.
|
|
head_resources: Dict[str, float] = head_node.get("Resources", {})
|
|
# Map: resource name -> (required, available).
|
|
insufficient: Dict[str, Tuple[float, float]] = {}
|
|
for name, req in required_resources.items():
|
|
avail = head_resources.get(name, 0.0)
|
|
if avail < req:
|
|
insufficient[name] = (req, avail)
|
|
|
|
# If nothing is below the required amount, we are good to proceed.
|
|
if not insufficient:
|
|
return
|
|
|
|
details = "; ".join(
|
|
f"{name} required {req:g} but head has {avail:g}"
|
|
for name, (req, avail) in insufficient.items()
|
|
)
|
|
|
|
raise ValueError(
|
|
f"{op_description} must run on the head node (e.g., for local:// paths), "
|
|
f"but the head node doesn't have enough resources: {details}. "
|
|
"Add resources to the head node, switch to a shared filesystem instead "
|
|
"of local://, or set the resource requests on this operation to 0 "
|
|
"(for example, num_cpus=0) so it can run without head resources."
|
|
)
|