Update app.py

app.py

@@ -191,14 +191,88 @@ def fuse_stdout(fragments):
 
 
 # ─── Fragmentation ─────────────────────────────────────────────────────────
+# Add this to your head node fragmentation logic
+
+def analyze_gpu_requirements(code):
+    """
+    Detect if code needs GPU by looking for:
+    1. @gpu_required decorator
+    2. torch.cuda usage
+    3. device="cuda" patterns
+    """
+    
+    # Check for explicit decorator
+    if "@gpu_required" in code or "@requires_gpu" in code:
+        return True
+    
+    # Check for torch GPU patterns
+    gpu_patterns = [
+        "torch.cuda",
+        'device="cuda"',
+        "device='cuda'",
+        ".cuda()",
+        "torch.device('cuda')",
+        'torch.device("cuda")',
+    ]
+    
+    return any(pattern in code for pattern in gpu_patterns)
+
+
+def is_parallelizable_safely(code, loop_info):
+    """
+    Determine if a loop can be safely parallelized.
+    Returns: (can_parallelize: bool, reason: str)
+    """
+    
+    # Don't parallelize GPU training loops
+    gpu_training_keywords = [
+        "model.train()",
+        "optimizer.step()",
+        "loss.backward()",
+        ".backward()",
+        "torch.nn",
+        "nn.Module",
+    ]
+    
+    if any(kw in code for kw in gpu_training_keywords):
+        return False, "GPU training loop (stateful)"
+    
+    # Don't parallelize if it has shared state
+    stateful_keywords = [
+        "global ",
+        "nonlocal ",
+        ".append(",  # modifying shared list
+        "self.",     # class methods
+    ]
+    
+    if any(kw in code for kw in stateful_keywords):
+        return False, "Stateful operations detected"
+    
+    # Safe to parallelize: map-reduce style loops
+    safe_patterns = [
+        "for i in range",
+        "for idx in range",
+        "for chunk in range",
+    ]
+    
+    return True, "Independent iterations"
+
+
 def fragment_code(code, cfg, task_id, libs):
+    """
+    IMPROVED fragmentation with GPU awareness
+    """
     frags = []
     cpus = max(1, cfg.get("cpus", 1))
-    needs_gpu = cfg.get("gpus", 0) > 0
+    
+    # Detect GPU requirements
+    needs_gpu = cfg.get("gpus", 0) > 0 or analyze_gpu_requirements(code)
     ram = cfg.get("ram_gb", 4)
-    target = max(2, min(cpus * 2, 8))
-
+    
+    # Find loops
     loop = None
+    loop_safe = False
+    
     try:
         tree = ast.parse(code)
         for node in ast.walk(tree):
@@ -207,62 +281,102 @@ def fragment_code(code, cfg, task_id, libs):
                 if isinstance(fn, ast.Name) and fn.id == "range":
                     args = node.iter.args
                     try:
-                        if len(args) == 1: loop = {"s": 0, "e": ast.literal_eval(args[0])}
-                        elif len(args) >= 2: loop = {"s": ast.literal_eval(args[0]), "e": ast.literal_eval(args[1])}
+                        if len(args) == 1:
+                            loop = {"s": 0, "e": ast.literal_eval(args[0])}
+                        elif len(args) >= 2:
+                            loop = {"s": ast.literal_eval(args[0]), "e": ast.literal_eval(args[1])}
+                        
                         if loop:
                             loop["var"] = node.target.id if isinstance(node.target, ast.Name) else "i"
+                            
+                            # Check if safe to parallelize
+                            can_parallelize, reason = is_parallelizable_safely(code, loop)
+                            if can_parallelize:
+                                loop_safe = True
+                                print(f"[HEAD] Loop is safe to parallelize: {reason}")
+                            else:
+                                print(f"[HEAD] Loop NOT safe to parallelize: {reason}")
+                                loop = None
                             break
