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