triangular_arbitrage_bot/oe_em/__main__.py

223 lines
7.6 KiB
Python

"""
Opportunity engine entry point.
Initialises the order-book consumer, triangle index, and signal socket client;
starts background tasks for book consumption and periodic stats logging; shuts
down cleanly on SIGTERM/SIGINT.
"""
import asyncio
import signal
from pathlib import Path
import aiohttp
import structlog
from common.log import configure_logging
from oe_em.book_consumer import BookConsumer
from oe_em.config import Settings
from oe_em.kucoin_api import KuCoinAPI
from oe_em.opportunity import OpportunityEngine
from oe_em.risk import RiskManager
from oe_em.socket_client import SignalSocketClient
from oe_em.triangle_enum import TradingPair, enumerate_triangles
async def sync_symbols_with_fh_ob(
fh_ob_url: str,
needed_symbols: set[str],
http_session: aiohttp.ClientSession,
log,
) -> set[str]:
"""
Ensure fh_ob is subscribed to every symbol needed by the triangle index.
Fetches the current subscription list from fh_ob, posts any missing symbols,
and returns the full set of subscribed symbols.
"""
get_url = f"{fh_ob_url}/symbols"
async with http_session.get(get_url) as resp:
resp.raise_for_status()
data = await resp.json()
current_symbols = set(data.get("symbols", []))
missing = needed_symbols - current_symbols
if missing:
log.info("syncing_symbols", missing=len(missing), current=len(current_symbols))
for sym in missing:
post_url = f"{fh_ob_url}/symbols"
payload = {"symbol": sym}
async with http_session.post(post_url, json=payload) as post_resp:
if post_resp.status == 400:
log.warning("symbol_cannot_be_subscribed", symbol=sym)
else:
log.debug("symbol_subscribed", symbol=sym)
return current_symbols | missing
async def main() -> None:
config_path = Path("config.yaml")
settings = await Settings.from_yaml(config_path) if config_path.exists() else Settings()
configure_logging(settings.oe_em.log_level, settings.oe_em.log_file)
log = structlog.get_logger().bind(component="oe_em")
log.info("oe_em_starting")
api = KuCoinAPI()
await api.fetch_pairs_and_fees()
pair_list = [
TradingPair(
symbol=p["symbol"],
base=p["base"],
quote=p["quote"],
fee_currency=p.get("fee_currency", ""),
)
for p in api.get_all_pairs()
]
excluded = set(settings.oe_em.excluded_currencies)
if excluded:
pair_list = [p for p in pair_list if p.base not in excluded and p.quote not in excluded]
log.info("pairs_loaded", count=len(pair_list), excluded=len(excluded))
fee_table = api._fee_table
triangle_index = enumerate_triangles(
pair_list,
fee_table,
hold_currencies=settings.oe_em.hold_currencies,
)
log.info("triangles_enumerated", count=len(triangle_index.triangles))
needed_symbols = set()
for tri in triangle_index.triangles:
needed_symbols.update(tri.pair_symbols)
async with aiohttp.ClientSession() as http_session:
subscribed = await sync_symbols_with_fh_ob(
settings.oe_em.fh_ob_url,
needed_symbols,
http_session,
log,
)
book_consumer = BookConsumer(
socket_path=settings.oe_em.socket_path,
on_update=lambda symbol, book: None,
)
signal_client: SignalSocketClient | None = None
signal_reconnect_task: asyncio.Task | None = None
if settings.oe_em.send_signals:
signal_client = SignalSocketClient(socket_path=settings.oe_em.executor_socket_path)
async def send_signal(signal_payload: dict) -> None:
"""Forward a signal payload to the executor's Unix socket."""
if signal_client:
await signal_client.send_signal(signal_payload)
opp_engine = OpportunityEngine(
book_consumer=book_consumer,
triangle_index=triangle_index,
signal_threshold_bps=settings.oe_em.signal_threshold_bps,
log_path=settings.oe_em.opportunity_log_path,
kcs_discount_active=settings.oe_em.kcs_discount_active,
cooldown_seconds=settings.oe_em.cooldown_seconds,
on_signal=send_signal if settings.oe_em.send_signals else None,
)
risk_mgr = RiskManager()
async def on_book_update(symbol: str, book) -> None:
"""Callback invoked by BookConsumer whenever a subscribed book is refreshed."""
if risk_mgr.should_continue():
opp_engine.evaluate_triangles_for_pair(symbol)
book_consumer.set_on_update(on_book_update)
fh_ob_url = settings.oe_em.fh_ob_url
async with aiohttp.ClientSession() as http_session:
async with http_session.get(f"{fh_ob_url}/symbols") as resp:
resp.raise_for_status()
data = await resp.json()
symbols_now = set(data.get("symbols", []))
missing = needed_symbols - symbols_now
if missing:
log.warning("symbols_not_subscribed_after_sync", count=len(missing))
for sym in missing:
log.warning("unavailable_symbol", symbol=sym)
if signal_client:
signal_reconnect_task = await signal_client.start()
log.info("signal_client_connecting", socket_path=str(settings.oe_em.executor_socket_path))
log.info(
"oe_em_ready",
triangles=len(triangle_index.triangles),
subscribed=len(symbols_now),
threshold_bps=settings.oe_em.signal_threshold_bps,
hold_currencies=settings.oe_em.hold_currencies,
send_signals=settings.oe_em.send_signals,
)
consumer_task = asyncio.create_task(book_consumer.start())
async def stats_loop() -> None:
"""
Periodically log evaluation stats so the operator can monitor the engine.
Runs until cancelled. Suppressed by setting stats_interval_seconds <= 0.
"""
interval = settings.oe_em.stats_interval_seconds
if interval <= 0:
return
while True:
await asyncio.sleep(interval)
try:
s = opp_engine.get_stats()
books_tracked = sum(
1 for t in triangle_index.triangles
if book_consumer.get_book(t.legs[0].pair.symbol) is not None
)
except Exception as e:
log.error("stats_error", error=str(e))
continue
log.info("stats", **{
"triangles_evaluated": s.triangles_evaluated,
"signals_fired": s.signals_fired,
"books_tracked": books_tracked,
"subscribed": len(book_consumer._books),
"best_net_bps": f"{s.best_net_bps:.2f}",
"best_legs": s.best_legs,
})
stats_task = asyncio.create_task(stats_loop())
def shutdown(sig: signal.Signals) -> None:
"""Begin graceful shutdown: stop book consumer, cancel stats, close signal client."""
log.info("shutdown_signal_received", signal=sig.name)
asyncio.create_task(book_consumer.stop())
stats_task.cancel()
if signal_reconnect_task:
signal_reconnect_task.cancel()
if signal_client:
asyncio.create_task(signal_client.close())
loop = asyncio.get_running_loop()
for sig in (signal.SIGTERM, signal.SIGINT):
loop.add_signal_handler(sig, lambda s=sig: shutdown(s))
tasks = [consumer_task, stats_task]
if signal_reconnect_task:
tasks.append(signal_reconnect_task)
try:
await asyncio.gather(*tasks)
except asyncio.CancelledError:
log.info("oe_em_cancelled")
if __name__ == "__main__":
asyncio.run(main())