[Data][1/N] add external shuffle runtime library

[ShockYoungCHN]

Adds the file-transport runtime for a new external (on-disk) shuffle variant: TCP wire protocol, per-node ShuffleManager actor, connection / fetch primitives, and error classification.

This PR lands only the runtime library and its testers required by external shuffle tasks. It's not wired into the plan yet and there's no user-visible behavior change.

description

• Reducers open one keep-alive TCP connection per source ShuffleManager, authenticate with a per-shuffle token, and stream FETCH responses straight to disk (fetch_into via _PwriteSink).
• ShuffleManager will be started by map task to serve reducer's file read request. Upon init, ShuffleManager will bind file serving function to a random available port.
• Fault tolerance is driven by Ray's actor state (ActorDied / ActorUnavailable / ActorUnschedulable). On mid-fetch TCP failure, the reducer compares endpoints: if the manager was restarted (where the endpoint are likely to be different) we retry in place; if the endpoint is unchanged (which is Ray RPC still works but our TCP path doesn't) it's almost always a network-configuration problem (NetworkPolicy, firewall, routing) that retries to the same manager can't fix, so we surface a terminal ShuffleManagerAnomalyError.
• ShuffleManager wraps its serve loop with os._exit(1) so a server-thread crash surfaces as ActorUnavailableError, thus Ray auto-restart.

Additional information

The complete implementation of shuffle is in #64733, including runtime libraries, external shuffle task definition, plan wiring and a user-visible flag. I plan to do three PRs (1 is runtime libraries, 2 is task definition and wiring, and 3 is the user-visible flag) to gradually merge the complete implementation.

[gemini-code-assist]

Code Review

This pull request introduces an external-shuffle file-transport runtime (external_shuffle_runtime.py) implementing a TCP wire protocol, a per-node ShuffleManager actor, prefetch layout, and associated unit and integration tests. The review feedback focuses on improving robustness, security, and resource management. Key recommendations include reading socket data in chunks and implementing a discard helper to prevent uncontrolled memory allocation, validating incoming lengths and counts (such as num_ranges and msg_len) to avoid denial-of-service vectors, setting appropriate socket timeouts on both client and server sides, adding a retry mechanism with exponential backoff for transient network errors, and dynamically supporting IPv6 environments.

[gemini-code-assist]

Set a timeout on the accepted socket to prevent leaked or idle connections from hanging indefinitely and exhausting server resources. Also, catch general Exception to avoid unhandled exceptions bubbling up and printing noisy tracebacks to stderr.

Suggested change

	def handle(self):
	srv = self.server
	sock = self.request
	_tune_shuffle_socket(sock)
	try:
	if not self._handshake(sock, srv):
	return
	self._serve_loop(sock, srv)
	except (ConnectionError, OSError):
	# Peer closed mid-read, or socket dead
	pass
	def handle(self):
	srv = self.server
	sock = self.request
	_tune_shuffle_socket(sock)
	sock.settimeout(300.0) # Prevent leaked/idle connections from hanging indefinitely
	try:
	if not self._handshake(sock, srv):
	return
	self._serve_loop(sock, srv)
	except (ConnectionError, OSError):
	# Peer closed mid-read, or socket dead
	pass
	except Exception:
	logger.exception(f"Unexpected error handling connection from {self.client_address}")

[ShockYoungCHN]

good bot

[gemini-code-assist]

Validate num_ranges to prevent extremely large values from causing CPU/memory exhaustion via the loop.

Suggested change

	num_ranges = _recv_u32(sock)
	ranges = []
	for _ in range(num_ranges):
	offset = _recv_u64(sock)
	length = _recv_u64(sock)
	ranges.append((offset, length))
	num_ranges = _recv_u32(sock)
	if num_ranges > 1000000:
	raise ValueError(f"Too many ranges requested: {num_ranges}")
	ranges = []
	for _ in range(num_ranges):
	offset = _recv_u64(sock)
	length = _recv_u64(sock)
	ranges.append((offset, length))

[ShockYoungCHN]

redundant design, because the num_ranges won't be so large that it causes a problem

[gemini-code-assist]

Allow a limited number of retries (e.g., 3) with exponential backoff for the same endpoint before raising a terminal ShuffleManagerAnomalyError. This prevents transient network drops or temporary server overloads from immediately failing the entire shuffle.

    retries = 0
    max_retries = 3
    while True:
        try:
            endpoint = _resolve()
            with open_shuffle_connection(endpoint, token) as conn:
                for batch in _chunk_members_by_bytes(members, max_bytes_per_fetch):
                    sources = [(m.path, m.ranges) for m in batch]
                    conn.fetch_into(sources, out_file_obj)
            return
        except PermissionError:
            raise
        except (ConnectionError, TimeoutError) as e:
            out_file_obj.reset()
            with _ENDPOINT_CACHE_LOCK:
                _ENDPOINT_CACHE.pop(key, None)
            fresh = _resolve()
            if fresh == endpoint:
                retries += 1
                if retries >= max_retries:
                    raise ShuffleManagerAnomalyError(
                        f"TCP fetch from node {node_id} failed {max_retries} times ({e}) but "
                        f"ShuffleManager at {fresh} is still reachable via Ray. "
                        f"Likely a network configuration issue (NetworkPolicy, "
                        f"firewall, routing) between reducer and manager."
                    ) from e
                logger.warning(
                    f"TCP fetch from node {node_id} failed ({e}); endpoint unchanged. "
                    f"Retrying ({retries}/{max_retries}) after a brief delay..."
                )
                time.sleep(2.0 * retries)
                continue
            else:
                retries = 0  # Reset retry count if endpoint changed (manager restarted)
                logger.warning(
                    f"TCP fetch from node {node_id} failed ({e}); ShuffleManager "
                    f"restarted (endpoint {endpoint} → {fresh}). Retrying in place."
                )
                continue

[cursor]

Infinite loop if os.pwrite returns zero

Medium Severity

_PwriteSink.write unconditionally accumulates the return value of os.pwrite into total. If os.pwrite ever returns 0 (which can happen on certain filesystems or under unusual kernel conditions), the loop spins forever because total never advances past n_total. This is a known real-world pitfall — the coreweave/tensorizer project encountered the same pattern and shipped a dedicated fix. A guard checking for a non-positive return and either raising or breaking would prevent the hang.

 

^{Reviewed by Cursor Bugbot for commit affc97d. Configure here.}

[cursor]

Partial prefetch on fetch failure

Medium Severity

_prefetch_node_into calls _PwriteSink.reset() only on ConnectionError or TimeoutError. When max_bytes_per_fetch splits work into multiple FETCH batches, a later batch failing with FileNotFoundError, PermissionError, or other errors leaves earlier batches written in the node’s prefetch region without reset.

 

^{Reviewed by Cursor Bugbot for commit 637fce1. Configure here.}

[cursor]

Cursor Bugbot has reviewed your changes using default effort and found 1 potential issue.

There are 3 total unresolved issues (including 2 from previous reviews).

^{Reviewed by Cursor Bugbot for commit 13b4661. Configure here.}