Assign specific worker URLs to specific tasks

[gabotechs]

Today, worker URL assignation is random, and users have no say over which URL is assigned to each worker.

How it works today

The process since there's a valid single-node plan until there's a distributed plan with worker URLs assigned is the following:

1. Users tell the distributed planner what is an appropriate amount of tasks to execute each node (typically only done for leaf nodes)

          ┌──────────────────┐          
          │CoalescePartitions│          
          └──────────────────┘          
          ┌──────────────────┐          
          │ Aggregate(final) │          
          └──────────────────┘          
          ┌──────────────────┐          
          │   Repartition    │          
          └──────────────────┘          
          ┌──────────────────┐          
          │Aggregate(partial)│          
          └──────────────────┘          
          ┌──────────────────┐          
          │      Filter      │          
          └──────────────────┘          
          ┌──────────────────┐          
          │      Union       │          
          └──────────────────┘          
┌──────────────────┐┌──────────────────┐
│   DataSource1    ││   DataSource2    │
└──────────────────┘└──────────────────┘
    TaskEstimator:       TaskEstimator: 
    I want 2 tasks       I want 3 tasks 

2. The distributed planner performs some task count reconciliation in each stage, as there might be multiple nodes contributing an estimated task count for the same stage:

                                            
                   ...                      
┌ ─ ─ ─ ─ ─                     ─Same stage 
           ┌──────────────────┐            │
│          │   Repartition    │ 3 tasks     
           └──────────────────┘            │
│          ┌──────────────────┐             
           │Aggregate(partial)│ 3 tasks    │
│          └──────────────────┘             
           ┌──────────────────┐            │
│          │      Filter      │ 3 tasks     
           └──────────────────┘            │
│          ┌──────────────────┐             
           │      Union       │ 3 tasks    │
│          └──────────────────┘             
 ┌──────────────────┐┌──────────────────┐  │
││   DataSource1    ││   DataSource2    │   
 └──────────────────┘└──────────────────┘  │
│       3 tasks            3 tasks          
 ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ┘

3. Network boundaries with specific input task counts are injected, but with empty slots for worker URLs

                               ┏━━━━━━━━━━━━━━━━━━┓                                   
                               ┃ DistributedExec  ┃                                   
                               ┗━━━━━━━━━━━━━━━━━━┛                                   
                               ┌──────────────────┐                                   
                               │CoalescePartitions│                                   
                               └──────────────────┘                                   
                               ┏━━━━━━━━━━━━━━━━━━┓                                   
                               ┃ NetworkCoalesce  ┃                                   
                               ┗━━━━━━━━━━━━━━━━━━┛                                   
                                                                                      
                   ┌ ─ ─ ─ ─ ─ ─URL: None  ┌ ─ ─ ─ ─ ─ URL: None                      
                    ┌──────────────────┐│   ┌──────────────────┐│                     
                   ││ Aggregate(final) │   ││ Aggregate(final) │                      
                    └──────────────────┘│   └──────────────────┘│                     
                   │┏━━━━━━━━━━━━━━━━━━┓   │┏━━━━━━━━━━━━━━━━━━┓                      
                    ┃  NetworkShuffle  ┃│   ┃  NetworkShuffle  ┃│                     
                   │┗━━━━━━━━━━━━━━━━━━┛   │┗━━━━━━━━━━━━━━━━━━┛                      
                    ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ┘   ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ┘                     
                                                                                      
┌ ─ ─ ─ ─ ─ ─ ─ ─ URL: None  ┌ ─ ─ ─ ─ ─ ─ ─ ─ URL: None  ┌ ─ ─ ─ ─ ─ ─ ─ ─ URL: None 
    ┌──────────────────┐   │     ┌──────────────────┐   │     ┌──────────────────┐   │
│   │   Repartition    │     │   │   Repartition    │     │   │   Repartition    │    
    └──────────────────┘   │     └──────────────────┘   │     └──────────────────┘   │
│   ┌──────────────────┐     │   ┌──────────────────┐     │   ┌──────────────────┐    
    │Aggregate(partial)│   │     │Aggregate(partial)│   │     │Aggregate(partial)│   │
