Dask-cuda: Communication timeout on DGX-2 with UCX+NVLink

[pentschev]

We've been experiencing a communication timeout on a DGX-2 with UCX+NVlink for a particular problem. The code is below, but I can't share data due to its size (210GB):

import time
import sys
import dask_cudf
import cudf
import cupy
import numpy as np
import cupy as cp
from distributed import wait

sys.path.append("../../tools/")
from readers import build_reader

from dask.distributed import Client
from dask_cuda import DGX
from dask_cuda.initialize import initialize

from dask_cuda import LocalCUDACluster

# ON/OFF settings for various devices
enable_tcp_over_ucx = True
enable_nvlink = True
enable_infiniband = False
interface="enp134s0f1"
protocol="ucx"


if __name__ == "__main__":
    # initialize client with the same settings as workers
    initialize(create_cuda_context=True,
               enable_tcp_over_ucx=enable_tcp_over_ucx,
               enable_infiniband=enable_infiniband,
               enable_nvlink=enable_nvlink)

    if protocol == "tcp":
        cluster = LocalCUDACluster(protocol='tcp',
                                   silence_logs=False,
                                   interface=interface,
                                   CUDA_VISIBLE_DEVICES=list(range(16)))
    elif protocol == "ucx":
        cluster = LocalCUDACluster(protocol='ucx',
                                   silence_logs=False,
                                   enable_tcp_over_ucx=enable_tcp_over_ucx,
                                   enable_infiniband=enable_infiniband,
                                   enable_nvlink=enable_nvlink,
                                   interface=interface,
                                   ucx_net_devices=interface,
                                   CUDA_VISIBLE_DEVICES=list(range(16)))

    client = Client(cluster)
    client.run(cudf.set_allocator, "default", pool=True)

    data_dir = '/data/parquet/'
    q30_session_timeout_inSec = 3600
    q30_limit = 1000
    file_format = "parquet"

    from dask.distributed import performance_report
    with performance_report(filename="dask_report.html"):

        t = time.time()
        def vec_arange(end_sr):
            """
            Returns flattented arange output with start=0, and for end<end_sr[i] in series
            """
            ar = cp.arange(end_sr.max())

            ### get flag matrix for values that are < end
            m = ar < end_sr.values[:, None]

            # ranges in a flattened 1d array
            ranges = (ar * m)[m]
            return ranges


        def get_session_id_from_session_boundry(session_change_df, last_session_len):
            """
                This function returns session starts given a session change df
            """

            user_val_counts = session_change_df.wcs_user_sk.value_counts(sort=False)
            user_val_counts = user_val_counts.reset_index(drop=False)
            user_val_counts = user_val_counts.rename(
                {"index": "wcs_user_sk", "wcs_user_sk": "user_count"}
            )

            ### sort again by user_sk because we want our starts to be aligned
            user_val_counts = user_val_counts.sort_values(by="wcs_user_sk").reset_index(
                drop=True
            )

            end_range = user_val_counts["user_count"]

            user_session_ids = vec_arange(end_range)

            ### up shift the session length df
            session_len = session_change_df["t_index"].diff().reset_index(drop=True)
            session_len = session_len.shift(-1)
            session_len.iloc[-1] = last_session_len

            session_id_final_series = (
                cudf.Series(user_session_ids).repeat(session_len).reset_index(drop=True)
            )
            return session_id_final_series


        def get_session_id(df, time_out):
            """
                This function creates a session id column for each click
                The session id grows in incremeant for each user's susbequent session
                Session boundry is defined by the time_out 
            """

            df["user_change_flag"] = df["wcs_user_sk"].diff(periods=1) != 0
            df["time_delta"] = df["tstamp_inSec"].diff(periods=1)
            df["session_timeout_flag"] = df["tstamp_inSec"].diff(periods=1) > time_out

            df["session_change_flag"] = df["session_timeout_flag"] | df["user_change_flag"]

            # print(f"Total session change = {df['session_change_flag'].sum():,}")

            cols_keep = ["wcs_user_sk", "i_category_id", "session_change_flag"]
            df = df[cols_keep]

            df = df.reset_index(drop=True)
            df["t_index"] = cudf.utils.cudautils.arange(start=0, stop=len(df), dtype=np.int32)

            session_change_df = df[df["session_change_flag"]]
            last_session_len = len(df) - session_change_df["t_index"].iloc[-1]

            session_ids = get_session_id_from_session_boundry(
                session_change_df, last_session_len
            )

            assert len(session_ids) == len(df)
            return session_ids

        table_reader = build_reader(file_format, basepath=data_dir)

        wcs_cols = ["wcs_user_sk", "wcs_item_sk", "wcs_click_date_sk", "wcs_click_time_sk"]
        wcs_df = table_reader.read("web_clickstreams", relevant_cols=wcs_cols)

        item_cols = ["i_category_id", "i_item_sk"]
        item_df = table_reader.read("item", relevant_cols=item_cols)

        f_wcs_df = wcs_df[wcs_df["wcs_user_sk"].notnull()]
        f_item_df = item_df[item_df["i_category_id"].notnull()]

        merged_df = f_wcs_df.merge(f_item_df, left_on=["wcs_item_sk"], right_on=["i_item_sk"])

        merged_df["tstamp_inSec"] = (
            merged_df["wcs_click_date_sk"] * 24 * 60 * 60 + merged_df["wcs_click_time_sk"]
        )
        cols_keep = ["wcs_user_sk", "tstamp_inSec", "i_category_id"]
        merged_df = merged_df[cols_keep]

        ### that the click for each user ends up at the same partition
        merged_df = merged_df.set_index("wcs_user_sk")
        merged_df = merged_df.reset_index(drop=False)

        def get_sessions(df):
            df = df.sort_values(by=["wcs_user_sk", "tstamp_inSec"]).reset_index(drop=True)
            df["session_id"] = get_session_id(df, q30_session_timeout_inSec)
            df = df[["wcs_user_sk", "i_category_id", "session_id"]]
            return df

        session_df = merged_df.map_partitions(get_sessions)
        del merged_df


        def get_distinct_sessions(df):
            df = df.drop_duplicates().reset_index(drop=True)
            return df

        distinct_session_df = session_df.map_partitions(get_distinct_sessions)

        ### get_pair_helper
        def get_pairs(
            df,
            merge_col=["session_id", "wcs_user_sk"],
            pair_col="i_category_id",
            output_col_1="category_id_1",
            output_col_2="category_id_2",
        ):
            """
                Gets pair after doing a inner merge
            """
            pair_df = df.merge(df, on=merge_col, suffixes=["_t1", "_t2"], how="inner")
            pair_df = pair_df[[f"{pair_col}_t1", f"{pair_col}_t2"]]
            pair_df = pair_df[pair_df[f"{pair_col}_t1"] < pair_df[f"{pair_col}_t2"]]
            pair_df = pair_df.rename(
                columns={f"{pair_col}_t1": output_col_1, f"{pair_col}_t2": output_col_2}
            )
            return pair_df

        pair_df = distinct_session_df.map_partitions(get_pairs)
        del distinct_session_df

        print("Time:", time.time() - t)

