Enable GPU acceleration for ML pipeline

Author: nikolay-e

Dockerfile

@@ -23,6 +23,7 @@ RUN pip install --no-cache-dir --upgrade pip && \
 
 # --- ML Builder Stage ---
 FROM builder AS ml-builder
+RUN pip install --no-cache-dir torch --index-url https://download.pytorch.org/whl/cu121
 RUN pip install --no-cache-dir '.[ml]'
 
 # --- Bot Builder Stage ---

src/ml/config.py

@@ -1,3 +1,4 @@
+import os
 from dataclasses import dataclass, field
 from pathlib import Path
 
@@ -6,6 +7,7 @@
 class MLConfig:
     model_dir: Path = Path("/app/models")
     chronos_base_model: str = "amazon/chronos-bolt-base"
+    device: str = field(default_factory=lambda: os.environ.get("DEVICE", "cpu"))
     forecast_horizons: list[int] = field(default_factory=lambda: [1, 7, 14])
     metrics: list[str] = field(
         default_factory=lambda: ["weight", "hrv", "rhr", "sleep_total", "steps"]

src/ml/predict.py

@@ -21,12 +21,12 @@
 
 def _extract_quantiles(forecast: torch.Tensor, pipeline, horizon_idx: int):
     if isinstance(pipeline, ChronosBoltPipeline):
-        quantiles = forecast.numpy()[0]
+        quantiles = forecast.cpu().numpy()[0]
         p10 = float(quantiles[BOLT_P10_IDX, horizon_idx])
         p50 = float(quantiles[BOLT_P50_IDX, horizon_idx])
         p90 = float(quantiles[BOLT_P90_IDX, horizon_idx])
     else:
-        samples = forecast.numpy()[0]
+        samples = forecast.cpu().numpy()[0]
         p10 = float(np.percentile(samples[:, horizon_idx], 10))
         p50 = float(np.percentile(samples[:, horizon_idx], 50))
         p90 = float(np.percentile(samples[:, horizon_idx], 90))

src/ml/run.py

@@ -42,7 +42,7 @@ def run_pipeline(user_id: int, do_train: bool = False) -> None:
             chronos_pipeline = load_chronos(config)
             save_chronos(chronos_pipeline, chronos_path)
         else:
-            chronos_pipeline = load_chronos_from_disk(chronos_path)
+            chronos_pipeline = load_chronos_from_disk(chronos_path, config)
 
         n_forecasts = generate_forecasts(
             chronos_pipeline,

src/ml/train.py

@@ -14,10 +14,12 @@
 def load_chronos(config: MLConfig) -> BaseChronosPipeline:
     pipeline = BaseChronosPipeline.from_pretrained(
         config.chronos_base_model,
-        device_map="cpu",
+        device_map=config.device,
         torch_dtype=torch.float32,
     )
-    logger.info("chronos_model_loaded", model=config.chronos_base_model)
+    logger.info(
+        "chronos_model_loaded", model=config.chronos_base_model, device=config.device
+    )
     return pipeline
 
 
@@ -27,13 +29,13 @@ def save_chronos(pipeline: BaseChronosPipeline, path: Path) -> None:
     logger.info("chronos_saved", path=str(path))
 
 
-def load_chronos_from_disk(path: Path) -> BaseChronosPipeline:
+def load_chronos_from_disk(path: Path, config: MLConfig) -> BaseChronosPipeline:
     pipeline = BaseChronosPipeline.from_pretrained(
         str(path),
-        device_map="cpu",
+        device_map=config.device,
         torch_dtype=torch.float32,
     )
-    logger.info("chronos_loaded_from_disk", path=str(path))
+    logger.info("chronos_loaded_from_disk", path=str(path), device=config.device)
     return pipeline