Reintroduce AzureMLCluster

[lostmygithubaccount]

Feature request: AzureMLCluster

Background: long ago, AzureMLCluster was added as the second "cloud provider" in this repository. There were some problems described in #206 and issues. The implementation was deprecated after the introduction of AzureVMCluster.

Since, azureml-core has stabilized and now contains everything needed to implement AzureMLCluster.

It is trivial to use dask-mpi to startup the cluster on Azure ML, which can be seen XGBoost tutorial: https://github.com/Azure/azureml-examples/blob/main/tutorials/using-xgboost/src/run.py

Implementation

Some open questions:

• implementation: dask-mpi?
  • do we need to specify # of CPUs, etc.?
  • does this work with GPUs?
  • can this scale up/down?
• implementation: custom start dask.Scheduler and dask.Workers based on IP
  • prototype here: https://github.com/Azure/azureml-v2-preview/blob/main/examples/dask/src/startDask.py
  • uses PyTorch backend to assign/get IPs then startup processes

Outline of work needed

This part is less clear to me, now that I'm looking at it - I don't want to further delay opening this issue. But some things we should ensure are:

• only depends on azureml-core
• this is properly scoped and in line with other cloud providers, i.e. not starting up Jupyter
• scaling up/down works (old implementation had issues; concerns with dask-mpi proposal)
• this is tested
  • tests in this repo
  • we can also test large workloads/GPUs in https://github.com/Azure/azureml-examples
• works with azureml.core.Environment
• need existing aml compute target, or create on behalf of user?

Additional details

I cannot personally contribute or maintain this work. A few people at Microsoft have indicated a willingness to do so - I am opening this issue publicly, after which I will start the discussion internally.

This would not be an "officially supported" part of the Azure Machine Learning service - it would be an open source contribution, provided "as-is" and without official SLAs or support, to the Dask community.

• Cody

[jacobtomlinson]

Thanks @lostmygithubaccount. I would be very much in favour of having a new implementation for AzureML with better stability.

I know of some organisations that only give their staff access to AzureML and not the rest of Azure, so they have been left behind in the switch to AzureVMCluster.

We can definitely work through the problems you raised and find appropriate solutions. I expect some work needs to be done in dask-mpi and we would need to understand more about how AzureML works in order to get scaling working.

I do ask that whoever writes this new cluster manager will need to help maintain it going forwards. I appreciate that it cannot be officially supported and have anything like a formal SLA, but contributing open source code is like giving someone a puppy and we would need some help to look after it. Perhaps a day a month of someone's time to do bug fixes would be an appropriate minimum.

[DPeterK]

As I remember it, I was one of a few proponent voices for the original AzureMLCluster, so I'd very much be in support of a reimplementation of the functionality! Particularly if the new implementation is able to be more independent of the wider stack used to implement AzureML.

What I particularly liked about AzureMLCluster was that it was all self-contained within AzureML. This seemed to be a neat extension of the concept behind AzureML, in that researchers could be given access to AzureML and no other part of the Azure estate, and still have everything they need for the research they're engaged in. As @jacobtomlinson says, the AzureVMCluster breaks this paradigm somewhat.

I reckon I can justify some time contributing to developing a new version of AzureMLCluster - I think there will be value to this cluster manager to at least the Informatics Lab. Of course, ongoing maintenance may be a bit harder to fit in 😑 @jacobtomlinson / @lostmygithubaccount happy to discuss further!

[TomAugspurger]

Cody and Jacob, do you have a sense for whether an AzureMLCluster (v2) would be easier to build and maintain now that we have the Base VMCluster class in place? With the talk of using MPI (just for setup? not for comms I assume?) I'm not too hopeful that that's the case.

[lostmygithubaccount]

What's your concern? With using MPI?

[jacobtomlinson]

@TomAugspurger I'm not sure VMCluster is the best base class for AzureMLCluster. I think part of the problem is to do with the architecture of SpecCluster.

SpecCluster has a 1:1 mapping between scheduler/worker class instances and the real resources. It starts up the scheduler, connects the comm and then starts up workers. This works well for systems where resources are created independently of each other.

