-
Notifications
You must be signed in to change notification settings - Fork 116
/
flintrock.py
1583 lines (1330 loc) · 52.8 KB
/
flintrock.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
858
859
860
861
862
863
864
865
866
867
868
869
870
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
891
892
893
894
895
896
897
898
899
900
901
902
903
904
905
906
907
908
909
910
911
912
913
914
915
916
917
918
919
920
921
922
923
924
925
926
927
928
929
930
931
932
933
934
935
936
937
938
939
940
941
942
943
944
945
946
947
948
949
950
951
952
953
954
955
956
957
958
959
960
961
962
963
964
965
966
967
968
969
970
971
972
973
974
975
976
977
978
979
980
981
982
983
984
985
986
987
988
989
990
991
992
993
994
995
996
997
998
999
1000
"""
Flintrock
A command-line tool and library for launching Apache Spark clusters.
Major TODOs:
* Handling of exceptions / reporting of issues during cluster launch.
- Spark install goes wrong
- Spark version is invalid
- Current exception output is quite ugly. Related to thread executor / asyncio.
* "Fix" Hadoop 2.6 S3 setup by installing appropriate Hadoop libraries
See: https://issues.apache.org/jira/browse/SPARK-7442
* ClusterInfo namedtuple -> FlintrockCluster class
- Platform-specific (e.g. EC2) implementations of class add methods to
stop, start, describe (with YAML output) etc. clusters
- Implement method that takes cluster name and returns FlintrockCluster
* Support submit command for Spark applications. Like a wrapper around spark-submit. (?)
* Check that EC2 enhanced networking is enabled.
Other TODOs:
* Instance type <-> AMI type validation/lookup.
- Maybe this can be automated.
- Otherwise have a separate YAML file with this info.
- Maybe support HVM only. AWS seems to position it as the future.
- Show friendly error when user tries to launch PV instance type.
* Use IAM roles to launch instead of AWS keys.
* Setup and teardown VPC, routes, gateway, etc. from scratch.
* Use SSHAgent instead of .pem files (?).
* Automatically replace failed instances during launch, perhaps up to a
certain limit (1-2 instances).
* Upgrade check -- Is a newer version of Flintrock available on PyPI?
* Credits command, for crediting contributors. (?)
Distant future:
* Local provider
"""
import os
import posixpath
import errno
import sys
import shlex
import subprocess
import pprint
import asyncio
import functools
import itertools
import socket
import json
import time
import urllib.request
import tempfile
import textwrap
from datetime import datetime
from collections import namedtuple
# External modules.
import boto
import boto.ec2
import click
import paramiko
import yaml
_SCRIPT_DIR = os.path.dirname(os.path.realpath(__file__))
def timeit(func):
@functools.wraps(func)
def wrapper(*args, **kwargs):
start = datetime.now().replace(microsecond=0)
res = func(*args, **kwargs)
end = datetime.now().replace(microsecond=0)
print("{f} finished in {t}.".format(f=func.__name__, t=(end - start)))
return res
return wrapper
def generate_ssh_key_pair() -> namedtuple('KeyPair', ['public', 'private']):
"""
Generate an SSH key pair that the cluster can use for intra-cluster
communication.
"""
with tempfile.TemporaryDirectory() as tempdir:
ret = subprocess.check_call(
"""
ssh-keygen -q -t rsa -N '' -f {key_file} -C flintrock
""".format(
key_file=shlex.quote(tempdir + "/flintrock_rsa")),
shell=True)
with open(file=tempdir + "/flintrock_rsa") as private_key_file:
private_key = private_key_file.read()
with open(file=tempdir + "/flintrock_rsa.pub") as public_key_file:
public_key = public_key_file.read()
return namedtuple('KeyPair', ['public', 'private'])(public_key, private_key)
# TODO: Think about extending this to represent everything that defines a cluster.
# * name
# * installed modules (?)
# * etc.
#
# Convert it into a class with variations (implementations?) for the specific
# providers.
#
# Add class methods to start, stop, destroy, and describe clusters.
ClusterInfo = namedtuple(
'ClusterInfo', [
'name',
'ssh_key_pair',
'master_host',
'slave_hosts',
'spark_scratch_dir',
'spark_master_opts'
])
# TODO: Cache these files. (?) They are being read potentially tens or
# hundreds of times. Maybe it doesn't matter because the files
# are so small.
# NOTE: functools.lru_cache() doesn't work here because the mapping is
# not hashable.
def get_formatted_template(path: str, mapping: dict) -> str:
with open(path) as f:
formatted = f.read().format(**mapping)
return formatted
# TODO: Turn this into an implementation of an abstract FlintrockModule class. (?)
class Spark:
def __init__(self, version):
self.version = version
def install(
self,
ssh_client: paramiko.client.SSHClient,
cluster_info: ClusterInfo):
"""
Downloads and installs Spark on a given node.
"""
# TODO: Allow users to specify the Spark "distribution".
distribution = 'hadoop1'
print("[{h}] Installing Spark...".format(
h=ssh_client.get_transport().getpeername()[0]))
try:
with ssh_client.open_sftp() as sftp:
sftp.put(
localpath='./install-spark.sh',
remotepath='/tmp/install-spark.sh')
sftp.chmod(path='/tmp/install-spark.sh', mode=0o755)
ssh_check_output(
client=ssh_client,
command="""
set -e
/tmp/install-spark.sh {spark_version} {distribution} {spark_scratch_dir}
rm -f /tmp/install-spark.sh
""".format(
spark_version=shlex.quote(self.version),
distribution=shlex.quote(distribution),
spark_scratch_dir=shlex.quote(cluster_info.spark_scratch_dir)))
except Exception as e:
# TODO: This should be a more specific exception.
print(
"Could not find package for Spark {s} / {d}.".format(
s=self.version,
d=distribution),
file=sys.stderr)
raise
def configure(
self,
ssh_client: paramiko.client.SSHClient,
cluster_info: ClusterInfo):
"""
Configures Spark after it's installed.
This method is master/slave-agnostic.
"""
template_path = "./spark/conf/spark-env.sh"
ssh_check_output(
client=ssh_client,
command="""
echo {f} > {p}
""".format(
f=shlex.quote(
get_formatted_template(
path="templates/" + template_path,
mapping=vars(cluster_info))),
p=shlex.quote(template_path)))
def configure_master(
self,
ssh_client: paramiko.client.SSHClient,
cluster_info: ClusterInfo):
"""
Configures the Spark master and starts both the master and slaves.
"""
host = ssh_client.get_transport().getpeername()[0]
print("[{h}] Configuring Spark master...".format(h=host))
# TODO: Maybe move this shell script out to some separate file/folder
# for the Spark module.
# TODO: Add some timeout for waiting on master UI to come up.
ssh_check_output(
client=ssh_client,
command="""
set -e
echo {s} > spark/conf/slaves
spark/sbin/start-master.sh
set +e
master_ui_response_code=0
while [ "$master_ui_response_code" -ne 200 ]; do
sleep 1
master_ui_response_code="$(
curl --head --silent --output /dev/null \
--write-out "%{{http_code}}" {m}:8080
)"
done
set -e
spark/sbin/start-slaves.sh
""".format(
s=shlex.quote('\n'.join(cluster_info.slave_hosts)),
m=shlex.quote(cluster_info.master_host)))
def configure_slave(self):
pass
def health_check(self, master_host: str):
"""
Check that Spark is functioning.
"""
spark_master_ui = 'http://{m}:8080/json/'.format(m=master_host)
try:
spark_ui_info = json.loads(
urllib.request.urlopen(spark_master_ui).read().decode('utf-8'))
except Exception as e:
# TODO: Catch a more specific problem known to be related to Spark not
# being up; provide a slightly better error message, and don't
# dump a large stack trace on the user.
print("Spark health check failed.", file=sys.stderr)
raise
print(textwrap.dedent(
"""\
Spark Health Report:
* Master: {status}
* Workers: {workers}
* Cores: {cores}
* Memory: {memory:.1f} GB\
""".format(
status=spark_ui_info['status'],
workers=len(spark_ui_info['workers']),
cores=spark_ui_info['cores'],
memory=spark_ui_info['memory'] / 1024)))
@click.group()
@click.option('--config', default=_SCRIPT_DIR + '/config.yaml')
@click.option('--provider', default='ec2', type=click.Choice(['ec2']))
@click.version_option(version='dev') # TODO: Replace with setuptools auto-detect.
@click.pass_context
def cli(cli_context, config, provider):
"""
Flintrock
A command-line tool and library for launching Apache Spark clusters.
"""
cli_context.obj['provider'] = provider
if os.path.exists(config):
with open(config) as f:
raw_config = yaml.safe_load(f)
config_map = config_to_click(normalize_keys(raw_config))
cli_context.default_map = config_map
else:
if config != (_SCRIPT_DIR + '/config.yaml'):
raise FileNotFoundError(errno.ENOENT, 'No such file or directory', config)
@cli.command()
@click.argument('cluster-name')
@click.option('--num-slaves', type=int, required=True)
@click.option('--install-spark/--no-install-spark', default=True)
@click.option('--spark-version')
@click.option('--ec2-key-name')
@click.option('--ec2-identity-file',
type=click.Path(exists=True, dir_okay=False),
help="Path to SSH .pem file for accessing nodes.")
@click.option('--ec2-instance-type', default='m3.medium', show_default=True)
@click.option('--ec2-region', default='us-east-1', show_default=True)
@click.option('--ec2-availability-zone')
@click.option('--ec2-ami')
@click.option('--ec2-user')
@click.option('--ec2-spot-price', type=float)
@click.option('--ec2-vpc-id')
@click.option('--ec2-subnet-id')
@click.option('--ec2-placement-group')
@click.option('--ec2-tenancy', default='default')
@click.option('--ec2-ebs-optimized/--no-ec2-ebs-optimized', default=False)
@click.option('--ec2-instance-initiated-shutdown-behavior', default='stop',
type=click.Choice(['stop', 'terminate']))
@click.pass_context
def launch(
cli_context,
cluster_name, num_slaves,
install_spark,
spark_version,
ec2_key_name,
ec2_identity_file,
ec2_instance_type,
ec2_region,
ec2_availability_zone,
ec2_ami,
ec2_user,
ec2_spot_price,
ec2_vpc_id,
ec2_subnet_id,
ec2_placement_group,
ec2_tenancy,
ec2_ebs_optimized,
ec2_instance_initiated_shutdown_behavior):
"""
Launch a new cluster.
"""
modules = []
if install_spark:
spark = Spark(version=spark_version)
modules += [spark]
if cli_context.obj['provider'] == 'ec2':
return launch_ec2(
cluster_name=cluster_name, num_slaves=num_slaves, modules=modules,
key_name=ec2_key_name,
identity_file=ec2_identity_file,
instance_type=ec2_instance_type,
region=ec2_region,
availability_zone=ec2_availability_zone,
ami=ec2_ami,
user=ec2_user,
spot_price=ec2_spot_price,
vpc_id=ec2_vpc_id,
subnet_id=ec2_subnet_id,
placement_group=ec2_placement_group,
tenancy=ec2_tenancy,
ebs_optimized=ec2_ebs_optimized,
instance_initiated_shutdown_behavior=ec2_instance_initiated_shutdown_behavior)
else:
raise Exception("This provider is not supported: {p}".format(p=cli_context.obj['provider']))
def get_or_create_ec2_security_groups(
*,
cluster_name,
vpc_id,
region) -> 'List[boto.ec2.securitygroup.SecurityGroup]':
"""
If they do not already exist, create all the security groups needed for a
Flintrock cluster.
"""
connection = boto.ec2.connect_to_region(region_name=region)
SecurityGroupRule = namedtuple(
'SecurityGroupRule', [
'ip_protocol',
'from_port',
'to_port',
'src_group',
'cidr_ip'])
# TODO: Make these into methods, since we need this logic (though simple)
# in multiple places. (?)
flintrock_group_name = 'flintrock'
cluster_group_name = 'flintrock-' + cluster_name
search_results = connection.get_all_security_groups(
filters={
'group-name': [flintrock_group_name, cluster_group_name]
})
# The Flintrock group is common to all Flintrock clusters and authorizes client traffic
# to them.
flintrock_group = next((sg for sg in search_results if sg.name == flintrock_group_name), None)
# The cluster group is specific to one Flintrock cluster and authorizes intra-cluster
# communication.
cluster_group = next((sg for sg in search_results if sg.name == cluster_group_name), None)
if not flintrock_group:
flintrock_group = connection.create_security_group(
name=flintrock_group_name,
description="flintrock base group",
vpc_id=vpc_id)
# Rules for the client interacting with the cluster.
flintrock_client_ip = (
urllib.request.urlopen('http://checkip.amazonaws.com/')
.read().decode('utf-8').strip())
flintrock_client_cidr = '{ip}/32'.format(ip=flintrock_client_ip)
client_rules = [
SecurityGroupRule(
ip_protocol='tcp',
from_port=22,
to_port=22,
cidr_ip=flintrock_client_cidr,
src_group=None),
SecurityGroupRule(
ip_protocol='tcp',
from_port=8080,
to_port=8081,
cidr_ip=flintrock_client_cidr,
src_group=None),
SecurityGroupRule(
ip_protocol='tcp',
from_port=4040,
to_port=4040,
cidr_ip=flintrock_client_cidr,
src_group=None)
]
# TODO: Don't try adding rules that already exist.
# TODO: Add rules in one shot.
for rule in client_rules:
try:
flintrock_group.authorize(**vars(rule))
except boto.exception.EC2ResponseError as e:
if e.error_code != 'InvalidPermission.Duplicate':
print("Error adding rule: {r}".format(r=rule))
raise
# Rules for internal cluster communication.
if not cluster_group:
cluster_group = connection.create_security_group(
name=cluster_group_name,
description="Flintrock cluster group",
vpc_id=vpc_id)
cluster_rules = [
SecurityGroupRule(
ip_protocol='icmp',
from_port=-1,
to_port=-1,
src_group=cluster_group,
cidr_ip=None),
SecurityGroupRule(
ip_protocol='tcp',
from_port=0,
to_port=65535,
src_group=cluster_group,
cidr_ip=None),
SecurityGroupRule(
ip_protocol='udp',
from_port=0,
to_port=65535,
src_group=cluster_group,
cidr_ip=None)
]
# TODO: Don't try adding rules that already exist.
# TODO: Add rules in one shot.
for rule in cluster_rules:
try:
cluster_group.authorize(**vars(rule))
except boto.exception.EC2ResponseError as e:
if e.error_code != 'InvalidPermission.Duplicate':
print("Error adding rule: {r}".format(r=rule))
raise
return [flintrock_group, cluster_group]
# Move to ec2 module and call as ec2.launch()?
@timeit
def launch_ec2(
*,
cluster_name, num_slaves, modules,
key_name, identity_file,
instance_type,
region,
availability_zone,
ami,
user,
spot_price=None,
vpc_id, subnet_id, placement_group,
tenancy="default", ebs_optimized=False,
instance_initiated_shutdown_behavior="stop"):
"""
Launch a fully functional cluster on EC2 with the specified configuration
and installed modules.
"""
try:
get_cluster_instances_ec2(
cluster_name=cluster_name,
region=region)
except ClusterNotFound as e:
pass
else:
print("Cluster already exists: {c}".format(c=cluster_name), file=sys.stderr)
sys.exit(1)
security_groups = get_or_create_ec2_security_groups(
cluster_name=cluster_name,
vpc_id=vpc_id,
region=region)
num_instances = num_slaves + 1
spot_requests = []
cluster_instances = []
try:
connection = boto.ec2.connect_to_region(region_name=region)
if spot_price:
print("Requesting {c} spot instances at a max price of ${p}...".format(
c=num_instances, p=spot_price))
spot_requests = connection.request_spot_instances(
price=spot_price,
image_id=ami,
count=num_instances,
key_name=key_name,
instance_type=instance_type,
placement=availability_zone,
security_group_ids=[sg.id for sg in security_groups],
subnet_id=subnet_id,
placement_group=placement_group,
ebs_optimized=ebs_optimized)
request_ids = [r.id for r in spot_requests]
pending_request_ids = request_ids
while pending_request_ids:
print("{grant} of {req} instances granted. Waiting...".format(
grant=num_instances - len(pending_request_ids),
req=num_instances))
time.sleep(30)
spot_requests = connection.get_all_spot_instance_requests(request_ids=request_ids)
pending_request_ids = [r.id for r in spot_requests if r.state != 'active']
print("All {c} instances granted.".format(c=num_instances))
cluster_instances = connection.get_only_instances(
instance_ids=[r.instance_id for r in spot_requests])
else:
print("Launching {c} instances...".format(c=num_instances))
reservation = connection.run_instances(
image_id=ami,
min_count=num_instances,
max_count=num_instances,
key_name=key_name,
instance_type=instance_type,
placement=availability_zone,
security_group_ids=[sg.id for sg in security_groups],
subnet_id=subnet_id,
placement_group=placement_group,
tenancy=tenancy,
ebs_optimized=ebs_optimized,
instance_initiated_shutdown_behavior=instance_initiated_shutdown_behavior)
cluster_instances = reservation.instances
time.sleep(10) # AWS metadata eventual consistency tax.
wait_for_cluster_state_ec2(
cluster_instances=cluster_instances,
state='running')
master_instance = cluster_instances[0]
slave_instances = cluster_instances[1:]
connection.create_tags(
resource_ids=[master_instance.id],
tags={
'flintrock-role': 'master',
'Name': '{c}-master'.format(c=cluster_name)})
connection.create_tags(
resource_ids=[i.id for i in slave_instances],
tags={
'flintrock-role': 'slave',
'Name': '{c}-slave'.format(c=cluster_name)})
cluster_info = ClusterInfo(
name=cluster_name,
ssh_key_pair=generate_ssh_key_pair(),
master_host=master_instance.public_dns_name,
slave_hosts=[instance.public_dns_name for instance in slave_instances],
spark_scratch_dir='/mnt/spark',
spark_master_opts="")
# TODO: Abstract away. No-one wants to see this async shite here.
loop = asyncio.get_event_loop()
tasks = []
for instance in cluster_instances:
# TODO: Use parameter names for run_in_executor() once Python 3.4.4 is released.
# Until then, we leave them out to maintain compatibility across Python 3.4
# and 3.5.
# See: http://stackoverflow.com/q/32873974/
task = loop.run_in_executor(
None,
functools.partial(
provision_node,
modules=modules,
user=user,
host=instance.ip_address,
identity_file=identity_file,
cluster_info=cluster_info))
tasks.append(task)
done, _ = loop.run_until_complete(asyncio.wait(tasks))
# Is this the right way to make sure no coroutine failed?
for future in done:
future.result()
loop.close()
print("All {c} instances provisioned.".format(
c=len(cluster_instances)))
master_ssh_client = get_ssh_client(
user=user,
host=master_instance.ip_address,
identity_file=identity_file)
with master_ssh_client:
for module in modules:
module.configure_master(
ssh_client=master_ssh_client,
cluster_info=cluster_info)
# NOTE: We sleep here so that Spark (currently the only supported module)
# has time to spin up all its slaves.
# TODO: Spark module methods to start_slave() and start_master() which are
# separate from configure_master() and which block until Spark services
# on that node are fully running.
time.sleep(30)
module.health_check(master_host=cluster_info.master_host)
except (Exception, KeyboardInterrupt) as e:
print(e, file=sys.stderr)
if spot_requests:
# TODO: Do this only if there are pending requests.
print("Canceling spot instance requests...", file=sys.stderr)
request_ids = [r.id for r in spot_requests]
connection.cancel_spot_instance_requests(
request_ids=request_ids)
# Make sure we have the latest information on any launched spot instances.
spot_requests = connection.get_all_spot_instance_requests(
request_ids=request_ids)
instance_ids = [r.instance_id for r in spot_requests if r.instance_id]
if instance_ids:
cluster_instances = connection.get_only_instances(
instance_ids=instance_ids)
if cluster_instances:
yes = click.confirm(
text="Do you want to terminate the {c} instances created by this operation?"
.format(c=len(cluster_instances)),
err=True,
default=True)
if yes:
print("Terminating instances...", file=sys.stderr)
connection.terminate_instances(
instance_ids=[instance.id for instance in cluster_instances])
sys.exit(1)
# finally:
# print("Terminating all {c} instances...".format(
# c=len(cluster_instances)))
# connection.terminate_instances(
# instance_ids=[instance.id for instance in cluster_instances])
def get_ssh_client(
*,
user: str,
host: str,
identity_file: str,
print_status: bool=False) -> paramiko.client.SSHClient:
"""
Get an SSH client for the provided host, waiting as necessary for SSH to become available.
"""
# paramiko.common.logging.basicConfig(level=paramiko.common.DEBUG)
client = paramiko.client.SSHClient()
client.load_system_host_keys()
client.set_missing_host_key_policy(paramiko.client.AutoAddPolicy())
while True:
try:
client.connect(
username=user,
hostname=host,
key_filename=identity_file,
look_for_keys=False,
timeout=3)
if print_status:
print("[{h}] SSH online.".format(h=host))
break
# TODO: Somehow rationalize these expected exceptions.
# TODO: Add some kind of limit on number of failures.
except socket.timeout as e:
time.sleep(5)
except socket.error as e:
if e.errno != 61:
raise
time.sleep(5)
# We get this exception during startup with CentOS but not Amazon Linux,
# for some reason.
except paramiko.ssh_exception.AuthenticationException as e:
time.sleep(5)
return client
def provision_node(
*,
modules: list,
user: str,
host: str,
identity_file: str,
cluster_info: ClusterInfo):
"""
Connect to a freshly launched node, set it up for SSH access, and
install the specified modules.
This function is intended to be called on all cluster nodes in parallel.
No master- or slave-specific logic should be in this method.
"""
client = get_ssh_client(
user=user,
host=host,
identity_file=identity_file,
print_status=True)
with client:
ssh_check_output(
client=client,
command="""
set -e
echo {private_key} > ~/.ssh/id_rsa
echo {public_key} >> ~/.ssh/authorized_keys
chmod 400 ~/.ssh/id_rsa
""".format(
private_key=shlex.quote(cluster_info.ssh_key_pair.private),
public_key=shlex.quote(cluster_info.ssh_key_pair.public)))
# The default CentOS AMIs on EC2 don't come with Java installed.
java_home = ssh_check_output(
client=client,
command="""
echo "$JAVA_HOME"
""")
if not java_home.strip():
print("[{h}] Installing Java...".format(h=host))
ssh_check_output(
client=client,
command="""
set -e
sudo yum install -y java-1.7.0-openjdk
sudo sh -c "echo export JAVA_HOME=/usr/lib/jvm/jre >> /etc/environment"
source /etc/environment
""")
for module in modules:
module.install(
ssh_client=client,
cluster_info=cluster_info)
module.configure(
ssh_client=client,
cluster_info=cluster_info)
def ssh_check_output(client: paramiko.client.SSHClient, command: str):
"""
Run a command via the provided SSH client and return the output captured
on stdout.
Raise an exception if the command returns a non-zero code.
"""
stdin, stdout, stderr = client.exec_command(command, get_pty=True)
exit_status = stdout.channel.recv_exit_status()
if exit_status:
# TODO: Return a custom exception that includes the return code.
# See: https://docs.python.org/3/library/subprocess.html#subprocess.check_output
# NOTE: We are losing the output order here since output from stdout and stderr
# may be interleaved.
raise Exception(
stdout.read().decode("utf8").rstrip('\n') +
stderr.read().decode("utf8").rstrip('\n'))
return stdout.read().decode("utf8").rstrip('\n')
class ClusterNotFound(Exception):
pass
def get_cluster_instances_ec2(
*,
cluster_name: str,
region: str) -> (boto.ec2.instance.Instance, list):
"""
Get the instances for an EC2 cluster.
"""
connection = boto.ec2.connect_to_region(region_name=region)
cluster_instances = connection.get_only_instances(
filters={
'instance.group-name': 'flintrock-' + cluster_name
})
if not cluster_instances:
raise ClusterNotFound("No such cluster: {c}".format(c=cluster_name))
# TODO: Raise errors if a cluster has no master or no slaves.
master_instance = list(filter(
lambda i: i.tags['flintrock-role'] == 'master',
cluster_instances))[0]
slave_instances = list(filter(
lambda i: i.tags['flintrock-role'] != 'master',
cluster_instances))
return master_instance, slave_instances
@cli.command()
@click.argument('cluster-name')
# @click.confirmation_option(help="Are you sure you want to destroy this cluster?")
@click.option('--assume-yes/--no-assume-yes', default=False)
@click.option('--ec2-region', default='us-east-1', show_default=True)
# TODO: Always delete cluster security group. People shouldn't be adding stuff to it.
# Instead, provide option for cluster to be assigned to additional, pre-existing
# security groups.
@click.pass_context
def destroy(cli_context, cluster_name, assume_yes, ec2_region):
"""
Destroy a cluster.
"""
if cli_context.obj['provider'] == 'ec2':
destroy_ec2(
cluster_name=cluster_name,
assume_yes=assume_yes,
region=ec2_region)
else:
# TODO: Create UnsupportedProviderException. (?)
raise Exception("This provider is not supported: {p}".format(p=cli_context.obj['provider']))
# assume_yes defaults to True here for library use (as opposed to command-line use,
# where the default is configured via Click).
def destroy_ec2(*, cluster_name, assume_yes=True, region):
try:
master_instance, slave_instances = get_cluster_instances_ec2(
cluster_name=cluster_name,
region=region)
cluster_instances = [master_instance] + slave_instances
except ClusterNotFound as e:
print(e, file=sys.stderr)
sys.exit(1)
if not assume_yes:
print_cluster_info_ec2(
cluster_name=cluster_name,
cluster_instances=cluster_instances)
print('---')
click.confirm(
text="Are you sure you want to destroy this cluster?",
abort=True)
# TODO: Figure out if we want to use "node" instead of "instance" when
# communicating with the user, even if we're talking about doing things
# to EC2 instances. Spark docs definitely favor "node".
print("Terminating {c} instances...".format(c=len(cluster_instances)))
connection = boto.ec2.connect_to_region(region_name=region)
connection.terminate_instances(
instance_ids=[instance.id for instance in cluster_instances])
# TODO: Destroy cluster security group. We're not reusing it.
def add_slaves(provider, cluster_name, num_slaves, provider_options):
# Need concept of cluster state so we can add slaves with the same config.
# Otherwise we must ask unreliable user to respecify slave config.
pass
def add_slaves_ec2(cluster_name, num_slaves, identity_file):
pass
def remove_slaves(provider, cluster_name, num_slaves, provider_options, assume_yes=False):
pass
def remove_slaves_ec2(cluster_name, num_slaves, assume_yes=True):
pass
def wait_for_cluster_state_ec2(*, cluster_instances: list, state: str):
"""
Wait for all the instances in a cluster to reach a specific state.
"""
while True:
for instance in cluster_instances:
if instance.state == state:
continue
else:
time.sleep(3)
instance.update()
break
else:
break
def get_cluster_state_ec2(cluster_instances: list) -> str:
"""
Get the state of an EC2 cluster.
This is distinct from the state of Spark on the cluster. At some point the two
concepts should be rationalized somehow.
"""
instance_states = set(instance.state for instance in cluster_instances)
if len(instance_states) == 1:
return next(iter(instance_states))
else:
return 'inconsistent'
def print_cluster_info_ec2(cluster_name: str, cluster_instances: list):
"""
Print information about an EC2 cluster to screen in a YAML-compatible format.
This is the current solution until cluster methods are centralized under a
FlintrockCluster class, or something similar.
"""
print(cluster_name + ':')
print(' state: {s}'.format(s=get_cluster_state_ec2(cluster_instances=cluster_instances)))
print(' node-count: {nc}'.format(nc=len(cluster_instances)))
if get_cluster_state_ec2(cluster_instances=cluster_instances) == 'running':
print('\n - '.join([' nodes:'] + [i.public_dns_name for i in cluster_instances]))
@cli.command()
@click.argument('cluster-name', required=False)
@click.option('--master-hostname-only', is_flag=True, default=False)
@click.option('--ec2-region')
@click.pass_context
def describe(
cli_context,
cluster_name,
master_hostname_only,
ec2_region):
"""
Describe an existing cluster.
Leave out the cluster name to find all Flintrock-managed clusters.
"""
if cli_context.obj['provider'] == 'ec2':
describe_ec2(
cluster_name=cluster_name,
master_hostname_only=master_hostname_only,
region=ec2_region)
else:
# TODO: Create UnsupportedProviderException. (?)
raise Exception("This provider is not supported: {p}".format(p=cli_context.obj['provider']))
def describe_ec2(*, cluster_name, master_hostname_only=False, region):
if cluster_name:
try:
master_instance, slave_instances = get_cluster_instances_ec2(