feat: configurable batch mode

src/Terrabuild.Configuration.Tests/Project.fs

@@ -54,19 +54,22 @@ let parseProject() =
               TargetBlock.Build = None
               TargetBlock.Outputs = None
               TargetBlock.Cache = None
+              TargetBlock.Batch = Expr.Enum "partition" |> Some
               TargetBlock.Steps = [ { Extension = "@dotnet"; Command = "build"; Parameters = Map.empty } ] }
         let targetDist =
             { TargetBlock.DependsOn = None
               TargetBlock.Build = None
               TargetBlock.Outputs = None
               TargetBlock.Cache = None
+              TargetBlock.Batch = None
               TargetBlock.Steps = [ { Extension = "@dotnet"; Command = "build"; Parameters = Map.empty }
                                     { Extension = "@dotnet"; Command = "publish"; Parameters = Map.empty } ] }
         let targetDocker =
             { TargetBlock.DependsOn = None
               TargetBlock.Build = "auto" |> Expr.Enum |> Some
               TargetBlock.Outputs = None
               TargetBlock.Cache = "remote" |> Expr.Enum |> Some
+              TargetBlock.Batch = None
               TargetBlock.Steps = [ { Extension = "@shell"; Command = "echo"
                                       Parameters = Map [ "arguments", Expr.Function (Function.Trim,
                                                                                      [ Expr.Function (Function.Plus,
@@ -125,6 +128,7 @@ let parseProject2() =
                                                                  Expr.String ".dll" ])] |> Some
               TargetBlock.DependsOn = None
               TargetBlock.Cache = None
+              TargetBlock.Batch = None
               TargetBlock.Steps = [ { Extension = "@dotnet"; Command = "build"; Parameters = Map.empty } ] }
 
         let locals = 

src/Terrabuild.Configuration.Tests/TestFiles/Success_PROJECT

@@ -36,6 +36,7 @@ project id {
 
 target build {
   depends_on = [ target.dist ]
+  batch = ~partition
   @dotnet build { }
 }
 

src/Terrabuild.Configuration.Tests/TestFiles/Success_WORKSPACE

@@ -18,6 +18,7 @@ target dist {
   depends_on = [ target.build ]
   build = ~auto
   artifacts = ~none
+  batch = ~none
 }
 
 target dummy { }

src/Terrabuild.Configuration.Tests/Workspace.fs

@@ -14,15 +14,18 @@ let parseWorkspace() =
         let targetBuild = 
             { TargetBlock.DependsOn = Set [ "install"; "^build" ] |> Some
               TargetBlock.Build = None
-              TargetBlock.Cache = None }
+              TargetBlock.Cache = None
+              TargetBlock.Batch = None }
         let targetDist =
             { TargetBlock.DependsOn = Set [ "build" ] |> Some
               TargetBlock.Build = "auto" |> Expr.Enum |> Some
-              TargetBlock.Cache = "none" |> Expr.Enum |> Some }
+              TargetBlock.Cache = "none" |> Expr.Enum |> Some
+              TargetBlock.Batch = "none" |> Expr.Enum |> Some }
         let targetDummy =
             { TargetBlock.DependsOn = None
               TargetBlock.Build = None
-              TargetBlock.Cache = None }
+              TargetBlock.Cache = None
+              TargetBlock.Batch = None }
 
         let extDotnet =
             { Image = Some (Expr.String "mcr.microsoft.com/dotnet/sdk:8.0.101")
@@ -77,15 +80,18 @@ let parseWorkspace2() =
         let targetBuild = 
             { TargetBlock.DependsOn = Set [ "^build" ] |> Some
               TargetBlock.Build = None
-              TargetBlock.Cache = None }
+              TargetBlock.Cache = None
+              TargetBlock.Batch = None }
         let targetDist =
             { TargetBlock.DependsOn = Set [ "build" ] |> Some
               TargetBlock.Build = None
-              TargetBlock.Cache = None }
+              TargetBlock.Cache = None
+              TargetBlock.Batch = None }
         let targetDummy =
             { TargetBlock.DependsOn = None
               TargetBlock.Build = None
-              TargetBlock.Cache = None }
+              TargetBlock.Cache = None
+              TargetBlock.Batch = None }
 
         let extDotnet =
             { Image = Expr.String "mcr.microsoft.com/dotnet/sdk:8.0.101" |> Some

src/Terrabuild.Configuration/AST/Project.fs

@@ -27,6 +27,7 @@ type TargetBlock =
       DependsOn: Set<string> option
       Build: Expr option
       Cache: Expr option
+      Batch: Expr option
       Steps: Step list }
 
 [<RequireQualifiedAccess>]

