test: add fail_flow and split_in_branch UX tests for windmill QA findings

- fail_flow.py: flow that raises RuntimeError mid-step; used to verify
  schedulers report FAILED status (not RUNNING) when a step crashes
  before reaching end (catches A03-1 in metaflow-windmill QA)
- split_in_branch_flow.py: outer split with a nested inner split inside
  branch_a; verifies the outer join is compiled correctly (catches A02-2
  where _find_join_step returned inner_join instead of outer_join)
- test_basic.py: test_fail_flow_reports_failed_status and
  test_split_in_branch_deployer, both marked scheduler_only + deployer

Author: Nissan Pow

test/ux/core/flows/basic/fail_flow.py

@@ -0,0 +1,25 @@
+"""Flow that deliberately fails mid-step.
+
+Used to verify that schedulers report FAILED status (not RUNNING or PENDING)
+when a flow step raises an unhandled exception.
+"""
+
+from metaflow import FlowSpec, step
+
+
+class FailFlow(FlowSpec):
+    @step
+    def start(self):
+        self.next(self.failing_step)
+
+    @step
+    def failing_step(self):
+        raise RuntimeError("Deliberate failure for testing scheduler status reporting.")
+
+    @step
+    def end(self):
+        pass
+
+
+if __name__ == "__main__":
+    FailFlow()

test/ux/core/flows/basic/split_in_branch_flow.py

@@ -0,0 +1,56 @@
+"""Flow with a nested split inside a branch.
+
+Graph:
+  start → [branch_a, branch_b] → outer_join → end
+  branch_a → [inner_x, inner_y] → inner_join → outer_join
+
+Used to verify that schedulers correctly compile and execute complex graph
+topologies where a split exists inside one branch of an outer split.
+Without correct join detection, outer_join and end are silently dropped.
+"""
+
+from metaflow import FlowSpec, step
+
+
+class SplitInBranchFlow(FlowSpec):
+    @step
+    def start(self):
+        self.next(self.branch_a, self.branch_b)
+
+    @step
+    def branch_a(self):
+        self.label = "a"
+        self.next(self.inner_x, self.inner_y)
+
+    @step
+    def inner_x(self):
+        self.sub_label = "x"
+        self.next(self.inner_join)
+
+    @step
+    def inner_y(self):
+        self.sub_label = "y"
+        self.next(self.inner_join)
+
+    @step
+    def inner_join(self, inputs):
+        self.sub_labels = sorted(i.sub_label for i in inputs)
+        self.next(self.outer_join)
+
+    @step
+    def branch_b(self):
+        self.label = "b"
+        self.next(self.outer_join)
+
+    @step
+    def outer_join(self, inputs):
+        self.labels = sorted(i.label for i in inputs)
+        self.next(self.end)
+
+    @step
+    def end(self):
+        pass
+
+
+if __name__ == "__main__":
+    SplitInBranchFlow()

test/ux/core/test_basic.py

@@ -201,6 +201,85 @@ def test_resources_cpu(exec_mode, decospecs, compute_env, tag, scheduler_config)
     ), "Message didn't match"
 
 
+@pytest.mark.scheduler_only
+@pytest.mark.deployer
+def test_fail_flow_reports_failed_status(
+    exec_mode, decospecs, compute_env, tag, scheduler_config
+):
+    """Verify schedulers report FAILED (not RUNNING/PENDING) when a step raises.
+
+    Catches A03-1: _check_sysroot_completion returns RUNNING forever for flows
+    that crash before reaching the end step, because end/ dir never appears.
+    """
+    import time
+    from .test_utils import deploy_flow_to_scheduler
+
+    scheduler_type = scheduler_config.scheduler_type
+    if scheduler_type is None:
+        pytest.skip("No scheduler configured — requires a scheduler_type")
+
+    deployed_flow = deploy_flow_to_scheduler(
+        flow_name="basic/fail_flow.py",
+        tl_args={"decospecs": decospecs, "env": compute_env},
+        scheduler_args={"cluster": scheduler_config.cluster},
+        deploy_args={"tags": tag, **(scheduler_config.deploy_args or {})},
+        scheduler_type=scheduler_type,
+    )
+
+    triggered = deployed_flow.trigger()
+
+    deadline = time.time() + 120
+    final_status = None
+    while time.time() < deadline:
+        s = triggered.status
+        if s in ("FAILED", "SUCCEEDED"):
+            final_status = s
+            break
+        time.sleep(5)
+
+    assert final_status == "FAILED", (
+        "A flow that raises RuntimeError mid-step should report FAILED, got %r"
+        % final_status
+    )
+
+
+@pytest.mark.scheduler_only
+@pytest.mark.deployer
+def test_split_in_branch_deployer(
+    exec_mode, decospecs, compute_env, tag, scheduler_config
+):
+    """Verify a split nested inside a branch compiles and executes correctly.
+
+    Catches A02-2: _find_join_step's while loop follows only out_funcs[0],
+    causing it to return the inner join instead of the outer join. Without the
+    fix, outer_join and end are silently dropped from the compiled flow.
+    """
+    from .test_utils import deploy_flow_to_scheduler, wait_for_deployed_run
+
+    scheduler_type = scheduler_config.scheduler_type
+    if scheduler_type is None:
+        pytest.skip("No scheduler configured — requires a scheduler_type")
+
+    deployed_flow = deploy_flow_to_scheduler(
+        flow_name="basic/split_in_branch_flow.py",
+        tl_args={"decospecs": decospecs, "env": compute_env},
+        scheduler_args={"cluster": scheduler_config.cluster},
+        deploy_args={"tags": tag, **(scheduler_config.deploy_args or {})},
+        scheduler_type=scheduler_type,
+    )
+
+    run = wait_for_deployed_run(deployed_flow)
+    assert run.successful, "SplitInBranchFlow should complete successfully"
+    assert sorted(run["outer_join"].task.data.labels) == [
+        "a",
+        "b",
+    ], "outer_join should receive results from both branch_a and branch_b"
+    assert sorted(run["inner_join"].task.data.sub_labels) == [
+        "x",
+        "y",
+    ], "inner_join should receive results from inner_x and inner_y"
+
+
 @pytest.mark.scheduler_only
 @pytest.mark.skip(reason="devstack has no GPU nodes")
 def test_resources_gpu(exec_mode, decospecs, compute_env, tag, scheduler_config):