Skip to content
This repository was archived by the owner on Jan 12, 2026. It is now read-only.

Commit f1b8fde

Browse files
pabloempabloem
authored
Fixes #14: Refactoring ray portable runner (#18)
* Refactoring ray portable runner * Supporting SDF-initiated checkpoint * Supporting SDF with SDF-initiated splitting. * Fix formatting * fixing smaller issues * Adding portability tests to CI
1 parent 95f7cc0 commit f1b8fde

9 files changed

Lines changed: 2973 additions & 2838 deletions

File tree

‎.flake8‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -33,4 +33,5 @@ ignore =
3333
I
3434
N
3535
avoid-escape = no
36-
36+
per-file-ignores =
37+
*ray_runner_test.py: B008

‎.github/workflows/ci.yml‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -40,6 +40,12 @@ jobs:
4040
- name: Format
4141
run: |
4242
bash scripts/format.sh
43+
- name: Install Ray Beam Runner
44+
run: |
45+
pip install -e .[test]
46+
- name: Run Portability tests
47+
run: |
48+
pytest -r A ray_beam_runner/portability/ray_runner_test.py ray_beam_runner/portability/execution_test.py
4349
4450
LicenseCheck:
4551
name: License Check

‎ray_beam_runner/portability/context_management.py‎

Lines changed: 154 additions & 118 deletions
Original file line numberDiff line numberDiff line change
@@ -15,137 +15,173 @@
1515
# limitations under the License.
1616
#
1717
import typing
18+
from typing import Dict
1819
from typing import List
1920
from typing import Optional
21+
from typing import Tuple
2022

2123
from apache_beam.portability.api import beam_fn_api_pb2
2224
from apache_beam.portability.api import beam_runner_api_pb2
2325
from apache_beam.portability.api import endpoints_pb2
2426
from apache_beam.runners.portability.fn_api_runner import execution as fn_execution
2527
from apache_beam.runners.portability.fn_api_runner import translations
2628
from apache_beam.runners.portability.fn_api_runner import worker_handlers
29+
from apache_beam.runners.portability.fn_api_runner.execution import PartitionableBuffer
30+
from apache_beam.runners.portability.fn_api_runner.fn_runner import OutputTimers
31+
from apache_beam.runners.portability.fn_api_runner.translations import DataOutput
32+
from apache_beam.runners.portability.fn_api_runner.translations import TimerFamilyId
2733
from apache_beam.runners.worker import bundle_processor
2834
from apache_beam.utils import proto_utils
2935

3036
import ray
3137
from ray_beam_runner.portability.execution import RayRunnerExecutionContext
3238

33-
class RayBundleContextManager:
39+
ENCODED_IMPULSE_REFERENCE = ray.put([fn_execution.ENCODED_IMPULSE_VALUE])
40+
3441

