Coverage for arrakis/publish.py: 88.5%

1# Copyright (c) 2022, California Institute of Technology and contributors

2#

3# You should have received a copy of the licensing terms for this

4# software included in the file "LICENSE" located in the top-level

5# directory of this package. If you did not, you can view a copy at

6# https://git.ligo.org/ngdd/arrakis-python/-/raw/main/LICENSE

7

8"""Publisher API."""

9

10 from __future__ import annotations

11

12 import contextlib

13 import logging

14 from typing import TYPE_CHECKING, Literal

15

16 import pyarrow

17 from pyarrow.flight import connect

18

19 from . import constants

20 from .client import Client

21 from .flight import (

22 MultiEndpointStream,

23 RequestType,

24 RequestValidator,

25 create_descriptor,

26 parse_url,

27)

28

29 try:

30 from confluent_kafka import Producer

31 except ImportError:

32 HAS_KAFKA = False

33 else:

34 HAS_KAFKA = True

35

36 if TYPE_CHECKING:

37 from collections.abc import Iterable

38 from datetime import timedelta

39

40 from .block import SeriesBlock

41 from .channel import Channel

42

43

44 logger = logging.getLogger("arrakis")

45

46

47 def serialize_batch(batch: pyarrow.RecordBatch):

48 """Serialize a record batch to bytes.

49

50 Parameters

51 ----------

52 batch : pyarrow.RecordBatch

53 The batch to serialize.

54

55 Returns

56 -------

57 bytes

58 The serialized buffer.

59

60 """

61 sink = pyarrow.BufferOutputStream()

62 with pyarrow.ipc.new_stream(sink, batch.schema) as writer:

63 writer.write_batch(batch)

64 return sink.getvalue()

65

66

67 def channel_to_dtype_name(channel: Channel) -> str:

68 """Given a channel, return the data type's name."""

69 assert channel.data_type is not None

70 if isinstance(channel.data_type, str):

71 return channel.data_type

72 return channel.data_type.name

73

74

75 class Publisher:

76 """Publisher to publish timeseries to Arrakis service.

77

78 Parameters

79 ----------

80 id : str

81 Publisher ID string.

82 url : str

83 Initial Flight URL to connect to.

84

85 """

86

87 def __init__(self, publisher_id: str, url: str | None = None):

88 if not HAS_KAFKA:

89 msg = (

90 "Publishing requires confluent-kafka to be installed."

91 "This is provided by the 'publish' extra or it can be "

92 "installed manually through pip or conda."

93 )

94 raise ImportError(msg)

95

96 self.publisher_id = publisher_id

97 self.initial_url = parse_url(url)

98

99 self.channels: dict[str, Channel] = {}

100

101 self._producer: Producer

102 self._partitions: dict[str, str]

103 self._registered = False

104 self._validator = RequestValidator()

105

106 def register(self):

107 assert not self._registered, "has already registered"

108

109 self.channels = {

110 channel.name: channel

111 for channel in Client(self.initial_url).find(publisher=self.publisher_id)

112 }

113 if not self.channels:

114 msg = f"unknown publisher ID '{self.publisher_id}'."

115 raise ValueError(msg)

116

117 # extract the channel partition map

118 self._partitions = {}

119 for channel in self.channels.values():

120 if not channel.partition_id:

121 msg = f"could not determine partition_id for channel {channel}."

122 raise ValueError(msg)

123 self._partitions[channel.name] = channel.partition_id

124

125 self._registered = True

126

127 return self

128

129 def enter(self):

130 if not self._registered:

131 msg = "must register publisher interface before publishing."

132 raise RuntimeError(msg)

133

134 # get connection properties

135 descriptor = create_descriptor(

136 RequestType.Publish,

137 publisher_id=self.publisher_id,

138 validator=self._validator,

139 )

140 properties: dict[str, str] = {}

141 with connect(self.initial_url) as client:

142 flight_info = client.get_flight_info(descriptor)

143 with MultiEndpointStream(flight_info.endpoints, client) as stream:

144 for data in stream.unpack():

145 kv_pairs = data["properties"]