src/Terrabuild.Configuration/AST/Workspace.fs

@@ -17,7 +17,8 @@ type WorkspaceBlock =
 type TargetBlock =
     { DependsOn: Set<string> option
       Build: Expr option
-      Cache: Expr option }
+      Cache: Expr option
+      Batch: Expr option }
 
 [<RequireQualifiedAccess>]
 type WorkspaceFile =

src/Terrabuild.Configuration/Transpiler/Project.fs

@@ -75,7 +75,7 @@ let toProject (block: Block) =
 
 let toTarget (block: Block) =
     block
-    |> checkAllowedAttributes ["outputs"; "depends_on"; "build"; "artifacts"]
+    |> checkAllowedAttributes ["outputs"; "depends_on"; "build"; "artifacts"; "batch"]
     |> ignore
 
     let outputs = block |> tryFindAttribute "outputs"
@@ -89,6 +89,7 @@ let toTarget (block: Block) =
                 | _ -> raiseInvalidArg $"Invalid target dependency '{dependency}'"))
     let build = block |> tryFindAttribute "build"
     let cache = block |> tryFindAttribute "artifacts"
+    let batch = block |> tryFindAttribute "batch"
     let steps =
         block.Blocks
         |> List.map (fun step ->
@@ -114,6 +115,7 @@ let toTarget (block: Block) =
       TargetBlock.DependsOn = dependsOn
       TargetBlock.Build = build
       TargetBlock.Cache = cache
+      TargetBlock.Batch = batch
       TargetBlock.Steps = steps }
 
 

src/Terrabuild.Configuration/Transpiler/Workspace.fs

@@ -65,7 +65,7 @@ let toWorkspace (block: Block) =
 
 let toTarget (block: Block) =
     block
-    |> checkAllowedAttributes ["depends_on"; "build"; "artifacts"]
+    |> checkAllowedAttributes ["depends_on"; "build"; "artifacts"; "batch"]
     |> checkNoNestedBlocks
     |> ignore
 
@@ -79,9 +79,11 @@ let toTarget (block: Block) =
                 | _ -> raiseInvalidArg $"Invalid target dependency '{dependency}'"))
     let build = block |> tryFindAttribute "build"
     let cache = block |> tryFindAttribute "artifacts"
+    let batch = block |> tryFindAttribute "batch"
     { TargetBlock.DependsOn = dependsOn
       TargetBlock.Build = build
-      TargetBlock.Cache = cache }
+      TargetBlock.Cache = cache
+      TargetBlock.Batch = batch }
 
 
 let toVariable (block: Block) =

src/Terrabuild.Tests/Core/GraphPipeline/Batch.fs

@@ -5,7 +5,7 @@ open GraphDef
 open GraphPipeline.Batch
 
 [<Test>]
-let ``check batch computation``() =
+let ``check partition computation``() =
     let buildNode id clusterHash action deps =
         { Node.Id = id
           Node.ProjectId = id
@@ -20,7 +20,8 @@ let ``check batch computation``() =
           Node.Operations = []
           Node.Artifacts = Artifacts.Workspace
           Node.Action = action
-          Node.Build = Build.Auto }
+          Node.Build = Build.Auto
+          Node.Batch = Group.Partition }
 
     let addNode (node: Node) nodes = nodes |> Map.add node.Id node
 
@@ -75,3 +76,72 @@ let ``check batch computation``() =
         expected |> List.map (fun b -> b.BatchId, b.ClusterHash, (b.Nodes |> List.map (fun n -> n.Id) |> Set.ofList))
                  |> Set.ofList
     )
