Coverage for adrs/execution/executor.py: 36%
277 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-08-29 05:07 +0800
1import time
2import json
3import math
4import signal
5import asyncio
6import logging
7import traceback
8import numpy as np
9import polars as pl
11from nats_client import Msg
12from functools import reduce
13from typing import NamedTuple, TypedDict, cast
14from datetime import datetime, timezone, timedelta
15from decimal import Decimal, ROUND_HALF_UP
17from aion import Scheduler, Trigger
19from adrs.alpha import Alpha
20from adrs.portfolio import Portfolio
21from adrs.logging import setup_logger
22from adrs.io.event import Event, EventType
23from adrs.types import Topic, CollectedData, Data
24from adrs.data import (
25 MetricStream,
26 Datamap,
27 DataInfo,
28 DataLoader,
29 DatasourceStream,
30 MetricBuilder,
31)
32from adrs.data.connector import DEFAULT_METRIC_NAMESPACE
33from adrs.subjects import (
34 alpha_signal_subject,
35 alpha_signal_subscription,
36 parse_alpha_id,
37 portfolio_signal_subject,
38)
41logger = logging.getLogger(__name__)
43_EPOCH = datetime(1970, 1, 1, tzinfo=timezone.utc)
46class LatestSignal(NamedTuple):
47 """The signal an alpha just produced, paired with the bar it belongs to."""
49 signal: str
50 bar_time_ns: int
53def _latest_signal(signal_df: pl.DataFrame) -> LatestSignal:
54 """Read the last row of `signal_df` — value and bar time together, so the
55 two can never end up coming from different bars.
57 Raises ValueError unless that row carries both a signal and the bar it
58 belongs to. `Alpha.next` is user code, and these are the same invariants
59 `Alpha.backtest` already enforces, so a live emit should fail the same way
60 instead of publishing a `"nan"` signal or one dated by the wall clock.
61 """
62 if missing := [c for c in ("start_time", "signal") if c not in signal_df.columns]:
63 raise ValueError(
64 f"DataFrame returned from `next()` must have the column(s): {missing}"
65 )
66 if not signal_df.schema["signal"].is_numeric():
67 raise ValueError(
68 "DataFrame returned from `next()` must have a 'signal' column that is numeric"
69 )
71 row = signal_df.tail(1)
72 if (value := row["signal"][0]) is None:
73 raise ValueError("`next()` returned a null 'signal' for the latest bar")
74 if (ts := row["start_time"][0]) is None:
75 raise ValueError("`next()` returned a null 'start_time' for the latest signal")
77 if ts.tzinfo is None:
78 # storage convention is UTC (see adrs.types._normalise); an unlabelled
79 # datetime is UTC, not local.
80 ts = ts.replace(tzinfo=timezone.utc)
81 return LatestSignal(
82 np.format_float_positional(value, precision=2, unique=False),
83 # integer math, so nothing is lost rounding through a float timestamp
84 (ts - _EPOCH) // timedelta(microseconds=1) * 1_000,
85 )
88def flat_map(f, xs):
89 ys = []
90 for x in xs:
91 ys.extend(f(x))
92 return ys
95def acq_info(topic: str, infos: list[DataInfo]) -> list[DataInfo]:
96 return cast(
97 list[DataInfo],
98 list(
99 filter(
100 lambda di: Topic.from_str(di.topic) == Topic.from_str(topic),
101 infos,
102 )
103 ),
104 )
107def flatten_json(value, _prefix: str = "") -> dict:
108 flattened_map = {}
110 if not isinstance(value, dict):
111 return flattened_map
113 for key, val in value.items():
114 if isinstance(val, dict):
115 nested_map = flatten_json(val)
116 for nested_key, nested_val in nested_map.items():
117 if key == "o":
118 flattened_map[nested_key] = nested_val
119 else:
120 flattened_map[f"{key}_{nested_key}"] = nested_val
121 else:
122 flattened_map[key] = val
124 return flattened_map
127class Signal(TypedDict):
128 signal: Decimal
129 weight: Decimal
132class PortfolioExecutor:
133 def __init__(
134 self,
135 portfolio: Portfolio,
136 metric_stream: MetricStream,
137 aggregate_window: Trigger = Trigger.Cron("*/5 * * * *"),
138 max_signal_age: timedelta = timedelta(hours=2),
139 signal_namespace: str | None = None,
140 insert_prefix: str = DEFAULT_METRIC_NAMESPACE,
141 ):
142 self.portfolio = portfolio
143 self.metric_builder = MetricBuilder(metric_stream, insert_prefix)
144 self.aggregate_window = aggregate_window
145 self.scheduler = Scheduler()
146 self.max_signal_age: timedelta = max_signal_age
147 self.signal_namespace = signal_namespace
148 # Alpha ids this portfolio actually holds (post any alpha_id_map rename,
149 # since run_portfolio remaps the frames too) — used to drop stray signals
150 # in on_signal instead of polluting lastest_signal.
151 self.alpha_ids: set[str] = set(portfolio.metadata_df["custom_id"].to_list())
153 last_signal_time = self.portfolio.signal_df["start_time"][-1]
154 if isinstance(last_signal_time, datetime):
155 last_signal_time = (
156 last_signal_time.replace(tzinfo=timezone.utc)
157 if last_signal_time.tzinfo is None
158 else last_signal_time
159 )
160 age = datetime.now(tz=timezone.utc) - last_signal_time
161 if age > max_signal_age:
162 raise Exception(
163 f"signal_df is stale: last row is {age} old (>{max_signal_age}). "
164 "Ensure signal_df covers current time"
165 )
167 setup_logger()
169 def handle_signal(signum, _):
170 match signum:
171 case signal.SIGINT | signal.SIGTERM:
172 self.on_shutdown()
173 exit(0)
174 case _:
175 logger.warning(f"Handling signal ({signum}) by doing nothing.")
177 signal.signal(signal.SIGINT, handle_signal)
178 signal.signal(signal.SIGTERM, handle_signal)
180 def on_shutdown(self):
181 logger.warning("[on_shutdown] received SIGINT / SIGTERM, shutdown process")
183 async def on_signal(self, msg: Msg):
184 alpha_id = parse_alpha_id(msg.subject, self.signal_namespace)
185 if alpha_id not in self.alpha_ids:
186 # Not one of this portfolio's alphas — ignore rather than store a
187 # stray id (e.g. another tenant's signal that slipped through).
188 logger.debug(f"[on_signal] ignoring signal for unknown alpha {alpha_id}")
189 return
190 payload = json.loads(msg.data.decode())
191 signal = float(payload["signal"])
192 if not math.isfinite(signal):
193 logger.error(
194 f"[on_signal] ignoring non-finite signal from {alpha_id}: {signal}"
195 )
196 return
198 self.portfolio.update_signal(alpha_id, signal)
200 logger.info(f"[on_signal] received signal from {alpha_id} signal: {signal}")
202 async def on_aggregate(self):
203 try:
204 target_assets_df = self.portfolio.get_signal()
205 # check if bounds are satisfied
206 if not target_assets_df["weighted_signal"].abs().sum() <= 1:
207 logger.error("Sum of fund percentage is out of bounds")
208 raise Exception("Total assets percentage cannot exceed ±1.0")
210 target_assets = dict(
211 zip(target_assets_df["base_asset"], target_assets_df["weighted_signal"])
212 )
213 logger.info(f"[on_aggregate] latest target assets: {target_assets}")
215 # create portfolio signal (aegis)
216 await self.metric_builder.create_portfolio_signal(
217 portfolio_id=self.portfolio.id, signals=target_assets
218 )
220 # send signal to upstream
221 payload = json.dumps(
222 {
223 "assets": {
224 asset: str(
225 Decimal(str(signal)).quantize(
226 Decimal("0.01"), rounding=ROUND_HALF_UP
227 )
228 )
229 for asset, signal in target_assets.items()
230 },
231 "timestamp": time.time_ns(),
232 }
233 ).encode()
234 await self.metric_builder.metric_stream.publish(
235 portfolio_signal_subject(self.portfolio.id, self.signal_namespace),
236 payload,
237 )
238 except Exception as e:
239 logger.warning(f"[on_aggregate] an exception has been raised: {e}")
240 await self.metric_builder.create_portfolio_alert(
241 portfolio_id=self.portfolio.id,
242 title="Exception occured in on_aggregate",
243 description=str(e),
244 priority=1,
245 )
247 async def start(self):
248 # List of jobs to schedule in the background
249 await self.scheduler.schedule(
250 handler=self.on_aggregate, trigger=self.aggregate_window
251 )
253 await self.metric_builder.metric_stream.subscribe(
254 alpha_signal_subscription(self.signal_namespace), callback=self.on_signal
255 )
256 await asyncio.gather(self.scheduler.start())
259class AlphaExecutor:
260 def __init__(
261 self,
262 alphas: list[Alpha],
263 dataloader: DataLoader,
264 metric_stream: MetricStream,
265 datasource_stream: DatasourceStream,
266 init_batch_size: int = 50,
267 signal_namespace: str | None = None,
268 insert_prefix: str = DEFAULT_METRIC_NAMESPACE,
269 resync_interval: timedelta | None = None,
270 ):
271 for alpha in alphas:
272 if len(alpha.data_infos) == 0:
273 raise ValueError(f"Alpha {alpha.id} has 0 data info")
275 self.alphas = alphas
276 # Publish signals under this namespace token so only same-namespace
277 # PortfolioExecutors receive them (see alpha_signal_subject).
278 self.signal_namespace = signal_namespace
279 self.scheduler = Scheduler()
280 self.datamap = Datamap(
281 data_infos=list(flat_map(lambda a: a.data_infos, self.alphas))
282 )
283 self.dataloader = dataloader
284 self.aegis = MetricBuilder(metric_stream, insert_prefix)
285 self.stream = datasource_stream
286 self.init_batch_size = init_batch_size
287 # Cadence of the REST-based fallback that re-triggers every alpha when
288 # the WS feed is unreliable. Default: half the finest signal interval, so
289 # we self-fetch at least once between consecutive candle closes. Pass 0/
290 # None-equivalent to disable by passing an explicit huge interval.
291 self.resync_interval = resync_interval
293 setup_logger()
295 def handle_signal(signum, _):
296 match signum:
297 case signal.SIGINT | signal.SIGTERM:
298 self.on_shutdown()
299 exit(0)
300 case _:
301 logger.warning(f"Handling signal ({signum}) by doing nothing.")
303 signal.signal(signal.SIGINT, handle_signal)
304 signal.signal(signal.SIGTERM, handle_signal)
306 async def on_event(self, event: Event):
307 match event.event_type:
308 case EventType.DatasourceSubscribed:
309 logger.info(f"[on_event] subscribed to datasource: {event.data}")
310 case EventType.DatasourceUpdate:
311 logger.info(f"[on_event] receive data: {event.data}")
312 data = cast(CollectedData, event.data)
313 topic = Topic.from_str(data["topic"])
315 # update datamap with latest data
316 for _data in data["data"]:
317 try:
318 self.datamap.update(topic, cast(Data, flatten_json(_data)))
319 except Exception as e:
320 logger.error(f"Failed to update Datamap due to: {e}")
322 # resync data by fetching from rest again
323 before_df = self.datamap.map[topic].to_df()
324 before_df = before_df.select(sorted(before_df.columns))
325 await self.datamap.resync(topic=topic, dataloader=self.dataloader)
327 after_df = self.datamap.map[topic].to_df()
328 # after_df = self.data_manager.map[topic].to_df()
329 after_df = after_df.select(sorted(after_df.columns))
330 if not before_df.equals(after_df):
331 for alpha in self.alphas:
332 if topic in alpha.data_infos:
333 await self.aegis.create_alpha_alert(
334 alpha_id=alpha.id,
335 title=f"[resync] [{topic}] different data prior to resyncing",
336 description=f"before resync: {before_df}\nafter resync: {after_df}",
337 priority=3,
338 )
339 logger.warning(
340 f"[resync] [{topic}] different data prior to resyncing"
341 )
342 logger.warning(f"[resync] [{topic}] before resync: {before_df}")
343 logger.warning(f"[resync] [{topic}] after resync: {after_df}")
345 # process the data for each alpha that uses that data
346 for alpha in self.alphas:
347 if Topic.from_str(data["topic"]) not in map(
348 lambda di: Topic.from_str(di.topic), alpha.data_infos
349 ):
350 continue
352 infos = acq_info(data["topic"], alpha.data_infos)
353 if len(infos) == 0:
354 continue
356 # insert alpha_trigger
357 await self.aegis.create_alpha_trigger(
358 alpha_id=alpha.id, topic=infos[0].topic
359 )
361 for info in infos:
362 last_closed_time = Topic.from_str(info.topic).last_closed_time(
363 True
364 )
365 if last_closed_time is None:
366 logger.error(f"{info.topic} has no interval")
367 continue
368 await self._emit_signal(alpha, last_closed_time)
370 case _:
371 logger.info(
372 f"[on_event] got {event.event_type} event with data: {event.data}"
373 )
375 def on_shutdown(self):
376 logger.warning("[on_shutdown] received SIGINT / SIGTERM, shutdown process")
378 async def _emit_signal(self, alpha: Alpha, last_closed_time: datetime) -> bool:
379 """Process the current datamap for `alpha`, compute its signal and
380 publish it (dashboard insert + portfolio routing). Returns True if a
381 signal was emitted. Shared by the WS path (on_event) and the periodic
382 resync fallback so both produce identical signals."""
383 df = alpha.data_processor.process(
384 datamap=self.datamap,
385 last_closed_time=last_closed_time,
386 )
387 if df is None:
388 return False
390 signal_df = alpha.next(df)
391 if len(signal_df) == 0:
392 logger.error(f"{alpha.id} couldn't generate signal")
393 logger.error(f"{alpha.id} df: {df}")
394 return False
396 logging.info(
397 f"[latest_signal] {alpha.id} data_df: {df}\nsignal_df: {signal_df}"
398 )
399 # stamp the signal with ITS OWN bar, not the wall clock: the emit can lag
400 # the close (ws delivery, resync retries) and a resync replay re-emits an
401 # older bar, both of which skew a now() timestamp.
402 signal, bar_time_ns = _latest_signal(signal_df)
404 payload = json.dumps({"signal": signal, "timestamp": bar_time_ns}).encode()
405 # insert alpha_signal (aegis) — dashboard/clickhouse metric. Best-effort:
406 # a metric-insert failure (e.g. a missing/unhealthy JetStream stream)
407 # must NOT block portfolio routing below.
408 try:
409 await self.aegis.create_alpha_signal(
410 alpha_id=alpha.id, signal=signal, timestamp=bar_time_ns
411 )
412 except Exception as e:
413 logger.error(
414 f"[_emit_signal] {alpha.id} alpha_signal metric insert failed: {e}"
415 )
416 # send signal to portfolios (core NATS, independent of the metric path)
417 subject = alpha_signal_subject(alpha.id, self.signal_namespace)
418 await self.aegis.metric_stream.publish(subject, payload)
419 logger.info(
420 f"[_emit_signal] published {alpha.id} -> {subject} signal: {signal}"
421 )
422 return True
424 def _finest_interval(self) -> timedelta | None:
425 """Smallest signal interval across all alphas' data topics."""
426 intervals = [
427 iv
428 for di in self.datamap.data_infos
429 if (iv := Topic.from_str(di.topic).interval()) is not None
430 ]
431 return min(intervals) if intervals else None
433 async def _resync_topic(
434 self,
435 topic: Topic,
436 max_retries: int = 3,
437 retry_delay: timedelta = timedelta(seconds=10),
438 ) -> bool:
439 """Pull the latest data for one topic via the REST dataloader, retrying
440 transient failures. Returns True if it resynced. Retries here only delay
441 THIS topic (callers run topics concurrently, so a slow/failing topic
442 never blocks the others)."""
443 for attempt in range(1, max_retries + 1):
444 try:
445 await self.datamap.resync(topic=topic, dataloader=self.dataloader)
446 return True
447 except Exception as e:
448 logger.warning(
449 f"[periodic_resync] fetch {topic} failed "
450 f"(attempt {attempt}/{max_retries}): {e}"
451 )
452 if attempt == max_retries:
453 return False
454 await asyncio.sleep(retry_delay.total_seconds())
455 return False
457 async def on_periodic_resync(self):
458 """Fallback signal source, independent of the flow-NATS WS feed: pull the
459 latest data over REST and re-trigger every alpha, so signals keep flowing
460 even when the live feed drops messages.
462 Topics shared across alphas are resynced ONCE per cycle (not once per
463 alpha — that re-downloaded a shared topic N times), then every alpha is
464 emitted from the shared datamap. Each unit runs in its own task so one
465 slow/failing fetch never stalls the rest."""
466 # 1. Resync each unique topic once across all alphas.
467 topics = {
468 Topic.from_str(di.topic) for alpha in self.alphas for di in alpha.data_infos
469 }
470 await asyncio.gather(*(self._resync_topic(topic) for topic in topics))
472 # 2. Emit every alpha from the now-fresh shared datamap.
473 async def emit(alpha: Alpha) -> None:
474 try:
475 last_closed_times = [
476 lct
477 for di in alpha.data_infos
478 if (lct := Topic.from_str(di.topic).last_closed_time(True))
479 is not None
480 ]
481 if not last_closed_times:
482 logger.error(
483 f"[periodic_resync] [{alpha.id}] no topic has an interval"
484 )
485 return
486 # Use the earliest close so every topic the alpha needs has a
487 # closed candle by then.
488 await self._emit_signal(alpha, min(last_closed_times))
489 except Exception as e:
490 logger.warning(f"[periodic_resync] [{alpha.id}] failed: {e}")
492 await asyncio.gather(*(emit(alpha) for alpha in self.alphas))
493 await self.on_health_check()
495 async def on_health_check(self):
496 try:
497 for alpha in self.alphas:
498 # Check whether data is not arrived for more than some time
499 topics = list(
500 map(lambda di: Topic.from_str(di.topic), alpha.data_infos)
501 )
502 dfs = list(
503 map(
504 lambda d: d[1].to_df(),
505 filter(
506 lambda d: d[0] in topics,
507 self.datamap.map.items(),
508 ),
509 )
510 )
511 df = reduce(
512 lambda acc, x: acc.select("start_time").join(
513 x.select("start_time"), on="start_time"
514 ),
515 dfs,
516 )
517 last_joined_data_time: datetime = df["start_time"][-1]
518 last_closed_time = min(
519 map(
520 lambda info: cast(
521 datetime, Topic.from_str(info.topic).last_closed_time(True)
522 ),
523 self.datamap.data_infos,
524 )
525 )
526 interval = max(
527 map(
528 lambda info: cast(
529 timedelta, Topic.from_str(info.topic).interval()
530 ),
531 self.datamap.data_infos,
532 )
533 )
535 if last_closed_time - last_joined_data_time >= 2 * interval:
536 logger.warning(
537 f"[on_health_check] [{alpha.id}] Data not arrived for more than {2 * interval}, latest_datapoint: {last_joined_data_time}, should_have: {last_closed_time}"
538 )
539 await self.aegis.create_alpha_alert(
540 alpha_id=alpha.id,
541 title=f"Data not arrived for more than {2 * interval}",
542 description=f"latest_datapoint: {last_joined_data_time}, should_have: {last_closed_time}",
543 priority=1,
544 )
545 except Exception as e:
546 logging.warning(f"[on_health_check] failed to health check: {e}")
548 async def _start_datasource(self):
549 topics = list(
550 set(
551 map(
552 lambda di: Topic.from_str(di.topic),
553 flat_map(lambda a: a.data_infos, self.alphas),
554 )
555 )
556 )
558 data_stream = await self.stream.connect(topics)
559 if data_stream is None:
560 raise Exception("Failed to connect to Datasource WS after max_retries")
562 async for data in data_stream:
563 try:
564 if "data" in data:
565 await self.on_event(
566 Event(
567 event_type=EventType.DatasourceUpdate,
568 orig=data,
569 data=data,
570 )
571 )
572 else:
573 await self.on_event(
574 Event(
575 event_type=EventType.DatasourceSubscribed,
576 orig=data,
577 data=data,
578 )
579 )
580 except Exception as e:
581 logger.warning(f"Datasource WS encountered an exception: {e}")
582 traceback.print_exc()
584 if "data" in data:
585 data = cast(CollectedData, data)
586 for alpha in self.alphas:
587 if Topic.from_str(data["topic"]) in map(
588 lambda di: Topic.from_str(di.topic), alpha.data_infos
589 ):
590 await self.aegis.create_alpha_alert(
591 alpha_id=alpha.id,
592 title="Exception raised in `on_event`",
593 description=str(e),
594 priority=2,
595 )
596 else:
597 for alpha in self.alphas:
598 await self.aegis.create_alpha_alert(
599 alpha_id=alpha.id,
600 title="Exception raised in `on_event`",
601 description=str(e),
602 priority=2,
603 )
605 async def start(self):
606 logging.info("[init] initializing Datamap")
608 topics = {
609 Topic.from_str(data_info.topic) for data_info in self.datamap.data_infos
610 }
611 async with asyncio.TaskGroup() as tg:
612 for topic in topics:
613 tg.create_task(
614 self.datamap.resync(topic=topic, dataloader=self.dataloader)
615 )
617 # REST resync fallback: re-trigger every alpha on a timer so signals keep
618 # flowing if the flow-NATS WS feed drops. Default cadence = finest signal
619 # interval / 2 (e.g. a 1h alpha resyncs every 30m).
620 resync_interval = self.resync_interval or (
621 (finest := self._finest_interval()) and finest / 2
622 )
623 if resync_interval is not None:
624 logging.info(f"[init] scheduling periodic resync every {resync_interval}")
625 await self.scheduler.schedule(
626 handler=self.on_periodic_resync,
627 trigger=Trigger.Interval(resync_interval),
628 )
629 else:
630 logging.warning(
631 "[init] no topic declares an interval — periodic resync disabled"
632 )
634 await asyncio.gather(
635 self.scheduler.start(),
636 self._start_datasource(),
637 )