From 116df9fed1cec0b89b132e351b2f48576e343913 Mon Sep 17 00:00:00 2001
From: Jack McCluskey <34928439+jrmccluskey@users.noreply.github.com>
Date: Wed, 18 Dec 2024 13:17:00 -0500
Subject: [PATCH] Fix env variable loading in Cost Benchmark workflow (#33404)

* Fix env variable loading in Cost Benchmark workflow

* fix output file for tf mnist

* add load test requirements file arg

* update mnist args

* revert how args are passed

* assign result correctly
---
 .github/workflows/beam_Python_CostBenchmarks_Dataflow.yml    | 5 +++--
 .../tensorflow_mnist_classification_cost_benchmark.py        | 2 +-
 2 files changed, 4 insertions(+), 3 deletions(-)

diff --git a/.github/workflows/beam_Python_CostBenchmarks_Dataflow.yml b/.github/workflows/beam_Python_CostBenchmarks_Dataflow.yml
index 18fe37e142ac..209325c429a1 100644
--- a/.github/workflows/beam_Python_CostBenchmarks_Dataflow.yml
+++ b/.github/workflows/beam_Python_CostBenchmarks_Dataflow.yml
@@ -92,7 +92,7 @@ jobs:
             -PloadTest.mainClass=apache_beam.testing.benchmarks.wordcount.wordcount \
             -Prunner=DataflowRunner \
             -PpythonVersion=3.10 \
-            '-PloadTest.args=${{ env.beam_Inference_Python_Benchmarks_Dataflow_test_arguments_1 }} --job_name=benchmark-tests-wordcount-python-${{env.NOW_UTC}} --output_file=gs://temp-storage-for-end-to-end-tests/wordcount/result_wordcount-${{env.NOW_UTC}}.txt' \
+            '-PloadTest.args=${{ env.beam_Python_Cost_Benchmarks_Dataflow_test_arguments_1 }} --job_name=benchmark-tests-wordcount-python-${{env.NOW_UTC}} --output_file=gs://temp-storage-for-end-to-end-tests/wordcount/result_wordcount-${{env.NOW_UTC}}.txt' \
       - name: Run Tensorflow MNIST Image Classification on Dataflow
         uses: ./.github/actions/gradle-command-self-hosted-action
         timeout-minutes: 30
@@ -102,4 +102,5 @@ jobs:
             -PloadTest.mainClass=apache_beam.testing.benchmarks.inference.tensorflow_mnist_classification_cost_benchmark \
             -Prunner=DataflowRunner \
             -PpythonVersion=3.10 \
-            '-PloadTest.args=${{ env.beam_Inference_Python_Benchmarks_Dataflow_test_arguments_2 }} --job_name=benchmark-tests-tf-mnist-classification-python-${{env.NOW_UTC}} --input_file=gs://apache-beam-ml/testing/inputs/it_mnist_data.csv --output_file=gs://temp-storage-for-end-to-end-tests/wordcount/result_tf_mnist-${{env.NOW_UTC}}.txt --model=gs://apache-beam-ml/models/tensorflow/mnist/' \
\ No newline at end of file
+            -PloadTest.requirementsTxtFile=apache_beam/ml/inference/tensorflow_tests_requirements.txt \
+            '-PloadTest.args=${{ env.beam_Python_Cost_Benchmarks_Dataflow_test_arguments_2 }} --job_name=benchmark-tests-tf-mnist-classification-python-${{env.NOW_UTC}} --input_file=gs://apache-beam-ml/testing/inputs/it_mnist_data.csv --output_file=gs://temp-storage-for-end-to-end-tests/inference/result_tf_mnist-${{env.NOW_UTC}}.txt --model=gs://apache-beam-ml/models/tensorflow/mnist/' \
\ No newline at end of file
diff --git a/sdks/python/apache_beam/testing/benchmarks/inference/tensorflow_mnist_classification_cost_benchmark.py b/sdks/python/apache_beam/testing/benchmarks/inference/tensorflow_mnist_classification_cost_benchmark.py
index f7e12dcead03..223b973e5fbe 100644
--- a/sdks/python/apache_beam/testing/benchmarks/inference/tensorflow_mnist_classification_cost_benchmark.py
+++ b/sdks/python/apache_beam/testing/benchmarks/inference/tensorflow_mnist_classification_cost_benchmark.py
@@ -31,7 +31,7 @@ def test(self):
     extra_opts['input'] = self.pipeline.get_option('input_file')
     extra_opts['output'] = self.pipeline.get_option('output_file')
     extra_opts['model_path'] = self.pipeline.get_option('model')
-    tensorflow_mnist_classification.run(
+    self.result = tensorflow_mnist_classification.run(
         self.pipeline.get_full_options_as_args(**extra_opts),
         save_main_session=False)