│   └──────────────────┘     │   └──────────────────┘     │   └──────────────────┘    
    ┌──────────────────┐   │     ┌──────────────────┐   │     ┌──────────────────┐   │
│   │      Filter      │     │   │      Filter      │     │   │      Filter      │    
    └──────────────────┘   │     └──────────────────┘   │     └──────────────────┘   │
│   ┌──────────────────┐     │   ┌──────────────────┐     │   ┌──────────────────┐    
    │      Union       │   │     │      Union       │   │     │      Union       │   │
│   └──────────────────┘     │   └──────────────────┘     │   └──────────────────┘    
 ┌───────────┐┌───────────┐│  ┌───────────┐┌───────────┐│  ┌───────────┐┌───────────┐│
││ModDataSrc1││ModDataSrc2│  ││ModDataSrc1││ModDataSrc2│  ││ModDataSrc1││ModDataSrc2│ 
 └───────────┘└───────────┘│  └───────────┘└───────────┘│  └───────────┘└───────────┘│
└ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─  └ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─  └ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ 

4. The DistributedExec head node lazily modifies the nodes below, assigning worker URLs to the different tasks right before execution:

                               ┏━━━━━━━━━━━━━━━━━━┓                                   
                               ┃ DistributedExec  ┃                                   
                               ┗━━━━━━━━━━━━━━━━━━┛                                   
                               ┌──────────────────┐                                   
                               │CoalescePartitions│                                   
                               └──────────────────┘                                   
                               ┏━━━━━━━━━━━━━━━━━━┓                                   
                               ┃ NetworkCoalesce  ┃                                   
                               ┗━━━━━━━━━━━━━━━━━━┛                                   
                                                                                      
                   ┌ ─ ─ ─ URL: http://4   ┌ ─ ─ ─ URL: http://5                      
                    ┌──────────────────┐│   ┌──────────────────┐│                     
                   ││ Aggregate(final) │   ││ Aggregate(final) │                      
                    └──────────────────┘│   └──────────────────┘│                     
                   │┏━━━━━━━━━━━━━━━━━━┓   │┏━━━━━━━━━━━━━━━━━━┓                      
                    ┃  NetworkShuffle  ┃│   ┃  NetworkShuffle  ┃│                     
                   │┗━━━━━━━━━━━━━━━━━━┛   │┗━━━━━━━━━━━━━━━━━━┛                      
                    ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ┘   ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ┘                     
                                                                                      
┌ ─ ─ ─ ─ ─ ─ URL: http://1  ┌ ─ ─ ─ ─ ─ ─ URL: http://2  ┌ ─ ─ ─ ─ ─ ─ URL: http://3 
    ┌──────────────────┐   │     ┌──────────────────┐   │     ┌──────────────────┐   │
│   │   Repartition    │     │   │   Repartition    │     │   │   Repartition    │    
    └──────────────────┘   │     └──────────────────┘   │     └──────────────────┘   │
│   ┌──────────────────┐     │   ┌──────────────────┐     │   ┌──────────────────┐    
    │Aggregate(partial)│   │     │Aggregate(partial)│   │     │Aggregate(partial)│   │
│   └──────────────────┘     │   └──────────────────┘     │   └──────────────────┘    
    ┌──────────────────┐   │     ┌──────────────────┐   │     ┌──────────────────┐   │
│   │      Filter      │     │   │      Filter      │     │   │      Filter      │    
    └──────────────────┘   │     └──────────────────┘   │     └──────────────────┘   │
│   ┌──────────────────┐     │   ┌──────────────────┐     │   ┌──────────────────┐    
    │      Union       │   │     │      Union       │   │     │      Union       │   │
│   └──────────────────┘     │   └──────────────────┘     │   └──────────────────┘    
 ┌───────────┐┌───────────┐│  ┌───────────┐┌───────────┐│  ┌───────────┐┌───────────┐│
││ModDataSrc1││ModDataSrc2│  ││ModDataSrc1││ModDataSrc2│  ││ModDataSrc1││ModDataSrc2│ 
 └───────────┘└───────────┘│  └───────────┘└───────────┘│  └───────────┘└───────────┘│
└ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─  └ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─  └ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ 

How it could work in the future

That last step is random, but at the moment it happens, there's a lot of information available for the users to decide what's an appropriate worker URL for executing that task:

• Task index and task count (what the project calls DistributedTaskContext)
• The subplan that is assigned to that task
• The overall query plan

This should be sufficient information for users to assign themselves a specific worker URL to each task.

Note that #374 is also proposing an extension to the API for allowing providing more things rather than just a number (the task count estimation), so any future evolution of this API will need to take into account the two efforts:

• this in this issue
• the one in Stream metadata from DistributedExec to leaf nodes #374

[sandugood]

I know that in the overall architecture of the project we do not have the "master" coordinator, that coordinates tasks between the workers (like in YARN). But isn't this issue somehow related to distributing the tasks based on the load factor of workers?

In the production env of a large-scale system I imagine that it would be nice to assign workers based on their load (current) and the potential load that would be by performing the query.

[jayshrivastava]

Concretely, I think these changes need to be made

• The WorkerResolver returns a list of URLs. We should allow the implementor to return worker IDs.
• The DistributedExt should have a with_task_router<T: TaskRouter>() where TaskRouter is a trait you implement with a method map_task_to_worker_with_id(task) -> Option<worker_id>. The user can implement consistent routing or load balancing if they want.

If the user doesn't have an opinion, then we can fallback to random?

Here's the problem though - what if a user wants to route consistently using files / partitions / file groups?

I wonder if we can extend this API to solve this problem:

pub trait TaskEstimator {
    fn scale_up_leaf_node(
        &self,
        plan: &Arc<dyn ExecutionPlan>,
        task_count: usize,
        cfg: &ConfigOptions,
    ) -> Option<Arc<dyn ExecutionPlan>>;
}

It would be nice if this also got a task_idx and was called for each leaf task, so I could edit each leaf node of each leaf task. It would also be nice if TaskEstimator was stateful (instantiated once per query) to be able to move files / file groups around.

[kurtvolmar]

I'm super excited to have the ability to route tasks to specific workers. My team has a use case where we would like specific partitions in a data source to be executed on specific workers, since this optimizes our caching on workers. This would be a consistent hash of partition data to workers.

Building on top of this idea, it would be helpful for users to also have a clear way to express the mapping of partition to task. An idea would be putting this in the TaskEstimator. The TaskEstimation provides the task_count, but it could additionally return partition_groups, which would be a grouping of partitions which should be executed within the same task.

For example, users may want to map partitions to specific workers, even if it produces an "uneven" distribution of partitions to tasks:

Partition	Task	Worker URL
0	0	worker-a.com
1	1	worker-b.com
2	0	worker-a.com
3	0	worker-a.com
4	0	worker-a.com

[gabotechs]

The WorkerResolver returns a list of URLs. We should allow the implementor to return worker IDs.

URLs act as unique identifiers to worker nodes (no two worker nodes can have the same URL). If nodes need to be uniquely identified, what do you think the limitations are to just identifying them as the URL?

The DistributedExt should have a with_task_router<T: TaskRouter>() where TaskRouter is a trait you implement with a method map_task_to_worker_with_id(task) -> Option<worker_id>. The user can implement consistent routing or load balancing if they want.

I'm hoping that we can ship this with a much simpler implementation. The TaskEstimator sounds like a good place to provide this extension. For example:

• Right now, once a task count has been decided, the planner prompts the node to "scale up", returning a modified version of the node account for it being executed in T tasks:

    /// After a final task_count is decided, taking into account all the leaf nodes in the [Stage],
    /// this allows performing a transformation in the leaf nodes for accounting for the fact that
    /// they are going to run in multiple tasks.
    fn scale_up_leaf_node(
        &self,
        plan: &Arc<dyn ExecutionPlan>,
        task_count: usize,
        cfg: &ConfigOptions,
    ) -> Option<Arc<dyn ExecutionPlan>>;

