#278: EKS Node Group Variable and Airflow Worker IAM Role

airflow/dags/cwl_dag.py

@@ -4,12 +4,14 @@
 # Parameter cwl_workflow: the URL of the CWL workflow to execute.
 # Parameter args_as_json: JSON string contained the specific values for the workflow specific inputs.
 import json
+import uuid
 from datetime import datetime
-from airflow import DAG
-from kubernetes.client import models as k8s
-from airflow.providers.cncf.kubernetes.operators.kubernetes_pod import KubernetesPodOperator
+
 from airflow.models.param import Param
-import uuid
+from airflow.providers.cncf.kubernetes.operators.kubernetes_pod import KubernetesPodOperator
+from kubernetes.client import models as k8s
+
+from airflow import DAG
 
 # The Kubernetes Pod that executes the CWL-Docker container
 # Must use elevated privileges to start/stop the Docker engine
@@ -19,35 +21,37 @@
 POD_NAMESPACE = "airflow"
 
 # Example arguments
-default_cwl_workflow = "https://raw.githubusercontent.com/unity-sds/sbg-workflows/main/preprocess/sbg-preprocess-workflow.cwl"
-default_args_as_json_dict = {"input_processing_labels": ["label1", "label2"],
-                             "input_cmr_stac": "https://cmr.earthdata.nasa.gov/search/granules.stac?collection_concept_id=C2408009906-LPCLOUD&temporal[]=2023-08-10T03:41:03.000Z,2023-08-10T03:41:03.000Z",
-                             "input_unity_dapa_client": "40c2s0ulbhp9i0fmaph3su9jch",
-                             "input_unity_dapa_api": "https://d3vc8w9zcq658.cloudfront.net",
-                             "input_crid": "001",
-                             "output_collection_id": "urn:nasa:unity:unity:dev:SBG-L1B_PRE___1",
-                             "output_data_bucket": "sps-dev-ds-storage"}
+default_cwl_workflow = (
+    "https://raw.githubusercontent.com/unity-sds/sbg-workflows/main/preprocess/sbg-preprocess-workflow.cwl"
+)
+default_args_as_json_dict = {
+    "input_processing_labels": ["label1", "label2"],
+    "input_cmr_stac": "https://cmr.earthdata.nasa.gov/search/granules.stac?collection_concept_id=C2408009906-LPCLOUD&temporal[]=2023-08-10T03:41:03.000Z,2023-08-10T03:41:03.000Z",
+    "input_unity_dapa_client": "40c2s0ulbhp9i0fmaph3su9jch",
+    "input_unity_dapa_api": "https://d3vc8w9zcq658.cloudfront.net",
+    "input_crid": "001",
+    "output_collection_id": "urn:nasa:unity:unity:dev:SBG-L1B_PRE___1",
+    "output_data_bucket": "sps-dev-ds-storage",
+}
 
 # Default DAG configuration
-dag_default_args = {
-    'owner': 'airflow',
-    'depends_on_past': False,
-    'start_date': datetime(2024, 1, 1, 0, 0)
-}
+dag_default_args = {"owner": "unity-sps", "depends_on_past": False, "start_date": datetime(2024, 1, 1, 0, 0)}
 
 # The DAG
-dag = DAG(dag_id='cwl-dag',
-          description='DAG to execute a generic CWL workflow',
-          tags=['cwl', 'unity-sps', "docker"],
-          is_paused_upon_creation=True,
-          catchup=False,
-          schedule_interval=None,
-          max_active_runs=1,
-          default_args=dag_default_args,
-          params={
-              "cwl_workflow": Param(default_cwl_workflow, type="string"),
-              "args_as_json": Param(json.dumps(default_args_as_json_dict), type="string"),
-          })
+dag = DAG(
+    dag_id="cwl-dag",
+    description="DAG to execute a generic CWL workflow",
+    tags=["cwl", "unity-sps", "docker"],
+    is_paused_upon_creation=True,
+    catchup=False,
+    schedule_interval=None,
+    max_active_runs=1,
+    default_args=dag_default_args,
+    params={
+        "cwl_workflow": Param(default_cwl_workflow, type="string"),
+        "args_as_json": Param(json.dumps(default_args_as_json_dict), type="string"),
+    },
+)
 
 # Environment variables
 default_env_vars = {}
