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.

.dockerignore

@@ -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

.env.example

@@ -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

Dockerfile

@@ -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"]

benchmark.py

@@ -106,11 +106,11 @@ def benchmark_model():
 
     for batch_size in batch_sizes:
         for seq_len in sequence_lengths:
-            # Create data for this configuration
+            # Create data for this configuration; batch_size drives num nodes
-            bx = torch.randn(num_stocks, seq_len, num_features).to(config.DEVICE)
+            bx = torch.randn(batch_size, seq_len, num_features).to(config.DEVICE)
-            bei = torch.randint(0, num_stocks, (2, num_edges)).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
+        bench_model = nn.Sequential(
-        model = IntradayGNN(num_features, config.SEQUENCE_LENGTH)
-        model.feature_processor = 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)
+        estimated_memory = memory_manager.estimate_model_memory(bench_model)
-
-        # Optimize model
-        model = amd_optimizer.optimize_model(model)
-
-        # Estimate memory usage
-        estimated_memory = memory_manager.estimate_model_memory(model)
         logger.info(
             f"Hidden Channels: {hidden_channels}, Estimated Memory: {estimated_memory / 1024**3:.2f}GB"
         )
-
+        del bench_model
-        # Clean up
-        del model
         memory_manager.empty_cache()
 
     # Final memory stats

docker-compose.yml

@@ -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

src/data/pipeline.py

@@ -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():
+            for price_df in self.price_data.values():
-                if not ticker.empty and candidate in ticker.index:
+                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
+                    ticker in self.price_data
-                    and 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:
+                    y_vals.append(ret)
-                    ret = 0.0
+                    valid_seqs_for_y.append(node_sequences[i])
-                y_vals.append(ret)
+                    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)
 

src/trading/broker.py

@@ -33,6 +33,6 @@ class Broker(ABC):
         pass
 
     @abstractmethod
-    def get_account_summary(self) -> Dict:
+    def get_account_summary(self, prices: Optional[Dict[str, float]] = None) -> Dict:
-        """Get account summary."""
+        """Get account summary. prices: optional mark-to-market prices keyed by ticker."""
         pass

src/trading/ib_broker.py

@@ -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 {}

src/trading/paper_broker.py

@@ -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]

src/web/api/dashboard.py

@@ -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)."""

src/web/api/trading_endpoints.py

@@ -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",
         }

src/web/app.py

@@ -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