-                    except: loop = None
-    except: pass
-
-    if loop and (loop["e"] - loop["s"]) >= 20:
+                    except:
+                        loop = None
+    except:
+        pass
+    
+    # Only fragment if safe AND big enough
+    if loop and loop_safe and (loop["e"] - loop["s"]) >= 50:  # higher threshold
         total = loop["e"] - loop["s"]
-        n = min(target, total)
-        chunk = math.ceil(total / n)
-        patterns = ([f"range({loop['e']})"] if loop["s"] == 0 else []) + [
-            f"range({loop['s']}, {loop['e']})", f"range({loop['s']},{loop['e']})"]
-        if loop["s"] == 0: patterns.append(f"range(0, {loop['e']})")
-
+        n = min(cpus * 2, total // 10)  # at least 10 iterations per fragment
+        chunk = max(10, total // n)
+        
+        print(f"[HEAD] Fragmenting loop: {total} iterations → {n} fragments ({chunk} each)")
+        
+        # Build pattern matching
+        patterns = []
+        if loop["s"] == 0:
+            patterns.append(f"range({loop['e']})")
+            patterns.append(f"range(0, {loop['e']})")
+        patterns.append(f"range({loop['s']}, {loop['e']})")
+        patterns.append(f"range({loop['s']},{loop['e']})")
+        
         for idx in range(n):
             cs = loop["s"] + idx * chunk
             ce = min(loop["s"] + (idx + 1) * chunk, loop["e"])
-            if cs >= loop["e"]: break
+            if cs >= loop["e"]:
+                break
+            
             mod = code
             for pat in patterns:
                 if pat in mod:
                     mod = mod.replace(pat, f"range({cs}, {ce})", 1)
                     break
-
-            # Inject SACCP metadata
+            
+            # Inject metadata
             header = f"""# ═══ SACCP Fragment {idx}/{n} ═══
 __saccp_rank__ = {idx}
 __saccp_world_size__ = {n}
 __saccp_chunk__ = ({cs}, {ce})
+__saccp_is_fragment__ = True
 """
+            
             frags.append({
-                "fragment_id": f"{task_id}_frag_{idx}", "task_id": task_id,
-                "fragment_index": idx, "fragment_type": "compute",
+                "fragment_id": f"{task_id}_frag_{idx}",
+                "task_id": task_id,
+                "fragment_index": idx,
+                "fragment_type": "compute",
                 "code": header + mod,
                 "input_data": json.dumps({"rank": idx, "world_size": n, "chunk": [cs, ce]}),
-                "required_libs": libs, "required_gpu": needs_gpu,
-                "min_ram_gb": max(1, ram // n), "timeout_seconds": 600,
+                "required_libs": libs,
+                "required_gpu": needs_gpu,
+                "min_ram_gb": max(1, ram // n),
+                "timeout_seconds": 600,
             })
-
+    
+    # Fallback: single fragment
     if not frags:
+        reason = "single" if not loop else ("unsafe" if not loop_safe else "too small")
+        print(f"[HEAD] No fragmentation ({reason}) — single fragment")
+        
         header = """# ═══ SACCP Single Fragment ═══
 __saccp_rank__ = 0
 __saccp_world_size__ = 1
+__saccp_is_fragment__ = False
 """
+        
         frags.append({
-            "fragment_id": f"{task_id}_frag_0", "task_id": task_id,
-            "fragment_index": 0, "fragment_type": "compute",
+            "fragment_id": f"{task_id}_frag_0",
+            "task_id": task_id,
+            "fragment_index": 0,
+            "fragment_type": "compute",
             "code": header + code,
             "input_data": json.dumps({"rank": 0, "world_size": 1}),
-            "required_libs": libs, "required_gpu": needs_gpu,
-            "min_ram_gb": ram, "timeout_seconds": 600,
+            "required_libs": libs,
+            "required_gpu": needs_gpu,
+            "min_ram_gb": ram,
+            "timeout_seconds": 600,
         })
-
-    print(f"[HEAD] Task {task_id[:8]} → {len(frags)} fragments")
+    
+    print(f"[HEAD] Task {task_id[:8]} → {len(frags)} fragments (GPU required: {needs_gpu})")
     return frags