The problem with this code is that communication times out, but that can be circumvented by increasing distributed.comm.timeouts.connect, still I believe this is just a consequence of a problem with how data transfer behaves. By watching the Dask taskstream I noticed that data transfer takes longer as time passes, the first transfers are of hundreds of milliseconds, then gradually start increasing to few seconds, tens and eventually hundreds of seconds. I've uploaded a Dask performance report of the code above, it was extracted on a DGX-2 running with 16 GPUs and UCX+NVLink. It's worth noting this doesn't happen with TCP communication.

It seems that this issue is created by a combination of factors, certainly small data transfers over NVLink are not efficient without caching CUDA IPC handles, and that probably contributes. I believe there's some sort of issue with a blocking task somewhere that eventually causes transfer times to be computed as if they were longer than they actually are, when they are actually waiting for some blocking task to finish.

Any ideas here @mrocklin @quasiben @jakirkham @madsbk ?

cc @randerzander @beckernick @VibhuJawa

[mrocklin]

Thanks for the upload @pentschev ! That's super interesting to see.

Next time you run this, can I ask you to include the following configuration setting?

distributed:
  comm:
    offload: false

I think that this should help give us more visibility into what is happening during communication.

[pentschev]

New report with distributed.comm.offload: False here.

[mrocklin]

Did something else change? A number of things just got better:

1. Bandwidth went fro 8MB/s to 60MB/s
2. The RMM finalize cost in the "Worker profile (administrative)" profile is mostly gone
3. Runtime generally went down

[mrocklin]

So we're only spending 30-40 cpu-seconds (probably 3s wall time assuming 16 works) in dask serialization or in handing off data with ucx-py. My guess is that the slowdown isn't visible from the Python side.

Does UCX-Py employ a separate Python thread to do any work, or does it hand things off to UCX which handles it itself?

[mrocklin]

For context, settiing distributed.comm.offload = False stops Dask from handling serialization/deserialization in a separate thread. This keeps all of the serialization code (and probably the RMM code as a result) on the same thread as ucx-py, and also the same thread that we're watching with the "Worker profile (administrative)" tab.

[pentschev]

Did something else change?

Yes, sorry about that. I was in a different conda environment before where I had applied some changes to use Numba allocator rather than RMM in dask/distributed. I'm generating a new report, will post soon.

Does UCX-Py employ a separate Python thread to do any work, or does it hand things off to UCX which handles it itself?

I'm not sure about this, I assume it just hands off to UCX. Can you confirm/deny this @quasiben ?

[mrocklin]

It seems to be a good change. Previously we were seeing 1400s (divided by 16 workers) in librmm.rmm_free.

[pentschev]

It seems to be a good change. Previously we were seeing 1400s (divided by 16 workers) in librmm.rmm_free.

Yes, using RMM is the default, but I replaced it to see how would that behave. I guess we know the answer to that now.

[pentschev]

Sorry, there was another change I missed, CUDA IPC cache was enabled. Here are new reports for various combinations, for which I hope I didn't mix anything up:

• With offload, RMM's default allocator + pool, CUDA IPC cache off: 377s
• With offload, RMM's default allocator + pool, CUDA IPC cache on: 138s
• With offload, no allocator/no pool, CUDA IPC cache off: 827s
• With offload, no allocator/no pool, CUDA IPC cache on: OOM
• No offload, RMM's default allocator + pool, CUDA IPC cache off: 345s
• No offload, RMM's default allocator + pool, CUDA IPC cache on: 100s
• No offload, no allocator/no pool, CUDA IPC cache off: 894s
• No offload, no allocator/no pool, CUDA IPC cache on: OOM

As seen above, our most likely solution is to have a CUDA IPC cache. However, this isn't possible in practice today, a continuation of the code above (which was omitted because it isn't relevant for the primary issue reported here), eventually OOMs if we have CUDA IPC cache enabled, as I expected. To improve this I currently don't see a way other than solving https://github.com/openucx/ucx/issues/4410.

[jakirkham]

Was just wondering if we could have a table of options like these. Thanks for putting it together. 🙂

Do we know how much data is being transmitted here? Trying to get a sense of what this means in terms of theoretical vs. actual bandwidth. Am guessing the amount of data transferred is the same between different runs (excepting OOM cases), but please let us know if that is not true.

[pentschev]

No, I don't have a sense of data being transmitted unfortunately. The results above account for the entire run of timed block from the code in the description, so there's a lot of time where no transfer is happening, thus computing bandwidth based on that is probably not going to represent anything correct.

[mrocklin]

Well, that clearly demonstrates the value of the CUDA IPC cache :)

If that was magically resolved, and if we removed the parquet reading (perhaps if this was persisted separately) then communication is still a bit of an issue at 500 MB/s (although clearly way better than it is today). This is entirely speculative, but I'm curious, do you have a sense for what's next to work on after the UCX issue gets resolved?

[pentschev]

Yes, I suspect that the way bandwidth is computed is not entirely representative of the actual transfer bandwidth. I see two possible reasons for that:

1. There's still some blocking operation on UCX-Py that is preventing us from getting optimal bandwidth when there's a large amount of transfer operations; or
2. Dask is accounting for more than it should while computing transfer time.

Obviously, this is just speculation from my part, but I would target these two possibilities as a next step to identify if there's something that we can optimize further.

[mrocklin]

I believe that Dask includes all the time spent once it gives the message off to the Comm object.
This includes both serialization and communication time and potentially connection establishment time.

[mrocklin]

(although for connection establishment we do cache connections)

[quasiben]

@pentschev are you still seeing a timeout with the code posted in this PR given recent techniques of using RMM pool and setting UCX_CUDA_IPC_CACHE=y ?

[jakirkham]

If so, we might want to work with @MattBBaker on this one. He mentioned he was looking into doing some DGX-2 testing soon 🙂

[MattBBaker]

Literally just now got access. Need to write up some script and containers to work with the local ORNL weirdness.

[pentschev]

Thanks for the ping here @quasiben . No, this now works if the following two conditions are met:

1. UCX_CUDA_IPC_CACHE=y;
2. An RMM pool is used.

[pentschev]

I'm not sure whether this will be a good example for @MattBBaker, the issue with this is that I think the dataset is not available outside our network and is something like 200-300 GB.

[jakirkham]

Could we create something similar using random data? That might be easier for others to reproduce 😉

[pentschev]

We probably could but I've asked around in the past (when filing this issue) and it seems pretty complicated data with nobody really has the bandwidth for that at the moment, and is also much beyond my dataframe skills for me to do it.

That said, I would say at this time it's unlikely anyone will get to generate synthetic data for this particular problem. I would also rather see other use cases, we might have a better chance of tracking down new issues with cases we don't know of.

[pentschev]

I'm somewhat convinced that this problem will remain when users can't use an RMM pool or disable the IPC cache. In such situations, all transfers would require mapping the memory handle which has a cost in the range of 100 ms. For that reason, I think we should encourage users to always use an RMM pool and turn on IPC cache (in UCX-Py we've enabled it back by default), in principle it would work without those but will offer very little benefit (if any), therefore I suggest we keep things as is and tell users to rely on RMM pool+IPC cache as there isn't much we'll be able to do otherwise. Any objections to this proposal?

[jakirkham]

Was going to add the option to fallback to CuPy in Distributed as CuPy also has a memory pool and is likely in use by users working with Python on a GPU.

IOW the order would be RMM -> CuPy -> Numba

[pentschev]

Agreed @jakirkham , we should encourage usage of _A_ memory pool + IPC cache, rather than saying RMM pool only.

[github-actions[bot]]

This issue has been marked stale due to no recent activity in the past 30d. Please close this issue if no further response or action is needed. Otherwise, please respond with a comment indicating any updates or changes to the original issue and/or confirm this issue still needs to be addressed. This issue will be marked rotten if there is no activity in the next 60d.

[github-actions[bot]]

This issue has been labeled inactive-90d due to no recent activity in the past 90 days. Please close this issue if no further response or action is needed. Otherwise, please respond with a comment indicating any updates or changes to the original issue and/or confirm this issue still needs to be addressed.