Add Docker support and fix benchmark/data pipeline bugs

Containerize the application with ROCm GPU support for AMD Radeon R9700:
- Add Dockerfile with PyTorch/PyG ROCm 5.6 wheels
- Add docker-compose.yml with dashboard, live-trading, and train services
- Add .dockerignore and .env.example for configuration

Fix benchmark script to use batch_size instead of num_stocks for variable
dimensions, and replace fragile partial model surgery with a standalone MLP
for memory estimation.

Fix data pipeline to skip dates with no next trading date instead of
fabricating zero returns.

Add slippage to paper broker, optional mark-to-market prices to broker
interface, and health check endpoint for container orchestration.
This commit is contained in:
2026-05-26 15:07:28 +02:00
parent bc40a67180
commit 46657c7ffe
12 changed files with 471 additions and 40 deletions
+68
View File
@@ -0,0 +1,68 @@
# Git
.git
.gitignore
# Python
__pycache__
*.py[cod]
*$py.class
*.so
.Python
build/
develop-eggs/
dist/
downloads/
eggs/
.eggs/
lib/
lib64/
parts/
sdist/
var/
wheels/
*.egg-info/
.installed.cfg
*.egg
# Virtual environments
venv/
env/
ENV/
# IDE
.vscode
.idea
*.swp
*.swo
*~
# Jupyter
.ipynb_checkpoints
# OS
.DS_Store
Thumbs.db
# Logs
*.log
# Docker
Dockerfile
docker-compose*.yml
.dockerignore
# Documentation
docs/
# Notebooks (heavy, usually not needed in image)
notebooks/*.ipynb
# Keep data/model directories but ignore contents
# (they are mounted as volumes at runtime)
data/raw/*
data/processed/*
data/external/*
models/*
# Don't ignore .env.example so it gets copied into image for reference
!.env.example
+45
View File
@@ -0,0 +1,45 @@
# =============================================================================
# StockGNN R9700 — Environment Variables
# =============================================================================
# Copy this file to .env and fill in your API keys.
# docker compose will automatically load variables from .env at runtime.
# =============================================================================
# -----------------------------------------------------------------------------
# AMD ROCm GPU Settings
# -----------------------------------------------------------------------------
# Override the GFX architecture version if auto-detection fails.
# Common values:
# gfx1030 (Navi 21 / RX 6900 XT / Radeon Pro W6800)
# gfx1100 (Navi 31 / RX 7900 XTX / Radeon Pro W7900)
# gfx1101 (Navi 31 / RX 7900 XT / Radeon Pro W7800)
# gfx1102 (Navi 32 / RX 7700 XT)
# Default assumes gfx1030 for broad compatibility.
# HSA_OVERRIDE_GFX_VERSION=10.3.0
# -----------------------------------------------------------------------------
# Data Provider API Keys
# -----------------------------------------------------------------------------
POLYGON_API_KEY=
ALPHA_VANTAGE_API_KEY=
# -----------------------------------------------------------------------------
# Alternative Data (News & Social)
# -----------------------------------------------------------------------------
NEWS_API_KEY=
TWITTER_BEARER_TOKEN=
REDDIT_CLIENT_ID=
REDDIT_CLIENT_SECRET=
# -----------------------------------------------------------------------------
# Interactive Brokers (if using IB gateway for live trading)
# -----------------------------------------------------------------------------
# IB_HOST=127.0.0.1
# IB_PORT=7497
# IB_CLIENT_ID=1
# -----------------------------------------------------------------------------
# Application Behaviour
# -----------------------------------------------------------------------------
# STOCKGNN_ENV=production
# LOG_LEVEL=INFO
+102
View File
@@ -0,0 +1,102 @@
# =============================================================================
# StockGNN R9700 — ROCm Ubuntu Docker Image
# =============================================================================
# Base image: Ubuntu 22.04 with ROCm 5.6 pre-installed
# Provides AMD GPU support for PyTorch/ PyG on Radeon R9700 AI Pro
# =============================================================================
FROM rocm/dev-ubuntu-22.04:5.6
LABEL maintainer="StockGNN Team"
LABEL description="AMD-optimized GNN trading system (ROCm 5.6 / Ubuntu 22.04)"
# ---------------------------------------------------------------------------
# System dependencies
# ---------------------------------------------------------------------------
ENV DEBIAN_FRONTEND=noninteractive
RUN apt-get update && apt-get install -y --no-install-recommends \
python3.10 \
python3.10-dev \
python3-pip \
git \
wget \
curl \
ca-certificates \
libglib2.0-0 \
libsm6 \
libxext6 \
libxrender-dev \
libgomp1 \
libnuma1 \
&& rm -rf /var/lib/apt/lists/*
# Make python3.10 the default python3
RUN update-alternatives --install /usr/bin/python3 python3 /usr/bin/python3.10 1 \
&& update-alternatives --install /usr/bin/python python /usr/bin/python3.10 1
# Upgrade pip
RUN python -m pip install --no-cache-dir --upgrade pip setuptools wheel
# ---------------------------------------------------------------------------
# Install PyTorch with ROCm 5.6 support (separate from other deps)
# ---------------------------------------------------------------------------
RUN pip install --no-cache-dir \
torch==2.1.0+rocm5.6 \
torchvision==0.16.0+rocm5.6 \
torchaudio==2.1.0+rocm5.6 \
--index-url https://download.pytorch.org/whl/rocm5.6
# ---------------------------------------------------------------------------
# Install PyTorch Geometric and extensions with ROCm 5.6 wheels
# ---------------------------------------------------------------------------
RUN pip install --no-cache-dir \
torch-geometric==2.4.0 \
torch-scatter==2.1.2+pt21rocm5.6 \
torch-sparse==0.6.18+pt21rocm5.6 \
torch-cluster==1.6.2+pt21rocm5.6 \
torch-spline-conv==1.2.2+pt21rocm5.6 \
-f https://data.pyg.org/whl/torch-2.1.0+rocm5.6.html
# ---------------------------------------------------------------------------
# Install remaining Python dependencies from PyPI
# ---------------------------------------------------------------------------
COPY requirements.txt /tmp/requirements.txt
# Strip out the torch / pyg lines and ROCm-specific index directives so we
# install the rest of the packages from standard PyPI.
RUN grep -vE '^(torch|torchvision|torchaudio|torch-geometric|torch-scatter|torch-sparse|torch-cluster|torch-spline-conv|rocblas|hipblaslt|miopen-hip|rccl|--index-url|--find-links)' /tmp/requirements.txt \
> /tmp/requirements-clean.txt || true
RUN pip install --no-cache-dir -r /tmp/requirements-clean.txt
# ---------------------------------------------------------------------------
# ROCm runtime tuning for AMD Radeon R9700 AI Pro (gfx1030 / gfx1100)
# ---------------------------------------------------------------------------
ENV HSA_OVERRIDE_GFX_VERSION=10.3.0
ENV PYTORCH_HIP_ALLOC_CONF=expandable_segments:True
ENV ROCM_PATH=/opt/rocm
ENV PATH="${ROCM_PATH}/bin:${PATH}"
ENV LD_LIBRARY_PATH="${ROCM_PATH}/lib:${LD_LIBRARY_PATH}"
# ---------------------------------------------------------------------------
# Application setup
# ---------------------------------------------------------------------------
WORKDIR /app
# Create directories that the app expects
RUN mkdir -p /app/data/raw /app/data/processed /app/data/external /app/models /app/logs
# Copy source code
COPY . /app/
# Ensure Python can find local src/ modules
ENV PYTHONPATH="/app:${PYTHONPATH}"
# Create non-root user for runtime security
RUN groupadd -r trader && useradd -r -g trader -d /app trader \
&& chown -R trader:trader /app
USER trader
# Default command (can be overridden per-service in docker-compose.yml)
CMD ["python", "-m", "uvicorn", "src.web.app:app", "--host", "0.0.0.0", "--port", "8000"]
+11 -19
View File
@@ -106,11 +106,11 @@ def benchmark_model():
for batch_size in batch_sizes:
for seq_len in sequence_lengths:
# Create data for this configuration
bx = torch.randn(num_stocks, seq_len, num_features).to(config.DEVICE)
bei = torch.randint(0, num_stocks, (2, num_edges)).to(config.DEVICE)
# Create data for this configuration; batch_size drives num nodes
bx = torch.randn(batch_size, seq_len, num_features).to(config.DEVICE)
bei = torch.randint(0, batch_size, (2, num_edges)).to(config.DEVICE)
bea = torch.randn(num_edges, 1).to(config.DEVICE)
by = torch.randn(num_stocks, 1).to(config.DEVICE)
by = torch.randn(batch_size, 1).to(config.DEVICE)
bdata = Data(x=bx, edge_index=bei, edge_attr=bea, y=by)
# Benchmark inference
@@ -136,34 +136,26 @@ def benchmark_model():
f"Inf Tput: {1 / inf_time:.2f} samples/s, Train Tput: {1 / train_time:.2f} samples/s"
)
# Memory benchmark
# Memory benchmark — estimate how memory scales with hidden width.
# Uses a standalone MLP matching IntradayGNN's feature_processor + output layer
# to avoid the dimension mismatches that arise from partial model surgery.
logger.info("\nMemory benchmark:")
# Test different model sizes
hidden_channels_list = [64, 128, 256, 512]
for hidden_channels in hidden_channels_list:
# Create a model with this configuration
model = IntradayGNN(num_features, config.SEQUENCE_LENGTH)
model.feature_processor = nn.Sequential(
bench_model = nn.Sequential(
nn.Linear(num_features, hidden_channels),
nn.SiLU(),
nn.Linear(hidden_channels, hidden_channels),
nn.LayerNorm(hidden_channels),
nn.Linear(hidden_channels, 1),
)
model.linear = nn.Linear(hidden_channels, 1)
# Optimize model
model = amd_optimizer.optimize_model(model)
# Estimate memory usage
estimated_memory = memory_manager.estimate_model_memory(model)
estimated_memory = memory_manager.estimate_model_memory(bench_model)
logger.info(
f"Hidden Channels: {hidden_channels}, Estimated Memory: {estimated_memory / 1024**3:.2f}GB"
)
# Clean up
del model
del bench_model
memory_manager.empty_cache()
# Final memory stats
+183
View File
@@ -0,0 +1,183 @@
# =============================================================================
# StockGNN R9700 — Docker Compose Stack
# =============================================================================
# Services:
# - dashboard : FastAPI web frontend (http://localhost:8000)
# - live-trading: Continuous paper/live trading engine
# - train : One-off model training & backtesting
#
# Requirements:
# - Docker 20.10+ with Compose v2
# - AMD GPU + ROCm 5.6+ drivers on host
# - docker-compose run --rm train # manual training
# =============================================================================
services:
# -------------------------------------------------------------------------
# Base image build (shared by all services)
# -------------------------------------------------------------------------
dashboard:
build:
context: .
dockerfile: Dockerfile
image: stockgnn:r9700
container_name: stockgnn-dashboard
restart: unless-stopped
# Override default CMD for the web dashboard
command: >
python -m uvicorn src.web.app:app
--host 0.0.0.0
--port 8000
--reload
ports:
- "8000:8000"
volumes:
# Persist data & model artefacts across restarts
- ./data:/app/data
- ./models:/app/models
- ./logs:/app/logs
# Optional: mount source for live-reload during development
# - ./src:/app/src
env_file:
- .env
environment:
# ROCm / AMD GPU tuning
HSA_OVERRIDE_GFX_VERSION: "${HSA_OVERRIDE_GFX_VERSION:-10.3.0}"
PYTORCH_HIP_ALLOC_CONF: "expandable_segments:True"
# Ensure the app knows it is inside a container
STOCKGNN_ENV: docker
# AMD GPU device access
devices:
- /dev/kfd
- /dev/dri
# Required groups for GPU access inside the container
group_add:
- video
- render
# Security / capability settings for ROCm
security_opt:
- seccomp:unconfined
# Shared memory size for PyTorch DataLoader multiprocessing
shm_size: "8gb"
# Health-check for the web dashboard
healthcheck:
test: ["CMD", "curl", "-f", "http://localhost:8000/api/dashboard/health"]
interval: 30s
timeout: 10s
retries: 3
start_period: 40s
networks:
- stockgnn-net
# -------------------------------------------------------------------------
# Live Trading Engine (paper or live — controlled via config.py)
# -------------------------------------------------------------------------
live-trading:
build:
context: .
dockerfile: Dockerfile
image: stockgnn:r9700
container_name: stockgnn-live-trading
restart: unless-stopped
command: >
python live_trading.py
volumes:
- ./data:/app/data
- ./models:/app/models
- ./logs:/app/logs
env_file:
- .env
environment:
HSA_OVERRIDE_GFX_VERSION: "${HSA_OVERRIDE_GFX_VERSION:-10.3.0}"
PYTORCH_HIP_ALLOC_CONF: "expandable_segments:True"
STOCKGNN_ENV: docker
devices:
- /dev/kfd
- /dev/dri
group_add:
- video
- render
security_opt:
- seccomp:unconfined
shm_size: "8gb"
networks:
- stockgnn-net
# Do not auto-start live trading until the user explicitly brings it up
profiles:
- live
# -------------------------------------------------------------------------
# Model Training & Backtesting (run manually)
# -------------------------------------------------------------------------
# Usage:
# docker compose run --rm train
# -------------------------------------------------------------------------
train:
build:
context: .
dockerfile: Dockerfile
image: stockgnn:r9700
container_name: stockgnn-train
command: >
python main.py
volumes:
- ./data:/app/data
- ./models:/app/models
- ./logs:/app/logs
env_file:
- .env
environment:
HSA_OVERRIDE_GFX_VERSION: "${HSA_OVERRIDE_GFX_VERSION:-10.3.0}"
PYTORCH_HIP_ALLOC_CONF: "expandable_segments:True"
STOCKGNN_ENV: docker
devices:
- /dev/kfd
- /dev/dri
group_add:
- video
- render
security_opt:
- seccomp:unconfined
shm_size: "16gb"
networks:
- stockgnn-net
profiles:
- train
# -----------------------------------------------------------------------------
# Shared network
# -----------------------------------------------------------------------------
networks:
stockgnn-net:
driver: bridge
+29 -9
View File
@@ -673,8 +673,8 @@ class StockDataPipeline:
"""Return the next date that has price data for at least one ticker."""
candidate = date + timedelta(days=1)
for _ in range(7):
for ticker in self.price_data.values():
if not ticker.empty and candidate in ticker.index:
for price_df in self.price_data.values():
if not price_df.empty and candidate in price_df.index:
return candidate
candidate += timedelta(days=1)
return None
@@ -742,13 +742,19 @@ class StockDataPipeline:
# Build graph edges
edge_index, edge_weight = self._build_edges(node_tickers, date, pit_cache)
# Build targets: next trading day's return for each stock
# Build targets: next trading day's return for each stock.
# Skip dates where no next trading date exists (e.g. tail end of dataset)
# to avoid polluting training labels with fabricated zero returns.
next_date = self._next_trading_date(date)
if next_date is None:
continue
y_vals: List[float] = []
for ticker in node_tickers:
valid_seqs_for_y: List[np.ndarray] = []
valid_tickers_for_y: List[str] = []
for i, ticker in enumerate(node_tickers):
if (
next_date is not None
and ticker in self.price_data
ticker in self.price_data
and next_date in self.price_data[ticker].index
and date in self.price_data[ticker].index
):
@@ -757,9 +763,23 @@ class StockDataPipeline:
/ self.price_data[ticker].loc[date]["Close"]
- 1
)
else:
ret = 0.0
y_vals.append(ret)
y_vals.append(ret)
valid_seqs_for_y.append(node_sequences[i])
valid_tickers_for_y.append(ticker)
if not y_vals:
continue
if len(valid_tickers_for_y) < len(node_tickers):
# Rebuild x and edges using only stocks with valid targets
x = torch.tensor(np.array(valid_seqs_for_y), dtype=torch.float32)
mean = x.mean(dim=0, keepdim=True)
std = x.std(dim=0, keepdim=True).clamp(min=1e-8)
x = (x - mean) / std
edge_index, edge_weight = self._build_edges(
valid_tickers_for_y, date, pit_cache
)
node_tickers = valid_tickers_for_y
y = torch.tensor(y_vals, dtype=torch.float32).unsqueeze(1)
+2 -2
View File
@@ -33,6 +33,6 @@ class Broker(ABC):
pass
@abstractmethod
def get_account_summary(self) -> Dict:
"""Get account summary."""
def get_account_summary(self, prices: Optional[Dict[str, float]] = None) -> Dict:
"""Get account summary. prices: optional mark-to-market prices keyed by ticker."""
pass
+1 -1
View File
@@ -41,7 +41,7 @@ class InteractiveBrokersBroker(Broker):
# Placeholder: implement IB positions retrieval
return {}
def get_account_summary(self) -> Dict:
def get_account_summary(self, prices: Optional[Dict[str, float]] = None) -> Dict:
"""Get account summary from Interactive Brokers."""
# Placeholder: implement IB account summary retrieval
return {}
+8 -3
View File
@@ -21,6 +21,7 @@ class PaperTradingBroker(Broker):
self.position_prices = {} # fill price per ticker for valuation
self.orders = {}
self.transaction_cost = config.TRANSACTION_COST
self.slippage_rate = config.SLIPPAGE_RATE
def submit_order(self, order: Dict) -> Optional[str]:
"""Submit a simulated order."""
@@ -31,22 +32,26 @@ class PaperTradingBroker(Broker):
quantity = order["quantity"]
price = order["price"]
action = order["action"]
cost = quantity * price * (1 + self.transaction_cost)
if action == "buy":
# Slippage pushes fill price up on buys
fill_price = price * (1 + self.slippage_rate)
cost = quantity * fill_price * (1 + self.transaction_cost)
if cost > self.cash:
logger.warning(f"Insufficient cash for buy order: {order_id}")
self.orders[order_id]["status"] = "rejected"
return None
self.cash -= cost
self.positions[ticker] = self.positions.get(ticker, 0) + quantity
self.position_prices[ticker] = price
self.position_prices[ticker] = fill_price
elif action == "sell":
if self.positions.get(ticker, 0) < quantity:
logger.warning(f"Insufficient shares for sell order: {order_id}")
self.orders[order_id]["status"] = "rejected"
return None
self.cash += quantity * price * (1 - self.transaction_cost)
# Slippage pushes fill price down on sells
fill_price = price * (1 - self.slippage_rate)
self.cash += quantity * fill_price * (1 - self.transaction_cost)
self.positions[ticker] -= quantity
if self.positions[ticker] == 0:
del self.positions[ticker]
+6
View File
@@ -23,6 +23,12 @@ def _get_state() -> AppState:
return _state
@router.get("/health")
async def health_check() -> Dict:
"""Health check endpoint for Docker / load-balancers."""
return {"status": "ok", "project": config.PROJECT_NAME, "version": config.VERSION}
@router.get("/metrics")
async def get_metrics() -> Dict:
"""Get current dashboard metrics (memory, account, model)."""
+15 -2
View File
@@ -96,6 +96,17 @@ async def cancel_order(order_id: str) -> Dict:
return {"status": "cancelled" if success else "failed"}
def _last_known_price(state, ticker: str) -> float:
"""Return the most recent daily close from the pipeline, or 0 if unavailable."""
pipeline = getattr(state, "pipeline", None)
if pipeline is None:
return 0.0
price_df = getattr(pipeline, "price_data", {}).get(ticker)
if price_df is not None and not price_df.empty:
return float(price_df.iloc[-1]["Close"])
return 0.0
@router.post("/positions/close/{ticker}")
async def close_position(ticker: str) -> Dict:
"""Close a position for a given ticker."""
@@ -104,12 +115,13 @@ async def close_position(ticker: str) -> Dict:
if not broker or ticker not in broker.get_positions():
return {"status": "error", "message": "No position found"}
price = _last_known_price(state, ticker)
quantity = broker.get_positions()[ticker]
order = {
"ticker": ticker,
"action": "sell",
"quantity": quantity,
"price": 0,
"price": price,
"timestamp": datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
"type": "market",
}
@@ -130,11 +142,12 @@ async def close_all_positions() -> Dict:
positions = broker.get_positions()
results = []
for ticker in list(positions.keys()):
price = _last_known_price(state, ticker)
order = {
"ticker": ticker,
"action": "sell",
"quantity": positions[ticker],
"price": 0,
"price": price,
"timestamp": datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
"type": "market",
}
+1 -4
View File
@@ -7,6 +7,7 @@ import json
import logging
import os
from contextlib import asynccontextmanager
from datetime import datetime
from typing import Any, Dict, List
from fastapi import FastAPI, Request, WebSocket, WebSocketDisconnect
@@ -140,8 +141,6 @@ async def _handle_ws_message(ws: WebSocket, data: Dict[str, Any]):
async def _broadcast_loop():
"""Background task that periodically broadcasts metrics to all WS clients."""
from datetime import datetime
while True:
try:
await asyncio.sleep(2) # broadcast every 2 seconds
@@ -163,8 +162,6 @@ async def _broadcast_loop():
def _build_broadcast_payload() -> Dict[str, Any]:
"""Build the real-time metrics payload."""
from datetime import datetime
mm = app_state.memory_manager
broker = app_state.broker