Data Preprocessing with tf.data.service not working

Question

I use TPU in colab:

try: # detect TPUs
  tpu = tf.distribute.cluster_resolver.TPUClusterResolver.connect() # TPU detection
  strategy = tf.distribute.TPUStrategy(tpu)
except ValueError: # detect GPUs
  strategy = tf.distribute.MirroredStrategy() # for GPU or multi-GPU machines

Then start DispatchServer and WorkerServer:

d_config = tf.data.experimental.service.DispatcherConfig(port=5000)
dispatcher = tf.data.experimental.service.DispatchServer()  #d_config
#w_config = tf.data.experimental.service.WorkerConfig(port=5001,dispatcher_address=dispatcher.target.split("://")[1], worker_address='localhost:5001')
w_config =tf.data.experimental.service.WorkerConfig(dispatcher_address=dispatcher.target.split("://")[1])
worker = tf.data.experimental.service.WorkerServer(w_config)
dispatcher.start()

And use it in apply in dataset loaded through TFRecordDataset:

return augmented.repeat().apply(tf.data.experimental.service.distribute(processing_mode="parallel_epochs", service=dispatcher.target)).batch(batch_size).prefetch(AUTO) #repeat().shuffle(2048) #augmented

But, when create model with strategy.scope() it hangs in an infinite loop:

with strategy.scope(): # creating the model in the TPUStrategy scope places the model on the TPU
   model = create_model()
model.summary()

What am i doing wrong? I want to prepare a dataset outside the TPU in 40 CPUs in colab.