KeyError : 'callable_inputs' when trying to save a Tensorflow model in S3 bucket

[hm-haitham]

Describe the bug
I'm getting KeyError : 'callable_inputs' when trying to save a Tensorflow model in S3 bucket. What's weird is that the exact same code used to work, even rolling back to previous versions doesn't solve the problem

To reproduce
parser.add_argument("--model_dir", type=str, default=os.environ.get("SM_MODEL_DIR"))

model.save(args.model_dir)

Expected behavior
The model is saved in S3 bucket

Screenshots or logs
Traceback (most recent call last): File "entrypoint_train.py", line 40, in <module> train_model(parse_args(cmd_args)[0]) File "/opt/ml/code/ciacoml/change_detector/training/train_tf_model.py", line 88, in train_model model.save(args.model_dir) File "/opt/ml/code/ciacoml/change_detector/domain/tf_model_wrapper.py", line 106, in save tf.keras.models.save_model(self.model, save_path) File "/usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/saving/save.py", line 138, in save_model signatures, options) File "/usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/saving/saved_model/save.py", line 78, in save save_lib.save(model, filepath, signatures, options
...
File "/usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/saving/saved_model/save_impl.py", line 453, in get_input_arg_value self._input_arg_name, args, kwargs, inputs_in_args=True) File "/usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/engine/base_layer.py", line 2291, in _get_call_arg_value return args_dict[arg_name] KeyError: 'callable_inputs'

System information
A description of your system. Please provide:

• SageMaker Python SDK version: 2.5.1
• Framework name (eg. PyTorch) or algorithm (eg. KMeans): Tensorflow
• Framework version: 2.3.0
• Python version: 3.7
• CPU or GPU: GPU
• Custom Docker image (Y/N): Y

[inderpartap]

I am getting the same issue as well. I am saving my model using:
model.save(os.path.join(model_dir, 'model/1'))

The estimator was created with the following arguments:

from sagemaker.tensorflow import TensorFlow

tf_estimator = TensorFlow(entry_point='autoencoder-model.py', 
                          role=role,
                          train_instance_count=1, 
                          train_instance_type='ml.m5.large',
                          image_name='763104351884.dkr.ecr.us-west-2.amazonaws.com/tensorflow-training:2.3.0-cpu-py37-ubuntu18.04',
                          script_mode=True,
                          hyperparameters={'epochs': 20},
                          subnets=['foo', 'bar'],
                          security_group_ids=['foo', 'bar'])

training_input_path   = '<s3 path>'
validation_input_path = '<s3 path>'

tf_estimator.fit({'training': training_input_path, 'validation': validation_input_path})

Here are the cloudwatch logs:

