From dcebfe718aaff304fb655d04316fbfbb3db50735 Mon Sep 17 00:00:00 2001
From: Roland Tannous <rolandtannous@gravityq.ai>
Date: Sat, 7 Mar 2026 17:53:36 +0000
Subject: [PATCH] fix: prevent training hang on Windows by adding
 triton-windows support

---
 setup.ps1                               |  5 +++++
 studio/backend/core/export/worker.py    | 12 ++++++++++++
 studio/backend/core/inference/worker.py | 12 ++++++++++++
 studio/backend/core/training/trainer.py |  6 ++++++
 studio/backend/core/training/worker.py  | 12 ++++++++++++
 5 files changed, 47 insertions(+)

diff --git a/setup.ps1 b/setup.ps1
index d2bfed3a4c..591310c249 100644
--- a/setup.ps1
+++ b/setup.ps1
@@ -685,6 +685,11 @@ $CuTag = Get-PytorchCudaTag
 Write-Host "   Installing PyTorch with CUDA support ($CuTag)..." -ForegroundColor Cyan
 pip install torch torchvision torchaudio --index-url "https://download.pytorch.org/whl/$CuTag" 2>&1 | Out-Null
 
+# Install Triton for Windows (enables torch.compile — without it training can hang)
+Write-Host "   Installing Triton for Windows..." -ForegroundColor Cyan
+pip install "triton-windows<3.7" 2>&1 | Out-Null
+Write-Host "[OK] Triton for Windows installed (enables torch.compile)" -ForegroundColor Green
+
 # Ordered heavy dependency installation — shared cross-platform script
 Write-Host "   Running ordered dependency installation..." -ForegroundColor Cyan
 python "$PSScriptRoot\install_python_stack.py"
diff --git a/studio/backend/core/export/worker.py b/studio/backend/core/export/worker.py
index adc197c335..9e7e72e9dd 100644
--- a/studio/backend/core/export/worker.py
+++ b/studio/backend/core/export/worker.py
@@ -228,6 +228,18 @@ def run_export_process(
         })
         return
 
+    # ── 1b. On Windows, check Triton availability (must be before import torch) ──
+    if sys.platform == "win32":
+        try:
+            import triton  # noqa: F401
+            logger.info("Triton available — torch.compile enabled")
+        except ImportError:
+            os.environ["TORCHDYNAMO_DISABLE"] = "1"
+            logger.warning(
+                "Triton not found on Windows — torch.compile disabled. "
+                'Install for better performance: pip install "triton-windows<3.7"'
+            )
+
     # ── 2. Import ML libraries (fresh in this clean process) ──
     try:
         _send_response(resp_queue, {
diff --git a/studio/backend/core/inference/worker.py b/studio/backend/core/inference/worker.py
index 04a93b8f2a..7683d5e314 100644
--- a/studio/backend/core/inference/worker.py
+++ b/studio/backend/core/inference/worker.py
@@ -323,6 +323,18 @@ def run_inference_process(
         })
         return
 
+    # ── 1b. On Windows, check Triton availability (must be before import torch) ──
+    if sys.platform == "win32":
+        try:
+            import triton  # noqa: F401
+            logger.info("Triton available — torch.compile enabled")
+        except ImportError:
+            os.environ["TORCHDYNAMO_DISABLE"] = "1"
+            logger.warning(
+                "Triton not found on Windows — torch.compile disabled. "
+                'Install for better performance: pip install "triton-windows<3.7"'
+            )
+
     # ── 2. Import ML libraries (fresh in this clean process) ──
     try:
         _send_response(resp_queue, {
diff --git a/studio/backend/core/training/trainer.py b/studio/backend/core/training/trainer.py
index 2dd2e6eeeb..c2ce009879 100644
--- a/studio/backend/core/training/trainer.py
+++ b/studio/backend/core/training/trainer.py
@@ -3,6 +3,7 @@ Unsloth Training Backend
 Integrates Unsloth training capabilities with the FastAPI backend
 """
 import os
+import sys
 # Prevent tokenizer parallelism deadlocks when datasets uses multiprocessing fork
 os.environ["TOKENIZERS_PARALLELISM"] = "false"
 
@@ -752,6 +753,11 @@ class UnslothTrainer:
                 "include_num_input_tokens_seen": True,  # Enable token counting
                 "dataset_num_proc": safe_num_proc(max(1, os.cpu_count() // 4)),
             }
+
+            # On Windows, disable DataLoader multiprocessing to avoid
+            # issues with modified sys.path in spawned subprocesses.
+            if sys.platform == "win32":
+                config_args["dataloader_num_workers"] = 0
             
             # Add warmup parameter - use warmup_ratio if provided, otherwise warmup_steps
             if warmup_ratio_val is not None:
diff --git a/studio/backend/core/training/worker.py b/studio/backend/core/training/worker.py
index a0ecb8d4e3..e4b833cb9a 100644
--- a/studio/backend/core/training/worker.py
+++ b/studio/backend/core/training/worker.py
@@ -97,6 +97,18 @@ def run_training_process(
         })
         return
 
+    # ── 1b. On Windows, check Triton availability (must be before import torch) ──
+    if sys.platform == "win32":
+        try:
+            import triton  # noqa: F401
+            logger.info("Triton available — torch.compile enabled")
+        except ImportError:
+            os.environ["TORCHDYNAMO_DISABLE"] = "1"
+            logger.warning(
+                "Triton not found on Windows — torch.compile disabled. "
+                'Install for better performance: pip install "triton-windows<3.7"'
+            )
+
     # ── 2. Now import ML libraries (fresh in this clean process) ──
     try:
         _send_status(event_queue, "Importing ML libraries...")