+
+
+
+[<Test>]
+let ``check none/partition computation``() =
+    let buildNode id clusterHash action deps group =
+        { Node.Id = id
+          Node.ProjectId = id
+          Node.ProjectName = None
+          Node.ProjectDir = $"/src/project{id}"
+          Node.Target = "build"
+          Node.Dependencies = deps
+          Node.Outputs = Set.empty
+          Node.ProjectHash = ""
+          Node.TargetHash = ""
+          Node.ClusterHash = clusterHash
+          Node.Operations = []
+          Node.Artifacts = Artifacts.Workspace
+          Node.Action = action
+          Node.Build = Build.Auto
+          Node.Batch = group }
+
+    let addNode (node: Node) nodes = nodes |> Map.add node.Id node
+
+    // Bucket hash-A: connected via A1 -> A2 (in-bucket edge)
+    let nodeA1 = buildNode "A1" (Some "hash-A") NodeAction.Build (Set ["A2"; "B1"]) Group.Partition
+    let nodeA2 = buildNode "A2" (Some "hash-A") NodeAction.Restore Set.empty Group.Partition
+
+    // Bucket hash-B: connected via B1 -> B2 (in-bucket edge)
+    let nodeB1 = buildNode "B1" (Some "hash-B") NodeAction.Build (Set ["B2"]) Group.None
+    let nodeB2 = buildNode "B2" (Some "hash-B") NodeAction.Build Set.empty Group.None
+    let nodeC1 = buildNode "C1" (Some "hash-B") NodeAction.Build (Set ["C2"]) Group.None
+    let nodeC2 = buildNode "C2" (Some "hash-B") NodeAction.Build Set.empty Group.None
+
+    // Not batchable
+    let nodeD1 = buildNode "D1" None NodeAction.Build Set.empty Group.Partition
+
+    let nodes =
+        Map.empty
+        |> addNode nodeA1 |> addNode nodeA2
+        |> addNode nodeB1 |> addNode nodeB2
+        |> addNode nodeC1 |> addNode nodeC2
+        |> addNode nodeD1
+
+    let graph =
+        { Graph.Nodes = nodes
+          Graph.RootNodes = Set [ "A1"; "B1"; "D1" ]
+          Graph.Batches = Map.empty }
+
+    let batches = computeBatches graph
+
+    let expectedBatchIdA = Hash.sha256strings ("hash-A" :: [ "A1"; "A2" ])
+    let expectedBatchIdB = Hash.sha256strings ("hash-B" :: [ "B1"; "B2"; "C1"; "C2" ])
+
+    let expected =
+        [ { BatchId = expectedBatchIdA
+            ClusterHash = "hash-A"
+            Nodes = [ nodeA1; nodeA2 ] }
+          { BatchId = expectedBatchIdB
+            ClusterHash = "hash-B"
+            Nodes = [ nodeB1; nodeB2; nodeC1; nodeC2 ] } ]
+
+    // Order is not guaranteed; compare as sets
+    batches |> List.map (fun b -> b.BatchId, b.ClusterHash, (b.Nodes |> List.map (fun n -> n.Id) |> Set.ofList))
+           |> Set.ofList
+    |> should equal (
+        expected |> List.map (fun b -> b.BatchId, b.ClusterHash, (b.Nodes |> List.map (fun n -> n.Id) |> Set.ofList))
+                 |> Set.ofList
+    )

src/Terrabuild/Core/Configuration.fs

@@ -30,7 +30,7 @@ type TargetOperation = {
 type Target = {
     Hash: string
     Build: Build option
-    Batch: bool
+    Batch: Group
     DependsOn: string set
     Outputs: string set
     Cache: Artifacts option
@@ -346,10 +346,12 @@ let private loadProjectDef (options: ConfigOptions.Options) (workspaceConfig: AS
                 let build = targetBlock.Build |> Option.orElseWith (fun () -> workspaceTarget |> Option.bind _.Build)
                 let dependsOn = targetBlock.DependsOn |> Option.orElseWith (fun () -> workspaceTarget |> Option.bind _.DependsOn)
                 let cache = targetBlock.Cache |> Option.orElseWith (fun () -> workspaceTarget |> Option.bind _.Cache)
+                let group = targetBlock.Batch |> Option.orElseWith (fun () -> workspaceTarget |> Option.bind _.Batch)
                 { targetBlock with 
                     Build = build
                     DependsOn = dependsOn
-                    Cache = cache })
+                    Cache = cache
+                    Batch = group })
         let environments =
             projectConfig.Project.Environments
             |> Option.bind (Eval.asStringSetOption << Eval.eval evaluationContext)
