We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
Fix flaky spark integration tests
Several Spark tests have been disabled in https://github.com/horovod/horovod/pull/3259, we should address the underlying flakiness and re-enable.

--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `examples/spark/pytorch/pytorch_lightning_spark_mnist.py`
Content:
```
1 import argparse
2 import os
3 import subprocess
4 import sys
5 from distutils.version import LooseVersion
6 
7 import numpy as np
8 
9 import pyspark
10 import pyspark.sql.types as T
11 from pyspark import SparkConf
12 from pyspark.ml.evaluation import MulticlassClassificationEvaluator
13 if LooseVersion(pyspark.__version__) < LooseVersion('3.0.0'):
14     from pyspark.ml.feature import OneHotEncoderEstimator as OneHotEncoder
15 else:
16     from pyspark.ml.feature import OneHotEncoder
17 from pyspark.sql import SparkSession
18 from pyspark.sql.functions import udf
19 
20 from pytorch_lightning import LightningModule
21 
22 import torch
23 import torch.nn as nn
24 import torch.nn.functional as F
25 import torch.optim as optim
26 
27 import horovod.spark.lightning as hvd
28 from horovod.spark.lightning.estimator import MIN_PL_VERSION
29 from horovod.spark.common.backend import SparkBackend
30 from horovod.spark.common.store import Store
31 
32 parser = argparse.ArgumentParser(description='PyTorch Spark MNIST Example',
33                                  formatter_class=argparse.ArgumentDefaultsHelpFormatter)
34 parser.add_argument('--master',
35                     help='spark master to connect to')
36 parser.add_argument('--num-proc', type=int,
37                     help='number of worker processes for training, default: `spark.default.parallelism`')
38 parser.add_argument('--batch-size', type=int, default=64,
39                     help='input batch size for training')
40 parser.add_argument('--epochs', type=int, default=12,
41                     help='number of epochs to train')
42 parser.add_argument('--work-dir', default='/tmp',
43                     help='temporary working directory to write intermediate files (prefix with hdfs:// to use HDFS)')
44 parser.add_argument('--data-dir', default='/tmp',
45                     help='location of the training dataset in the local filesystem (will be downloaded if needed)')
46 parser.add_argument('--enable-profiler', action='store_true',
47                     help='Enable profiler')
48 
49 
50 def train_model(args):
51     # do not run this test for pytorch lightning below min supported verson
52     import pytorch_lightning as pl
53     if LooseVersion(pl.__version__) < LooseVersion(MIN_PL_VERSION):
54         print("Skip test for pytorch_ligthning=={}, min support version is {}".format(pl.__version__, MIN_PL_VERSION))
55         return
56 
57     # Initialize SparkSession
58     conf = SparkConf().setAppName('pytorch_spark_mnist').set('spark.sql.shuffle.partitions', '16')
59     if args.master:
60         conf.setMaster(args.master)
61     elif args.num_proc:
62         conf.setMaster('local[{}]'.format(args.num_proc))
63     spark = SparkSession.builder.config(conf=conf).getOrCreate()
64 
65     # Setup our store for intermediate data
66     store = Store.create(args.work_dir)
67 
68     # Download MNIST dataset
69     data_url = 'https://www.csie.ntu.edu.tw/~cjlin/libsvmtools/datasets/multiclass/mnist.bz2'
70     libsvm_path = os.path.join(args.data_dir, 'mnist.bz2')
71     if not os.path.exists(libsvm_path):
72         subprocess.check_output(['wget', data_url, '-O', libsvm_path])
73 
74     # Load dataset into a Spark DataFrame
75     df = spark.read.format('libsvm') \
76         .option('numFeatures', '784') \
77         .load(libsvm_path)
78 
79     # One-hot encode labels into SparseVectors
80     encoder = OneHotEncoder(inputCols=['label'],
81                             outputCols=['label_vec'],
82                             dropLast=False)
83     model = encoder.fit(df)
84     train_df = model.transform(df)
85 
86     # Train/test split
87     train_df, test_df = train_df.randomSplit([0.9, 0.1])
88 
89     # Define the PyTorch model without any Horovod-specific parameters
90     class Net(LightningModule):
91         def __init__(self):
92             super(Net, self).__init__()
93             self.conv1 = nn.Conv2d(1, 10, kernel_size=5)
94             self.conv2 = nn.Conv2d(10, 20, kernel_size=5)
95             self.conv2_drop = nn.Dropout2d()
96             self.fc1 = nn.Linear(320, 50)
97             self.fc2 = nn.Linear(50, 10)
98 
99         def forward(self, x):
100             x = x.float().reshape((-1, 1, 28, 28))
101             x = F.relu(F.max_pool2d(self.conv1(x), 2))
102             x = F.relu(F.max_pool2d(self.conv2_drop(self.conv2(x)), 2))
103             x = x.view(-1, 320)
104             x = F.relu(self.fc1(x))
105             x = F.dropout(x, training=self.training)
106             x = self.fc2(x)
107             return F.log_softmax(x, -1)
108 
109         def configure_optimizers(self):
110             return optim.SGD(self.parameters(), lr=0.01, momentum=0.5)
111 
112         def training_step(self, batch, batch_idx):
113             if batch_idx == 0:
114                 print(f"training data batch size: {batch['label'].shape}")
115             x, y = batch['features'], batch['label']
116             y_hat = self(x)
117             loss = F.nll_loss(y_hat, y.long())
118             self.log('train_loss', loss)
119             return loss
120 
121         def validation_step(self, batch, batch_idx):
122             if batch_idx == 0:
123                 print(f"validation data batch size: {batch['label'].shape}")
124             x, y = batch['features'], batch['label']
125             y_hat = self(x)
126             loss = F.nll_loss(y_hat, y.long())
127             self.log('val_loss', loss)
128 
129         def validation_epoch_end(self, outputs):
130             avg_loss = torch.stack([x['val_loss'] for x in outputs]).mean() if len(outputs) > 0 else float('inf')
131             self.log('avg_val_loss', avg_loss)
132 
133     model = Net()
134 
135     # Train a Horovod Spark Estimator on the DataFrame
136     backend = SparkBackend(num_proc=args.num_proc,
137                            stdout=sys.stdout, stderr=sys.stderr,
138                            prefix_output_with_timestamp=True)
139 
140     from pytorch_lightning.callbacks import Callback
141 
142     epochs = args.epochs
143 
144     class MyDummyCallback(Callback):
145         def __init__(self):
146             self.epcoh_end_counter = 0
147             self.train_epcoh_end_counter = 0
148             self.validation_epoch_end_counter = 0
149 
150         def on_init_start(self, trainer):
151             print('Starting to init trainer!')
152 
153         def on_init_end(self, trainer):
154             print('Trainer is initialized.')
155 
156         def on_epoch_end(self, trainer, model):
157             print('A train or eval epoch ended.')
158             self.epcoh_end_counter += 1
159 
160         def on_train_epoch_end(self, trainer, model, unused=None):
161             print('A train epoch ended.')
162             self.train_epcoh_end_counter += 1
163 
164         def on_validation_epoch_end(self, trainer, model, unused=None):
165             print('A val epoch ended.')
166             self.validation_epoch_end_counter += 1
167 
168         def on_train_end(self, trainer, model):
169             print("Training ends:"
170                   f"epcoh_end_counter={self.epcoh_end_counter}, "
171                   f"train_epcoh_end_counter={self.train_epcoh_end_counter}, "
172                   f"validation_epoch_end_counter={self.validation_epoch_end_counter} \n")
173             assert self.train_epcoh_end_counter <= epochs
174             assert self.epcoh_end_counter == self.train_epcoh_end_counter + self.validation_epoch_end_counter
175 
176     callbacks = [MyDummyCallback()]
177 
178     # added EarlyStopping and ModelCheckpoint
179     from pytorch_lightning.callbacks.model_checkpoint import ModelCheckpoint
180     callbacks.append(ModelCheckpoint(monitor='val_loss', mode="min",
181                                      save_top_k=1, verbose=True))
182 
183     from pytorch_lightning.callbacks.early_stopping import EarlyStopping
184     callbacks.append(EarlyStopping(monitor='val_loss',
185                                    min_delta=0.001,
186                                    patience=3,
187                                    verbose=True,
188                                    mode='min'))
189 
190     torch_estimator = hvd.TorchEstimator(backend=backend,
191                                          store=store,
192                                          model=model,
193                                          input_shapes=[[-1, 1, 28, 28]],
194                                          feature_cols=['features'],
195                                          label_cols=['label'],
196                                          batch_size=args.batch_size,
197                                          epochs=args.epochs,
198                                          validation=0.1,
199                                          verbose=1,
200                                          callbacks=callbacks,
201                                          profiler="simple" if args.enable_profiler else None)
202 
203     torch_model = torch_estimator.fit(train_df).setOutputCols(['label_prob'])
204 
205     # Evaluate the model on the held-out test DataFrame
206     pred_df = torch_model.transform(test_df)
207 
208     argmax = udf(lambda v: float(np.argmax(v)), returnType=T.DoubleType())
209     pred_df = pred_df.withColumn('label_pred', argmax(pred_df.label_prob))
210     evaluator = MulticlassClassificationEvaluator(predictionCol='label_pred', labelCol='label', metricName='accuracy')
211     print('Test accuracy:', evaluator.evaluate(pred_df))
212 
213     spark.stop()
214 
215 
216 if __name__ == '__main__':
217     args = parser.parse_args()
218     train_model(args)
219
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/examples/spark/pytorch/pytorch_lightning_spark_mnist.py b/examples/spark/pytorch/pytorch_lightning_spark_mnist.py
--- a/examples/spark/pytorch/pytorch_lightning_spark_mnist.py
+++ b/examples/spark/pytorch/pytorch_lightning_spark_mnist.py
@@ -17,6 +17,17 @@
 from pyspark.sql import SparkSession
 from pyspark.sql.functions import udf
 
+# Spark PyTorch Lightning tests conflict with Tensorflow 2.5.x: https://github.com/horovod/horovod/pull/3263
+try:
+    # tensorflow has to be imported BEFORE pytorch_lightning, otherwise we see the segfault right away
+    import tensorflow as tf
+    from distutils.version import LooseVersion
+    if LooseVersion('2.5.0') <= LooseVersion(tf.__version__) < LooseVersion('2.6.0'):
+        print('Skipping test as Pytorch Lightning conflicts with present Tensorflow 2.5.x', file=sys.stderr)
+        sys.exit(0)
+except ImportError:
+    pass
+
 from pytorch_lightning import LightningModule
 
 import torch

{"golden_diff": "diff --git a/examples/spark/pytorch/pytorch_lightning_spark_mnist.py b/examples/spark/pytorch/pytorch_lightning_spark_mnist.py\n--- a/examples/spark/pytorch/pytorch_lightning_spark_mnist.py\n+++ b/examples/spark/pytorch/pytorch_lightning_spark_mnist.py\n@@ -17,6 +17,17 @@\n from pyspark.sql import SparkSession\n from pyspark.sql.functions import udf\n \n+# Spark PyTorch Lightning tests conflict with Tensorflow 2.5.x: https://github.com/horovod/horovod/pull/3263\n+try:\n+    # tensorflow has to be imported BEFORE pytorch_lightning, otherwise we see the segfault right away\n+    import tensorflow as tf\n+    from distutils.version import LooseVersion\n+    if LooseVersion('2.5.0') <= LooseVersion(tf.__version__) < LooseVersion('2.6.0'):\n+        print('Skipping test as Pytorch Lightning conflicts with present Tensorflow 2.5.x', file=sys.stderr)\n+        sys.exit(0)\n+except ImportError:\n+    pass\n+\n from pytorch_lightning import LightningModule\n \n import torch\n", "issue": "Fix flaky spark integration tests\nSeveral Spark tests have been disabled in https://github.com/horovod/horovod/pull/3259, we should address the underlying flakiness and re-enable.\n", "before_files": [{"content": "import argparse\nimport os\nimport subprocess\nimport sys\nfrom distutils.version import LooseVersion\n\nimport numpy as np\n\nimport pyspark\nimport pyspark.sql.types as T\nfrom pyspark import SparkConf\nfrom pyspark.ml.evaluation import MulticlassClassificationEvaluator\nif LooseVersion(pyspark.__version__) < LooseVersion('3.0.0'):\n    from pyspark.ml.feature import OneHotEncoderEstimator as OneHotEncoder\nelse:\n    from pyspark.ml.feature import OneHotEncoder\nfrom pyspark.sql import SparkSession\nfrom pyspark.sql.functions import udf\n\nfrom pytorch_lightning import LightningModule\n\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport torch.optim as optim\n\nimport horovod.spark.lightning as hvd\nfrom horovod.spark.lightning.estimator import MIN_PL_VERSION\nfrom horovod.spark.common.backend import SparkBackend\nfrom horovod.spark.common.store import Store\n\nparser = argparse.ArgumentParser(description='PyTorch Spark MNIST Example',\n                                 formatter_class=argparse.ArgumentDefaultsHelpFormatter)\nparser.add_argument('--master',\n                    help='spark master to connect to')\nparser.add_argument('--num-proc', type=int,\n                    help='number of worker processes for training, default: `spark.default.parallelism`')\nparser.add_argument('--batch-size', type=int, default=64,\n                    help='input batch size for training')\nparser.add_argument('--epochs', type=int, default=12,\n                    help='number of epochs to train')\nparser.add_argument('--work-dir', default='/tmp',\n                    help='temporary working directory to write intermediate files (prefix with hdfs:// to use HDFS)')\nparser.add_argument('--data-dir', default='/tmp',\n                    help='location of the training dataset in the local filesystem (will be downloaded if needed)')\nparser.add_argument('--enable-profiler', action='store_true',\n                    help='Enable profiler')\n\n\ndef train_model(args):\n    # do not run this test for pytorch lightning below min supported verson\n    import pytorch_lightning as pl\n    if LooseVersion(pl.__version__) < LooseVersion(MIN_PL_VERSION):\n        print(\"Skip test for pytorch_ligthning=={}, min support version is {}\".format(pl.__version__, MIN_PL_VERSION))\n        return\n\n    # Initialize SparkSession\n    conf = SparkConf().setAppName('pytorch_spark_mnist').set('spark.sql.shuffle.partitions', '16')\n    if args.master:\n        conf.setMaster(args.master)\n    elif args.num_proc:\n        conf.setMaster('local[{}]'.format(args.num_proc))\n    spark = SparkSession.builder.config(conf=conf).getOrCreate()\n\n    # Setup our store for intermediate data\n    store = Store.create(args.work_dir)\n\n    # Download MNIST dataset\n    data_url = 'https://www.csie.ntu.edu.tw/~cjlin/libsvmtools/datasets/multiclass/mnist.bz2'\n    libsvm_path = os.path.join(args.data_dir, 'mnist.bz2')\n    if not os.path.exists(libsvm_path):\n        subprocess.check_output(['wget', data_url, '-O', libsvm_path])\n\n    # Load dataset into a Spark DataFrame\n    df = spark.read.format('libsvm') \\\n        .option('numFeatures', '784') \\\n        .load(libsvm_path)\n\n    # One-hot encode labels into SparseVectors\n    encoder = OneHotEncoder(inputCols=['label'],\n                            outputCols=['label_vec'],\n                            dropLast=False)\n    model = encoder.fit(df)\n    train_df = model.transform(df)\n\n    # Train/test split\n    train_df, test_df = train_df.randomSplit([0.9, 0.1])\n\n    # Define the PyTorch model without any Horovod-specific parameters\n    class Net(LightningModule):\n        def __init__(self):\n            super(Net, self).__init__()\n            self.conv1 = nn.Conv2d(1, 10, kernel_size=5)\n            self.conv2 = nn.Conv2d(10, 20, kernel_size=5)\n            self.conv2_drop = nn.Dropout2d()\n            self.fc1 = nn.Linear(320, 50)\n            self.fc2 = nn.Linear(50, 10)\n\n        def forward(self, x):\n            x = x.float().reshape((-1, 1, 28, 28))\n            x = F.relu(F.max_pool2d(self.conv1(x), 2))\n            x = F.relu(F.max_pool2d(self.conv2_drop(self.conv2(x)), 2))\n            x = x.view(-1, 320)\n            x = F.relu(self.fc1(x))\n            x = F.dropout(x, training=self.training)\n            x = self.fc2(x)\n            return F.log_softmax(x, -1)\n\n        def configure_optimizers(self):\n            return optim.SGD(self.parameters(), lr=0.01, momentum=0.5)\n\n        def training_step(self, batch, batch_idx):\n            if batch_idx == 0:\n                print(f\"training data batch size: {batch['label'].shape}\")\n            x, y = batch['features'], batch['label']\n            y_hat = self(x)\n            loss = F.nll_loss(y_hat, y.long())\n            self.log('train_loss', loss)\n            return loss\n\n        def validation_step(self, batch, batch_idx):\n            if batch_idx == 0:\n                print(f\"validation data batch size: {batch['label'].shape}\")\n            x, y = batch['features'], batch['label']\n            y_hat = self(x)\n            loss = F.nll_loss(y_hat, y.long())\n            self.log('val_loss', loss)\n\n        def validation_epoch_end(self, outputs):\n            avg_loss = torch.stack([x['val_loss'] for x in outputs]).mean() if len(outputs) > 0 else float('inf')\n            self.log('avg_val_loss', avg_loss)\n\n    model = Net()\n\n    # Train a Horovod Spark Estimator on the DataFrame\n    backend = SparkBackend(num_proc=args.num_proc,\n                           stdout=sys.stdout, stderr=sys.stderr,\n                           prefix_output_with_timestamp=True)\n\n    from pytorch_lightning.callbacks import Callback\n\n    epochs = args.epochs\n\n    class MyDummyCallback(Callback):\n        def __init__(self):\n            self.epcoh_end_counter = 0\n            self.train_epcoh_end_counter = 0\n            self.validation_epoch_end_counter = 0\n\n        def on_init_start(self, trainer):\n            print('Starting to init trainer!')\n\n        def on_init_end(self, trainer):\n            print('Trainer is initialized.')\n\n        def on_epoch_end(self, trainer, model):\n            print('A train or eval epoch ended.')\n            self.epcoh_end_counter += 1\n\n        def on_train_epoch_end(self, trainer, model, unused=None):\n            print('A train epoch ended.')\n            self.train_epcoh_end_counter += 1\n\n        def on_validation_epoch_end(self, trainer, model, unused=None):\n            print('A val epoch ended.')\n            self.validation_epoch_end_counter += 1\n\n        def on_train_end(self, trainer, model):\n            print(\"Training ends:\"\n                  f\"epcoh_end_counter={self.epcoh_end_counter}, \"\n                  f\"train_epcoh_end_counter={self.train_epcoh_end_counter}, \"\n                  f\"validation_epoch_end_counter={self.validation_epoch_end_counter} \\n\")\n            assert self.train_epcoh_end_counter <= epochs\n            assert self.epcoh_end_counter == self.train_epcoh_end_counter + self.validation_epoch_end_counter\n\n    callbacks = [MyDummyCallback()]\n\n    # added EarlyStopping and ModelCheckpoint\n    from pytorch_lightning.callbacks.model_checkpoint import ModelCheckpoint\n    callbacks.append(ModelCheckpoint(monitor='val_loss', mode=\"min\",\n                                     save_top_k=1, verbose=True))\n\n    from pytorch_lightning.callbacks.early_stopping import EarlyStopping\n    callbacks.append(EarlyStopping(monitor='val_loss',\n                                   min_delta=0.001,\n                                   patience=3,\n                                   verbose=True,\n                                   mode='min'))\n\n    torch_estimator = hvd.TorchEstimator(backend=backend,\n                                         store=store,\n                                         model=model,\n                                         input_shapes=[[-1, 1, 28, 28]],\n                                         feature_cols=['features'],\n                                         label_cols=['label'],\n                                         batch_size=args.batch_size,\n                                         epochs=args.epochs,\n                                         validation=0.1,\n                                         verbose=1,\n                                         callbacks=callbacks,\n                                         profiler=\"simple\" if args.enable_profiler else None)\n\n    torch_model = torch_estimator.fit(train_df).setOutputCols(['label_prob'])\n\n    # Evaluate the model on the held-out test DataFrame\n    pred_df = torch_model.transform(test_df)\n\n    argmax = udf(lambda v: float(np.argmax(v)), returnType=T.DoubleType())\n    pred_df = pred_df.withColumn('label_pred', argmax(pred_df.label_prob))\n    evaluator = MulticlassClassificationEvaluator(predictionCol='label_pred', labelCol='label', metricName='accuracy')\n    print('Test accuracy:', evaluator.evaluate(pred_df))\n\n    spark.stop()\n\n\nif __name__ == '__main__':\n    args = parser.parse_args()\n    train_model(args)\n", "path": "examples/spark/pytorch/pytorch_lightning_spark_mnist.py"}], "after_files": [{"content": "import argparse\nimport os\nimport subprocess\nimport sys\nfrom distutils.version import LooseVersion\n\nimport numpy as np\n\nimport pyspark\nimport pyspark.sql.types as T\nfrom pyspark import SparkConf\nfrom pyspark.ml.evaluation import MulticlassClassificationEvaluator\nif LooseVersion(pyspark.__version__) < LooseVersion('3.0.0'):\n    from pyspark.ml.feature import OneHotEncoderEstimator as OneHotEncoder\nelse:\n    from pyspark.ml.feature import OneHotEncoder\nfrom pyspark.sql import SparkSession\nfrom pyspark.sql.functions import udf\n\n# Spark PyTorch Lightning tests conflict with Tensorflow 2.5.x: https://github.com/horovod/horovod/pull/3263\ntry:\n    # tensorflow has to be imported BEFORE pytorch_lightning, otherwise we see the segfault right away\n    import tensorflow as tf\n    from distutils.version import LooseVersion\n    if LooseVersion('2.5.0') <= LooseVersion(tf.__version__) < LooseVersion('2.6.0'):\n        print('Skipping test as Pytorch Lightning conflicts with present Tensorflow 2.5.x', file=sys.stderr)\n        sys.exit(0)\nexcept ImportError:\n    pass\n\nfrom pytorch_lightning import LightningModule\n\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport torch.optim as optim\n\nimport horovod.spark.lightning as hvd\nfrom horovod.spark.lightning.estimator import MIN_PL_VERSION\nfrom horovod.spark.common.backend import SparkBackend\nfrom horovod.spark.common.store import Store\n\nparser = argparse.ArgumentParser(description='PyTorch Spark MNIST Example',\n                                 formatter_class=argparse.ArgumentDefaultsHelpFormatter)\nparser.add_argument('--master',\n                    help='spark master to connect to')\nparser.add_argument('--num-proc', type=int,\n                    help='number of worker processes for training, default: `spark.default.parallelism`')\nparser.add_argument('--batch-size', type=int, default=64,\n                    help='input batch size for training')\nparser.add_argument('--epochs', type=int, default=12,\n                    help='number of epochs to train')\nparser.add_argument('--work-dir', default='/tmp',\n                    help='temporary working directory to write intermediate files (prefix with hdfs:// to use HDFS)')\nparser.add_argument('--data-dir', default='/tmp',\n                    help='location of the training dataset in the local filesystem (will be downloaded if needed)')\nparser.add_argument('--enable-profiler', action='store_true',\n                    help='Enable profiler')\n\n\ndef train_model(args):\n    # do not run this test for pytorch lightning below min supported verson\n    import pytorch_lightning as pl\n    if LooseVersion(pl.__version__) < LooseVersion(MIN_PL_VERSION):\n        print(\"Skip test for pytorch_ligthning=={}, min support version is {}\".format(pl.__version__, MIN_PL_VERSION))\n        return\n\n    # Initialize SparkSession\n    conf = SparkConf().setAppName('pytorch_spark_mnist').set('spark.sql.shuffle.partitions', '16')\n    if args.master:\n        conf.setMaster(args.master)\n    elif args.num_proc:\n        conf.setMaster('local[{}]'.format(args.num_proc))\n    spark = SparkSession.builder.config(conf=conf).getOrCreate()\n\n    # Setup our store for intermediate data\n    store = Store.create(args.work_dir)\n\n    # Download MNIST dataset\n    data_url = 'https://www.csie.ntu.edu.tw/~cjlin/libsvmtools/datasets/multiclass/mnist.bz2'\n    libsvm_path = os.path.join(args.data_dir, 'mnist.bz2')\n    if not os.path.exists(libsvm_path):\n        subprocess.check_output(['wget', data_url, '-O', libsvm_path])\n\n    # Load dataset into a Spark DataFrame\n    df = spark.read.format('libsvm') \\\n        .option('numFeatures', '784') \\\n        .load(libsvm_path)\n\n    # One-hot encode labels into SparseVectors\n    encoder = OneHotEncoder(inputCols=['label'],\n                            outputCols=['label_vec'],\n                            dropLast=False)\n    model = encoder.fit(df)\n    train_df = model.transform(df)\n\n    # Train/test split\n    train_df, test_df = train_df.randomSplit([0.9, 0.1])\n\n    # Define the PyTorch model without any Horovod-specific parameters\n    class Net(LightningModule):\n        def __init__(self):\n            super(Net, self).__init__()\n            self.conv1 = nn.Conv2d(1, 10, kernel_size=5)\n            self.conv2 = nn.Conv2d(10, 20, kernel_size=5)\n            self.conv2_drop = nn.Dropout2d()\n            self.fc1 = nn.Linear(320, 50)\n            self.fc2 = nn.Linear(50, 10)\n\n        def forward(self, x):\n            x = x.float().reshape((-1, 1, 28, 28))\n            x = F.relu(F.max_pool2d(self.conv1(x), 2))\n            x = F.relu(F.max_pool2d(self.conv2_drop(self.conv2(x)), 2))\n            x = x.view(-1, 320)\n            x = F.relu(self.fc1(x))\n            x = F.dropout(x, training=self.training)\n            x = self.fc2(x)\n            return F.log_softmax(x, -1)\n\n        def configure_optimizers(self):\n            return optim.SGD(self.parameters(), lr=0.01, momentum=0.5)\n\n        def training_step(self, batch, batch_idx):\n            if batch_idx == 0:\n                print(f\"training data batch size: {batch['label'].shape}\")\n            x, y = batch['features'], batch['label']\n            y_hat = self(x)\n            loss = F.nll_loss(y_hat, y.long())\n            self.log('train_loss', loss)\n            return loss\n\n        def validation_step(self, batch, batch_idx):\n            if batch_idx == 0:\n                print(f\"validation data batch size: {batch['label'].shape}\")\n            x, y = batch['features'], batch['label']\n            y_hat = self(x)\n            loss = F.nll_loss(y_hat, y.long())\n            self.log('val_loss', loss)\n\n        def validation_epoch_end(self, outputs):\n            avg_loss = torch.stack([x['val_loss'] for x in outputs]).mean() if len(outputs) > 0 else float('inf')\n            self.log('avg_val_loss', avg_loss)\n\n    model = Net()\n\n    # Train a Horovod Spark Estimator on the DataFrame\n    backend = SparkBackend(num_proc=args.num_proc,\n                           stdout=sys.stdout, stderr=sys.stderr,\n                           prefix_output_with_timestamp=True)\n\n    from pytorch_lightning.callbacks import Callback\n\n    epochs = args.epochs\n\n    class MyDummyCallback(Callback):\n        def __init__(self):\n            self.epcoh_end_counter = 0\n            self.train_epcoh_end_counter = 0\n            self.validation_epoch_end_counter = 0\n\n        def on_init_start(self, trainer):\n            print('Starting to init trainer!')\n\n        def on_init_end(self, trainer):\n            print('Trainer is initialized.')\n\n        def on_epoch_end(self, trainer, model):\n            print('A train or eval epoch ended.')\n            self.epcoh_end_counter += 1\n\n        def on_train_epoch_end(self, trainer, model, unused=None):\n            print('A train epoch ended.')\n            self.train_epcoh_end_counter += 1\n\n        def on_validation_epoch_end(self, trainer, model, unused=None):\n            print('A val epoch ended.')\n            self.validation_epoch_end_counter += 1\n\n        def on_train_end(self, trainer, model):\n            print(\"Training ends:\"\n                  f\"epcoh_end_counter={self.epcoh_end_counter}, \"\n                  f\"train_epcoh_end_counter={self.train_epcoh_end_counter}, \"\n                  f\"validation_epoch_end_counter={self.validation_epoch_end_counter} \\n\")\n            assert self.train_epcoh_end_counter <= epochs\n            assert self.epcoh_end_counter == self.train_epcoh_end_counter + self.validation_epoch_end_counter\n\n    callbacks = [MyDummyCallback()]\n\n    # added EarlyStopping and ModelCheckpoint\n    from pytorch_lightning.callbacks.model_checkpoint import ModelCheckpoint\n    callbacks.append(ModelCheckpoint(monitor='val_loss', mode=\"min\",\n                                     save_top_k=1, verbose=True))\n\n    from pytorch_lightning.callbacks.early_stopping import EarlyStopping\n    callbacks.append(EarlyStopping(monitor='val_loss',\n                                   min_delta=0.001,\n                                   patience=3,\n                                   verbose=True,\n                                   mode='min'))\n\n    torch_estimator = hvd.TorchEstimator(backend=backend,\n                                         store=store,\n                                         model=model,\n                                         input_shapes=[[-1, 1, 28, 28]],\n                                         feature_cols=['features'],\n                                         label_cols=['label'],\n                                         batch_size=args.batch_size,\n                                         epochs=args.epochs,\n                                         validation=0.1,\n                                         verbose=1,\n                                         callbacks=callbacks,\n                                         profiler=\"simple\" if args.enable_profiler else None)\n\n    torch_model = torch_estimator.fit(train_df).setOutputCols(['label_prob'])\n\n    # Evaluate the model on the held-out test DataFrame\n    pred_df = torch_model.transform(test_df)\n\n    argmax = udf(lambda v: float(np.argmax(v)), returnType=T.DoubleType())\n    pred_df = pred_df.withColumn('label_pred', argmax(pred_df.label_prob))\n    evaluator = MulticlassClassificationEvaluator(predictionCol='label_pred', labelCol='label', metricName='accuracy')\n    print('Test accuracy:', evaluator.evaluate(pred_df))\n\n    spark.stop()\n\n\nif __name__ == '__main__':\n    args = parser.parse_args()\n    train_model(args)\n", "path": "examples/spark/pytorch/pytorch_lightning_spark_mnist.py"}]}