Optuna Integration

[mrocklin]

I took a look at what it would take to integrate Dask with Optuna, a hyper-parameter optimization library from the good folks at Preferred Networks (the same people that make CuPy).

Here is a tiny example that works, but slowly. It's informative.

Example

First, create a study

optuna create-study --study-name "distributed-example" --storage "sqlite:///example.db"

Python code

import optuna, time

def objective(trial):
    x = trial.suggest_uniform('x', -10, 10)
    time.sleep(0.050)
    return (x - 2) ** 2

from dask.distributed import Client, wait
client = Client()

def f():
    study = optuna.load_study(study_name='distributed-example', storage='sqlite:///example.db')
    study.optimize(objective, n_trials=100)

futures = [client.submit(f, pure=False) for _ in range(100)]
wait(futures)

# look at dashboard at client.dashboard_link

Problems

• I had to load the study each time (they don't seem to be easy to serialize)
• I had to create the study with the command line (but probably there was a way to handle this in Python that I didn't find)
• The SQLite engine is slow (more below) taking up pretty much all of the time
• I had to manually create some futures, set pure=False and so on, which might not be immediately obvious to new users

Storage

One thing we could do here is create our own storage backend for Optuna. This would either place the information in a Worker (probably with an Actor) or on the scheduler (similar to how we handle Lock/Variable/Queue/...). The Storage API isn't trivial, but is intended to be subclassed. My preference is to put this on the scheduler. Code here: https://github.com/optuna/optuna/tree/master/optuna/storages

This isn't quite as good as a long-term database (presumably it's nice to look at old runs) but it would be fast and easy for users.

Ideal user experience

How would we replace the futures stuff? Maybe something like this might be a target API?

# import optuna
import dask_optuna

study = dask_optuna.load_study(study_name='distributed-example')
study.optimize(objective, n_trials=10000)

With greater integration with Optuna we could imagine something else, perhaps like the following:

import optuna

study = optuna.load_study(study_name='distributed-example', use_dask=True)
study.optimize(objective, n_trials=10000)

This is what we see in projects like TPot. This helps to make the feature a bit more discoverable, but requires integration with an upstream library, which may not be ideal as a first step.

Who cares?

I came to this mostly from a technology perspective. These two libraries seem to both have some traction, and complement each other nicely. However, I don't know this space well enough to know if this is valuable, or if there are users who would find it interesting. I'd love to learn more about this.

[jakirkham]

cc @toshihikoyanase @kmaehashi @Nanthini10 @JohnZed (in case you have thoughts here 🙂)

[mrocklin]

For timing reference I think that making the Dask Storage class would probably take a day for someone who has done this before, and maybe a bit under a week for someone who understood async programming a bit, but hadn't yet played in Dask internals. In general I find peripheral features like this to be good onboarding exercises. The coordination primitives (see lock.py, queue.py, variable.py, semaphore.py) have I think generally been made by different people as onboarding activities, and those people seem to come out of the experience with more comfort.

[jakirkham]

For clarity I don't know what (if anything) people will do with this at this stage. Am just making people aware.

[mrocklin]

Understood. No expectations from my end. Just laying out some seeds to see if I can nerd snipe some people later.

…

On Fri, Aug 28, 2020 at 4:29 PM jakirkham ***@***.***> wrote: For clarity I don't what (if anything) people will do with this at this stage. Am just making people aware. — You are receiving this because you authored the thread. Reply to this email directly, view it on GitHub <#6571 (comment)>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AACKZTHCQX2IUIDBPELFAZDSDA4WVANCNFSM4QOSM6FA> .

[Nanthini10]

The method using futures crashed for me with 20M rows on two RTX 8000s with the following error

Traceback (most recent call last):
File "/opt/conda/envs/rapids/lib/python3.7/site-packages/distributed/comm/core.py", line 295, in _
raise CommClosedError()
distributed.comm.core.CommClosedError
Task exception was never retrieved

I use the following as a workaround with joblib which runs fairly quickly.

from dask_cuda import LocalCUDACluster
from dask.distributed import Client
import optuna
from joblib import parallel_backend

# Read X and y

# Set up client
cluster = LocalCUDACluster(threads_per_worker=1, ip="", dashboard_address="8081")
c = Client(cluster)

# Query the client for all connected workers
workers = c.has_what().keys()
n_workers = len(workers)


def objective(trial, X_param, y_param):
    # ... Add objective ...
    return score


study = optuna.create_study(study_name='joblib-example',
                            storage="sqlite:///example.db",
                            load_if_exists=True)
        

with parallel_backend("dask", n_jobs=n_workers, client=c, scatter=[X,y]):
    study.optimize(lambda trial: objective(trial, X, y),
                   n_jobs=n_workers)

[mrocklin]

@jrbourbeau and I are talking about general design. We thought that we'd put notes here.

class OptunaSchedulerExtension:
    def __init__(self, scheduler):
        self.storage = optuna._storages._in_memory.Storage()  # or whatever this is called
        self.scheduler = scheduler

        self.scheduler.handlers.update({
            "optuna_load_study": self.load_study,
            ...
        })

    def load_study(self, ...):
        self.storage.load_study(...)

class DaskStorage(optuna._storages._base.Storage):  
    """ Implements Optuna Storage API """
    def __init__(self, scheduler_address=None):
        self.scheduler_address = scheduler_address or default_client().scheduler_address

    def load_study(self, ...):
        self.client.scheduler.optuna_load_study(...)


from dask.distributed import Client
dask_client = Client()

study = optuna.load_study(..., storage=DaskStorage(dask_client.scheduler_address))

[jrbourbeau]

Just laying out some seeds to see if I can nerd snipe some people later

Thanks for laying out an example and initial thoughts on an implementation @mrocklin. I'm interested in improving integration between Optuna and Dask and will spend some time putting together a dask storage class 👍

[mrocklin]

I'm glad to hear it @Nanthini10 ! Also, I'm glad to see that adding in RAPIDS work to this doesn't add any extra complications :)

[JohnZed]

Not sure I fully understood the storage proposal. For me, being able to eliminate the need to fire up a DB for Dask-distributed Optuna HPO is a really nice usability benefit. I'm less concerned about the perf implications of using a standard DB, as I think @mrocklin 's demo simulated an optimization that completes in 50ms. For such a fast optimization task, many users would not need to distribute it. If we extend to more typical model optimization times that require distributed workers (several seconds per run), I expect that sqlite contention gets quite low even for a fair level of parallelism. (10 ops/s for sqlite should be not a big issue.)
That's also a long-winded way of asking - can we just use a sqlite backend for storage (attached to the scheduler via like OptunaSchedulerExtension as noted above) to get all the nice properties of on-disk persistence and an already-built API? On-disk is nice for very long-running jobs, like HPO for deep learning models that can take days of runtime and may be resumed later.

[mrocklin]

Sure, so what I'm hearing is that ...

1. wrapping up an Optuna Storage with Dask Comms is valuable from a usability standpoint,
2. but we should make sure that that wrapping doesn't hard-code the current in_memory Optuna storage, but can instead use any Optuna storage, like a sqlite database stored on the machine that holds the Scheduler.

That makes sense to me.

[jrbourbeau]

Thanks for the additional context @JohnZed! It sounds like using optuna.storages.get_storage(...) to determine the underlying Optuna storage class to back a Dask storage class could work

[jrbourbeau]

I've put together a dask-optuna package which has a Dask-compatible DaskStorage class that implements Optuna's BaseStorage API (roughly as outlined in #6571 (comment)). The DaskStorage class works well with Optuna's existing Joblib internals, so it can be used directly with Optuna's API.

For example:

import optuna
import joblib
from dask.distributed import Client
import dask_optuna

def objective(trial):
    x = trial.suggest_uniform("x", -10, 10)
    return (x - 2) ** 2

with Client() as client:
    print(f"Dask dashboard is available at {client.dashboard_link}")
    dask_storage = dask_optuna.DaskStorage("sqlite:///example.db")
    study = optuna.create_study(storage=dask_storage)
    with joblib.parallel_backend("dask"):
        study.optimize(objective, n_trials=200, n_jobs=-1)
    print(f"best_params = {study.best_params}")

I hope this is useful for current Optuna + Dask users and any feedback is very welcome! Here are some relevant links for Dask-Optuna:

GitHub repo: https://github.com/jrbourbeau/dask-optuna
Documentation: https://jrbourbeau.github.io/dask-optuna

cc @JohnZed @Nanthini10 and anyone else who may be interested

[pseudotensor]

Amazing work, thank you so much.

[jcrist]

I believe this can be closed now, since dask-optuna now exists.