Share your experiences with worker-saturation config to reduce memory usage

[gjoseph92]

If you have a workload that struggles with running out of memory, please try it with this new setting and report back!

(and if you don't, please try it anyway and report back too, whether it goes well or poorly! we want to hear as much feedback as possible.)

distributed>=2022.9.2 includes a new configuration option: distributed.scheduler.worker-saturation. This setting controls how many extra initial data-loading tasks workers will run. Full documentation is here.

It's currently set to inf by default, to be consistent with previous behavior. But setting it to 1.0-1.5 can enormously reduce memory usage in many cases: This setting is now on by default in version 2022.11.0 and later!

How do I set it?

Update: This setting is now on by default in version 2022.11.0 and later. You no longer need to follow these instructions to set it; just upgrade to the latest version.

Instructions

However you set dask configuration. Different deployment systems have different ways to do this.

The easiest is often to set an environment variable on the cluster (note this must be set before the scheduler process starts):

$ DASK_DISTRIBUTED__SCHEDULER__WORKER_SATURATION=1.0 dask-scheduler

dask-gateway / Pangeo

Note that different dask-gateway systems may use an option name besides environment (environment_vars, for example), or may not offer the option at all.

from dask_gateway import Gateway
gateway = Gateway()
options = gateway.cluster_options()
options.environment = dict(
    DASK_DISTRIBUTED__SCHEDULER__WORKER_SATURATION="1.0"
)
cluster = gateway.new_cluster(options)

Coiled

import dask
import coiled

with dask.config.set({"distributed.scheduler.worker-saturation": 1.0}):
    # coiled sends current dask config automatically
    cluster = coiled.Cluster(..., package_sync=True)

Local client

import dask
import distributed

with dask.config.set({"distributed.scheduler.worker-saturation": 1.0}):
    client = distributed.Client(...)

dask-kubernetes

from dask_kubernetes.operator import KubeCluster

cluster = KubeCluster(..., env=dict(DASK_DISTRIBUTED__SCHEDULER__WORKER_SATURATION="1.0"))

dask-cloudprovider

from dask_cloudprovider.aws import EC2Cluster

# Same syntax for any cluster type (not just `EC2Cluster`)
cluster = EC2Cluster(..., env_vars=dict(DASK_DISTRIBUTED__SCHEDULER__WORKER_SATURATION="1.0"))

YAML config file

See https://docs.dask.org/en/latest/configuration.html#yaml-files

distributed:
  scheduler:
    worker-saturation: 1.0

Last-resort option / modify a live cluster without restarting

If you have no other option for setting config, or want to change the value on a live cluster without re-creating it, you can do this. Only run this while the scheduler is idle (no tasks):

# enable queuing (new behavior)
client.run_on_scheduler(lambda dask_scheduler: setattr(dask_scheduler, "WORKER_SATURATION", 1.0))

# disable queuing (old behavior)
client.run_on_scheduler(lambda dask_scheduler: setattr(dask_scheduler, "WORKER_SATURATION", float("inf")))

Remember that in all cases, you must be using the latest version of distributed (>=2022.9.2), otherwise nothing will happen!

What does this magic actually do?

It prevents root task overproduction, a phenomenon where workers are too quick to load initial data and then start to run out of memory.

With worker-saturation: 1.0, each worker will never have more than nthreads input chunks in memory at once. So in many cases, you'll not only see lower memory usage, but more constant memory usage (see the flatter blue lines compared to more "peak-y" yellow lines above).

When can I expect it to work?

When you think, "this feels like it should be embarrassingly parallel / dask shouldn't need to load much data to compute each output, so why is it?"

When can't I expect it to work?

When your workload inherently requires having all data in memory at once, this won't help much. Those are generally computations that need to reorganize a dataset, like DataFrame.set_index, DataFrame.merge, and some instances of Array.rechunk.

What are the downsides?

In a number of cases, it's slower. (In every case we've seen that's been significantly slower, though, it's also reduced memory use a lot.)

If your workload is currently spilling a lot to disk (or just crashing), this will probably make things (much) faster.

If your workload is currently fast and happy, this may make it slower. We'd love to hear about that! We hope to make this setting the default soon, so all feedback is very helpful.

This also loses co-assignment #4892, so you may see more data transfer than before. However, benchmarks that we thought were dependent on co-assignment seem to work okay with worker-saturation: 1.0 (they still work and don't use much memory, they're just slower).

What value should I set the magic parameter to? What does it actually mean?

