dask / dask/distributed

[RFC] Worker groups and environments

Open
#495 6 comments 0 reactions 0 assignees View on GitHub
enhancement
Dominant language
Python
Stars
1.7k
Forks
778
Avg merge
2h 50m
Merged PRs (30d)
3

Description

This is a concrete proposal for solving #85. For preliminary discussion, see that issue.
### What is an environment:

An environment is defined as an object of type `Environment`, which currently has the following design:

``` python
class Environment(object):
def isinstance(self):
"""Return if worker it's run on is an instance of the environment"""
pass

def setup(self):
"""Setup the environment"""
pass

def teardown(self):
"""Teardown the environment"""
pass
```

The reason for making it a class is grouping the methods together, providing a consistent place to store state, and making the user facing signature consistent (even if more optional methods are added to the class). Using the examples from the original issue, example `Environment` classes might be:

``` python
class MemoryGreaterThan(Environment):
def __init__(self, threshold=30e9):
self.threshold = threshold

def isinstance(self):
import psutil
return psutil.virtual_memory.total > self.threshold

class DataBaseAccess(Environment):
def __init__(self, credentials):
self.credentials = credentials

def isinstance(self):
return can_connect_to_database(self.credentials)

def setup(self):
self.conn = connect_to_database(self.credentials)

def teardown(self):
self.conn.disconnect()
```
### How environments are registered

Environments are registered by passing _instances_ of the object to `Executor.register_environment`, along with a name. The decision to use instances instead of classes was made to allow environments to be parametrized (if needed), and contain internal state. An example of registering environments is:

``` python
# possibly use keywords instead (e.g. `env_name=env_object`)
exc.register_environment('medium', MemoryGreaterThan(16e9))
exc.register_environment('large', MemoryGreaterThan(32e9))
exc.register_environment('database', DataBaseAccess('mypassword'))
```
### Changes to the scheduler/worker state and transitions
#### When an environment is registered
1. After being registered, the environment is serialized and sent to the scheduler, where it is stored in a mapping `{name: environment_object}` called `environments`.
2. The scheduler sends the environment to each worker. There the `isinstance` method is ran to check if the worker is an instance of the environment. If so, the worker stores it in a mapping of `{name: environment_object}` called environments. The (optional) `setup` method on the environment is then run, which can initialize any additional needed state.
3. The worker replies to the scheduler, letting it know if the environment was added or not. If it was added, the scheduler updates a mapping of `{environment_name: {worker_ids}}` to reflect this. Call this `environment_workers`.
#### When a worker is added
1. The scheduler iterates through all of the current environments, and does the dance described above. There might be merit in batching these together (sending a list instead of many instances), unknown.
#### When a worker is removed
1. The worker runs the (optional) teardown method to cleanup any necessary state. Note that this won't happen if the worker fails in a hard manner (segfault, etc...).
2. The scheduler removes the worker from `environment_workers` as needed
#### When a task is submitted
1. The `workers` keyword to `submit` will be modified to take in environment names as well. This will be sent along in addition to `restrictions` and `loose_restrictions` to the scheduler's `update_graph` method as an `environments` keyword.
2. `update_graph` will update a mapping of `{key: {environment_names}}` called `environment_restrictions` on the state. The reason for not converting this immediately into `restrictions` is that the workers in the environment may change between task submission time and task run time (new workers may register for example).
#### When selecting tasks to run
1. `ensure_occupied` will be modified to also take into account `environment_restrictions` when picking tasks to run and where to run them.
### Summary of new state

**New worker state**
- `environments`: a mapping of `{environment_name: environment_object}`

**New scheduler state**
- `environments`: a mapping of `{environment_name: environment_object}`
- `environment_workers`: a mapping of `{environment_name: {workers}}`
- `environment_restrictions`: a mapping of `{key: {environment_names}}`
### Accessing environments from inside a task

Tasks can call the `distributed.get_environment('name')` function to get the environment object from inside the worker. The implementation of this might look like:

``` python
def get_environment(name):
worker = get_worker_this_is_running_on()
return worker.environments[name]
```

This allows tasks to access whatever state the environment may contain internally. For example, a user may do the following to pull data from a stored database connection:

``` python
def get_data(query):
env = distributed.get_environment('database')
return env.conn.run_query(query)

exc.submit(get_data, query, workers='database')
```

Contributor guide

Open the contributing guide

Research direction

Start by reading the preliminary discussion in issue #85 and the scheduler/worker entry points named here: register_environment, update_graph, and ensure_occupied. Compare the proposed Environment lifecycle and environment_workers and environment_restrictions state with the current design. Done requires an agreed design and implementation plan for registration, worker matching, task selection, and teardown.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
distributed-systems
Issue type
Feature
Difficulty
5/5
Estimated time
Over a week
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
25/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.