Ray integration

Ray is an optional transport. Domain processors and algorithm APIs are identical for inprocess, local processes, and Ray.

Configuration

params_defaults:
  runtime.parallelism.backend: ray
  runtime.parallelism.workers: 4
  runtime.parallelism.ray_address: ray://ray-head:10001

If the address is empty, RAY_ADDRESS is used; otherwise Ray starts locally. integration/ray/connect.py lazily imports Ray and calls ray.init once.

runtime.parallelism.* controls execution workers. compute.parallelism.* controls algorithm-internal resources and is separate.

partition_loop Ray backend

parallel/backends/ray_map.py:

1. connects to Ray;
2. creates a serializable WorkerRunSpec from manifest path, RunContext identity and params;
3. assigns a distinct write shard id per worker;
4. creates one actor per worker;
5. passes WorkerShardAssignment to each actor;
6. each actor constructs its own WorkerSession, storage session and processor context;
7. driver merges worker manifests and metric deltas.

Workers stream their assigned primary shards. The driver does not ship full datasets or live StorageSession objects through Ray.

Metrics channel

When metrics are enabled, workers send neutral metric deltas through a Ray channel. Driver-side telemetry owns Prometheus emission. Worker code does not start scrape servers.

Checkpoint/resume constraints

ParallelPlan.require_inprocess() rejects features that are not safe for multi-worker execution. Checkpoint/resume is currently inprocess-only. Run-global mutable registries cannot be merged in driver finalization; model them as partial rows followed by reduce_merge.

embed_train Ray backend

Embedding training has a distinct transport because it owns run-global mutable W.

Current production driver path:

• one EmbedVocabActor parameter server per material group;
• pass 1 loads/merges token counts and initializes vocabulary on the driver;
• pair batches are submitted asynchronously to that actor with bounded inflight;
• non-blocking sample polling collects train metrics;
• actor-side geometry diagnostics avoid shipping full W when possible;
• driver requests batched snapshots and writes them through StorageSession;
• actor is killed after group finalization.

runtime.parallelism.workers is recorded in result metadata, but the active Ray train driver does not currently spawn N shard train workers. A multi-worker helper API exists for tests/legacy use; do not document or rely on “actor + N workers” as the live path until the driver uses it.

The shared run_train_over_groups controls group order/result shape for inprocess, processes, and Ray backends. Production embed_train currently hard-requires compute.embedding_method: sgns_v1.

Relevant tuning params

Common execution params:

• runtime.parallelism.workers — shard workers for partition_loop; metadata for current Ray train
• runtime.parallelism.ray_address

Embed transport params are explicit compute.*/mode defaults, including pair batch size, max inflight updates, snapshot batch rows, vocabulary threshold, and geometry sampling sizes. Declare them in DRTML; do not hide them in step code.

Serialization rules

Safe worker payloads contain:

• manifest path;
• run/step/artifacts identity;
• merged params;
• shard assignment;
• simple DTOs.

Do not serialize open files, DB clients, StorageSession, metrics clients, or arbitrary domain closures. Reconstruct resources inside the worker.

Adding a Ray capability

• Connection/init → integration/ray.
• Partition dispatch → processing/partition_loop/parallel/backends.
• Mode-specific transport → that mode’s backends/ package.
• Pure computation → algorithms/.

Implement the same transport-neutral callable contract as existing backends, reuse shared outer loops, add bounded backpressure, explicit cleanup, failure propagation, and tests with Ray optional/skipped when unavailable.