def test_workflow_move_saved(tmpdir):
raw = """US>SC>519 US>CA>807 US>MI>505 US>CA>510 CA>NB US>CA>534""".split()
data = cudf.DataFrame({"geo": raw})
geo_location = ColumnGroup(["geo"])
state = geo_location >> (lambda col: col.str.slice(0, 5)) >> ops.Rename(
postfix="_state")
country = geo_location >> (lambda col: col.str.slice(0, 2)) >> ops.Rename(
postfix="_country")
geo_features = state + country + geo_location >> ops.Categorify()
# create the workflow and transform the input
workflow = Workflow(geo_features)
expected = workflow.fit_transform(Dataset(data)).to_ddf().compute()
# save the workflow (including categorical mapping parquet files)
# and then verify we can load the saved workflow after moving the directory
out_path = os.path.join(tmpdir, "output", "workflow")
workflow.save(out_path)
moved_path = os.path.join(tmpdir, "output", "workflow2")
shutil.move(out_path, moved_path)
workflow2 = Workflow.load(moved_path)
# also check that when transforming our input we get the same results after loading
transformed = workflow2.transform(Dataset(data)).to_ddf().compute()
assert_eq(expected, transformed)
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"]
norms = ops.Normalize()
cat_features = cat_names >> ops.Categorify()
if replace:
cont_features = cont_names >> ops.FillMissing() >> ops.LogOp >> norms
else:
fillmissing_logop = (cont_names >> ops.FillMissing() >> ops.LogOp >>
ops.Rename(postfix="_FillMissing_1_LogOp_1"))
cont_features = cont_names + fillmissing_logop >> norms
workflow = Workflow(cat_features + cont_features + label_name,
client=client)
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)
def get_norms(tar):
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; Clip, Log
concat_ops = "_FillMissing_1_LogOp_1"
if replace:
concat_ops = ""
assert math.isclose(get_norms(df.x).mean(),
norms.means["x" + concat_ops],
rel_tol=1e-1)
assert math.isclose(get_norms(df.y).mean(),
norms.means["y" + concat_ops],
rel_tol=1e-1)
assert math.isclose(get_norms(df.x).std(),
norms.stds["x" + concat_ops],
rel_tol=1e-1)
assert math.isclose(get_norms(df.y).std(),
norms.stds["y" + concat_ops],
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])
cats_expected1 = (df["name-string"].unique().values_host
if HAS_GPU else 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 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: 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
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
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(
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 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