Set it to 1.1. If it's still using too much memory, then set it to 1.0, and please let us know.

Specifically, ceil(worker_nthreads * worker_saturation) controls how many input chunks a worker will have in memory at once.

For the old scheduling behavior, set it to "inf"—this turns off the feature flag. This will also re-enable co-assignment.

We haven't seen larger values reliably make things much faster (but they do seem to reliably increase memory use). Feel free to play around with it though, and please share what you find.

What should I do if I read this far and have tried it out?

Please reply to this discussion (in a new thread) and share your experience! Please include the size of your cluster (# workers, worker CPU/memory), size of your data (total size, # chunks), and what worker-saturation value you used. And of course, how it went! Dashboard screenshots or performance reports are also welcome.

[dhruvbalwada]

I think there is a typo in the code blocks you suggested, I had to use:

with dask.config.set({"distributed.scheduler.worker-saturation":  1.0}): # notice 1. is  not a str
    client = distributed.Client()

or else I get a type error.

[gjoseph92]

What version of distributed are you using? I would expect it to work as a string for 2022.9.2 or later thanks to #7064.

[dhruvbalwada]

Oops, sorry the pangeo docker images are still using 2022.9.1. I have opened an issue pangeo-data/pangeo-docker-images#388 to fix this. Maybe someone who knows how to update these images can take a look.

[dcherian]

Epic stuff 👏🏾 👏🏾 👏🏾 @gjoseph92 See this repo for performance reports, csv of memory usage timeseries, and code.

Anecdotally, it seemed like for 1.0 there was more whitespace in the task stream, so workers kept waiting, but with 1.2 the workers seemed to not have much pauses. This is not at all obvious in those performance reports because there are so many tasks.

TEM

I found one ridiculous improvement (8x?! TEM for "transformed eulerian mean"):

        dims = ("time", "level", "lat", "lon")
        # 1 is number of years, adjust to make bigger, 
        # full dataset is 60-ish years.
        shape = (1 * 365 * 24, 37, 72, 1440)
        chunks = (24, 15, -1, -1)

        ds = xr.Dataset(
            {
                "U": (dims, dask.array.random.random(shape, chunks=chunks)),
                "V": (dims, dask.array.random.random(shape, chunks=chunks)),
                "W": (dims, dask.array.random.random(shape, chunks=chunks)),
                "T": (dims, dask.array.random.random(shape, chunks=chunks)),
            },
            coords={"time": pd.date_range("2001-01-01", periods=shape[0], freq="H")},
        )
        zonal_means = ds.mean("lon")
        anomaly = ds - zonal_means

        anomaly['uv'] = anomaly.U*anomaly.V
        anomaly['vt'] = anomaly.V*anomaly.T
        anomaly['uw'] = anomaly.U*anomaly.W

        temdiags = zonal_means.merge(anomaly[['uv','vt','uw']].mean("lon"))

        temdiags = temdiags.groupby('time.dayofyear').mean()
        temdiags = temdiags.rename({'dayofyear':'time'})

        with performance_report(f"tem-saturation-{val}.html"):
            for repeat in range(repeats):
                with ms.sample(f"{val}_{repeat}"):
                    temdiags.compute()
                client.restart()

uvmean

This one had no difference (uvmean from this xarray issue; this one I ran two repeats for each value),
[I now realize that this one is a mean over the only chunked dimension so it isn't really embarassingly parallel]

       ds = xr.Dataset(
            dict(
                anom_u=(["time", "face", "j", "i"], da.random.random((5000, 1, 987, 1920), chunks=(10, 1, -1, -1))),
                anom_v=(["time", "face", "j", "i"], da.random.random((5000, 1, 987, 1920), chunks=(10, 1, -1, -1))),
            )
        )

        mean = ds**2
        mean["uv"] = ds.anom_u * ds.anom_v
        mean = mean.mean("time")

        with performance_report(f"uvmean-saturation-{val}.html"):
            for repeat in range(repeats):
                with ms.sample(f"{val}_{repeat}"):
                    mean.compute()
                client.restart()

[dcherian]

[gjoseph92]

Thanks @dcherian, awesome results!

re: uvmean, as I mentioned on that xarray issue, I think something else is odd there not related to root task overproduction. I haven't fully grokked the problem to understand whether it's something that should even be doable in fixed memory, or whether it inherently requires most of the data in memory at once.

[Scartography]

This works fine, just kinda stalls the scheduler if the tasks take to long you have the scheduler_timeout for less, than the tasks takes to complete, the dask-gateway thinks it's idle and kills the cluster, so "mind the gap here". You always want to keep scheduler alive, longer than your tasks are processing (or add a fix to check on workers). This is a somewhat extreme scenario for us where a single task takes ~20-30[min].

[fjetter]

Are you talking about idle_timeout that shuts down a cluster if it's idling? I couldn't find anything in dask-gateway or dask/distributed that is called scheduler_timeout.
If you're talking about this, that should definitely not be happening. At least when reading the dask.distributed implementation I don't know what's going wrong but this sounds like a genuine bug. Idk if dask-gateway has a similar functionality.

Generally speaking, the scheduler is aware that tasks are queued and/or processing on the worker regardless of how long the tasks are running. It should not shut down due to a timeout.

[gjoseph92]

@Scartography if there are any queued tasks on the scheduler, it shouldn't shut down due to idle_timeout:

distributed/distributed/scheduler.py

Lines 7684 to 7690 in 6afce9c

	if (
	self.queued
	or self.unrunnable
	or any([ws.processing for ws in self.workers.values()])
	):
	self.idle_since = None
	return

Maybe there's some other system involved in your deployment that thinks the scheduler is idle, but I don't think anything in core distributed is doing this. If you can reproduce, that would be helpful.

[Scartography]

well the idle_timeout of the dask-gateway is given to the schedulers/workers: https://github.com/dask/distributed/blob/main/distributed/scheduler.py#L3405 here I would assume.

It only happens before the first tasks was finished, on a KubeCluster.

I will try to look out for it and try to report further, if this keeps happening.

[TomNicholas]

@gjoseph92 I finally tried this out on my vorticity problem (on the LEAP hub now that we have distributed=2022.10.0 there) - notebook here

Notes:

• This is run on 2% of our full problem (200/9000 timesteps), but with just 10 workers with 20GB each.
• I saw spilling to disk for inf but not for 1.0 or 1.2.
• This was actually the harder version of our problem, the one including grid metrics. That means the graph has every timestep dependent on a small number of inputs (the grid metric chunks). I'm very pleased that it looks like the saturation is working, but we will see if we can scale up the task graph to the full problem without faling afoul of Ignore widely-shared dependencies in decide_worker #5325.
• I'm not sure why the memory profile has those bumps in it even with worker saturation on. I might try a version without the metrics and see if that's flat like @dcherian 's TEM example.

Overall this looks like a huge improvement! 🥳

Next step is for me to try it out on the larger dataset. (I will also give it more workers and more memory per worker in that case though.)

[gjoseph92]

@TomNicholas great news, that looks nice! Please let us know when you try on the larger dataset.

I'm not sure why the memory profile has those bumps in it even with worker saturation on

I'm not too surprised to see that, since I don't think your workload is completely embarrassingly parallel. It's also possible that it's some other scheduling issue (like #5325) popping up, but without root task overproduction, the consequences are significantly reduced.

[gzt5142]

I believe there may be a typo in the gateway/pangeo codeblock:

     26 from dask_gateway import Gateway
     27 gateway = Gateway()
     28 options = gateway.cluster_options()
---> 29 options.environment = dict(
     30     DASK_DISTRIBUTED__SCHEDULER__WORKER_SATURATION="1.0"
     31 )

File /home/conda/noname/e5d2a0145f87c999f2dad147481e52423ffbfc9d8f8d53858ed58d8b5d80a922-20221024-212044-746186-222-blather/lib/python3.10/site-packages/dask_gateway/options.py:108, in Options._set(self, key, value, exc_cls)
    106     self._fields[key].set(value)
    107 except KeyError:
--> 108     raise exc_cls("No option %r available" % key) from None
AttributeError: No option 'environment' available

Versions:

• distributed.__version__ == '2022.10.0'
• dask.__version__ == '2022.10.0'
• uname -s -m == 'Linux x86_64'

??? Should the statement instead be ???:

options.environment_vars = dict(
    DASK_DISTRIBUTED__SCHEDULER__WORKER_SATURATION="1.0"
)

[gjoseph92]

@gzt5142 I just referenced the docs from https://pangeo.io/cloud.html#environment-variables-on-the-cluster. My guess is that different dask-gateway-based deployment systems use a different attribute name for environment variables. Pangeo happens to use environment, yours seems to use environment_vars.

I'll update the instructions to note that the option name may be different on different systems; thanks for pointing this out.

[jbusecke]

Hi @gjoseph92, thank you soooo much for this improvement. I was chatting with @dcherian yesterday and we coincidentally stumbled over an example for an improvement here.
In the past I have not been able to compute a linear trend over a chunked time dimension for CMIP6 models without rechunking.

The code to calculate this trend is a simple xarray one liner

trend = da.fillna(-1e10).polyfit('time', 1, skipna=False)

but computing this in the past would ultimately lead to spilling and eventually killed workers.

Setting the worker saturation fixes this without the need for rechunking the array first. This is a great improvement for this very common workflow in geosciences.

I made a little demo notebook(can be reproduced on the pangeo 2i2c hub - I used a Huge deployment) in case that is helpful!

I ran the computation for a range of saturation values, but interestingly 1.0 seems to work both faster and with lower memory useage.

Ill check this out on another real world example soon and report back.

[gjoseph92]

That's awesome, thanks for sharing @jbusecke!

[jbusecke]

I have not properly tested this yet, but it seems to me that each time I start a cluster with worker-saturation set to 1.0, the start up takes quite a bit longer. Is this something others here have experienced?

[gjoseph92]

Hm, I can't think of any reason within dask why that would happen. If anyone can reproduce this reliably, I'd be interested.

I'd wonder if it maybe has something to do with the deployment system—possibly setting the environment variable requires something different to happen, which takes longer? This feels very unlikely to me as well though. And for pangeo, I can't think of why that would possibly be.

[jbusecke]

Its not a big deal for me atm. Ill keep an eye out and maybe run some more isolated tests in a bit. Just thought Id put that out here if others have the same experience.

[jgdwyer]

Thanks so much for this - I've seen big improvements in reliability with this when using dask.array.tensordot to aggregate gridded data to regions. My workflow was to load gridded data from cloud storage (around 20 GiB), aggregate using tensordot, and save the resulting averaged data back to cloud storage. Without worker saturation, if my cluster was not sufficiently large (24 workers, 30 GiB each), workers would become overwhelmed with data and crash, taking the whole cluster down with it. With this new parameter, I've been able to reduce the number of workers by a factor of 3 and worker memory in half, and it now runs without crashing.

I had previously tried all sorts of different chunking schemes and other cluster settings (worker memory spill/target/terminate, Malloc trim threshold, etc), but none made much of a difference.

[jgdwyer]

If there is a tradeoff in terms of slower performance, I've not noticed it, as I'm just happy that I can run this operation on a smaller cluster.

[gjoseph92]

@jgdwyer just to clarify—you were using worker-saturation: 1.0, right? Not 1.1 or anything larger?

[jgdwyer]

Yes 1.0. I didn't try with 1.1 or any other values

[tasansal]

This solves a lot of our memory bound problems!

I just tried it with our workflow (structured binary serialization of Zarr arrays) and it keeps memory under check. In the past, we have been spilling to disk and having trouble with memory usage.

Should a reasonable number like 1.1 or 1.25 be the default instead of inf?

[gjoseph92]

@tasansal that's the plan, hopefully by the next release: #7213

If you've tried both 1.0 and 1.1 on your workload, I'd be curious to hear if you noticed any difference.

[tasansal]

Yea I can try that, but I need to figure out how to make that nice memory graph first :-)

[gjoseph92]

A common way to do that is with distributed.MemorySampler. You could use to_pandas to write the results to a CSV per run if you need to as well.

But a qualitative answer is also good :)

[gjoseph92]

FYI, 1.1 is now the default in the latest release! (2022.11.0)

[roe246]

Thanks to @gjoseph92 for the worker saturation upgrade and the support for our use case!

We've been struggling with high memory usage (spills, killed workers, timeouts...) along with inconsistent and unbalanced processing for our process that resulted in long runtime and underutilization of cloud resources. With worker saturation, we have immediately seen 30% reduction in runtime and thus cost, 60% reduction in peak memory usage, and significantly more reasonable and consistent task assignment and task processing among Dask workers!

This is what we used to see on our dashboard before -

[roe246]

And this is after!

[mrocklin]

🎉

…

On Tue, Nov 15, 2022, 6:08 PM Gabe Joseph ***@***.***> wrote: FYI, 1.1 is now the default in the latest release! (2022.11.0) — Reply to this email directly, view it on GitHub <#7128 (reply in thread)>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AACKZTGE7TTJOOOMMODCRCDWIQQYBANCNFSM6AAAAAAQ73VLKA> . You are receiving this because you are subscribed to this thread.Message ID: ***@***.***>

[odo2063]

Seems to setting saturation to 0.0 does not have the effect i wished for. Still 256 Tasks for 128 Workers with singlethreads :-/
dask 2023.1.1
distributed 2023.1.1

[gjoseph92]

@odo2063 setting it to 0.0 should actually have raised an error:

distributed/distributed/scheduler.py

Lines 1692 to 1695 in 70abff0

	raise ValueError( # pragma: nocover
	"`distributed.scheduler.worker-saturation` must be a float > 0; got "
	+ repr(self.WORKER_SATURATION)
	)

What's the effect you're hoping for? If you want 128 tasks on 128 threads, then you want 1.0, not 0.0.

[odo2063]

1.0 fills my 128 clients with 128 threads and fills the queue with another 128 threads.
So that tells me that saturation is the value of "queue filling".
So I expect, setting saturation to 0.0, that the queue is always 0.

[fjetter]

1.0 actually means that we are assigning as many tasks to the worker as there are threads. We are not distinguishing between a "queue" and the tasks on the threadpool (since the scheduler doesn't know about this).

What you want to achieve with 0.0 is the behavior of 1.0. The default behavior with 1.1 is very, very close to what you want to achieve

[NadimKawwaProphia]

Thanks for the solution! It solved a problem we had where we kept using a larger EC2 instance but the job would fail to finish. Eventually it dawned on us that it probably wasn't a memory issue but how we were handling it.

[vijayaditya]

client.get_dataset and worker-saturation

I am running out of memory when trying to execute a producer-consumer workflow via publish_dataset and get_dataset. The consumer is executing a trivially parallel UDF which just needs the data related to a single partition at a time(e.g. save partition to remote location). However the consumer worker-pool doesn't have sufficient memory/disk to host the entire published dataset. I noticed that the workers in this pool crash due to RAM / Disk exhaustion issues and do not saturate after the desired number of tasks (worker-saturation * nthreads) are transferred to them.

I am curious if the get_dataset when executed in a newly created client leads to creation of root-tasks which follow the WORKER_SATURATION setting ?

---

NOT A CONTRIBUTION

[fjetter]

If you are using publish_dataset the entire dataset has to fit onto the cluster, similar to when you are using persist. Dask doesn't support storing these datasets in a remote location. If you want to do this, I recommend looking into read/write_parquet or other formats that are appropriate for your data. If you need more help about this, I recommend posting your question in https://dask.discourse.group/

This is unrelated to the root-task scheduling topic that is being discussed here.

[vijayaditya]

Sorry if I was not clear. The question is when I am procuring some data as

with dask.distributed.Client(scheduler_address) as client
  dataset = client.get_dataset(dataset_id)
  .....

will the get_dataset call result in root-tasks, which will enjoy the benefits of worker saturation.

For some context, if I have have a heterogeneous Dask cluster with two sets of nodes in the same dask cluster and I am running publish_dataset for outputs persisted on one set of nodes (say Set1 which can host the entire dataset in-memory) and get_dataset from another set of nodes (say Set2 which has significantly lesser RAM than Set1) I see that Set2 nodes get overwhelmed (i.e., have OOMs or disk exhaustion errors due to spillage) due to the transfer of all the dataset instantaneously. I was curious if "worker-saturation" feature could be taken advantage of in this scenario to queue the data transfer.

My current reading of related code shows that this depends on whether these get_dataset tasks (?) can be considered rootish

distributed/distributed/scheduler.py

Lines 2877 to 2892 in 19c8bf9

	def is_rootish(self, ts: TaskState) -> bool:
	"""
	Whether ``ts`` is a root or root-like task.
	Root-ish tasks are part of a group that's much larger than the cluster,
	and have few or no dependencies.
	"""
	if ts.resource_restrictions or ts.worker_restrictions or ts.host_restrictions:
	return False
	tg = ts.group
	# TODO short-circuit to True if `not ts.dependencies`?
	return (
	len(tg) > self.total_nthreads * 2
	and len(tg.dependencies) < 5
	and sum(map(len, tg.dependencies)) < 5
	)

But given my observations with node failures in Set2 I assume these are not considered root tasks.

Would be happy to move the discussion to https://dask.discourse.group/ if you prefer.

---

NOT A CONTRIBUTION

[fjetter]

get_dataset is similar to persist and keeps the entire dataset on the cluster as long as it is not released. the root-ish scheduling will only benefit computations that are able to release tasks again