However on systems like Azure ML or other MPI based batch systems all of the resources are submitted as a single job. If you look at cluster managers like dask-yarn or tools like dask-mpi there is an assumption that the scheduler and a number of workers are created as a single unit. The first process in the job (MPI rank 0) starts up a scheduler, all other processes block until this is done and the comm info is shared, then workers are started.

There are some challenges with scaling on systems which work this way. It's not always possible to change the number of tasks in a job.

I've been discussing with @quasiben lately about creating an alternative base class in distributed which shares some code with SpecCluster but fits this paradigm more closely. I could imagine this would be useful for creating AzureMLCluster v2.

[TomAugspurger]

What's your concern? With using MPI?

Less of a concern, more of a lack of understanding about where MPI fits in the picture of forming the cluster.

Thanks for the background Jacob. A variant of SpecCluster that fits this model does seem like it's worth investigating.

[seanmcclure22]

Hello! Is this feature still being considered/worked-on? If not, do you have suggestions for setting up a dask cluster on an Azure ML Cluster?

[jacobtomlinson]

There hasn't been much activity here I'm afraid @seanmcclure22. As AzureML uses MPI you may want to look at dask-mpi.

[lostmygithubaccount]

hi @seanmcclure22! in addition to the dask-mpi suggestion, we do have a public example of how a Dask cluster can be setup on Azure ML compute here: https://github.com/Azure/azureml-examples/tree/main/cli/jobs/single-step/dask/nyctaxi

It uses the preview CLI (v2), but can also be adapted fairly easily to use the v1 Python SDK for spinning up the cluster.

CC: @danielsc

[jacobdanovitch]

Can vouch for using dask-mpi within AML jobs, it works flawlessly. Haven't had a single issue so far.

[jacobtomlinson]

@jacobdanovitch that's great to hear. Do you have any examples? It would be really nice to update the docs both here and in dask-mpi to help folks to do this.

[jacobdanovitch]

@jacobtomlinson For sure, what kind of example would be helpful? There wasn't anything extra I had to do to make it work, I just used the mpi distribution in AML and initialized the cluster as shown in the docs.

[jacobtomlinson]

@jacobdanovitch awesome! I think a step by step would be really useful. Perhaps you could point me to some AzureML docs for creating the MPI distribution and I can mash that together with an example from dask-mpi?

[jacobdanovitch]

@jacobtomlinson So sorry, just getting to this now. You can use MPI in AzureML using a commandJob and specifying the distribution.type as MPI (docs here).

You need the az cli and the az ml extension installed (instructions here). The first thing you have to do is create an environment with dask[distributed], dask_mpi, and mpi4py. You can use a conda file to create an AML environment. MWE:

# envs/dask.yaml
$schema: https://azuremlschemas.azureedge.net/latest/environment.schema.json
name: dask-mpi
image: mcr.microsoft.com/azureml/openmpi3.1.2-ubuntu18.04
conda_file: envs/conda/dask.yaml
description: Dask MPI environment.

# envs/conda/dask.yaml
name: dask-mpi
dependencies:
  - python=3.8.10
  - pip:
    - dask[distributed,dataframe,bag]
    - dask_mpi
    - mpi4py
    - pandas
    - tqdm

Then register the environment:

az ml environment create --file envs/dask.yaml

Then you can use it in a commandJob:

$schema: https://azuremlschemas.azureedge.net/latest/commandJob.schema.json
code: 
  local_path: src
command: >-
  python main.py --mpi
environment: azureml:dask-mpi

compute: azureml:cpu-cluster
distribution:
  type: mpi 
  process_count_per_instance: 2
resources:
  instance_count: 2
display_name: ...
experiment_name: ...
description: ...

I then have this snippet:

from dask.distributed import Client

def initialize_client(mpi: bool = True, n_workers: int = 8, scheduler_address: str = None) -> Client:
    if mpi:
        from dask_mpi import initialize
        initialize()
        return Client()

    client = Client(scheduler_address, n_workers=n_workers)
    return client

I just add an --mpi flag to my argparser that I include in the job yaml for AzureML as shown above and omit when running locally and then initialize the client using that snippet. It's really nice because my code stays the exact same running locally and remotely.

Hope this is helpful, let me know if anything else is needed.