Integrate predictionengine into the workflows

application/datamanager/src/datamanager/config.py

@@ -19,7 +19,8 @@ class Bucket(BaseModel):
     @computed_field
     def daily_bars_path(self) -> str:
         if self.name is None:
-            raise ValueError("DATA_BUCKET environment variable is required")
+            msg = "DATA_BUCKET environment variable is required"
+            raise ValueError(msg)
         return f"gs://{self.name}/equity/bars/"
 
 

application/datamanager/src/datamanager/main.py

@@ -137,8 +137,8 @@ async def fetch_equity_bars(request: Request, summary_date: SummaryDate) -> Bars
     polygon = request.app.state.settings.polygon
     bucket = request.app.state.settings.gcp.bucket
 
-    summary_date: str = summary_date.date.strftime("%Y-%m-%d")
-    url = f"{polygon.base_url}{polygon.daily_bars}{summary_date}"
+    request_summary_date: str = summary_date.date.strftime("%Y-%m-%d")
+    url = f"{polygon.base_url}{polygon.daily_bars}{request_summary_date}"
     logger.info(f"polygon_api_endpoint={url}")
 
     params = {"adjusted": "true", "apiKey": polygon.api_key}
@@ -178,7 +178,7 @@ async def fetch_equity_bars(request: Request, summary_date: SummaryDate) -> Bars
                 status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
                 detail="Failed to write data",
             ) from e
-    return BarsSummary(date=summary_date, count=count)
+    return BarsSummary(date=request_summary_date, count=count)
 
 
 @application.delete("/equity-bars")

application/positionmanager/src/positionmanager/clients.py

@@ -1,9 +1,7 @@
-import requests
-import polars as pl
-from typing import Dict, Any
-import pyarrow as pa
+from typing import Any
 
 import polars as pl
+import pyarrow as pa
 import requests
 from alpaca.trading.client import TradingClient
 from alpaca.trading.enums import OrderSide, TimeInForce
@@ -18,13 +16,11 @@ def __init__(
         *,
         api_key: str | None = "",
         api_secret: str | None = "",
-        api_key: str | None = None,
-        api_secret: str | None = None,
         paper: bool = True,
     ) -> None:
         if not api_key or not api_secret:
-            msg = "Alpaca API key and secret are required"
-            raise ValueError(msg)
+            message = "Alpaca API key and secret are required"
+            raise ValueError(message)
 
         self.trading_client = TradingClient(api_key, api_secret, paper=paper)
 
@@ -33,7 +29,8 @@ def get_cash_balance(self) -> Money:
         cash_balance = getattr(account, "cash", None)
 
         if cash_balance is None:
-            raise ValueError("Cash balance is not available")
+            message = "Cash balance is not available"
+            raise ValueError(message)
 
         return Money.from_float(float(cash_balance))
 
@@ -74,8 +71,8 @@ def get_data(
         date_range: DateRange,
     ) -> pl.DataFrame:
         if not self.datamanager_base_url:
-            msg = "Data manager URL is not configured"
-            raise ValueError(msg)
+            message = "Data manager URL is not configured"
+            raise ValueError(message)
 
         endpoint = f"{self.datamanager_base_url}/equity-bars"
 
