In anticipation for upcoming improvements to Flotilla, including streaming task execution (Flotilla 1.75) and moving off of Ray (Flotilla 2), we need to ensure that our intermediate data management and shuffle capabilities are on par with other engines and what Ray would provide us. This design doc consists of 3 parts:

• Adding task retries based on partition / data lineage
• Supporting shuffle backends for intermediate and result storage in a more generic fashion
• Making Flight Shuffle competitive with Ray-based shuffles
  • This will be WIP based on our investigations

Couple of potential dependencies

• Lineage will make improving shuffles easier, because we can safely store in memory
• Flotilla 1.75 will help with performance, because it’s easier to combine output partitions on the worker directly instead of going through the scheduler (pre / post-shuffle merge)
• I can imaging that some of the simplifications we make for generic shuffle backends will make lineage a bit cleaner, if not easier

Lineage Support

Right now, Flotilla can retry tasks, but it can’t handle when a task failed because an input is missing, such as when it was lost from a worker failure. Thus, it only works for single-stage / map-only pipelines with fixed inputs. To support generic task retries, we need to track how intermediate partitions are utilized, both as outputs of tasks and inputs of others.

State Management

Partitions should be globally tracked within the scheduler per query. For each partition, we should know its current status (alive or dead / gone), as well as how to recreate it if missing. Rough idea of the state is:

struct QueryState {
    ...
    completed_tasks: HashMap<TaskID, CompletedTask>,
    partition_context: HashMap<PartitionID, PartitionContext>,
}

struct CompletedTask {
    task: Arc<SwordfishTask>,
    results: Vec<PartitionID>,
}

// --------- These already exist, just for context ----------- //
struct SwordfishTask {
    task_context: TaskContext,  // Contains TaskID
    plan: LocalPhysicalPlanRef,
    inputs: HashMap<SourceId, Input>,
    ...
}

enum Input {
    ScanTasks(Vec<ScanTaskLikeRef>),
    GlobPaths(Vec<String>),
    InMemory(Vec<MicroPartitionRef>), // Change to PartitionID
		FlightShuffle(...) // Currently missing, will be added soon
}

// ----------------------------------------------------------- // 

struct PartitionContext {
    partition_id: PartitionID,
    created_task: Arc<SwordfishTask>,
    output_idx: usize,
    state: PartitionState,
}

enum PartitionState {
  Alive,
  Dead,
}

• Changes with Flotilla 1.75

Worker Handling

Right now, when an error occurs when running a task, the scheduler categorizes in the following categories.

pub(crate) enum TaskStatus {
    Success {
        result: MaterializedOutput,
        stats: ExecutionStats,
    },
    Failed {
        error: DaftError,
    },
    Cancelled,
    WorkerDied,
    WorkerUnavailable,
}

There are a couple of things we need to change: