title: Sagemaker Tensorflow TF Dataset Never Has Available Data for Training
body: <p>I have stripped down a TF script I am trying to run on AWS Sagemaker and I cant figure out why I can train a Keras model using a dataset as input locally but as soon as I run it with the Sagemaker framework I get errors that I dont have sufficient data and training doesnt start.</p>
  <p><strong>Code that processes the dataset:</strong></p>
  <pre><code>import tensorflow as tf
  import boto3
  import re 
  
  #parse the protobuffer
  def _parse_function(proto):
      # define tfrecord again. 
      keys_to_features = {'target': tf.io.FixedLenFeature([], tf.float32),
                          'feature1': tf.io.FixedLenFeature([], tf.float32),
                          'feature2': tf.io.FixedLenFeature([], tf.float32),
                          'feature3': tf.io.FixedLenFeature([], tf.float32),
                          'feature4': tf.io.FixedLenFeature([], tf.float32)
                          }
     
      # Load one example
      parsed_features = tf.io.parse_example(proto, keys_to_features)
      
      return(parsed_features)
  
  
  # means and standard deviations from the training set
  hold_meanstd={
      'Feature1':[12178.24,227789.9],
      'Feature2':[0.002085435,0.01342237],
      'Feature3':[0.002572172,0.01819553],
      'Feature4':[90480.89,352610]
      }
  
  def normalize(example):
          example['Feature1']=(example['Feature1']-hold_meanstd['Feature1'][0])/hold_meanstd['Feature1'][1]
          example['Feature2']=(example['Feature2']-hold_meanstd['Feature2'][0])/hold_meanstd['Feature2'][1]
          example['Feature3']=(example['Feature3']-hold_meanstd['Feature3'][0])/hold_meanstd['Feature3'][1]
          example['Feature4']=(example['Feature4']-hold_meanstd['Feature4'][0])/hold_meanstd['Feature4'][1]
          
          label=example.pop('target')
          
          # keras needs:  Should return a tuple of either (inputs, targets) or (inputs, targets, sample_weights).
          return([example['Feature1'],example['Feature2'],example['Feature3'],example['Feature4']],label)
  
      
  # get tfrecords shards
  def get_tfrecords(tfr_dir):
      
      s3 = boto3.resource('s3')
      _bucket = s3.Bucket(‘bucketname’)
      file_lst=[]
  
      prefix = tfr_dir[re.search(r&quot;[a-z]/&quot;, tfr_dir).start()+2:]
      for i, obj_summary in enumerate(_bucket.objects.filter(Prefix=prefix)):
          if not obj_summary.key.endswith('SUCCESS'):
              file_lst.append('S3-PATH/'+obj_summary.key)
      
      return(file_lst)
  
  
  # load and return a dataset    
  def load_ds(tfr_dir,batchsize):
      file_list = get_tfrecords(tfr_dir)
      dataset = tf.data.TFRecordDataset(file_list)
      dataset = dataset.map(_parse_function)
      dataset = dataset.map(normalize)
      dataset = dataset.shuffle(100)
      dataset= dataset.repeat()
      dataset = dataset.batch(batchsize,drop_remainder=True)
      
      
      return(dataset)
  </code></pre>
  <p>I create a simple Keras model (not shown, but its basic for testing purposes). Finally, I run the following. There are 5,000 records in tf records. This works fine.</p>
  <pre><code>tf_model = create_keras_model(0.001,4, None)
  ds = load_ds('PATH_TO_S3_tfrecords', 16)
  
  tf_model.fit(ds, 
               epochs = 5, 
               steps_per_epoch= 5000 // 16 
              )
  </code></pre>
  <p>When I include the same code as above, in Sagemaker:</p>
  <p><em>train_fn.py includes</em></p>
  <pre><code>def train_model(args):
      
      train_ds = utilities.load_ds(args.train_dir, args.batchsize)
      tf_model = model.create_keras_model(args.learning_rate,args.inputdim, None)
  
      tf_model.fit(train_ds, 
                       steps_per_epoch= 5000 // 16, 
                       epochs = 5
                  )
  
  
  
  if __name__ == '__main__':
      
      args = get_args()
      train_model(args)
  </code></pre>
  <p>and run:</p>
  <pre><code>tf_estimator = sm_tf(
      entry_point=&quot;train_fn.py&quot;,
      source_dir = './code_test',
      role=role,
      instance_count=1, 
      instance_type = &quot;ml.c4.8xlarge&quot;,
      framework_version = &quot;2.8&quot;,
      py_version = &quot;py39&quot;,
      script_mode= True,
      
      
      sagemaker_session = sagemaker_session,
      hyperparameters = hyperparams
  
  )
  
  
  
  tf_estimator.fit({'train'  : training_data_uri},                 
                   job_name = job_name
                   
                  )
  </code></pre>
  <p>Error:</p>
  <pre><code>WARNING:tensorflow:Your input ran out of data; interrupting training. Make sure that your dataset or generator can generate at least `steps_per_epoch * epochs` batches (in this case, 1560 batches). You may need to use the repeat() function when building your dataset.
  WARNING:tensorflow:Your input ran out of data; interrupting training. Make sure that your dataset or generator can generate at least `steps_per_epoch * epochs` batches (in this case, 1560 batches). You may need to use the repeat() function when building your dataset.
  Traceback (most recent call last):
    File &quot;/opt/ml/code/train_fn.py&quot;, line 68, in &lt;module&gt;
      train_model(args)
    File &quot;/opt/ml/code/train_fn.py&quot;, line 58, in train_model
  tf_model.fit(train_ds, 
    File &quot;/usr/local/lib/python3.9/site-packages/keras/utils/traceback_utils.py&quot;, line 67, in error_handler
  raise e.with_traceback(filtered_tb) from None
    File &quot;/usr/local/lib/python3.9/site-packages/keras/engine/training.py&quot;, line 1395, in fit
  raise ValueError('Unexpected result of `train_function` '
  ValueError: Unexpected result of `train_function` (Empty logs). Please use `Model.compile(..., run_eagerly=True)`, or `tf.config.run_functions_eagerly(True)` for more information of where went wrong, or file a issue/bug to `tf.keras`.
  2022-05-18 12:52:22,703 sagemaker-training-toolkit ERROR    Reporting training FAILURE
  2022-05-18 12:52:22,703 sagemaker-training-toolkit ERROR    ExecuteUserScriptError:
  ExitCode 1
  </code></pre>