@@ -62,9 +66,9 @@
     get_logs=True,
     task_id="docker-cwl-task",
     full_pod_spec=k8s.V1Pod(
-        metadata=k8s.V1ObjectMeta(
-            name='docker-cwl-pod-' + uuid.uuid4().hex),
+        metadata=k8s.V1ObjectMeta(name="docker-cwl-pod-" + uuid.uuid4().hex),
     ),
     pod_template_file=POD_TEMPLATE_FILE,
     arguments=["{{ params.cwl_workflow }}", "{{ params.args_as_json }}"],
-    dag=dag)
+    dag=dag,
+)

airflow/dags/sbg_preprocess_cwl_dag.py

@@ -42,10 +42,10 @@
     params={
         "cwl_workflow": Param(CWL_WORKFLOW, type="string"),
         "input_cmr_stac": Param(CMR_STAC, type="string"),
-        #"input_processing_labels": Param(["label1", "label2"], type="string[]"),
-        #"input_cmr_collection_name": Param("C2408009906-LPCLOUD", type="string"),
-        #"input_cmr_search_start_time": Param("2024-01-03T13:19:36.000Z", type="string"),
-        #"input_cmr_search_stop_time": Param("2024-01-03T13:19:36.000Z", type="string"),
+        # "input_processing_labels": Param(["label1", "label2"], type="string[]"),
+        # "input_cmr_collection_name": Param("C2408009906-LPCLOUD", type="string"),
+        # "input_cmr_search_start_time": Param("2024-01-03T13:19:36.000Z", type="string"),
+        # "input_cmr_search_stop_time": Param("2024-01-03T13:19:36.000Z", type="string"),
         "input_unity_dapa_api": Param("https://d3vc8w9zcq658.cloudfront.net", type="string"),
         "input_unity_dapa_client": Param("40c2s0ulbhp9i0fmaph3su9jch", type="string"),
         "input_crid": Param("001", type="string"),
@@ -54,6 +54,7 @@
     },
 )
 
+
 # Task that serializes the job arguments into a JSON string
 def setup(ti=None, **context):
     task_dict = {
@@ -90,4 +91,4 @@ def setup(ti=None, **context):
     dag=dag,
 )
 
-setup_task >> cwl_task
+setup_task >> cwl_task

airflow/dags/sbg_preprocess_no_cwl.py

@@ -1,88 +1,78 @@
 # DAG for SBG Workflow #1
 # See https://github.com/unity-sds/sbg-workflows/blob/main/preprocess/sbg-preprocess-workflow.cwl
-import os
+import json
+import uuid
 from datetime import datetime
-from airflow import DAG
-from kubernetes.client import models as k8s
-from airflow.providers.cncf.kubernetes.operators.kubernetes_pod import KubernetesPodOperator
-from airflow.operators.python import PythonOperator
+
 from airflow.models.param import Param
-import uuid
-import json
+from airflow.providers.cncf.kubernetes.operators.kubernetes_pod import KubernetesPodOperator
+from kubernetes.client import models as k8s
+
+from airflow import DAG
 
 # Fixed Parameters
 UNITY_DS_IMAGE = "ghcr.io/unity-sds/unity-data-services:6.4.3"
 SBG_PREPROCESS_IMAGE = "gangl/sbg-unity-preprocess:266e40d8"
-COGNITO_URL = 'https://cognito-idp.us-west-2.amazonaws.com'
-UNITY_USERNAME = ''
-UNITY_PASSWORD = ''
-UNITY_PASSWORD_TYPE = ''
-DOWNLOAD_DIR = '/scratch/granules'
-DOWNLOADING_KEYS = 'data, data1'
-GRANULES_DOWNLOAD_TYPE = 'DAAC'
+COGNITO_URL = "https://cognito-idp.us-west-2.amazonaws.com"
+UNITY_USERNAME = ""
+UNITY_PASSWORD = ""
+UNITY_PASSWORD_TYPE = ""
+DOWNLOAD_DIR = "/scratch/granules"
+DOWNLOADING_KEYS = "data, data1"
+GRANULES_DOWNLOAD_TYPE = "DAAC"
 
