[testing-on-gke part independent 6.8] Add gsheet utilities

perfmetrics/scripts/testing_on_gke/examples/dlio/run_tests.py

@@ -30,8 +30,8 @@
 
 # local imports from other directories
 sys.path.append(os.path.join(os.path.dirname(__file__), '..', 'utils'))
-from run_tests_common import escape_commas_in_string, parse_args, run_command, add_iam_role_for_buckets
-from utils import UnknownMachineTypeError, resource_limits
+from run_tests_common import escape_commas_in_string, parse_args, add_iam_role_for_buckets
+from utils import UnknownMachineTypeError, resource_limits, run_command
 
 # local imports from same directory
 import dlio_workload

perfmetrics/scripts/testing_on_gke/examples/fio/run_tests.py

@@ -29,8 +29,8 @@
 
 # local imports from other directories
 sys.path.append(os.path.join(os.path.dirname(__file__), '..', 'utils'))
-from run_tests_common import escape_commas_in_string, parse_args, run_command, add_iam_role_for_buckets
-from utils import UnknownMachineTypeError, resource_limits
+from run_tests_common import escape_commas_in_string, parse_args, add_iam_role_for_buckets
+from utils import UnknownMachineTypeError, resource_limits, run_command
 
 # local imports from same directory
 import fio_workload

perfmetrics/scripts/testing_on_gke/examples/utils/gsheet.py

@@ -0,0 +1,141 @@
+# Copyright 2024 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import os
+import tempfile
+from typing import Tuple
+from google.oauth2 import service_account
+from googleapiclient.discovery import build
+from utils.utils import run_command
+
+_SCOPES = ['https://www.googleapis.com/auth/spreadsheets']
+
+
+def _get_sheets_service_client(localServiceAccountKeyFile):
+  creds = service_account.Credentials.from_service_account_file(
+      localServiceAccountKeyFile, scopes=_SCOPES
+  )
+  # Alternatively, use from_service_account_info for client-creation,
+  # documented at
+  # https://google-auth.readthedocs.io/en/master/reference/google.oauth2.service_account.html
+  # .
+  client = build('sheets', 'v4', credentials=creds)
+  return client
+
+
+def download_gcs_object_locally(gcsObjectUri: str) -> str:
+  """Downloads the given gcs file-object to a temporary local file (collision-free) and returns the full-path of this local-file.
+
+  gcsObjectUri is of the form: gs://<bucket-name>/object-name .
+  On failure, returned int will be non-zero.
+
+  Caller has to delete the tempfile after usage if a proper tempfile was
+  returned, to avoid disk-leak.
+  """
+  if not gcsObjectUri.startswith('gs:'):
+    raise ValueError(
+        f'Passed input {gcsObjectUri} is not proper. Expected:'
+        ' gs://<bucket>/<object-name>'
+    )
+  with tempfile.NamedTemporaryFile(mode='w+b', dir='/tmp', delete=False) as fp:
+    returncode = run_command(f'gcloud storage cp {gcsObjectUri} {fp.name}')
+    if returncode == 0:
+      return fp.name
+    else:
+      raise Exception(
+          f'failed to copy gcs object {gcsObjectUri} to local-file {fp.name}:'
+          f' returncode={returncode}. Deleting tempfile {fp.name}...'
+      )
+      os.remove(fp.name)
+
+
+def append_data_to_gsheet(
+    serviceAccountKeyFile: str,
+    gsheet_id: str,
+    worksheet: str,
+    data: dict,
+) -> None:
+  """Calls the API to append the given data at the end of the given worksheet in the given gsheet.
+
+  If the passed header matches the first row of the file, then the
+  header is not inserted again.
+  Args:
+    serviceAccountKeyFile: Path of a service-account key-file for authentication
+      read/write from/to the given gsheet. This can be a local filepath or a GCS
+      path starting with `gs:`
+    worksheet: string, name of the worksheet to be edited appended by a "!"
+    data: Dictionary of {'header': tuple, 'values': list(tuples)}, to be added
+      to the worksheet.
+    gsheet_id: Unique ID to identify a gsheet.
+
+  Raises:
+    HttpError: For any Google Sheets API call related errors
+  """
+  for arg in [serviceAccountKeyFile, worksheet, gsheet_id]:
+    if not arg or not arg.strip():
+      raise ValueError(
+          f"Passed  argument '{serviceAccountKeyFile}' is not proper"
+      )
+
+  def _using_local_service_key_file(localServiceAccountKeyFile: str):
+    # Open a read-write gsheet client.
+    client = _get_sheets_service_client(localServiceAccountKeyFile)
+
+    def _read_from_range(cell_range: str):
+      """Returns a list of list of values for the given range in the worksheet."""
+      gsheet_response = (
+          client.spreadsheets()
+          .values()
+          .get(spreadsheetId=gsheet_id, range=f'{worksheet}!{cell_range}')
+          .execute()
+      )
+      return gsheet_response['values'] if 'values' in gsheet_response else []
+
+    def _write_at_address(cell_address: str, data):
+      """Writes a list of tuple of values at the given cell_cell_address in the worksheet."""
+      client.spreadsheets().values().update(
+          spreadsheetId=gsheet_id,
+          valueInputOption='USER_ENTERED',
+          body={'majorDimension': 'ROWS', 'values': data},
+          range=f'{worksheet}!{cell_address}',
+      ).execute()
+
+    data_in_first_column = _read_from_range('A1:A')
+    num_rows = len(data_in_first_column)
+    data_in_first_row = _read_from_range('A1:1')
+    original_header = tuple(data_in_first_row[0]) if data_in_first_row else ()
+    new_header = data['header']
+
+    # Insert header in the file, if needed.
+    if not original_header or not original_header == new_header:
+      # Append header after last row.
+      _write_at_address(f'A{num_rows+1}', [new_header])
+      num_rows = num_rows + 1
+
+    # Append given values after the last row.
+    _write_at_address(f'A{num_rows+1}', data['values'])
+    num_rows = num_rows + 1
+
+  if serviceAccountKeyFile.startswith('gs://'):
+    localServiceAccountKeyFile = download_gcs_object_locally(
+        serviceAccountKeyFile
+    )
+    _using_local_service_key_file(localServiceAccountKeyFile)
+    os.remove(localServiceAccountKeyFile)
+  else:
+    _using_local_service_key_file(serviceAccountKeyFile)
+
+
+def url(gsheet_id: str) -> str:
+  return f'https://docs.google.com/spreadsheets/d/{gsheet_id}'

