Fix kubernetespodoperator resources param (GoogleCloudPlatform#2506)

leahecole · web-flow · commit cc1f10a38b4c · 2019-11-01T16:25:35.000-07:00
* call new subpackage format for Airflow

* Fix resources param in kubernetespodoperator, add clarifying comment
diff --git a/composer/workflows/kubernetes_pod_operator.py b/composer/workflows/kubernetes_pod_operator.py
@@ -17,7 +17,6 @@
 # [START composer_kubernetespodoperator]
 import datetime
 from airflow import models
-from airflow.contrib.kubernetes import pod
 from airflow.contrib.kubernetes import secret
 from airflow.contrib.operators import kubernetes_pod_operator
 
@@ -200,7 +199,14 @@
         annotations={'key1': 'value1'},
         # Resource specifications for Pod, this will allow you to set both cpu
         # and memory limits and requirements.
-        resources=pod.Resources().__dict__,
+        # Prior to Airflow 1.10.4, resource specifications were
+        # passed as a Pod Resources Class object,
+        # If using this example on a version of Airflow prior to 1.10.4,
+        # import the "pod" package from airflow.contrib.kubernetes and use
+        # resources = pod.Resources() instead passing a dict
+        # For more info see:
+        # https://github.com/apache/airflow/pull/4551
+        resources={'limit_memory': 1, 'limit_cpu': 1},
         # Specifies path to kubernetes config. If no config is specified will
         # default to '~/.kube/config'. The config_file is templated.
         config_file='/home/airflow/composer_kube_config',
diff --git a/composer/workflows/requirements.txt b/composer/workflows/requirements.txt
@@ -1,4 +1,4 @@
-apache-airflow[gcp_api]==1.10.5
+apache-airflow[gcp]==1.10.5
 kubernetes==10.0.1
 scipy==1.3.1
 numpy==1.17.2
