Feature/create refactored surrogate training

.cursor/rules/hatchet-docs.mdc

@@ -0,0 +1,12 @@
+---
+description: Hatchet documentation MCP server
+alwaysApply: true
+---
+
+When working with Hatchet (task queues, workflows, durable execution), use the
+Hatchet MCP docs server for accurate, up-to-date API reference and examples.
+
+MCP server URL: https://docs.hatchet.run/api/mcp
+
+Use the search_docs tool to find relevant documentation pages, or get_full_docs
+for comprehensive context. Documentation covers Python, TypeScript, and Go SDKs.

.env.scythe.training

@@ -0,0 +1,9 @@
+SCYTHE_WORKER_SLOTS=1
+SCYTHE_WORKER_DOES_FAN=False
+SCYTHE_WORKER_DOES_LEAF=True
+SCYTHE_WORKER_HAS_GPU=True
+
+SCYTHE_TIMEOUT_EXPERIMENT_SCHEDULE=2h
+SCYTHE_TIMEOUT_EXPERIMENT_EXECUTION=1h
+SCYTHE_TIMEOUT_SCATTER_GATHER_SCHEDULE=10h
+SCYTHE_TIMEOUT_SCATTER_GATHER_EXECUTION=10h

.gitignore

@@ -217,3 +217,5 @@ inputs/
 
 .env.local.hatchet
 .env.local.host.hatchet
+
+scratch/

Makefile

@@ -71,6 +71,10 @@ simulations-native: ## Run the simulations
 fanouts-native: ## Run the fanouts
 	@uv run --env-file .env.$(AWS_ENV).aws --env-file .env.$(HATCHET_ENV).hatchet --env-file .env.scythe.storage --env-file .env.scythe.fanouts worker
 
+.PHONY: training-native
+training-native: ## Run the training
+	@uv run --env-file .env.$(AWS_ENV).aws --env-file .env.$(HATCHET_ENV).hatchet --env-file .env.scythe.storage --env-file .env.scythe.training worker
+
 .PHONY: viz-native
 viz-native: ## Run the visualization tool # TODO: possibly add env vars to the command
 	@uv run streamlit run src/globi/tools/visualization/main.py

docker-compose.yml

@@ -20,6 +20,31 @@ services:
       - ./inputs:/code/inputs
       - ./outputs:/code/outputs
       - ./tests/data/e2e:/code/tests/data/e2e
+  training:
+    image: ${AWS_ACCOUNT_ID:-123456789012}.dkr.ecr.${AWS_REGION:-us-east-1}.amazonaws.com/hatchet/globi:${IMAGE_TAG:-latest}
+    build:
+      context: .
+      dockerfile: src/globi/worker/Dockerfile
+      args:
+        EP_VERSION: ${EP_VERSION:-25.2.0}
+        PYTHON_VERSION: ${PYTHON_VERSION:-3.12}
+    env_file:
+      - .env
+      - .env.${AWS_ENV:-local}.aws
+      - .env.${HATCHET_ENV:-local}.hatchet
+      - .env.scythe.storage
+      - .env.scythe.training
+    deploy:
+      mode: replicated
+      replicas: ${TRAINING_REPLICAS:-0}
+      resources:
+        reservations:
+          devices:
+            - capabilities: [gpu] # Requests access to all GPUs
+    volumes:
+      - ./inputs:/code/inputs
+      - ./outputs:/code/outputs
+      - ./tests/data/e2e:/code/tests/data/e2e
   fanouts:
     image: ${AWS_ACCOUNT_ID:-123456789012}.dkr.ecr.${AWS_REGION:-us-east-1}.amazonaws.com/hatchet/globi:${IMAGE_TAG:-latest}
     build:

pyproject.toml

@@ -46,6 +46,25 @@ visualization = [
     "playwright>=1.40.0",
 ]
 
+# ml = [
+#     "lightgbm>=4.6.0",
+#     "xgboost>=3.2.0",
+#     "pytorch-tabular>=1.2.0",
+#     "torch>=2.5.0",
+#     "tensorboard>=2.20.0",
+#     "wandb>=0.25.0",
+# ]
+
+ml-gpu = [
+    "lightgbm>=4.6.0",
+    "xgboost>=3.2.0",
+    "numba>=0.63.1",
+    "pytorch-tabular>=1.2.0",
+    "torch>=2.5.0",
+    "tensorboard>=2.20.0",
+    "wandb>=0.25.0",
+]
+
 cli = [
     "click>=8.1.7",
     "xlsxwriter>=3.2.9",
@@ -78,9 +97,29 @@ docs = [
 worker = "globi.worker.main:main"
 globi = "globi.tools.cli.main:cli"
 
+[[tool.uv.index]]
+name = "pytorch-cu128"
+url = "https://download.pytorch.org/whl/cu128"
+explicit = true
+
+[[tool.uv.index]]
+name = "pytorch-cpu"
+url = "https://download.pytorch.org/whl/cpu"
+explicit = true
+
+[[tool.uv.index]]
+name = "pypi"
+url = "https://pypi.org/simple"
+explicit = true
+
 [tool.uv.sources]
+# PyTorch: CUDA 12.8 on Linux/Windows (where builds exist), PyPI (CPU) on macOS
+torch = [
+  { index = "pytorch-cu128", marker = "sys_platform != 'darwin'", extra = "ml-gpu" },
+  { index = "pypi", marker = "sys_platform == 'darwin'", extra = "ml-gpu" },
+]
 # scythe-engine = {git = "https://github.com/szvsw/scythe", branch = "feature/allow-optional-filerefs"}
-# scythe-engine = {git = "https://github.com/szvsw/scythe", branch = "feature/update-hatchet"}
+scythe-engine = {git = "https://github.com/szvsw/scythe", branch = "feature/allow-versioning-workflows"}
 # scythe-engine = {path = "../scythe", editable = true}
 # epinterface = {path = "../epinterface", editable = true}
 # epinterface = {path = "epinterface", editable = true}

src/globi/allocate.py

@@ -102,7 +102,7 @@ def allocate_globi_experiment(
         raise ValueError(msg)
 
     experiment = BaseExperiment[ExperimentInputSpec, ExperimentOutputSpec](
-        experiment=simulate_globi_building, run_name=name
+        runnable=simulate_globi_building, run_name=name
     )
     print(f"Submitting {len(buildings_gdf)} buildings for experiment {name}")
     min_branches_required, _, _ = calculate_branching_factor(specs)
@@ -182,7 +182,7 @@ def allocate_globi_dryrun(
         raise ValueError(msg)
 
     experiment = BaseExperiment[ExperimentInputSpec, ExperimentOutputSpec](
-        experiment=simulate_globi_building,
+        runnable=simulate_globi_building,
         run_name=f"{config.name}/dryrun/{config.scenario}",
     )
 

src/globi/models/surrogate/__init__.py

@@ -0,0 +1 @@
+"""Models used for the surrogate pipeline."""

src/globi/models/surrogate/configs/__init__.py

@@ -0,0 +1 @@
+"""Configs for the surrogate model pipeline."""