Skip to content

Commit bb43c71

Browse files
exployAlice Berard
authored and
Alice Berard
committed
[AIRFLOW-2797] Create Google Dataproc cluster with custom image (apache#3871)
1 parent 9701f83 commit bb43c71

File tree

2 files changed

+48
-0
lines changed

2 files changed

+48
-0
lines changed

airflow/contrib/operators/dataproc_operator.py

+14
Original file line numberDiff line numberDiff line change
@@ -68,6 +68,9 @@ class DataprocClusterCreateOperator(BaseOperator):
6868
:type metadata: dict
6969
:param image_version: the version of software inside the Dataproc cluster
7070
:type image_version: str
71+
:param custom_image: custom Dataproc image for more info see
72+
https://cloud.google.com/dataproc/docs/guides/dataproc-images
73+
:type: custom_image: str
7174
:param properties: dict of properties to set on
7275
config files (e.g. spark-defaults.conf), see
7376
https://cloud.google.com/dataproc/docs/reference/rest/v1/projects.regions.clusters#SoftwareConfig
@@ -148,6 +151,7 @@ def __init__(self,
148151
init_actions_uris=None,
149152
init_action_timeout="10m",
150153
metadata=None,
154+
custom_image=None,
151155
image_version=None,
152156
properties=None,
153157
master_machine_type='n1-standard-4',
@@ -180,6 +184,7 @@ def __init__(self,
180184
self.init_actions_uris = init_actions_uris
181185
self.init_action_timeout = init_action_timeout
182186
self.metadata = metadata
187+
self.custom_image = custom_image
183188
self.image_version = image_version
184189
self.properties = properties
185190
self.master_machine_type = master_machine_type
@@ -201,6 +206,9 @@ def __init__(self,
201206
self.auto_delete_time = auto_delete_time
202207
self.auto_delete_ttl = auto_delete_ttl
203208

209+
assert not (self.custom_image and self.image_version), \
210+
"custom_image and image_version can't be both set"
211+
204212
def _get_cluster_list_for_project(self, service):
205213
result = service.projects().regions().clusters().list(
206214
projectId=self.project_id,
@@ -338,6 +346,12 @@ def _build_cluster_data(self):
338346
cluster_data['config']['gceClusterConfig']['tags'] = self.tags
339347
if self.image_version:
340348
cluster_data['config']['softwareConfig']['imageVersion'] = self.image_version
349+
elif self.custom_image:
350+
custom_image_url = 'https://www.googleapis.com/compute/beta/projects/' \
351+
'{}/global/images/{}'.format(self.project_id,
352+
self.custom_image)
353+
cluster_data['config']['masterConfig']['imageUri'] = custom_image_url
354+
cluster_data['config']['workerConfig']['imageUri'] = custom_image_url
341355
if self.properties:
342356
cluster_data['config']['softwareConfig']['properties'] = self.properties
343357
if self.idle_delete_ttl:

tests/contrib/operators/test_dataproc_operator.py

+34
Original file line numberDiff line numberDiff line change
@@ -59,6 +59,7 @@
5959
TAGS = ['tag1', 'tag2']
6060
STORAGE_BUCKET = 'gs://airflow-test-bucket/'
6161
IMAGE_VERSION = '1.1'
62+
CUSTOM_IMAGE = 'test-custom-image'
6263
MASTER_MACHINE_TYPE = 'n1-standard-2'
6364
MASTER_DISK_SIZE = 100
6465
MASTER_DISK_TYPE = 'pd-standard'
@@ -264,6 +265,39 @@ def test_build_cluster_data_with_autoDeleteTime_and_autoDeleteTtl(self):
264265
self.assertEqual(cluster_data['config']['lifecycleConfig']['autoDeleteTime'],
265266
"2017-06-07T00:00:00.000000Z")
266267

268+
def test_init_with_image_version_and_custom_image_both_set(self):
269+
with self.assertRaises(AssertionError):
270+
DataprocClusterCreateOperator(
271+
task_id=TASK_ID,
272+
cluster_name=CLUSTER_NAME,
273+
project_id=PROJECT_ID,
274+
num_workers=NUM_WORKERS,
275+
zone=ZONE,
276+
dag=self.dag,
277+
image_version=IMAGE_VERSION,
278+
custom_image=CUSTOM_IMAGE
279+
)
280+
281+
def test_init_with_custom_image(self):
282+
dataproc_operator = DataprocClusterCreateOperator(
283+
task_id=TASK_ID,
284+
cluster_name=CLUSTER_NAME,
285+
project_id=PROJECT_ID,
286+
num_workers=NUM_WORKERS,
287+
zone=ZONE,
288+
dag=self.dag,
289+
custom_image=CUSTOM_IMAGE
290+
)
291+
292+
cluster_data = dataproc_operator._build_cluster_data()
293+
expected_custom_image_url = \
294+
'https://www.googleapis.com/compute/beta/projects/' \
295+
'{}/global/images/{}'.format(PROJECT_ID, CUSTOM_IMAGE)
296+
self.assertEqual(cluster_data['config']['masterConfig']['imageUri'],
297+
expected_custom_image_url)
298+
self.assertEqual(cluster_data['config']['workerConfig']['imageUri'],
299+
expected_custom_image_url)
300+
267301
def test_cluster_name_log_no_sub(self):
268302
with patch('airflow.contrib.operators.dataproc_operator.DataProcHook') \
269303
as mock_hook:

0 commit comments

Comments
 (0)