perfmetrics/scripts/testing_on_gke/examples/utils/gsheet_test.py

@@ -0,0 +1,76 @@
+"""This file defines unit tests for functionalities in utils.py"""
+
+# Copyright 2018 The Kubernetes Authors.
+# Copyright 2024 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import os
+from os import open
+from os.path import isfile
+import random
+from random import choices
+import string
+import unittest
+from gsheet import append_data_to_gsheet, download_gcs_object_locally, url
+
+
+class GsheetTest(unittest.TestCase):
+
+  def test_append_data_to_gsheet(self):
+    _DEFAULT_GSHEET_ID = '1s9DCis6XZ_oHRIFTy0F8yVN93EGA2Koks_pzpCqAIS4'
+
+    def _default_service_account_key_file(project_id: str) -> str:
+      if project_id in ['gcs-fuse-test', 'gcs-fuse-test-ml']:
+        return f'gs://gcsfuse-aiml-test-outputs/creds/{project_id}.json'
+      else:
+        raise Exception(f'Unknown project-id: {project_id}')
+
+    for project_id in ['gcs-fuse-test', 'gcs-fuse-test-ml']:
+      for worksheet in ['fio-test', 'dlio-test']:
+        serviceAccountKeyFile = _default_service_account_key_file(project_id)
+        append_data_to_gsheet(
+            worksheet=worksheet,
+            data={
+                'header': ('Column1', 'Column2'),
+                'values': [(
+                    ''.join(random.choices(string.ascii_letters, k=9)),
+                    random.random(),
+                )],
+            },
+            serviceAccountKeyFile=serviceAccountKeyFile,
+            gsheet_id=_DEFAULT_GSHEET_ID,
+        )
+
+  def test_gsheet_url(self):
+    gsheet_id = ''.join(random.choices(string.ascii_letters, k=20))
+    gsheet_url = url(gsheet_id)
+    self.assertTrue(gsheet_id in gsheet_url)
+    self.assertTrue(len(gsheet_id) < len(gsheet_url))
+
+  def test_download_gcs_object_locally(self):
+    gcs_object = 'gs://gcsfuse-aiml-test-outputs/creds/gcs-fuse-test.json'
+    localfile = download_gcs_object_locally(gcs_object)
+    self.assertTrue(localfile)
+    self.assertTrue(localfile.strip())
+    os.stat(localfile)
+    os.remove(localfile)
+
+  def test_download_gcs_object_locally_nonexistent(self):
+    gcs_object = 'gs://non/existing/gcs/object'
+    with self.assertRaises(Exception):
+      localfile = download_gcs_object_locally(gcs_object)
+
+
+if __name__ == '__main__':
+  unittest.main()

perfmetrics/scripts/testing_on_gke/examples/utils/parse_logs_common.py

@@ -19,6 +19,7 @@
 import os
 import subprocess
 from typing import Tuple
+from utils.utils import run_command
 
 SUPPORTED_SCENARIOS = [
     "local-ssd",

perfmetrics/scripts/testing_on_gke/examples/utils/run_tests_common.py

@@ -18,21 +18,7 @@
 import argparse
 import subprocess
 import sys
-
-
-def run_command(command: str) -> int:
-  """Runs the given string command as a subprocess.
-
-  Returns exit-code which would be non-zero for error.
-  """
-  result = subprocess.run(
-      [word for word in command.split(' ') if (word and not str.isspace(word))],
-      capture_output=True,
-      text=True,
-  )
-  print(result.stdout)
-  print(result.stderr)
-  return result.returncode
+from utils import run_command
 
 
 def escape_commas_in_string(unescapedStr: str) -> str:

perfmetrics/scripts/testing_on_gke/examples/utils/utils.py

@@ -362,3 +362,18 @@ def get_cpu_from_monitoring_api(
       ),
       5,  # round up to 5 decimal places.
   )
+
+
+def run_command(command: str) -> int:
+  """Runs the given string command as a subprocess.
+
+  Returns exit-code which would be non-zero for error.
+  """
+  result = subprocess.run(
+      [word for word in command.split(" ") if (word and word.strip())],
+      capture_output=True,
+      text=True,
+  )
+  print(result.stdout)
+  print(result.stderr)
+  return result.returncode