-EDL_USERNAME = '/sps/processing/workflows/edl_username'
-EDL_PASSWORD = '/sps/processing/workflows/edl_password'
-EDL_PASSWORD_TYPE = 'PARAM_STORE'
-EDL_BASE_URL = 'https://urs.earthdata.nasa.gov/'
-LOG_LEVEL = '20'
+EDL_USERNAME = "/sps/processing/workflows/edl_username"
+EDL_PASSWORD = "/sps/processing/workflows/edl_password"
+EDL_PASSWORD_TYPE = "PARAM_STORE"
+EDL_BASE_URL = "https://urs.earthdata.nasa.gov/"
+LOG_LEVEL = "20"
 
 
 # This path must be inside the shared Persistent Volume
 STAC_JSON_PATH = "/scratch/search_results.json"
 STAGE_IN_RESULTS = "/scratch/granules/stage-in-results.json"
 
 # Venue dependent parameters
-CLIENT_ID = '40c2s0ulbhp9i0fmaph3su9jch'
-DAPA_API = 'https://d3vc8w9zcq658.cloudfront.net'
-STAGING_BUCKET = 'sps-dev-ds-storage'
+CLIENT_ID = "40c2s0ulbhp9i0fmaph3su9jch"
+DAPA_API = "https://d3vc8w9zcq658.cloudfront.net"
+STAGING_BUCKET = "sps-dev-ds-storage"
 
 
-#CWL_URL = "http://awslbdockstorestack-lb-1429770210.us-west-2.elb.amazonaws.com:9998/api/ga4gh/trs/v2/tools/%23workflow%2Fdockstore.org%2Fmike-gangl%2FSBG-unity-preprocess/versions/16/PLAIN-CWL/descriptor/%2Fprocess.cwl"
+# CWL_URL = "http://awslbdockstorestack-lb-1429770210.us-west-2.elb.amazonaws.com:9998/api/ga4gh/trs/v2/tools/%23workflow%2Fdockstore.org%2Fmike-gangl%2FSBG-unity-preprocess/versions/16/PLAIN-CWL/descriptor/%2Fprocess.cwl"
 CWL_URL = "https://raw.githubusercontent.com/unity-sds/unity-sps-workflows/sbg/sbg/process.cwl"
 # YAML_FILE = "/scratch/process.yaml"
-ARGS = {
-  "download_dir": {
-    "class": "Directory",
-    "path": "/scratch/granules"
-  }
-}
+ARGS = {"download_dir": {"class": "Directory", "path": "/scratch/granules"}}
 
 # Default DAG configuration
-dag_default_args = {
-    'owner': 'airflow',
-    'depends_on_past': False,
-    'start_date': datetime(2024, 1, 1, 0, 0)
-}
+dag_default_args = {"owner": "unity-sps", "depends_on_past": False, "start_date": datetime(2024, 1, 1, 0, 0)}
 
 volume = k8s.V1Volume(
-    name='unity-sps-airflow-pv',
-    persistent_volume_claim=k8s.V1PersistentVolumeClaimVolumeSource(claim_name='unity-sps-airflow-pvc')
+    name="unity-sps-airflow-pv",
+    persistent_volume_claim=k8s.V1PersistentVolumeClaimVolumeSource(claim_name="unity-sps-airflow-pvc"),
 )
 
 volume_mount = k8s.V1VolumeMount(
-    name='unity-sps-airflow-pv',
-    mount_path='/scratch',
-    sub_path=None,
-    read_only=False
+    name="unity-sps-airflow-pv", mount_path="/scratch", sub_path=None, read_only=False
 )
 
