lint

Author: edknv

merlin/dataloader/utils/tf/tf_trainer.py

@@ -1,13 +1,9 @@
 # External dependencies
 import argparse
-import glob
 import logging
 import os
 
 from merlin.core.compat import cupy, numpy
-from merlin.io import Dataset
-from nvtabular.ops import TagAsItemID, TagAsUserID, JoinExternal
-from merlin.core.dispatch import get_lib
 from merlin.schema import Tags
 
 # we can control how much memory to give tensorflow with this environment variable
@@ -76,8 +72,8 @@ def seed_fn():
 EMBEDDING_TABLE_SHAPES, MH_EMBEDDING_TABLE_SHAPES = nvt.ops.get_embedding_sizes(proc)
 EMBEDDING_TABLE_SHAPES.update(MH_EMBEDDING_TABLE_SHAPES)
 
-train_ds = nvt.Dataset(f'{BASE_DIR}/train', engine='parquet', dtypes={'rating': xp.int8})
-train_ds.schema = train_ds.schema.remove_col('genres')
+train_ds = nvt.Dataset(f"{BASE_DIR}/train", engine="parquet", dtypes={"rating": xp.int8})
+train_ds.schema = train_ds.schema.remove_col("genres")
 
 target_column = train_ds.schema.select_by_tag(Tags.TARGET).column_names[0]
 

tests/unit/dataloader/test_tf_dataloader.py

@@ -536,10 +536,6 @@ def test_multigpu_partitioning(dataset, batch_size, global_rank):
 
 
 @pytest.mark.multigpu
-@pytest.mark.skipif(
-    os.environ.get("NR_USER") is not None,
-    reason="not working correctly in ci environment",
-)
 @pytest.mark.skipif(importlib.util.find_spec("horovod") is None, reason="needs horovod")
 @pytest.mark.skipif(
     HAS_GPU and cupy and cupy.cuda.runtime.getDeviceCount() <= 1,
@@ -722,7 +718,7 @@ def test_wrong_batch_size_raises_warning():
             _ = tf_loader(dataset, batch_size=batch_size)
 
     for power in range(4, 10):
-        batch_size = 2**power
+        batch_size = 2 ** power
         # warning not raised for power of two
         with warnings.catch_warnings():
             warnings.simplefilter("error")

tox.ini

@@ -82,7 +82,6 @@ commands =
 passenv =
     CUDA_VISIBLE_DEVICES
     OPAL_PREFIX
-    NR_USER
 sitepackages=true
 ; Runs in: 1GPU Github Actions runners.
 ; Runs GPU-based tests.