def test_s3_dataset(s3, paths, engine, df):
# create a mocked out bucket here
bucket = "testbucket"
s3.create_bucket(Bucket=bucket)
s3_paths = []
for path in paths:
s3_path = f"s3://{bucket}/{path}"
with fsspec.open(s3_path, "wb") as f:
f.write(open(path, "rb").read())
s3_paths.append(s3_path)
# create a basic s3 dataset
dataset = nvt.Dataset(s3_paths)
# make sure the iteration API works
columns = mycols_pq if engine == "parquet" else mycols_csv
gdf = cudf.concat(list(dataset.to_iter()))[columns]
assert_eq(gdf.reset_index(drop=True), df.reset_index(drop=True))
cat_names = ["name-cat", "name-string"] if engine == "parquet" else ["name-string"]
cont_names = ["x", "y", "id"]
label_name = ["label"]
processor = nvt.Workflow(cat_names=cat_names, cont_names=cont_names, label_name=label_name)
processor.add_feature([ops.FillMissing(), ops.Clip(min_value=0), ops.LogOp()])
processor.add_preprocess(ops.Normalize())
processor.add_preprocess(ops.Categorify(cat_cache="host"))
processor.finalize()
processor.update_stats(dataset)
def test_concatenate_dataframe(tmpdir, output_model):
# we were seeing an issue in the rossmann workflow where we dropped certain columns,
# https://github.com/NVIDIA/NVTabular/issues/961
df = _make_df({
"cat": ["aaaa", "bbbb", "cccc", "aaaa", "bbbb", "aaaa"],
"cont": [0.0, 1.0, 2.0, 3.0, 4.0, 5],
})
# this bug only happened with a dataframe representation: force this by using a lambda
cats = ["cat"] >> ops.LambdaOp(lambda col: _hash_series(col) % 1000)
conts = ["cont"] >> ops.Normalize() >> ops.FillMissing() >> ops.LogOp()
dataset = Dataset(df)
workflow = nvt.Workflow(cats + conts).fit_schema(dataset.infer_schema())
if output_model == "pytorch":
model_info = {
"cat": {
"columns": ["cat"],
"dtype": "int32"
},
"cont": {
"columns": ["cont"],
"dtype": "float32"
},
}
else:
model_info = None
_verify_workflow_on_tritonserver(tmpdir, workflow, df,
"test_concatenate_dataframe",
output_model, model_info)
def test_log(tmpdir, df, dataset, gpu_memory_frac, engine, op_columns):
cont_names = ["x", "y", "id"]
log_op = ops.LogOp(columns=op_columns)
columns_ctx = {}
columns_ctx["continuous"] = {}
columns_ctx["continuous"]["base"] = cont_names
for gdf in dataset.to_iter():
new_gdf = log_op.apply_op(gdf, columns_ctx, "continuous")
assert new_gdf[cont_names] == np.log(gdf[cont_names].astype(np.float32))
def test_target_encode(tmpdir, cat_groups, kfold, fold_seed):
df = cudf.DataFrame({
"Author": list(string.ascii_uppercase),
"Engaging-User": list(string.ascii_lowercase),
"Cost": range(26),
"Post": [0, 1] * 13,
})
df = dask_cudf.from_cudf(df, npartitions=3)
cat_names = ["Author", "Engaging-User"]
cont_names = ["Cost"]
label_name = ["Post"]
processor = nvt.Workflow(cat_names=cat_names,
cont_names=cont_names,
label_name=label_name)
processor.add_feature(
[ops.FillMissing(),
ops.Clip(min_value=0),
ops.LogOp()])
processor.add_preprocess(
ops.TargetEncoding(
cat_groups,
"Cost", # cont_target
out_path=str(tmpdir),
kfold=kfold,
out_col="test_name",
out_dtype="float32",
fold_seed=fold_seed,
drop_folds=False, # Keep folds to validate
))
processor.finalize()
processor.apply(nvt.Dataset(df), output_format=None)
df_out = processor.get_ddf().compute(scheduler="synchronous")
assert "test_name" in df_out.columns
assert df_out["test_name"].dtype == "float32"
if kfold > 1:
# Cat columns are unique.
# Make sure __fold__ mapping is correct
if cat_groups == "Author":
name = "__fold___Author"
cols = ["__fold__", "Author"]
else:
name = "__fold___Author_Engaging-User"
cols = ["__fold__", "Author", "Engaging-User"]
check = cudf.io.read_parquet(processor.stats["te_stats"][name])
check = check[cols].sort_values(cols).reset_index(drop=True)
df_out_check = df_out[cols].sort_values(cols).reset_index(drop=True)
assert_eq(check, df_out_check)
def test_s3_dataset(s3_base, s3so, paths, datasets, engine, df):
# Copy files to mock s3 bucket
files = {}
for i, path in enumerate(paths):
with open(path, "rb") as f:
fbytes = f.read()
fn = path.split(os.path.sep)[-1]
files[fn] = BytesIO()
files[fn].write(fbytes)
files[fn].seek(0)
if engine == "parquet":
# Workaround for nvt#539. In order to avoid the
# bug in Dask's `create_metadata_file`, we need
# to manually generate a "_metadata" file here.
# This can be removed after dask#7295 is merged
# (see https://github.com/dask/dask/pull/7295)
fn = "_metadata"
files[fn] = BytesIO()
meta = create_metadata_file(
paths,
engine="pyarrow",
out_dir=False,
)
meta.write_metadata_file(files[fn])
files[fn].seek(0)
with s3_context(s3_base=s3_base, bucket=engine, files=files):
# Create nvt.Dataset from mock s3 paths
url = f"s3://{engine}" if engine == "parquet" else f"s3://{engine}/*"
dataset = nvt.Dataset(url, engine=engine, storage_options=s3so)
# Check that the iteration API works
columns = mycols_pq if engine == "parquet" else mycols_csv
gdf = cudf.concat(list(dataset.to_iter()))[columns]
assert_eq(gdf.reset_index(drop=True), df.reset_index(drop=True))
cat_names = ["name-cat", "name-string"
] if engine == "parquet" else ["name-string"]
cont_names = ["x", "y", "id"]
label_name = ["label"]
conts = cont_names >> ops.FillMissing() >> ops.Clip(
min_value=0) >> ops.LogOp()
cats = cat_names >> ops.Categorify(cat_cache="host")
processor = nvt.Workflow(conts + cats + label_name)
processor.fit(dataset)
def test_target_encode(tmpdir, cat_groups, kfold, fold_seed, cpu):
df = dispatch._make_df({
"Author": list(string.ascii_uppercase),
"Engaging-User": list(string.ascii_lowercase),
"Cost": range(26),
"Post": [0, 1] * 13,
})
if cpu:
df = dd.from_pandas(
df if isinstance(df, pd.DataFrame) else df.to_pandas(),
npartitions=3)
else:
df = dask_cudf.from_cudf(df, npartitions=3)
cont_names = ["Cost"]
te_features = cat_groups >> ops.TargetEncoding(
cont_names,
out_path=str(tmpdir),
kfold=kfold,
out_dtype="float32",
fold_seed=fold_seed,
drop_folds=False, # Keep folds to validate
)
cont_features = cont_names >> ops.FillMissing() >> ops.Clip(
min_value=0) >> ops.LogOp()
workflow = nvt.Workflow(te_features + cont_features +
["Author", "Engaging-User"])
df_out = workflow.fit_transform(
nvt.Dataset(df)).to_ddf().compute(scheduler="synchronous")
df_lib = dispatch.get_lib()
if kfold > 1:
# Cat columns are unique.
# Make sure __fold__ mapping is correct
if cat_groups == "Author":
name = "__fold___Author"
cols = ["__fold__", "Author"]
else:
name = "__fold___Author_Engaging-User"
cols = ["__fold__", "Author", "Engaging-User"]
check = df_lib.read_parquet(te_features.op.stats[name])
check = check[cols].sort_values(cols).reset_index(drop=True)
df_out_check = df_out[cols].sort_values(cols).reset_index(drop=True)
assert_eq(check, df_out_check, check_dtype=False)
def test_hugectr(tmpdir, df, dataset, output_format, engine, op_columns):
cat_names = ["name-cat", "name-string"] if engine == "parquet" else ["name-string"]
cont_names = ["x", "y"]
label_names = ["label"]
# set variables
nfiles = 10
ext = ""
outdir = tmpdir + "/hugectr"
os.mkdir(outdir)
# process data
processor = nvt.Workflow(cat_names=cat_names, cont_names=cont_names, label_name=label_names)
processor.add_feature([ops.ZeroFill(columns=op_columns), ops.LogOp()])
processor.add_preprocess(ops.Normalize())
processor.add_preprocess(ops.Categorify())
processor.finalize()
# Need to collect statistics first (for now)
processor.update_stats(dataset)
# Second "online" pass to write HugeCTR output
processor.apply(
dataset,
apply_offline=False,
record_stats=False,
output_path=outdir,
out_files_per_proc=nfiles,
output_format=output_format,
shuffle=False,
)
# Check files
ext = ""
if output_format == "parquet":
ext = "parquet"
assert os.path.isfile(outdir + "/metadata.json")
elif output_format == "hugectr":
ext = "data"
assert os.path.isfile(outdir + "/file_list.txt")
for n in range(nfiles):
assert os.path.isfile(os.path.join(outdir, str(n) + "." + ext))
def test_log(tmpdir, datasets, gpu_memory_frac, engine, op_columns):
paths = glob.glob(str(datasets[engine]) + "/*." + engine.split("-")[0])
if engine == "parquet":
df1 = cudf.read_parquet(paths[0])[mycols_pq]
df2 = cudf.read_parquet(paths[1])[mycols_pq]
else:
df1 = cudf.read_csv(paths[0], header=False,
names=allcols_csv)[mycols_csv]
df2 = cudf.read_csv(paths[1], header=False,
names=allcols_csv)[mycols_csv]
df = cudf.concat([df1, df2], axis=0)
df["id"] = df["id"].astype("int64")
if engine == "parquet":
columns = mycols_pq
else:
columns = mycols_csv
cont_names = ["x", "y", "id"]
data_itr = nvtabular.io.GPUDatasetIterator(
paths,
columns=columns,
use_row_groups=True,
gpu_memory_frac=gpu_memory_frac,
names=allcols_csv,
)
log_op = ops.LogOp(columns=op_columns)
columns_ctx = {}
columns_ctx["continuous"] = {}
columns_ctx["continuous"]["base"] = cont_names
for gdf in data_itr:
new_gdf = log_op.apply_op(gdf, columns_ctx, "continuous")
assert new_gdf[cont_names] == np.log(gdf[cont_names].astype(
np.float32))
def test_dask_workflow_api_dlrm(
client, tmpdir, datasets, freq_threshold, part_mem_fraction, engine, cat_cache, on_host, shuffle
):
paths = glob.glob(str(datasets[engine]) + "/*." + engine.split("-")[0])
if engine == "parquet":
df1 = cudf.read_parquet(paths[0])[mycols_pq]
df2 = cudf.read_parquet(paths[1])[mycols_pq]
elif engine == "csv":
df1 = cudf.read_csv(paths[0], header=0)[mycols_csv]
df2 = cudf.read_csv(paths[1], header=0)[mycols_csv]
else:
df1 = cudf.read_csv(paths[0], names=allcols_csv)[mycols_csv]
df2 = cudf.read_csv(paths[1], names=allcols_csv)[mycols_csv]
df0 = cudf.concat([df1, df2], axis=0)
if engine == "parquet":
cat_names = ["name-cat", "name-string"]
else:
cat_names = ["name-string"]
cont_names = ["x", "y", "id"]
label_name = ["label"]
cats = cat_names >> ops.Categorify(
freq_threshold=freq_threshold, out_path=str(tmpdir), cat_cache=cat_cache, on_host=on_host
)
conts = cont_names >> ops.FillMissing() >> ops.Clip(min_value=0) >> ops.LogOp()
workflow = Workflow(cats + conts + label_name, client=client)
if engine in ("parquet", "csv"):
dataset = Dataset(paths, part_mem_fraction=part_mem_fraction)
else:
dataset = Dataset(paths, names=allcols_csv, part_mem_fraction=part_mem_fraction)
output_path = os.path.join(tmpdir, "processed")
transformed = workflow.fit_transform(dataset)
transformed.to_parquet(output_path=output_path, shuffle=shuffle)
# Can still access the final ddf if we didn't shuffle
if not shuffle:
result = transformed.to_ddf().compute()
assert len(df0) == len(result)
assert result["x"].min() == 0.0
assert result["x"].isna().sum() == 0
assert result["y"].min() == 0.0
assert result["y"].isna().sum() == 0
# Check category counts
cat_expect = df0.groupby("name-string").agg({"name-string": "count"}).reset_index(drop=True)
cat_result = (
result.groupby("name-string").agg({"name-string": "count"}).reset_index(drop=True)
)
if freq_threshold:
cat_expect = cat_expect[cat_expect["name-string"] >= freq_threshold]
# Note that we may need to skip the 0th element in result (null mapping)
assert_eq(
cat_expect,
cat_result.iloc[1:] if len(cat_result) > len(cat_expect) else cat_result,
check_index=False,
)
else:
assert_eq(cat_expect, cat_result)
# Read back from disk
df_disk = dask_cudf.read_parquet(output_path, index=False).compute()
for col in df_disk:
assert_eq(result[col], df_disk[col])
else:
df_disk = dask_cudf.read_parquet(output_path, index=False).compute()
assert len(df0) == len(df_disk)
def test_gpu_preproc(tmpdir, df, dataset, dump, gpu_memory_frac, engine,
preprocessing):
cat_names = ["name-cat", "name-string"
] if engine == "parquet" else ["name-string"]
cont_names = ["x", "y", "id"]
label_name = ["label"]
processor = nvt.Workflow(cat_names=cat_names,
cont_names=cont_names,
label_name=label_name)
processor.add_feature(
[ops.FillMedian(),
ops.LogOp(preprocessing=preprocessing)])
processor.add_preprocess(ops.Normalize())
processor.add_preprocess(ops.Categorify())
processor.finalize()
processor.update_stats(dataset)
if dump:
config_file = tmpdir + "/temp.yaml"
processor.save_stats(config_file)
processor.clear_stats()
processor.load_stats(config_file)
def get_norms(tar: cudf.Series):
ser_median = tar.dropna().quantile(0.5, interpolation="linear")
gdf = tar.fillna(ser_median)
gdf = np.log(gdf + 1)
return gdf
# Check mean and std - No good right now we have to add all other changes; Zerofill, Log
x_col = "x" if preprocessing else "x_LogOp"
y_col = "y" if preprocessing else "y_LogOp"
assert math.isclose(get_norms(df.x).mean(),
processor.stats["means"][x_col],
rel_tol=1e-2)
assert math.isclose(get_norms(df.y).mean(),
processor.stats["means"][y_col],
rel_tol=1e-2)
assert math.isclose(get_norms(df.x).std(),
processor.stats["stds"][x_col],
rel_tol=1e-2)
assert math.isclose(get_norms(df.y).std(),
processor.stats["stds"][y_col],
rel_tol=1e-2)
# Check median (TODO: Improve the accuracy)
x_median = df.x.dropna().quantile(0.5, interpolation="linear")
y_median = df.y.dropna().quantile(0.5, interpolation="linear")
id_median = df.id.dropna().quantile(0.5, interpolation="linear")
assert math.isclose(x_median, processor.stats["medians"]["x"], rel_tol=1e1)
assert math.isclose(y_median, processor.stats["medians"]["y"], rel_tol=1e1)
assert math.isclose(id_median,
processor.stats["medians"]["id"],
rel_tol=1e1)
# Check that categories match
if engine == "parquet":
cats_expected0 = df["name-cat"].unique().values_host
cats0 = get_cats(processor, "name-cat")
assert cats0.tolist() == [None] + cats_expected0.tolist()
cats_expected1 = df["name-string"].unique().values_host
cats1 = get_cats(processor, "name-string")
assert cats1.tolist() == [None] + cats_expected1.tolist()
# Write to new "shuffled" and "processed" dataset
processor.write_to_dataset(tmpdir,
dataset,
nfiles=10,
shuffle=True,
apply_ops=True)
processor.create_final_cols()
# if preprocessing
if not preprocessing:
for col in cont_names:
assert f"{col}_LogOp" in processor.columns_ctx["final"]["cols"][
"continuous"]
dlc = torch_dataloader.DLCollator(preproc=processor, apply_ops=False)
data_files = [
torch_dataloader.FileItrDataset(x,
use_row_groups=True,
gpu_memory_frac=gpu_memory_frac,
names=allcols_csv)
for x in glob.glob(str(tmpdir) + "/*.parquet")
]
data_itr = torch.utils.data.ChainDataset(data_files)
dl = torch_dataloader.DLDataLoader(data_itr,
collate_fn=dlc.gdf_col,
pin_memory=False,
num_workers=0)
len_df_pp = 0
for chunk in dl:
len_df_pp += len(chunk[0][0])
dataset = Dataset(glob.glob(str(tmpdir) + "/*.parquet"),
part_mem_fraction=gpu_memory_frac)
x = processor.ds_to_tensors(dataset.to_iter(), apply_ops=False)
num_rows, num_row_groups, col_names = cudf.io.read_parquet_metadata(
str(tmpdir) + "/_metadata")
assert len(x[0]) == len_df_pp
itr_ds = torch_dataloader.TensorItrDataset([x[0], x[1], x[2]],
batch_size=512000)
count_tens_itr = 0
for data_gd in itr_ds:
count_tens_itr += len(data_gd[1])
assert data_gd[0].shape[1] > 0
assert data_gd[1].shape[1] > 0
assert len_df_pp == count_tens_itr
def test_dask_workflow_api_dlrm(dask_cluster, tmpdir, datasets, freq_threshold,
part_mem_fraction, engine):
paths = glob.glob(str(datasets[engine]) + "/*." + engine.split("-")[0])
if engine == "parquet":
df1 = cudf.read_parquet(paths[0])[mycols_pq]
df2 = cudf.read_parquet(paths[1])[mycols_pq]
elif engine == "csv":
df1 = cudf.read_csv(paths[0], header=0)[mycols_csv]
df2 = cudf.read_csv(paths[1], header=0)[mycols_csv]
else:
df1 = cudf.read_csv(paths[0], names=allcols_csv)[mycols_csv]
df2 = cudf.read_csv(paths[1], names=allcols_csv)[mycols_csv]
df0 = cudf.concat([df1, df2], axis=0)
if engine == "parquet":
cat_names = ["name-cat", "name-string"]
else:
cat_names = ["name-string"]
cont_names = ["x", "y", "id"]
label_name = ["label"]
processor = Workflow(client=client,
cat_names=cat_names,
cont_names=cont_names,
label_name=label_name)
processor.add_feature([ops.ZeroFill(), ops.LogOp()])
processor.add_preprocess(
ops.Categorify(freq_threshold=freq_threshold,
out_path=str(tmpdir),
split_out=2))
processor.finalize()
if engine in ("parquet", "csv"):
dataset = DaskDataset(paths, part_mem_fraction=part_mem_fraction)
else:
dataset = DaskDataset(paths,
names=allcols_csv,
part_mem_fraction=part_mem_fraction)
processor.apply(dataset, output_path=str(tmpdir))
result = processor.get_ddf().compute()
assert len(df0) == len(result)
assert result["x"].min() == 0.0
assert result["x"].isna().sum() == 0
assert result["y"].min() == 0.0
assert result["y"].isna().sum() == 0
# Check category counts
cat_expect = df0.groupby("name-string").agg({
"name-string": "count"
}).reset_index(drop=True)
cat_result = result.groupby("name-string").agg({
"name-string": "count"
}).reset_index(drop=True)
if freq_threshold:
cat_expect = cat_expect[cat_expect["name-string"] >= freq_threshold]
# Note that we may need to skip the 0th element in result (null mapping)
assert_eq(
cat_expect,
cat_result.iloc[1:]
if len(cat_result) > len(cat_expect) else cat_result,
check_index=False,
)
else:
assert_eq(cat_expect, cat_result)
# Read back from disk
df_disk = dask_cudf.read_parquet("/".join([str(tmpdir), "processed"]),
index=False).compute()
for col in df_disk:
assert_eq(result[col], df_disk[col])
def main(args):
# Input
data_path = args.data_path
out_path = args.out_path
freq_limit = args.freq_limit
out_files_per_proc = args.splits
if args.protocol == "ucx":
os.environ["UCX_TLS"] = "tcp,cuda_copy,cuda_ipc,sockcm"
# Use Criteo dataset by default (for now)
cont_names = (args.cont_names.split(",")
if args.cont_names else ["I" + str(x) for x in range(1, 14)])
cat_names = (args.cat_names.split(",")
if args.cat_names else ["C" + str(x) for x in range(1, 27)])
label_name = ["label"]
if args.cat_splits:
tree_width = {
name: int(s)
for name, s in zip(cat_names, args.cat_splits.split(","))
}
else:
tree_width = {col: 1 for col in cat_names}
if args.cat_names is None:
# Using Criteo... Use more hash partitions for
# known high-cardinality columns
tree_width["C20"] = 8
tree_width["C1"] = 8
tree_width["C22"] = 4
tree_width["C10"] = 4
tree_width["C21"] = 2
tree_width["C11"] = 2
tree_width["C23"] = 2
tree_width["C12"] = 2
# Specify categorical caching location
cat_cache = None
if args.cat_cache:
cat_cache = args.cat_cache.split(",")
if len(cat_cache) == 1:
cat_cache = cat_cache[0]
else:
# If user is specifying a list of options,
# they must specify an option for every cat column
assert len(cat_names) == len(cat_cache)
if isinstance(cat_cache, str):
cat_cache = {col: cat_cache for col in cat_names}
elif isinstance(cat_cache, list):
cat_cache = {name: c for name, c in zip(cat_names, cat_cache)}
else:
# Criteo/DLRM Defaults
cat_cache = {col: "device" for col in cat_names}
if args.cat_names is None:
cat_cache["C20"] = "host"
cat_cache["C1"] = "host"
# Only need to cache the largest two on a dgx-2
if args.n_workers < 16:
cat_cache["C22"] = "host"
cat_cache["C10"] = "host"
# Use total device size to calculate args.device_limit_frac
device_size = device_mem_size(kind="total")
device_limit = int(args.device_limit_frac * device_size)
device_pool_size = int(args.device_pool_frac * device_size)
part_size = int(args.part_mem_frac * device_size)
# Setup LocalCUDACluster
if args.protocol == "tcp":
cluster = LocalCUDACluster(
protocol=args.protocol,
n_workers=args.n_workers,
CUDA_VISIBLE_DEVICES=args.devs,
device_memory_limit=device_limit,
local_directory=args.dask_workspace,
dashboard_address=":3787",
)
else:
cluster = LocalCUDACluster(
protocol=args.protocol,
n_workers=args.n_workers,
CUDA_VISIBLE_DEVICES=args.devs,
enable_nvlink=True,
device_memory_limit=device_limit,
local_directory=args.dask_workspace,
dashboard_address=":3787",
)
client = Client(cluster)
# Setup RMM pool
if not args.no_rmm_pool:
setup_rmm_pool(client, device_pool_size)
# Define Dask NVTabular "Workflow"
processor = Workflow(cat_names=cat_names,
cont_names=cont_names,
label_name=label_name,
client=client)
processor.add_feature([ops.ZeroFill(), ops.LogOp()])
processor.add_preprocess(
ops.Categorify(
out_path=out_path,
tree_width=tree_width,
cat_cache=cat_cache,
freq_threshold=freq_limit,
on_host=args.cat_on_host,
))
processor.finalize()
dataset = Dataset(data_path, "parquet", part_size=part_size)
# Execute the dask graph
runtime = time.time()
if args.profile is not None:
with performance_report(filename=args.profile):
processor.apply(
dataset,
shuffle="full" if args.worker_shuffle else "partial",
out_files_per_proc=out_files_per_proc,
output_path=out_path,
)
else:
processor.apply(
dataset,
shuffle="full" if args.worker_shuffle else "partial",
out_files_per_proc=out_files_per_proc,
output_path=out_path,
)
runtime = time.time() - runtime
print("\nDask-NVTabular DLRM/Criteo benchmark")
print("--------------------------------------")
print(f"partition size | {part_size}")
print(f"protocol | {args.protocol}")
print(f"device(s) | {args.devs}")
print(f"rmm-pool | {(not args.no_rmm_pool)}")
print(f"out_files_per_proc | {args.splits}")
print(f"worker-shuffle | {args.worker_shuffle}")
print("======================================")
print(f"Runtime[s] | {runtime}")
print("======================================\n")
client.close()
def test_gpu_workflow_api(tmpdir, df, dataset, gpu_memory_frac, engine, dump,
op_columns):
cat_names = ["name-cat", "name-string"
] if engine == "parquet" else ["name-string"]
cont_names = ["x", "y", "id"]
label_name = ["label"]
processor = nvt.Workflow(
cat_names=cat_names,
cont_names=cont_names,
label_name=label_name,
)
processor.add_feature([ops.ZeroFill(columns=op_columns), ops.LogOp()])
processor.add_preprocess(ops.Normalize())
processor.add_preprocess(ops.Categorify())
processor.finalize()
processor.update_stats(dataset)
if dump:
config_file = tmpdir + "/temp.yaml"
processor.save_stats(config_file)
processor.clear_stats()
processor.load_stats(config_file)
def get_norms(tar: cudf.Series):
gdf = tar.fillna(0)
gdf = gdf * (gdf >= 0).astype("int")
gdf = np.log(gdf + 1)
return gdf
# Check mean and std - No good right now we have to add all other changes; Zerofill, Log
if not op_columns:
assert math.isclose(
get_norms(df.y).mean(),
processor.stats["means"]["y"],
rel_tol=1e-1,
)
assert math.isclose(
get_norms(df.y).std(),
processor.stats["stds"]["y"],
rel_tol=1e-1,
)
assert math.isclose(
get_norms(df.x).mean(),
processor.stats["means"]["x"],
rel_tol=1e-1,
)
assert math.isclose(
get_norms(df.x).std(),
processor.stats["stds"]["x"],
rel_tol=1e-1,
)
# Check that categories match
if engine == "parquet":
cats_expected0 = df["name-cat"].unique().values_to_string()
cats0 = processor.stats["encoders"]["name-cat"].get_cats(
).values_to_string()
# adding the None entry as a string because of move from gpu
assert cats0 == ["None"] + cats_expected0
cats_expected1 = df["name-string"].unique().values_to_string()
cats1 = processor.stats["encoders"]["name-string"].get_cats(
).values_to_string()
# adding the None entry as a string because of move from gpu
assert cats1 == ["None"] + cats_expected1
# Write to new "shuffled" and "processed" dataset
processor.write_to_dataset(tmpdir,
dataset,
nfiles=10,
shuffle=True,
apply_ops=True)
data_itr_2 = nvtabular.io.GPUDatasetIterator(
glob.glob(str(tmpdir) + "/ds_part.*.parquet"),
use_row_groups=True,
gpu_memory_frac=gpu_memory_frac,
)
df_pp = None
for chunk in data_itr_2:
df_pp = cudf.concat([df_pp, chunk], axis=0) if df_pp else chunk
if engine == "parquet":
assert df_pp["name-cat"].dtype == "int64"
assert df_pp["name-string"].dtype == "int64"
num_rows, num_row_groups, col_names = cudf.io.read_parquet_metadata(
str(tmpdir) + "/_metadata")
assert num_rows == len(df_pp)
return processor.ds_exports
def main(args):
"""Multi-GPU Criteo/DLRM Preprocessing Benchmark
This benchmark is designed to measure the time required to preprocess
the Criteo (1TB) dataset for Facebook’s DLRM model. The user must specify
the path of the raw dataset (using the `--data-path` flag), as well as the
output directory for all temporary/final data (using the `--out-path` flag)
Example Usage
-------------
python dask-nvtabular-criteo-benchmark.py
--data-path /path/to/criteo_parquet --out-path /out/dir/`
Dataset Requirements (Parquet)
------------------------------
This benchmark is designed with a parquet-formatted dataset in mind.
While a CSV-formatted dataset can be processed by NVTabular, converting
to parquet will yield significantly better performance. To convert your
dataset, try using the `optimize_criteo.ipynb` notebook (also located
in `NVTabular/examples/`)
For a detailed parameter overview see `NVTabular/examples/MultiGPUBench.md`
"""
# Input
data_path = args.data_path
freq_limit = args.freq_limit
out_files_per_proc = args.out_files_per_proc
high_card_columns = args.high_cards.split(",")
dashboard_port = args.dashboard_port
if args.protocol == "ucx":
UCX_TLS = os.environ.get("UCX_TLS", "tcp,cuda_copy,cuda_ipc,sockcm")
os.environ["UCX_TLS"] = UCX_TLS
# Cleanup output directory
BASE_DIR = args.out_path
dask_workdir = os.path.join(BASE_DIR, "workdir")
output_path = os.path.join(BASE_DIR, "output")
stats_path = os.path.join(BASE_DIR, "stats")
if not os.path.isdir(BASE_DIR):
os.mkdir(BASE_DIR)
for dir_path in (dask_workdir, output_path, stats_path):
if os.path.isdir(dir_path):
shutil.rmtree(dir_path)
os.mkdir(dir_path)
# Use Criteo dataset by default (for now)
cont_names = (args.cont_names.split(",")
if args.cont_names else ["I" + str(x) for x in range(1, 14)])
cat_names = (args.cat_names.split(",")
if args.cat_names else ["C" + str(x) for x in range(1, 27)])
label_name = ["label"]
# Specify Categorify/GroupbyStatistics options
tree_width = {}
cat_cache = {}
for col in cat_names:
if col in high_card_columns:
tree_width[col] = args.tree_width
cat_cache[col] = args.cat_cache_high
else:
tree_width[col] = 1
cat_cache[col] = args.cat_cache_low
# Use total device size to calculate args.device_limit_frac
device_size = device_mem_size(kind="total")
device_limit = int(args.device_limit_frac * device_size)
device_pool_size = int(args.device_pool_frac * device_size)
part_size = int(args.part_mem_frac * device_size)
# Parse shuffle option
shuffle = None
if args.shuffle == "PER_WORKER":
shuffle = nvt_io.Shuffle.PER_WORKER
elif args.shuffle == "PER_PARTITION":
shuffle = nvt_io.Shuffle.PER_PARTITION
# Check if any device memory is already occupied
for dev in args.devices.split(","):
fmem = _pynvml_mem_size(kind="free", index=int(dev))
used = (device_size - fmem) / 1e9
if used > 1.0:
warnings.warn(
f"BEWARE - {used} GB is already occupied on device {int(dev)}!"
)
# Setup LocalCUDACluster
if args.protocol == "tcp":
cluster = LocalCUDACluster(
protocol=args.protocol,
n_workers=args.n_workers,
CUDA_VISIBLE_DEVICES=args.devices,
device_memory_limit=device_limit,
local_directory=dask_workdir,
dashboard_address=":" + dashboard_port,
)
else:
cluster = LocalCUDACluster(
protocol=args.protocol,
n_workers=args.n_workers,
CUDA_VISIBLE_DEVICES=args.devices,
enable_nvlink=True,
device_memory_limit=device_limit,
local_directory=dask_workdir,
dashboard_address=":" + dashboard_port,
)
client = Client(cluster)
# Setup RMM pool
if args.device_pool_frac > 0.01:
setup_rmm_pool(client, device_pool_size)
# Define Dask NVTabular "Workflow"
processor = Workflow(cat_names=cat_names,
cont_names=cont_names,
label_name=label_name,
client=client)
if args.normalize:
processor.add_feature([ops.FillMissing(), ops.Normalize()])
else:
processor.add_feature(
[ops.FillMissing(),
ops.Clip(min_value=0),
ops.LogOp()])
processor.add_preprocess(
ops.Categorify(
out_path=stats_path,
tree_width=tree_width,
cat_cache=cat_cache,
freq_threshold=freq_limit,
search_sorted=not freq_limit,
on_host=not args.cats_on_device,
))
processor.finalize()
dataset = Dataset(data_path, "parquet", part_size=part_size)
# Execute the dask graph
runtime = time.time()
if args.profile is not None:
with performance_report(filename=args.profile):
processor.apply(
dataset,
shuffle=shuffle,
out_files_per_proc=out_files_per_proc,
output_path=output_path,
num_io_threads=args.num_io_threads,
)
else:
processor.apply(
dataset,
num_io_threads=args.num_io_threads,
shuffle=shuffle,
out_files_per_proc=out_files_per_proc,
output_path=output_path,
)
runtime = time.time() - runtime
print("\nDask-NVTabular DLRM/Criteo benchmark")
print("--------------------------------------")
print(f"partition size | {part_size}")
print(f"protocol | {args.protocol}")
print(f"device(s) | {args.devices}")
print(f"rmm-pool-frac | {(args.device_pool_frac)}")
print(f"out-files-per-proc | {args.out_files_per_proc}")
print(f"num_io_threads | {args.num_io_threads}")
print(f"shuffle | {args.shuffle}")
print(f"cats-on-device | {args.cats_on_device}")
print("======================================")
print(f"Runtime[s] | {runtime}")
print("======================================\n")
client.close()
def test_dask_workflow_api_dlrm(
client,
tmpdir,
datasets,
freq_threshold,
part_mem_fraction,
engine,
cat_cache,
on_host,
shuffle,
cpu,
):
paths = glob.glob(str(datasets[engine]) + "/*." + engine.split("-")[0])
paths = sorted(paths)
if engine == "parquet":
df1 = cudf.read_parquet(paths[0])[mycols_pq]
df2 = cudf.read_parquet(paths[1])[mycols_pq]
elif engine == "csv":
df1 = cudf.read_csv(paths[0], header=0)[mycols_csv]
df2 = cudf.read_csv(paths[1], header=0)[mycols_csv]
else:
df1 = cudf.read_csv(paths[0], names=allcols_csv)[mycols_csv]
df2 = cudf.read_csv(paths[1], names=allcols_csv)[mycols_csv]
df0 = cudf.concat([df1, df2], axis=0)
df0 = df0.to_pandas() if cpu else df0
if engine == "parquet":
cat_names = ["name-cat", "name-string"]
else:
cat_names = ["name-string"]
cont_names = ["x", "y", "id"]
label_name = ["label"]
cats = cat_names >> ops.Categorify(freq_threshold=freq_threshold,
out_path=str(tmpdir),
cat_cache=cat_cache,
on_host=on_host)
conts = cont_names >> ops.FillMissing() >> ops.Clip(
min_value=0) >> ops.LogOp()
workflow = Workflow(cats + conts + label_name, client=client)
if engine in ("parquet", "csv"):
dataset = Dataset(paths, cpu=cpu, part_mem_fraction=part_mem_fraction)
else:
dataset = Dataset(paths,
cpu=cpu,
names=allcols_csv,
part_mem_fraction=part_mem_fraction)
output_path = os.path.join(tmpdir, "processed")
transformed = workflow.fit_transform(dataset)
transformed.to_parquet(output_path=output_path,
shuffle=shuffle,
out_files_per_proc=1)
result = transformed.to_ddf().compute()
assert len(df0) == len(result)
assert result["x"].min() == 0.0
assert result["x"].isna().sum() == 0
assert result["y"].min() == 0.0
assert result["y"].isna().sum() == 0
# Check categories. Need to sort first to make sure we are comparing
# "apples to apples"
expect = df0.sort_values(["label", "x", "y",
"id"]).reset_index(drop=True).reset_index()
got = result.sort_values(["label", "x", "y",
"id"]).reset_index(drop=True).reset_index()
dfm = expect.merge(got, on="index",
how="inner")[["name-string_x", "name-string_y"]]
dfm_gb = dfm.groupby(["name-string_x", "name-string_y"]).agg({
"name-string_x":
"count",
"name-string_y":
"count"
})
if freq_threshold:
dfm_gb = dfm_gb[dfm_gb["name-string_x"] >= freq_threshold]
assert_eq(dfm_gb["name-string_x"],
dfm_gb["name-string_y"],
check_names=False)
# Read back from disk
if cpu:
df_disk = dd_read_parquet(output_path).compute()
else:
df_disk = dask_cudf.read_parquet(output_path).compute()
# we don't have a deterministic ordering here, especially when using
# a dask client with multiple workers - so we need to sort the values here
columns = ["label", "x", "y", "id"] + cat_names
got = result.sort_values(columns).reset_index(drop=True)
expect = df_disk.sort_values(columns).reset_index(drop=True)
assert_eq(got, expect)
@pytest.mark.parametrize(
"op",
[
ops.Bucketize([1]),
ops.Rename(postfix="_trim"),
ops.Categorify(),
ops.Categorify(encode_type="combo"),
ops.Clip(0),
ops.DifferenceLag("col1"),
ops.FillMissing(),
ops.Groupby("col1"),
ops.HashBucket(1),
ops.HashedCross(1),
ops.JoinGroupby("col1"),
ops.ListSlice(0),
ops.LogOp(),
ops.Normalize(),
ops.TargetEncoding("col1"),
],
)
def test_workflow_select_by_tags(op):
schema1 = ColumnSchema("col1", tags=["b", "c", "d"])
schema2 = ColumnSchema("col2", tags=["c", "d"])
schema3 = ColumnSchema("col3", tags=["d"])
schema = Schema([schema1, schema2, schema3])
cont_features = ColumnSelector(tags=["c"]) >> op
workflow = Workflow(cont_features)
workflow.fit_schema(schema)
output_cols = op.output_column_names(ColumnSelector(["col1", "col2"]))
def test_s3_dataset(s3_base, s3so, paths, datasets, engine, df,
patch_aiobotocore):
# Copy files to mock s3 bucket
files = {}
for i, path in enumerate(paths):
with open(path, "rb") as f:
fbytes = f.read()
fn = path.split(os.path.sep)[-1]
files[fn] = BytesIO()
files[fn].write(fbytes)
files[fn].seek(0)
if engine == "parquet":
# Workaround for nvt#539. In order to avoid the
# bug in Dask's `create_metadata_file`, we need
# to manually generate a "_metadata" file here.
# This can be removed after dask#7295 is merged
# (see https://github.com/dask/dask/pull/7295)
fn = "_metadata"
files[fn] = BytesIO()
meta = create_metadata_file(
paths,
engine="pyarrow",
out_dir=False,
)
meta.write_metadata_file(files[fn])
files[fn].seek(0)
with s3_context(s3_base=s3_base, bucket=engine, files=files) as s3fs:
# Create nvt.Dataset from mock s3 paths
url = f"s3://{engine}" if engine == "parquet" else f"s3://{engine}/*"
dataset = nvt.Dataset(url, engine=engine, storage_options=s3so)
# Check that the iteration API works
columns = mycols_pq if engine == "parquet" else mycols_csv
gdf = nvt.dispatch._concat(list(dataset.to_iter()))[columns]
assert_eq(gdf.reset_index(drop=True), df.reset_index(drop=True))
cat_names = ["name-cat", "name-string"
] if engine == "parquet" else ["name-string"]
cont_names = ["x", "y", "id"]
label_name = ["label"]
conts = cont_names >> ops.FillMissing() >> ops.Clip(
min_value=0) >> ops.LogOp()
cats = cat_names >> ops.Categorify(cat_cache="host")
processor = nvt.Workflow(conts + cats + label_name)
processor.fit(dataset)
# make sure we can write out the dataset back to S3
# (https://github.com/NVIDIA-Merlin/NVTabular/issues/1214)
processor.transform(dataset).to_parquet(f"s3://{engine}/output")
expected = processor.transform(dataset).to_ddf().compute()
# make sure we can write out the workflow to s3
processor.save(f"s3://{engine}/saved_workflow/")
# make sure the workflow got saved to the right spot in S3
workflow_files = s3fs.ls(f"/{engine}/saved_workflow/")
assert workflow_files
# finally make sure we can read in the workflow from S3, and use it
# to transform values and get the same result as on the local fs
reloaded = nvt.Workflow.load(f"s3://{engine}/saved_workflow/")
from_s3 = reloaded.transform(dataset).to_ddf().compute()
assert_eq(expected, from_s3)
def test_hugectr(tmpdir, df, dataset, output_format, engine, op_columns, num_io_threads):
cat_names = ["name-cat", "name-string"] if engine == "parquet" else ["name-string"]
cont_names = ["x", "y"]
label_names = ["label"]
# set variables
nfiles = 10
ext = ""
outdir = tmpdir + "/hugectr"
os.mkdir(outdir)
# process data
processor = nvt.Workflow(cat_names=cat_names, cont_names=cont_names, label_name=label_names)
processor.add_feature([ops.ZeroFill(columns=op_columns), ops.LogOp()])
processor.add_preprocess(ops.Normalize())
processor.add_preprocess(ops.Categorify())
processor.finalize()
# Need to collect statistics first (for now)
processor.update_stats(dataset)
# Second "online" pass to write HugeCTR output
processor.apply(
dataset,
apply_offline=False,
record_stats=False,
output_path=outdir,
out_files_per_proc=nfiles,
output_format=output_format,
shuffle=False,
num_io_threads=num_io_threads,
)
# Check for _file_list.txt
assert os.path.isfile(outdir + "/_file_list.txt")
# Check for _metadata.json
assert os.path.isfile(outdir + "/_metadata.json")
# Check contents of _metadata.json
data = {}
col_summary = {}
with open(outdir + "/_metadata.json", "r") as fil:
for k, v in json.load(fil).items():
data[k] = v
assert "cats" in data
assert "conts" in data
assert "labels" in data
assert "file_stats" in data
assert len(data["file_stats"]) == nfiles
for cdata in data["cats"] + data["conts"] + data["labels"]:
col_summary[cdata["index"]] = cdata["col_name"]
# Check that data files exist
ext = ""
if output_format == "parquet":
ext = "parquet"
elif output_format == "hugectr":
ext = "data"
for n in range(nfiles):
assert os.path.isfile(os.path.join(outdir, str(n) + "." + ext))
# Make sure the columns in "_metadata.json" make sense
if output_format == "parquet":
df_check = cudf.read_parquet(os.path.join(outdir, "0.parquet"))
for i, name in enumerate(df_check.columns):
if i in col_summary:
assert col_summary[i] == name
def test_workflow_apply(client, use_client, tmpdir, shuffle, apply_offline):
out_files_per_proc = 2
out_path = str(tmpdir.mkdir("processed"))
path = str(tmpdir.join("simple.parquet"))
size = 25
row_group_size = 5
cont_columns = ["cont1", "cont2"]
cat_columns = ["cat1", "cat2"]
label_column = ["label"]
df = pd.DataFrame(
{
"cont1": np.arange(size, dtype=np.float64),
"cont2": np.arange(size, dtype=np.float64),
"cat1": np.arange(size, dtype=np.int32),
"cat2": np.arange(size, dtype=np.int32),
"label": np.arange(size, dtype=np.float64),
}
)
df.to_parquet(path, row_group_size=row_group_size, engine="pyarrow")
dataset = nvt.Dataset(path, engine="parquet", row_groups_per_part=1)
processor = nvt.Workflow(
cat_names=cat_columns,
cont_names=cont_columns,
label_name=label_column,
client=client if use_client else None,
)
processor.add_cont_feature([ops.FillMissing(), ops.Clip(min_value=0), ops.LogOp()])
processor.add_cat_preprocess(ops.Categorify())
processor.finalize()
# Force dtypes
dict_dtypes = {}
for col in cont_columns:
dict_dtypes[col] = np.float32
for col in cat_columns:
dict_dtypes[col] = np.float32
for col in label_column:
dict_dtypes[col] = np.int64
if not apply_offline:
processor.apply(
dataset,
output_format=None,
record_stats=True,
)
processor.apply(
dataset,
apply_offline=apply_offline,
record_stats=apply_offline,
output_path=out_path,
shuffle=shuffle,
out_files_per_proc=out_files_per_proc,
dtypes=dict_dtypes,
)
# Check dtypes
for filename in glob.glob(os.path.join(out_path, "*.parquet")):
gdf = cudf.io.read_parquet(filename)
assert dict(gdf.dtypes) == dict_dtypes
def test_gpu_preproc(tmpdir, datasets, dump, gpu_memory_frac, engine,
preprocessing):
paths = glob.glob(str(datasets[engine]) + "/*." + engine.split("-")[0])
if engine == "parquet":
df1 = cudf.read_parquet(paths[0])[mycols_pq]
df2 = cudf.read_parquet(paths[1])[mycols_pq]
else:
df1 = cudf.read_csv(paths[0], header=False,
names=allcols_csv)[mycols_csv]
df2 = cudf.read_csv(paths[1], header=False,
names=allcols_csv)[mycols_csv]
df = cudf.concat([df1, df2], axis=0)
df["id"] = df["id"].astype("int64")
if engine == "parquet":
cat_names = ["name-cat", "name-string"]
columns = mycols_pq
else:
cat_names = ["name-string"]
columns = mycols_csv
cont_names = ["x", "y", "id"]
label_name = ["label"]
processor = nvt.Workflow(
cat_names=cat_names,
cont_names=cont_names,
label_name=label_name,
to_cpu=True,
)
processor.add_feature(
[ops.FillMissing(),
ops.LogOp(preprocessing=preprocessing)])
processor.add_preprocess(ops.Normalize())
processor.add_preprocess(ops.Categorify())
processor.finalize()
data_itr = nvtabular.io.GPUDatasetIterator(
paths,
columns=columns,
use_row_groups=True,
gpu_memory_frac=gpu_memory_frac,
names=allcols_csv,
)
processor.update_stats(data_itr)
if dump:
config_file = tmpdir + "/temp.yaml"
processor.save_stats(config_file)
processor.clear_stats()
processor.load_stats(config_file)
def get_norms(tar: cudf.Series):
ser_median = tar.dropna().quantile(0.5, interpolation="linear")
gdf = tar.fillna(ser_median)
gdf = np.log(gdf + 1)
return gdf
# Check mean and std - No good right now we have to add all other changes; Zerofill, Log
x_col = "x" if preprocessing else "x_LogOp"
y_col = "y" if preprocessing else "y_LogOp"
assert math.isclose(
get_norms(df.x).mean(),
processor.stats["means"][x_col],
rel_tol=1e-2,
)
assert math.isclose(
get_norms(df.y).mean(),
processor.stats["means"][y_col],
rel_tol=1e-2,
)
assert math.isclose(
get_norms(df.x).std(),
processor.stats["stds"][x_col],
rel_tol=1e-2,
)
assert math.isclose(
get_norms(df.y).std(),
processor.stats["stds"][y_col],
rel_tol=1e-2,
)
# Check median (TODO: Improve the accuracy)
x_median = df.x.dropna().quantile(0.5, interpolation="linear")
y_median = df.y.dropna().quantile(0.5, interpolation="linear")
id_median = df.id.dropna().quantile(0.5, interpolation="linear")
assert math.isclose(x_median, processor.stats["medians"]["x"], rel_tol=1e1)
assert math.isclose(y_median, processor.stats["medians"]["y"], rel_tol=1e1)
assert math.isclose(id_median,
processor.stats["medians"]["id"],
rel_tol=1e1)
# Check that categories match
if engine == "parquet":
cats_expected0 = df["name-cat"].unique().values_to_string()
cats0 = processor.stats["encoders"]["name-cat"].get_cats(
).values_to_string()
assert cats0 == ["None"] + cats_expected0
cats_expected1 = df["name-string"].unique().values_to_string()
cats1 = processor.stats["encoders"]["name-string"].get_cats(
).values_to_string()
print(cats1)
assert cats1 == ["None"] + cats_expected1
# Write to new "shuffled" and "processed" dataset
processor.write_to_dataset(tmpdir,
data_itr,
nfiles=10,
shuffle=True,
apply_ops=True)
processor.create_final_cols()
# if preprocessing
if not preprocessing:
for col in cont_names:
assert f"{col}_LogOp" in processor.columns_ctx["final"]["cols"][
"continuous"]
dlc = nvtabular.torch_dataloader.DLCollator(preproc=processor,
apply_ops=False)
data_files = [
nvtabular.torch_dataloader.FileItrDataset(
x,
use_row_groups=True,
gpu_memory_frac=gpu_memory_frac,
names=allcols_csv,
) for x in glob.glob(str(tmpdir) + "/ds_part.*.parquet")
]
data_itr = torch.utils.data.ChainDataset(data_files)
dl = nvtabular.torch_dataloader.DLDataLoader(data_itr,
collate_fn=dlc.gdf_col,
pin_memory=False,
num_workers=0)
len_df_pp = 0
for chunk in dl:
len_df_pp += len(chunk[0][0])
data_itr = nvtabular.io.GPUDatasetIterator(
glob.glob(str(tmpdir) + "/ds_part.*.parquet"),
use_row_groups=True,
gpu_memory_frac=gpu_memory_frac,
names=allcols_csv,
)
x = processor.ds_to_tensors(data_itr, apply_ops=False)
num_rows, num_row_groups, col_names = cudf.io.read_parquet_metadata(
str(tmpdir) + "/_metadata")
assert len(x[0]) == len_df_pp
itr_ds = nvtabular.torch_dataloader.TensorItrDataset([x[0], x[1], x[2]],
batch_size=512000)
count_tens_itr = 0
for data_gd in itr_ds:
count_tens_itr += len(data_gd[1])
assert data_gd[0][0].shape[1] > 0
assert data_gd[0][1].shape[1] > 0
assert len_df_pp == count_tens_itr
if os.path.exists(processor.ds_exports):
shutil.rmtree(processor.ds_exports)
def test_gpu_workflow_api(tmpdir, client, df, dataset, gpu_memory_frac, engine,
dump, op_columns, use_client):
cat_names = ["name-cat", "name-string"
] if engine == "parquet" else ["name-string"]
cont_names = ["x", "y", "id"]
label_name = ["label"]
processor = nvt.Workflow(
cat_names=cat_names,
cont_names=cont_names,
label_name=label_name,
client=client if use_client else None,
)
processor.add_feature([ops.ZeroFill(columns=op_columns), ops.LogOp()])
processor.add_preprocess(ops.Normalize())
processor.add_preprocess(ops.Categorify(cat_cache="host"))
processor.finalize()
processor.update_stats(dataset)
if dump:
config_file = tmpdir + "/temp.yaml"
processor.save_stats(config_file)
processor.clear_stats()
processor.load_stats(config_file)
def get_norms(tar: cudf.Series):
gdf = tar.fillna(0)
gdf = gdf * (gdf >= 0).astype("int")
gdf = np.log(gdf + 1)
return gdf
# Check mean and std - No good right now we have to add all other changes; Zerofill, Log
if not op_columns:
assert math.isclose(get_norms(df.y).mean(),
processor.stats["means"]["y"],
rel_tol=1e-1)
assert math.isclose(get_norms(df.y).std(),
processor.stats["stds"]["y"],
rel_tol=1e-1)
assert math.isclose(get_norms(df.x).mean(),
processor.stats["means"]["x"],
rel_tol=1e-1)
assert math.isclose(get_norms(df.x).std(),
processor.stats["stds"]["x"],
rel_tol=1e-1)
# Check that categories match
if engine == "parquet":
cats_expected0 = df["name-cat"].unique().values_host
cats0 = get_cats(processor, "name-cat")
# adding the None entry as a string because of move from gpu
assert cats0.tolist() == [None] + cats_expected0.tolist()
cats_expected1 = df["name-string"].unique().values_host
cats1 = get_cats(processor, "name-string")
# adding the None entry as a string because of move from gpu
assert cats1.tolist() == [None] + cats_expected1.tolist()
# Write to new "shuffled" and "processed" dataset
processor.write_to_dataset(tmpdir,
dataset,
out_files_per_proc=10,
shuffle="partial",
apply_ops=True)
dataset_2 = Dataset(glob.glob(str(tmpdir) + "/*.parquet"),
part_mem_fraction=gpu_memory_frac)
df_pp = cudf.concat(list(dataset_2.to_iter()), axis=0)
if engine == "parquet":
assert is_integer_dtype(df_pp["name-cat"].dtype)
assert is_integer_dtype(df_pp["name-string"].dtype)
num_rows, num_row_groups, col_names = cudf.io.read_parquet_metadata(
str(tmpdir) + "/_metadata")
assert num_rows == len(df_pp)
def test_hugectr(tmpdir, client, df, dataset, output_format, engine,
op_columns, num_io_threads, use_client):
client = client if use_client else None
cat_names = ["name-cat", "name-string"
] if engine == "parquet" else ["name-string"]
cont_names = ["x", "y"]
label_names = ["label"]
# set variables
nfiles = 10
ext = ""
outdir = tmpdir + "/hugectr"
os.mkdir(outdir)
# process data
processor = nvt.Workflow(client=client,
cat_names=cat_names,
cont_names=cont_names,
label_name=label_names)
processor.add_feature([
ops.FillMissing(columns=op_columns),
ops.Clip(min_value=0, columns=op_columns),
ops.LogOp(),
])
processor.add_preprocess(ops.Normalize())
processor.add_preprocess(ops.Categorify())
processor.finalize()
# apply the workflow and write out the dataset
processor.apply(
dataset,
output_path=outdir,
out_files_per_proc=nfiles,
output_format=output_format,
shuffle=None,
num_io_threads=num_io_threads,
)
# Check for _file_list.txt
assert os.path.isfile(outdir + "/_file_list.txt")
# Check for _metadata.json
assert os.path.isfile(outdir + "/_metadata.json")
# Check contents of _metadata.json
data = {}
col_summary = {}
with open(outdir + "/_metadata.json", "r") as fil:
for k, v in json.load(fil).items():
data[k] = v
assert "cats" in data
assert "conts" in data
assert "labels" in data
assert "file_stats" in data
assert len(data["file_stats"]) == nfiles if not client else nfiles * len(
client.cluster.workers)
for cdata in data["cats"] + data["conts"] + data["labels"]:
col_summary[cdata["index"]] = cdata["col_name"]
# Check that data files exist
ext = ""
if output_format == "parquet":
ext = "parquet"
elif output_format == "hugectr":
ext = "data"
data_files = [
os.path.join(outdir, filename) for filename in os.listdir(outdir)
if filename.endswith(ext)
]
# Make sure the columns in "_metadata.json" make sense
if output_format == "parquet":
df_check = cudf.read_parquet(os.path.join(outdir, data_files[0]))
for i, name in enumerate(df_check.columns):
if i in col_summary:
assert col_summary[i] == name
def test_gpu_workflow_config(tmpdir, datasets, dump, gpu_memory_frac, engine,
replace):
paths = glob.glob(str(datasets[engine]) + "/*." + engine.split("-")[0])
if engine == "parquet":
df1 = cudf.read_parquet(paths[0])[mycols_pq]
df2 = cudf.read_parquet(paths[1])[mycols_pq]
else:
df1 = cudf.read_csv(paths[0], header=False,
names=allcols_csv)[mycols_csv]
df2 = cudf.read_csv(paths[1], header=False,
names=allcols_csv)[mycols_csv]
df = cudf.concat([df1, df2], axis=0)
df["id"] = df["id"].astype("int64")
if engine == "parquet":
cat_names = ["name-cat", "name-string"]
columns = mycols_pq
else:
cat_names = ["name-string"]
columns = mycols_csv
cont_names = ["x", "y", "id"]
label_name = ["label"]
config = nvt.workflow.get_new_config()
# add operators with dependencies
config["FE"]["continuous"] = [[
ops.FillMissing(replace=replace),
ops.LogOp()
]]
config["PP"]["continuous"] = [[
ops.LogOp(replace=replace),
ops.Normalize()
]]
config["PP"]["categorical"] = [ops.Categorify()]
processor = nvt.Workflow(
cat_names=cat_names,
cont_names=cont_names,
label_name=label_name,
config=config,
to_cpu=False,
)
data_itr = nvt.io.GPUDatasetIterator(
paths,
columns=columns,
use_row_groups=True,
gpu_memory_frac=gpu_memory_frac,
names=allcols_csv,
)
processor.update_stats(data_itr)
if dump:
config_file = tmpdir + "/temp.yaml"
processor.save_stats(config_file)
processor.clear_stats()
processor.load_stats(config_file)
def get_norms(tar: cudf.Series):
ser_median = tar.dropna().quantile(0.5, interpolation="linear")
gdf = tar.fillna(ser_median)
gdf = np.log(gdf + 1)
return gdf
# Check mean and std - No good right now we have to add all other changes; Zerofill, Log
concat_ops = "_FillMissing_LogOp"
if replace:
concat_ops = ""
assert math.isclose(
get_norms(df.x).mean(),
processor.stats["means"]["x" + concat_ops],
rel_tol=1e-1,
)
assert math.isclose(
get_norms(df.y).mean(),
processor.stats["means"]["y" + concat_ops],
rel_tol=1e-1,
)
assert math.isclose(
get_norms(df.x).std(),
processor.stats["stds"]["x" + concat_ops],
rel_tol=1e-1,
)
assert math.isclose(
get_norms(df.y).std(),
processor.stats["stds"]["y" + concat_ops],
rel_tol=1e-1,
)
# Check that categories match
if engine == "parquet":
cats_expected0 = df["name-cat"].unique().values_to_string()
cats0 = processor.stats["encoders"]["name-cat"].get_cats(
).values_to_string()
# adding the None entry as a string because of move from gpu
assert cats0 == ["None"] + cats_expected0
cats_expected1 = df["name-string"].unique().values_to_string()
cats1 = processor.stats["encoders"]["name-string"].get_cats(
).values_to_string()
# adding the None entry as a string because of move from gpu
assert cats1 == ["None"] + cats_expected1
# Write to new "shuffled" and "processed" dataset
processor.write_to_dataset(tmpdir,
data_itr,
nfiles=10,
shuffle=True,
apply_ops=True)
data_itr_2 = nvtabular.io.GPUDatasetIterator(
glob.glob(str(tmpdir) + "/ds_part.*.parquet"),
use_row_groups=True,
gpu_memory_frac=gpu_memory_frac,
)
df_pp = None
for chunk in data_itr_2:
df_pp = cudf.concat([df_pp, chunk], axis=0) if df_pp else chunk
if engine == "parquet":
assert df_pp["name-cat"].dtype == "int64"
assert df_pp["name-string"].dtype == "int64"
num_rows, num_row_groups, col_names = cudf.io.read_parquet_metadata(
str(tmpdir) + "/_metadata")
assert num_rows == len(df_pp)
return processor.ds_exports
def test_gpu_workflow_config(tmpdir, client, df, dataset, gpu_memory_frac,
engine, dump, replace):
cat_names = ["name-cat", "name-string"
] if engine == "parquet" else ["name-string"]
cont_names = ["x", "y", "id"]
label_name = ["label"]
config = nvt.workflow.get_new_config()
# add operators with dependencies
config["FE"]["continuous"] = [[
ops.FillMissing(replace=replace),
ops.LogOp(replace=replace)
]]
config["PP"]["continuous"] = [[
ops.LogOp(replace=replace),
ops.Normalize()
]]
config["PP"]["categorical"] = [ops.Categorify()]
processor = nvt.Workflow(
cat_names=cat_names,
cont_names=cont_names,
label_name=label_name,
config=config,
client=client,
)
processor.update_stats(dataset)
if dump:
config_file = tmpdir + "/temp.yaml"
processor.save_stats(config_file)
processor.clear_stats()
processor.load_stats(config_file)
def get_norms(tar: cudf.Series):
ser_median = tar.dropna().quantile(0.5, interpolation="linear")
gdf = tar.fillna(ser_median)
gdf = np.log(gdf + 1)
return gdf
# Check mean and std - No good right now we have to add all other changes; Zerofill, Log
concat_ops = "_FillMissing_LogOp"
if replace:
concat_ops = ""
assert math.isclose(get_norms(df.x).mean(),
processor.stats["means"]["x" + concat_ops],
rel_tol=1e-1)
assert math.isclose(get_norms(df.y).mean(),
processor.stats["means"]["y" + concat_ops],
rel_tol=1e-1)
assert math.isclose(get_norms(df.x).std(),
processor.stats["stds"]["x" + concat_ops],
rel_tol=1e-1)
assert math.isclose(get_norms(df.y).std(),
processor.stats["stds"]["y" + concat_ops],
rel_tol=1e-1)
# Check that categories match
if engine == "parquet":
cats_expected0 = df["name-cat"].unique().values_host
cats0 = get_cats(processor, "name-cat")
# adding the None entry as a string because of move from gpu
assert cats0.tolist() == [None] + cats_expected0.tolist()
cats_expected1 = df["name-string"].unique().values_host
cats1 = get_cats(processor, "name-string")
# adding the None entry as a string because of move from gpu
assert cats1.tolist() == [None] + cats_expected1.tolist()
# Write to new "shuffled" and "processed" dataset
processor.write_to_dataset(tmpdir,
dataset,
out_files_per_proc=10,
shuffle="partial",
apply_ops=True)
dataset_2 = Dataset(glob.glob(str(tmpdir) + "/*.parquet"),
part_mem_fraction=gpu_memory_frac)
df_pp = cudf.concat(list(dataset_2.to_iter()), axis=0)
if engine == "parquet":
assert is_integer_dtype(df_pp["name-cat"].dtype)
assert is_integer_dtype(df_pp["name-string"].dtype)
num_rows, num_row_groups, col_names = cudf.io.read_parquet_metadata(
str(tmpdir) + "/_metadata")
assert num_rows == len(df_pp)
def nvt_etl(
data_path,
out_path,
devices,
protocol,
device_limit_frac,
device_pool_frac,
part_mem_frac,
cats,
conts,
labels,
out_files_per_proc,
):
# Set up data paths
input_path = data_path[:-1] if data_path[-1] == "/" else data_path
base_dir = out_path[:-1] if out_path[-1] == "/" else out_path
dask_workdir = os.path.join(base_dir, "workdir")
output_path = os.path.join(base_dir, "output")
stats_path = os.path.join(base_dir, "stats")
output_train_dir = os.path.join(output_path, "train/")
output_valid_dir = os.path.join(output_path, "valid/")
# Make sure we have a clean worker space for Dask
if os.path.isdir(dask_workdir):
shutil.rmtree(dask_workdir)
os.makedirs(dask_workdir)
# Make sure we have a clean stats space for Dask
if os.path.isdir(stats_path):
shutil.rmtree(stats_path)
os.mkdir(stats_path)
# Make sure we have a clean output path
if os.path.isdir(output_path):
shutil.rmtree(output_path)
os.mkdir(output_path)
os.mkdir(output_train_dir)
os.mkdir(output_valid_dir)
# Get train/valid files
train_paths = [
os.path.join(input_path, f) for f in os.listdir(input_path)
if os.path.isfile(os.path.join(input_path, f))
]
n_files = int(len(train_paths) * 0.9)
valid_paths = train_paths[n_files:]
train_paths = train_paths[:n_files]
# Force dtypes for HugeCTR usage
dict_dtypes = {}
for col in cats:
dict_dtypes[col] = np.int64
for col in conts:
dict_dtypes[col] = np.float32
for col in labels:
dict_dtypes[col] = np.float32
# Use total device size to calculate args.device_limit_frac
device_size = device_mem_size(kind="total")
device_limit = int(device_limit_frac * device_size)
device_pool_size = int(device_pool_frac * device_size)
part_size = int(part_mem_frac * device_size)
# Check if any device memory is already occupied
for dev in devices.split(","):
fmem = _pynvml_mem_size(kind="free", index=int(dev))
used = (device_size - fmem) / 1e9
if used > 1.0:
warnings.warn(
f"BEWARE - {used} GB is already occupied on device {int(dev)}!"
)
# Setup dask cluster and perform ETL
with managed_client(dask_workdir, devices, device_limit,
protocol) as client:
# Setup RMM pool
if device_pool_frac > 0.01:
setup_rmm_pool(client, device_pool_size)
# Define Dask NVTabular "Workflow"
cont_features = conts >> ops.FillMissing() >> ops.Clip(
min_value=0) >> ops.LogOp()
cat_features = cats >> ops.Categorify(out_path=stats_path,
max_size=10000000)
workflow = Workflow(cat_features + cont_features + labels,
client=client)
train_dataset = Dataset(train_paths,
engine="parquet",
part_size=part_size)
valid_dataset = Dataset(valid_paths,
engine="parquet",
part_size=part_size)
workflow.fit(train_dataset)
workflow.transform(train_dataset).to_parquet(
output_path=output_train_dir,
shuffle=nvt_io.Shuffle.PER_WORKER,
dtypes=dict_dtypes,
cats=cats,
conts=conts,
labels=labels,
out_files_per_proc=out_files_per_proc,
)
workflow.transform(valid_dataset).to_parquet(
output_path=output_valid_dir,
shuffle=nvt_io.Shuffle.PER_WORKER,
dtypes=dict_dtypes,
cats=cats,
conts=conts,
labels=labels,
out_files_per_proc=out_files_per_proc,
)
workflow.save(os.path.join(output_path, "workflow"))
return workflow