-dag = DAG(dag_id='sbg-preprocess-no-cwl',
-          description='SBG Preprocess Workflow',
-          tags=["SBG", "Unity", "SPS", "NASA", "JPL"],
-          is_paused_upon_creation=True,
-          catchup=False,
-          schedule=None,
-          max_active_runs=1,
-          default_args=dag_default_args,
-          params={
-              "input_cmr_collection_name": Param("C2408009906-LPCLOUD", type="string"),
-              "input_cmr_search_start_time": Param("2023-08-10T03:41:03.000Z", type="string"),
-              "input_cmr_search_stop_time": Param("2023-08-10T03:41:03.000Z", type="string"),
-              "input_crid": Param("001", type="string"),
-              "output_collection_id": Param("urn:nasa:unity:unity:dev:SBG-L1B_PRE___1", type="string"),
-              "output_data_bucket": Param("sps-dev-ds-storage", type="string")
-          }, )
+dag = DAG(
+    dag_id="sbg-preprocess-no-cwl",
+    description="SBG Preprocess Workflow",
+    tags=["SBG", "Unity", "SPS", "NASA", "JPL"],
+    is_paused_upon_creation=True,
+    catchup=False,
+    schedule=None,
+    max_active_runs=1,
+    default_args=dag_default_args,
+    params={
+        "input_cmr_collection_name": Param("C2408009906-LPCLOUD", type="string"),
+        "input_cmr_search_start_time": Param("2023-08-10T03:41:03.000Z", type="string"),
+        "input_cmr_search_stop_time": Param("2023-08-10T03:41:03.000Z", type="string"),
+        "input_crid": Param("001", type="string"),
+        "output_collection_id": Param("urn:nasa:unity:unity:dev:SBG-L1B_PRE___1", type="string"),
+        "output_data_bucket": Param("sps-dev-ds-storage", type="string"),
+    },
+)
 
 cmr_query_env_vars = [
     # k8s.V1EnvVar(name="AWS_ACCESS_KEY_ID", value=""),
@@ -120,18 +110,18 @@
     startup_timeout_seconds=1000,
     get_logs=True,
     task_id="CMR_Query",
-    full_pod_spec=k8s.V1Pod(
-        k8s.V1ObjectMeta(name=('cmr-query-pod-' + uuid.uuid4().hex))),
+    full_pod_spec=k8s.V1Pod(k8s.V1ObjectMeta(name=("cmr-query-pod-" + uuid.uuid4().hex))),
     do_xcom_push=True,
     volumes=[volume],
     volume_mounts=[volume_mount],
-    dag=dag)
+    dag=dag,
+)
 
 stage_in_env_vars = [
-    #k8s.V1EnvVar(name="AWS_ACCESS_KEY_ID", value=AWS_ACCESS_KEY_ID),
-    #k8s.V1EnvVar(name="AWS_SECRET_ACCESS_KEY", value=AWS_SECRET_ACCESS_KEY),
-    #k8s.V1EnvVar(name="AWS_SESSION_TOKEN", value=AWS_SESSION_TOKEN),
-    #k8s.V1EnvVar(name="AWS_REGION", value=AWS_REGION),
+    # k8s.V1EnvVar(name="AWS_ACCESS_KEY_ID", value=AWS_ACCESS_KEY_ID),
+    # k8s.V1EnvVar(name="AWS_SECRET_ACCESS_KEY", value=AWS_SECRET_ACCESS_KEY),
+    # k8s.V1EnvVar(name="AWS_SESSION_TOKEN", value=AWS_SESSION_TOKEN),
+    # k8s.V1EnvVar(name="AWS_REGION", value=AWS_REGION),
     k8s.V1EnvVar(name="PYTHONUNBUFFERED", value="1"),
     # k8s.V1EnvVar(name="USERNAME", value=UNITY_USERNAME),
     # k8s.V1EnvVar(name="PASSWORD", value=UNITY_PASSWORD),
@@ -166,12 +156,12 @@
     startup_timeout_seconds=1000,
     get_logs=True,
     task_id="Stage_In",
-    full_pod_spec=k8s.V1Pod(
-        k8s.V1ObjectMeta(name=('stage-in-pod-' + uuid.uuid4().hex))),
+    full_pod_spec=k8s.V1Pod(k8s.V1ObjectMeta(name=("stage-in-pod-" + uuid.uuid4().hex))),
     # do_xcom_push=True,
     volumes=[volume],
     volume_mounts=[volume_mount],
-    dag=dag)
+    dag=dag,
+)
 
 
 # ref:http://awslbdockstorestack-lb-1429770210.us-west-2.elb.amazonaws.com:9998/api/ga4gh/trs/v2/tools/%23workflow%2Fdockstore.org%2Fmike-gangl%2FSBG-unity-preprocess/versions/16/PLAIN-CWL/descriptor/%2Fprocess.cwl
@@ -191,12 +181,12 @@
     startup_timeout_seconds=1000,
     get_logs=True,
     task_id="Process",