146 properties.update(dict(kv_pairs))

147

148 # set up producer

149 self._producer = Producer(

150 {

151 "message.max.bytes": 10_000_000, # 10 MB

152 "enable.idempotence": True,

153 **properties,

154 }

155 )

156

157 def __enter__(self) -> Publisher:

158 self.enter()

159 return self

160

161 def publish(

162 self,

163 block: SeriesBlock,

164 timeout: timedelta = constants.DEFAULT_TIMEOUT,

165 ) -> None:

166 """Publish timeseries data

167

168 Parameters

169 ----------

170 block : SeriesBlock

171 A data block with all channels to publish.

172 timeout : timedelta, optional

173 The maximum time to wait to publish before timing out.

174 Default is 2 seconds.

175

176 """

177 if not hasattr(self, "_producer") or not self._producer:

178 msg = (

179 "publication interface not initialized, "

180 "please use context manager when publishing."

181 )

182 raise RuntimeError(msg)

183

184 for name, channel in block.channels.items():

185 if channel != self.channels[name]:

186 msg = f"invalid channel for this publisher: {channel}"

187 raise ValueError(msg)

188

189 # FIXME: updating partitions should only be allowed for

190 # special blessed publishers, that are currently not using

191 # this interface, so we're disabling this functionality for

192 # the time being, until we have a better way to manage it.

193 #

194 # # check for new metadata changes

195 # changed = set(block.channels.values()) - set(self.channels.values())

196 # # exchange to transfer metadata and get new/updated partition IDs

197 # if changed:

198 # self._update_partitions(changed)

199

200 # publish data for each data type, splitting into

201 # subblocks based on a maximum channel maximum

202 for partition_id, batch in block.to_row_batches(self._partitions):

203 topic = f"arrakis-{partition_id}"

204 logger.debug("publishing to topic %s: %s", topic, batch)

205 self._producer.produce(topic=topic, value=serialize_batch(batch))

206 self._producer.flush()

207

208 def _update_partitions(

209 self, channels: Iterable[Channel]

210 ) -> None: # pragma: no cover

211 # set up flight

212 assert self._registered, "has not registered yet"

213 descriptor = create_descriptor(

214 RequestType.Partition,

215 publisher_id=self.publisher_id,

216 validator=self._validator,

217 )

218 # FIXME: should we not get FlightInfo first?

219 with connect(self.initial_url) as client:

220 writer, reader = client.do_exchange(descriptor)

221

222 # send over list of channels to map new/updated partitions for

223 dtypes = [channel_to_dtype_name(channel) for channel in channels]

224 schema = pyarrow.schema(

225 [

226 pyarrow.field("channel", pyarrow.string(), nullable=False),

227 pyarrow.field("data_type", pyarrow.string(), nullable=False),

228 pyarrow.field("sample_rate", pyarrow.int32(), nullable=False),

229 pyarrow.field("partition_id", pyarrow.string()),

230 ]

231 )

232 batch = pyarrow.RecordBatch.from_arrays(

233 [

234 pyarrow.array(

235 [str(channel) for channel in channels],

236 type=schema.field("channel").type,

237 ),

238 pyarrow.array(dtypes, type=schema.field("data_type").type),

239 pyarrow.array(

240 [channel.sample_rate for channel in channels],

241 type=schema.field("sample_rate").type,

242 ),

243 pyarrow.array(

244 [None for _ in channels],

245 type=schema.field("partition_id").type,

246 ),

247 ],

248 schema=schema,

249 )

250

251 # send over the partitions

252 writer.begin(schema)

253 writer.write_batch(batch)

254 writer.done_writing()

255

256 # get back the partition IDs and update

257 partitions = reader.read_all().to_pydict()

258 for channel, id_ in zip(partitions["channel"], partitions["partition_id"]):

259 self._partitions[channel] = id_

260

261 def close(self) -> None:

262 logger.info("closing kafka producer...")

263 with contextlib.suppress(Exception):

264 self._producer.flush()

265

266 def __exit__(self, *exc) -> Literal[False]:

267 self.close()

268 return False