@@ -521,7 +523,7 @@ let private finalizeProject workspaceDir projectDir evaluationContext (projectDe
                     | Ok x -> raiseParseError $"Invalid build value '{x}'"
                     | Error error -> raiseParseError error
 
-            let targetBatch, targetOperations =
+            let canBatch, targetOperations =
                 target.Steps |> List.fold (fun (targetBatch, targetOperations) step ->
                     let extension = 
                         match projectDef.Extensions |> Map.tryFind step.Extension with
@@ -628,6 +630,19 @@ let private finalizeProject workspaceDir projectDir evaluationContext (projectDe
                 |> List.map (fun ope -> ope.Hash)
                 |> Hash.sha256strings
 
+            let targetBatch = 
+                let targetGroup =
+                    target.Batch
+                    |> Option.map (fun batch -> batch |> Eval.eval evaluationContext |> Eval.asEnum)
+                match targetGroup with
+                | Some group ->
+                    match group with
+                    | Ok "none" -> Group.None
+                    | Ok "partition" -> Group.Partition
+                    | Ok x -> raiseParseError $"Invalid group value '{x}'"
+                    | Error error -> raiseParseError error
+                | _ -> Group.None
+
             let target =
                 { Target.Hash = targetHash
                   Target.Build = targetBuild

src/Terrabuild/Core/GraphDef.fs

@@ -13,6 +13,11 @@ type ContaineredShellOperation = {
     ErrorLevel: int
 }
 
+[<RequireQualifiedAccess>]
+type Group =
+    | None
+    | Partition
+
 [<RequireQualifiedAccess>]
 type Artifacts =
     | None
@@ -53,6 +58,7 @@ type Node = {
     Operations: ContaineredShellOperation list
     Artifacts: Artifacts
     Build: Build
+    Batch: Group
 
     Action: NodeAction
 }

src/Terrabuild/Core/GraphPipeline/Batch.fs

@@ -84,8 +84,17 @@ let computeBatches (graph: Graph) =
         // if fewer than 2, no possible batch
         if bucketNodes.Length <= 1 then Seq.empty
         else
-            bucketNodes
+            let batchModes = 
+                bucketNodes
+                |> List.groupBy (fun node -> node.Batch)
+                |> Map.ofSeq
+
+            let noneGroup = batchModes |> Map.tryFind Group.None |> Option.defaultValue []  
+            let partitionGroups = batchModes |> Map.tryFind Group.Partition |> Option.defaultValue []  
+
+            partitionGroups
             |> partitionByDependencies
+            |> (fun partitionGroups -> noneGroup :: partitionGroups)
             |> Seq.choose (fun comp ->
                 // only batch if > 1 node and at least one member is actually executing
                 if comp.Length <= 1 then None
@@ -179,7 +188,8 @@ let private createBatchNodes (options: ConfigOptions.Options) (configuration: Co
                   GraphDef.Node.ProjectHash = batch.BatchId
                   GraphDef.Node.TargetHash = headNode.TargetHash
                   GraphDef.Node.Action = NodeAction.Build
-                  GraphDef.Node.Build = headNode.Build }
+                  GraphDef.Node.Build = headNode.Build
+                  GraphDef.Node.Batch = headNode.Batch }
 
             Some (batch.BatchId, batchNode)
     )

src/Terrabuild/Core/GraphPipeline/Node.fs

@@ -127,7 +127,7 @@ let build (options: ConfigOptions.Options) (configuration: Configuration.Workspa
                         | _ -> false
 
                     cacheability, batchable, ops @ newops
-                ) (Artifacts.Managed, targetConfig.Batch, [])
+                ) (Artifacts.Managed, true, [])
 
             let opsCmds = ops |> List.map Json.Serialize
 
@@ -155,15 +155,16 @@ let build (options: ConfigOptions.Options) (configuration: Configuration.Workspa
                 if cache = Artifacts.None then Set.empty
                 else targetConfig.Outputs
 
-            let batchContent = [
-                targetConfig.Hash
-                $"{buildAction}"
-            ]
-            let batchHash = batchContent |> Hash.sha256strings
-
             let targetClusterHash =
-                if targetConfig.Batch && batchable then Some batchHash
-                else None
+                if batchable then
+                    let batchContent = [
+                        targetConfig.Hash
+                        $"{buildAction}"
+                    ]
+                    let batchHash = batchContent |> Hash.sha256strings
+                    Some batchHash
+                else
+                    None
 
             let node =
                 { Node.Id = nodeId
@@ -175,6 +176,7 @@ let build (options: ConfigOptions.Options) (configuration: Configuration.Workspa
                   Node.Operations = ops
                   Node.Artifacts = cache
                   Node.Build = build
+                  Node.Batch = targetConfig.Batch
 
                   Node.Dependencies = children
                   Node.Outputs = targetOutput

tests/basic/results/terrabuild-debug.action-graph.json

@@ -24,6 +24,7 @@
       ],
       "artifacts": "none",
       "build": "always",
+      "batch": "none",
       "action": "build"
     }
   },