Skip to content

Add logic to broadcast current scheduler info to workers in CLI #46

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 8 commits into from
Dec 19, 2019
Merged

Add logic to broadcast current scheduler info to workers in CLI #46

Changes from all commits
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
600d696
add scheduler bcast logic from core into cli
Dec 14, 2019
4111d77
add scheduler option
Dec 14, 2019
ae4bbe9
forgot import sys, move mpi comm assign inside main
Dec 14, 2019
0d95cba
try to get dask-mpi to run in background to prevent unit test hangs
Dec 16, 2019
9381dae
general cleanup suggested by kmpaul
Dec 19, 2019
205f32e
changes from black and flake8
Dec 19, 2019
c9ca1d5
solve this issue with barriers instead of bcast
Dec 19, 2019
da75a6d
changes from black and flake8
Dec 19, 2019
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
23 changes: 14 additions & 9 deletions dask_mpi/cli.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,12 +1,11 @@
import click

import asyncio
from mpi4py import MPI

from dask.distributed import Scheduler, Worker, Nanny
from distributed.cli.utils import check_python_3

comm = MPI.COMM_WORLD
rank = comm.Get_rank()
from mpi4py import MPI


@click.command()
Expand Down Expand Up @@ -72,34 +71,40 @@ def main(
protocol,
):

comm = MPI.COMM_WORLD
rank = comm.Get_rank()

if rank == 0 and scheduler:

async def run():
async def run_scheduler():
async with Scheduler(
interface=interface,
protocol=protocol,
scheduler_file=scheduler_file,
dashboard_address=dashboard_address,
port=scheduler_port,
scheduler_file=scheduler_file,
) as s:
comm.Barrier()
await s.finished()

asyncio.get_event_loop().run_until_complete(run_scheduler())

else:
comm.Barrier()

async def run():
async def run_worker():
WorkerType = Nanny if nanny else Worker
async with WorkerType(
scheduler_file=scheduler_file,
interface=interface,
protocol=protocol,
nthreads=nthreads,
memory_limit=memory_limit,
local_directory=local_directory,
name=rank,
scheduler_file=scheduler_file,
) as worker:
await worker.finished()

asyncio.get_event_loop().run_until_complete(run())
asyncio.get_event_loop().run_until_complete(run_worker())


def go():
Expand Down