Skip to content

Commit 1ca6249

Browse files
kohlisidkohlisid
committed
add TCP conns
Signed-off-by: Sidhant Kohli <sidhant.kohli@gmail.com>
1 parent b596cc4 commit 1ca6249

10 files changed

Lines changed: 241 additions & 25 deletions

File tree

‎examples/map/async_multiproc_map/Dockerfile‎

Lines changed: 55 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,55 @@
1+
####################################################################################################
2+
# builder: install needed dependencies
3+
####################################################################################################
4+
5+
FROM python:3.10-slim-bullseye AS builder
6+
7+
ENV PYTHONFAULTHANDLER=1 \
8+
PYTHONUNBUFFERED=1 \
9+
PYTHONHASHSEED=random \
10+
PIP_NO_CACHE_DIR=on \
11+
PIP_DISABLE_PIP_VERSION_CHECK=on \
12+
PIP_DEFAULT_TIMEOUT=100 \
13+
POETRY_VERSION=1.2.2 \
14+
POETRY_HOME="/opt/poetry" \
15+
POETRY_VIRTUALENVS_IN_PROJECT=true \
16+
POETRY_NO_INTERACTION=1 \
17+
PYSETUP_PATH="/opt/pysetup"
18+
19+
ENV EXAMPLE_PATH="$PYSETUP_PATH/examples/map/async_multiproc_map"
20+
ENV VENV_PATH="$EXAMPLE_PATH/.venv"
21+
ENV PATH="$POETRY_HOME/bin:$VENV_PATH/bin:$PATH"
22+
23+
RUN apt-get update \
24+
&& apt-get install --no-install-recommends -y \
25+
curl \
26+
wget \
27+
# deps for building python deps
28+
build-essential \
29+
&& apt-get install -y git \
30+
&& apt-get clean && rm -rf /var/lib/apt/lists/* \
31+
\
32+
# install dumb-init
33+
&& wget -O /dumb-init https://github.com/Yelp/dumb-init/releases/download/v1.2.5/dumb-init_1.2.5_x86_64 \
34+
&& chmod +x /dumb-init \
35+
&& curl -sSL https://install.python-poetry.org | python3 -
36+
37+
####################################################################################################
38+
# udf: used for running the udf vertices
39+
####################################################################################################
40+
FROM builder AS udf
41+
42+
WORKDIR $PYSETUP_PATH
43+
COPY ./ ./
44+
45+
WORKDIR $EXAMPLE_PATH
46+
RUN poetry lock
47+
RUN poetry install --no-cache --no-root && \
48+
rm -rf ~/.cache/pypoetry/
49+
50+
RUN chmod +x entry.sh
51+
52+
ENTRYPOINT ["/dumb-init", "--"]
53+
CMD ["sh", "-c", "$EXAMPLE_PATH/entry.sh"]
54+
55+
EXPOSE 5000

‎examples/map/async_multiproc_map/Makefile‎

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,22 @@
1+
TAG ?= v4
2+
PUSH ?= false
3+
IMAGE_REGISTRY = quay.io/skohli/numaflow-python/async-multiproc:${TAG}
4+
DOCKER_FILE_PATH = examples/map/async_multiproc_map/Dockerfile
5+
6+
.PHONY: update
7+
update:
8+
poetry update -vv
9+
10+
.PHONY: image-push
11+
image-push: update
12+
cd ../../../ && docker buildx build \
13+
-f ${DOCKER_FILE_PATH} \
14+
-t ${IMAGE_REGISTRY} \
15+
--platform linux/amd64,linux/arm64 . --push
16+
17+
.PHONY: image
18+
image: update
19+
cd ../../../ && docker build \
20+
-f ${DOCKER_FILE_PATH} \
21+
-t ${IMAGE_REGISTRY} .
22+
@if [ "$(PUSH)" = "true" ]; then docker push ${IMAGE_REGISTRY}; fi

‎examples/map/async_multiproc_map/README.md‎

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,20 @@
1+
# Multiprocessing Map
2+
3+
`pynumaflow` supports only asyncio based Reduce UDFs because we found that procedural Python is not able to handle
4+
any substantial traffic.
5+
6+
This features enables the `pynumaflow` developer to utilise multiprocessing capabilities while
7+
writing UDFs using the map function. These are particularly useful for CPU intensive operations,
8+
as it allows for better resource utilisation.
9+
10+
In this mode we would spawn N number (N = Cpu count) of grpc servers in different processes, where each of them are
11+
listening on multiple TCP sockets.
12+
13+
To enable multiprocessing mode start the multiproc server in the UDF using the following command,
14+
providing the optional argument `server_count` to specify the number of
15+
servers to be forked (defaults to `os.cpu_count` if not provided):
16+
```python
17+
if __name__ == "__main__":
18+
grpc_server = MapMultiProcServer(handler, server_count = 3)
19+
grpc_server.start()
20+
```

‎examples/map/async_multiproc_map/entry.sh‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,4 @@
1+
#!/bin/sh
2+
set -eux
3+
4+
python example.py

‎examples/map/async_multiproc_map/example.py‎

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,34 @@
1+
import os
2+
3+
from pynumaflow.mapper import Messages, Message, Datum, Mapper, AsyncMapMultiprocServer
4+
from pynumaflow._constants import _LOGGER
5+
6+
7+
class FlatMap(Mapper):
8+
"""
9+
This class needs to be of type Mapper class to be used
10+
as a handler for the MapServer class.
11+
Example of a mapper that calculates if a number is prime.
12+
"""
13+
14+
async def handler(self, keys: list[str], datum: Datum) -> Messages:
15+
val = datum.value
16+
_ = datum.event_time
17+
_ = datum.watermark
18+
messages = Messages()
19+
messages.append(Message(val, keys=keys))
20+
_LOGGER.info(f"MY PID {os.getpid()}")
21+
return messages
22+
23+
24+
if __name__ == "__main__":
25+
"""
26+
Example of starting a multiprocessing map vertex.
27+
"""
28+
# To set the env server_count value set the env variable
29+
# NUM_CPU_MULTIPROC="N"
30+
server_count = int(os.getenv("NUM_CPU_MULTIPROC", "2"))
31+
_class = FlatMap()
32+
# Server count is the number of server processes to start
33+
grpc_server = AsyncMapMultiprocServer(_class, server_count=server_count, use_tcp=True)
34+
grpc_server.start()

‎examples/map/async_multiproc_map/pipeline.yaml‎

Lines changed: 40 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,40 @@
1+
apiVersion: numaflow.numaproj.io/v1alpha1
2+
kind: Pipeline
3+
metadata:
4+
name: simple-pipeline
5+
spec:
6+
limits:
7+
readBatchSize: 10
8+
vertices:
9+
- name: in
10+
source:
11+
# A self data generating source
12+
generator:
13+
rpu: 200
14+
duration: 1s
15+
- name: mult
16+
udf:
17+
container:
18+
image: quay.io/skohli/numaflow-python/async-multiproc:v4
19+
# imagePullPolicy: Always
20+
env:
21+
- name: PYTHONDEBUG
22+
value: "true"
23+
- name: NUM_CPU_MULTIPROC
24+
value: "3" # DO NOT forget the double quotes!!!
25+
containerTemplate:
26+
env:
27+
- name: NUMAFLOW_RUNTIME
28+
value: "rust"
29+
- name: NUMAFLOW_DEBUG
30+
value: "true" # DO NOT forget the double quotes!!!
31+
32+
- name: out
33+
sink:
34+
# A simple log printing sink
35+
log: {}
36+
edges:
37+
- from: in
38+
to: mult
39+
- from: mult
40+
to: out

‎examples/map/async_multiproc_map/pyproject.toml‎

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,15 @@
1+
[tool.poetry]
2+
name = "async-multiproc-forward-message"
3+
version = "0.2.4"
4+
description = ""
5+
authors = ["Numaflow developers"]
6+
7+
[tool.poetry.dependencies]
8+
python = ">=3.10,<3.13"
9+
pynumaflow = { path = "../../../"}
10+
11+
[tool.poetry.dev-dependencies]
12+
13+
[build-system]
14+
requires = ["poetry-core>=1.0.0"]
15+
build-backend = "poetry.core.masonry.api"

‎pynumaflow/mapper/_servicer/_async_servicer.py‎

Lines changed: 44 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
import asyncio
2+
import contextlib
23
from collections.abc import AsyncIterable
34

45
from google.protobuf import empty_pb2 as _empty_pb2
@@ -35,6 +36,7 @@ async def MapFn(
3536
"""
3637
# proto repeated field(keys) is of type google._upb._message.RepeatedScalarContainer
3738
# we need to explicitly convert it to list
39+
producer = None
3840
try:
3941
# The first message to be received should be a valid handshake
4042
req = await request_iterator.__anext__()
@@ -62,37 +64,57 @@ async def MapFn(
6264
yield msg
6365
# wait for the producer task to complete
6466
await producer
67+
except GeneratorExit:
68+
_LOGGER.info("Client disconnected, generator closed.")
69+
raise
6570
except BaseException as e:
6671
_LOGGER.critical("UDFError, re-raising the error", exc_info=True)
6772
await handle_async_error(context, e, ERR_UDF_EXCEPTION_STRING, self.multiproc)
6873
return
74+
finally:
75+
if producer and not producer.done():
76+
producer.cancel()
77+
with contextlib.suppress(asyncio.CancelledError):
78+
await producer
6979

70-
async def _process_inputs(
71-
self,
72-
request_iterator: AsyncIterable[map_pb2.MapRequest],
73-
result_queue: NonBlockingIterator,
74-
):
75-
"""
76-
Utility function for processing incoming MapRequests
77-
"""
80+
async def _process_inputs(self, request_iterator, result_queue):
7881
try:
79-
# for each incoming request, create a background task to execute the
80-
# UDF code
8182
async for req in request_iterator:
82-
msg_task = asyncio.create_task(self._invoke_map(req, result_queue))
83-
# save a reference to a set to store active tasks
84-
self.background_tasks.add(msg_task)
85-
msg_task.add_done_callback(self.background_tasks.discard)
86-
87-
# wait for all tasks to complete
88-
for task in self.background_tasks:
89-
await task
90-
91-
# send an EOF to result queue to indicate that all tasks have completed
92-
await result_queue.put(STREAM_EOF)
83+
task = asyncio.create_task(self._invoke_map(req, result_queue))
84+
self.background_tasks.add(task)
85+
task.add_done_callback(self.background_tasks.discard)
9386

87+
await asyncio.gather(*self.background_tasks)
9488
except BaseException:
95-
_LOGGER.critical("MapFn Error, re-raising the error", exc_info=True)
89+
_LOGGER.critical("MapFn Error in _process_inputs", exc_info=True)
90+
finally:
91+
await result_queue.put(STREAM_EOF)
92+
# async def _process_inputs(
93+
# self,
94+
# request_iterator: AsyncIterable[map_pb2.MapRequest],
95+
# result_queue: NonBlockingIterator,
96+
# ):
97+
# """
98+
# Utility function for processing incoming MapRequests
99+
# """
100+
# try:
101+
# # for each incoming request, create a background task to execute the
102+
# # UDF code
103+
# async for req in request_iterator:
104+
# msg_task = asyncio.create_task(self._invoke_map(req, result_queue))
105+
# # save a reference to a set to store active tasks
106+
# self.background_tasks.add(msg_task)
107+
# msg_task.add_done_callback(self.background_tasks.discard)
108+
#
109+
# # wait for all tasks to complete
110+
# for task in self.background_tasks:
111+
# await task
112+
#
113+
# # send an EOF to result queue to indicate that all tasks have completed
114+
# await result_queue.put(STREAM_EOF)
115+
#
116+
# except BaseException:
117+
# _LOGGER.critical("MapFn Error, re-raising the error", exc_info=True)
96118

97119
async def _invoke_map(self, req: map_pb2.MapRequest, result_queue: NonBlockingIterator):
98120
"""

‎pynumaflow/mapper/async_multiproc_server.py‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -113,6 +113,8 @@ async def run_server():
113113
server_info = ServerInfo.get_default_server_info()
114114
server_info.minimum_numaflow_version = MINIMUM_NUMAFLOW_VERSION[ContainerType.Mapper]
115115
server_info.metadata = get_metadata_env(envs=METADATA_ENVS)
116+
if self.use_tcp:
117+
server_info.protocol = Protocol.TCP
116118
# Add the MULTIPROC metadata using the number of servers to use
117119
server_info.metadata[MULTIPROC_KEY] = str(self._process_count)
118120
# Add the MAP_MODE metadata to the server info for the correct map mode
@@ -123,7 +125,7 @@ async def run_server():
123125
sock_path=bind_address,
124126
max_threads=self.max_threads,
125127
cleanup_coroutines=list(),
126-
server_info_file=self.server_info_file,
128+
server_info_file=None,
127129
server_info=server_info,
128130
)
129131

‎pynumaflow/shared/server.py‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -180,7 +180,7 @@ async def start_async_server(
180180
sock_path: str,
181181
max_threads: int,
182182
cleanup_coroutines: list,
183-
server_info_file: str,
183+
server_info_file: Optional[str] = None,
184184
server_info: Optional[ServerInfo] = None,
185185
):
186186
"""
@@ -194,7 +194,9 @@ async def start_async_server(
194194
# Create the server info file if not provided
195195
server_info = ServerInfo.get_default_server_info()
196196
# Add the server information to the server info file
197-
info_server_write(server_info=server_info, info_file=server_info_file)
197+
198+
if server_info_file:
199+
info_server_write(server_info=server_info, info_file=server_info_file)
198200

199201
# Log the server start
200202
_LOGGER.info(

0 commit comments

Comments
 (0)