35-
def __init__(self,
36-
execution_context: RayRunnerExecutionContext,
37-
stage: translations.Stage,
38-
) -> None:
39-
self.execution_context = execution_context
40-
self.stage = stage
41-
# self.extract_bundle_inputs_and_outputs()
42-
self.bundle_uid = self.execution_context.next_uid()
43-
44-
# Properties that are lazily initialized
45-
self._process_bundle_descriptor = None # type: Optional[beam_fn_api_pb2.ProcessBundleDescriptor]
46-
self._worker_handlers = None # type: Optional[List[worker_handlers.WorkerHandler]]
47-
# a mapping of {(transform_id, timer_family_id): timer_coder_id}. The map
48-
# is built after self._process_bundle_descriptor is initialized.
49-
# This field can be used to tell whether current bundle has timers.
50-
self._timer_coder_ids = None # type: Optional[Dict[Tuple[str, str], str]]
51-
52-
def __reduce__(self):
53-
data = (self.execution_context,
54-
self.stage)
55-
deserializer = lambda args: RayBundleContextManager(args[0], args[1])
56-
return (deserializer, data)
57-
58-
@property
59-
def worker_handlers(self) -> List[worker_handlers.WorkerHandler]:
60-
return []
61-
62-
def data_api_service_descriptor(self) -> Optional[endpoints_pb2.ApiServiceDescriptor]:
63-
return endpoints_pb2.ApiServiceDescriptor(url='fake')
64-
65-
def state_api_service_descriptor(self) -> Optional[endpoints_pb2.ApiServiceDescriptor]:
66-
return None
67-
68-
@property
69-
def process_bundle_descriptor(self):
70-
# type: () -> beam_fn_api_pb2.ProcessBundleDescriptor
71-
if self._process_bundle_descriptor is None:
72-
self._process_bundle_descriptor = beam_fn_api_pb2.ProcessBundleDescriptor.FromString(
73-
self._build_process_bundle_descriptor())
74-
self._timer_coder_ids = fn_execution.BundleContextManager._build_timer_coders_id_map(self)
75-
return self._process_bundle_descriptor
76-
77-
def _build_process_bundle_descriptor(self):
78-
# Cannot be invoked until *after* _extract_endpoints is called.
79-
# Always populate the timer_api_service_descriptor.
80-
pbd = beam_fn_api_pb2.ProcessBundleDescriptor(
81-
id=self.bundle_uid,
82-
transforms={
83-
transform.unique_name: transform
84-
for transform in self.stage.transforms
85-
},
86-
pcollections=dict(
87-
self.execution_context.pipeline_components.pcollections.items()),
88-
coders=dict(self.execution_context.pipeline_components.coders.items()),
89-
windowing_strategies=dict(
90-
self.execution_context.pipeline_components.windowing_strategies.
91-
items()),
92-
environments=dict(
93-
self.execution_context.pipeline_components.environments.items()),
94-
state_api_service_descriptor=self.state_api_service_descriptor(),
95-
timer_api_service_descriptor=self.data_api_service_descriptor())
96-
97-
return pbd.SerializeToString()
98-
99-
def extract_bundle_inputs_and_outputs(self):
100-
# type: () -> Tuple[Dict[str, PartitionableBuffer], DataOutput, Dict[TimerFamilyId, bytes]]
101-
102-
"""Returns maps of transform names to PCollection identifiers.
103-
104-
Also mutates IO stages to point to the data ApiServiceDescriptor.
105-
106-
Returns:
107-
A tuple of (data_input, data_output, expected_timer_output) dictionaries.
108-
`data_input` is a dictionary mapping (transform_name, output_name) to a
109-
PCollection buffer; `data_output` is a dictionary mapping
110-
(transform_name, output_name) to a PCollection ID.
111-
`expected_timer_output` is a dictionary mapping transform_id and
112-
timer family ID to a buffer id for timers.
113-
"""
114-
transform_to_buffer_coder: typing.Dict[str, typing.Tuple[bytes, str]] = {}
115-
data_output = {} # type: DataOutput
116-
expected_timer_output = {} # type: OutputTimers
117-
for transform in self.stage.transforms:
118-
if transform.spec.urn in (bundle_processor.DATA_INPUT_URN,
119-
bundle_processor.DATA_OUTPUT_URN):
120-
pcoll_id = transform.spec.payload
121-
if transform.spec.urn == bundle_processor.DATA_INPUT_URN:
122-
coder_id = self.execution_context.data_channel_coders[translations.only_element(
123-
transform.outputs.values())]
124-
if pcoll_id == translations.IMPULSE_BUFFER:
125-
buffer_actor = ray.get(self.execution_context.pcollection_buffers.get.remote(
126-
transform.unique_name))
127-
ray.get(buffer_actor.append.remote(fn_execution.ENCODED_IMPULSE_VALUE))
128-
pcoll_id = transform.unique_name.encode('utf8')
129-
else:
130-
pass
131-
transform_to_buffer_coder[transform.unique_name] = (
132-
pcoll_id,
133-
self.execution_context.safe_coders.get(coder_id, coder_id)
134-
)
135-
elif transform.spec.urn == bundle_processor.DATA_OUTPUT_URN:
136-
data_output[transform.unique_name] = pcoll_id
137-
coder_id = self.execution_context.data_channel_coders[translations.only_element(
138-
transform.inputs.values())]
139-
else:
140-
raise NotImplementedError
141-
# TODO(pabloem): Figure out when we DO and we DONT need this particular rewrite of coders.
142-
data_spec = beam_fn_api_pb2.RemoteGrpcPort(coder_id=coder_id)
143-
# data_spec.api_service_descriptor.url = 'fake'
144-
transform.spec.payload = data_spec.SerializeToString()
145-
elif transform.spec.urn in translations.PAR_DO_URNS:
146-
payload = proto_utils.parse_Bytes(
147-
transform.spec.payload, beam_runner_api_pb2.ParDoPayload)
148-
for timer_family_id in payload.timer_family_specs.keys():
149-
expected_timer_output[(transform.unique_name, timer_family_id)] = (
150-
translations.create_buffer_id(timer_family_id, 'timers'))
151-
return transform_to_buffer_coder, data_output, expected_timer_output
42+
class RayBundleContextManager:
43+
def __init__(
44+
self,
45+
execution_context: RayRunnerExecutionContext,
46+
stage: translations.Stage,
47+
) -> None:
48+
self.execution_context = execution_context
49+
self.stage = stage
50+
# self.extract_bundle_inputs_and_outputs()
51+
self.bundle_uid = self.execution_context.next_uid()
52+
53+
# Properties that are lazily initialized
54+
self._process_bundle_descriptor = (
55+
None
56+
) # type: Optional[beam_fn_api_pb2.ProcessBundleDescriptor]
57+
self._worker_handlers = (
58+
None
59+
) # type: Optional[List[worker_handlers.WorkerHandler]]
60+
# a mapping of {(transform_id, timer_family_id): timer_coder_id}. The map
61+
# is built after self._process_bundle_descriptor is initialized.
62+
# This field can be used to tell whether current bundle has timers.
63+
self._timer_coder_ids = None # type: Optional[Dict[Tuple[str, str], str]]
64+
65+
def __reduce__(self):
66+
data = (self.execution_context, self.stage)
67+
68+
def deserializer(args):
69+
RayBundleContextManager(args[0], args[1])
70+
71+
return (deserializer, data)
72+
73+
@property
74+
def worker_handlers(self) -> List[worker_handlers.WorkerHandler]:
75+
return []
76+
77+
def data_api_service_descriptor(
78+
self,
79+
) -> Optional[endpoints_pb2.ApiServiceDescriptor]:
80+
return endpoints_pb2.ApiServiceDescriptor(url="fake")
81+
82+
def state_api_service_descriptor(
83+
self,
84+
) -> Optional[endpoints_pb2.ApiServiceDescriptor]:
85+
return None
86+
87+
@property
88+
def process_bundle_descriptor(self) -> beam_fn_api_pb2.ProcessBundleDescriptor:
89+
if self._process_bundle_descriptor is None:
90+
self._process_bundle_descriptor = (
91+
beam_fn_api_pb2.ProcessBundleDescriptor.FromString(
92+
self._build_process_bundle_descriptor()
93+
)
94+
)
95+
self._timer_coder_ids = (
96+
fn_execution.BundleContextManager._build_timer_coders_id_map(self)
97+
)
98+
return self._process_bundle_descriptor
99+
100+
def _build_process_bundle_descriptor(self):
101+
# Cannot be invoked until *after* _extract_endpoints is called.
102+
# Always populate the timer_api_service_descriptor.
103+
pbd = beam_fn_api_pb2.ProcessBundleDescriptor(
104+
id=self.bundle_uid,
105+
transforms={
106+
transform.unique_name: transform for transform in self.stage.transforms
107+
},
108+
pcollections=dict(
109+
self.execution_context.pipeline_components.pcollections.items()
110+
),
111+
coders=dict(self.execution_context.pipeline_components.coders.items()),
112+
windowing_strategies=dict(
113+
self.execution_context.pipeline_components.windowing_strategies.items()
114+
),
115+
environments=dict(
116+
self.execution_context.pipeline_components.environments.items()
117+
),
118+
state_api_service_descriptor=self.state_api_service_descriptor(),
119+
timer_api_service_descriptor=self.data_api_service_descriptor(),
120+
)
121+
122+
return pbd.SerializeToString()
123+
124+
def get_bundle_inputs_and_outputs(
125+
self,
126+
) -> Tuple[Dict[str, PartitionableBuffer], DataOutput, Dict[TimerFamilyId, bytes]]:
127+
"""Returns maps of transform names to PCollection identifiers.
128+
129+
Also mutates IO stages to point to the data ApiServiceDescriptor.
130+
131+
Returns:
132+
A tuple of (data_input, data_output, expected_timer_output) dictionaries.
133+
`data_input` is a dictionary mapping (transform_name, output_name) to a
134+
PCollection buffer; `data_output` is a dictionary mapping
135+
(transform_name, output_name) to a PCollection ID.
136+
`expected_timer_output` is a dictionary mapping transform_id and
137+
timer family ID to a buffer id for timers.
138+
"""
139+
return self.transform_to_buffer_coder, self.data_output, self.stage_timers
140+
141+
def setup(self):
142+
transform_to_buffer_coder: typing.Dict[str, typing.Tuple[bytes, str]] = {}
143+
data_output = {} # type: DataOutput
144+
expected_timer_output = {} # type: OutputTimers
145+
for transform in self.stage.transforms:
146+
if transform.spec.urn in (
147+
bundle_processor.DATA_INPUT_URN,
148+
bundle_processor.DATA_OUTPUT_URN,
149+
):
150+
pcoll_id = transform.spec.payload
151+
if transform.spec.urn == bundle_processor.DATA_INPUT_URN:
152+
coder_id = self.execution_context.data_channel_coders[
153+
translations.only_element(transform.outputs.values())
154+
]
155+
if pcoll_id == translations.IMPULSE_BUFFER:
156+
pcoll_id = transform.unique_name.encode("utf8")
157+
self.execution_context.pcollection_buffers.put.remote(
158+
pcoll_id, [ENCODED_IMPULSE_REFERENCE]
159+
)
160+
else:
161+
pass
162+
transform_to_buffer_coder[transform.unique_name] = (
163+
pcoll_id,
164+
self.execution_context.safe_coders.get(coder_id, coder_id),
165+
)
166+
elif transform.spec.urn == bundle_processor.DATA_OUTPUT_URN:
167+
data_output[transform.unique_name] = pcoll_id
168+
coder_id = self.execution_context.data_channel_coders[
169+
translations.only_element(transform.inputs.values())
170+
]
171+
else:
172+
raise NotImplementedError
173+
data_spec = beam_fn_api_pb2.RemoteGrpcPort(coder_id=coder_id)
174+
transform.spec.payload = data_spec.SerializeToString()
175+
elif transform.spec.urn in translations.PAR_DO_URNS:
176+
payload = proto_utils.parse_Bytes(
177+
transform.spec.payload, beam_runner_api_pb2.ParDoPayload
178+
)
179+
for timer_family_id in payload.timer_family_specs.keys():
180+
expected_timer_output[
181+
(transform.unique_name, timer_family_id)
182+
] = translations.create_buffer_id(timer_family_id, "timers")
183+
self.transform_to_buffer_coder, self.data_output, self.stage_timers = (
184+
transform_to_buffer_coder,
185+
data_output,
186+
expected_timer_output,
187+
)

0 commit comments

Comments
 (0)