Traceback (most recent call last):
  File "autoencoder-model.py", line 214, in <module>
    model.save(os.path.join(model_dir, 'model/1'))
  File "/usr/local/lib/python3.7/site-packages/tensorflow/python/keras/engine/training.py", line 2005, in save
    signatures, options)
  File "/usr/local/lib/python3.7/site-packages/tensorflow/python/keras/saving/save.py", line 134, in save_model
    signatures, options)
  File "/usr/local/lib/python3.7/site-packages/tensorflow/python/keras/saving/saved_model/save.py", line 80, in save
    save_lib.save(model, filepath, signatures, options)
  File "/usr/local/lib/python3.7/site-packages/tensorflow/python/saved_model/save.py", line 976, in save
    obj, export_dir, signatures, options, meta_graph_def)
  File "/usr/local/lib/python3.7/site-packages/tensorflow/python/saved_model/save.py", line 1047, in _build_meta_graph
    checkpoint_graph_view)
  File "/usr/local/lib/python3.7/site-packages/tensorflow/python/saved_model/signature_serialization.py", line 75, in find_function_to_export
    functions = saveable_view.list_functions(saveable_view.root)
  File "/usr/local/lib/python3.7/site-packages/tensorflow/python/saved_model/save.py", line 145, in list_functions
    self._serialization_cache)
  File "/usr/local/lib/python3.7/site-packages/tensorflow/python/keras/engine/training.py", line 2616, in _list_functions_for_serialization
    Model, self)._list_functions_for_serialization(serialization_cache)
  File "/usr/local/lib/python3.7/site-packages/tensorflow/python/keras/engine/base_layer.py", line 3019, in _list_functions_for_serialization
    .list_functions_for_serialization(serialization_cache))
  File "/usr/local/lib/python3.7/site-packages/tensorflow/python/keras/saving/saved_model/base_serialization.py", line 87, in list_functions_for_serialization
    fns = self.functions_to_serialize(serialization_cache)
  File "/usr/local/lib/python3.7/site-packages/tensorflow/python/keras/saving/saved_model/layer_serialization.py", line 79, in functions_to_serialize
    serialization_cache).functions_to_serialize)
  File "/usr/local/lib/python3.7/site-packages/tensorflow/python/keras/saving/saved_model/layer_serialization.py", line 95, in _get_serialized_attributes
    serialization_cache)
  File "/usr/local/lib/python3.7/site-packages/tensorflow/python/keras/saving/saved_model/model_serialization.py", line 57, in _get_serialized_attributes_internal
    serialization_cache))
  File "/usr/local/lib/python3.7/site-packages/tensorflow/python/keras/saving/saved_model/layer_serialization.py", line 104, in _get_serialized_attributes_internal
    functions = save_impl.wrap_layer_functions(self.obj, serialization_cache)
  File "/usr/local/lib/python3.7/site-packages/tensorflow/python/keras/saving/saved_model/save_impl.py", line 155, in wrap_layer_functions
    original_fns = _replace_child_layer_functions(layer, serialization_cache)
  File "/usr/local/lib/python3.7/site-packages/tensorflow/python/keras/saving/saved_model/save_impl.py", line 274, in _replace_child_layer_functions
    serialization_cache).functions)
  File "/usr/local/lib/python3.7/site-packages/tensorflow/python/keras/saving/saved_model/layer_serialization.py", line 95, in _get_serialized_attributes
    serialization_cache)
  File "/usr/local/lib/python3.7/site-packages/tensorflow/python/keras/saving/saved_model/layer_serialization.py", line 104, in _get_serialized_attributes_internal
    functions = save_impl.wrap_layer_functions(self.obj, serialization_cache)
  File "/usr/local/lib/python3.7/site-packages/tensorflow/python/keras/saving/saved_model/save_impl.py", line 165, in wrap_layer_functions
    '{}_layer_call_and_return_conditional_losses'.format(layer.name))
  File "/usr/local/lib/python3.7/site-packages/tensorflow/python/keras/saving/saved_model/save_impl.py", line 505, in add_function
    self.add_trace(*self._input_signature)
  File "/usr/local/lib/python3.7/site-packages/tensorflow/python/keras/saving/saved_model/save_impl.py", line 423, in add_trace
    fn.get_concrete_function(*args, **kwargs)
  File "/usr/local/lib/python3.7/site-packages/tensorflow/python/keras/saving/saved_model/save_impl.py", line 549, in get_concrete_function
    return super(LayerCall, self).get_concrete_function(*args, **kwargs)
  File "/usr/local/lib/python3.7/site-packages/tensorflow/python/eager/def_function.py", line 1167, in get_concrete_function
    concrete = self._get_concrete_function_garbage_collected(*args, **kwargs)
  File "/usr/local/lib/python3.7/site-packages/tensorflow/python/eager/def_function.py", line 1073, in _get_concrete_function_garbage_collected
    self._initialize(args, kwargs, add_initializers_to=initializers)
  File "/usr/local/lib/python3.7/site-packages/tensorflow/python/eager/def_function.py", line 697, in _initialize
    *args, **kwds))
  File "/usr/local/lib/python3.7/site-packages/tensorflow/python/eager/function.py", line 2855, in _get_concrete_function_internal_garbage_collected
    graph_function, _, _ = self._maybe_define_function(args, kwargs)
  File "/usr/local/lib/python3.7/site-packages/tensorflow/python/eager/function.py", line 3213, in _maybe_define_function
    graph_function = self._create_graph_function(args, kwargs)
  File "/usr/local/lib/python3.7/site-packages/tensorflow/python/eager/function.py", line 3075, in _create_graph_function
    capture_by_value=self._capture_by_value),
  File "/usr/local/lib/python3.7/site-packages/tensorflow/python/framework/func_graph.py", line 986, in func_graph_from_py_func
    func_outputs = python_func(*func_args, **func_kwargs)
  File "/usr/local/lib/python3.7/site-packages/tensorflow/python/eager/def_function.py", line 600, in wrapped_fn
    return weak_wrapped_fn().__wrapped__(*args, **kwds)
  File "/usr/local/lib/python3.7/site-packages/tensorflow/python/keras/saving/saved_model/save_impl.py", line 515, in wrapper
    inputs = call_collection.get_input_arg_value(args, kwargs)
  File "/usr/local/lib/python3.7/site-packages/tensorflow/python/keras/saving/saved_model/save_impl.py", line 455, in get_input_arg_value
    self._input_arg_name, args, kwargs, inputs_in_args=True)
  File "/usr/local/lib/python3.7/site-packages/tensorflow/python/keras/engine/base_layer.py", line 2536, in _get_call_arg_value
    return args_dict[arg_name]
