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_schema_write_read_dataset(tmpdir, dataset, engine):
cat_names = ["name-cat", "name-string"
] if engine == "parquet" else ["name-string"]
cont_names = ["x", "y", "id"]
label_name = ["label"]
norms = ops.Normalize()
cat_features = cat_names >> ops.Categorify(cat_cache="host")
cont_features = cont_names >> ops.FillMissing() >> ops.Clip(
min_value=0) >> ops.LogOp >> norms
workflow = Workflow(cat_features + cont_features + label_name)
workflow.fit(dataset)
workflow.transform(dataset).to_parquet(
tmpdir,
out_files_per_proc=10,
)
schema_path = Path(tmpdir)
proto_schema = PbTxt_SchemaWriter._read(schema_path / "schema.pbtxt")
new_dataset = Dataset(glob.glob(str(tmpdir) + "/*.parquet"))
assert """name: "name-cat"\n min: 0\n max: 27\n""" in str(
proto_schema)
assert new_dataset.schema == workflow.output_schema
def test_gpu_workflow(tmpdir, df, dataset, gpu_memory_frac, engine, dump):
cat_names = ["name-cat", "name-string"
] if engine == "parquet" else ["name-string"]
cont_names = ["x", "y", "id"]
label_name = ["label"]
norms = ops.Normalize()
conts = cont_names >> ops.FillMissing() >> ops.Clip(min_value=0) >> norms
cats = cat_names >> ops.Categorify()
workflow = nvt.Workflow(conts + cats + label_name)
workflow.fit(dataset)
if dump:
workflow_dir = os.path.join(tmpdir, "workflow")
workflow.save(workflow_dir)
workflow = None
workflow = Workflow.load(workflow_dir)
def get_norms(tar: cudf.Series):
gdf = tar.fillna(0)
gdf = gdf * (gdf >= 0).astype("int")
return gdf
assert math.isclose(get_norms(df.x).mean(), norms.means["x"], rel_tol=1e-4)
assert math.isclose(get_norms(df.y).mean(), norms.means["y"], rel_tol=1e-4)
assert math.isclose(get_norms(df.x).std(), norms.stds["x"], rel_tol=1e-3)
assert math.isclose(get_norms(df.y).std(), norms.stds["y"], rel_tol=1e-3)
# Check that categories match
if engine == "parquet":
cats_expected0 = df["name-cat"].unique().values_host
cats0 = get_cats(workflow, "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(workflow, "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
workflow.transform(dataset).to_parquet(
output_path=tmpdir,
out_files_per_proc=10,
shuffle=nvt.io.Shuffle.PER_PARTITION)
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_fit_schema_works_when_subtracting_column_names():
schema = Schema(["x", "y", "id"])
cont_features = (ColumnSelector(
["x", "y"]) >> ops.FillMissing() >> ops.Clip(min_value=0) >> ops.LogOp
>> ops.Normalize() >> ops.Rename(postfix="_renamed"))
workflow1 = Workflow(cont_features - "y_renamed")
workflow1.fit_schema(schema)
assert workflow1.output_schema.column_names == ["x_renamed"]
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_names = ["cont1", "cont2"]
cat_names = ["cat1", "cat2"]
label_name = ["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)
cat_features = cat_names >> ops.Categorify()
cont_features = cont_names >> ops.FillMissing() >> ops.Clip(
min_value=0) >> ops.LogOp
workflow = Workflow(cat_features + cont_features + label_name,
client=client if use_client else None)
workflow.fit(dataset)
# Force dtypes
dict_dtypes = {}
for col in cont_names:
dict_dtypes[col] = np.float32
for col in cat_names:
dict_dtypes[col] = np.float32
for col in label_name:
dict_dtypes[col] = np.int64
workflow.transform(dataset).to_parquet(
# apply_offline=apply_offline, Not any more?
# record_stats=apply_offline, Not any more?
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_fit_schema():
schema = Schema(["x", "y", "id"])
cont_features = (ColumnSelector(schema.column_names) >> ops.FillMissing()
>> ops.Clip(min_value=0) >> ops.LogOp >> ops.Normalize()
>> ops.Rename(postfix="_renamed"))
workflow = Workflow(cont_features)
workflow.fit_schema(schema)
assert workflow.output_schema.column_names == [
"x_renamed", "y_renamed", "id_renamed"
]
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_grab_additional_input_columns(dataset, engine):
schema = Schema(["x", "y"])
node1 = ["x"] >> ops.FillMissing()
node2 = node1 >> ops.Clip(min_value=0)
add_node = node2 + ["y"]
workflow = Workflow(add_node).fit_schema(schema)
output_df = workflow.transform(dataset).to_ddf().compute()
assert len(workflow.output_node.input_columns.names) == 2
assert workflow.output_node.input_columns.names == ["x", "y"]
assert len(workflow.output_node.output_columns.names) == 2
assert workflow.output_node.output_columns.names == ["x", "y"]
assert len(output_df.columns) == 2
assert output_df.columns.tolist() == ["x", "y"]
import nvtabular as nvt
from nvtabular import ColumnSchema, ColumnSelector, Schema, dispatch, ops
from nvtabular.dispatch import HAS_GPU
@pytest.mark.parametrize("properties", [{}, {"p1": "1"}])
@pytest.mark.parametrize("tags", [[], ["TAG1", "TAG2"]])
@pytest.mark.parametrize(
"op",
[
ops.Bucketize([1]),
ops.Rename(postfix="_trim"),
ops.Categorify(),
ops.Categorify(encode_type="combo"),
ops.Clip(0),
ops.DifferenceLag("1"),
ops.FillMissing(),
ops.Groupby(["1"]),
ops.HashBucket(1),
ops.HashedCross(1),
ops.JoinGroupby(["1"]),
ops.ListSlice(0),
ops.LogOp(),
ops.Normalize(),
ops.TargetEncoding(["1"]),
ops.AddMetadata(tags=["excellent"], properties={"domain": {"min": 0, "max": 20}}),
ops.ValueCount(),
],
)
@pytest.mark.parametrize("selection", [["1"], ["2", "3"], ["1", "2", "3", "4"]])
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.FillMissing(), ops.Clip(min_value=0, 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; Clip, 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=nvt.io.Shuffle.PER_PARTITION, 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_gpu_workflow(tmpdir, client, df, dataset, gpu_memory_frac, engine, dump):
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()
config["FE"]["continuous"] = [ops.FillMissing(), ops.Clip(min_value=0)]
config["PP"]["continuous"] = [[ops.FillMissing(), ops.Clip(min_value=0), 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):
gdf = tar.fillna(0)
gdf = gdf * (gdf >= 0).astype("int")
return gdf
assert math.isclose(get_norms(df.x).mean(), processor.stats["means"]["x"], rel_tol=1e-4)
assert math.isclose(get_norms(df.y).mean(), processor.stats["means"]["y"], rel_tol=1e-4)
assert math.isclose(get_norms(df.x).std(), processor.stats["stds"]["x"], rel_tol=1e-3)
assert math.isclose(get_norms(df.y).std(), processor.stats["stds"]["y"], rel_tol=1e-3)
# 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=nvt.io.Shuffle.PER_PARTITION, 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_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 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_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_cpu_workflow(tmpdir, df, dataset, cpu, engine, dump):
# Make sure we are in cpu formats
if cudf and isinstance(df, cudf.DataFrame):
df = df.to_pandas()
if cpu:
dataset.to_cpu()
cat_names = ["name-cat", "name-string"] if engine == "parquet" else ["name-string"]
cont_names = ["x", "y", "id"]
label_name = ["label"]
norms = ops.Normalize()
conts = cont_names >> ops.FillMissing() >> ops.Clip(min_value=0) >> norms
cats = cat_names >> ops.Categorify()
workflow = nvt.Workflow(conts + cats + label_name)
workflow.fit(dataset)
if dump:
workflow_dir = os.path.join(tmpdir, "workflow")
workflow.save(workflow_dir)
workflow = None
workflow = Workflow.load(workflow_dir)
def get_norms(tar: pd.Series):
df = tar.fillna(0)
df = df * (df >= 0).astype("int")
return df
assert math.isclose(get_norms(df.x).mean(), norms.means["x"], rel_tol=1e-4)
assert math.isclose(get_norms(df.y).mean(), norms.means["y"], rel_tol=1e-4)
assert math.isclose(get_norms(df.x).std(), norms.stds["x"], rel_tol=1e-3)
assert math.isclose(get_norms(df.y).std(), norms.stds["y"], rel_tol=1e-3)
# Check that categories match
if engine == "parquet":
cats_expected0 = df["name-cat"].unique()
cats0 = get_cats(workflow, "name-cat", cpu=True)
# adding the None entry as a string because of move from gpu
assert cats0.tolist() == [None] + sorted(cats_expected0.tolist())
cats_expected1 = df["name-string"].unique()
cats1 = get_cats(workflow, "name-string", cpu=True)
# adding the None entry as a string because of move from gpu
assert cats1.tolist() == [None] + sorted(cats_expected1.tolist())
# Write to new "shuffled" and "processed" dataset
workflow.transform(dataset).to_parquet(
output_path=tmpdir, out_files_per_proc=10, shuffle=nvt.io.Shuffle.PER_PARTITION
)
dataset_2 = Dataset(glob.glob(str(tmpdir) + "/*.parquet"), cpu=cpu)
df_pp = pd.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)
metadata = pq.read_metadata(str(tmpdir) + "/_metadata")
assert metadata.num_rows == len(df_pp)
def processing(
self,
X_pd,
y_names=[],
encode_categor_type=None,
#'categorify', 'onehotencoding',
outliers_detection_technique=None,
#'iqr_proximity_rule', 'gaussian_approximation','quantiles'
fill_with_value=None,
#'extreme_values', 'zeros','mean-median'
targetencoding=False,
file_path=None,
):
X = dd.from_pandas(X_pd, npartitions=self.n_gpus)
X = X.replace(np.nan, None)
try:
self.time_columns
except AttributeError:
try:
self.initialize_types(
X,
n_unique_val_th=n_unique_val_th_,
categor_columns_keep=categor_columns_keep_,
numer_columns_keep=numer_columns_keep_)
except NameError:
self.initialize_types(X)
workflow = nvt.Workflow(cat_names=self.categor_columns,
cont_names=self.numer_columns,
label_name=y_names,
client=self.client)
# Operators: https://nvidia.github.io/NVTabular/main/api/ops/index.html
# Categorify https://nvidia.github.io/NVTabular/main/api/ops/categorify.html
if encode_categor_type == 'categorify':
if len(self.categor_columns) != 0:
workflow.add_preprocess(
ops.Categorify(columns=self.categor_columns,
out_path='./'))
if encode_categor_type == 'onehotencoding':
#OneHotEncoder().get_feature_names(input_features=<list of features encoded>) does not work
#lengths=True - chunk sizes can be computed
for column in self.categor_columns:
#X[column] = X[column].astype(str)
X_cat_encoded = OneHotEncoder().fit_transform(
X[column].to_dask_array(lengths=True).reshape(-1, 1))
uvs = X[column].unique().compute().values
X = X.drop([column], axis=1)
X_cat_encoded = dd.from_array(
X_cat_encoded.compute().todense())
X_cat_encoded.columns = [
column + '_{}'.format(uv) for uv in uvs
]
X = dd.concat([X, X_cat_encoded], axis=1)
X = X.repartition(npartitions=2)
for column in X.columns:
if any(str(column)[-4:] == t
for t in ['_nan', 'None']): # What else?
X = X.drop([column], axis=1)
self.initialize_types(X)
print('Retyping:', self.initialize_types(X))
# Reinitialize workflow
workflow = nvt.Workflow(cat_names=self.categor_columns,
cont_names=self.numer_columns,
label_name=y_names,
client=self.client)
# OutlDetect https://nvidia.github.io/NVTabular/main/api/ops/clip.html
if (len(self.numer_columns) != 0) and (outliers_detection_technique !=
None):
lower, upper = self.outldetect(outliers_detection_technique,
X[self.numer_columns])
for i in range(len(self.numer_columns)):
logging.info(
f'column: {self.numer_columns[i]}, lower: {lower[i]}, upper: {upper[i]}'
)
print(
f'column: {self.numer_columns[i]}, lower: {lower[i]}, upper: {upper[i]}'
)
workflow.add_preprocess(
ops.Clip(min_value=lower[i],
max_value=upper[i],
columns=[self.numer_columns[i]]))
# FillMissing https://nvidia.github.io/NVTabular/main/api/ops/fillmissing.html
if fill_with_value == 'zeros':
workflow.add_preprocess(
ops.FillMissing(fill_val=0,
columns=self.categor_columns +
self.numer_columns))
if fill_with_value == 'extreme_values':
extrim_values = {}
if len(self.numer_columns) != 0:
extrim_values.update(
self.extrvalsdetect(X[self.numer_columns],
'numer_columns'))
if len(self.categor_columns) != 0:
extrim_values.update(
self.extrvalsdetect(X[self.categor_columns],
'categor_columns'))
logging.info(f'extrim_values: {extrim_values}')
output = open('extrim_values', 'wb')
pickle.dump(extrim_values, output)
output.close()
for fill_val, column in zip(list(extrim_values.values()),
list(extrim_values.keys())):
workflow.add_preprocess(
ops.FillMissing(fill_val=fill_val, columns=[column]))
if fill_with_value == 'mean-median':
if len(self.categor_columns) != 0:
workflow.add_preprocess(
ops.FillMedian(columns=self.categor_columns,
preprocessing=True,
replace=True))
if len(self.numer_columns) != 0:
means = list(
dd.from_pandas(
X[self.numer_columns],
npartitions=self.n_gpus).mean().compute().values)
for fill_val, column in zip(means, self.numer_columns):
workflow.add_preprocess(
ops.FillMissing(fill_val=fill_val, columns=[column]))
if targetencoding:
#https://nvidia.github.io/NVTabular/main/api/ops/targetencoding.html
if len(self.y_names) != 0:
if len(self.cat_groups) == 0:
print(
'\n Target encoding will be applied to all categorical columns'
)
workflow.add_preprocess(
ops.TargetEncoding(cat_groups=self.categor_columns,
cont_target=self.y_names))
else:
workflow.add_preprocess(
ops.TargetEncoding(cat_groups=self.cat_groups,
cont_target=self.y_names))
#-----------------------------------------------------------------------------------------
workflow.finalize()
dataset = nvt.Dataset(X)
tmp_output_path = "./parquet_data_tmp"
workflow.apply(
dataset,
output_format="parquet",
output_path=tmp_output_path,
shuffle=Shuffle.PER_WORKER, # Shuffle algorithm
out_files_per_proc=1, # Number of output files per worker
)
files = glob.glob(tmp_output_path + "/*.parquet")
X_final = cudf.read_parquet(files[0])
for i in range(1, len(files)):
X_final = X_final.append(cudf.read_parquet(files[i]))
# Delete temporary files
shutil.rmtree(tmp_output_path, ignore_errors=True)
# if len(self.rest_col_names) != 0:
# print(1)
# X_final = pd.concat([X_final.to_pandas(), X_pd[self.rest_col_names]], axis=1)
if file_path is not None:
X_final.to_csv(file_path, index=False)
return X_final
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_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)
def test_gpu_workflow_api(tmpdir, client, df, dataset, gpu_memory_frac, engine,
dump, use_client):
cat_names = ["name-cat", "name-string"
] if engine == "parquet" else ["name-string"]
cont_names = ["x", "y", "id"]
label_name = ["label"]
norms = ops.Normalize()
cat_features = cat_names >> ops.Categorify(cat_cache="host")
cont_features = cont_names >> ops.FillMissing() >> ops.Clip(
min_value=0) >> ops.LogOp >> norms
workflow = Workflow(cat_features + cont_features + label_name,
client=client if use_client else None)
workflow.fit(dataset)
if dump:
workflow_dir = os.path.join(tmpdir, "workflow")
workflow.save(workflow_dir)
workflow = None
workflow = Workflow.load(workflow_dir,
client=client if use_client else None)
def get_norms(tar):
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; Clip, Log
assert math.isclose(get_norms(df.y).mean(), norms.means["y"], rel_tol=1e-1)
assert math.isclose(get_norms(df.y).std(), norms.stds["y"], rel_tol=1e-1)
assert math.isclose(get_norms(df.x).mean(), norms.means["x"], rel_tol=1e-1)
assert math.isclose(get_norms(df.x).std(), norms.stds["x"], rel_tol=1e-1)
# Check that categories match
if engine == "parquet":
cats_expected0 = df["name-cat"].unique(
).values_host if HAS_GPU else df["name-cat"].unique()
cats0 = get_cats(workflow, "name-cat")
# adding the None entry as a string because of move from gpu
assert all(cat in [None] + sorted(cats_expected0.tolist())
for cat in cats0.tolist())
assert len(cats0.tolist()) == len(cats_expected0.tolist() + [None])
if HAS_GPU:
cats_expected1 = (df["name-string"].unique().values_host
if HAS_GPU else df["name-string"].unique())
else:
cats_expected1 = df["name-string"].unique()
cats1 = get_cats(workflow, "name-string")
# adding the None entry as a string because of move from gpu
assert all(cat in [None] + sorted(cats_expected1.tolist())
for cat in cats1.tolist())
assert len(cats1.tolist()) == len(cats_expected1.tolist() + [None])
# Write to new "shuffled" and "processed" dataset
workflow.transform(dataset).to_parquet(
tmpdir,
out_files_per_proc=10,
shuffle=nvt.io.Shuffle.PER_PARTITION,
)
dataset_2 = Dataset(glob.glob(str(tmpdir) + "/*.parquet"),
part_mem_fraction=gpu_memory_frac)
df_pp = nvt.dispatch._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 = nvt.dispatch._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
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)
assert workflow1.output_schema.column_names == [
"TE_x_cost_renamed", "TE_y_cost_renamed"
]
# initial column selector works with tags
# filter within the workflow by tags
# test tags correct at output
@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"])
