title: Sagemaker python sdk: accessing custom_attributes in inference job
body: <p>I am using Sagemaker python sdk for my inference job and <a href="https://sagemaker.readthedocs.io/en/stable/frameworks/tensorflow/using_tf.html#create-python-scripts-for-custom-input-and-output-formats" rel="nofollow noreferrer">following this guide</a>. I am triggering my sagemaker inference job from Airflow with below python callable:</p>
  <pre><code>def transform(sage_role, inference_file_local_path, **kwargs):
      &quot;&quot;&quot;
      Python callable to execute Sagemaker SDK train job. It takes infer_batch_output, infer_batch_input, model_artifact,
      instance_type and infer_file_name as run time parameter.
      :param inference_file_local_path: Local entry_point path for Inference file.
      :param sage_role: Sagemaker execution role.
      &quot;&quot;&quot;
      model = TensorFlowModel(entry_point=infer_file_name,
                              source_dir=inference_file_local_path,
                              model_data=model_artifact,
                              role=sage_role,
                              framework_version=&quot;2.5.1&quot;)
  
      tensorflow_serving_transformer = model.transformer(
          instance_count=1,
          instance_type=instance_type,
          accept=&quot;text/csv&quot;,
          strategy=&quot;SingleRecord&quot;,
          max_payload=10,
          max_concurrent_transforms=10,
          output_path=batch_output)
  
      return tensorflow_serving_transformer.transform(data=batch_input, content_type='text/csv')
  </code></pre>
  <p>and my simply <code>inference.py</code> looks like:</p>
  <pre><code>
  def input_handler(data, context):
      &quot;&quot;&quot; Pre-process request input before it is sent to TensorFlow Serving REST API
      Args:
          data (obj): the request data, in format of dict or string
          context (Context): an object containing request and configuration details
      Returns:
          (dict): a JSON-serializable dict that contains request body and headers
      &quot;&quot;&quot;
      if context.request_content_type == 'application/x-npy':
          # very simple numpy handler
          payload = np.load(data.read().decode('utf-8'))
          x_user_feature = np.asarray(payload.item().get('test').get('feature_a_list'))
          x_channel_feature = np.asarray(payload.item().get('test').get('feature_b_list'))
          examples = []
          for index, elem in enumerate(x_user_feature):
              examples.append({'feature_a_list': elem, 'feature_b_list': x_channel_feature[index]})
          return json.dumps({'instances': examples})
  
      if context.request_content_type == 'text/csv':
          payload = pd.read_csv(data)
          print(&quot;Model name is ..............&quot;)
          model_name = context.model_name
          print(model_name)
          examples = []
          row_ch = []
          if config_exists(model_bucket, &quot;{}{}&quot;.format(config_path, model_name)):
              config_keys = get_s3_json_file(model_bucket, &quot;{}{}&quot;.format(config_path, model_name))
              feature_b_list = config_keys[&quot;feature_b_list&quot;].split(&quot;,&quot;)
              row_ch = [float(ch_feature_str) for ch_feature_str in feature_b_list]
              if &quot;column_names&quot; in config_keys.keys():
                  cols = config_keys[&quot;column_names&quot;].split(&quot;,&quot;)
                  payload.columns = cols
          for index, row in payload.iterrows():
              row_user = row['feature_a_list'].replace('[', '').replace(']', '').split()
              row_user = [float(x) for x in row_user]
              if not row_ch:
                  row_ch = row['feature_b_list'].replace('[', '').replace(']', '').split()
                  row_ch = [float(x) for x in row_ch]
              example = {'feature_a_list': row_user, 'feature_b_list': row_ch}
              examples.append(example)
  
      raise ValueError('{{&quot;error&quot;: &quot;unsupported content type {}&quot;}}'.format(
          context.request_content_type or &quot;unknown&quot;))
  
  
  def output_handler(data, context):
      &quot;&quot;&quot;Post-process TensorFlow Serving output before it is returned to the client.
      Args:
          data (obj): the TensorFlow serving response
          context (Context): an object containing request and configuration details
      Returns:
          (bytes, string): data to return to client, response content type
      &quot;&quot;&quot;
      if data.status_code != 200:
          raise ValueError(data.content.decode('utf-8'))
  
      response_content_type = context.accept_header
      prediction = data.content
      return prediction, response_content_type
  
  </code></pre>
  <p>It is working fine however I want to pass custom arguments to <code>inference.py</code> so that I can modify the input data accordingly based on requirement. I thought of using a config file per requirement and download it from s3 based on model name but as I am using <code>model_data</code> and passes <code>model.tar.gz</code> at runtime <code>context.model_name</code> is always <code>None</code>.</p>
  <p>Is there a way I can pass run time argument to <code>inference.py</code> that I can use for customization?
  In the docs I see sagemaker provides <code>custom_attributes</code> but I don't see any example of it on how to use it and access it in <code>inference.py</code>.</p>
  <pre><code>custom_attributes (string): content of ‘X-Amzn-SageMaker-Custom-Attributes’ header from the original request. For example, ‘tfs-model-name=half*plus*three,tfs-method=predict’
  </code></pre>