KeyError: 'callable_inputs'

System information
A description of your system. Please provide:

• Framework name (eg. PyTorch) or algorithm (eg. KMeans): Tensorflow (Keras Autoencoder)
• Framework version: 2.3.0
• Python version: 3.7
• CPU or GPU: CPU
• Custom Docker image (Y/N): N

[inderpartap]

FYI: seting the debugger_hook_config to False solves the issue.

tf_estimator = TensorFlow(entry_point='autoencoder-model.py', 
                          role=role,
                          train_instance_count=1, 
                          train_instance_type='ml.m5.large',
                          image_name='763104351884.dkr.ecr.us-west-2.amazonaws.com/tensorflow-training:2.3.0-cpu-py37-ubuntu18.04',
                          script_mode=True,
                          hyperparameters={'epochs': 20},
                          subnets=['foo', 'bar'],
                          security_group_ids=['foo', 'bar'],
                          debugger_hook_config=False,)

Take a look at this Stackoverflow Discussion

[aggieKevin]

The issue is resolved in smdebug 0.9.3 awslabs/sagemaker-debugger#333 . I have tested that the tf 2.3 and tf 2.2 are able to save the model successfully, regardless "debugger_hook_config=False" is set or not.

[Othmane796]

Hi team,

I came across this bug when updating my code to Tensorflow 2.3.1, given it has been fixed in smdebug 0.9.3, I just changed my docker image to a newer version which had smdebug 0.9.4, but the error was still present. I ended up forcing the uninstall of 0.9.4 and installing 0.9.3 and saving the model works now.

