Expand initial Grafana metrics resources and upload script

application/datamanager/src/datamanager/main.py

@@ -17,6 +17,7 @@
 from google.cloud import storage  # type: ignore
 from loguru import logger
 from polars.exceptions import ComputeError
+from prometheus_client import Gauge
 from prometheus_fastapi_instrumentator import Instrumentator
 
 from .config import Settings
@@ -75,12 +76,42 @@ async def lifespan(app: FastAPI) -> AsyncGenerator[None, None]:
 application = FastAPI(lifespan=lifespan)
 Instrumentator().instrument(application).expose(application)
 
+equity_bars_total_rows = Gauge(
+    "equity_bars_total_rows",
+    "Total number of rows in equity bars bucket",
+)
+
 
 @application.get("/health")
 async def health_check() -> Response:
     return Response(status_code=status.HTTP_200_OK)
 
 
+@application.get("/metrics")
+async def update_metrics(request: Request) -> dict[str, int]:
+    settings: Settings = request.app.state.settings
+
+    count_query = f"""
+        SELECT COUNT(*) as total_rows
+        FROM read_parquet(
+            'gs://{settings.gcp.bucket.name}/equity/bars/*/*/*/*', 
+            HIVE_PARTITIONING=1
+        )
+    """  # noqa: S608
+
+    try:
+        result = request.app.state.connection.execute(count_query).fetchone()
+        total_rows = result[0] if result else 0
+        equity_bars_total_rows.set(total_rows)
+
+        logger.info(f"Updated equity_bars_total_rows metric: {total_rows}")
+        return {"total_rows": total_rows}  # noqa: TRY300
+
+    except requests.RequestException as e:
+        logger.error(f"Error updating metrics: {e}")
+        return {"total_rows": 0}
+
+
 @application.get("/equity-bars")
 async def get_equity_bars(
     request: Request,

application/positionmanager/src/positionmanager/clients.py

@@ -1,4 +1,4 @@
-from typing import Any
+from typing import TYPE_CHECKING, Any, cast
 
 import polars as pl
 import pyarrow as pa
@@ -9,6 +9,9 @@
 
 from .models import DateRange, Money
 
+if TYPE_CHECKING:
+    from alpaca.trading.models import Position, TradeAccount
+
 
 class AlpacaClient:
     def __init__(
@@ -23,7 +26,10 @@ def __init__(
             raise ValueError(message)
 
         self.trading_client: TradingClient = TradingClient(
-            api_key, api_secret, paper=paper
+            api_key,
+            api_secret,
+            paper=paper,
+            raw_data=False,
         )
 
     def get_cash_balance(self) -> Money:
@@ -63,6 +69,39 @@ def clear_positions(self) -> dict[str, Any]:
             "message": "All positions have been closed",
         }
 
+    def get_account_information(self) -> dict[str, Any]:
+        account: TradeAccount = cast("TradeAccount", self.trading_client.get_account())
+        return {
+            "portfolio_value": float(account.portfolio_value or 0),
+            "cash": float(account.cash or 0),
+            "buying_power": float(account.buying_power or 0),
+            "equity": float(account.equity or 0),
+        }
+
+    def get_positions(self) -> list[dict[str, Any]]:
+        positions: list[Position] = cast(
+            "list[Position]",
+            self.trading_client.get_all_positions(),
+        )
+        position_list = []
+
+        for position in positions:
+            position_data = {
+                "symbol": position.symbol,
+                "quantity": float(position.qty or 0),
+                "market_value": float(position.market_value or 0),
+                "cost_basis": float(position.cost_basis or 0),
+                "unrealized_profit_and_loss": float(position.unrealized_pl or 0),
+                "unrealized_profit_and_loss_percent": float(
+                    position.unrealized_plpc or 0
+                ),
+                "current_price": float(position.current_price or 0),
+                "side": position.side.value,
+            }
+            position_list.append(position_data)
+
+        return position_list
+
 
 class DataClient:
     def __init__(self, datamanager_base_url: str | None) -> None:

application/positionmanager/src/positionmanager/main.py

@@ -7,6 +7,7 @@
 import requests
 from alpaca.common.exceptions import APIError
 from fastapi import FastAPI, HTTPException
+from prometheus_client import Gauge
 from prometheus_fastapi_instrumentator import Instrumentator
 from pydantic import ValidationError
 
@@ -20,12 +21,91 @@
 application = FastAPI()
 Instrumentator().instrument(application).expose(application)
 
+portfolio_value_gauge = Gauge(
+    "portfolio_total_value",
+    "Current total portfolio value from Alpaca",
+)
+
+portfolio_cash_balance_gauge = Gauge(
+    "portfolio_cash_balance",
+    "Current cash balance in portfolio",
+)
+
+portfolio_positions_count_gauge = Gauge(
+    "portfolio_positions_count",
+    "Number of current positions in portfolio",
+)
+
+portfolio_position_value_gauge = Gauge(
+    "portfolio_position_value",
+    "Value of specific position",
+    ["symbol"],
+)
+
+portfolio_position_profit_and_loss_gauge = Gauge(
+    "portfolio_position_profit_and_loss",
+    "Unrealized P&L for specific position",
+    ["symbol"],
+)
+
 
 @application.get("/health")
 def get_health() -> dict[str, str]:
     return {"status": "healthy"}
 
 
+@application.get("/metrics")
+def update_metrics() -> dict[str, Any]:
+    alpaca_client = AlpacaClient(
+        api_key=os.getenv("ALPACA_API_KEY", ""),
+        api_secret=os.getenv("ALPACA_API_SECRET", ""),
+        paper=os.getenv("ALPACA_PAPER", "true").lower() == "true",
+    )
+
+    try:
+        account_information = alpaca_client.get_account_information()
+        positions = alpaca_client.get_positions()
+
+        portfolio_value_gauge.set(account_information["portfolio_value"])
+        portfolio_cash_balance_gauge.set(account_information["cash"])
+        portfolio_positions_count_gauge.set(len(positions))
+
+        position_metrics = []
+        for position in positions:
+            symbol = position["symbol"]
+            portfolio_position_value_gauge.labels(symbol=symbol).set(
+                position["market_value"]
+            )
+            portfolio_position_profit_and_loss_gauge.labels(symbol=symbol).set(
+                position["unrealized_profit_and_loss"]
+            )
+
+            position_metrics.append(
+                {
+                    "symbol": symbol,
+                    "market_value": position["market_value"],
+                    "unrealized_profit_and_loss": position[
+                        "unrealized_profit_and_loss"
+                    ],
+                    "cost_basis": position["cost_basis"],
+                    "current_price": position["current_price"],
+                }
+            )
+
+        return {
+            "portfolio_value": account_information["portfolio_value"],
+            "cash_balance": account_information["cash"],
+            "positions_count": len(positions),
+            "positions": position_metrics,
+        }
+
+    except (requests.RequestException, APIError, ValidationError) as e:
+        raise HTTPException(
+            status_code=500,
+            detail=f"Error updating metrics: {e!r}",
+        ) from e
+
+
 @application.post("/positions")
 def create_position(payload: PredictionPayload) -> dict[str, Any]:
     alpaca_client = AlpacaClient(

application/predictionengine/src/predictionengine/dataset.py

@@ -208,7 +208,8 @@ def batches(self) -> Generator[tuple[Tensor, Tensor, Tensor], None, None]:
                 )
 
             targets = batch_data[: self.batch_size, close_price_idx].reshape(
-                self.batch_size, 1
+                self.batch_size,
+                1,
             )
 
             yield tickers, historical_features, targets

infrastructure/__main__.py

@@ -72,3 +72,13 @@
 
 
 export("DATAMANAGER_BASE_URL", datamanager_service.statuses[0].url)
+
+export(
+    "DATAMANAGER_METRICS_URL",
+    datamanager_service.statuses[0].url.apply(lambda url: f"{url}/metrics"),
+)
+
+export(
+    "POSITIONMANAGER_METRICS_URL",
+    positionmanager_service.statuses[0].url.apply(lambda url: f"{url}/metrics"),
+)