[Python] Dataset from generator in tensorflow error GetNextFromShard

I have a train step for a Vertex AI pipeline that uses a template code for setting up distributed training. When fitting any model I get a message in the first epoch indicating the following message:

025-07-01 17:55:35.569062: I tensorflow/core/framework/local_rendezvous.cc:407] Local rendezvous is aborting with status: CANCELLED: GetNextFromShard was cancelled [[{{node MultiDeviceIteratorGetNextFromShard}}]] [[RemoteCall]] [type.googleapis.com/tensorflow.DerivedStatus='']


I have tried using ds = strategy.experimental_distribute_dataset(ds) but this doesnt solve this issue. despite this message the model is being trained, and doesnt appear in other epochs is this really an issue?

Sharing the code below:

# Single Machine, single compute device
if args.distribute == 'single':
if tf.test.is_gpu_available():
strategy = tf.distribute.OneDeviceStrategy(device="/gpu:0")
else:
strategy = tf.distribute.OneDeviceStrategy(device="/cpu:0")
# Single Machine, multiple compute device
elif args.distribute == 'mirror':
strategy = tf.distribute.MirroredStrategy()
# Multiple Machine, multiple compute device
elif args.distribute == 'multi':
strategy = tf.distribute.experimental.MultiWorkerMirroredStrategy()



def get_data_set(generator_fn)->tuple[tf.data.Dataset,int]:

data_tmp=generator_fn()

for sample_X_train,sample_y_train in data_tmp:
break
total=sum(1 for _ in data_tmp)+1


ds= tf.data.Dataset.from_generator(
generator_fn,
output_signature=(
tf.TensorSpec(shape=(sample_X_train.shape), dtype=tf.float16),
tf.TensorSpec(shape=(sample_y_train.shape), dtype=tf.float16)
)
)
ds=ds.batch(BATCH_SIZE).repeat()
#ds = strategy.experimental_distribute_dataset(ds)
return ds,total


and fit step is:

with strategy.scope():
train_generator,train_samples=get_data_set(lambda:multi_window.train)
val_generator,val_samples=get_data_set(lambda:multi_window.val)
print(train_samples,val_samples)
train_steps = train_samples // BATCH_SIZE
val_steps = val_samples //BATCH_SIZE

model.fit(
train_generator,
validation_data=val_generator,
epochs=epochs,
validation_steps=val_steps,
steps_per_epoch=train_steps,
)



Tried strategy.experimental_distribute_dataset(ds) unsuccessfully

Continue reading...