@@ -87,14 +84,15 @@ def get_data(
         try:
             response = requests.get(endpoint, params=params, timeout=30)
         except requests.RequestException as err:
-            msg = f"Data manager service call error: {err}"
-            raise RuntimeError(msg) from err
+            message = f"Data manager service call error: {err}"
+            raise RuntimeError(message) from err
 
-        if response.status_code == 404:
+        if response.status_code == requests.codes.not_found:
             return pl.DataFrame()
-        if response.status_code != 200:
+        if response.status_code != requests.codes.ok:
+            message = f"Data service error: {response.text}, status code: {response.status_code}"  # noqa: E501
             raise requests.HTTPError(
-                f"Data service error: {response.text}, status code: {response.status_code}",
+                message,
             )
 
         buffer = pa.py_buffer(response.content)

application/positionmanager/src/positionmanager/main.py

@@ -1,15 +1,10 @@
 import os
-from datetime import datetime, timedelta, timezone
-import polars as pl
-from typing import Dict, Any
-from .models import Money, DateRange, PredictionPayload
-from .clients import AlpacaClient, DataClient
-from .portfolio import PortfolioOptimizer
-from prometheus_fastapi_instrumentator import Instrumentator
+from datetime import UTC, datetime, timedelta
+from typing import Any
 
 import polars as pl
 import requests
-from alpaca.common.rest import APIError
+from alpaca.common.exceptions import APIError
 from fastapi import FastAPI, HTTPException
 from prometheus_fastapi_instrumentator import Instrumentator
 from pydantic import ValidationError

application/predictionengine/src/predictionengine/dataset.py

@@ -1,8 +1,9 @@
-from typing import Dict, List, Any, Tuple, Generator
-from tinygrad.tensor import Tensor
+from collections.abc import Generator
+from typing import Any
+
 import polars as pl
 from category_encoders import OrdinalEncoder
-
+from tinygrad.tensor import Tensor
 
 continuous_variable_columns = [
     "open_price",
@@ -20,13 +21,15 @@ def __init__(
         batch_size: int,
         sequence_length: int,
         sample_count: int,
-        scalers: Dict[str, Dict[str, Tensor]] = {},
+        scalers: dict[str, dict[str, Tensor]] | None = None,
     ) -> None:
+        if scalers is None:
+            scalers = {}
         self.batch_size = batch_size
         self.sequence_length = sequence_length
         self.sample_count = sample_count
         self.scalers = scalers if scalers is not None else {}
-        self.preprocessors: Dict[str, Any] = {}
+        self.preprocessors: dict[str, Any] = {}
 
     def __len__(self) -> int:
         return (self.sample_count + self.batch_size - 1) // self.batch_size
@@ -106,7 +109,7 @@ def _encode_tickers(self, data: pl.DataFrame) -> pl.DataFrame:
 
     def _compute_scalers(self, data: pl.DataFrame) -> None:
         if len(self.scalers) == 0:
-            self.scalers: Dict[str, Dict[str, Tensor]] = {}
+            self.scalers: dict[str, dict[str, Tensor]] = {}
             for ticker_key, group in data.group_by("ticker"):
                 ticker = ticker_key[0]
                 means = group[continuous_variable_columns].mean()
@@ -118,7 +121,7 @@ def _compute_scalers(self, data: pl.DataFrame) -> None:
                 }
 
     def _scale_data(self, data: pl.DataFrame) -> Tensor:
-        groups: List[Tensor] = []
+        groups: list[Tensor] = []
         for ticker_key, group in data.group_by("ticker"):
             ticker = ticker_key[0]
             means = self.scalers[str(ticker)]["means"]
@@ -133,7 +136,8 @@ def _scale_data(self, data: pl.DataFrame) -> Tensor:
             groups.append(combined_group)
 
         if not groups:
-            raise ValueError("No data available after preprocessing")
+            message = "No data available after preprocessing"
+            raise ValueError(message)
 
         output_data = Tensor.empty(groups[0].shape)
         return output_data.cat(*groups, dim=0)
@@ -150,9 +154,10 @@ def load_data(self, data: pl.DataFrame) -> None:
         self._compute_scalers(data)
         self.data = self._scale_data(data)
 
-    def get_preprocessors(self) -> Dict[str, Any]:
+    def get_preprocessors(self) -> dict[str, Any]:
         if not self.preprocessors:
-            raise ValueError("Preprocessors have not been initialized.")
+            message = "Preprocessors have not been initialized."
+            raise ValueError(message)
 
         means_by_ticker = {
             ticker: values["means"] for ticker, values in self.scalers.items()
@@ -169,7 +174,7 @@ def get_preprocessors(self) -> Dict[str, Any]:
             "indices": self.preprocessors["indices"],
         }
 
-    def batches(self) -> Generator[Tuple[Tensor, Tensor, Tensor], None, None]:
+    def batches(self) -> Generator[tuple[Tensor, Tensor, Tensor], None, None]:
         close_price_idx = self.preprocessors["indices"]["close_price"]
 
         for i in range(0, self.sample_count, self.batch_size):
@@ -193,9 +198,8 @@ def batches(self) -> Generator[Tuple[Tensor, Tensor, Tensor], None, None]:
             ]
 
             if not batch_tensors:
-                raise ValueError(
-                    "Cannot stack empty batch tensors (batch_size must be ≥ 1)"
-                )
+                message = "Cannot stack empty batch tensors (batch_size must be ≥ 1)"
+                raise ValueError(message)
             if len(batch_tensors) == 1:
                 historical_features = batch_tensors[0].unsqueeze(0)
             else:

application/predictionengine/src/predictionengine/gated_residual_network.py

@@ -1,15 +1,15 @@
 from typing import cast
+
+from tinygrad.nn import LayerNorm, Linear
 from tinygrad.tensor import Tensor