Docker image used : '763104351884.dkr.ecr.eu-west-1.amazonaws.com/tensorflow-training:2.3.1-gpu-py37-cu110-ubuntu18.04' taken from (https://github.com/aws/deep-learning-containers/blob/master/available_images.md)
Code launching the training job :

training_image_uri = '763104351884.dkr.ecr.eu-west-1.amazonaws.com/tensorflow-training:2.3.1-gpu-py37-cu110-ubuntu18.04'
estimator = TensorFlow(base_job_name=JOB_NAME,
                       source_dir='../code',
                       entry_point='train.py', 
                       role=role,
                       image_uri=training_image_uri,
                       hyperparameters=hyperparameters,
                       instance_count=1,
                       script_mode=True,
                       metric_definitions=[
                           {'Name': 'train loss', 'Regex': 'loss: (.*?) -'},
                           {'Name': 'train accuracy', 'Regex': ' accuracy: (.*?) -'},
                           {'Name': 'val loss', 'Regex': 'val_loss: (.*?) -'},
                           {'Name': 'val accuracy', 'Regex': 'val_accuracy: (.*?)$'}
                       ],
                       output_path=OUTPUT_PATH,
                       instance_type=INSTANCE_TYPE)

Additional libraries : Added a requirements.txt file within the source_dir folder:

tensorflow-gpu==2.3.1
tensorflow-hub==0.10.0
bert-for-tf2==0.14.7
bert-tensorflow==1.0.4
pydot==1.4.1
graphviz==0.15
boto3==1.16.17
scikit-learn==0.22
pandas==1.0.1

Error log :

Traceback (most recent call last):
  File "train.py", line 207, in <module>
    main(args)
  File "train.py", line 168, in main
    model.save(export_dir)
  File "/usr/local/lib/python3.7/site-packages/tensorflow/python/keras/engine/training.py", line 2032, in save
    signatures, options)
  File "/usr/local/lib/python3.7/site-packages/tensorflow/python/keras/saving/save.py", line 134, in save_model
    signatures, options)
  File "/usr/local/lib/python3.7/site-packages/tensorflow/python/keras/saving/saved_model/save.py", line 80, in save
    save_lib.save(model, filepath, signatures, options)
  File "/usr/local/lib/python3.7/site-packages/tensorflow/python/saved_model/save.py", line 976, in save
    obj, export_dir, signatures, options, meta_graph_def)
  File "/usr/local/lib/python3.7/site-packages/tensorflow/python/saved_model/save.py", line 1047, in _build_meta_graph
    checkpoint_graph_view)
  File "/usr/local/lib/python3.7/site-packages/tensorflow/python/saved_model/signature_serialization.py", line 75, in find_function_to_export
    functions = saveable_view.list_functions(saveable_view.root)
  File "/usr/local/lib/python3.7/site-packages/tensorflow/python/saved_model/save.py", line 145, in list_functions
    self._serialization_cache)
  File "/usr/local/lib/python3.7/site-packages/tensorflow/python/keras/engine/training.py", line 2644, in _list_functions_for_serialization
    Model, self)._list_functions_for_serialization(serialization_cache)
  File "/usr/local/lib/python3.7/site-packages/tensorflow/python/keras/engine/base_layer.py", line 3064, in _list_functions_for_serialization
    .list_functions_for_serialization(serialization_cache))
  File "/usr/local/lib/python3.7/site-packages/tensorflow/python/keras/saving/saved_model/base_serialization.py", line 87, in list_functions_for_serialization
    fns = self.functions_to_serialize(serialization_cache)
  File "/usr/local/lib/python3.7/site-packages/tensorflow/python/keras/saving/saved_model/layer_serialization.py", line 79, in functions_to_serialize
    serialization_cache).functions_to_serialize)
  File "/usr/local/lib/python3.7/site-packages/tensorflow/python/keras/saving/saved_model/layer_serialization.py", line 95, in _get_serialized_attributes
    serialization_cache)
  File "/usr/local/lib/python3.7/site-packages/tensorflow/python/keras/saving/saved_model/model_serialization.py", line 57, in _get_serialized_attributes_internal
    serialization_cache))
  File "/usr/local/lib/python3.7/site-packages/tensorflow/python/keras/saving/saved_model/layer_serialization.py", line 104, in _get_serialized_attributes_internal
    functions = save_impl.wrap_layer_functions(self.obj, serialization_cache)
  File "/usr/local/lib/python3.7/site-packages/tensorflow/python/keras/saving/saved_model/save_impl.py", line 165, in wrap_layer_functions
    '
{}
_layer_call_and_return_conditional_losses'.format(layer.name))
  File "/usr/local/lib/python3.7/site-packages/tensorflow/python/keras/saving/saved_model/save_impl.py", line 505, in add_function
    self.add_trace(*self._input_signature)
  File "/usr/local/lib/python3.7/site-packages/tensorflow/python/keras/saving/saved_model/save_impl.py", line 420, in add_trace
    trace_with_training(True)
  File "/usr/local/lib/python3.7/site-packages/tensorflow/python/keras/saving/saved_model/save_impl.py", line 418, in trace_with_training
    fn.get_concrete_function(*args, **kwargs)
  File "/usr/local/lib/python3.7/site-packages/tensorflow/python/keras/saving/saved_model/save_impl.py", line 549, in get_concrete_function
    return super(LayerCall, self).get_concrete_function(*args, **kwargs)
  File "/usr/local/lib/python3.7/site-packages/tensorflow/python/eager/def_function.py", line 1167, in get_concrete_function
    concrete = self._get_concrete_function_garbage_collected(*args, **kwargs)
  File "/usr/local/lib/python3.7/site-packages/tensorflow/python/eager/def_function.py", line 1073, in _get_concrete_function_garbage_collected
    self._initialize(args, kwargs, add_initializers_to=initializers)
  File "/usr/local/lib/python3.7/site-packages/tensorflow/python/eager/def_function.py", line 697, in _initialize
    *args, **kwds))
  File "/usr/local/lib/python3.7/site-packages/tensorflow/python/eager/function.py", line 2855, in _get_concrete_function_internal_garbage_collected
    graph_function, _, _ = self._maybe_define_function(args, kwargs)
  File "/usr/local/lib/python3.7/site-packages/tensorflow/python/eager/function.py", line 3213, in _maybe_define_function
    graph_function = self._create_graph_function(args, kwargs)
  File "/usr/local/lib/python3.7/site-packages/tensorflow/python/eager/function.py", line 3075, in _create_graph_function
    capture_by_value=self._capture_by_value),
  File "/usr/local/lib/python3.7/site-packages/tensorflow/python/framework/func_graph.py", line 986, in func_graph_from_py_func
    func_outputs = python_func(*func_args, **func_kwargs)
  File "/usr/local/lib/python3.7/site-packages/tensorflow/python/eager/def_function.py", line 600, in wrapped_fn
    return weak_wrapped_fn().__wrapped__(*args, **kwds)
  File "/usr/local/lib/python3.7/site-packages/tensorflow/python/keras/saving/saved_model/save_impl.py", line 526, in wrapper
    ret = method(*args, **kwargs)
  File "/usr/local/lib/python3.7/site-packages/tensorflow/python/keras/saving/saved_model/utils.py", line 170, in wrap_with_training_arg
    lambda: replace_training_and_call(False))
  File "/usr/local/lib/python3.7/site-packages/tensorflow/python/keras/utils/tf_utils.py", line 65, in smart_cond
    pred, true_fn=true_fn, false_fn=false_fn, name=name)
  File "/usr/local/lib/python3.7/site-packages/tensorflow/python/framework/smart_cond.py", line 54, in smart_cond
    return true_fn()
  File "/usr/local/lib/python3.7/site-packages/tensorflow/python/keras/saving/saved_model/utils.py", line 169, in <lambda>
    lambda: replace_training_and_call(True),
  File "/usr/local/lib/python3.7/site-packages/tensorflow/python/keras/saving/saved_model/utils.py", line 165, in replace_training_and_call
    return wrapped_call(*args, **kwargs)
  File "/usr/local/lib/python3.7/site-packages/tensorflow/python/keras/saving/saved_model/save_impl.py", line 569, in call_and_return_conditional_losses
    call_output = layer_call(inputs, *args, **kwargs)
  File "/usr/local/lib/python3.7/site-packages/tensorflow/python/keras/engine/functional.py", line 389, in call
    inputs, training=training, mask=mask)
  File "/usr/local/lib/python3.7/site-packages/tensorflow/python/keras/engine/functional.py", line 512, in _run_internal_graph
    outputs = node.layer(*args, **kwargs)
  File "/usr/local/lib/python3.7/site-packages/tensorflow/python/keras/engine/base_layer.py", line 999, in __call__
    outputs = call_fn(inputs, *args, **kwargs)
  File "/usr/local/lib/python3.7/site-packages/tensorflow/python/keras/saving/saved_model/utils.py", line 71, in return_outputs_and_add_losses
    outputs, losses = fn(inputs, *args, **kwargs)
  File "/usr/local/lib/python3.7/site-packages/tensorflow/python/keras/saving/saved_model/utils.py", line 170, in wrap_with_training_arg
    lambda: replace_training_and_call(False))
  File "/usr/local/lib/python3.7/site-packages/tensorflow/python/keras/utils/tf_utils.py", line 65, in smart_cond
    pred, true_fn=true_fn, false_fn=false_fn, name=name)
  File "/usr/local/lib/python3.7/site-packages/tensorflow/python/framework/smart_cond.py", line 54, in smart_cond
    return true_fn()
  File "/usr/local/lib/python3.7/site-packages/tensorflow/python/keras/saving/saved_model/utils.py", line 169, in <lambda>
    lambda: replace_training_and_call(True),
  File "/usr/local/lib/python3.7/site-packages/tensorflow/python/keras/saving/saved_model/utils.py", line 165, in replace_training_and_call
    return wrapped_call(*args, **kwargs)
  File "/usr/local/lib/python3.7/site-packages/tensorflow/python/keras/saving/saved_model/save_impl.py", line 543, in __call__
    self.call_collection.add_trace(*args, **kwargs)
  File "/usr/local/lib/python3.7/site-packages/tensorflow/python/keras/saving/saved_model/save_impl.py", line 420, in add_trace
    trace_with_training(True)
  File "/usr/local/lib/python3.7/site-packages/tensorflow/python/keras/saving/saved_model/save_impl.py", line 418, in trace_with_training
    fn.get_concrete_function(*args, **kwargs)
  File "/usr/local/lib/python3.7/site-packages/tensorflow/python/keras/saving/saved_model/save_impl.py", line 549, in get_concrete_function
    return super(LayerCall, self).get_concrete_function(*args, **kwargs)
  File "/usr/local/lib/python3.7/site-packages/tensorflow/python/eager/def_function.py", line 1167, in get_concrete_function
    concrete = self._get_concrete_function_garbage_collected(*args, **kwargs)
  File "/usr/local/lib/python3.7/site-packages/tensorflow/python/eager/def_function.py", line 1073, in _get_concrete_function_garbage_collected
    self._initialize(args, kwargs, add_initializers_to=initializers)
  File "/usr/local/lib/python3.7/site-packages/tensorflow/python/eager/def_function.py", line 697, in _initialize
    *args, **kwds))
  File "/usr/local/lib/python3.7/site-packages/tensorflow/python/eager/function.py", line 2855, in _get_concrete_function_internal_garbage_collected
    graph_function, _, _ = self._maybe_define_function(args, kwargs)
  File "/usr/local/lib/python3.7/site-packages/tensorflow/python/eager/function.py", line 3213, in _maybe_define_function
    graph_function = self._create_graph_function(args, kwargs)
  File "/usr/local/lib/python3.7/site-packages/tensorflow/python/eager/function.py", line 3075, in _create_graph_function
    capture_by_value=self._capture_by_value),
  File "/usr/local/lib/python3.7/site-packages/tensorflow/python/framework/func_graph.py", line 986, in func_graph_from_py_func
    func_outputs = python_func(*func_args, **func_kwargs)
  File "/usr/local/lib/python3.7/site-packages/tensorflow/python/eager/def_function.py", line 600, in wrapped_fn
    return weak_wrapped_fn().__wrapped__(*args, **kwds)
  File "/usr/local/lib/python3.7/site-packages/tensorflow/python/keras/saving/saved_model/save_impl.py", line 515, in wrapper
    inputs = call_collection.get_input_arg_value(args, kwargs)
  File "/usr/local/lib/python3.7/site-packages/tensorflow/python/keras/saving/saved_model/save_impl.py", line 455, in get_input_arg_value
    self._input_arg_name, args, kwargs, inputs_in_args=True)
  File "/usr/local/lib/python3.7/site-packages/tensorflow/python/keras/engine/base_layer.py", line 2581, in _get_call_arg_value
    return args_dict[arg_name]

KeyError: 'inputs'

Everything works fine in smdebug 0.9.3 but not sure why it's breaking with 0.9.4

[neda60]

I have the exactly the same issue, with the exact same logs! Any solution for it?

[gbordyugov]

The fix with debugger_hook_config described above helped me with a KeyError: 'inputs' problem, too.