TFX custom config argument in trainer not working

This question is based on the TFX recommender tutorial. Please note that the code is being orchestrated by LocalDagRunner rather than run interactively in a notebook.

In the Trainer, we pass in a custom_config with the transformed ratings and movies:

trainer = tfx.components.Trainer(

The problem is that all of the outputs passed into custom_config seem to be empty. This results in errors, for example

class MovielensModel(tfrs.Model):

  def __init__(self, user_model, movie_model, tf_transform_output, movies_uri):
    self.movie_model: tf.keras.Model = movie_model
    self.user_model: tf.keras.Model = user_model

    movies_artifact = movies_uri.get()[0]

complains that movie_uri.get() is empty. The same is true for ratings. Ratings passed in through the examples parameter however are not empty (the artefact uri is available), so it seems as though this custom_config is ‘breaking things’.

I have tried debugging it but to no avail. I did notice that arguments in custom_config are serialised and deserialised, but this didn’t seem to be the cause of the problem. Does anyone know why this happens and how to resolve this?

That seems really odd. I’ve reached out to see if anyone has any ideas, but there are a couple of things you can try. You can add print statements (, etc) to inspect the outputs from the components in the custom config. And you can also set a breakpoint (maybe using pdb) and inspect. If they’re empty then it suggests that the trainer isn’t waiting for them to finish, and there’s a race condition. It will also help to insert print statements in those components to verify when they run.