added appgen dag

Author: grallewellyn

airflow/dags/appgen_dag.py

@@ -6,6 +6,7 @@
 import logging
 import os
 from datetime import datetime
+import boto3
 
 from airflow.models.baseoperator import chain
 from airflow.models.param import Param
@@ -32,6 +33,11 @@
     }
 )
 
+# AWS SSM parameter paths for credentials 
+DOCKERHUB_USERNAME = "/unity/ads/app_gen/development/dockerhub_username"
+DOCKERHUB_TOKEN = "/unity/ads/app_gen/development/dockerhub_api_key"
+DOCKSTORE_TOKEN = "/unity/ads/app_gen/development/dockstore_token"
+
 # >>> This part will be removed once the parameters can be imported from unity_sps_plugins.py
 DEFAULT_LOG_LEVEL = 20
 EC2_TYPES = {
@@ -116,7 +122,7 @@ def build_ec2_type_label(key):
 LOG_LEVEL_TYPE = {10: "DEBUG", 20: "INFO"}
 
 # Change this to the Docker image that contains the Application Package Generator
-DOCKER_IMAGE = "docker.io/busybox"
+DOCKER_IMAGE = "jplmdps/unity-app-gen:v1.0.0"
 
 # Default DAG configuration
 dag_default_args = {
@@ -137,7 +143,7 @@ def build_ec2_type_label(key):
     max_active_tasks=30,
     default_args=dag_default_args,
     params={
-        "message": Param("Hello World", type="string", title="Message", description="The greeting message"),
+        "repository": Param("https://github.com/unity-sds/unity-example-application", type="string", title="Repository", description="Repository to build from"),
         "log_level": Param(
             DEFAULT_LOG_LEVEL,
             type="integer",
@@ -160,11 +166,45 @@ def build_ec2_type_label(key):
     },
 )
 
+app_gen_env_vars = [
+    k8s.V1EnvVar(name="DOCKERHUB_USERNAME", value="{{ ti.xcom_pull(task_ids='Setup', key='dockerhub_username') }}"),
+    k8s.V1EnvVar(name="DOCKERHUB_TOKEN", value= "{{ ti.xcom_pull(task_ids='Setup', key='dockerhub_token') }}"),
+    k8s.V1EnvVar(name="DOCKSTORE_TOKEN", value="{{ ti.xcom_pull(task_ids='Setup', key='dockstore_token') }}"),
+    k8s.V1EnvVar(name="GITHUB_REPO", value="{{ params.repository }}"),
+]
 
 def setup(ti=None, **context):
     """
     Task that selects the proper Karpenter Node Pool depending on the user requested resources.
     """
+
+    ## Retrieve the docker credentials and DockStore token
+    ssm_client = boto3.client("ssm", region_name="us-west-2")
+    ssm_response = ssm_client.get_parameters(
+        Names=[DOCKERHUB_USERNAME, DOCKERHUB_TOKEN, DOCKSTORE_TOKEN], WithDecryption=True
+    )
+    logging.info(ssm_response)
+
+    # Somehow get the correct variables from SSM here 
+    credentials_dict = {}
+    for param in ssm_response["Parameters"]:
+        if param["Name"] == DOCKERHUB_USERNAME:
+            credentials_dict["dockerhub_username"] = param["Value"]
+        elif param["Name"] == DOCKERHUB_TOKEN:
+            credentials_dict["dockerhub_token"] = param["Value"]
+        elif param["Name"] == DOCKSTORE_TOKEN:
+            credentials_dict["dockstore_token"] = param["Value"]
+
+    required_credentials = ["dockerhub_username", "dockerhub_token", "dockstore_token"]
+    # make sure all required credentials are provided
+    if (not set(required_credentials).issubset(list(credentials_dict.keys()))):
+        logging.error(f"Expected all of credentials to run mdps app generator {required_credentials}")
+    
+    # use xcom to push to avoid putting credentials to the logs 
+    ti.xcom_push(key="dockerhub_username", value=credentials_dict["dockerhub_username"])
+    ti.xcom_push(key="dockerhub_token", value=credentials_dict["dockerhub_token"])
+    ti.xcom_push(key="dockstore_token", value=credentials_dict["dockstore_token"])
+
     context = get_current_context()
     logging.info(f"DAG Run parameters: {json.dumps(context['params'], sort_keys=True, indent=4)}")
 
@@ -204,13 +244,21 @@ def setup(ti=None, **context):
     retries=1,
     task_id="appgen_task",
     namespace=POD_NAMESPACE,
+    env_vars=app_gen_env_vars,
     name="appgen-task-pod",
     image=DOCKER_IMAGE,
     service_account_name="airflow-worker",
     in_cluster=True,
     get_logs=True,
     startup_timeout_seconds=600,
-    arguments=["echo", "{{ti.xcom_pull(task_ids='Setup', key='message')}}"],
+    arguments=[
+        "-r",
+        "{{ params.repository }}",
+        "-l",
+        "{{ params.log_level }}",
+        "-e",
+        "{{ ti.xcom_pull(task_ids='Setup', key='ecr_login') }}",
+    ],
     container_security_context={"privileged": True},
     container_resources=k8s.V1ResourceRequirements(
         requests={