Silent failure with --runner=DataflowRunner

[feelingsonice]

System information

• Environment: Google Colab, Vertex AI (KubeflowV2DagRunner)
• TensorFlow version: 2.8.0
• TFX Version: 1.7.0
• Python version: 3.7.12

---

Here's my preprocessing_fn (redacted for clarity):

_FEATURES = [# list of str
]
_SPECIAL_IMPUTE = {
 'special_foo': 1,
}
HOURS = [1, 2, 3, 4]
TABLE_KEYS = {
 'XXX': ['XXX_1', 'XXX_2', 'XXX_3'],
 'YYY': ['YYY_1', 'YYY_2', 'YYY_3'],
}

@tf.function
def _divide(a, b):
 return tf.math.divide_no_nan(tf.cast(a, tf.float32), tf.cast(b, tf.float32))

def preprocessing_fn(inputs):
 x = {}

 for name, tensor in sorted(inputs.items()):
 if tensor.dtype == tf.bool:
 tensor = tf.cast(tensor, tf.int64)

 if isinstance(tensor, tf.sparse.SparseTensor):
 default_value = '' if tensor.dtype == tf.string else 0
 tensor = tft.sparse_tensor_to_dense_with_shape(tensor, [None, 1], default_value)
 
 x[name] = tensor

 x['foo'] = _divide((x['foo1'] - x['foo2']), x['foo_denom'])
 x['bar'] = tf.cast(x['bar'] > 0, tf.int64)

 for hour in HOURS:
 total = tf.constant(0, dtype=tf.int64)
 for device_type in DEVICE_TYPES.keys():
 total = total + x[f'some_device_{device_type}_{hour}h']

 # one hot encode categorical values
 for name, keys in TABLE_KEYS.items():
 with tf.init_scope():
 initializer = tf.lookup.KeyValueTensorInitializer(
 tf.constant(keys), 
 tf.constant([i for i in range(len(keys))]))
 table = tf.lookup.StaticHashTable(initializer, default_value=-1)
 
 indices = table.lookup(tf.squeeze(x[name], axis=1))
 one_hot = tf.one_hot(indices, len(keys), dtype=tf.int64)

 for i, _tensor in enumerate(tf.split(one_hot, num_or_size_splits=len(keys), axis=1)):
 x[f'{name}_{keys[i]}'] = _tensor

 return {name: tft.scale_to_0_1(x[name]) for name in _FEATURES}

Here's the beam_pipeline_args:

BIG_QUERY_WITH_DIRECT_RUNNER_BEAM_PIPELINE_ARGS = [
 '--project=' + GOOGLE_CLOUD_PROJECT,
 '--temp_location=' + os.path.join('gs://', GCS_BUCKET_NAME, 'tmp'),
 '--runner=DataflowRunner',
 '--region=us-central1',
 '--experiments=upload_graph', # must be enabled, otherwise fails with 413
 '--dataflow_service_options=enable_prime',
 '--autoscaling_algorithm=THROUGHPUT_BASED',
 ]

Not sure if related but with the above preprocessing_fn, my transform first failed with the error:

RuntimeError: The order of analyzers in your `preprocessing_fn` appears to be non-deterministic. This can be fixed either by changing your `preprocessing_fn` such that tf.Transform analyzers are encountered in a deterministic order or by passing a unique name to each analyzer API call.

I then added names to the tft.scale_to_0_1 analyzers:

return {name: tft.scale_to_0_1(x[name], name=f'{name}_scale_to_0_1') for name in _FEATURES}

After which my transform just silently failed without logs (see first screenshot). I check the worker logs but there's nothing substantial, only warnings (see second screenshot).

It's worth noting that I have the enable_prime flag.

[pindinagesh]

Hi @bli00

Could you please take a look at this issue and see if it helps in resolving your issue? Also you can refer to this doc which discusses about the error you mentioned. Thanks!

[feelingsonice]

@pindinagesh After I removed the --dataflow_service_options=enable_prime flag it worked fine. Are you guys planning to support dataflow prime?

[bharatjetti]

Hi @feelingsonice Thank you so much for showing interest! We truly appreciate your enthusiasm and support.
We regret to inform you that we are not accepting new implementation requests or feature additions at this time. Please feel free to fork the repository and submit pull requests if you have any suggestions for enhancements, problem fixes, or ideas. We’d be more than happy to review and collaborate! If you would like to contribute , please follow our CONTRIBUTION GUIDELINES here and submit a pull request.
Thank you again for your interest, and we look forward to seeing your contributions!