Proxy class that executes Ray Data operations remotely on cluster workers.
| 8 | |
| 9 | |
| 10 | class RemoteDatasetProxy: |
| 11 | """Proxy class that executes Ray Data operations remotely on cluster workers.""" |
| 12 | |
| 13 | def __init__(self, dataset_ref: Any): |
| 14 | """Initialize with a reference to the remote dataset.""" |
| 15 | self._dataset_ref = dataset_ref |
| 16 | |
| 17 | def map_batches( |
| 18 | self, |
| 19 | func, |
| 20 | num_gpus: float = 0, |
| 21 | worker_task_options: Optional[Dict[str, Any]] = None, |
| 22 | **kwargs, |
| 23 | ) -> "RemoteDatasetProxy": |
| 24 | """Execute map_batches remotely on cluster workers. |
| 25 | |
| 26 | Resource options are applied at two levels: |
| 27 | |
| 28 | 1. **Orchestration task** (@ray.remote wrapper) — receives only the |
| 29 | non-compute keys from worker_task_options (runtime_env, max_retries, |
| 30 | scheduling_strategy, memory, …). Compute-scheduling keys |
| 31 | (num_gpus, num_cpus, accelerator_type, resources) are intentionally |
| 32 | excluded: the orchestration task only calls dataset.map_batches() |
| 33 | and holds no GPU/CPU work itself. Including num_gpus here would |
| 34 | waste a GPU slot for the entire operation duration and, on |
| 35 | GPU-constrained clusters, could cause deadlock where the orchestrator |
| 36 | holds a GPU while the data workers queue waiting for the same slots. |
| 37 | |
| 38 | 2. **Data workers** (inside dataset.map_batches) — receives the full |
| 39 | compute-scheduling subset (num_gpus, num_cpus, accelerator_type, |
| 40 | resources). Ray Data propagates these to the actual processing tasks. |
| 41 | |
| 42 | Args: |
| 43 | func: Batch transformation function. |
| 44 | num_gpus: Shorthand GPU count (merged into worker_task_options, |
| 45 | takes precedence). Kept first-class because it also drives |
| 46 | gpu_batch_format selection in the compute engine. |
| 47 | worker_task_options: Arbitrary Ray .options() kwargs (num_cpus, |
| 48 | memory, accelerator_type, resources, runtime_env, |
| 49 | max_retries, …). |
| 50 | **kwargs: Additional map_batches kwargs (batch_format, concurrency). |
| 51 | """ |
| 52 | # Merge num_gpus into worker_task_options; dedicated field takes precedence |
| 53 | opts: Dict[str, Any] = dict(worker_task_options or {}) |
| 54 | if num_gpus: |
| 55 | opts["num_gpus"] = num_gpus |
| 56 | |
| 57 | # Keys accepted by Ray Data's map_batches for per-worker scheduling |
| 58 | _MAP_BATCHES_RESOURCE_KEYS = { |
| 59 | "num_gpus", |
| 60 | "num_cpus", |
| 61 | "accelerator_type", |
| 62 | "resources", |
| 63 | } |
| 64 | |
| 65 | # Data workers get the compute-scheduling subset only |
| 66 | map_resource_kwargs = { |
| 67 | k: v for k, v in opts.items() if k in _MAP_BATCHES_RESOURCE_KEYS |
no outgoing calls