diff --git a/custom_image_utils/args_parser.py b/custom_image_utils/args_parser.py index 05408e4..8433f48 100644 --- a/custom_image_utils/args_parser.py +++ b/custom_image_utils/args_parser.py @@ -47,7 +47,7 @@ def _full_image_family_uri_regex_type(s): """Check if the partial image family uri string matches regex.""" if not _FULL_IMAGE_FAMILY_URI.match(s): raise argparse.ArgumentTypeError("Invalid image family URI: {}.".format(s)) - return s + return s def parse_args(args): """Parses command-line arguments.""" @@ -75,7 +75,7 @@ def parse_args(args): "--base-image-family", type=_full_image_family_uri_regex_type, help="""The source image family URI. The latest non-depracated image associated with the family will be used. - """) + """) required_args.add_argument( "--customization-script", type=str, @@ -154,6 +154,13 @@ def parse_args(args): that builds the custom image. A full subnetwork URL is required. Default subnetwork is None. For shared VPC only provide this parameter and do not use the --network argument.""") + parser.add_argument( + "--smoke-test-tags", + type=str, + required=False, + default="", + help="""(Optional) Network tags used to launch the Dataproc cluster used + for the smoke test.""") parser.add_argument( "--no-external-ip", action="store_true", diff --git a/custom_image_utils/smoke_test_runner.py b/custom_image_utils/smoke_test_runner.py index f6ba417..13b665a 100644 --- a/custom_image_utils/smoke_test_runner.py +++ b/custom_image_utils/smoke_test_runner.py @@ -24,7 +24,7 @@ _LOG.setLevel(logging.INFO) def _create_workflow_template(workflow_name, image_name, project_id, zone, region, - network, subnet, no_external_ip): + network, subnet, no_external_ip, tags): """Create a Dataproc workflow template for testing.""" create_command = [ "gcloud", "dataproc", "workflow-templates", "create", @@ -41,6 +41,8 @@ def _create_workflow_template(workflow_name, image_name, project_id, zone, regio set_cluster_command.extend(["--subnet", subnet]) if no_external_ip: set_cluster_command.extend(["--no-address"]) + if tags: + set_cluster_command.extend(["--tags", tags]) add_job_command = [ "gcloud", "dataproc", "workflow-templates", "add-job", "spark", "--workflow-template", workflow_name, "--project", project_id, "--region", region, @@ -92,7 +94,7 @@ def _delete_workflow_template(workflow_name, project_id, region): raise RuntimeError("Error deleting workfloe template %s.", workflow_name) -def _verify_custom_image(image_name, project_id, zone, network, subnetwork, no_external_ip): +def _verify_custom_image(image_name, project_id, zone, network, subnetwork, no_external_ip, tags): """Verifies if custom image works with Dataproc.""" region = zone[:-2] date = datetime.datetime.now().strftime("%Y%m%d%H%M%S") @@ -103,7 +105,7 @@ def _verify_custom_image(image_name, project_id, zone, network, subnetwork, no_e _LOG.info("Creating Dataproc workflow-template %s with image %s...", workflow_name, image_name) _create_workflow_template(workflow_name, image_name, project_id, zone, region, - network, subnetwork, no_external_ip) + network, subnetwork, no_external_ip, tags) _LOG.info( "Successfully created Dataproc workflow-template %s with image %s...", workflow_name, image_name) @@ -133,7 +135,8 @@ def run(args): if not args.no_smoke_test: _LOG.info("Verifying the custom image...") _verify_custom_image(args.image_name, args.project_id, args.zone, - args.network, args.subnetwork, args.no_external_ip) + args.network, args.subnetwork, args.no_external_ip, + args.smoke_test_tags) _LOG.info("Successfully verified the custom image...") else: _LOG.info("Skip running smoke test (dry run).")