diff --git a/perfmetrics/scripts/custom_vm_perf_test/custom_vm_perf_script.py b/perfmetrics/scripts/custom_vm_perf_test/custom_vm_perf_script.py new file mode 100644 index 0000000000..220d2bc915 --- /dev/null +++ b/perfmetrics/scripts/custom_vm_perf_test/custom_vm_perf_script.py @@ -0,0 +1,115 @@ +# Copyright 2024 Google Inc. All Rights Reserved. +# +# 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. + +# To automate perf tests, this script creates a VM instance based on the +# flags passed.The command for running the script: +# python3 custom_vm_perf_script.py --vm_name= --machine_type= +# image_family= --image_project= --zone= +# --startup_script= + +# For changing the bucket/spreadsheet id, change values in custom_vm_startup_script.sh + +import argparse +import sys +import subprocess + +DEFAULT_VM_NAME = "perf-tests-vm" +DEFAULT_MACHINE_TYPE = "n2-standard-96" +DEFAULT_IMAGE_FAMILY = "ubuntu-2004-lts" +DEFAULT_IMAGE_PROJECT = "ubuntu-os-cloud" +DEFAULT_ZONE = "us-west1-b" +DEFAULT_STARTUP_SCRIPT = "custom_vm_startup_script.sh" +BOOT_DISK_SIZE = "100GiB" + +def _parse_arguments(argv): + """Parses the arguments provided to the script via command line. + + Args: + argv: List of arguments received by the script. + + Returns: + A class containing the parsed arguments. + """ + + if argv is None: + argv = sys.argv[1:] + + parser = argparse.ArgumentParser() + + parser.add_argument( + '--vm_name', + help='Provide name of the vm instance', + action='store', + default=DEFAULT_VM_NAME, + required=False, + ) + + parser.add_argument( + '--machine_type', + help='Provide machine type of the vm instance', + action='store', + default=DEFAULT_MACHINE_TYPE, + required=False, + ) + + parser.add_argument( + '--image_family', + help='Provide image family of the vm instance', + action='store', + default=DEFAULT_IMAGE_FAMILY, + required=False, + ) + + parser.add_argument( + '--image_project', + help='Provide image project of the vm instance', + action='store', + default=DEFAULT_IMAGE_PROJECT, + required=False, + ) + + parser.add_argument( + '--zone', + help='Provide zone of the vm instance', + action='store', + default=DEFAULT_ZONE, + required=False, + ) + + parser.add_argument( + '--startup_script', + help='Provide startup script for the vm instance', + action='store', + default=DEFAULT_STARTUP_SCRIPT, + required=False, + ) + + return parser.parse_args(argv[1:]) + + +if __name__ == '__main__': + argv = sys.argv + args = _parse_arguments(argv) + # creating vm using gcloud command + try: + subprocess.check_output(f"gcloud compute instances create {args.vm_name}\ + --machine-type={args.machine_type}\ + --image-family={args.image_family}\ + --image-project={args.image_project}\ + --boot-disk-size={BOOT_DISK_SIZE}\ + --zone={args.zone}\ + --metadata-from-file=startup-script={args.startup_script}", + shell=True) + except subprocess.CalledProcessError as e: + print(e.output) diff --git a/perfmetrics/scripts/custom_vm_perf_test/custom_vm_perf_script_test.py b/perfmetrics/scripts/custom_vm_perf_test/custom_vm_perf_script_test.py new file mode 100644 index 0000000000..001d0ae2a1 --- /dev/null +++ b/perfmetrics/scripts/custom_vm_perf_test/custom_vm_perf_script_test.py @@ -0,0 +1,60 @@ +# Copyright 2024 Google Inc. All Rights Reserved. +# +# 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. + +# unit tests for custom_vm_perf_script + +import unittest +import custom_vm_perf_script + +class TestParseArguments(unittest.TestCase): + def test_explicit_values(self): + vm_name = "custom-vm" + machine_type = "n2-standard-32" + image_family = "debian-10" + image_project = "debian-cloud" + zone = "us-west1-a" + startup_script = "custom_startup_script.sh" + + args = custom_vm_perf_script._parse_arguments( + [ + "script", + "--vm_name", vm_name, + "--machine_type", machine_type, + "--image_family", image_family, + "--image_project", image_project, + "--zone", zone, + "--startup_script", startup_script + ], + ) + + self.assertEqual(args.vm_name, vm_name) + self.assertEqual(args.machine_type, machine_type) + self.assertEqual(args.image_family, image_family) + self.assertEqual(args.image_project, image_project) + self.assertEqual(args.zone, zone) + self.assertEqual(args.startup_script, startup_script) + + def test_default_values(self): + args = custom_vm_perf_script._parse_arguments(["script"]) + + self.assertEqual(args.vm_name, custom_vm_perf_script.DEFAULT_VM_NAME) + self.assertEqual(args.machine_type, custom_vm_perf_script.DEFAULT_MACHINE_TYPE) + self.assertEqual(args.image_family, custom_vm_perf_script.DEFAULT_IMAGE_FAMILY) + self.assertEqual(args.image_project, custom_vm_perf_script.DEFAULT_IMAGE_PROJECT) + self.assertEqual(args.zone, custom_vm_perf_script.DEFAULT_ZONE) + self.assertEqual(args.startup_script, custom_vm_perf_script.DEFAULT_STARTUP_SCRIPT) + + +if __name__ == '__main__': + unittest.main()