docstr: InputMode.SPARK TFEstimator w/ explicit saved_model export for TFModel inferencing
docstr_tok: ["InputMode.SPARK","TFEstimator","w\/","explicit","saved_model","export","for","TFModel","inferencing"]
code: def test_spark_saved_model(self):
      """InputMode.SPARK TFEstimator w/ explicit saved_model export for TFModel inferencing"""
  
      def _spark_train(args, ctx):
          """Basic linear regression in a distributed TF cluster using InputMode.SPARK"""
          import tensorflow as tf
          from tensorflowonspark import TFNode
  
          tf.compat.v1.reset_default_graph()
          strategy = (
              tf.distribute.experimental.MultiWorkerMirroredStrategy()
          )
  
          with strategy.scope():
              model = Sequential()
              model.add(Dense(1, activation="linear", input_shape=[2]))
              model.compile(
                  optimizer=tf.keras.optimizers.Adam(lr=0.2),
                  loss="mse",
                  metrics=["mse"],
              )
              model.summary()
  
          tf_feed = TFNode.DataFeed(
              ctx.mgr, input_mapping=args.input_mapping
          )
  
          def rdd_generator():
              while not tf_feed.should_stop():
                  batch = tf_feed.next_batch(1)
                  if len(batch["x"]) > 0:
                      features = batch["x"][0]
                      label = batch["y_"][0]
                      yield (features, label)
                  else:
                      return
  
          ds = tf.data.Dataset.from_generator(
              rdd_generator,
              (tf.float32, tf.float32),
              (tf.TensorShape([2]), tf.TensorShape([1])),
          )
          # disable auto-sharding since we're feeding from an RDD generator
          options = tf.data.Options()
          compat.disable_auto_shard(options)
          ds = ds.with_options(options)
          ds = ds.batch(args.batch_size)
  
          # only train 90% of each epoch to account for uneven RDD partition sizes
          steps_per_epoch = (
              1000 * 0.9 // (args.batch_size * ctx.num_workers)
          )
  
          tf.io.gfile.makedirs(args.model_dir)
          filepath = args.model_dir + "/weights-{epoch:04d}"
          callbacks = [
              tf.keras.callbacks.ModelCheckpoint(
                  filepath=filepath,
                  verbose=1,
                  load_weights_on_restart=True,
                  save_weights_only=True,
              )
          ]
  
          model.fit(
              ds,
              epochs=args.epochs,
              steps_per_epoch=steps_per_epoch,
              callbacks=callbacks,
          )
  
          # This fails with: "NotImplementedError: `fit_generator` is not supported for models compiled with tf.distribute.Strategy"
          # model.fit_generator(ds, epochs=args.epochs, steps_per_epoch=steps_per_epoch, callbacks=callbacks)
  
          if args.export_dir:
              print("exporting model to: {}".format(args.export_dir))
              compat.export_saved_model(
                  model, args.export_dir, ctx.job_name == "chief"
              )
  
          tf_feed.terminate()
  
      # create a Spark DataFrame of training examples (features, labels)
      rdd = self.sc.parallelize(self.train_examples, 2)
      trainDF = rdd.toDF(["col1", "col2"])
  
      # train and export model
      args = {}
      estimator = (
          TFEstimator(_spark_train, args)
          .setInputMapping({"col1": "x", "col2": "y_"})
          .setModelDir(self.model_dir)
          .setExportDir(self.export_dir)
          .setClusterSize(self.num_workers)
          .setMasterNode("chief")
          .setNumPS(0)
          .setBatchSize(1)
          .setEpochs(1)
      )
      model = estimator.fit(trainDF)
      self.assertTrue(os.path.isdir(self.export_dir))
  
      # create a Spark DataFrame of test examples (features, labels)
      testDF = self.spark.createDataFrame(
          self.test_examples, ["c1", "c2"]
      )
  
      # test saved_model using exported signature
      model.setTagSet("serve").setSignatureDefKey(
          "serving_default"
      ).setInputMapping({"c1": "dense_input"}).setOutputMapping(
          {"dense": "cout"}
      )
      preds = model.transform(testDF).head()  # take first/only result
      pred = preds.cout[0]  # unpack scalar from tensor
      expected = np.sum(self.weights)
      self.assertAlmostEqual(pred, expected, 2)
  
file_path: TensorFlowOnSpark/tests/test_pipeline.py
imports: import numpy
  import os
  import shutil
  import test
  import unittest
  from tensorflowonspark import compat
  from tensorflowonspark.pipeline import HasBatchSize, HasSteps, Namespace, TFEstimator, TFParams
  from tensorflow.keras import Sequential
  from tensorflow.keras.layers import Dense
label_desc: Use image node_id 7 for calling the PipelineTest obj's underlying member method code with example usage: obj.test_spark_saved_model() without return types