• In the future, just returning the list of worker URLs in which the different tasks should be scheduled sounds like a simple enough thing that should do the trick:

    /// After a final task_count is decided, taking into account all the leaf nodes in the [Stage],
    /// this allows performing a transformation in the leaf nodes for accounting for the fact that
    /// they are going to run in multiple tasks. Additionally, a list of URLs with length equal to
    /// `task_count` can be provided, and the different tasks will be located in those URLs
    fn scale_up_leaf_node(
        &self,
        plan: &Arc<dyn ExecutionPlan>,
        task_count: usize,
        cfg: &ConfigOptions,
    ) -> Option<{ plan: Arc<dyn ExecutionPlan>, urls: Vec<Url> }>;

The nice thing about something like this is that it does not need any new concept, the existing ones are enough for allowing people to choose URLs for their tasks.

[JSOD11]

I like the task router idea suggested by Jay and have been working on an implementation in #397. Consistent routing will be valuable to us for similar reasons to what you described @kurtvolmar.

Putting the urls in the return value of scale_up_leaf_node may work, but it also potentially collapses scaling and routing/scheduling into one interface in a way that might be unnecessary or confusing for users.

Alternatively, it is straightforward for us to provide them with a dedicated routing interface: The task router creates a clear separation of concerns between scaling/routing and a simple interface for a user to define their routing. This way, scale_up_leaf_node is clearly scoped to scaling up the leaf nodes while task_router.route is clearly scoped to routing tasks to URLs.

In the PR I linked, the idea is that rather than returning:

plan: Arc<dyn ExecutionPlan>, urls: Vec<Url>

We instead return the plan and a struct containing the affinity keys, maybe something like this, which will be more general:

pub struct ScaledLeafNode {
    pub plan: Arc<dyn ExecutionPlan>,
    pub affinity_keys: Option<Vec<Option<Vec<u8>>>>,
}

Users can set the affinity key bytes here as they wish (we choose bytes to leave this interface very unopinionated regarding what users may want to route on, e.g. a combination of strings and numbers) and then routing is handled according to their custom defined routing function later closer to execution.

This pattern makes matching tasks to URLs simple in prepare_plan where the tasks are actually matched to URLs:

            let tasks = stage
                .tasks
                .iter()
                .enumerate()
                .map(|(i, task)| {
                    let url = task_router.route(RouterInfo {
                        task_number: i,
                        urls: urls.clone(),
                        stage_seed: start_idx,
                        affinity_key: task.affinity_key.clone(),
                    });

[gabotechs]

We instead return the plan and a struct containing the affinity keys, maybe something like this, which will be more general

The issue with the affinity keys mechanism is that it biases towards the "cache affinity" use case, leaving behind other valid routing use cases like assigning tasks to the "more available worker" based on load information.

For this reason, letting the users define the URLs themselves in which the tasks should be executed accounts for all possible routing policies, without biasing towards the "affinity" use case.

Putting the urls in the return value of scale_up_leaf_node may work, but it also potentially collapses scaling and routing/scheduling into one interface in a way that might be unnecessary or confusing for users.

I'd say that scaling + routing is not necessarily a different problem. For example:

• A user is about to "scale up" a node
• In order to know the amount of tasks to "scale up", they resolve in which machines their data is scattered
• The result of the resolution is a vector that says vec!["http://machine1", "http://machine2", "http://machine3"]
• Then they select 3 tasks
• Upon scaling up, they already know there the number of machines there, they should not need a different abstraction for wiring up those tasks

The task router creates a clear separation of concerns between scaling/routing and a simple interface for a user to define their routing. This way, scale_up_leaf_node is clearly scoped to scaling up the leaf nodes while task_router.route is clearly scoped to routing tasks to URLs.

My bet is that we do not want impose that separation of concerns ourselves in this project, whether they are different concerns or not is not really up to us here. I think we should start from the simplest thing, see if we face some limitations, and just if we see those limitations go to the complex solution involving a new abstraction.

[JSOD11]

Yes, that's a fair assessment. I've opened #409 moving towards the approach you presented, and I also see this as a good opportunity to resolve #331 since both modify the task estimator interface.