-    full_pod_spec=k8s.V1Pod(
-        k8s.V1ObjectMeta(name=('process-pod-' + uuid.uuid4().hex))),
+    full_pod_spec=k8s.V1Pod(k8s.V1ObjectMeta(name=("process-pod-" + uuid.uuid4().hex))),
     # do_xcom_push=True,
     volumes=[volume],
     volume_mounts=[volume_mount],
-    dag=dag)
+    dag=dag,
+)
 
 stage_out_env_vars = [
     # k8s.V1EnvVar(name="AWS_ACCESS_KEY_ID", value=AWS_ACCESS_KEY_ID),
@@ -232,15 +222,15 @@
     startup_timeout_seconds=1000,
     get_logs=True,
     task_id="Stage_Out",
-    full_pod_spec=k8s.V1Pod(
-        k8s.V1ObjectMeta(name=('stage-out-pod-' + uuid.uuid4().hex))),
+    full_pod_spec=k8s.V1Pod(k8s.V1ObjectMeta(name=("stage-out-pod-" + uuid.uuid4().hex))),
     # do_xcom_push=True,
     volumes=[volume],
     volume_mounts=[volume_mount],
-    dag=dag)
+    dag=dag,
+)
 
 
-'''
+"""
 def preprocess(ti=None, **context):
     # cmr_query = ti.xcom_pull(task_ids=['CMR_Query'])[0]
     #print(cmr_query)
@@ -250,6 +240,6 @@ def preprocess(ti=None, **context):
 preprocess_task = PythonOperator(task_id="Preprocess",
                                  python_callable=preprocess,
                                  dag=dag)
-'''
+"""
 
 cmr_query_task >> stage_in_task >> process_task >> stage_out_task

terraform-unity/modules/terraform-eks-cluster/README.md

@@ -19,7 +19,7 @@
 
 | Name | Source | Version |
 |------|--------|---------|
-| <a name="module_unity-eks"></a> [unity-eks](#module\_unity-eks) | git@github.com:unity-sds/unity-cs-infra.git//terraform-unity-eks_module | 0.1.3 |
+| <a name="module_unity-eks"></a> [unity-eks](#module\_unity-eks) | git@github.com:unity-sds/unity-cs-infra.git//terraform-unity-eks_module | main |
 
 ## Resources
 
@@ -35,6 +35,7 @@
 | Name | Description | Type | Default | Required |
 |------|-------------|------|---------|:--------:|
 | <a name="input_cluster_name"></a> [cluster\_name](#input\_cluster\_name) | n/a | `string` | n/a | yes |
+| <a name="input_nodegroups"></a> [nodegroups](#input\_nodegroups) | A map of node group configurations | <pre>map(object({<br>    create_iam_role            = optional(bool)<br>    iam_role_arn               = optional(string)<br>    ami_id                     = optional(string)<br>    min_size                   = optional(number)<br>    max_size                   = optional(number)<br>    desired_size               = optional(number)<br>    instance_types             = optional(list(string))<br>    capacity_type              = optional(string)<br>    enable_bootstrap_user_data = optional(bool)<br>    metadata_options           = optional(map(any))<br>  }))</pre> | <pre>{<br>  "defaultGroup": {<br>    "desired_size": 1,<br>    "instance_types": [<br>      "m5.xlarge"<br>    ],<br>    "max_size": 1,<br>    "metadata_options": {<br>      "http_endpoint": "enabled",<br>      "http_put_response_hop_limit": 3<br>    },<br>    "min_size": 1<br>  }<br>}</pre> | no |
 
 ## Outputs
 

terraform-unity/modules/terraform-eks-cluster/main.tf

@@ -1,15 +1,9 @@
 module "unity-eks" {
-  source          = "git@github.com:unity-sds/unity-cs-infra.git//terraform-unity-eks_module?ref=0.1.3"
+  source          = "git@github.com:unity-sds/unity-cs-infra.git//terraform-unity-eks_module?ref=main"
   deployment_name = var.cluster_name
 
-  nodegroups = {
-    defaultGroup = {
-      instance_types = ["m5.xlarge"]
-      min_size       = 1
-      max_size       = 1
-      desired_size   = 1
-    }
-  }
+  nodegroups = var.nodegroups
+
   aws_auth_roles = [{
     rolearn  = "arn:aws:iam::${data.aws_caller_identity.current.account_id}:role/mcp-tenantOperator"
     username = "admin"