[
https://issues.apache.org/jira/browse/AIRFLOW-2797?page=com.atlassian.jira.plugin.system.issuetabpanels:comment-tabpanel&focusedCommentId=16611752#comment-16611752
]
ASF GitHub Bot commented on AIRFLOW-2797:
-----------------------------------------
Fokko closed pull request #3871: [AIRFLOW-2797] Create Google Dataproc cluster
with custom image
URL: https://github.com/apache/incubator-airflow/pull/3871
This is a PR merged from a forked repository.
As GitHub hides the original diff on merge, it is displayed below for
the sake of provenance:
As this is a foreign pull request (from a fork), the diff is supplied
below (as it won't show otherwise due to GitHub magic):
diff --git a/airflow/contrib/operators/dataproc_operator.py
b/airflow/contrib/operators/dataproc_operator.py
index ad74c36b37..61f3895bc0 100644
--- a/airflow/contrib/operators/dataproc_operator.py
+++ b/airflow/contrib/operators/dataproc_operator.py
@@ -68,6 +68,9 @@ class DataprocClusterCreateOperator(BaseOperator):
:type metadata: dict
:param image_version: the version of software inside the Dataproc cluster
:type image_version: str
+ :param custom_image: custom Dataproc image for more info see
+ https://cloud.google.com/dataproc/docs/guides/dataproc-images
+ :type: custom_image: str
:param properties: dict of properties to set on
config files (e.g. spark-defaults.conf), see
https://cloud.google.com/dataproc/docs/reference/rest/v1/projects.regions.clusters#SoftwareConfig
@@ -148,6 +151,7 @@ def __init__(self,
init_actions_uris=None,
init_action_timeout="10m",
metadata=None,
+ custom_image=None,
image_version=None,
properties=None,
master_machine_type='n1-standard-4',
@@ -180,6 +184,7 @@ def __init__(self,
self.init_actions_uris = init_actions_uris
self.init_action_timeout = init_action_timeout
self.metadata = metadata
+ self.custom_image = custom_image
self.image_version = image_version
self.properties = properties
self.master_machine_type = master_machine_type
@@ -201,6 +206,9 @@ def __init__(self,
self.auto_delete_time = auto_delete_time
self.auto_delete_ttl = auto_delete_ttl
+ assert not (self.custom_image and self.image_version), \
+ "custom_image and image_version can't be both set"
+
def _get_cluster_list_for_project(self, service):
result = service.projects().regions().clusters().list(
projectId=self.project_id,
@@ -338,6 +346,12 @@ def _build_cluster_data(self):
cluster_data['config']['gceClusterConfig']['tags'] = self.tags
if self.image_version:
cluster_data['config']['softwareConfig']['imageVersion'] =
self.image_version
+ elif self.custom_image:
+ custom_image_url =
'https://www.googleapis.com/compute/beta/projects/' \
+ '{}/global/images/{}'.format(self.project_id,
+ self.custom_image)
+ cluster_data['config']['masterConfig']['imageUri'] =
custom_image_url
+ cluster_data['config']['workerConfig']['imageUri'] =
custom_image_url
if self.properties:
cluster_data['config']['softwareConfig']['properties'] =
self.properties
if self.idle_delete_ttl:
diff --git a/tests/contrib/operators/test_dataproc_operator.py
b/tests/contrib/operators/test_dataproc_operator.py
index 5b403a86ba..4095ca2068 100644
--- a/tests/contrib/operators/test_dataproc_operator.py
+++ b/tests/contrib/operators/test_dataproc_operator.py
@@ -59,6 +59,7 @@
TAGS = ['tag1', 'tag2']
STORAGE_BUCKET = 'gs://airflow-test-bucket/'
IMAGE_VERSION = '1.1'
+CUSTOM_IMAGE = 'test-custom-image'
MASTER_MACHINE_TYPE = 'n1-standard-2'
MASTER_DISK_SIZE = 100
MASTER_DISK_TYPE = 'pd-standard'
@@ -264,6 +265,39 @@ def
test_build_cluster_data_with_autoDeleteTime_and_autoDeleteTtl(self):
self.assertEqual(cluster_data['config']['lifecycleConfig']['autoDeleteTime'],
"2017-06-07T00:00:00.000000Z")
+ def test_init_with_image_version_and_custom_image_both_set(self):
+ with self.assertRaises(AssertionError):
+ DataprocClusterCreateOperator(
+ task_id=TASK_ID,
+ cluster_name=CLUSTER_NAME,
+ project_id=PROJECT_ID,
+ num_workers=NUM_WORKERS,
+ zone=ZONE,
+ dag=self.dag,
+ image_version=IMAGE_VERSION,
+ custom_image=CUSTOM_IMAGE
+ )
+
+ def test_init_with_custom_image(self):
+ dataproc_operator = DataprocClusterCreateOperator(
+ task_id=TASK_ID,
+ cluster_name=CLUSTER_NAME,
+ project_id=PROJECT_ID,
+ num_workers=NUM_WORKERS,
+ zone=ZONE,
+ dag=self.dag,
+ custom_image=CUSTOM_IMAGE
+ )
+
+ cluster_data = dataproc_operator._build_cluster_data()
+ expected_custom_image_url = \
+ 'https://www.googleapis.com/compute/beta/projects/' \
+ '{}/global/images/{}'.format(PROJECT_ID, CUSTOM_IMAGE)
+ self.assertEqual(cluster_data['config']['masterConfig']['imageUri'],
+ expected_custom_image_url)
+ self.assertEqual(cluster_data['config']['workerConfig']['imageUri'],
+ expected_custom_image_url)
+
def test_cluster_name_log_no_sub(self):
with patch('airflow.contrib.operators.dataproc_operator.DataProcHook')
\
as mock_hook:
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
[email protected]
> Add ability to create Google Dataproc cluster with custom image
> ---------------------------------------------------------------
>
> Key: AIRFLOW-2797
> URL: https://issues.apache.org/jira/browse/AIRFLOW-2797
> Project: Apache Airflow
> Issue Type: Improvement
> Components: contrib, gcp, operators
> Affects Versions: 2.0.0
> Reporter: Jarosław Śmietanka
> Assignee: Jarosław Śmietanka
> Priority: Minor
>
> In GCP, it is possible to create Dataproc cluster with a [custom
> image|https://cloud.google.com/dataproc/docs/guides/dataproc-images] that
> includes user's pre-installed packages. It significantly reduces the startup
> time of the cluster.
>
> Since I already have a code which does that, I volunteer to bring it to the
> community :)
> This improvement won't change many components and should not require
> groundbreaking changes to DataprocClusterCreateOperator.
--
This message was sent by Atlassian JIRA
(v7.6.3#76005)