-from tinygrad.nn import Linear, LayerNorm
-from typing import Optional
 
 
 class GatedResidualNetwork:
     def __init__(
         self,
         input_size: int,
         hidden_size: int,
-        output_size: Optional[int] = None,
+        output_size: int | None = None,
     ) -> None:
         output_size = output_size if output_size is not None else input_size
 
@@ -30,18 +30,18 @@ def __init__(
 
     def forward(
         self,
-        input: Tensor,
+        features: Tensor,
     ) -> Tensor:
-        hidden_state = self.dense_input(input).relu()
+        hidden_state = self.dense_input(features).relu()
 
         output_state = self.dense_output(hidden_state)
 
         gate_state = self.gate(hidden_state).sigmoid()
 
         if self.residual_projection is not None:
-            residual = self.residual_projection(input)
+            residual = self.residual_projection(features)
         else:
-            residual = input
+            residual = features
 
-        gated_output = cast(Tensor, gate_state * output_state + residual)
+        gated_output = cast("Tensor", gate_state * output_state + residual)
         return self.layer_normalizer(gated_output)

application/predictionengine/src/predictionengine/long_short_term_memory.py

@@ -1,6 +1,5 @@
-from typing import List, Tuple
-from tinygrad.tensor import Tensor
 from tinygrad.nn import LSTMCell
+from tinygrad.tensor import Tensor
 
 
 class LongShortTermMemory:
@@ -15,16 +14,16 @@ def __init__(
         self.layer_count = layer_count
         self.dropout_rate = dropout_rate
 
-        self.layers: List[LSTMCell] = []
+        self.layers: list[LSTMCell] = []
         for index in range(layer_count):
             input_size = input_size if index == 0 else self.hidden_size
             self.layers.append(LSTMCell(input_size, self.hidden_size))
 
     def forward(
         self,
-        input: Tensor,
-    ) -> Tuple[Tensor, Tuple[Tensor, Tensor]]:
-        batch_size, sequence_length, _ = input.shape
+        features: Tensor,
+    ) -> tuple[Tensor, tuple[Tensor, Tensor]]:
+        batch_size, sequence_length, _ = features.shape
 
         hidden_state = Tensor.zeros(
             self.layer_count, batch_size, self.hidden_size
@@ -36,7 +35,7 @@ def forward(
         outputs = []
 
         for t in range(int(sequence_length)):
-            layer_input = input[:, t]
+            layer_input = features[:, t]
 
             for index, layer in enumerate(self.layers):
                 layer_hidden_state, layer_cell_state = layer(
@@ -59,8 +58,10 @@ def forward(
             outputs.append(hidden_state[-1])
 
         if not outputs:
-            raise ValueError("Cannot stack empty outputs list")
-        elif len(outputs) == 1:
+            message = "Cannot stack empty outputs list"
+            raise ValueError(message)
+
+        if len(outputs) == 1:
             output_tensor = outputs[0].unsqueeze(1)
         else:
             output_tensor = Tensor.stack(outputs[0], *outputs[1:], dim=1)

application/predictionengine/src/predictionengine/loss_function.py

@@ -1,6 +1,7 @@
-from tinygrad.tensor import Tensor
 from typing import cast
 
+from tinygrad.tensor import Tensor
+
 Quantiles = tuple[float, float, float] | tuple[float, float, float, float, float]
 
 
@@ -11,18 +12,18 @@ def quantile_loss(
         quantiles = (0.25, 0.5, 0.75)
 
     if y_pred.shape != y_true.shape:
-        raise ValueError(
-            f"Shape mismatch: y_pred {y_pred.shape} vs y_true {y_true.shape}"
-        )
+        message = f"Shape mismatch: y_pred {y_pred.shape} vs y_true {y_true.shape}"
+        raise ValueError(message)
 
     if not all(0 <= q <= 1 for q in quantiles):
-        raise ValueError("All quantiles must be between 0 and 1")
+        message = "All quantiles must be between 0 and 1"
+        raise ValueError(message)
 
     loss: Tensor = Tensor.zeros(1)
-    error = cast(Tensor, y_true - y_pred)
+    error = cast("Tensor", y_true - y_pred)
     for quantile in quantiles:
-        quantile_error = cast(Tensor, quantile * error)
-        quantile_minus_one_error = cast(Tensor, (quantile - 1) * error)
+        quantile_error = cast("Tensor", quantile * error)
+        quantile_minus_one_error = cast("Tensor", (quantile - 1) * error)
         loss += Tensor.maximum(quantile_error, quantile_minus_one_error).mean()
 
     return loss