From 492bb9683ad23307ce7460c97d07010a1bcfe787 Mon Sep 17 00:00:00 2001 From: David Olaru Date: Wed, 27 Apr 2022 19:40:56 +0100 Subject: [PATCH] Add support for Cloud ID to hub model import script The Cloud ID simplifies sending data to a cluster on Elastic Cloud. With this change, the user will have the option specify a Cloud ID using the `--cloud-id` argument as an alternative to an Elasticsearch URL (`--url` argument). `--cloud-id` and `--url` are mutually exclusive arguments. --- bin/eland_import_hub_model | 19 ++++++++++++++++--- 1 file changed, 16 insertions(+), 3 deletions(-) diff --git a/bin/eland_import_hub_model b/bin/eland_import_hub_model index fc044ff..e234007 100755 --- a/bin/eland_import_hub_model +++ b/bin/eland_import_hub_model @@ -39,12 +39,17 @@ MODEL_HUB_URL = "https://huggingface.co" def get_arg_parser(): parser = argparse.ArgumentParser() - parser.add_argument( + location_args = parser.add_mutually_exclusive_group(required=True) + location_args.add_argument( "--url", - required=True, default=os.environ.get("ES_URL"), help="An Elasticsearch connection URL, e.g. http://localhost:9200", ) + location_args.add_argument( + "--cloud-id", + default=os.environ.get("CLOUD_ID"), + help="Cloud ID as found in the 'Manage Deployment' page of an Elastic Cloud deployment", + ) parser.add_argument( "--hub-model-id", required=True, @@ -125,6 +130,14 @@ def get_es_client(cli_args): 'ca_certs': cli_args.ca_certs } + # Deployment location + if cli_args.url: + es_args['hosts'] = cli_args.url + + if cli_args.cloud_id: + es_args['cloud_id'] = cli_args.cloud_id + + # Authentication if cli_args.es_api_key: es_args['api_key'] = cli_args.es_api_key elif cli_args.es_username: @@ -134,7 +147,7 @@ def get_es_client(cli_args): es_args['basic_auth'] = (cli_args.es_username, cli_args.es_password) - es_client = Elasticsearch(args.url, **es_args) + es_client = Elasticsearch(**es_args) es_info = es_client.info() logger.info(f"Connected to cluster named '{es_info['cluster_name']}' (version: {es_info['version']['number']})")